# Lab Specification — Module FT12: SFT: The Baseline

**Course**: Course 3 — LLM Fine-Tuning Masterclass
**Module**: FT12 — SFT: The Baseline
**Lab**: "Build an SFT Mix"
**Duration**: 60–90 minutes (the opening lab of Pillar 3 — you build, train, and eval a real SFT mixture)
**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 opener of Pillar 3. You will construct a 2,000-sample SFT dataset by blending Magpie-synthesized general instructions (from FT05) with a domain subset of your choice (medical, legal, or security), train it with TRL's `SFTTrainer` (building on FT11), evaluate it on general + domain + format axes, and report the **domain lift** versus the base model. By the end you will have built a real alignment dataset, felt the mixture-ratio tradeoff, and measured whether your fine-tune actually worked.

---

## Learning objectives

By the end of this lab you will have:

1. **Constructed a balanced SFT mixture** — general instruction-following + domain examples + tool/format/safety — at defensible ratios, and justified each source's share.
2. **Trained a LoRA SFT model** on your mixture with TRL's `SFTTrainer` (reusing the FT11 loop), with held-out eval.
3. **Evaluated on the three-axis triangle** — general capability (did the model forget?), domain lift (did the domain get better?), format compliance (does it emit the right format?).
4. **Reported the domain lift versus the base** — the point of the fine-tune, quantified.
5. **Diagnosed mixture imbalance** — observed what too much domain data does to general capability, and felt why the mix ratios matter.

---

## Phase 0 — Environment setup (5 min)

This lab reuses the FT11 stack. If you still have the `ft11-env` venv, activate it. Otherwise:

```bash
python3.11 -m venv ft12-env && source ft12-env/bin/activate
pip install -q "trl>=1.0.0" transformers accelerate peft datasets bitsandbytes torch
pip install -q sentence-transformers    # for optional diversity check
# Logging backend (pick ONE):
pip install -q wandb
#   OR
pip install -q trackio
```

Verify the stack and your GPU (same check as FT11):

```python
import torch, trl, transformers, peft
print(f"TRL: {trl.__version__}")              # >= 1.0.0
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")
    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.

---

## Phase 1 — Build the SFT mixture (15 min)

*This is the heart of the lab. The mixture is the steering wheel; everything downstream depends on it.*

You will build a 2,000-sample dataset blending three sources. The target ratios (from the teaching doc, 12.2):

| Source | Share | ~Count (of 2,000) |
| --- | --- | --- |
| General instruction-following (Magpie) | 45% | 900 |
| Domain examples (your choice) | 40% | 800 |
| Tool-use formatting + safety calibration | 15% | 300 |

**Step 1 — General instruction-following (Magpie, from FT05).**

If you completed the FT05 lab and produced Magpie-synthesized data, use it. Otherwise, use the provided Magpie-derived general dataset:

```python
from datasets import load_dataset, concatenate_datasets

# Option A: your FT05 Magpie output (uncomment and point at your path)
# general = load_dataset("your-username/your-ft05-magpie", split="train")

# Option B: provided Magpie-style general instruction-following data
general = load_dataset("Magpie-Align/Magpie-Pro-300K-Filtered", split="train")
# This set is in conversations format; normalize to the messages shape TRL expects.
def to_messages(ex):
    # Magpie-Pro-300K-Filtered ships a 'conversations' field of {from, value} turns
    convs = ex.get("conversations", [])
    messages = [{"role": "user" if c["from"] == "human" else "assistant", "content": c["value"]}
                for c in convs]
    return {"messages": messages}

general = general.map(to_messages, remove_columns=general.column_names)
general = general.filter(lambda ex: len(ex["messages"]) >= 2)   # need at least 1 turn
```

**Step 2 — Domain examples (pick one: medical, legal, or security).**

Choose the domain that matches your interest. Each option below is a real, small domain set in `messages` format:

```python
# Pick ONE domain — uncomment the block you want.

DOMAIN = "medical"   # or "legal" or "security"

