LLM Training & Alignment — Engineering Implementations

ZeRO: How Memory Partitioning Trains Trillion-Parameter Models Across GPUs

Rajbhandari, S., Rasley, J., Ruwase, O., He, Y. (Microsoft)

arXiv 1910.02054 / SC20

An engineering walkthrough of ZeRO: the GPU memory problem on real hardware, why data parallelism replicates training state, and how ZeRO-1/2/3 partition optimizer states, gradients, and parameters across GPU memory.

Model states (parameters + gradients + optimizer states) dominate GPU memory: 16 bytes per parameter under mixed-precision Adam
Traditional data parallelism replicates the entire training state on every GPU, wasting HBM
ZeRO partitions training state across data-parallel GPUs in three cumulative stages (optimizer, gradients, parameters)
ZeRO-3 reduces per-device model-state memory linearly with GPU count, at a 1.5x communication cost

ZeRO DeepSpeed Distributed Training GPU Memory Data Parallelism Model Training A100 Optimizer State Partitioning Large Language Models

Read Implementation → View Original Paper ↗

GRPO: Group Relative Policy Optimization in DeepSeekMath

Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y. K., Wu, Y., Guo, D.

arXiv 2402.03300

A technical dissection of Group Relative Policy Optimization (GRPO), the reinforcement learning algorithm introduced in DeepSeekMath. Covers why the paper replaced PPO's value model with a group-relative baseline, how the advantage is estimated from sampled outputs, the GRPO objective, outcome vs. process supervision, and iterative RL — with an engineering interpretation of why it works.

GRPO removes PPO's separate value/critic model, cutting the memory and compute overhead of RL for large language models
The baseline is estimated from a group of sampled outputs per question rather than a learned value function
Group-relative reward normalization turns each sampled group into its own baseline, producing the advantage directly
Supports both outcome supervision and process supervision, and extends to iterative RL as the reward model is refreshed

GRPO Reinforcement Learning PPO Policy Optimization LLM Mathematical Reasoning DeepSeekMath RLHF

Read Implementation → View Original Paper ↗

DeepSeek-R1: Incentivizing LLM Reasoning with Pure Reinforcement Learning (GRPO, Cold-Start, and Distillation)

DeepSeek-AI

arXiv 2501.12948

An engineering dissection of DeepSeek-R1 — the first open work to show that reasoning can be incentivized in an LLM by pure reinforcement learning, with no supervised fine-tuning first. Covers DeepSeek-R1-Zero's critic-free GRPO loop, rule-based accuracy and format rewards, the emergent 'aha moment' and self-evolution, the four-stage cold-start pipeline that fixes readability and language mixing, distillation into small dense models, and the honest account of what did not work (PRM and MCTS).

DeepSeek-R1-Zero is trained by pure RL directly on DeepSeek-V3-Base with no supervised fine-tuning — the first open result showing reasoning behaviour (self-verification, reflection, long chains) can emerge from reward alone, lifting AIME 2024 pass@1 from 15.6% to 71.0%
The RL algorithm is GRPO: it deletes PPO's separate critic/value network and instead scores a group of sampled answers against their own average, using that group-relative advantage to update the policy — cheaper RL that makes reasoning training affordable
Rewards are rule-based (accuracy + a format reward that forces reasoning inside ...), deliberately avoiding a neural reward model because a learned reward gets hacked in large-scale RL
DeepSeek-R1 adds a four-stage pipeline around the RL — cold-start CoT fine-tuning, reasoning RL with a language-consistency reward, rejection-sampling + SFT on 800K samples, then a final RL for all scenarios — reaching parity with OpenAI-o1-1217; the same 800K samples distil that reasoning into 1.5B–70B dense models with SFT alone

DeepSeek-R1 Reasoning Reinforcement Learning GRPO Chain-of-Thought Test-Time Compute Reward Modeling Distillation Cold Start LLM Reasoning OpenAI o1 Self-Evolution

Read Implementation → View Original Paper ↗

Scaling Laws for Neural Language Models: Predicting Loss from Model, Data, and Compute

Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., Amodei, D.

arXiv 2001.08361

An engineering dissection of the Kaplan et al. scaling laws — the empirical power laws relating language-model cross-entropy loss to model size, dataset size, and compute. Covers the three basic power laws and their exponents, why scale dominates architecture shape, the combined L(N,D) overfitting law, the L(N,S) training-curve law, the compute-optimal allocation that says to train very large models and stop short of convergence, the critical batch size, and how the paper reframed how a compute budget should be spent.

Language-model loss follows clean power laws in model size N, dataset size D, and compute C — spanning up to eight orders of magnitude — while depending only weakly on depth/width/head-count
A single combined law L(N,D) predicts overfitting from the ratio N^0.74/D: every 8× in parameters needs only ~5× more data to avoid a penalty
For a fixed compute budget, optimal training means a very large model trained on modest data and stopped well before convergence — N grows as C^0.73, data only as C^0.27
Larger models are more sample-efficient; a training curve's early portion predicts its converged loss via a per-step power law

Scaling Laws Compute-Optimal Power Laws Pretraining Sample Efficiency Model Size Language Models GPT Kaplan

Read Implementation → View Original Paper ↗

LoRA: Low-Rank Adaptation — Fine-Tuning Large Models Without Paying for Them Twice

Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., Chen, W.

ICLR 2022 · arXiv 2106.09685

