flowchart TD
  A["1. Load dataset\n(chat-format, train/eval split)"] --> B["2. Load model\n(bf16, attn_implementation)"]
  B --> C["3. Attach PEFT config\n(LoraConfig) — or None for full FT"]
  C --> D["4. Build SFTConfig\n(TrainingArguments levers)"]
  D --> E["5. trainer.train()\n(loss curve, eval every N steps)"]
  E --> F["6. Save / merge / load\n(adapter or full)"]

from datasets import load_dataset

dataset = load_dataset("trl-lib/Capybara", split="train")
# Standard split: hold out 5-10% for eval — never train on everything.
split = dataset.train_test_split(test_size=0.05, seed=42)
train_ds, eval_ds = split["train"], split["test"]

import torch
# TRL loads the model from a string — bf16 + FlashAttention 2 are passed via model_init_kwargs
# (shown explicitly here; the trainer accepts a string model id and these kwargs)

from peft import LoraConfig

peft_config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type="CAUSAL_LM",
)

from trl import SFTConfig, SFTTrainer

training_args = SFTConfig(
    output_dir="./sft-out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,        # effective batch = 2 * 8 = 16
    learning_rate=2e-4,                   # LoRA needs higher LR than full FT
    num_train_epochs=1,
    warmup_ratio=0.05,
    lr_scheduler_type="cosine",
    bf16=True,                            # use bf16, not fp16
    gradient_checkpointing=True,
    logging_steps=10,
    eval_strategy="steps",
    eval_steps=50,
    save_strategy="steps",
    save_steps=100,
    report_to="wandb",                    # or "trackio", or "none"
    packing=True,
    max_length=2048,
)

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-3B-Instruct",   # string id, or a loaded model object
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=eval_ds,
    peft_config=peft_config,
)
trainer.train()

metrics = trainer.evaluate()
print(metrics)                      # eval_loss, eval_runtime, etc.
trainer.save_model("./sft-final")   # saves adapter (PEFT) or full model

from peft import PeftModel
from transformers import AutoModelForCausalLM

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B-Instruct", torch_dtype=torch.bfloat16)
model = PeftModel.from_pretrained(base, "./sft-final")
merged = model.merge_and_unload()          # adapter folded into base weights
merged.save_pretrained("./sft-merged")     # a single standalone model

from peft import PeftModel
model = PeftModel.from_pretrained(base, "./sft-final")
# Use model; swap to a different adapter with model.load_adapter(other_path)

# Module FT11 — The Training Loop with TRL

**Course**: Course 3 — LLM Fine-Tuning Masterclass
**Module**: FT11 — The Training Loop with TRL
**Duration**: 90 minutes
**Level**: Senior Engineer and above
**Prerequisites**: FT08 (LoRA/QLoRA), FT10 (Full FT vs PEFT decision)

---

## Learning Objectives

After completing this module, you will be able to:

1. Run a complete, instrumented SFT job end-to-end with TRL's `SFTTrainer` — dataset loading, model loading, PEFT attachment, `TrainingArguments`, `trainer.train()`, evaluation, and save — and explain what each stage does.
2. Read a training loss curve and diagnose the three failure modes (loss not decreasing, loss exploding/NaN, loss plateauing), mapping each to its probable cause (LR, data, template, numerical instability).
3. Choose the standard `TrainingArguments` levers — learning rate, batch size + gradient accumulation, warmup ratio, LR scheduler, gradient checkpointing, FlashAttention 2/3, optimizer — for a given hardware budget and method.
4. Save, merge, and load a fine-tuned model — adapter-only save, `merge_and_unload` for deployment, `PeftModel.from_pretrained` for adapter loading — and explain when each is appropriate.
5. Wire up logging (Weights & Biases, Trackio, or JSONL) and know what to watch (loss, eval loss, LR, grad norm, throughput) and which anti-patterns (no eval, wrong LR for the method, FP16 instead of BF16) will silently ruin a run.

---

# 11.1 — TRL: The Canonical Post-Training Library

*The library that turned "I need to fine-tune a model" from a week of glue code into a CLI command.*

