🏋️ Training Pipeline

10 Training Scripts Covering Pretrain to RLHF

📖 Estimated reading time: 9 min

Files

Training Modes

~2800

Lines

Python

Language

🔵 Pretrain🔵 SFT 🟢 LoRA🟢 DPO 🟡 PPO🟡 GRPO🟡 SPO 🔴 Distillation🔴 Reasoning

💡 Analogy — A Multi-Stage School Education

Think of the training pipeline as a student progressing through school. Pretraining is elementary school—learning the basics of language from a massive textbook. SFT is high school—learning to follow instructions from a tutor. DPO/PPO/GRPO/SPO are university seminars—refining judgment by comparing good and bad answers under expert feedback. LoRA is a crash course—quickly adapting knowledge without rewriting everything. Distillation is mentorship—a senior model teaching a junior one. Reasoning is the final thesis—learning to show your work step by step.

1. Architecture Overview

The training pipeline is organized as a progressive curriculum. Each script shares a common loop (data → forward → loss → backward → optimize) but specializes in its objective and data format. The dependency graph below shows the recommended training order: pretrain first, then branch into SFT, LoRA, or distillation, and finally RLHF alignment via DPO, PPO, GRPO, or SPO.

🔍 Click to zoom

flowchart TD TU[trainer_utils.py
Shared Utils] TK[train_tokenizer.py
BPE Tokenizer] PT[train_pretrain.py
Pretraining] SFT[train_full_sft.py
Full SFT] LORA[train_lora.py
LoRA SFT] DPO[train_dpo.py
DPO] PPO[train_ppo.py
PPO] GRPO[train_grpo.py
GRPO] SPO[train_spo.py
SPO] DIST[train_distillation.py
Distillation] REAS[train_reason.py
Reasoning] TK --> PT PT --> SFT PT --> LORA SFT --> DPO DPO --> REAS SFT --> PPO SFT --> GRPO SFT --> SPO PT --> DIST TU -.->|init, LR, ckpt| PT TU -.-> SFT TU -.-> LORA TU -.-> DPO TU -.-> PPO TU -.-> GRPO TU -.-> SPO TU -.-> DIST TU -.-> REAS style TU fill:#1f6feb,stroke:#58a6ff,color:#fff style PT fill:#238636,stroke:#3fb950,color:#fff style SFT fill:#238636,stroke:#3fb950,color:#fff style LORA fill:#2ea043,stroke:#3fb950,color:#fff style DPO fill:#9a6700,stroke:#d29922,color:#fff style PPO fill:#9a6700,stroke:#d29922,color:#fff style GRPO fill:#9a6700,stroke:#d29922,color:#fff style SPO fill:#9a6700,stroke:#d29922,color:#fff style DIST fill:#b62324,stroke:#f85149,color:#fff style REAS fill:#b62324,stroke:#f85149,color:#fff style TK fill:#1f6feb,stroke:#58a6ff,color:#fff

2. Key Classes & Functions

Symbol	File	Type	Description
`get_lr(step, total, lr)`	trainer_utils.py	func	Cosine learning-rate schedule: `lr(0.1+0.45(1+cos(π·step/total)))`
`init_distributed_mode()`	trainer_utils.py	func	Initializes DDP with NCCL backend & sets local rank
`lm_checkpoint()`	trainer_utils.py	func	Save/load checkpoints (model, optimizer, epoch, step, wandb_id)
`init_model(cfg, weight, tok, dir, dev)`	trainer_utils.py	func	Builds model + tokenizer from config, loads optional pretrained weights
`SkipBatchSampler`	trainer_utils.py	class	Sampler that skips first N batches for training resume
`train_pretrain.main()`	train_pretrain.py	func	Next-token prediction loop with AdamW, cosine LR, grad accumulation
`train_full_sft.main()`	train_full_sft.py	func	SFT on instruction data with chat template formatting
`save_lora()`	train_lora.py	func	Extracts & saves only LoRA adapter weights (~0.5% of params)
`dpo_loss()`	train_dpo.py	func	Bradley-Terry preference loss with beta-scaled log-ratio
`logits_to_log_probs()`	train_dpo.py	func	Converts model logits to per-token log probabilities
`CriticModel`	train_ppo.py	class	Value head on top of MiniMindForCausalLM for PPO critic
`distillation_loss()`	train_distillation.py	func	KL divergence with temperature: `T²·KL(student ‖ teacher)`
`AutoAdaptiveValueTracker`	train_spo.py	class	Bayesian baseline with KL-adaptive ρ via Beta distribution
`compute_rho()`	train_spo.py	func	Decay rate: `2^(−kl / D_half)` for adaptive baseline

3. Data-Flow Diagram

Every training script follows this core loop. RLHF variants add reward / reference-model steps between Loss and Backward.

🔍 Click to zoom

sequenceDiagram participant DL as DataLoader participant M as Model participant L as Loss Function participant O as Optimizer participant C as Checkpoint DL->>M: batch (input_ids, labels) M->>L: logits Note over L: CE / DPO / PPO / KL-distill L->>M: loss.backward() Note over M: Gradient Accumulation (N steps) M->>O: accumulated gradients Note over O: AdamW + grad clip (1.0) O->>M: updated parameters O->>C: periodic save (epoch, step, optimizer) Note over C: lm_checkpoint()

