# Lab Specification — Module FT11: The Training Loop with TRL

**Course**: Course 3 — LLM Fine-Tuning Masterclass
**Module**: FT11 — The Training Loop with TRL
**Lab**: "The Full SFT Loop"
**Duration**: 60–90 minutes (the capstone lab of Pillar 2 — you run a real SFT job end-to-end)
**Environment**: Python 3.11+. A consumer NVIDIA GPU (RTX 4090 / 24GB recommended; RTX 3090 24GB works) OR free Google Colab T4/A100. ~20GB free disk for the model + checkpoints.

> This is the hands-on capstone of Pillar 2. You will run a complete, instrumented SFT job with TRL's `SFTTrainer`, log it to W&B or Trackio, evaluate on a held-out split every N steps, and produce a merged, deployable model. By the end you will have a loss curve, an eval curve, and a model you can serve.

---

## Learning objectives

By the end of this lab you will have:

1. **Run a complete SFT job** with TRL `SFTTrainer` — dataset, model, PEFT config, `SFTConfig`, `trainer.train()`, eval, save — and explain what each stage did.
2. **Read the resulting loss curve** and classify it as healthy descent, plateau, or failure mode (and explain the cause if it failed).
3. **Evaluated on a held-out split** during training and identified the best checkpoint (lowest eval loss).
4. **Merged the LoRA adapter** into the base and saved a standalone deployable model.
5. **Generated samples before and after** fine-tuning to see the steering effect with your own eyes.

---

## Phase 0 — Environment setup (5 min)

```bash
# Create a clean venv
python3.11 -m venv ft11-env && source ft11-env/bin/activate

# The full TRL stack. TRL 1.0+ pulls in transformers, accelerate, peft, datasets.
pip install -q "trl>=1.0.0" transformers accelerate peft datasets bitsandbytes torch
# Logging backend — pick ONE:
pip install -q wandb        # Weights & Biases
#   OR
pip install -q trackio      # HuggingFace's lighter, Spaces-backed tracker
```

Verify the stack and your GPU:

```python
import torch, trl, transformers, peft
print(f"PyTorch: {torch.__version__}")
print(f"TRL: {trl.__version__}")              # must be >= 1.0.0
print(f"Transformers: {transformers.__version__}")
print(f"PEFT: {peft.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    # BF16 support check — Ampere (RTX 30xx+) and later
    print(f"BF16 supported: {torch.cuda.is_bf16_supported()}")
```

You need CUDA (Colab T4 or a consumer NVIDIA GPU). Apple Silicon MPS works for inference but TRL training on MPS is not well-supported — use Colab if you lack an NVIDIA GPU. BF16 support requires Ampere or later (RTX 30xx, A100, etc.); the Colab T4 lacks BF16 and the script falls back to FP16 with clipping in that case.

---

## Phase 1 — Choose your model and data (10 min)

**Model (pick one):**

- `Qwen/Qwen2.5-3B-Instruct` — recommended; strong base, clean chat template, fits 24GB with QLoRA.
- `openbmb/MiniCPM3-4B` — open-data alternative (auditable, per FT02); slightly larger.

**Data (pick one):**

- **Use your FT05/FT06 output.** If you completed the Pillar 1 labs and produced a chat-format dataset (a `messages` column of role/content turns), point `load_dataset` at it. This is the preferred path — you steer on data you built.
- **Use the provided dataset.** `trl-lib/Capybara` — a multi-turn conversational dataset in the exact `messages` format TRL expects. Good for a first run.

```python
from datasets import load_dataset

# Option A: your Pillar 1 data (uncomment and point at your path/Hub repo)
# dataset = load_dataset("your-username/your-ft06-dataset", split="train")

# Option B: provided dataset
dataset = load_dataset("trl-lib/Capybara", split="train")

# Hold out 5% for eval — NEVER train on everything (flying blind is an anti-pattern).
split = dataset.train_test_split(test_size=0.05, seed=42)
train_ds, eval_ds = split["train"], split["test"]
print(f"Train: {len(train_ds)}  Eval: {len(eval_ds)}")
print("\nSample example:")
print(train_ds[0])   # confirm it has a 'messages' field with role/content turns
```