## What TRL is

TRL — Transformers Reinforcement Learning — is the canonical post-training library for the HuggingFace stack. As of its 1.0 release on March 31, 2026, it is the single library you reach for when you want to steer a transformer model: SFT, DPO, KTO, RLOO, GRPO, reward modeling, and the async RL trainers that power frontier reasoning research.

The numbers that make it the default: 75-plus training methods and trainer variants, roughly three million downloads a month, and a Stability Contract that pins the public API so production pipelines don't break on a minor bump. It ships a production CLI — `trl sft`, `trl dpo`, `trl grpo` — that runs a complete, reproducible training job from a YAML config with no Python glue. For anything beyond a one-liner, the Python API (`SFTTrainer`, `DPOTrainer`, `GRPOTrainer`) gives you the full control surface.

The thesis from FT00 — *fine-tuning steers behavior* — is the design principle of the library. Every trainer in TRL is a steering tool. `SFTTrainer` steers format and instruction-following. `DPOTrainer` steers preference. `GRPOTrainer` steers reasoning toward verifiable rewards. None of them inject knowledge. The library's job is to make the steering reliable, observable, and cheap.

## Why TRL and not raw Transformers

You *can* write a training loop with `transformers`' `Trainer` and a custom data collator. People did, for years. The reason they stopped is that SFT is not just "language modeling on a new dataset." Three things make it different, and TRL handles all three:

1. **Chat templates.** A base model trained on raw text expects raw text. An instruct model expects its chat template applied — `<|im_start|>user\n...<|im_end|>\n<|im_start|>assistant\n...`. TRL applies the model's tokenizer chat template automatically and trains only on the completion (masking the prompt), so you learn to *respond*, not to parrot the user. Get this wrong in a hand-rolled loop and you train the model to imitate the user — the classic FT07 template bug.
2. **PEFT integration.** Attach a `LoraConfig`, pass it to the trainer, and the model is wrapped, the base frozen, the adapter trained, and the adapter saved — without you writing a line of PEFT glue. (Module FT08.)
3. **Packing and sequence handling.** Short sequences waste compute. TRL can *pack* multiple examples into a single sequence (with proper attention masking) to fill the context window, often doubling throughput. Doing this by hand is error-prone.

The cost of using TRL: you must learn its config surface (`SFTConfig`, which extends `transformers.TrainingArguments`). The benefit: you stop debugging glue code and start debugging your data and your hyperparameters — which is where the actual work is.

---

# 11.2 — The SFTTrainer, End to End

*The full pipeline, in the order it actually runs on your machine.*

## The pipeline

The `SFTTrainer` flow is six stages. You will write this loop, or a close variant of it, in every SFT job for the rest of the course:

```mermaid
flowchart TD
  A["1. Load dataset\n(chat-format, train/eval split)"] --> B["2. Load model\n(bf16, attn_implementation)"]
  B --> C["3. Attach PEFT config\n(LoraConfig) — or None for full FT"]
  C --> D["4. Build SFTConfig\n(TrainingArguments levers)"]
  D --> E["5. trainer.train()\n(loss curve, eval every N steps)"]
  E --> F["6. Save / merge / load\n(adapter or full)"]
```

### 1. Dataset loading

TRL accepts a HuggingFace `Dataset` (or a dataset name string it will load for you). For SFT, the expected shape is a `messages` field — a list of `{"role": ..., "content": ...}` turns — which TRL renders through the model's chat template. If you produced data in FT05/FT06 (Pillar 1), this is where it plugs in.

```python
from datasets import load_dataset

dataset = load_dataset("trl-lib/Capybara", split="train")
# Standard split: hold out 5-10% for eval — never train on everything.
split = dataset.train_test_split(test_size=0.05, seed=42)
train_ds, eval_ds = split["train"], split["test"]
```

The eval split is non-negotiable. Without it you are flying blind (see Anti-Patterns). It does not need to be large — a few hundred examples is enough to detect overfitting and to compare checkpoints.

### 2. Model loading