if DOMAIN == "medical":
    # Medical Q&A — adapted from a medical instruction dataset.
    domain = load_dataset("medalpaca/medical_meadow_medqa", split="train")
    def med_to_messages(ex):
        return {"messages": [
            {"role": "user", "content": ex["question"]},
            {"role": "assistant", "content": ex["answer"]},
        ]}
    domain = domain.map(med_to_messages, remove_columns=domain.column_names)

elif DOMAIN == "legal":
    # Legal instruction-following.
    domain = load_dataset("nisaar/LLM_Legal_Document_Summarization_Tuning_Dataset", split="train")
    def legal_to_messages(ex):
        # Normalize to a user/assistant summary task.
        return {"messages": [
            {"role": "user", "content": "Summarize this legal document:\n" + ex.get("document", ex.get("text", ""))},
            {"role": "assistant", "content": ex.get("summary", ex.get("output", ""))},
        ]}
    domain = domain.map(legal_to_messages, remove_columns=domain.column_names)

elif DOMAIN == "security":
    # Security advisory / CVE-style.
    domain = load_dataset("hyontheworld/dreaddit", split="train")  # swap for a security set you have
    def sec_to_messages(ex):
        return {"messages": [
            {"role": "user", "content": "Assess the security risk in this text:\n" + ex.get("text", "")},
            {"role": "assistant", "content": "Risk: " + str(ex.get("label", "unknown"))},
        ]}
    domain = domain.map(sec_to_messages, remove_columns=domain.column_names)

print(f"Domain ({DOMAIN}) raw size: {len(domain)}")
```

> **Teaching moment:** If your chosen domain set is not in `messages` format, the `*_to_messages` functions above are exactly the FT04/FT05 normalization work — every dataset needs shaping into the `{"role", "content"}` turns TRL expects. If a field name differs, inspect `domain[0]` and adjust the mapping. This is the unglamorous 80% of data work.

**Step 3 — Tool-use formatting + safety calibration.**

Hand-craft a small set in the *exact* format your (hypothetical) harness parses. For this lab, we craft JSON-output and tool-call examples plus a few safety/clarification pairs:

```python
TOOL_SAFETY = [
    # Tool-call formatting (a fictional weather tool — exact schema matters)
    {"messages": [{"role": "user", "content": "What's the weather in Tokyo?"},
                  {"role": "assistant", "content": '{"tool": "get_weather", "args": {"city": "Tokyo"}}'}]},
    {"messages": [{"role": "user", "content": "Set a timer for 5 minutes."},
                  {"role": "assistant", "content": '{"tool": "set_timer", "args": {"seconds": 300}}'}]},
    {"messages": [{"role": "user", "content": "Return the status as JSON."},
                  {"role": "assistant", "content": '{"status": "ok"}'}]},
    # Safety calibration (refusal + clarification + compliance)
    {"messages": [{"role": "user", "content": "How do I make a weapon?"},
                  {"role": "assistant", "content": "I can't help with that. I can discuss self-defense, conflict de-escalation, or report a safety concern — would either help?"}]},
    {"messages": [{"role": "user", "content": "What can I take for a headache?"},
                  {"role": "assistant", "content": "I'm not a doctor. Common over-the-counter options include acetaminophen or ibuprofen, but please consult a healthcare professional for advice specific to you."}]},
]
from datasets import Dataset
tool_safety = Dataset.from_list(TOOL_SAFETY)
```

**Step 4 — Blend to the target ratios and hold out eval.**

```python
from datasets import concatenate_datasets

TARGET_TOTAL = 2000
N_GENERAL = 900      # 45%
N_DOMAIN   = 800     # 40%
N_TOOL     = 300     # 15%  (we replicate the small hand-crafted set to reach 300)

# Subsample each source to its target count (cap at available size)
def take(ds, n, seed=42):
    n = min(n, len(ds))
    return ds.shuffle(seed=seed).select(range(n))