4. Module Interactions

Source	Target	Interaction
All train_*.py	trainer_utils.py	Import `get_lr`, `init_model`, `lm_checkpoint`, `Logger`
train_pretrain.py	model/dataset	Loads `PretrainDataset`, creates `MiniMindForCausalLM`
train_full_sft.py	model/dataset	Loads `SFTDataset`, resumes from pretrain checkpoint
train_lora.py	model	Calls `apply_lora()` then freezes non-LoRA params
train_dpo.py	model	Frozen reference model + trainable policy; `DPODataset`
train_ppo.py	model	4 models: policy, ref, critic, reward; GAE advantage
train_grpo.py	reward model	External `InternLM2-1.8B-reward` scores generations
train_spo.py	reward model	Same reward model; `AutoAdaptiveValueTracker` baseline
train_distillation.py	teacher model	Frozen teacher provides soft targets; student trained on mixed loss
train_reason.py	model	10× loss weight on think/answer special tokens; starts from DPO weights

5. Deep-Dive Hub

🎯 RLHF Pipeline

DPO uses Bradley-Terry preference model (β=0.1). PPO adds critic + GAE with clipped objective. GRPO generates 8 responses and normalizes group-relative advantages. SPO uses a single generation with Bayesian adaptive baseline (ρ = 2^(−kl/D_half)).

📉 Cosine LR & Grad Accumulation

get_lr() implements cosine annealing: lr × (0.1 + 0.45 × (1 + cos(π·step/total))). Gradient accumulation allows effective batch sizes N× physical batch. Gradient clipping at norm 1.0 prevents explosion.

💾 Checkpoint & Resume

lm_checkpoint() saves model, optimizer, epoch, step, wandb_id, and world_size. SkipBatchSampler skips already-processed batches on resume. Supports multi-GPU DDP state restoration.

🧬 Knowledge Distillation

Student-teacher framework: total_loss = α·CE + (1−α)·T²·KL(student ‖ teacher). Default α=0.5, T=1.5. Teacher is frozen; only student gradients flow. Temperature softens probability distributions for richer supervision signal.

6. Directory Map

All 11 files live in the trainer/ directory. Shared infrastructure is in trainer_utils.py; each training paradigm gets its own self-contained script that imports from the shared utilities module.

trainer/ ├── trainer_utils.py # Shared utils: LR, DDP, checkpoint, init_model ├── train_tokenizer.py # BPE tokenizer training (vocab=6400) ├── train_pretrain.py # Next-token prediction pretraining ├── train_full_sft.py # Full supervised fine-tuning ├── train_lora.py # LoRA adapter fine-tuning ├── train_dpo.py # Direct Preference Optimization ├── train_ppo.py # Proximal Policy Optimization + CriticModel ├── train_grpo.py # Group Relative Policy Optimization ├── train_spo.py # Self-Play Optimization + AdaptiveTracker ├── train_distillation.py # KL-divergence knowledge distillation └── train_reason.py # Reasoning with weighted special tokens

7. Configuration Reference

Knob	Default	Script(s)	Notes
`batch_size`	32	pretrain	16 for SFT, 4 for DPO/PPO
`learning_rate`	5e-4	pretrain	1e-6 SFT, 1e-4 LoRA, 4e-8 DPO
`max_seq_len`	340	pretrain	720 for reasoning
`epochs`	1–2	all	50 for LoRA
`grad_accumulation`	varies	all	Effective batch = batch_size × accumulation
`grad_clip`	1.0	all	Max gradient norm
`from_weight`	'pretrain'	SFT, DPO, …	'dpo' for reasoning
`lora_rank`	8	LoRA	LoRA decomposition rank
`beta`	0.1	DPO	KL constraint strength
`num_generations`	8	GRPO	Responses per prompt for group scoring
`alpha`	0.5	distillation	CE vs KL loss weight balance
`temperature`	1.5	distillation	Softmax temperature for soft targets
`vocab_size`	6400	tokenizer	ByteLevel BPE vocabulary size

8. Code Examples

Below are annotated excerpts from the most important functions in the training pipeline. These snippets are simplified for clarity — see the source files for full implementations.

Cosine Learning Rate Schedule

trainer_utils.py — get_lr()

import math

def get_lr(current_step, total_steps, lr):
    """Cosine annealing with warm floor at 10% of peak LR."""
    return lr * (0.1 + 0.45 * (1 + math.cos(math.pi * current_step / total_steps)))

DPO Loss with Bradley-Terry Model

train_dpo.py — dpo_loss()

import torch
import torch.nn.functional as F

def dpo_loss(policy_chosen_logps, policy_rejected_logps,
             ref_chosen_logps, ref_rejected_logps, beta=0.1):
    """Direct Preference Optimization with frozen reference model."""
    chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps)
    rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps)
    return -F.logsigmoid(chosen_rewards - rejected_rewards).mean()

Knowledge Distillation Loss

train_distillation.py — distillation_loss()

def distillation_loss(student_logits, teacher_logits, temperature=1.5):
    """KL divergence scaled by T² for knowledge transfer."""
    s_probs = F.log_softmax(student_logits / temperature, dim=-1)
    t_probs = F.softmax(teacher_logits / temperature, dim=-1)
    return (temperature ** 2) * F.kl_div(s_probs, t_probs, reduction='batchmean')