Pass a model ID string and TRL loads it for you, or load it yourself and pass the object. The two knobs that matter for memory and speed: `torch_dtype` (use `bfloat16` on Ampere-or-later GPUs — never `float16` for training unless you enjoy NaNs) and `attn_implementation` (FlashAttention 2 or 3, which is effectively mandatory for speed).

```python
import torch
# TRL loads the model from a string — bf16 + FlashAttention 2 are passed via model_init_kwargs
# (shown explicitly here; the trainer accepts a string model id and these kwargs)
```

### 3. PEFT config attachment

Pass a `peft_config` (a `LoraConfig`) to the trainer and it wraps the model automatically — base frozen, adapter trainable. Pass `None` (omit it) and you get full-parameter fine-tuning. The decision between the two is FT10; the *mechanics* are FT08 and this module.

```python
from peft import LoraConfig

peft_config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type="CAUSAL_LM",
)
```

### 4. SFTConfig (the TrainingArguments levers)

`SFTConfig` extends `transformers.TrainingArguments` and adds SFT-specific fields (packing, `max_length`, completion-only loss masking). This is where you set the levers that determine whether your run converges, diverges, or explodes. Section 11.4 covers them in detail.

```python
from trl import SFTConfig, SFTTrainer

training_args = SFTConfig(
    output_dir="./sft-out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,        # effective batch = 2 * 8 = 16
    learning_rate=2e-4,                   # LoRA needs higher LR than full FT
    num_train_epochs=1,
    warmup_ratio=0.05,
    lr_scheduler_type="cosine",
    bf16=True,                            # use bf16, not fp16
    gradient_checkpointing=True,
    logging_steps=10,
    eval_strategy="steps",
    eval_steps=50,
    save_strategy="steps",
    save_steps=100,
    report_to="wandb",                    # or "trackio", or "none"
    packing=True,
    max_length=2048,
)
```

### 5. trainer.train()

```python
trainer = SFTTrainer(
    model="Qwen/Qwen2.5-3B-Instruct",   # string id, or a loaded model object
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=eval_ds,
    peft_config=peft_config,
)
trainer.train()
```

This is the loop. Forward pass, loss, backward, optimizer step, repeat — with logging every `logging_steps` and evaluation every `eval_steps`. The loss curve is your primary instrument. Read it (11.3).

### 6. Eval and save

```python
metrics = trainer.evaluate()
print(metrics)                      # eval_loss, eval_runtime, etc.
trainer.save_model("./sft-final")   # saves adapter (PEFT) or full model
```

For deployment, merge the adapter into the base (11.5). For experimentation, save the adapter and hot-swap it.

---

# 11.3 — Reading a Training Loss Curve

*The loss curve is the EKG of your training run. Learn to read it before you touch a hyperparameter.*

A healthy SFT loss curve descends from its initial value (often 1.5–2.5 for a base model on a new format), drops steeply for the first few hundred steps, then flattens into a gentle decline. The eval loss tracks it, ideally staying below or near the training loss. Three failure modes account for almost every bad run you will ever see.

## Failure mode 1 — Loss not decreasing

The loss is flat or barely dropping from step one. The model is not learning. Three usual suspects, in order of likelihood:

1. **Learning rate too low.** The most common cause. For full FT, `1e-5`–`5e-5`; for LoRA, `1e-4`–`5e-4` — an order of magnitude higher, because adapters need to learn faster against a frozen base. If you copied a full-FT LR into a LoRA run, the loss will barely move.
2. **Data problem.** The dataset is misformatted, the chat template isn't being applied, or the completion-masking is wrong so the model is training on the prompts. (FT07.) Inspect a tokenized batch — print `input_ids` and the loss mask — before assuming it's the LR.
3. **Template bug.** The tokenizer's chat template doesn't match what the model expects. The model sees garbage formatting and can't latch onto the pattern. This is the FT07 failure; the loss curve looks identical to "LR too low."

## Failure mode 2 — Loss exploding (NaN or spike)

The loss jumps to `inf` or `nan`, or spikes violently after a stable start. Causes:

1. **Learning rate too high.** The optimizer overshoots, weights blow up, gradients explode, loss goes to NaN. Halve the LR and retry.
2. **Numerical instability / FP16 overflow.** This is the classic `fp16` training bug. FP16 has a narrow dynamic range; gradients can overflow to `inf` during backward. **Use BF16 instead** — same memory, wider range, no overflow on Ampere-or-later hardware. If you are on an older GPU that lacks BF16, you are stuck with FP16 plus gradient clipping, and you should plan to upgrade.
3. **EOS or template bug producing degenerate targets.** If the completion is empty or the EOS token is mis-set, the loss target is undefined and you get NaN. Check that every completion ends with the right EOS token and that the loss mask isn't all-zero on any example.

## Failure mode 3 — Loss plateaus

The loss descended, then goes flat. Two interpretations, and the eval loss tells you which:

1. **Converged.** Eval loss is also flat and low. The model has learned what this dataset can teach it. Stop. (Diminishing returns past this point; you are burning compute.)
2. **Stuck.** Eval loss is flat but high, or training loss is flat while eval loss rises (overfitting). The model is not improving. Try: a higher LR with warmup to escape the local minimum, a different LR schedule (cosine with decay), or — most often — more and better data. A plateau on bad data is not a hyperparameter problem; it's a data problem.

## The grad norm signal

Beyond loss, watch the **gradient norm** (logged as `grad_norm` by default). A healthy run has a stable, slowly-decreasing grad norm. A sudden spike in grad norm — even if loss hasn't NaN'd yet — is the early-warning system for instability. If you see grad norm jump 10x between two logging steps, that run is about to explode. Kill it, lower the LR, add gradient clipping (`max_grad_norm`, default 1.0 in TRL — usually fine), and retry.

---

# 11.4 — The TrainingArguments Levers

*The knobs on the front panel. Know what each one does and what the safe range is.*

## Learning rate (LR)

The single most important hyperparameter. Typical ranges:

| Method | LR range | Why |
| --- | --- | --- |
| Full FT | `1e-5` – `5e-5` | Updating all weights; small steps to avoid catastrophic forgetting |
| LoRA / QLoRA | `1e-4` – `5e-4` | Adapters are small and start random; they need faster learning to catch up to the frozen base |

The cardinal error: using a full-FT LR for LoRA (loss barely moves) or a LoRA LR for full FT (loss explodes). FT10's decision determines the method; the method determines the LR band.

## Batch size + gradient accumulation

`effective_batch = per_device_batch_size × gradient_accumulation_steps × world_size`

When VRAM-limited (you usually are), you can't fit a large `per_device_batch_size`. Gradient accumulation simulates a large batch by accumulating gradients across several forward/backward passes before stepping the optimizer. The effective batch — what the optimizer "sees" — is the product. A batch of 2 with 8 accumulation steps on 1 GPU is an effective batch of 16. This is how you get stable optimization on a 24GB card.

The tradeoff: accumulation steps slow wall-clock time (you do more forward/backward per optimizer step) but don't change the math of what the optimizer sees. Use it freely.

## Warmup ratio

Typically `0.03`–`0.1`. Warmup ramps the LR from zero to the target over the first fraction of steps, stabilizing early training when gradients are noisy and the model is far from any good solution. Skip it and you risk an early spike (especially with a high LR). For LoRA with its higher LRs, warmup is near-mandatory.

## LR scheduler

- **cosine** — the common default. Smooth decay following a cosine curve to near-zero by the end. Gentle, reliable.
- **linear** — straight-line decay to zero. Slightly more aggressive at the end.
- **constant_with_warmup** — warmup then flat. Used when you plan to train for a fixed budget and don't want late-training decay (e.g., you'll checkpoint and pick the best).

Cosine is the safe default for SFT. Reach for the others when you have a reason.

## Gradient checkpointing

Trade compute for memory: recompute activations during backward instead of storing them. Cuts activation memory roughly 60–70% for about a 30% slowdown in wall-clock time. This is the FT01 VRAM-math lever — when you can't fit the model, turn this on before you reduce batch size. On by default in many TRL configs for a reason.

## FlashAttention 2/3