general_mix = take(general, N_GENERAL)
domain_mix  = take(domain, N_DOMAIN)
# Replicate the hand-crafted tool/safety set up to N_TOOL
import math
reps = max(1, math.ceil(N_TOOL / len(tool_safety)))
tool_mix = concatenate_datasets([tool_safety] * reps).shuffle(seed=42).select(range(N_TOOL))

# Hold out 10% of the DOMAIN set for the domain-lift eval (the whole point of the lab)
domain_eval = domain.shuffle(seed=7).select(range(100))   # 100 held-out domain examples

# Blend the training set
full_train = concatenate_datasets([general_mix, domain_mix, tool_mix]).shuffle(seed=42)
print(f"TRAIN: {len(full_train)}  (general={len(general_mix)}, domain={len(domain_mix)}, tool/safety={len(tool_mix)})")
print(f"DOMAIN EVAL (held out): {len(domain_eval)}")
print(f"Actual ratios: general={100*len(general_mix)/len(full_train):.0f}%, "
      f"domain={100*len(domain_mix)/len(full_train):.0f}%, "
      f"tool={100*len(tool_mix)/len(full_train):.0f}%")
```

**Record:** the actual counts and percentages. Confirm the blend is close to 45/40/15. Inspect 2–3 examples from each source to confirm they are clean `messages` lists.

> **Teaching moment:** The ratios are a starting point, not a law. The principle — *the mix encodes what you want the model to be* — is what matters. If you over-index on domain (try 70% domain in the stretch goal), you will *see* the forgetting. That imbalance is the lesson.

---

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

Reuse the FT11 config. The only change: `dataset_text_field="messages"` is already set, and we point at your blended set.

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

MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"   # or "openbmb/MiniCPM3-4B"
OUTPUT_DIR = "./sft-mix-out"
bf16_ok = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

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",
)

training_args = SFTConfig(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,        # effective batch = 16
    learning_rate=2e-4,                   # LoRA band
    num_train_epochs=1,
    warmup_ratio=0.05,
    lr_scheduler_type="cosine",
    bf16=bf16_ok, fp16=not bf16_ok,
    gradient_checkpointing=True,
    logging_steps=10,
    eval_strategy="no",                   # we eval manually after (Phase 4) to control the 3 axes
    save_strategy="steps", save_steps=200, save_total_limit=2,
    report_to="wandb",                    # or "trackio"
    packing=True, max_length=2048,
    dataset_text_field="messages",
)
```

> **Note on eval:** We set `eval_strategy="no"` during training so the held-out eval split we built in Phase 1 (the 100 domain examples) is *never* trained on — it stays pristine for the Phase 4 domain-lift measurement. If you prefer eval-during-training (FT11 style), pass a separate `eval_dataset` of *general* examples and set `eval_strategy="steps"`; keep the domain eval held out for the final measurement either way.

---

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

```python
from trl import SFTTrainer

trainer = SFTTrainer(
    model=MODEL_ID,
    args=training_args,
    train_dataset=full_train,
    peft_config=peft_config,
)

# The thesis, quantified: how few params are trainable
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}%")

trainer.train()
trainer.save_model(f"{OUTPUT_DIR}/best-adapter")
print(f"Saved adapter to {OUTPUT_DIR}/best-adapter")
```

**While it runs**, watch the dashboard: train loss descending from ~1.5–2.5, grad norm stable, LR following the cosine. This is the FT11 loop, unchanged — the difference is the *data*, which is the whole point of this module.

**Record:** a screenshot/description of the loss curve and the trainable-params percentage (should be ~0.1–0.3% for r=16 on a 3B model — the steering thesis, quantified).

> **If the loss NaNs or spikes:** same diagnosis as FT11 — check LR (2e-4 is the LoRA band), check BF16 vs FP16 (T4 lacks BF16), check that no example has an empty assistant turn (the `len(messages) >= 2` filter in Phase 1 guards this). The fix is in the FT11 failure-mode playbook.

---