**Record**: the train/eval sizes and the shape of one example. Confirm the `messages` field exists and each turn has `role` and `content`.

> **Teaching moment:** This is the steering wheel (FT00). The dataset is what points the model. TRL renders each example through the model's chat template and trains only on the assistant's completion — so the model learns to *respond*, not to parrot the user. If your data lacks the `messages` field, TRL also accepts a plain `text` field, but you lose completion-only masking (you'd train on the prompts too).

---

## Phase 2 — Configure PEFT and SFTConfig (10 min)

This is where the FT10 decision (LoRA vs full FT) becomes concrete. We use LoRA — it fits a 3–4B model on a 24GB card with headroom, and the LR band (`1e-4`–`5e-4`) is the one this lab teaches.

```python
from peft import LoraConfig
from trl import SFTConfig

# --- PEFT config (FT08): the swappable steer ---
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,                # common rule of thumb: alpha = 2 * r
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type="CAUSAL_LM",
)

# --- SFTConfig: the TrainingArguments levers (11.4) ---
MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"   # or "openbmb/MiniCPM3-4B"
OUTPUT_DIR = "./sft-out"

import torch
bf16_ok = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

training_args = SFTConfig(
    output_dir=OUTPUT_DIR,
    # --- Batch + accumulation: effective batch = 2 * 8 = 16 ---
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    # --- LR: LoRA band (1e-4 to 5e-4). Full FT would be 1e-5 to 5e-5. ---
    learning_rate=2e-4,
    num_train_epochs=1,
    # --- Warmup + scheduler ---
    warmup_ratio=0.05,                    # 0.03-0.1 range
    lr_scheduler_type="cosine",           # the common default
    # --- Precision: BF16 if supported, else FP16 + clipping ---
    bf16=bf16_ok,
    fp16=not bf16_ok,
    # --- Memory + speed ---
    gradient_checkpointing=True,          # ~60-70% activation mem cut, ~30% slower
    # --- Logging + eval ---
    logging_steps=10,
    eval_strategy="steps",
    eval_steps=50,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=3,                   # keep last 3 checkpoints
    load_best_model_at_end=True,          # ship the best (lowest eval loss)
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    # --- Experiment tracking: pick ONE ---
    report_to="wandb",                    # or "trackio", or "tensorboard"
    # --- SFT specifics ---
    packing=True,                         # pack short examples → ~2x throughput
    max_length=2048,
    dataset_text_field="messages",        # TRL applies the chat template
)
```

**Record**: your effective batch size (`per_device_train_batch_size × gradient_accumulation_steps`), your LR, and whether you're on BF16 or FP16. If you're on FP16 (Colab T4), note that you may see instability — the lab is designed to complete anyway, but watch the grad norm.

> **Teaching moment:** Every lever here answers a question from the teaching doc. `gradient_checkpointing=True` because we're VRAM-limited (FT01). `bf16=True` because FP16 overflows (11.3, failure mode 2). `learning_rate=2e-4` because LoRA needs the higher band (11.4). `load_best_model_at_end=True` so we ship the lowest-eval-loss checkpoint, not the last one. None of these are arbitrary.

---

## Phase 3 — Run the training loop (20–40 min)

```python
from trl import SFTTrainer

# Log in to your tracking backend (do this once in your shell, not in the script):
#   wandb:  export WANDB_API_KEY=...   (or run `wandb login`)
#   trackio: no login needed; logs to a HF Space

trainer = SFTTrainer(
    model=MODEL_ID,                # string id — TRL loads it with bf16 + FlashAttn
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=eval_ds,
    peft_config=peft_config,
)

# Count how few parameters are trainable (the thesis, quantified)
trainable = sum(p.numel() for p in trainer.model.parameters() if p.requires_grad)
total = sum(p.numel() for p in trainer.model.parameters())
print(f"Trainable params: {trainable:,} / {total:,} = {100*trainable/total:.3f}%")

# THE LOOP
trainer.train()
```

**While it runs**, open your tracking dashboard (W&B run URL is printed at startup; Trackio prints a Spaces URL). Watch:

- **train_loss** — should descend from ~1.5–2.5 and flatten.
- **eval_loss** (every 50 steps) — should track train_loss; if it rises while train drops, you're overfitting (stop at the minimum).
- **learning_rate** — should follow the cosine schedule.
- **grad_norm** — should be stable; a 10x spike is the early warning for explosion.

**Record**: a screenshot or text description of the loss curve (train + eval on the same chart). Note the final train loss, the best eval loss, and the step it occurred at.

> **If the loss NaNs or spikes:** you hit failure mode 2. Kill the run. Check: (a) are you on FP16 on a T4? Try `per_device_train_batch_size=1` and `gradient_accumulation_steps=16`. (b) Is the LR too high? Drop to `1e-4`. (c) Inspect a tokenized batch for empty completions (a template/EOS bug). This diagnosis IS the lab — the failure teaches you to read the curve.

---

## Phase 4 — Evaluate and save the best model (5 min)

```python
# Final eval on the held-out split
metrics = trainer.evaluate()
print("=== FINAL EVAL METRICS ===")
print(metrics)
# e.g. {'eval_loss': 1.12, 'eval_runtime': 45.2, 'eval_samples_per_second': 8.8}

# Save the best model (adapter only — PEFT)
trainer.save_model(f"{OUTPUT_DIR}/best-adapter")
print(f"Saved adapter to {OUTPUT_DIR}/best-adapter")
```

Because we set `load_best_model_at_end=True`, the model in memory is already the best (lowest eval loss) checkpoint. `save_model` writes the adapter — a few hundred MB, not the multi-GB base.

---

## Phase 5 — Merge the adapter and save a standalone model (5 min)

For deployment, merge the adapter into the base. The result is a single `transformers` model with no adapter layer — ready to quantize (FT19) or serve (FT20).

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# Reload the base fresh, in BF16
base = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto"
)
# Load the adapter on top
model = PeftModel.from_pretrained(base, f"{OUTPUT_DIR}/best-adapter")

