Course: Course 3 — LLM Fine-Tuning Masterclass Module: FT15 — CoT Distillation and Rejection-Sampling FT Duration: 30–45 minutes of GPU time (plus ~10 min setup) on a consumer GPU or Colab T4 Environment: Python 3.10+, a CUDA GPU with ≥16 GB VRAM (RTX 4090 / 3090 / A10G / Colab T4 all work). Apple Silicon (M-series, ≥16 GB) works for the smaller base via MPS — see the fallback note in Phase 6.
This is the distillation lab. You will generate chain-of-thought traces from a strong reasoning teacher on 200 GSM8K math problems, verify and curate them, SFT a small student (MiniCPM5-1B or Qwen2.5-1.5B) on those traces via QLoRA, merge, and measure the GSM8K accuracy lift versus the un-distilled base. By the end you will have built a reasoning model with SFT alone — no RL — and felt the FT15 thesis: distillation is the cheapest path that works.
python3 -m venv ft15-env && source ft15-env/bin/activate
pip install -q -U "torch>=2.3" "transformers>=4.46" "peft>=0.13" \
"trl>=0.12" "bitsandbytes>=0.43" "accelerate>=0.34" "datasets>=3.0" \
"openai>=1.50" "vllm>=0.6"
# On Apple Silicon (no CUDA), drop bitsandbytes and vllm — use the MPS fallback (Phase 6).
On Colab: Runtime → Change runtime type → T4 GPU. Then run the same pip install (drop the venv lines).
You need a strong reasoning teacher to generate traces. Two options:
| Teacher | How | Cost | Notes |
|---|---|---|---|
| API (recommended for quality) | OpenAI o1/o3-mini, or Claude/GPT with extended thinking, via the openai SDK or Anthropic SDK |
~$0.50–2 for 200 problems | Best traces. Set OPENAI_API_KEY (or equivalent). |
| Local reasoning model | A local model that emits <think> traces — e.g., deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B served via vLLM |
Free (GPU) | Lower ceiling than a frontier API, but demonstrates the full local pipeline. |
You do not need a Hugging Face token for the default student (Qwen/Qwen2.5-1.5B-Instruct is open) or for openbmb/MiniCPM5-1B.
By the end of this lab you will have:
<think>…</think> tags.The default student is Qwen/Qwen2.5-1.5B-Instruct — open, ChatML, fits a free Colab T4. The default teacher is an API with thinking enabled; a local fallback is a small R1-Distill model.
# config.py — set once, reuse across phases
STUDENT_ID = "Qwen/Qwen2.5-1.5B-Instruct" # the student we will SFT
# Fallbacks: "openbmb/MiniCPM5-1B"
# Teacher choice — set ONE of these
TEACHER_MODE = "api" # "api" (recommended) or "local"
TEACHER_API = "openai" # "openai" (o1/o3-mini) — or adapt for Anthropic
TEACHER_API_MODEL = "o3-mini" # or whatever thinking-enabled model you have access to
TEACHER_LOCAL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" # for TEACHER_MODE="local"
N_PROBLEMS = 200 # GSM8K problems to distill (200 is enough for a measurable lift)
N_EVAL = 200 # GSM8K test samples for the eval
Why a small R1-Distill as a local teacher is legitimate: it is itself a distilled student, but it already emits structured
<think>traces. Distilling from it demonstrates the pipeline end-to-end without an API key. The ceiling is lower than a frontier API (it is a 1.5B teacher), so expect a smaller lift — but the workflow is identical.
We take 200 GSM8K problems, ask the teacher to solve each with an explicit <think> trace, and keep only the traces whose final answer is correct.
Create make_distill_data.py:
# make_distill_data.py — generate & curate CoT traces from a strong teacher
import json, re, os
from datasets import load_dataset
# --- load 200 GSM8K train problems ---
gsm = load_dataset("openai/gsm8k", "main", split="train")
problems = gsm.select(range(200)) # first 200 train problems
def extract_groundtruth(answer_str: str) -> str:
# GSM8K answers end with "#### <number>"
m = re.search(r"####\s*(.+)$", answer_str.strip())
return m.group(1).strip().replace(",", "") if m else answer_str.strip()
def extract_final_answer(text: str) -> str:
# The teacher's final answer — look for \boxed{...} or a trailing "#### "
m = re.search(r"\\boxed\{([^}]+)\}", text)
if m:
return m.group(1).strip().replace(",", "")
m = re.search(r"####\s*(.+)$", text.strip())
if m:
return m.group(1).strip().replace(",", "")
# fallback: last number in the text
nums = re.findall(r"-?\d[\d,]*\.?\d*", text)
return nums[-1].replace(",", "") if nums else ""
# --- teacher: API or local ---
def teacher_solve_api(question: str) -> str:
from openai import OpenAI
client = OpenAI() # uses OPENAI_API_KEY env var
# o3-mini / o1 reasoning models emit a reasoning trace then the answer.
# For non-reasoning models, prompt explicitly for a <think> block.
resp = client.chat.completions.create(
model=os.environ.get("TEACHER_API_MODEL", "o3-mini"),
messages=[
{"role": "system", "content": "Solve the math problem. First reason step by step inside <think>...</think> tags, then give the final answer as \\boxed{N}."},
{"role": "user", "content": question},
],
)
return resp.choices[0].message.content
def teacher_solve_local(question: str, model, tokenizer) -> str:
import torch
messages = [
{"role": "system", "content": "Solve the math problem. First reason step by step inside <think>...</think> tags, then give the final answer as \\boxed{N}."},
{"role": "user", "content": question},
]
inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
with torch.no_grad():
out = model.generate(inputs, max_new_tokens=1024, do_sample=False, temperature=1.0)
return tokenizer.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True)
# --- generate + curate ---
def main():
traces = []
if os.environ.get("TEACHER_MODE", "api") == "local":
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
tid = os.environ.get("TEACHER_LOCAL_ID", "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")
tok = AutoTokenizer.from_pretrained(tid)
mdl = AutoModelForCausalLM.from_pretrained(tid, torch_dtype=torch.bfloat16, device_map="auto")
solve = lambda q: teacher_solve_local(q, mdl, tok)
else:
solve = teacher_solve_api
kept = 0
for i, ex in enumerate(problems):
q, gt = ex["question"], extract_groundtruth(ex["answer"])
try:
teacher_out = solve(q)
except Exception as e:
print(f"[{i}] teacher error: {e}")
continue
pred = extract_final_answer(teacher_out)
if pred == gt: # CURATION: keep only correct traces
traces.append({"question": q, "trace": teacher_out, "answer": gt})
kept += 1
if (i + 1) % 25 == 0:
print(f"[{i+1}/200] kept {kept} correct traces")
with open("distill_traces.jsonl", "w") as f:
for t in traces:
f.write(json.dumps(t) + "\n")
print(f"Wrote {kept} verified traces to distill_traces.jsonl (kept {kept}/{len(problems)})")
if __name__ == "__main__":
main()
Run it:
export TEACHER_MODE=api # or "local"
export OPENAI_API_KEY=sk-... # if using api
python make_distill_data.py
head -1 distill_traces.jsonl | python -m json.tool
Read one trace. You should see a <think>…</think> block followed by a \boxed{N} answer. The kept traces are your distillation corpus — verified correct, the same curation step R1-Distill applied at 800K scale.
If the API is unavailable: set
TEACHER_MODE=localand use a small R1-Distill model as teacher. The lift will be smaller (the teacher's ceiling is lower), but the full pipeline runs locally and free.
The student must learn to emit the trace. We format each example so the target text is the teacher's full trace (including the <think> block and the final answer). Create format_data.py:
# format_data.py — convert traces into SFT (messages) format
import json
SYSTEM = (
"Solve the math problem. First reason step by step inside <think>...</think> tags, "
"then give the final answer as \\boxed{N}."
)
rows = []
with open("distill_traces.jsonl") as f:
for line in f:
t = json.loads(line)
rows.append({
"messages": [
{"role": "system", "content": SYSTEM},
{"role": "user", "content": t["question"]},
{"role": "assistant", "content": t["trace"]}, # the trace IS the target
]
})
with open("distill_sft.jsonl", "w") as f:
for r in rows:
f.write(json.dumps(r) + "\n")
print(f"Wrote {len(rows)} SFT examples to distill_sft.jsonl")
Run it and inspect one example:
python format_data.py
head -1 distill_sft.jsonl | python -m json.tool
FT07 discipline checkpoint. Before training, run one example through
apply_chat_template(..., return_assistant_tokens_mask=True)and decode it. Confirm the ChatML role tokens and EOS are present and the assistant mask covers the trace. You already learned why (FT07) — do not skip it.
The FT08 QLoRA workflow, applied to distillation. Create train.py:
# train.py — QLoRA distillation SFT (Phases 3–4)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
STUDENT_ID = "Qwen/Qwen2.5-1.5B-Instruct"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained(STUDENT_ID)
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
STUDENT_ID,
quantization_config=bnb_config,
device_map="auto",
attn_implementation="flash_attention_2", # or "sdpa" on older GPUs / Colab T4
)
# prepare + attach (the FT08 workflow)
model = prepare_model_for_kbit_training(model)
lora_config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters() # expect <1% trainable
You should see trainable params: ~9.4M || all params: ~1.58B || trainable%: ~0.60%. Under 1% — the LoRA promise, applied to distillation.
Standard TRL SFTTrainer. The chat template is applied automatically because we use the messages format.
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
dataset = load_dataset("json", data_files="distill_sft.jsonl", split="train")
training_args = SFTConfig(
output_dir="./distilled-student",
num_train_epochs=3,
per_device_train_batch_size=2,
gradient_accumulation_steps=4, # effective batch = 8
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_ratio=0.03,
logging_steps=5,
save_strategy="epoch",
bf16=True, # fp16=True on older cards (V100)
gradient_checkpointing=True,
optim="paged_adamw_8bit",
max_seq_length=1024, # traces are longer than pirate-style SFT
report_to="none",
)
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=dataset,
processing_class=tokenizer,
)
trainer.train()
# merge for deployment + eval (Option B from FT08)
merged = model.merge_and_unload()
merged.save_pretrained("./distilled-student/merged", safe_serialization=True)
tokenizer.save_pretrained("./distilled-student/merged")
Watch the loss. It should descend from ~1.2–1.8 to ~0.6–0.9 over the epochs. The traces are longer than the FT08 pirate examples, so max_seq_length=1024 (raise to 2048 if you see truncation warnings). If loss is NaN, that is an FT07 template/EOS bug, not a distillation bug — re-run the inspection loop.
Watch VRAM with nvidia-smi -l 2. Peak ~7–10 GB on a 1.5B QLoRA at 1024 context. On a Colab T4 you have headroom.
Already done at the end of Phase 4 (the merge_and_unload call). The merged model at ./distilled-student/merged is a single self-contained model — load directly with AutoModelForCausalLM.from_pretrained, no PEFT, no 4-bit at serve time. This is what we eval.
Caveat (FT08):
merge_and_unload()on a 4-bit base dequantizes — the merged model is larger than the 4-bit base. For production you re-quantize (GGUF/AWQ) afterward. For this lab, eval the merged FP16 model directly.
The moment of truth. We sample N GSM8K test problems, run both the distilled model and the un-distilled base, and compare accuracy.
# eval.py — measure the GSM8K lift
import torch, re, json
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
STUDENT_ID = "Qwen/Qwen2.5-1.5B-Instruct"
DISTILLED_PATH = "./distilled-student/merged"
N_EVAL = 200
gsm = load_dataset("openai/gsm8k", "main", split="test")
eval_set = gsm.select(range(N_EVAL))
def extract_gt(answer_str):
m = re.search(r"####\s*(.+)$", answer_str.strip())
return m.group(1).strip().replace(",", "") if m else answer_str.strip()
def extract_pred(text):
m = re.search(r"\\boxed\{([^}]+)\}", text)
if m: return m.group(1).strip().replace(",", "")
m = re.search(r"####\s*(.+)$", text.strip())
if m: return m.group(1).strip().replace(",", "")
nums = re.findall(r"-?\d[\d,]*\.?\d*", text)
return nums[-1].replace(",", "") if nums else ""
SYSTEM = ("Solve the math problem. First reason step by step inside <think>...</think> tags, "
"then give the final answer as \\boxed{N}.")
def evaluate(model_id_or_path, label):
tok = AutoTokenizer.from_pretrained(model_id_or_path)
mdl = AutoModelForCausalLM.from_pretrained(model_id_or_path, torch_dtype=torch.bfloat16, device_map="auto")
correct = 0
for i, ex in enumerate(eval_set):
msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": ex["question"]}]
inputs = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(mdl.device)
with torch.no_grad():
out = mdl.generate(inputs, max_new_tokens=512, do_sample=False) # greedy for reproducibility
pred = extract_pred(tok.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True))
if pred == extract_gt(ex["answer"]):
correct += 1
if (i + 1) % 50 == 0:
print(f" [{label}] {i+1}/{N_EVAL} acc so far: {correct/(i+1):.3f}")
acc = correct / N_EVAL
print(f"=== {label}: GSM8K accuracy = {acc:.3f} ({correct}/{N_EVAL}) ===")
return acc
base_acc = evaluate(STUDENT_ID, "UN-DISTILLED BASE")
distilled_acc = evaluate(DISTILLED_PATH, "DISTILLED STUDENT")
print(f"\nLIFT from distillation: {distilled_acc - base_acc:+.3f} "
f"({base_acc:.3f} -> {distilled_acc:.3f})")
Run it:
python eval.py
Expected: a measurable GSM8K accuracy gain — often +0.10 to +0.25 (10–25 percentage points) from 200 traces alone, depending on teacher quality. The base already knows arithmetic (FT00: the capability is latent); distillation steered it to reason reliably and in a structured trace. The lift is your distillation.
Control check: load the un-distilled base with the same system prompt asking for a <think> trace. It will often produce a shorter, less reliable trace — or skip the trace and answer directly, getting more wrong. The distilled student reliably emits the structured trace. The difference between the two accuracies is what 200 curated traces bought you.
Submit ft15-lab-report.md containing:
kept/total curation rate (how many of the 200 traces were verified correct and kept).distill_traces.jsonl — show the <think>…</think> block and the \boxed{N} answer. This is your distillation target.print_trainable_parameters() output — confirm trainable% < 1.nvidia-smi).These are defensible answers, not the only wording.
With a frontier API teacher (o3-mini / equivalent) on 200 GSM8K train problems, expect to keep ~150–190 of 200 (75–95%) — the teacher is strong and GSM8K train is not adversarial. With a local 1.5B R1-Distill teacher, expect ~100–160 of 200 (50–80%) — weaker teacher, more traces rejected. Either is a usable corpus; the lift scales with teacher quality.
Expect start ~1.2–1.8, smooth cosine descent, end ~0.6–0.9 after 3 epochs. Traces are longer than the FT08 pirate SFT, so max_seq_length=1024 is the right floor. If loss plateaus above ~1.1, the rank may be too low or the LR too low; if it crashes below ~0.5, you are overfitting (the corpus is small — raise dropout or reduce epochs to 2). A healthy run lands ~0.7–0.9 and generalizes (Phase 6 confirms).
Peak ~7–10 GB on a 1.5B QLoRA at 1024 context, batch 2 + grad accum 4. The FT08 rule of thumb holds; the longer context adds a bit of activation overhead. On a Colab T4 (16 GB) you have comfortable headroom. If you OOM, lower per_device_train_batch_size to 1 and raise gradient_accumulation_steps to 8.
With an API teacher: expect the un-distilled Qwen2.5-1.5B base to score ~0.30–0.45 on GSM8K (it can do arithmetic but does not reliably emit a structured trace), and the distilled student to score 0.50–0.70 — a lift of +0.10 to +0.25. With a local 1.5B R1-Distill teacher, expect a smaller lift (+0.05 to +0.15) because the teacher's ceiling is lower. Either way, the lift is real and attributable to the traces: the base already knew the arithmetic (FT00); distillation routed it through a reliable reasoning pathway.
SFT on the teacher's traces produced a reasoning lift because distillation is steering, not teaching (FT00). The Qwen2.5-1.5B base already saw arithmetic during pretraining — the capability was latent. The teacher's <think> traces are a steering wheel: they route the student's probability mass through a structured reasoning pathway reliably, so it emits a chain of thought and reaches the correct answer more often. No RL was needed because the teacher already did the expensive search at data-generation time; the student imitates. This is the R1-Distill result in miniature: R1-Distill-Qwen-32B beat o1-mini via SFT on ~800K traces; we got a measurable lift from 200. The ceiling of this approach is the teacher's quality — the student cannot exceed the teacher (imitation). I would escalate to RFT (generate the student's own candidates, filter by math reward, retrain on winners) if the distilled model falls short of the target and I have no stronger teacher; to GRPO only if RFT plateaus and I have the budget and stability tolerance for full RL.
llama.cpp's convert_hf_to_gguf.py, then load in Ollama. You now have a local, served reasoning model distilled from a teacher — the full FT15 → FT19 deploy path.# Lab Specification — Module FT15: CoT Distillation and Rejection-Sampling FT
**Course**: Course 3 — LLM Fine-Tuning Masterclass
**Module**: FT15 — CoT Distillation and Rejection-Sampling FT
**Duration**: 30–45 minutes of GPU time (plus ~10 min setup) on a consumer GPU or Colab T4
**Environment**: Python 3.10+, a CUDA GPU with ≥16 GB VRAM (RTX 4090 / 3090 / A10G / Colab T4 all work). Apple Silicon (M-series, ≥16 GB) works for the smaller base via MPS — see the fallback note in Phase 6.
> **This is the distillation lab.** You will generate chain-of-thought traces from a strong reasoning teacher on 200 GSM8K math problems, verify and curate them, SFT a small student (MiniCPM5-1B or Qwen2.5-1.5B) on those traces via QLoRA, merge, and measure the GSM8K accuracy lift versus the un-distilled base. By the end you will have built a reasoning model with SFT alone — no RL — and felt the FT15 thesis: distillation is the cheapest path that works.
---
## Setup (one time)
```bash
python3 -m venv ft15-env && source ft15-env/bin/activate
pip install -q -U "torch>=2.3" "transformers>=4.46" "peft>=0.13" \
"trl>=0.12" "bitsandbytes>=0.43" "accelerate>=0.34" "datasets>=3.0" \
"openai>=1.50" "vllm>=0.6"
# On Apple Silicon (no CUDA), drop bitsandbytes and vllm — use the MPS fallback (Phase 6).
```
On Colab: Runtime → Change runtime type → T4 GPU. Then run the same `pip install` (drop the venv lines).
You need a **strong reasoning teacher** to generate traces. Two options:
| Teacher | How | Cost | Notes |
| --- | --- | --- | --- |
| **API (recommended for quality)** | OpenAI o1/o3-mini, or Claude/GPT with extended thinking, via the `openai` SDK or Anthropic SDK | ~$0.50–2 for 200 problems | Best traces. Set `OPENAI_API_KEY` (or equivalent). |
| **Local reasoning model** | A local model that emits `<think>` traces — e.g., `deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B` served via vLLM | Free (GPU) | Lower ceiling than a frontier API, but demonstrates the full local pipeline. |
You do **not** need a Hugging Face token for the default student (`Qwen/Qwen2.5-1.5B-Instruct` is open) or for `openbmb/MiniCPM5-1B`.
---
## Learning objectives
By the end of this lab you will have:
1. **Generated a distillation corpus** — CoT traces from a strong teacher on 200 GSM8K problems, with the reasoning in `<think>…</think>` tags.
2. **Curated the traces** — verified the final answer against ground truth, kept only correct traces (the FT15 curation step that R1-Distill used at 800K scale).
3. **SFT'd a small student on the traces** via QLoRA (the FT08 workflow, applied to distillation), producing a merged reasoning model.
4. **Measured the GSM8K lift** — the distilled model vs the un-distilled base — and felt that distillation is steering (the base already knew arithmetic; you routed it through a reliable reasoning trace).
5. **Internalized the FT15 thesis**: you built a reasoning model with SFT alone, no RL — the cheapest path that works at scale.
---
## Phase 0 — Pick your base and teacher (3 min)
The default **student** is **`Qwen/Qwen2.5-1.5B-Instruct`** — open, ChatML, fits a free Colab T4. The default **teacher** is an API with thinking enabled; a local fallback is a small R1-Distill model.
```python
# config.py — set once, reuse across phases
STUDENT_ID = "Qwen/Qwen2.5-1.5B-Instruct" # the student we will SFT
# Fallbacks: "openbmb/MiniCPM5-1B"
# Teacher choice — set ONE of these
TEACHER_MODE = "api" # "api" (recommended) or "local"
TEACHER_API = "openai" # "openai" (o1/o3-mini) — or adapt for Anthropic
TEACHER_API_MODEL = "o3-mini" # or whatever thinking-enabled model you have access to
TEACHER_LOCAL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" # for TEACHER_MODE="local"
N_PROBLEMS = 200 # GSM8K problems to distill (200 is enough for a measurable lift)
N_EVAL = 200 # GSM8K test samples for the eval
```
> **Why a small R1-Distill as a local teacher is legitimate:** it is itself a distilled student, but it already emits structured `<think>` traces. Distilling from it demonstrates the pipeline end-to-end without an API key. The ceiling is lower than a frontier API (it is a 1.5B teacher), so expect a smaller lift — but the workflow is identical.
---
## Phase 1 — Build the distillation corpus (10–15 min)
We take 200 GSM8K problems, ask the teacher to solve each with an explicit `<think>` trace, and keep only the traces whose final answer is correct.
Create `make_distill_data.py`:
```python
# make_distill_data.py — generate & curate CoT traces from a strong teacher
import json, re, os
from datasets import load_dataset
# --- load 200 GSM8K train problems ---
gsm = load_dataset("openai/gsm8k", "main", split="train")
problems = gsm.select(range(200)) # first 200 train problems
def extract_groundtruth(answer_str: str) -> str:
# GSM8K answers end with "#### <number>"
m = re.search(r"####\s*(.+)$", answer_str.strip())
return m.group(1).strip().replace(",", "") if m else answer_str.strip()
def extract_final_answer(text: str) -> str:
# The teacher's final answer — look for \boxed{...} or a trailing "#### "
m = re.search(r"\\boxed\{([^}]+)\}", text)
if m:
return m.group(1).strip().replace(",", "")
m = re.search(r"####\s*(.+)$", text.strip())
if m:
return m.group(1).strip().replace(",", "")
# fallback: last number in the text
nums = re.findall(r"-?\d[\d,]*\.?\d*", text)
return nums[-1].replace(",", "") if nums else ""
# --- teacher: API or local ---
def teacher_solve_api(question: str) -> str:
from openai import OpenAI
client = OpenAI() # uses OPENAI_API_KEY env var
# o3-mini / o1 reasoning models emit a reasoning trace then the answer.
# For non-reasoning models, prompt explicitly for a <think> block.
resp = client.chat.completions.create(
model=os.environ.get("TEACHER_API_MODEL", "o3-mini"),
messages=[
{"role": "system", "content": "Solve the math problem. First reason step by step inside <think>...</think> tags, then give the final answer as \\boxed{N}."},
{"role": "user", "content": question},
],
)
return resp.choices[0].message.content
def teacher_solve_local(question: str, model, tokenizer) -> str:
import torch
messages = [
{"role": "system", "content": "Solve the math problem. First reason step by step inside <think>...</think> tags, then give the final answer as \\boxed{N}."},
{"role": "user", "content": question},
]
inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
with torch.no_grad():
out = model.generate(inputs, max_new_tokens=1024, do_sample=False, temperature=1.0)
return tokenizer.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True)
# --- generate + curate ---
def main():
traces = []
if os.environ.get("TEACHER_MODE", "api") == "local":
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
tid = os.environ.get("TEACHER_LOCAL_ID", "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")
tok = AutoTokenizer.from_pretrained(tid)
mdl = AutoModelForCausalLM.from_pretrained(tid, torch_dtype=torch.bfloat16, device_map="auto")
solve = lambda q: teacher_solve_local(q, mdl, tok)
else:
solve = teacher_solve_api
kept = 0
for i, ex in enumerate(problems):
q, gt = ex["question"], extract_groundtruth(ex["answer"])
try:
teacher_out = solve(q)
except Exception as e:
print(f"[{i}] teacher error: {e}")
continue
pred = extract_final_answer(teacher_out)
if pred == gt: # CURATION: keep only correct traces
traces.append({"question": q, "trace": teacher_out, "answer": gt})
kept += 1
if (i + 1) % 25 == 0:
print(f"[{i+1}/200] kept {kept} correct traces")
with open("distill_traces.jsonl", "w") as f:
for t in traces:
f.write(json.dumps(t) + "\n")
print(f"Wrote {kept} verified traces to distill_traces.jsonl (kept {kept}/{len(problems)})")
if __name__ == "__main__":
main()
```
Run it:
```bash
export TEACHER_MODE=api # or "local"
export OPENAI_API_KEY=sk-... # if using api
python make_distill_data.py
head -1 distill_traces.jsonl | python -m json.tool
```
Read one trace. You should see a `<think>…</think>` block followed by a `\boxed{N}` answer. **The kept traces are your distillation corpus** — verified correct, the same curation step R1-Distill applied at 800K scale.
> **If the API is unavailable:** set `TEACHER_MODE=local` and use a small R1-Distill model as teacher. The lift will be smaller (the teacher's ceiling is lower), but the full pipeline runs locally and free.
---
## Phase 2 — Format the traces for student SFT (3 min)
The student must learn to *emit* the trace. We format each example so the target text is the teacher's full trace (including the `<think>` block and the final answer). Create `format_data.py`:
```python
# format_data.py — convert traces into SFT (messages) format
import json
SYSTEM = (
"Solve the math problem. First reason step by step inside <think>...</think> tags, "
"then give the final answer as \\boxed{N}."
)
rows = []
with open("distill_traces.jsonl") as f:
for line in f:
t = json.loads(line)
rows.append({
"messages": [
{"role": "system", "content": SYSTEM},
{"role": "user", "content": t["question"]},
{"role": "assistant", "content": t["trace"]}, # the trace IS the target
]
})
with open("distill_sft.jsonl", "w") as f:
for r in rows:
f.write(json.dumps(r) + "\n")
print(f"Wrote {len(rows)} SFT examples to distill_sft.jsonl")
```
Run it and inspect one example:
```bash
python format_data.py
head -1 distill_sft.jsonl | python -m json.tool
```
> **FT07 discipline checkpoint.** Before training, run one example through `apply_chat_template(..., return_assistant_tokens_mask=True)` and decode it. Confirm the ChatML role tokens and EOS are present and the assistant mask covers the trace. You already learned why (FT07) — do not skip it.
---
## Phase 3 — Load the student in 4-bit + attach LoRA (3 min)
The FT08 QLoRA workflow, applied to distillation. Create `train.py`:
```python
# train.py — QLoRA distillation SFT (Phases 3–4)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
STUDENT_ID = "Qwen/Qwen2.5-1.5B-Instruct"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained(STUDENT_ID)
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
STUDENT_ID,
quantization_config=bnb_config,
device_map="auto",
attn_implementation="flash_attention_2", # or "sdpa" on older GPUs / Colab T4
)
# prepare + attach (the FT08 workflow)
model = prepare_model_for_kbit_training(model)
lora_config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters() # expect <1% trainable
```
You should see `trainable params: ~9.4M || all params: ~1.58B || trainable%: ~0.60%`. **Under 1%** — the LoRA promise, applied to distillation.
---
## Phase 4 — Train the student on the traces (8–15 min GPU)
Standard TRL `SFTTrainer`. The chat template is applied automatically because we use the `messages` format.
```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
dataset = load_dataset("json", data_files="distill_sft.jsonl", split="train")
training_args = SFTConfig(
output_dir="./distilled-student",
num_train_epochs=3,
per_device_train_batch_size=2,
gradient_accumulation_steps=4, # effective batch = 8
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_ratio=0.03,
logging_steps=5,
save_strategy="epoch",
bf16=True, # fp16=True on older cards (V100)
gradient_checkpointing=True,
optim="paged_adamw_8bit",
max_seq_length=1024, # traces are longer than pirate-style SFT
report_to="none",
)
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=dataset,
processing_class=tokenizer,
)
trainer.train()
# merge for deployment + eval (Option B from FT08)
merged = model.merge_and_unload()
merged.save_pretrained("./distilled-student/merged", safe_serialization=True)
tokenizer.save_pretrained("./distilled-student/merged")
```
**Watch the loss.** It should descend from ~1.2–1.8 to ~0.6–0.9 over the epochs. The traces are longer than the FT08 pirate examples, so `max_seq_length=1024` (raise to 2048 if you see truncation warnings). If loss is NaN, that is an FT07 template/EOS bug, not a distillation bug — re-run the inspection loop.
**Watch VRAM** with `nvidia-smi -l 2`. Peak ~7–10 GB on a 1.5B QLoRA at 1024 context. On a Colab T4 you have headroom.
---
## Phase 5 — Merge and save (2 min)
Already done at the end of Phase 4 (the `merge_and_unload` call). The merged model at `./distilled-student/merged` is a single self-contained model — load directly with `AutoModelForCausalLM.from_pretrained`, no PEFT, no 4-bit at serve time. This is what we eval.
> **Caveat (FT08):** `merge_and_unload()` on a 4-bit base dequantizes — the merged model is larger than the 4-bit base. For production you re-quantize (GGUF/AWQ) afterward. For this lab, eval the merged FP16 model directly.
---
## Phase 6 — Eval: distilled vs un-distilled on GSM8K (5 min)
The moment of truth. We sample N GSM8K test problems, run both the distilled model and the un-distilled base, and compare accuracy.
```python
# eval.py — measure the GSM8K lift
import torch, re, json
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
STUDENT_ID = "Qwen/Qwen2.5-1.5B-Instruct"
DISTILLED_PATH = "./distilled-student/merged"
N_EVAL = 200
gsm = load_dataset("openai/gsm8k", "main", split="test")
eval_set = gsm.select(range(N_EVAL))
def extract_gt(answer_str):
m = re.search(r"####\s*(.+)$", answer_str.strip())
return m.group(1).strip().replace(",", "") if m else answer_str.strip()
def extract_pred(text):
m = re.search(r"\\boxed\{([^}]+)\}", text)
if m: return m.group(1).strip().replace(",", "")
m = re.search(r"####\s*(.+)$", text.strip())
if m: return m.group(1).strip().replace(",", "")
nums = re.findall(r"-?\d[\d,]*\.?\d*", text)
return nums[-1].replace(",", "") if nums else ""
SYSTEM = ("Solve the math problem. First reason step by step inside <think>...</think> tags, "
"then give the final answer as \\boxed{N}.")
def evaluate(model_id_or_path, label):
tok = AutoTokenizer.from_pretrained(model_id_or_path)
mdl = AutoModelForCausalLM.from_pretrained(model_id_or_path, torch_dtype=torch.bfloat16, device_map="auto")
correct = 0
for i, ex in enumerate(eval_set):
msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": ex["question"]}]
inputs = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(mdl.device)
with torch.no_grad():
out = mdl.generate(inputs, max_new_tokens=512, do_sample=False) # greedy for reproducibility
pred = extract_pred(tok.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True))
if pred == extract_gt(ex["answer"]):
correct += 1
if (i + 1) % 50 == 0:
print(f" [{label}] {i+1}/{N_EVAL} acc so far: {correct/(i+1):.3f}")
acc = correct / N_EVAL
print(f"=== {label}: GSM8K accuracy = {acc:.3f} ({correct}/{N_EVAL}) ===")
return acc
base_acc = evaluate(STUDENT_ID, "UN-DISTILLED BASE")
distilled_acc = evaluate(DISTILLED_PATH, "DISTILLED STUDENT")
print(f"\nLIFT from distillation: {distilled_acc - base_acc:+.3f} "
f"({base_acc:.3f} -> {distilled_acc:.3f})")
```
Run it:
```bash
python eval.py
```
**Expected:** a measurable GSM8K accuracy gain — often **+0.10 to +0.25** (10–25 percentage points) from 200 traces alone, depending on teacher quality. The base already knows arithmetic (FT00: the capability is latent); distillation steered it to reason reliably and in a structured trace. **The lift is your distillation.**
**Control check:** load the un-distilled base with the *same* system prompt asking for a `<think>` trace. It will often produce a shorter, less reliable trace — or skip the trace and answer directly, getting more wrong. The distilled student reliably emits the structured trace. The difference between the two accuracies is what 200 curated traces bought you.
---
## Deliverables
Submit `ft15-lab-report.md` containing:
- [ ] **Your teacher choice** (API model or local R1-Distill) and the `kept/total` curation rate (how many of the 200 traces were verified correct and kept).
- [ ] **One sample trace** from `distill_traces.jsonl` — show the `<think>…</think>` block and the `\boxed{N}` answer. This is your distillation target.
- [ ] **The FT07 inspection output** for one formatted SFT example (decoded, EOS present, assistant mask covers the trace).
- [ ] **`print_trainable_parameters()` output** — confirm trainable% < 1.
- [ ] **The training loss curve** (start loss, end loss, step count).
- [ ] **Peak VRAM observed** (`nvidia-smi`).
- [ ] **The eval results:** un-distilled base accuracy, distilled student accuracy, and the lift. Include 2–3 example outputs from each model on the same problem, showing the distilled model's structured trace vs the base's.
- [ ] A 4–6 sentence reflection: why did SFT on traces (no RL) produce a reasoning lift? Tie it to the FT00 thesis (steering, not teaching) and the R1-Distill result. What is the ceiling of this approach, and when would you escalate to RFT or GRPO?
---
## Solution key
These are defensible answers, not the only wording.
### Curation rate
With a frontier API teacher (o3-mini / equivalent) on 200 GSM8K train problems, expect to keep **~150–190 of 200** (75–95%) — the teacher is strong and GSM8K train is not adversarial. With a local 1.5B R1-Distill teacher, expect **~100–160 of 200** (50–80%) — weaker teacher, more traces rejected. Either is a usable corpus; the lift scales with teacher quality.
### Loss trajectory
Expect start ~1.2–1.8, smooth cosine descent, end ~0.6–0.9 after 3 epochs. Traces are longer than the FT08 pirate SFT, so `max_seq_length=1024` is the right floor. If loss plateaus above ~1.1, the rank may be too low or the LR too low; if it crashes below ~0.5, you are overfitting (the corpus is small — raise dropout or reduce epochs to 2). A healthy run lands ~0.7–0.9 and generalizes (Phase 6 confirms).
### VRAM
Peak ~7–10 GB on a 1.5B QLoRA at 1024 context, batch 2 + grad accum 4. The FT08 rule of thumb holds; the longer context adds a bit of activation overhead. On a Colab T4 (16 GB) you have comfortable headroom. If you OOM, lower `per_device_train_batch_size` to 1 and raise `gradient_accumulation_steps` to 8.
### Eval lift
With an API teacher: expect the un-distilled Qwen2.5-1.5B base to score ~0.30–0.45 on GSM8K (it can do arithmetic but does not reliably emit a structured trace), and the distilled student to score ~0.50–0.70 — a lift of **+0.10 to +0.25**. With a local 1.5B R1-Distill teacher, expect a smaller lift (~+0.05 to +0.15) because the teacher's ceiling is lower. Either way, the lift is real and attributable to the traces: the base already knew the arithmetic (FT00); distillation routed it through a reliable reasoning pathway.
### Reflection (model answer)
SFT on the teacher's traces produced a reasoning lift because distillation is steering, not teaching (FT00). The Qwen2.5-1.5B base already saw arithmetic during pretraining — the capability was latent. The teacher's `<think>` traces are a steering wheel: they route the student's probability mass through a structured reasoning pathway reliably, so it emits a chain of thought and reaches the correct answer more often. No RL was needed because the teacher already did the expensive search at data-generation time; the student imitates. This is the R1-Distill result in miniature: R1-Distill-Qwen-32B beat o1-mini via SFT on ~800K traces; we got a measurable lift from 200. The ceiling of this approach is the teacher's quality — the student cannot exceed the teacher (imitation). I would escalate to RFT (generate the student's own candidates, filter by math reward, retrain on winners) if the distilled model falls short of the target and I have no stronger teacher; to GRPO only if RFT plateaus and I have the budget and stability tolerance for full RL.
---
## Stretch goals
1. **RFT variant.** After distillation, run one RFT iteration: generate 4 candidates per problem (temperature 0.8) from the *distilled* model on the GSM8K train set, filter by math reward (keep only correct), and SFT the model on the winners. Measure the additional lift. This is the FT15 RFT loop, felt directly — and it tests whether RFT improves on the distilled baseline without RL.
2. **Teacher-quality sweep.** Run the full pipeline with (a) a frontier API teacher and (b) the local 1.5B R1-Distill teacher. Compare the lift. This demonstrates the FT15 ceiling-is-teacher property empirically: the stronger teacher produces a stronger student.
3. **Trace-length curation.** Some teacher traces may be excessively long. Re-run with a length filter (keep traces under 512 tokens) and compare the lift. This tests the "distilling traces the student cannot imitate" anti-pattern — do overly long traces hurt the small student?
4. **Non-thinking control.** Eval the distilled model with a system prompt that says "answer directly, do not use a think block." Compare to the thinking-mode eval. This previews the hybrid-model tradeoff (FT15.5): thinking costs latency but buys accuracy.
5. **Scale the corpus.** Re-run with 500 and 1000 problems (if API budget allows). Plot the lift vs corpus size. Where does it saturate? This is the R1-Distill scaling question in miniature — they used ~800K; you will see diminishing returns well before that.
6. **Export to GGUF.** Take the merged distilled model and convert to GGUF (FT19 preview) via `llama.cpp`'s `convert_hf_to_gguf.py`, then load in Ollama. You now have a local, served reasoning model distilled from a teacher — the full FT15 → FT19 deploy path.