1. The Problem​

Picture asking Qwen3-0.6B something like "Under the Office of the Prime Minister's regulations, when is procurement by the specific-method route permitted?" The model will answer confidently, and it will be wrong, because it has never seen enough Thai government documents.

Plenty of people try to fix this with fine-tuning on a few thousand question-answer pairs, and find it doesn't work. The reason is that SFT teaches the shape of an answer, not the knowledge behind it. If the knowledge isn't in the weights, teaching the model to answer in the right tone of voice just makes it confident while it lies.

New knowledge reaches a model through three routes, and picking the wrong one is why most LLM projects fail:

This article is about the second route.

2. What We're Going to Do​

We take a base model (one that hasn't been through instruction tuning) and keep training it with the exact same objective used during pretraining — next-token prediction — on raw Thai text from the domain we care about. No labels, no question-answer pairs, just plain text.

But the heart of this article isn't "train it and it gets better." It's what you give up in exchange:

The phenomenon of a model forgetting what it used to be able to do is called catastrophic forgetting. We won't just gesture at it — we'll measure it as a number and then find a point we can live with.

3. The Equations​

3.1 The CPT objective​

• $x_t$ = the token at position $t$
• $x_{<t}$ = every token before it
• $p_\theta$ = the probability the model predicts

This equation is identical to the one used during pretraining. The only thing that changes is the data. That's why CPT needs no labels — the text is its own answer key.

3.2 Perplexity: our unit of measurement​

In plain language: "on average, how many options is the model torn between?" PPL = 20 means it's wavering among roughly 20 candidates; PPL = 5 means it's far more certain. Lower is better.

3.3 The most important equation in this chapter — replay mixing​

$\lambda$ is the share of domain data in each batch.

• $\lambda = 1.0$ → pure domain data → fastest domain gains and the fastest forgetting
• $\lambda = 0.5$ → an even mix → slower, but far less forgetting

Don't hardcode this value. Sweep it, then pick the point you can accept.

4. Seeing the Equations​

What perplexity actually tells us​

The right-hand plot shows what people usually get wrong: if someone tells you they "cut perplexity by 5 points" without saying where they started, the sentence is nearly meaningless. 80 → 75 is a 6% improvement; 10 → 5 is 50%.

The trade-off the replay ratio controls​

5. Setting Up the Environment​

Open Colab and pick Runtime → Change runtime type → T4 GPU (the free tier is enough).

torch_dtype=torch.float16      # not bfloat16
attn_implementation="sdpa"     # not flash_attention_2
fp16=True                      # in TrainingArguments (not bf16=True)

The first cell of every notebook in this series prints this line so you can see it with your own eyes:

cap = torch.cuda.get_device_capability(0)
print("compute capability:", cap)                    # T4 = (7, 5)
print("native bf16:", cap[0] >= 8)                   # T4 -> False
print("torch says   :", torch.cuda.is_bf16_supported())  # T4 -> True (counts emulation!)

How much VRAM you have, and where it goes​

Notice that the optimizer state takes up more room than the model itself — Adam keeps one copy each of $m$ and $v$ for every parameter. Switching to adamw_bnb_8bit saves 3.3 GB, which means a lot more headroom for batch size or sequence length.

Play with the VRAM budget yourself — change the numbers and see where it OOMs:

6. Preparing the Data​

We use pythainlp/thaigov-v2-corpus-22032023 — a corpus of Thai government news and official documents (public domain). It stands in for "specialized knowledge the model has never seen enough of."

We also use a second set of general Thai text as replay data to hold back the forgetting.

from datasets import load_dataset

domain = load_dataset("pythainlp/thaigov-v2-corpus-22032023", split="train")
domain = domain.shuffle(seed=42).select(range(8000))

Packing: don't let padding eat your budget​

If you pad every document out to the same length, you burn an enormous amount of compute on <pad>. The right approach is to concatenate every document and slice the stream into equal 512-token blocks.

def pack(examples, block_size=512):
    ids = []
    for text in examples["context"]:
        ids.extend(tokenizer(text + tokenizer.eos_token).input_ids)
    n = (len(ids) // block_size) * block_size
    return {"input_ids": [ids[i:i+block_size] for i in range(0, n, block_size)]}

7. The Main Code​

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B-Base",           # base, not instruct — CPT must start from a base model
    torch_dtype=torch.float16,        # the T4 has no bf16
    attn_implementation="sdpa",       # the T4 has no FlashAttention-2
).cuda()

model = model.float()                 # must cast to fp32 before training — see the box below

args = TrainingArguments(
    output_dir="cpt-out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,    # effective batch = 16
    num_train_epochs=1,
    learning_rate=2e-5,               # 10x lower than SFT — see the warning below
    lr_scheduler_type="cosine",
    warmup_steps=50,
    optim="adamw_bnb_8bit",           # saves 3.3 GB
    gradient_checkpointing=True,
    max_grad_norm=1.0,                # keeps fp16 from blowing up
    fp16=True,                        # not bf16
    logging_steps=10,
)

ValueError: Attempting to unscale FP16 gradients.

8. Results​

The notebook measures three things before and after training and writes them to results.json:

1. Domain held-out PPL — should drop clearly (this is what we're paying for)
2. General held-out PPL — should rise somewhat (this is the price)
3. TH-KNOW accuracy from the KobEval-TH benchmark, with a Wilson 95% CI

9. Comparison​

The notebook trains three variants on the same data so you can see the trade-off in numbers:

Reading this table correctly is the point of the chapter:

• λ = 1.0 wins on domain PPL (4.07, the lowest) but general PPL gets worse, 5.88 → 6.72. That is catastrophic forgetting, measured rather than asserted.
• λ = 0.5 gives up a little domain accuracy (4.29) and general perplexity improves to 4.98. Replay does not merely prevent forgetting here; it leaves the model better at Thai overall.

10. Summary​

• CPT puts knowledge into the weights using the same objective as pretraining, with no labels required
• It is always a trade — domain accuracy is bought with general capability you lose
• The replay ratio $\lambda$ is the dial that sets the exchange rate — sweep it, don't guess
• A low learning rate is the line between doing CPT and destroying the model
• Every number needs a confidence interval

Next chapter: SFT and LoRA — once the model has the knowledge, how do we teach it to answer, and why does training just 1.7% of the parameters get you almost as far as training all of them?

References​

1. Gururangan et al. (2020). Don't Stop Pretraining: Adapt Language Models to Domains and Tasks — the domain-adaptive pretraining recipe this chapter follows
2. Ibrahim et al. (2024). Simple and Scalable Strategies to Continually Pre-train Large Language Models — the replay and LR strategies that keep CPT from destroying the model
3. Gupta et al. (2023). Continual Pre-Training of Large Language Models: How to (re)warm your model? — why LR warmup matters so much when continuing pretraining
4. Luo et al. (2023). An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning — a systematic measurement of catastrophic forgetting
5. Kaplan et al. (2020). Scaling Laws for Neural Language Models — scaling laws -- the basis for calling 8,000 documents far too little
6. Hoffmann et al. (2022). Training Compute-Optimal Large Language Models — Chinchilla: the compute-optimal data-to-parameter ratio
7. Yuenyong et al. (2025). OpenThaiGPT 1.6 and R1: Thai-Centric Open Source and Reasoning Large Language Models — Thai CPT at real scale, for contrast with this chapter's toy run
8. Lowphansirikul et al. (2021). WangchanBERTa: Pretraining transformer-based Thai Language Models — the pioneering Thai LM and its corpus preparation

The writing, code and notebooks in this series are licensed under CC BY-NC-SA 4.0 — reuse and adapt them freely with attribution, for non-commercial purposes, and share your adaptations under the same terms. Third-party models and datasets referenced here keep their own licences.

1. The Problem​

Take a true base model like Qwen3-0.6B-Base from the last chapter and type in (in Thai) "Could you recommend some Thai dishes, please?" What comes back is usually not an answer but a continuation — it may compose three more questions, keep going as a travel article, or switch into English partway through, because the only thing it was ever trained on is "on the internet, what usually follows text like this?"

The ability to answer — take an instruction, respond to the point, then stop — does not come with pretraining. It comes from SFT: continued training on (instruction, good answer) pairs, thousands to millions of them. Every instruct model you have ever used has been through this stage.

But the moment you try it yourself, you hit two problems stacked on top of each other:

Layer one — the cost of full fine-tuning. Train every parameter and you get an entire new model (~1.2 GB per task for a 0.6B model). An organization with ten tasks — document summaries, letter drafting, customer replies, complaint triage — has to store ten copies. And moving every single weight with a high learning rate is a recipe for erasing the knowledge you just added in chapter 1 (remember the learning-rate box?).

Layer two — Thai. Even the post-trained Qwen3-0.6B shows the symptom we will see all series long: ask in Thai, and the answer drifts into English mid-sentence (this is where the series' th_ratio metric comes from), because the SFT data it saw was overwhelmingly English.

This chapter fixes both layers at once: SFT on Thai instruction data, done through LoRA instead of full fine-tuning.

2. What We're Going to Do​

We take Qwen3-0.6B and train it on 4,000 Thai instruction-answer pairs, using the exact same loss as chapter 1 plus two new pieces:

1. Completion mask — compute the loss only on the answer tokens, never the question (section 3.1 explains the comically broken failure you get if you skip this)
2. LoRA (Low-Rank Adaptation) — freeze all the original weights and instead train two small matrices laid over each layer

3. The Equations​

3.1 The SFT loss and the completion mask​

• $(x, y)$ = one training pair — $x$ is the instruction part (chat template included) and $y$ is the token sequence of the example the model actually sees during training
• $y_t$ = the token at position $t$, and $y_{<t}$ = every token before it
• $p_\theta$ = the probability predicted by the model with parameters $\theta$
• $m_t \in \{0,1\}$ = the completion mask — 1 only on answer-side tokens, and 0 on prompt tokens

Remove $m_t$ (i.e. set $m_t \equiv 1$ everywhere) and this equation instantly becomes the CPT objective from chapter 1. SFT is CPT on text staged as a conversation, plus one mask — nothing more.

But that one mask is half the battle, because $m_t \equiv 1$ is the easiest default to fall into (many pipelines, SFTTrainer included, will silently train this way if the collator isn't wired correctly). And in typical Thai instruction data, prompt-side tokens make up around 60% of each example — meaning most of your gradient is busy teaching the model to write the user's questions, not to answer them. The side effect this causes will make its appearance in section 9.

3.2 LoRA: training the correction, not the weights​

Instead of updating the weight matrix $W_0$ directly, LoRA pins $W_0$ in place and learns a difference that is the product of two small matrices:

• $W_0 \in \mathbb{R}^{d\times k}$ = the layer's original weights, frozen — they receive no gradient at all
• $B \in \mathbb{R}^{d\times r}$ and $A \in \mathbb{R}^{r\times k}$ = the two adapter matrices we actually train
• $r \ll \min(d, k)$ = the rank of the correction — LoRA's main dial (this chapter uses $r = 16$)
• $\alpha$ = a scale multiplier — the product $BA$ is always multiplied by $\alpha/r$ (here $\alpha = 32$, so $\alpha/r = 2$)

Two details in this definition matter more than they look:

$B$ is initialized to all zeros, so $\Delta W = 0$ at the first step — training starts from exactly the base model, with no period where random weights disturb it. ($A$ is random Gaussian — set both to zero and both gradients stay zero forever, since each is multiplied by the other's zeros.)

The divisor $r$ in $\alpha/r$ makes the update magnitude independent of the rank — double $r$ and the sum $BA$ has twice as many terms, but it gets divided right back. So you can sweep $r$ without re-tuning the learning rate every time.

And when training finishes, you have two choices: merge ($W' = W_0 + \frac{\alpha}{r}BA$, giving a single model with no added latency) or keep it separate — the second is the road this series takes, because a detachable adapter is what chapter 4 uses to build a reference model for free.

3.3 The fraction of parameters trained — a number to check, not to recite​

For a single $d \times k$ matrix, the adapter has $rd + rk$ parameters, a fraction of

• $d, k$ = the dimensions of the original weight matrix
• $r$ = the adapter's rank

Plug in the real values for Qwen3-0.6B (hidden 1024, intermediate 3072, 28 layers, adapters on all 7 matrices: q, k, v, o, gate, up, down) and you get 10,092,544 trainable parameters on a base of 596,049,920 = 1.69%.

4. Seeing the Equations​

One mask redirects the entire gradient​

There is nothing deeper in this figure than counting tokens — but the counting is exactly what people skip: a Thai prompt (system message and chat template included) is often longer than its answer, so training unmasked spends most of your compute teaching something you never wanted. The notebook also prints the actual ratio for the sampled dataset — the 60/40 in the figure stands in for a typical example; it is not a sacred constant.

Why rank 16 is "enough" — the low-rank hypothesis​

LoRA's hypothesis (Hu et al., 2021) is that fine-tuning moves the weights along only a few important directions relative to the full dimensionality of the matrix. This figure builds a fake $\Delta W$ with exactly that structure (a rank-16 signal + noise) and lets SVD dig it back out — if real updates look like this, a rank-16 adapter captures nearly all of the content. What the figure does not prove is that real updates always look like this. That is an empirical finding of the research, and it is why LoRA "usually" comes close to full fine-tuning — not "always."

Check the number, don't recite it​

A straight line on log-log axes confirms what the equation says: the fraction grows exactly linearly with $r$. And on this model, "under 1%" is only true when $r \le 9$ — which is not a value anyone actually uses for language tasks. The point we use ($r=16$) is 1.69%, or about 40 MB stored as fp32.

5. Setting Up the Environment​

Open Colab and pick Runtime → Change runtime type → T4 GPU (the free tier is enough).

torch_dtype=torch.float16      # not bfloat16
attn_implementation="sdpa"     # not flash_attention_2
fp16=True                      # in SFTConfig (not bf16=True)

cap = torch.cuda.get_device_capability(0)
print("compute capability:", cap)                    # T4 = (7, 5)
print("native bf16:", cap[0] >= 8)                   # T4 -> False
print("torch says   :", torch.cuda.is_bf16_supported())  # T4 -> True (counts emulation!)

for p in model.parameters():
    if p.requires_grad:
        p.data = p.data.float()      # cast the adapter only — not the whole model

The VRAM budget LoRA transforms​

Remember how in chapter 1 Adam's optimizer state took more room than the model itself (about 4.4 GB for $m$ and $v$)? Training only the adapter, Adam keeps state for just 10.1 million parameters — about 0.08 GB — to the point that adamw_bnb_8bit is no longer needed. The big remaining budget items are the base weights (fp16, ~1.2 GB) and activations.

Try switching between full fine-tuning and various LoRA ranks in the calculator and watch where the budget moves:

6. Preparing the Data​

We use airesearch/wangchanx-seed-free-synthetic-instruct-thai-120k — 120,000 synthetic Thai instruction-answer pairs from the WangchanX team, a Thai open-source NLP research effort (we sample 4,000).

from datasets import load_dataset

ds = load_dataset("airesearch/wangchanx-seed-free-synthetic-instruct-thai-120k", split="train")
ds = ds.shuffle(seed=42).select(range(4000))

def to_text(ex):
    messages = [
        {"role": "user", "content": ex["instruction"]},   # always check column names against the dataset card
        {"role": "assistant", "content": ex["output"]},
    ]
    return {"text": tok.apply_chat_template(messages, tokenize=False,
                                            enable_thinking=False)}

train_ds = ds.map(to_text, remove_columns=ds.column_names)

Two lines worth reading slowly:

• apply_chat_template assembles the messages into the format Qwen3 was trained on (<|im_start|>user … <|im_end|> …). We call it once, here and nowhere else — the reason lives in trap 2 of section 9.
• enable_thinking=False turns off Qwen3's thinking mode, keeping this chapter's examples simple and the mask straightforward.

Then perform the ritual that should become a habit: always decode the first example and look at it with your own eyes.

print(train_ds[0]["text"][:400])
# you must see <|im_start|>user ... <|im_end|> ... <|im_start|>assistant ... exactly once each per turn

The notebook also prints token-length statistics for the sampled set: the median on the prompt side versus the answer side, and the fraction of examples longer than 768 tokens — that last number will come back to collect its debt in trap 3.

7. The Main Code​

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer, SFTConfig, DataCollatorForCompletionOnlyLM

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",                # the post-trained model — this chapter fixes behavior, not knowledge
    torch_dtype=torch.float16,        # the T4 has no bf16
    attn_implementation="sdpa",       # the T4 has no FlashAttention-2
).cuda()

lora = LoraConfig(
    r=16,
    lora_alpha=32,                    # α/r = 2 — see equation 3.2
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora)
model.print_trainable_parameters()    # trust this line, not the number in any blog

for p in model.parameters():          # the fp16 box from section 5
    if p.requires_grad:
        p.data = p.data.float()

Next comes the piece that turns equation 3.1 into code — the collator that plays the role of $m_t$:

collator = DataCollatorForCompletionOnlyLM(
    response_template="<|im_start|>assistant\n",   # everything before this = prompt
    tokenizer=tok,
)

This collator sets the label of every token before <|im_start|>assistant to -100, the value PyTorch's loss function skips — the same value chapter 4 uses to mask out the prompt side in its seq_logp function. Same mask, showing up in a different chapter.

cfg = SFTConfig(
    output_dir="sft-out",
    dataset_text_field="text",
    max_seq_length=768,               # the notebook prints how many examples get truncated — trap 3
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,    # effective batch = 16
    num_train_epochs=1,
    learning_rate=2e-4,               # fine for LoRA — catastrophic for full FT (chapter 1)
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    packing=False,                    # the completion mask can't be used with packing directly
    fp16=True,                        # T4: not bf16
    logging_steps=10,
)

trainer = SFTTrainer(
    model=model,
    args=cfg,
    train_dataset=train_ds,
    data_collator=collator,
    processing_class=tok,
)
trainer.train()                       # ~15 minutes on a T4 (measured: 14.8)

Training done, save only the correction:

model.save_pretrained("qwen3-0.6b-th-sft-lora")   # ~40 MB — not 1.2 GB

This is the very adapter chapter 4 loads under the name kobkrit/qwen3-0.6b-th-sft-lora, serving as both the starting policy and (with the adapter switched off) the reference model — a cross-chapter gift priced at 40 MB.

8. Results​

The notebook measures 3 things before and after training and writes them to results.json:

1. TH-INSTR — the fraction of answers passing an instruction-following rubric (answers the question, stops on its own, stays in role) from the KobEval-TH benchmark, with a Wilson 95% CI
2. th_ratio — the fraction of Thai characters in the answers, the series' standing metric. The stock Qwen3-0.6B is notorious for drifting into English on Thai prompts; SFT on 4,000 all-Thai examples should move this number visibly — and if it doesn't, that is your signal to go hunt through the mask and the template before anything else.
3. The fraction of answers ending with eos within the token budget — the detector for the "won't stop" symptom from traps 1 and 3.

9. Comparison​

The notebook trains 3 variants on the same data, measured against the starting model:

SFT improves on every axis: TH-INSTR +10 points (73.3% → 83.3%), th_ratio up (more answers stay in Thai instead of drifting to English), and perplexity down from 21.3 to 17.7 — training only 1.69% of the parameters for a ~40 MB adapter.

On full fine-tuning and the no-mask row: the notebook does train a no-mask LoRA to show the behaviour where the model finishes its answer and then invents the user's next question (samples in §8), but it does not score full metrics for it, and it does not run a 100%-parameter full fine-tune — that exceeds the free-Colab budget. We report only what was measured and do not fill in numbers for runs we did not do.

10. Summary​

• SFT is CPT on conversations + a completion mask — the same loss as chapter 1; what changes is that the data is staged, and the mask chooses which tokens count
• The mask is half the battle — $m_t \equiv 1$ sends 60% of the gradient into practicing questions, and yields a model that answers, then writes the next question itself
• LoRA trains a "correction," not the weights: $W' = W_0 + \frac{\alpha}{r}BA$ with $B$ starting at zero, so training departs from exactly the base — and $\alpha/r$ lets you change rank without re-tuning the lr
• 1.69% is not "under 1%" — adapters grow linearly with hidden size while the base grows quadratically, so the smaller the model, the bigger the adapter's share; always check with print_trainable_parameters()
• ~40 MB of adapter per task: one base plus many adapters — and it is the reason chapter 4's reference model comes free
• lr 2e-4 is safe because the base is frozen — the same value destroys the model when training everything
• pad ≠ eos, template once, never truncate mid-answer — three traps whose symptoms surface at inference but whose causes live in the data

Next chapter: RLHF and PPO — the model can answer now, but "answers well" cannot be written as a loss function. We will have humans compare answers, train a reward model from those preferences, then use reinforcement learning to push the model toward it — with this chapter's 40 MB adapter as the starting point.

References​

1. Hu et al. (2021). LoRA: Low-Rank Adaptation of Large Language Models — the source of the W' = W₀ + (α/r)BA equation in section 3
2. Aghajanyan et al. (2020). Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning — evidence that fine-tuning has low intrinsic dimension -- why LoRA works at all
3. Dettmers et al. (2023). QLoRA: Efficient Finetuning of Quantized LLMs — extends LoRA with 4-bit so larger models fit one GPU
4. Biderman et al. (2024). LoRA Learns Less and Forgets Less — LoRA trades less learning for less forgetting -- read alongside section 9
5. Ouyang et al. (2022). Training language models to follow instructions with human feedback — InstructGPT: the origin of the whole SFT -> RM -> PPO pipeline
6. Wang et al. (2022). Self-Instruct: Aligning Language Models with Self-Generated Instructions — generating instruction data from the model itself
7. Zhou et al. (2023). LIMA: Less Is More for Alignment — a few thousand high-quality examples suffice -- why we use only 4,000

The writing, code and notebooks in this series are licensed under CC BY-NC-SA 4.0 — reuse and adapt them freely with attribution, for non-commercial purposes, and share your adaptations under the same terms. Third-party models and datasets referenced here keep their own licences.

1. The Problem​

SFT in chapter 2 carries one hidden assumption: there must be an answer key to imitate. But think about what we actually want — say, "solve the math problem correctly, and explain it in readable Thai." That sentence has no single answer key: good answers come in a hundred shapes, and "readable" cannot be written as an equation.

Try to optimize these things directly and you always run into two walls:

Wall one — quality cannot be written as a loss. "Better" cannot be defined as a function, but humans are very good at comparing: show them two answers and ask which one they prefer, and they can do it immediately and fairly consistently. So the data you can actually collect comes in threes: a prompt $x$, the chosen answer $y_w$, the rejected answer $y_l$.

Wall two — even with a score, you cannot backprop. Suppose a magical function $r(x,y)$ scored every answer. You still couldn't train supervised, because the answer $y$ comes from sampling tokens one at a time. The score arrives after the sampling is done, and derivatives cannot travel back through sampling — the path from $r$ back to the weights $\theta$ breaks exactly there.

Hence the chapter's title: we are about to optimize a reward we cannot differentiate. The tool that can do that is called reinforcement learning.

2. What We're Going to Do​

RLHF gets past both walls with a two-step walk:

• Stage A — Reward Model: train a model $r_\phi(x,y)$ to imitate human comparisons from preference pairs (defeats wall one, and stands in for humans who can't score fast enough)
• Stage B — PPO: use policy-gradient RL to push the policy toward $r_\phi$'s scores without differentiating through the sampling (defeats wall two), with a KL leash holding it back from fleeing the starting model

The price you pay is complexity: during training there are four models resident in VRAM at once — the policy $\pi_\theta$ (the one being trained), the reference $\pi_{\text{ref}}$ (the frozen starting model), the reward model $r_\phi$, and a value network $V_\psi$ that hasn't introduced itself yet (wait for section 3.4).

And one more sentence to hold on to for the whole series: this chapter's objective is the master equation of the second half of the series. Chapter 4 (DPO) solves it in closed form until the reward model and the RL loop cancel out. Chapter 5 (GRPO) changes how the advantage is estimated until the value network disappears. Understand this one chapter, and the next two become "deleting parts" — readable at sight.

3. The Equations​

3.1 The reward model: Bradley–Terry​

Stage A trains $r_\phi$ with one short loss:

• $r_\phi(x,y)$ = a single scalar score per text — in practice a language model whose head is swapped for a single linear layer (num_labels=1)
• $\sigma$ = the sigmoid, turning the score difference into the probability a human picks $y_w$ (the Bradley–Terry model)
• the wider the gap $r_\phi(x,y_w) - r_\phi(x,y_l)$, the lower the loss

The point people overlook and pay for later: this loss sees only the difference of the scores. Replace $r_\phi$ with $r_\phi + c$ for any constant $c$ — the loss doesn't change at all. Which means the absolute scale of a reward model is meaningless and is not pinned down by training. Two runs can yield mean scores of 3.7 and −12.4 that rank identically. This is why you must always standardize rewards before feeding them into PPO (subtract the mean, divide by the std) — hold onto this; it returns in sections 7 and 9.

3.2 The master equation: the RLHF objective​

If you memorize a single equation from this whole series, memorize this one:

In plain language: "collect as much reward as you can, but every step you walk away from the starting model costs a fine."

• $\pi_\theta$ = the policy, the model being trained — note that $y$ is sampled from $\pi_\theta$ itself; this is the structural difference from SFT, which learns from static data sitting in a file
• $\pi_{\text{ref}}$ = the reference, the starting model (the post-SFT model from chapter 2), frozen throughout training
• $\beta$ = the price per nat of wandering off — the tightness of the leash
• $\mathbb{D}_{\text{KL}}$ = the distributional distance between the policy and the reference

3.3 The PPO clipped surrogate: Stage B's engine​

Raw policy gradient (REINFORCE) can use a batch of rollouts for exactly one update before throwing it away — very expensive, because generation is the bottleneck. PPO wants to squeeze the same rollouts for several epochs, so it needs a correction factor (the importance sampling ratio):

• $s_t$ = the state at position $t$: the prompt plus every token sampled so far
• $a_t$ = the "action": the next token, already sampled during the rollout
• $\pi_{\theta_{\text{old}}}$ = a snapshot of the policy at rollout time — computed once and frozen

Then $\rho_t$ gets pinched with a clip:

• $\hat A_t$ = the advantage: "how much better than expected was this token" (defined in the next subsection)
• $\epsilon$ = the width of the trust region (standard value 0.2)

The heart is that min + clip work together with deliberate pessimism: if $\hat A_t$ is positive (a good token), the payoff from pushing $\rho_t$ is capped at $1+\epsilon$ — pushing beyond that earns nothing, gradient zero. But if $\hat A_t$ is negative (a bad token), the min always picks the worse branch — the penalty has no floor. One-sentence summary: gains capped, losses uncapped. The policy therefore moves in small steps that stay close to where it already is.

3.4 GAE: computing the advantage without drowning in noise​

The advantage is built from the TD error of the value network $V_\psi$:

• $V_\psi(s_t)$ = the value network's prediction of "from here to the end, how much more reward will be collected" — this is model number four
• $r_t$ = the per-token reward (in our task: a KL penalty at every position, plus the task score at the final token)
• $\gamma$ = the discount factor (LLM work usually uses 1.0)
• $\lambda$ = the bias–variance dial: $\lambda = 0$ trusts $V_\psi$ wholeheartedly (high bias if $V_\psi$ mispredicts), $\lambda = 1$ doesn't trust it at all and waits for the real outcome to the very end (high variance, carrying the noise of the whole trajectory); the popular value is 0.95

3.5 The full PPO loss: three terms, two models​

Assemble every piece into the one loss the optimizer actually sees (written as a minimization):

• first term = the clipped surrogate from 3.3 (negated because we want to maximize it)
• second term = the value loss teaching $V_\psi$ to predict close to the actual return $\hat R_t$; $c_1$ is usually 0.5
• third term = the entropy bonus $\mathcal{H}$, keeping the distribution from collapsing too fast; $c_2$ is usually 0.01
• as for the KL leash from equation 3.2, standard practice folds it into the per-token reward: $r_t \leftarrow r_t - \beta\,(\log\pi_\theta - \log\pi_{\text{ref}})$, which is the approach we take in section 7

Count all the toys that need tuning: 4 models, plus $\epsilon, \beta, \gamma, \lambda, c_1, c_2$, and two separate learning rates. This is why PPO is famous for "run it twice with different seeds, get two completely different stories" — and it is the reason for chapter 4's entire existence.

4. Seeing the Equations​

Bradley–Terry: the gradient piles onto the pairs still ranked wrong​

Pick the Bradley-Terry mode in the tool below and drag the margin: pairs the reward model already ranks confidently right (large positive margin) have almost no gradient left — Stage A training automatically spends its budget on the pairs it still ranks wrong.

If this plot reminds you of DPO's gradient in chapter 4 — that is no coincidence. DPO lifts this Bradley–Terry model wholesale, merely changing what plays the role of the reward.

The trust region made visible: min + clip​

Study the asymmetry closely — it is PPO's entire personality: moving toward something good earns at most 20% per step, but hand too much probability to a bad token and you are hauled back at full force, every time. The "flat region" in the figure is the trust region that lets PPO reuse old rollouts for several epochs without exploding.

GAE's bias–variance dial​

Read the lines from the bottom up: at $\lambda = 0$ (green) the advantage hugs zero nearly everywhere — the end-of-episode score never reaches the early tokens, because everything is filtered through a $V_\psi$ that mispredicts. At $\lambda = 1$ (red), every token gets full credit from the ending but carries the accumulated noise of the whole trajectory with it. $\lambda = 0.95$ (blue) is the middle ground the whole field has settled on — the signal travels far, but the noise is damped.

Before moving on, build your intuition for the word advantage — "how much better than expected" — by hand: this tool uses the group mean as its baseline instead of $V_\psi$ (a full-blown spoiler of chapter 5). Press the All correct preset and watch what happens when every answer earns the same reward:

Advantage zero across the whole group = no signal to learn from — hold onto that feeling for section 9.

The KL leash: equation 3.2 drawn as a picture​

The green line's stopping point is not a guess — it is the mathematics of equation 3.2: optimization stops exactly where the reward gained per nat equals β; walking further loses money. When $\beta = 0$ that stopping condition doesn't exist — every nat of wandering that buys even a sliver of reward is "profit," so the model keeps running out of natural language for as long as the reward number still twitches upward.

5. Setting Up the Environment​

Open Colab and pick Runtime → Change runtime type → T4 GPU (the free tier is enough).

torch_dtype=torch.float16      # not bfloat16
attn_implementation="sdpa"     # not flash_attention_2

cap = torch.cuda.get_device_capability(0)
print("compute capability:", cap)                    # T4 = (7, 5)
print("native bf16:", cap[0] >= 8)                   # T4 -> False
print("torch says   :", torch.cuda.is_bf16_supported())  # T4 -> True (counts emulation!)

Why RLHF is expensive, in one picture​

The numbers in the figure are just the weights — before activations, the KV cache during generation, gradients, and optimizer state. And at 0.6B everything still looks tiny, but the ×4 multiplier goes nowhere as you scale: at 7B it is 56 GB before you've done anything at all.

6. Preparing the Data​

Stage A — preference pairs for the reward model​

We use iapp/dpo_thai_tutorial (100 pairs, Apache-2.0) — a Thai preference dataset I built myself for this series and released for free reuse. Each row has prompt, chosen, rejected, hand-curated with an emphasis on politeness and natural-sounding Thai.

Split 80/20: train on 80 pairs, hold 20 pairs out — never touched during training. Stage A's passing bar is pairwise ranking accuracy on the 20 held-out pairs, with a Wilson 95% CI that does not straddle 0.5 — "significantly better than a coin flip" already proves the mechanism, because 20 pairs genuinely cannot make the CI any tighter.

Stage B — problems a rule can grade​

For the PPO loop we use 64 math problems from VISAI-AI/gsm8k-thai (a Thai translation of GSM8K) and score them with a verifiable rule instead of the Stage A reward model:

import re

def rule_reward(response: str, gold: int) -> float:
    """+1.0 if the last integer in the answer is correct, +0.2 if the answer is actually in Thai"""
    nums = re.findall(r"-?\d+", response.replace(",", ""))
    correct = 1.0 if nums and int(nums[-1]) == gold else 0.0
    thai = sum("\u0e01" <= ch <= "\u0e5b" for ch in response)   # the Thai Unicode block, U+0E01..U+0E5B
    thai_bonus = 0.2 if thai / max(len(response), 1) >= 0.5 else 0.0
    return correct + thai_bonus

7. The Main Code​

7.1 Stage A — training a real reward model (~6 minutes)​

Turn a language model into a scoring machine: the LM head is replaced with a single linear layer returning one scalar.

import torch, torch.nn.functional as F
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from peft import LoraConfig, get_peft_model

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
rm = AutoModelForSequenceClassification.from_pretrained(
    "Qwen/Qwen3-0.6B",
    num_labels=1,                          # scalar head: one score per text
    torch_dtype=torch.float16,             # the T4 has no bf16
    attn_implementation="sdpa",
).cuda()
rm.config.pad_token_id = tok.pad_token_id

rm = get_peft_model(rm, LoraConfig(
    task_type="SEQ_CLS", r=8, lora_alpha=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    modules_to_save=["score"],             # the score head starts random — must be trained in full
))
for p in rm.parameters():
    if p.requires_grad:
        p.data = p.data.float()            # the fp16 lesson from chapter 1: train in fp32

def rm_loss(chosen, rejected):
    s_w = rm(**chosen).logits.squeeze(-1)     # r_φ(x, y_w)
    s_l = rm(**rejected).logits.squeeze(-1)   # r_φ(x, y_l)
    return -F.logsigmoid(s_w - s_l).mean()    # equation 3.1, literally

Train 3 epochs on 80 pairs, then measure pairwise accuracy on the 20 held-out pairs with a Wilson CI. If the CI clears 0.5 — you have just trained the first reward model of your life, on 80 rows of data.

7.2 Stage B — PPO written from scratch, ~120 lines (~10 minutes)​

We will not use TRL's PPOTrainer, and this is a deliberate decision, not stubbornness: TRL has moved PPOTrainer into trl.experimental and announced plans to remove it in 0.29.0 — code taught against that library will stop running within a few months. A self-written PPO loop of about 120 lines will run for as long as PyTorch exists. And more important still: write it yourself and you will know what every line does, just like the 25-line DPO loss in chapter 4.

The pieces on the board: policy = base + LoRA adapter (from chapter 2, is_trainable=True), reference = the same base with the adapter off, and a value head — a two-layer MLP plugged into the final hidden state:

value_head = torch.nn.Sequential(
    torch.nn.Linear(1024, 1024), torch.nn.Tanh(),
    torch.nn.Linear(1024, 1),
).float().cuda()                           # ~1M parameters — tiny next to the other three

The rollout: sample answers 8 prompts at a time (max_new_tokens=200, do_sample=True), score with rule_reward, standardize the scores within the batch, and store the rollout-time log-probs, detached, immediately. Then enter this update loop — the 40 lines that are the heart of the whole chapter (the full loop is in the notebook):

def compute_gae(rewards, values, gamma=1.0, lam=0.95):
    """rewards: [T] for one response, values: [T+1] (last slot = 0 after the end)"""
    adv, acc = torch.zeros_like(rewards), 0.0
    for t in reversed(range(rewards.shape[0])):
        delta = rewards[t] + gamma * values[t + 1] - values[t]   # equation 3.4
        acc = delta + gamma * lam * acc                          # Â_t = δ_t + γλ Â_{t+1}
        adv[t] = acc
    return adv

def ppo_update(rollout, eps=0.2, beta=0.05):
    out = policy(rollout.ids, output_hidden_states=True)         # one forward, two outputs
    logp = gather_logprobs(out.logits, rollout.ids)              # [B, T], carries gradient
    with torch.no_grad(), policy.disable_adapter():
        logp_ref = gather_logprobs(policy(rollout.ids).logits, rollout.ids)

    # per-token reward = KL penalty at every position + task score at the last token (equation 3.2)
    rew = -beta * (logp.detach() - logp_ref)
    rew[:, -1] += rollout.scores_std                 # the standardized rule score

    hidden = out.hidden_states[-1].detach().float()  # detach = cut the gradient into the trunk
    values = value_head(hidden).squeeze(-1)
    adv = torch.stack([compute_gae(r, F.pad(v, (0, 1))) for r, v in zip(rew, values)])
    adv = (adv - adv.mean()) / (adv.std() + 1e-8)    # standardize the advantage once more

    log_ratio = (logp - rollout.logp_old).clamp(-10, 10)   # prevents fp16 overflow!
    ratio = log_ratio.exp()                                # ρ_t (equation 3.3)
    surr = torch.min(ratio * adv.detach(),
                     ratio.clamp(1 - eps, 1 + eps) * adv.detach())

    returns = (adv + values).detach()                # the value head's target
    m = rollout.resp_mask                            # count response tokens only
    loss = (-(surr * m).sum() + 0.5 * ((values - returns).pow(2) * m).sum()) / m.sum()
    return loss                                      # (the entropy term is in the notebook)

Settings: $\epsilon = 0.2$, $\beta = 0.05$, $\gamma = 1.0$, $\lambda = 0.95$, 4 epochs per rollout on 64 prompts, adapter LR 1e-5, value head LR 1e-4. Roughly 10 minutes per run on a T4 (the notebook runs twice: $\beta = 0.05$ and $\beta = 0$).

8. Results​

The notebook tracks three curves simultaneously at every update and writes them to results.json:

1. Mean reward per rollout — should climb (this is what we're buying)
2. Per-token KL against $\pi_{\text{ref}}$ — should grow and then saturate under the ceiling that $\beta$ sets (this is the price we pay)
3. Mean answer length — the disease detector: length spiking or cratering abnormally is the first sign of a policy going strange

Catching reward hacking red-handed​

The notebook's second run sets $\beta = 0$ and touches absolutely nothing else — no leash. The pattern expected: reward climbs as fast or faster, but KL soars with no ceiling, and the answers start to degenerate — repetitive, oddly short, or collapsing into a fixed formula with a number stuffed at the end of the sentence. Because rule_reward sees only the final number and the fraction of Thai characters, everything it cannot see is free for the model to throw away.

Samples of degenerate answers from the $\beta = 0$ run are printed by the notebook's final cell:

[This space is filled only from the notebook's real run — I will not invent degenerate
 examples myself, because this whole series stands on the rule that no number or output
 is ever made up. Run the notebook and the "hacking exhibits" cell will show 2-3 real
 answers along with their KL.]

9. Comparison​

The notebook runs the same battery over three systems on held-out Thai math problems (TH-MATH):

The pattern you should expect: the $\beta = 0.05$ row nudges accuracy up slightly or holds even, with KL steady, while the $\beta = 0$ row earns the highest training-time rule reward but should not do better on held-out accuracy — and the language breaks. Cheated points do not transfer to new problems.

If you see something else, read it this way:

• All three rows nearly identical → the advantage was near zero throughout; check whether rewards within each batch have any spread (if the model gets every problem wrong, the post-standardization advantage is pure noise — the same feeling as the All wrong preset in section 4's tool)
• β = 0 but KL doesn't soar → 4 epochs times 64 prompts is too short for the hack to ripen — extend the run and look again; don't conclude "no hacking" just yet
• β = 0.05 but KL soars anyway → almost certainly the scores weren't standardized, letting the reward's scale swamp β — or the LR is too high

Traps to watch for​

1. Not standardizing the reward The reward's scale is undefined per equation 3.1 — two RMs that rank identically can differ in scale by a factor of ten. That scale multiplies straight into the advantage and the gradient: a run that used to be stable explodes the moment you swap RMs.

2. Recomputing logp_old inside the update loop The quietest bug of the chapter: no error, the loss looks normal, but $\rho_t = 1$ forever. The clip never fires, the trust region doesn't exist — you are running REINFORCE while believing you are running PPO.

3. Sharing the trunk between policy and value without a stop-gradient The value loss has a large scale and flows back into the backbone, overwriting language ability. This notebook guards against it with .detach() on the hidden state (visible in the section 7.2 code) — a one-line price.

4. fp16 ratio overflow exp(log_ratio) can blow past the fp16 ceiling starting around log-ratio 11.1. Always .clamp(-10, 10) before .exp() — and if you see log-ratios reaching that level often, it is a sign the policy is fleeing the old rollout too fast (lower the LR or the number of epochs per rollout).

10. Summary​

• RLHF = a two-step detour to optimize what cannot be differentiated: train a judge ($r_\phi$), then use RL to run toward the judge's score
• Bradley–Terry sees only differences — a reward's absolute scale is undefined, so always standardize before use
• The master equation $\max\ \mathbb{E}[r] - \beta\,\mathbb{D}_{\text{KL}}$ is the one equation to memorize — chapters 4 and 5 are this same equation solved by other means
• KL is not a regularizer — it is the system's only stopping condition; the moment you remove it, Goodhart goes to work
• min + clip = pessimism by design: gains capped, losses uncapped — the trust region that makes old rollouts reusable
• GAE is the bias–variance dial, and the $V_\psi$ it leans on is the fourth model — the one GRPO deletes in chapter 5
• PPO is expensive by structure, not by badly written code: 4 models + around ten hyperparameters is the sticker price

Next chapter: DPO — DPO deletes both the reward model and the RL loop with pure algebra, starting from master equation 3.2, the one you just memorized.

References​

1. Schulman et al. (2017). Proximal Policy Optimization Algorithms — the original PPO paper: the clipped surrogate in section 3
2. Schulman et al. (2015). High-Dimensional Continuous Control Using Generalized Advantage Estimation — GAE: the advantage estimator PPO uses
3. Christiano et al. (2017). Deep reinforcement learning from human preferences — the paper that started RL from human preferences
4. Stiennon et al. (2020). Learning to summarize from human feedback — the first convincingly working RLHF, on summarization
5. Ouyang et al. (2022). Training language models to follow instructions with human feedback — InstructGPT: the origin of the whole SFT -> RM -> PPO pipeline
6. Bai et al. (2022). Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback — the helpful/harmless dataset and its lessons about KL
7. Zheng et al. (2023). Secrets of RLHF in Large Language Models Part I: PPO — the PPO implementation details other papers omit
8. Bradley & Terry (1952). Rank Analysis of Incomplete Block Designs: I. The Method of Paired Comparisons — the Bradley-Terry model every reward model rests on

The writing, code and notebooks in this series are licensed under CC BY-NC-SA 4.0 — reuse and adapt them freely with attribution, for non-commercial purposes, and share your adaptations under the same terms. Third-party models and datasets referenced here keep their own licences.

1. The Problem​

Say you want your AI assistant to "always answer in Thai" — sounds simple. Now try writing that as a loss function. You can't.

This is the central problem of alignment: answer quality cannot be written as an equation. "More polite," "more natural," "doesn't drift into English" — there is no single correct answer key. There is only comparison: show a human two answers and ask which one they prefer. So the data comes in threes: a prompt $x$, a preferred answer $y_w$ (chosen), and a dispreferred one $y_l$ (rejected).

PPO-style RLHF solves this by taking a two-step detour:

1. Train a reward model $r_\phi(x,y)$ to imitate human preferences
2. Use RL to push the policy toward high scores from that reward model

The detour has a price:

So this chapter asks one short question: can we skip steps 1 and 2 entirely?

2. What We're Going to Do​

Yes, we can — and the reason is beautiful.

The starting point is the observation that the KL-constrained RLHF objective has a closed-form solution. We already know exactly what the optimal policy looks like, without running a single step of RL. Knowing that, we invert the equation — instead of asking "what policy does this reward produce?" we ask "what reward does this policy imply?"

This is DPO (Direct Preference Optimization), proposed by Rafailov et al. (2023). The "Direct" comes from optimizing on preference data directly, with no intermediary.

3. The Equations​

3.1 Setting up: the RLHF objective​

In plain language: "maximize the reward, but don't wander too far from where you started."

• $\pi$ = the policy, i.e. the model we're training
• $\pi_{\text{ref}}$ = the reference policy, i.e. the starting model (here, the post-SFT model from chapter 2)
• $r(x,y)$ = the reward for answer $y$ given prompt $x$
• $\beta$ = how tight the leash is; higher values pull harder back toward $\pi_{\text{ref}}$

The KL term is not decoration. Without it the model runs off to wherever reward is highest and the language falls apart.

3.2 Step 1 — the closed-form solution​

That problem can be solved by hand (it amounts to finding the distribution that minimizes KL against a target distribution), giving

• $Z(x) = \sum_{y}\pi_{\text{ref}}(y|x)\exp\!\big(r(x,y)/\beta\big)$ is the partition function, the denominator that makes everything sum to 1
• Note that $Z(x)$ depends only on $x$, not on $y$ — hold on to that sentence, it's about to become the hero of the story

The intuition: the optimal policy is the original model, reweighted by $\exp(r/\beta)$. High-reward answers get their probability amplified, low-reward answers get pushed down, but everything starts from the original shape of $\pi_{\text{ref}}$.

In practice we can't compute $Z(x)$, because it requires summing over every possible answer in the universe. That's why people reach for RL — and it's exactly why DPO doesn't have to.

3.3 Step 2 — invert the equation to get the reward​

Take the log of both sides and rearrange:

This line is the crux: any reward function can be rewritten in terms of the optimal policy and the starting policy. Which means that if we have two models, we can compute their implicit reward immediately, without ever training a reward model.

3.4 Step 3 — substitute into Bradley-Terry and $Z(x)$ vanishes​

The standard model of preference is Bradley-Terry: the probability a human picks $y_w$ over $y_l$ is

where $\sigma$ is the sigmoid. Notice that reward appears in this equation only as a difference. Substitute equation 3.3 — $\beta\log Z(x)$ appears identically on both sides because it's the same $x$ — so it cancels out.

3.5 The gradient — where the intuition lives​

where $\hat r = \beta\log\big(\pi_\theta/\pi_{\text{ref}}\big)$ is called the implicit reward.

Read it piece by piece:

• The right-hand parenthesis = the direction: push the log-prob of $y_w$ up and the log-prob of $y_l$ down, simultaneously
• $\sigma(\hat r_l - \hat r_w)$ = the weight: "how badly is the model ranking this pair?"

That weight is the most important teaching point here. If the model already ranks the pair correctly ($\hat r_w$ clearly greater than $\hat r_l$), then $\sigma(\hat r_l - \hat r_w)$ approaches zero and the pair contributes almost no gradient. DPO therefore focuses on its own mistakes automatically, with nobody curating the data for it.

4. Seeing the Equations​

Bradley-Terry: from a reward difference to a probability​

Where the difference is zero the probability is exactly 0.5 — "the model has no opinion." And because the curve only cares about the difference, adding a constant to both rewards changes nothing at all. That is the geometric reason $Z(x)$ can cancel.

Loss and the gradient weight​

Let $\Delta = \log\frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)} - \log\frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)}$. Then the loss is $-\log\sigma(\beta\Delta)$ and the gradient weight is $\sigma(-\beta\Delta)$.

The right panel is the one to study. Once the margin is strongly positive, the weight converges to zero — that pair has "graduated." And the higher $\beta$ is, the steeper the curve: the model both learns fast and "stops learning" fast. At $\beta = 1.0$, pairs with a margin past 4 have essentially no gradient left. At $\beta = 0.1$ the curve is far flatter, so the model keeps collecting gradient from every pair — slower, but steadier.

$\beta$ controls how far the model can drift from where it started​

Read this one right to left: at $\beta = 0.1$ nearly all the probability mass collapses onto $y_5$, the highest-reward answer. That is mode collapse — a great score with all diversity wiped out. At $\beta = 10$ the bars nearly overlap the dashed line, meaning almost nothing was learned. $\beta$ isn't a hyperparameter you "tune for lowest loss." It is a choice of where to sit on the trade-off between following preferences and keeping your original identity.

Drag the slider yourself and watch how loss and gradient change shape:

5. Setting Up the Environment​

Open Colab and pick Runtime → Change runtime type → T4 GPU (the free tier suffices).

torch_dtype=torch.float16      # not bfloat16
attn_implementation="sdpa"     # not flash_attention_2
fp16=True                      # in DPOConfig (not bf16=True)

cap = torch.cuda.get_device_capability(0)
print("compute capability:", cap)                    # T4 = (7, 5)
print("native bf16:", cap[0] >= 8)                   # T4 -> False
print("torch says   :", torch.cuda.is_bf16_supported())  # T4 -> True (counts emulation!)

A reference model that costs zero extra VRAM​

DPO needs both $\pi_\theta$ and $\pi_{\text{ref}}$, which sounds like loading two models. But we're building on the LoRA adapter from chapter 2, which lets both share the same base weights.

from peft import PeftModel

base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    torch_dtype=torch.float16,
    attn_implementation="sdpa",
).cuda()

policy = PeftModel.from_pretrained(base, "kobkrit/qwen3-0.6b-th-sft-lora", is_trainable=True)

policy is base + adapter, while $\pi_{\text{ref}}$ is that same base with the adapter switched off. Just call it inside the policy.disable_adapter() context manager — nothing extra to load.

6. Preparing the Data​

DPO data needs exactly three columns: prompt, chosen, rejected.

Set 1 — iapp/dpo_thai_tutorial (100 pairs, Apache-2.0) A dataset I built myself for this series and released under Apache-2.0 so anyone can reuse it. Hand-curated Thai preference pairs, emphasizing politeness and natural-sounding language.

Set 2 — roughly 400 pairs built from scratch out of airesearch/wangchanx-seed-free-synthetic-instruct-thai-120k. The construction is about as direct as it gets:

• chosen = the Thai reference answer shipped with the dataset
• rejected = the base model's own output under greedy decoding

def build_rejected(prompt):
    with torch.no_grad(), policy.disable_adapter():
        out = policy.generate(**tok(prompt, return_tensors="pt").to("cuda"),
                              max_new_tokens=192, do_sample=False)
    return tok.decode(out[0], skip_special_tokens=True)[len(prompt):]

That gives roughly 500 pairs, with 15% held out for evaluation and never touched during training.

7. The Main Code​

The point of this section isn't calling a library. It's writing the DPO loss by hand and then proving it matches the real thing. If the equations in section 3 didn't convince you, this code will.

7.1 Twenty-five lines, written from scratch​

import torch, torch.nn.functional as F

def seq_logp(model, input_ids, attention_mask, labels):
    """Sum of log-probs over the ANSWER ONLY (prompt tokens are masked to -100)"""
    logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
    logits = logits[:, :-1, :]                   # position t predicts the token at t+1
    target = labels[:, 1:]                       # so the targets shift by 1
    mask = target.ne(-100)                       # count answer tokens only
    target = target.masked_fill(~mask, 0)        # keep gather from breaking on -100
    logp = torch.log_softmax(logits.float(), dim=-1)
    tokp = logp.gather(-1, target.unsqueeze(-1)).squeeze(-1)
    return (tokp * mask).sum(-1)                 # [B] — a sum, not a mean

def dpo_loss(policy, batch, beta=0.1):
    pi_w = seq_logp(policy, batch["chosen_ids"],   batch["chosen_mask"],   batch["chosen_labels"])
    pi_l = seq_logp(policy, batch["rejected_ids"], batch["rejected_mask"], batch["rejected_labels"])

    with torch.no_grad(), policy.disable_adapter():        # reference: adapter off + no gradient
        ref_w = seq_logp(policy, batch["chosen_ids"],   batch["chosen_mask"],   batch["chosen_labels"])
        ref_l = seq_logp(policy, batch["rejected_ids"], batch["rejected_mask"], batch["rejected_labels"])

    delta = (pi_w - ref_w) - (pi_l - ref_l)                # the Δ from section 4
    loss = -F.logsigmoid(beta * delta).mean()              # equation 3.4, literally

    r_w = beta * (pi_w - ref_w).detach()                   # implicit reward of chosen
    r_l = beta * (pi_l - ref_l).detach()                   # implicit reward of rejected
    return loss, r_w, r_l

All of DPO is right there. Nothing else is hiding. The delta line is equation 3.4 transcribed one-to-one, and -F.logsigmoid(beta * delta) is the entire loss.

7.2 Proving it matches TRL​

from trl import DPOTrainer, DPOConfig

cfg = DPOConfig(
    output_dir="dpo-out",
    beta=0.1,
    loss_type="sigmoid",              # must match the formula we wrote by hand
    label_smoothing=0.0,              # anything nonzero and it stops being equation 3.4
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,    # effective batch = 16
    num_train_epochs=2,
    learning_rate=5e-6,               # much lower than SFT — see the warning below
    lr_scheduler_type="cosine",
    warmup_ratio=0.1,
    max_length=768,
    max_prompt_length=256,
    fp16=True,                        # the T4 has no bf16
    logging_steps=5,
)

trainer = DPOTrainer(model=policy, args=cfg, train_dataset=train_ds, processing_class=tok)

# Critical: a freshly created LoRA has lora_B = 0, so the policy is exactly the
# reference. The margin is zero and BOTH implementations return ln 2 — even if the
# formula is wrong. Perturb the weights first or the assert proves nothing.
with torch.no_grad():
    for name, p in policy.named_parameters():
        if "lora_B" in name:
            p.add_(torch.randn_like(p) * 0.01)

loss_manual, _, _ = dpo_loss(policy, batch, beta=0.1)
loss_trl = trainer.compute_loss(policy, trl_batch)

assert torch.allclose(loss_manual, loss_trl, atol=1e-4)
print("match:", loss_manual.item(), loss_trl.item())