Effectively mandatory for speed and memory. FlashAttention computes the exact attention but in tiles that fit in SRAM, avoiding the materialization of the full attention matrix. For long sequences this is a 2–4x speedup and a large memory saving. Set `attn_implementation="flash_attention_2"` when loading the model. If your hardware doesn't support it (older GPUs, some CPU paths), fall back — but treat that as a constraint to work around, not a normal config.

## Optimizer

- **AdamW** — the default. Robust, well-understood. Uses the most memory (two momentum states per parameter).
- **AdamW 8-bit** (bitsandbytes) — quantizes the optimizer states to 8-bit. Saves memory at a small accuracy cost. Good when optimizer state is the bottleneck.
- **Adafactor** — factored second-moment. Low memory, but can be finicky; less commonly used for LLM SFT now.
- **PagedAdamW** / **PagedAdamW 32-bit** — the QLoRA default. Uses CUDA paging to offload optimizer states to CPU, transparently, avoiding OOM spikes. Pairs with 4-bit quantized bases.

---

# 11.5 — Saving, Loading, and Merging

*Three operations, three different purposes. Confuse them and you ship the wrong artifact.*

## Save

`trainer.save_model(path)` saves what was trained. With a `peft_config`, that's the adapter — a few hundred megabytes of LoRA weights, not the multi-gigabyte base. With no PEFT config (full FT), it's the entire model. The distinction matters for disk, for the Hub, and for loading later.

## Merge for deployment

Adapters are great for experimentation (swappable, small) but inference with a separate adapter has overhead. For deployment, merge the adapter into the base and serve a single model:

```python
from peft import PeftModel
from transformers import AutoModelForCausalLM

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B-Instruct", torch_dtype=torch.bfloat16)
model = PeftModel.from_pretrained(base, "./sft-final")
merged = model.merge_and_unload()          # adapter folded into base weights
merged.save_pretrained("./sft-merged")     # a single standalone model
```

After `merge_and_unload`, the result is a plain `transformers` model with no adapter layer — it can be quantized (FT19), served by vLLM/Ollama (FT20), or shipped as a standalone checkpoint. This is the standard path from "I trained a LoRA" to "I have a deployable model."

## Load an adapter

For hot-swapping at inference (one base, many adapters), don't merge — load the adapter on top of the base on the fly:

```python
from peft import PeftModel
model = PeftModel.from_pretrained(base, "./sft-final")
# Use model; swap to a different adapter with model.load_adapter(other_path)
```

This is the swappability property from FT00, realized in code: Layer 2 detaches from Layer 1 without touching it.

---

# 11.6 — Logging: What to Watch

*If it isn't logged, it didn't happen. If you can't see the loss curve, you are not training — you are hoping.*

## Backends

- **Weights & Biases (W&B)** — the long-standing default. Rich dashboards, experiment comparison, team sharing. Set `report_to="wandb"`.
- **Trackio** — HuggingFace's lighter, Spaces-backed tracker. Set `report_to="trackio"`. Good when you want a dashboard without a separate account.
- **JSONL / TensorBoard** — `report_to="tensorboard"` or `"none"` with manual logging. Zero external dependencies; the logs live next to your checkpoints.

## What to log