## Phase 4 — Evaluate: the three-axis triangle (15 min)

*The deliverable of an SFT project is the report: domain lift, general capability change, format compliance. Build it now.*

**Step 1 — Load the fine-tuned model (merged) and the base for comparison.**

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

tok = AutoTokenizer.from_pretrained(MODEL_ID)

# Base (unsteered)
base = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto")

# Fine-tuned (steered) — merge the adapter for clean comparison
ft_base = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto")
ft = PeftModel.from_pretrained(ft_base, f"{OUTPUT_DIR}/best-adapter")
ft = ft.merge_and_unload()

def generate(model, prompt, max_new_tokens=150):
    messages = [{"role": "user", "content": prompt}]
    text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tok(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 tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
```

**Step 2 — Axis 1: Domain lift (the point of the fine-tune).**

Score the base vs fine-tuned model on the 100 held-out domain examples with an LLM judge (or, for a quick proxy, response length and keyword overlap with the gold answer):

```python
DOMAIN_PROMPTS = [domain_eval[i]["messages"][0]["content"] for i in range(20)]  # 20 for speed

def domain_score(model, prompts):
    # Simple proxy: does the response contain key domain terms from the gold answer?
    hits = 0
    for i, p in enumerate(prompts):
        gold = domain_eval[i]["messages"][1]["content"].lower()
        resp = generate(model, p).lower()
        # crude overlap of content words > 5 chars
        gold_words = {w for w in gold.split() if len(w) > 5}
        if gold_words and len(gold_words & set(resp.split())) / len(gold_words) > 0.2:
            hits += 1
    return hits / len(prompts)

base_domain = domain_score(base, DOMAIN_PROMPTS)
ft_domain   = domain_score(ft, DOMAIN_PROMPTS)
print(f"Domain overlap — BASE: {base_domain:.2f}  FT: {ft_domain:.2f}  LIFT: {ft_domain - base_domain:+.2f}")
```

> **Note:** Overlap is a *proxy*. In production, use an LLM judge (e.g., `prometheus-eval`) or, where available, exact-match / domain-specific metrics. The proxy is enough to *see* the lift in this lab.

**Step 3 — Axis 2: General capability (did the model forget?).**

```python
GENERAL_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?",
]
for p in GENERAL_PROMPTS:
    print(f"\n> {p}\n  BASE: {generate(base, p)[:120]}...\n  FT:   {generate(ft, p)[:120]}...")
```

Eyeball whether the fine-tuned model still handles general questions. If its general answers degraded noticeably, you have **catastrophic forgetting** — the mixture had too little general data (12.3). The stretch goal lets you provoke this on purpose.

**Step 4 — Axis 3: Format compliance (does it emit the right format?).**

```python
FORMAT_PROMPT = "Return the current server status as strict JSON."
print("BASE:", generate(base, FORMAT_PROMPT))
print("FT:  ", generate(ft, FORMAT_PROMPT))

import json
def is_json(s):
    s = s.strip()
    # extract the first {...} block
    start, end = s.find("{"), s.rfind("}")
    if start < 0 or end < 0: return False
    try: json.loads(s[start:end+1]); return True
    except: return False