# MERGE — adapter folded into base weights
merged = model.merge_and_unload()
print("Merged. The adapter is now baked into the base weights.")

# Save the standalone model
MERGED_DIR = f"{OUTPUT_DIR}/merged-model"
merged.save_pretrained(MERGED_DIR)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.save_pretrained(MERGED_DIR)
print(f"Saved standalone merged model to {MERGED_DIR}")
```

**Record**: the size of the merged model directory vs the adapter directory. The adapter is a few hundred MB; the merged model is the full base size (several GB). This is the tradeoff: small + swappable (adapter) vs. large + standalone (merged).

---

## Phase 6 — Generate before/after samples (10 min)

The steering effect, visible. Generate from the base and from your fine-tuned model on the same prompts.

```python
def generate(model, tokenizer, prompt, max_new_tokens=150):
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    return tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

PROMPTS = [
    "Explain what a hash function is, in two sentences.",
    "Write a Python function that returns the n-th Fibonacci number.",
    "What are three signs of a phishing email?",
]

print("=" * 60)
print("BEFORE — base model (unsteered)")
print("=" * 60)
for p in PROMPTS:
    print(f"\n> {p}")
    print(generate(base, AutoTokenizer.from_pretrained(MODEL_ID), p))

print("\n" + "=" * 60)
print("AFTER — fine-tuned model (steered)")
print("=" * 60)
for p in PROMPTS:
    print(f"\n> {p}")
    print(generate(merged, tokenizer, p))