At minimum, every step you log: **train loss**, **learning rate** (so you can see the schedule), **grad norm** (the early-warning system). Every eval step: **eval loss**. Periodically: **throughput** (tokens/sec or samples/sec) so you can spot a slowdown (e.g., a data collator that's accidentally on CPU).

The single most useful plot is train loss and eval loss on the same chart. The gap between them is your overfitting signal: if train keeps dropping while eval rises, you are memorizing, not steering. Stop at the eval-loss minimum and use that checkpoint.

---

## Anti-Patterns

### No eval (flying blind)

Training without a held-out eval split. You cannot tell convergence from overfitting, you cannot compare checkpoints, and you will ship a memorizing model. The eval split is not optional. Even 200 examples is enough to catch the failure.

### Wrong LR for the method

Full-FT LR (`2e-5`) on a LoRA run — loss barely moves. LoRA LR (`2e-4`) on a full-FT run — loss explodes. The method (FT10) sets the LR band. Mismatching them is the most common "why isn't my model learning" ticket.

### No logging

Running a job blind and checking only at the end. A run that NaN'd at step 50 and a run that converged at step 50 look identical from the final checkpoint. Log loss, LR, grad norm, eval loss — every step.

### Ignoring grad norm spikes

The loss looks fine, but grad norm jumped 10x. You ignore it. Fifty steps later, NaN. Grad norm is the canary; watch it.

### FP16 instead of BF16

Using `fp16=True` on Ampere-or-later hardware that supports BF16. FP16 overflows; BF16 doesn't. Same memory, strictly better for training. If your GPU is pre-Ampere (no BF16), that's a hardware constraint — plan to upgrade, or use FP16 with aggressive gradient clipping and accept the instability.

---

## Key Terms

| Term | Definition |
| --- | --- |
| **TRL** | Transformers Reinforcement Learning; the canonical HuggingFace post-training library. 75+ methods, ~3M downloads/month, ships a Stability Contract and production CLI. v1.0 released March 31, 2026. |
| **SFTTrainer** | TRL's supervised fine-tuning trainer. Handles chat templates, completion masking, PEFT attachment, packing. The workhorse for Layer 3 steering. |
| **SFTConfig** | TRL config class extending `transformers.TrainingArguments`. Holds the LR, batch size, scheduler, eval, packing, and all training levers. |
| **Effective batch** | `per_device_batch_size × gradient_accumulation_steps × world_size`. The batch the optimizer "sees." Simulate large batches when VRAM-limited. |
| **Gradient checkpointing** | Recompute activations in backward instead of storing them. ~60-70% activation memory cut for ~30% slowdown. (FT01.) |
| **FlashAttention 2/3** | Tiled exact attention. ~2-4x speedup and large memory saving on long sequences. Effectively mandatory for serious SFT. |
| **Grad norm** | The norm of the gradient vector. Stable and decreasing = healthy; sudden spike = about to explode. The early-warning system. |
| **merge_and_unload** | Fold a PEFT adapter into the base weights, producing a standalone model for deployment. |
| **PagedAdamW** | Optimizer that uses CUDA paging to offload states to CPU, avoiding OOM. The QLoRA default. |
| **Trackio** | HuggingFace's Spaces-backed experiment tracker. A lighter alternative to W&B. |

---

## Lab Exercise

See `07-lab-spec.md`. "The Full SFT Loop": run a complete SFT job with TRL `SFTTrainer` on MiniCPM3-4B or Qwen2.5-3B (student choice), logging to W&B or Trackio, with held-out eval every N steps. Produce the loss curve, eval curve, and merged model. Consumer-GPU (RTX 4090 / 24GB) or Colab.

---

## References

1. **TRL documentation** — *TRL: Transformer Reinforcement Learning*. https://huggingface.co/docs/trl. The canonical reference for `SFTTrainer`, `SFTConfig`, and all trainers.
2. **HuggingFace SFTTrainer guide** — *Fine-tuning a language model with SFTTrainer*. https://huggingface.co/docs/trl/sft_trainer. End-to-end walkthrough, packing, completion-only loss.
3. **TrainingArguments docs** — *transformers.TrainingArguments*. https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments. The full lever surface (LR, scheduler, checkpointing, logging).
4. **PEFT documentation** — *Parameter-Efficient Fine-Tuning*. https://huggingface.co/docs/peft. `LoraConfig`, `PeftModel.from_pretrained`, `merge_and_unload`.
5. **FlashAttention** — Dao et al., *FlashAttention: Fast and Memory-Efficient Exact Attention*. The tiled-attention paper. The 2/3 implementations are in `transformers` via `attn_implementation`.
6. **Module FT08** — *LoRA / QLoRA*. The adapter mechanics this module trains.
7. **Module FT10** — *Full FT vs PEFT: The Decision*. Determines the method, which determines the LR band.
8. **Module FT01** — *VRAM Math*. The memory budget that gradient checkpointing and batch-size choices are made against.