{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=training-systems-md-1 fw=all -->\n",
    "\n",
    "# Model Training\n",
    "\n",
    "Throughout this book we trained on a single device, occasionally two. Real\n",
    "budgets are larger: a LoRA fine-tune fits on one consumer GPU, a full\n",
    "fine-tune of a 7B model wants a node of them, and pretraining anything\n",
    "competitive occupies hundreds to hundreds of thousands of accelerators for\n",
    "weeks. This section is a practical map of that territory: the parallelism\n",
    "concepts you need in order to read any framework's documentation, the\n",
    "actual libraries people use in 2026, and which tool fits which scale.\n",
    "Distributed training is not a switch that makes a program faster — it\n",
    "divides model state and work, adds communication, and changes failure\n",
    "recovery — so the golden rule is to start from a measured single-GPU\n",
    "baseline and add the *simplest* technique that removes a demonstrated\n",
    "constraint.\n",
    "\n",
    "## From One GPU to Many\n",
    "\n",
    "### The Scaling Ladder\n",
    "\n",
    "![A practical scaling ladder progresses from one accelerator to replication, sharding, and composed parallelism.](../img/tools-training-ladder.svg)\n",
    "\n",
    "1. **One accelerator** establishes the reference: loss curve, tokens or\n",
    "   examples per second, peak memory, and a checkpoint you can trust.\n",
    "1. **Data parallelism (DDP)** replicates the model on every GPU, feeds each\n",
    "   a different slice of the batch, and averages gradients after every\n",
    "   backward pass. It is the entry point whenever model, gradients, and\n",
    "   optimizer state fit on one device.\n",
    "1. **Fully sharded data parallelism (FSDP)** shards parameters, gradients,\n",
    "   and optimizer state across GPUs and gathers each layer only while it is\n",
    "   being used. Same idea as DeepSpeed's ZeRO stage 3; it buys roughly a\n",
    "   world-size reduction in state memory at the price of all-gather traffic.\n",
    "1. **Tensor, pipeline, context, and expert parallelism** split individual\n",
    "   layers, layer groups, the sequence dimension, and mixture-of-experts\n",
    "   experts, respectively. Large systems compose several of these because no\n",
    "   single axis suffices at scale.\n",
    "\n",
    "In PyTorch, the first two rungs are a launcher and a wrapper:\n",
    "\n",
    "```bash\n",
    "torchrun --standalone --nproc-per-node=4 train.py\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=training-systems-md-2 fw=all -->\n",
    "\n",
    "Inside `train.py`, DDP wraps the model\n",
    "(`DistributedDataParallel(model)`), while FSDP's current API applies\n",
    "`fully_shard(model)` — the rewritten \"FSDP2\", which shards each parameter\n",
    "as a `DTensor` and has displaced the original FSDP wrapper in PyTorch's\n",
    "documentation and in downstream libraries. JAX expresses the same ideas\n",
    "through sharded arrays and named meshes: `jax.jit` with sharding\n",
    "specifications compiles the collectives for you, which is why the JAX code\n",
    "in this book has barely mentioned devices at all.\n",
    "\n",
    "One bookkeeping identity matters on every rung. The batch size that\n",
    "optimization sees is\n",
    "\n",
    "$$\n",
    "B_{\\textrm{global}} = B_{\\textrm{device}} \\cdot N_{\\textrm{ranks}} \\cdot N_{\\textrm{accumulation}},\n",
    "$$\n",
    "\n",
    "so changing the number of GPUs without adjusting per-device batch size or\n",
    "learning rate silently changes the optimization problem. Distributed\n",
    "samplers must not duplicate examples, and evaluation must aggregate\n",
    "metrics with the right denominators — the two classic sources of\n",
    "\"multi-GPU accuracy is different\" bugs.\n",
    "\n",
    "### Where the Memory Goes\n",
    "\n",
    "![Sharding, checkpointing, accumulation, and offload reduce different memory terms and introduce different costs.](../img/tools-training-memory.svg)\n",
    "\n",
    "Peak training memory is the sum of parameters, gradients, optimizer state,\n",
    "activations, communication buffers, and allocator overhead — for a 7B-\n",
    "parameter model trained in BF16 with Adam, roughly 14 + 14 + 56 GB of\n",
    "state before a single activation is stored. The toy model below is crude\n",
    "but sorts strategies correctly, and needs no GPU to run:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "training-systems-memory-model"
   },
   "source": [
    "terms_gib = {\"parameters\": 14.0, \"gradients\": 14.0, \"optimizer\": 56.0,\n",
    "             \"activations\": 18.0, \"temporary\": 5.0}\n",
    "\n",
    "def per_device_gib(world_size=1, shard_state=False, ckpt_activations=False):\n",
    "    divisor = world_size if shard_state else 1\n",
    "    state = sum(terms_gib[k] / divisor\n",
    "                for k in (\"parameters\", \"gradients\", \"optimizer\"))\n",
    "    activations = terms_gib[\"activations\"] * (0.35 if ckpt_activations else 1)\n",
    "    return state + activations + terms_gib[\"temporary\"]\n",
    "\n",
    "for label, cfg in [(\"1 GPU, plain\", (1, False, False)),\n",
    "                   (\"8 GPU DDP + ckpt\", (8, False, True)),\n",
    "                   (\"8 GPU FSDP + ckpt\", (8, True, True))]:\n",
    "    print(f\"{label:>20s}: {per_device_gib(*cfg):6.1f} GiB/device\")"
   ],
   "execution_count": 1,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "        1 GPU, plain:  107.0 GiB/device\n    8 GPU DDP + ckpt:   95.3 GiB/device\n   8 GPU FSDP + ckpt:   21.8 GiB/device\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=training-systems-md-3 fw=all -->\n",
    "\n",
    "Plain training does not fit any single GPU; DDP does not help, because it\n",
    "replicates state; FSDP brings the same model under 25 GiB per device. Each\n",
    "standard memory technique attacks one term and charges a different price:\n",
    "**mixed precision** (BF16 compute, FP8 on recent hardware) halves compute\n",
    "bytes; **gradient accumulation** shrinks activation memory but not state;\n",
    "**activation checkpointing** trades ~30% recomputation for most of the\n",
    "activation term; **CPU/NVMe offload** (DeepSpeed's ZeRO-Offload and\n",
    "descendants) trades transfer latency for capacity; **LoRA and QLoRA**\n",
    "sidestep the problem by making the optimizer state tiny — which is why a\n",
    "7B fine-tune fits on a 16 GB card. Combine techniques deliberately: four\n",
    "individually sensible optimizations can interfere with compilation,\n",
    "overlap, or numerics.\n",
    "\n",
    "Two topology rules from the Hugging Face *Ultra-Scale Playbook* (a free,\n",
    "experiment-backed guide distilled from 4,000+ runs on up to 512 GPUs, and\n",
    "the best single thing to read after this section) compress a lot of\n",
    "practice: keep tensor parallelism *inside* a node, where NVLink can carry\n",
    "its dense traffic — letting TP cross nodes loses several times as much\n",
    "throughput as letting pipeline parallelism cross — and map the\n",
    "highest-volume collectives to the fastest links available.\n",
    "\n",
    "## The Library Landscape\n",
    "\n",
    "What should you actually type? As of mid-2026 the ecosystem has settled\n",
    "into layers.\n",
    "\n",
    "**PyTorch built-ins.** `torchrun` + DDP for replication; `fully_shard`\n",
    "(FSDP2) for sharding; `DTensor` underneath as the common abstraction for\n",
    "sharded tensors; tensor-, context-, and pipeline-parallel APIs exist but\n",
    "are still marked experimental. These primitives are what most higher-level\n",
    "tools now generate.\n",
    "\n",
    "**Hugging Face stack.** `accelerate` launches the same script on DDP,\n",
    "FSDP2, or DeepSpeed with a config file rather than code changes;\n",
    "`transformers.Trainer` sits on top of it; `peft` implements LoRA/QLoRA;\n",
    "and `trl` provides post-training — supervised fine-tuning, DPO, and the\n",
    "GRPO family of reinforcement-learning trainers that took off when\n",
    "DeepSeek-R1 showed pure RL could buy reasoning. For models up to roughly\n",
    "30B parameters on a handful of nodes, this stack is the default answer.\n",
    "\n",
    "**DeepSpeed.** ZeRO stages 1–3 (optimizer, gradient, parameter sharding)\n",
    "plus CPU/NVMe offload. It pioneered the sharding ideas that FSDP absorbed\n",
    "and remains actively maintained, but for new PyTorch projects it has\n",
    "ceded the default slot to FSDP2; reach for it when you need its offload\n",
    "machinery to squeeze a large model onto scarce memory, or because a tool\n",
    "you use (notably several RLHF frameworks) builds on it.\n",
    "\n",
    "**Megatron-Core and TorchTitan.** The serious pretraining tier. NVIDIA's\n",
    "Megatron-Core implements tensor/pipeline/sequence/expert parallelism with\n",
    "FP8 support and powers many industrial labs (and NVIDIA's own NeMo\n",
    "framework and Nemotron models). TorchTitan is the PyTorch-native\n",
    "equivalent: a clean reference stack composing FSDP, tensor, pipeline, and\n",
    "context parallelism plus expert parallelism for MoE, demonstrated at\n",
    "1,000-GPU scale on models from Llama 3 405B to DeepSeek-V3. Choose this\n",
    "tier only when you are genuinely pretraining — the configuration surface\n",
    "is proportional to its power.\n",
    "\n",
    "**Fine-tuning frontends.** Unsloth owns the single-GPU niche with fused\n",
    "kernels and quantized training (roughly 2× speed and large memory savings\n",
    "for LoRA/QLoRA work). Axolotl drives full and parameter-efficient\n",
    "fine-tunes across a node or several from one YAML file, with FSDP2 and\n",
    "DeepSpeed backends. LLaMA-Factory covers a similar space with a GUI and\n",
    "very broad model support. (Its former peer `torchtune` wound down in 2025\n",
    "— check a library's pulse before adopting it; this landscape churns.)\n",
    "\n",
    "**JAX.** MaxText is the reference for TPU (and GPU) pretraining and now\n",
    "post-training; Levanter/Marin demonstrated fully reproducible open\n",
    "pretraining. On TPUs the XLA compiler handles the parallelism that the\n",
    "PyTorch world configures by hand — the trade is less knob-turning for\n",
    "less low-level control.\n",
    "\n",
    "**RL post-training at scale.** Beyond TRL's single-node comfort zone, veRL\n",
    "is the most widely adopted open framework (FSDP or Megatron for training,\n",
    "vLLM or SGLang for rollouts), with OpenRLHF as a leaner Ray+DeepSpeed\n",
    "alternative. This corner moves fastest of all; expect its details to age\n",
    "first.\n",
    "\n",
    "### What to Use at Which Scale\n",
    "\n",
    "<!-- tbl-caption:Training tools by scale (mid-2026) -->\n",
    "\n",
    "| Scale | Typical job | Reach for |\n",
    "|---|---|---|\n",
    "| 1 GPU, ≥ 8 GB | LoRA/QLoRA fine-tune ≤ 8B | Unsloth, or PEFT + Trainer |\n",
    "| 1 node, 2–8 GPUs | full fine-tune ≤ 70B, DDP/FSDP2 | Accelerate or Axolotl |\n",
    "| few nodes | large fine-tune, small pretrain | FSDP2 + torchrun, DeepSpeed |\n",
    "| many nodes | serious pretraining, MoE | Megatron-Core, TorchTitan, MaxText |\n",
    "| RL post-training | GRPO/DPO pipelines | TRL (small), veRL/OpenRLHF (large) |\n",
    "\n",
    "The table is a starting point, not a verdict — but note its shape: you buy\n",
    "generality with configuration burden as you move down, and most projects\n",
    "live in the first two rows. If a workload fits a simpler row, use the\n",
    "simpler row.\n",
    "\n",
    "## Keeping a Long Run Alive\n",
    "\n",
    "Everything above buys throughput; what turns throughput into results is\n",
    "operational discipline, because at scale *something* fails every few\n",
    "hours.\n",
    "\n",
    "* **Checkpoint completely and atomically.** A resumable checkpoint holds\n",
    "  model and optimizer state, the learning-rate schedule and step count,\n",
    "  data-loader position, and RNG state. Write to a temporary path and\n",
    "  rename on completion, so a crash mid-write cannot destroy the previous\n",
    "  checkpoint; keep more than one until restore has been tested. Sharded\n",
    "  (per-rank) checkpoints avoid gathering a model larger than any single\n",
    "  host's memory.\n",
    "* **Drill the recovery.** Kill the job deliberately, restore on fresh\n",
    "  workers, and check that loss, step count, and data position continue as\n",
    "  if nothing happened. Untested recovery is a rumor, and spot-priced\n",
    "  training (that section) is only cheap if recovery\n",
    "  actually works.\n",
    "* **Feed the accelerators.** Profile data loading separately from compute:\n",
    "  storage reads, decoding and augmentation, tokenization, and host-to-\n",
    "  device transfer. A starved GPU shows high nominal utilization and low\n",
    "  throughput; adding more GPUs to a data-bound job makes it *worse*.\n",
    "* **Watch convergence, not just speed.** Tokens per second without a loss\n",
    "  curve rewards fast wrong runs. Log both, plus time spent in collectives\n",
    "  and checkpoint duration; debug distributed failures by reproducing on\n",
    "  one process, then one node, then many, with rank and host on every log\n",
    "  line.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* Scale in order — DDP while state fits, FSDP2/ZeRO when it does not,\n",
    "  tensor/pipeline/expert parallelism only for genuine pretraining — from a\n",
    "  measured single-GPU baseline.\n",
    "* Peak memory is a sum of terms; each technique (sharding, checkpointing,\n",
    "  accumulation, offload, LoRA) removes one term at a known price.\n",
    "* In 2026 the default stacks are: Unsloth/PEFT on one GPU, Accelerate or\n",
    "  Axolotl with FSDP2 on a node, Megatron-Core/TorchTitan/MaxText for\n",
    "  pretraining, TRL then veRL for RL post-training.\n",
    "* Keep tensor parallelism inside a node; map heavy collectives to fast\n",
    "  links.\n",
    "* Long runs survive on atomic checkpoints, tested recovery, a fed input\n",
    "  pipeline, and convergence metrics — not on optimism.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Extend the memory model with a communication-buffer term and a LoRA\n",
    "   configuration (frozen base weights, small trainable adapter). At what\n",
    "   world size does FSDP stop paying for a 7B model?\n",
    "1. For a fixed global batch of 4M tokens, enumerate three combinations of\n",
    "   device batch, world size, and accumulation steps, and reason about\n",
    "   their relative throughput and optimizer behavior.\n",
    "1. Take a training notebook of roughly the scale of\n",
    "   that section and run it once with\n",
    "   `torchrun --nproc-per-node=2`. Measure the actual speedup over one GPU\n",
    "   and explain the gap from 2×.\n",
    "1. Design (on paper) the checkpoint contents and restore protocol for a\n",
    "   sharded training job that must resume with a *different* number of\n",
    "   workers. Which parts of the state are per-rank, and which are global?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