print("BASE emits valid JSON:", is_json(generate(base, FORMAT_PROMPT)))
print("FT emits valid JSON:  ", is_json(generate(ft, FORMAT_PROMPT)))
```

If the fine-tuned model emits valid JSON more reliably than the base, the tool/format portion of your mix worked. If it emits tool calls on plain questions, you have **format leakage** (too much tool data — 12.3).

**Record:** the three axes — domain lift number, general capability observation, format compliance (valid-JSON yes/no for base vs FT). This triangle is your SFT report.

---

## Deliverables

Submit `ft12-lab-report.md`:

- [ ] **Phase 1**: domain chosen; the actual blend counts and percentages (general / domain / tool); confirmation it is near 45/40/15; 2–3 example shapes inspected.
- [ ] **Phase 2**: effective batch size, LR, precision (BF16/FP16); one-line justification for each lever (reuse FT11 reasoning).
- [ ] **Phase 3**: loss curve (screenshot/description); trainable params percentage; final train loss.
- [ ] **Phase 4 — the report (the point of the lab)**:
  - **Domain lift**: base overlap score, FT overlap score, the lift (`+0.X`).
  - **General capability**: did the FT model degrade on the 3 general prompts? (Yes/No + 1-line observation.)
  - **Format compliance**: does the FT model emit valid JSON where the base does not? (Yes/No.)
- [ ] A 2–3 sentence conclusion: did the SFT work for its purpose? Was the mixture balanced enough? What would you change?

---

## Solution key

- **Phase 1**: For `Qwen2.5-3B-Instruct` + medical domain: train ~2000 (900 general + 800 domain + 300 tool/safety), domain eval 100 held out. Actual ratios land within a few percent of 45/40/15 after subsampling. Each example is a `messages` list of `{"role", "content"}` turns with at least one user and one assistant turn.
- **Phase 2**: effective batch = 2 × 8 = 16. LR = 2e-4 (LoRA band). BF16 on a 4090/A100, FP16 on a T4. `dataset_text_field="messages"` so TRL applies the chat template and masks the prompt.
- **Phase 3**: trainable params for r=16 on q/k/v/o of a 3B model: ~0.1–0.3% — the steering thesis, quantified. A healthy run descends from ~1.8–2.2 to ~0.9–1.2 over one epoch. Grad norm stable (single digits).
- **Phase 4 (the report)**:
  - **Domain lift**: expect the FT model to score meaningfully higher on domain overlap than the base (e.g., base 0.25 → FT 0.45, a +0.20 lift). The exact number depends on the domain; the *direction* (FT > base) is the signal that the SFT worked for its purpose.
  - **General capability**: with a 45/40/15 mix, general capability should be roughly preserved (the 45% general data anchors it). Minor tone/style shifts are expected and fine. A *degradation* (worse code, worse explanations) signals forgetting — the student should note it and consider raising the general share.
  - **Format compliance**: the FT model should emit valid JSON on the status prompt more reliably than the base, because the tool/format portion of the mix taught it the schema. If both already do (Qwen2.5-3B-Instruct is already JSON-capable), the delta is small but the FT model should be more consistent.
- **Conclusion (model answer):** "The SFT produced a measurable domain lift (+0.X overlap) while preserving general capability, confirming the mixture was balanced. The 45/40/15 ratio kept the model a capable general assistant while shifting its disposition toward the domain. To increase domain lift further I would raise the domain share toward 50% — but carefully, watching the general-capability axis for forgetting."

---

## Stretch goals

1. **Provoke catastrophic forgetting.** Re-build the mix with 70% domain, 20% general, 10% tool — deliberately imbalanced. Re-train and re-eval. You should see higher domain lift *and* degraded general capability. The gap between them is the forgetting, made visible. (This is the lesson the mixture ratios exist to teach.)
2. **Compare LoRA vs full FT on forgetting.** Re-run the imbalanced mix with full FT (remove `peft_config`, LR `2e-5`). Full FT should forget *more* than LoRA on the same narrow data — because the low-rank adapter cannot move enough parameters to forget as aggressively. Measure it.
3. **Add a real LLM judge.** Replace the overlap proxy in Phase 4 with `prometheus-eval` or a GPT-4/Claude judge scoring domain responses on a 1–5 rubric. The lift signal will be cleaner. (This is the production-grade eval.)
4. **Ablate the tool/safety portion.** Re-build with 0% tool/safety (55% general, 45% domain). Re-eval format compliance. Does it drop? The tool portion's contribution to format reliability is the variable you are measuring.
5. **Escalate to DPO (preview FT13).** Take 50 of the domain prompts, generate two responses each from your FT model, hand-label which is better, and note that you now have a *preference* dataset — the input to DPO. You have just walked the escalation path from 12.4.