Instruction-Tuned Assistant
Transform a base pretrained language model into a multi-turn instruction-following assistant using ChatML formatting, assistant-masked SFT, LoRA from scratch, and DPO alignment
What You Will Build
A base pretrained language model (like `MyLLM v1` or a small open-weight base model) only knows how to continue documents—if you ask it a question, it often generates three more questions. In this project, you will build the complete post-training and alignment pipeline that turns a base model into a helpful assistant. You will implement ChatML turn serialization with special delimiters (`<|im_start|>`, `<|im_end|>`), build an SFT data collator that masks system and user prompt tokens with `-100` so cross-entropy loss trains strictly on assistant completions, implement `LoRALinear` ($W_0 x + \frac{\alpha}{r} BAx$) from scratch in PyTorch, run Supervised Fine-Tuning, optionally apply Direct Preference Optimization (DPO) on `(prompt, chosen, rejected)` pairs, and evaluate the model before and after post-training.
Why This Project Matters in the AI/LLM Stack
Post-training (SFT + LoRA + DPO) is the bridge between raw foundation models and real-world AI assistants. Implementing assistant-only loss masking, custom LoRA adapters, weight merging, and DPO log-ratio losses directly in PyTorch removes all mystery from modern alignment frameworks.
Learning Objectives
- Format single-turn and multi-turn conversation datasets into deterministic ChatML token sequences.
- Construct an SFT collator that masks `system` and `user` tokens with `-100` so gradients only update on `assistant` responses and `<|im_end|>` stop tokens.
- Implement `LoRALinear` from scratch in PyTorch, verify zero-initialization identity behavior, and merge trained adapters back into base weights.
- Implement the Direct Preference Optimization (DPO) sequence log-probability loss and compare base vs. SFT vs. DPO responses.
Prerequisites
- Phase 5 Project: Mini GPT Architecture (`milestone-10-mini-gpt-architecture`)
- Phase 6 Project: Pretrained Mini LLM (`milestone-11-pretrained-mini-llm`)
- Phase 7 Modules 7.1–7.3: Instruction Tuning, PEFT (LoRA/QLoRA), and Preference Alignment
Required Mathematical Concepts
- Assistant-only masked cross-entropy loss (`ignore_index = -100` on prompt tokens)
- Low-Rank Adaptation: h = W_0 x + (alpha / r) B (A x), with B initialized to 0
- LoRA weight merging: W_merged = W_0 + (alpha / r) B A
- DPO objective: L_DPO = -log sigma(beta * ((log pi_theta(y_w|x) - log pi_ref(y_w|x)) - (log pi_theta(y_l|x) - log pi_ref(y_l|x))))
Step-by-Step Implementation Tasks (12 Steps)
Check off each engineering task as your code passes its acceptance criteria. Progress is saved automatically in your browser.
ChatML Serialization & Assistant-Masked SFT Collator
Define chat control tokens, format multi-turn dialogues, and build the SFT batch collator with `-100` prompt masking.
Structured role delimiters (`<|im_start|>system`, `<|im_start|>user`, `<|im_start|>assistant`, `<|im_end|>`) give the model unambiguous turn boundaries.
During SFT, the model should be trained to predict the assistant's answer given the user's prompt—not penalized for failing to predict what question the user will ask.
Instruction examples vary in length; padding `input_ids` with `pad_id` and `labels` with `-100` up to the longest example in the batch prevents wasted compute.
Low-Rank Adaptation (`LoRALinear`) From Scratch
Implement `LoRALinear`, inject adapters into attention/MLP projections, verify VRAM/parameter savings, and implement adapter weight merging.
LoRA freezes the pretrained weight $W_0 \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$ and adds a trainable low-rank update $\frac{\alpha}{r} B A x$ where $A \sim \mathcal{N}(0, \sigma^2)$ and $B = 0$.
Replacing target linear layers (`q_proj`, `k_proj`, `v_proj`, `o_proj`, `up_proj`, `down_proj`) with `LoRALinear` freezes >98% of model parameters.
After training, folding $\Delta W = \frac{\alpha}{r} B A$ directly into $W_0$ eliminates the adapter branch and restores zero-overhead inference.
Supervised Fine-Tuning (SFT) & Multi-Turn Chat CLI
Curate a high-quality instruction dataset, run the SFT training loop, and build an interactive multi-turn chat loop that stops cleanly at `<|im_end|>`.
Following the LIMA principle, a clean, formatting-consistent dataset of 1,000–5,000 instruction-response pairs teaches reliable assistant behavior without degrading base fluency.
Fine-tuning for 1–3 epochs with a low learning rate (e.g., `2e-4` for LoRA or `2e-5` for full SFT) adapts the model to follow instructions and emit `<|im_end|>` at the end of its answer.
Multi-turn chat appends each user message and assistant reply to the conversation history, formatting the growing transcript with ChatML on each turn.
Direct Preference Optimization (DPO) & Before/After Evaluation
Implement the DPO loss function on pairwise preference data and evaluate base vs. SFT vs. DPO completions on a held-out test suite.
Direct Preference Optimization requires summing the per-token log-probabilities $\log \pi(y \mid x) = \sum_{t \in \text{response}} \log P(y_t \mid x, y_{<t})$ over the unmasked assistant tokens.
DPO aligns the policy model $\pi_\theta$ against a frozen reference SFT model $\pi_{\text{ref}}$ on `(prompt, chosen, rejected)` triples without training a separate reward model.
Side-by-side comparison on held-out prompts (question answering, summarization, formatting constraints, refusal of nonsensical premises) proves what each stage achieved.
Starter & Verification Code Examples
LoRALinear & Direct Preference Optimization (DPO) Loss in PyTorch
Clean from-scratch implementation of a mergeable LoRA linear layer and the DPO preference loss.
import torch
import torch.nn as nn
import torch.nn.functional as F
class LoRALinear(nn.Module):
def __init__(self, base: nn.Linear, r: int = 8, alpha: float = 16.0):
super().__init__()
self.base = base
self.base.weight.requires_grad_(False)
if self.base.bias is not None:
self.base.bias.requires_grad_(False)
self.r = r
self.scaling = alpha / r
self.lora_A = nn.Parameter(torch.randn(r, base.in_features) * 0.02)
self.lora_B = nn.Parameter(torch.zeros(base.out_features, r))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.base(x) + (x @ self.lora_A.T @ self.lora_B.T) * self.scaling
@torch.no_grad()
def merge(self) -> nn.Linear:
self.base.weight.add_((self.lora_B @ self.lora_A) * self.scaling)
return self.base
def dpo_loss(pi_win: torch.Tensor, pi_lose: torch.Tensor, ref_win: torch.Tensor, ref_lose: torch.Tensor, beta: float = 0.1):
chosen_logratios = pi_win - ref_win
rejected_logratios = pi_lose - ref_lose
logits = chosen_logratios - rejected_logratios
loss = -F.logsigmoid(beta * logits).mean()
acc = (logits > 0).float().mean()
return loss, accHow to Test Your Implementation
Verify prompt loss masking, LoRA zero-init and merge equivalence, and DPO step-0 analytical loss.
Prompt Loss Masking Invariance Test
Prove that corrupting logits on user prompt positions does not change the SFT loss.
Steps: Randomize logits at all positions where `labels == -100` and recompute cross-entropy loss.
Pass Criteria: Loss value and gradients on assistant tokens remain 100% identical.
LoRA Merge Equivalence Test
Confirm `W_merged = W_0 + (alpha/r) * B @ A` produces identical outputs to the unmerged branch.
Steps: Randomize `lora_A` and `lora_B`, compute `y_unmerged = lora_layer(x)`, call `merge_lora_weights()`, and compute `y_merged = merged_layer(x)`.
Pass Criteria: `torch.allclose(y_unmerged, y_merged, atol=1e-5)` is `True`.
DPO Reference Identity Test
Verify mathematical correctness of the DPO objective when policy equals reference.
Steps: Pass identical tensors for `policy_*_logps` and `ref_*_logps` into `dpo_loss`.
Pass Criteria: `loss.item()` equals `math.log(2)` (`0.693147`) to 6 decimal places.
Overall Project Acceptance Criteria
- SFT collator masks 100% of `system`, `user`, and `pad` tokens with `-100` so loss is computed exclusively on `assistant` tokens.
- `LoRALinear` produces exact zero delta at initialization (`B = 0`) and merges cleanly into `nn.Linear` with `< 1e-5` discrepancy.
- DPO loss equals `ln(2)` when policy and reference log-probabilities are identical and increases `reward_accuracy` during preference training.
- Repository includes side-by-side Base vs. SFT vs. DPO evaluation results on held-out prompts.
Expandable Hints & Debugging Tips
Try implementing each stage on your own first. Expand a hint below if you want to verify tensor shapes, mathematical identities, or common edge cases.
Optional Stretch Challenges
LoRA Rank & Target Module Ablation
Compare LoRA ranks `r = 4, 16, 64` applied to attention-only (`q_proj, v_proj`) vs. all linear layers (`q, k, v, o, up, down`) on validation SFT loss and peak VRAM.
What You Learn: Why targeting all linear layers with moderate rank outperforms high-rank attention-only adapters.
Verifiable Math Reasoning Mini-Loop (GRPO)
Implement a miniature Group Relative Policy Optimization (GRPO) loop on synthetic arithmetic word problems using a deterministic regex answer grader.
What You Learn: How Reinforcement Learning with Verifiable Rewards (RLVR) trains `<think>...</think>` reasoning traces.
What Your Completed Repository Should Contain
An instruction-tuned and DPO-aligned conversational assistant with from-scratch `LoRALinear` adapters, interactive `chat.py` CLI, and a Base vs. SFT vs. DPO evaluation report.
Deliverable Artifacts
- `sft/chat_template.py` & `sft/collator.py` with ChatML formatting and `-100` prompt masking
- `sft/lora.py` implementing `LoRALinear`, recursive injection, and weight merging
- `sft/dpo.py` implementing sequence log-probabilities and DPO preference loss
- Interactive multi-turn CLI assistant (`chat.py`)
- Before-and-after Base vs. SFT vs. DPO benchmark evaluation (`ALIGNMENT_REPORT.md`)
After Completing This Project, You Should Be Able To...
- Serialize multi-turn conversations into ChatML tokens and mask prompt tokens with `-100`.
- Implement, train, save, and merge Low-Rank Adaptation (`LoRALinear`) layers from scratch.
- Implement Direct Preference Optimization (DPO) and explain how it aligns model preferences without a separate reward model.
- Evaluate post-trained assistants systematically against base pretrained models.
GitHub Repository Submission
Link your public or private GitHub repository URL to keep a structured portfolio of your from-scratch AI builds. Stored locally in your browser (self-reported; no GitHub tokens or remote code execution).
Suggested Repository: instruction-tuned-assistant-lora-dpo
Recommended Directory Layout:
sft/__init__.py sft/chat_template.py sft/collator.py sft/lora.py sft/dpo.py train_sft.py train_dpo.py chat.py ALIGNMENT_REPORT.md tests/test_sft_lora_dpo.py README.md
Recommended README Documentation Checklist
- Include a table showing total parameters vs. trainable LoRA parameters.
- Show side-by-side Base vs. SFT vs. DPO completions on 5 example prompts.
- Provide a sample multi-turn terminal transcript from `python chat.py`.