An engineering dissection of Low-Rank Adaptation (LoRA). Covers why full fine-tuning of large language models is a deployment problem, why adapters and prefix-tuning trade quality for efficiency, the low-intrinsic-rank hypothesis, the BA reparametrization and its scaling, why merging the update means zero added inference latency, the memory and task-switching economics, and the paper's own evidence for why a rank as small as one can work.

Freezes the pretrained weights and injects a trainable low-rank update ΔW = BA into chosen weight matrices, cutting trainable parameters by up to 10,000× on GPT-3 175B
Because ΔW = BA can be merged into W0 at deploy time, LoRA adds no inference latency — unlike adapter layers, which must run sequentially
Reduces GPT-3 175B training VRAM from 1.2TB to 350GB and the per-task checkpoint from 350GB to 35MB, making task-switching a weight swap rather than a redeploy
The paper's subspace and amplification analysis shows the adaptation matrix has very low intrinsic rank and amplifies features already latent in the pretrained weights

LoRA Fine-Tuning PEFT Low-Rank Parameter-Efficient LLM Adapters GPT-3 Deployment

Read Implementation → View Original Paper ↗

InstructGPT: Aligning Language Models With Human Feedback (RLHF)

Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., Schulman, J., Hilton, J., Kelton, F., Miller, L., Simens, M., Askell, A., Welinder, P., Christiano, P., Leike, J., Lowe, R.

NeurIPS 2022 · arXiv 2203.02155

An engineering dissection of InstructGPT — OpenAI's three-stage reinforcement learning from human feedback (RLHF) pipeline. Covers why the language-modeling objective is misaligned with following instructions, the SFT → reward-model → PPO pipeline, the data-collection pipeline, the reward-model training objective and the batching trick that makes it efficient, the PPO objective with its KL penalty and pretraining mix (PPO-ptx), the alignment tax, and why a 1.3B aligned model beats a 175B unaligned one.

Three-stage pipeline: supervised fine-tuning on demonstrations, a reward model trained on ranked comparisons, then PPO against that reward model
A 1.3B InstructGPT model's outputs are preferred to the 175B GPT-3's — alignment buys more than 100× the parameters on this distribution
The reward model trains all C(K,2) comparisons per prompt as one batch element — one forward pass per completion instead of C(K,2), which both speeds training and stops overfitting
PPO adds a per-token KL penalty to the SFT model; PPO-ptx mixes pretraining gradients back in to pay down the alignment tax on public NLP benchmarks

InstructGPT RLHF Reward Model PPO Alignment SFT Preference Optimization GPT-3 Human Feedback

Read Implementation → View Original Paper ↗

DPO: Direct Preference Optimization — Your LM Is Secretly a Reward Model

Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C. D., Finn, C.

arXiv 2305.18290

A technical dissection of Direct Preference Optimization (DPO). The focus is the one thing that matters: how DPO turns a preference dataset directly into a policy-training loss, removing the separate reward model and the PPO rollout/value-model machinery — while keeping a frozen reference policy. Covers the preference data, the trainable policy versus the frozen reference, token-level log-probabilities, the DPO objective term-by-term, the implicit reward, the full KL-constrained-RL derivation, a detailed DPO-vs-PPO comparison, and why it works.

DPO replaces the two-stage RLHF pipeline (train a reward model, then run PPO) with a single closed-form loss computed directly on preference pairs
The trainable policy and a frozen reference policy are both required — DPO removes the reward model and the value/critic model, not the reference model
It optimizes a relative preference — make the chosen response more likely than the rejected one compared with the reference — not raw chosen-response likelihood
The implicit reward r(x,y) = beta * log(policy / reference) is the paper's title made precise: the language model already encodes a reward inside its own log-ratio

DPO Direct Preference Optimization RLHF PPO Preference Optimization Reward Model Bradley-Terry Reference Policy LLM Alignment KL Regularization

Read Implementation → View Original Paper ↗

TDPO: Token-level Direct Preference Optimization

Zeng, Y., Liu, G., Ma, W., Yang, N., Zhang, H., Wang, J.

arXiv 2404.11999

A technical dissection of Token-level Direct Preference Optimization (TDPO), building directly on the DPO entry. The focus is the one change that matters: DPO scores a whole response as a single sequence-level unit, but a language model generates and drifts token-by-token — so TDPO moves the preference optimization and, crucially, the KL regularization to the token level. Covers why DPO's sequence-level KL causes mode collapse and low diversity, the forward-KL-vs-reverse-KL distinction at the heart of the method, the token-level reward built from the same log-ratio (no separate reward model, no PPO), the TDPO objective term-by-term, the TDPO-1 vs TDPO-2 variants and their two KL coefficients, and why token-level forward-KL control improves the alignment-vs-diversity trade-off.

TDPO keeps DPO's direct-preference philosophy — no separate reward model, no PPO rollout loop, a frozen reference — and changes only where the KL constraint is applied: per token instead of per whole sequence
The core insight is forward KL vs reverse KL: DPO's implicit sequence-level constraint is mode-seeking (drives diversity down, causing repetitive collapse), while TDPO adds a token-level forward-KL term that is mass-covering and preserves diversity
TDPO-1 and TDPO-2 are two variants, not one: TDPO-2 adds a second coefficient alpha and a stop-gradient baseline on the chosen response, giving tighter KL control and a better reward-KL frontier than TDPO-1
The two KL coefficients beta and alpha, and the ratio between them, are what tune the alignment-versus-diversity balance — the single knob DPO does not have

TDPO Token-level DPO Direct Preference Optimization DPO Forward KL Reverse KL KL Regularization Preference Optimization LLM Alignment Diversity

Read Implementation → View Original Paper ↗