```

**Record**: the before/after outputs side by side. Depending on your dataset, the differences may be subtle (format, tone, conciseness) or pronounced (Capybara is multi-turn conversational, so expect more conversational, engaged responses). The point is to *see* that steering happened — the base moved.

---

## Deliverables

Submit `ft11-lab-report.md`:

- [ ] **Phase 1**: model chosen, dataset chosen, train/eval sizes, one example's shape.
- [ ] **Phase 2**: effective batch size, LR, precision (BF16/FP16), and a one-line justification for each lever.
- [ ] **Phase 3**: loss curve screenshot or description (train + eval on same chart); final train loss; best eval loss and the step it occurred; grad norm behavior (stable or spiked?).
- [ ] **Phase 4**: final eval metrics (`eval_loss`).
- [ ] **Phase 5**: adapter dir size vs merged model dir size.
- [ ] **Phase 6**: before/after generations for the three prompts, with a 1–2 sentence observation of what changed.

---

## Solution key

- **Phase 1**: For `Qwen2.5-3B-Instruct` on Capybara: train ~6600 / eval ~350 examples (after the 5% split). Each example has a `messages` list of `{"role": "user"/"assistant", "content": "..."}` turns.
- **Phase 2**: effective batch = 2 × 8 = 16. LR = 2e-4 (LoRA band). On a 4090 or A100: BF16. On a Colab T4: FP16 (no BF16 support). Trainable params for r=16 on q/k/v/o of a 3B model: ~0.1–0.3% of total — consistent with the FT00 thesis.
- **Phase 3**: A healthy run on Capybara descends from ~1.8–2.2 to ~0.9–1.2 over one epoch, with eval loss tracking slightly above train loss. Best eval loss typically appears in the last quarter of training (the cosine schedule has decayed). Grad norm should be stable (single digits, slowly decreasing). If it NaN'd: the student is likely on FP16 on a T4 — guide them to reduce batch size or drop the LR. The diagnosis is the lesson.
- **Phase 4**: `eval_loss` in the 0.9–1.3 range is healthy for this setup. Much higher (>2.0) suggests the run didn't converge (check LR, data).
- **Phase 5**: adapter ~50–150 MB; merged model ~6 GB (3B params × 2 bytes in BF16). The size difference is the entire point of PEFT.
- **Phase 6**: Capybara-trained model produces more conversational, multi-turn-aware responses (asks clarifying questions, engages). Base Qwen2.5-3B-Instruct is already conversational, so the delta is subtle — look for tone and engagement shifts. The before/after is more dramatic if the student used their own FT05/FT06 data with a strong format signal.

---

## Stretch goals

1. **Try a different LR.** Re-run with `learning_rate=1e-4` (low end of LoRA band) and `5e-4` (high end). Compare the loss curves. Which converged faster? Which had a lower best eval loss? (You are building intuition for the LR lever.)
2. **Try a different scheduler.** Re-run with `lr_scheduler_type="linear"` and `"constant_with_warmup"`. Compare. Does constant-with-warmup overfit more (no late-training decay)?
3. **Try full FT (no PEFT).** Remove the `peft_config` and set `learning_rate=2e-5` (the full-FT band). This needs more VRAM — you may need `per_device_train_batch_size=1`, `gradient_accumulation_steps=16`, and definitely `gradient_checkpointing=True`. Compare the loss curve and final eval loss to the LoRA run. (This is the FT10 decision, felt in compute.)
4. **Turn off packing.** Re-run with `packing=False`. Compare throughput (samples/sec in the logs). You should see packing roughly double it for short-sequence datasets. (You are building intuition for the packing lever.)
5. **Ablate gradient checkpointing.** On a 24GB GPU with a 3B model, you can fit without it. Re-run with `gradient_checkpointing=False` and compare wall-clock time. The ~30% slowdown is real — but so is the memory it saves when you scale up.