Distributed training systems: frameworks, parallelism, sharding, and the production details that make it work

Field notes on the question underneath this whole chapter: how do you train a model too large for one GPU, one node, or even one cluster, without losing correctness or efficiency? Same standard as the rest of this series -- intuition, then math traced through a concrete number, then the actual framework code a team runs, then the failure modes and how to debug them. Where the deep math for a mechanism is already fully derived in the companion documents (FSDP/ZeRO memory arithmetic, network topology, activation-checkpointing tradeoffs -- all in the LLM lifecycle notes, chapters 1 and 8), this chapter cross-references rather than re-deriving it, and spends its own depth on what those documents don’t cover: the actual framework APIs by name, sequence and expert parallelism, resharding, and -- the section with the least existing coverage anywhere -- how you actually debug a distributed training job when it’s silently wrong rather than loudly crashed.

00 · The production mindset#

❓ Why this framing. A single GPU has, at frontier scale, roughly four orders of magnitude less memory than a large model’s full training state needs (lifecycle notes, ch. 8: a 70B model’s weights+gradients+optimizer state alone is over 1TB). Every framework and technique in this chapter exists to answer one question in a slightly different way: which piece of the problem gets split across which boundary -- GPUs within a node, nodes within a cluster, or time (recompute now, save memory later) -- and what does that split cost in communication or compute to undo when the next operation needs the whole picture back?

$$\text{training a model this large} = \underbrace{\text{split it}}_{\text{parallelism, ch. 2}} + \underbrace{\text{split its state}}_{\text{sharding, ch. 3}} + \underbrace{\text{trade memory for compute where splitting isn't enough}}_{\text{ch. 4}} + \underbrace{\text{survive the cluster actually running for weeks}}_{\text{ch. 5–6}}$$

Nothing in this chapter is optional at frontier scale -- it’s not “use DeepSpeed if you like it,” it’s “the model doesn’t fit without some combination of these,” and the combination that’s right for a 7B model on 8 GPUs is usually wrong for a 405B model on 16,000.

If you’ve never done any of this before: the whole idea in one page, no jargon#

Skip this box if the framing above already made sense. If it didn’t, start here -- everything later in this document is a detailed, technical version of the following six sentences.

A GPU is a chip built to do one simple kind of math (multiply and add numbers) an enormous number of times at once, instead of doing complicated things one at a time the way your laptop’s regular processor does -- that’s why it’s the right tool for training a neural network, which really is just an enormous pile of multiply-and-add operations. “Parameters” are the individual adjustable numbers inside the model -- a “7 billion parameter model” just means there are 7 billion little dials, and “training” means repeatedly showing the model examples and nudging every single one of those dials very slightly, millions of times, until the model’s outputs get better. Those dials, plus some bookkeeping numbers the training process needs to decide how to nudge each dial, all have to sit in the GPU’s own fast memory (VRAM) at the same time -- think of VRAM as a desk: you can only fit so many papers on it before you run out of surface. A single high-end GPU today has somewhere around 80 gigabytes of desk space. A 7-billion-parameter model’s dials, plus its bookkeeping numbers, need well over 100GB just to get started -- it doesn’t fit on one desk. A 70-billion or 400-billion parameter model needs over a terabyte. There is no single GPU with a terabyte of memory, so there is no way to train a model that size except to somehow spread the work across many GPUs’ desks at once -- which is the entire subject of this document. Two different spreading strategies exist and they’re not the same thing, which is worth getting straight before anything else: parallelism (ch. 2) is splitting up the work -- different GPUs process different chunks of data, or different pieces of the calculation, at the same time, like several cooks each stirring a different pot. Sharding (ch. 3) is splitting up the storage -- instead of every GPU keeping a full copy of every dial and every bookkeeping number, each GPU keeps only its own slice, and they exchange pieces with each other exactly when needed, like several people sharing one filing cabinet by each keeping a different drawer at their own desk instead of everyone photocopying the whole cabinet. Real large-model training always does both at once, which is why the frameworks in chapter 1 (DeepSpeed, FSDP, and the rest) exist -- they’re the software that actually manages “which GPU has which piece of which dial, and how do we get the pieces to talk to each other fast enough that it’s still worth doing this at all.”

01 · Core frameworks#

❓ Why this topic. Every framework in this section solves the same underlying problem -- get gradients synchronized and memory sharded across many GPUs -- but they disagree on how much control you give up for how much boilerplate disappears, and that tradeoff is the actual decision a team is making when it picks one.

The decision table, before any code#

framework	what it actually gives you	you give up	pick it when
PyTorch DDP	the floor: each GPU holds a full model copy, gradients synced via all-reduce after backward	no memory sharding at all -- model must fit on one GPU	model fits on one GPU comfortably; you just want more throughput via more data in parallel
FSDP	full parameter/gradient/optimizer sharding (ZeRO-3-equivalent), native to PyTorch	more complex wrapping rules (which modules get sharded together) than DDP	you’re in the PyTorch ecosystem already and want ZeRO-3-class memory savings without adding a second framework
DeepSpeed	ZeRO stages 1–3, CPU/NVMe offload (ZeRO-Infinity), 3D parallelism, its own optimized fused kernels	a second framework’s config surface and its own launch tooling	you need offload to CPU/NVMe (fitting a model that doesn’t fit in GPU memory even after sharding) or Megatron-style 3D parallelism out of the box
Hugging Face Accelerate	a thin, unified wrapper letting the same script run under DDP, FSDP, or DeepSpeed depending on a config file, no code branching	some of the fine-grained control each backend offers directly	you want to write one training script and choose the backend at launch time, not at code-authoring time
PyTorch Lightning / Fabric	Accelerate’s idea plus a structured training loop (Lightning) or a minimal, “keep your own loop” wrapper (Fabric), both with DDP/FSDP/DeepSpeed strategies built in	Lightning specifically imposes its own training-loop structure (training_step, configure_optimizers)	you want a maintained, batteries-included training loop (Lightning) or the thinnest possible wrapper around a loop you already have (Fabric)
Ray Train	orchestration above any of the above -- launches many workers, each running your DDP/FSDP/DeepSpeed script, handles fault-tolerant restart and autoscaling across a cluster	it’s an orchestration layer, not a parallelism strategy itself -- you still pick DDP/FSDP/DeepSpeed underneath it	you need to run training across a heterogeneous or autoscaling cluster, or you’re already using Ray for data processing/serving and want one system for all of it
Horovod	framework-agnostic (PyTorch, TensorFlow, MXNet) all-reduce-based data parallelism, historically the fastest ring-all-reduce implementation before NCCL matured	most new PyTorch-only projects have less reason to reach for it now that DDP/NCCL cover the same ground natively	a mixed-framework organization, or infrastructure already standardized on Horovod’s launcher (horovodrun)
NCCL	not a training framework -- the actual communication library every one of the above eventually calls into for GPU-to-GPU collectives (all-reduce, all-gather, broadcast)	nothing; it’s the substrate, not a choice	always present under the hood; you interact with it directly only when debugging (ch. 6)
FairScale	Meta’s earlier sharding/parallelism primitives library (FullyShardedDataParallel, pipeline parallel, tensor parallel utilities) built directly on torch.distributed	largely superseded -- native PyTorch FSDP grew directly out of FairScale’s implementation and is now the maintained path	you’re reading or maintaining an older codebase that predates native FSDP; not the right choice for a new project today
Megatron-LM / NeMo	NVIDIA’s purpose-built tensor+pipeline+sequence parallelism implementation for transformer pretraining specifically, with fused, hand-optimized CUDA kernels for attention and the MLP block	narrower scope than DeepSpeed/FSDP -- built and tuned specifically for transformer-shaped pretraining at the largest scale, not general-purpose	you’re pretraining a large transformer from scratch and want the most heavily hardware-optimized TP/PP implementation available, ideally on NVIDIA hardware with NeMo’s higher-level recipe layer on top

Where Ray Train sits inside “Ray AIR.” Ray Train is one component of the broader Ray AI Runtime (Ray AIR) -- the umbrella that also includes Ray Tune (hyperparameter search across many trial runs), Ray Data (distributed data loading and preprocessing, directly relevant to chapter 9’s parallel-I/O problem), and Ray Serve (model serving, relevant to chapter 7). The practical reason this distinction matters: a team that picks Ray specifically for training orchestration often ends up using Ray Data for the data pipeline and Ray Serve for deployment too, precisely because they share the same underlying cluster and object store -- the appeal of “one system for all of it” from the table above is Ray AIR’s actual pitch, with Ray Train being only the training-specific slice of it.

The one-line version of this whole table, as a strategy, not a list. Start with DDP. The moment the model doesn’t fit, move to FSDP or DeepSpeed ZeRO -- whichever integrates more naturally with the rest of your stack (native PyTorch vs. a config-driven second framework). The moment sharding alone still doesn’t fit, or throughput plateaus below what the hardware should allow, add tensor and pipeline parallelism -- via DeepSpeed’s built-in 3D parallelism if you’re already there, or Megatron-LM/NeMo if you’re pretraining a transformer from scratch and want the most optimized kernels available. Reach for Ray Train only when the orchestration problem (not the parallelism problem) is what’s unsolved -- a cluster that needs to autoscale or tolerate heterogeneous hardware. This progression -- DDP → FSDP/ZeRO → 3D parallelism -- is worth internalizing as the default path, because skipping straight to the most complex tool before the simpler one has actually been shown to be insufficient is itself a common, avoidable production mistake.

The four communication primitives, defined plainly, once#

Every framework above is, underneath its own API, issuing some combination of exactly four operations. Worth having these straight before anything else, since “all-reduce” and “all-gather” get used constantly from here on without re-explaining each time:

• Broadcast: one GPU has a piece of data, every other GPU needs an identical copy of it. Used once at the start of DDP/Horovod training to make sure every worker starts from the same initial weights (hvd.broadcast_parameters in the Horovod code below is exactly this).
• All-reduce: every GPU has its own version of some tensor (e.g. its own locally-computed gradient), and every GPU needs the combined result (e.g. the sum or average across all of them) -- DDP’s gradient synchronization after backward() is an all-reduce: everyone contributes, everyone receives the same combined answer.
• All-gather: every GPU has a different piece of a whole (a shard), and every GPU needs the full, reassembled thing. FSDP’s parameter reconstruction right before a layer’s forward pass (ch. 3) is an all-gather: each GPU holds 1/8th of a parameter, all-gather hands every GPU the complete parameter, temporarily.
• Reduce-scatter: the reverse of all-gather, combined with a reduce -- every GPU has a full-sized tensor (e.g. a full gradient), and each GPU ends up owning only its own shard of the combined result, rather than the full thing. This is exactly how ZeRO-2/3 (ch. 1, 3) get a sharded gradient without ever materializing the full combined gradient on every GPU at once.

The practical rule that follows directly from these definitions: broadcast is cheap and rare (once, at startup); all-reduce and reduce-scatter cost scales with tensor size and happen every step, which is why they’re the ones chapter 4’s overlap techniques target; all-gather is what makes FSDP’s “the full parameter only exists briefly, right when needed” property possible at all.

PyTorch DDP: the floor everything else builds on#

import torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def train(rank, world_size, model, dataset):
    setup(rank, world_size)
    model = model.to(rank)
    model = DDP(model, device_ids=[rank])           # wraps forward/backward with automatic gradient all-reduce
    sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = torch.utils.data.DataLoader(dataset, sampler=sampler, batch_size=32)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
    for batch in loader:
        optimizer.zero_grad()
        loss = model(batch.to(rank))
        loss.backward()                              # DDP triggers all-reduce here, overlapped with backward compute
        optimizer.step()

torchrun --nproc_per_node=8 --nnodes=1 train_ddp.py

DeepSpeed: ZeRO stages via one JSON config, no model code changes#

{
  "train_batch_size": 256,
  "gradient_accumulation_steps": 4,
  "optimizer": {"type": "AdamW", "params": {"lr": 3e-5, "betas": [0.9, 0.95], "weight_decay": 0.1}},
  "fp16": {"enabled": false},
  "bf16": {"enabled": true},
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu", "pin_memory": true},
    "offload_param": {"device": "none"},
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e8
  }
}

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(), config="ds_config.json"
)
for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)     # DeepSpeed handles ZeRO-sharded gradient computation and reduction internally
    model_engine.step()

deepspeed --num_gpus=8 train_deepspeed.py --deepspeed --deepspeed_config ds_config.json

The offload_optimizer block is ZeRO-Offload/ZeRO-Infinity (ch. 3): once GPU sharding alone still doesn’t fit, this spills optimizer state (and, with offload_param, parameters too) to CPU RAM or NVMe -- the correctness-preserving last resort when even ZeRO-3’s 16N/P (lifecycle notes, ch. 8) doesn’t clear the GPU’s memory ceiling.

DeepSpeed’s own pipeline and tensor parallelism, not just ZeRO#

ZeRO (above) is a memory-sharding strategy layered on top of whatever architecture you hand DeepSpeed -- it doesn’t split the model across devices by layer or by weight matrix on its own. DeepSpeed separately ships its own pipeline parallelism (PipelineModule) and, more recently, tensor parallelism (AutoTP), usable instead of or alongside ZeRO:

from deepspeed.pipe import PipelineModule

# split an existing nn.Sequential-like model into pipeline stages, DeepSpeed handles the microbatch scheduling
model = PipelineModule(
    layers=[layer1, layer2, layer3, layer4, layer5, layer6, layer7, layer8],   # your model's layers, in order
    num_stages=4,               # 8 layers split across 4 pipeline stages -> 2 layers/stage
    loss_fn=my_loss_function,
    partition_method="parameters",   # balance stages by parameter count rather than naive equal layer count (ch. 3)
)

{
  "train_batch_size": 256,
  "pipeline": {"stages": 4},
  "tensor_parallel": {"autotp_size": 2},
  "zero_optimization": {"stage": 1}
}

That last config line is the real point: DeepSpeed’s own “3D parallelism” (the term in the framework table above) means running ZeRO (for optimizer/gradient sharding), PipelineModule (for layer-wise pipeline stages), and AutoTP (for within-stage tensor parallelism) together, all driven from one config -- the same DP × TP × PP combination Megatron-LM implements with its own kernels, available here as a config-driven option inside a framework you may already be using for ZeRO alone.

Hugging Face Accelerate: one script, backend chosen at launch time#

from accelerate import Accelerator

accelerator = Accelerator(mixed_precision="bf16", gradient_accumulation_steps=4)
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)   # backend-agnostic wrapping

for batch in dataloader:
    with accelerator.accumulate(model):
        loss = model(batch)
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

accelerate config     # interactive: choose DDP / FSDP / DeepSpeed, num GPUs, mixed precision -- writes a YAML
accelerate launch train_accelerate.py    # same script, whatever backend the config file specifies

The entire point: the training script above never mentions DDP, FSDP, or DeepSpeed by name -- switching backends is a accelerate config re-run, not a code change, which is exactly the “one script, pick the backend at launch time” tradeoff from the table above.

PyTorch Lightning and Fabric: structured loop vs. minimal wrapper#

# Lightning: the framework owns the loop, you fill in the steps
import lightning as L

class LanguageModel(L.LightningModule):
    def training_step(self, batch, batch_idx):
        loss = self.model(batch)
        self.log("train_loss", loss)                 # automatic, backend-agnostic metric logging
        return loss
    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters(), lr=3e-4)

trainer = L.Trainer(accelerator="cuda", devices=8, strategy="fsdp", precision="bf16-mixed")
trainer.fit(LanguageModel(), train_dataloader)

# Fabric: you keep your own loop, Fabric just handles device/precision/strategy plumbing
import lightning as L

fabric = L.Fabric(accelerator="cuda", devices=8, strategy="fsdp", precision="bf16-mixed")
fabric.launch()
model, optimizer = fabric.setup(model, optimizer)
dataloader = fabric.setup_dataloaders(dataloader)
for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch)
    fabric.backward(loss)                              # replaces loss.backward() -- Fabric handles strategy-specific sync
    optimizer.step()

Ray Train: orchestration above the parallelism strategy#

from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig, RunConfig, CheckpointConfig

def train_loop_per_worker(config):
    model, optimizer, dataloader = build_training_objects(config)   # your usual DDP/FSDP setup, runs on each worker
    for epoch in range(config["num_epochs"]):
        for batch in dataloader:
            loss = train_step(model, optimizer, batch)
        ray.train.report({"loss": loss.item()})         # metrics + fault-tolerant checkpoint reporting

trainer = TorchTrainer(
    train_loop_per_worker,
    train_loop_config={"num_epochs": 10, "lr": 3e-4},
    scaling_config=ScalingConfig(num_workers=8, use_gpu=True),
    run_config=RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=3)),
)
result = trainer.fit()     # Ray launches 8 worker processes, sets up the torch distributed group, handles restarts

Ray Train’s train_loop_per_worker is the exact same DDP/FSDP code from earlier in this chapter -- Ray’s contribution is entirely the layer around it: launching workers across a cluster (possibly autoscaling, possibly spanning multiple machines with heterogeneous hardware), and automatically retrying a worker that dies rather than failing the whole job, which is a direct answer to chapter 5’s failure-recovery requirement.

Horovod: framework-agnostic ring-allreduce#

import horovod.torch as hvd

hvd.init()
torch.cuda.set_device(hvd.local_rank())
optimizer = torch.optim.SGD(model.parameters(), lr=0.01 * hvd.size())    # scale LR with worker count, standard practice
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
hvd.broadcast_parameters(model.state_dict(), root_rank=0)                # ensure every worker starts identical

for data, target in train_loader:
    optimizer.zero_grad()
    loss = F.nll_loss(model(data), target)
    loss.backward()
    optimizer.step()          # DistributedOptimizer wraps step() to all-reduce gradients via Horovod's ring-allreduce

horovodrun -np 8 -H localhost:8 python train_horovod.py

NCCL: what’s actually running underneath every one of the above#

Every collective operation above -- DDP’s gradient all-reduce, FSDP’s parameter all-gather, DeepSpeed’s reduce-scatter -- eventually calls into NCCL, NVIDIA’s GPU-to-GPU communication library. You rarely call it directly, but its environment variables are the first thing to reach for when a distributed job hangs or underperforms (ch. 6):

NCCL_DEBUG=INFO           # print which communication backend/topology NCCL actually chose -- the first debug step
NCCL_DEBUG_SUBSYS=ALL     # verbose: which specific collective, which ring/tree algorithm, at what size
NCCL_SOCKET_IFNAME=eth0   # force the network interface NCCL uses -- fixes silent hangs on multi-NIC machines
NCCL_IB_DISABLE=1         # force NCCL off InfiniBand onto TCP/sockets -- a diagnostic step, not a real fix, to isolate whether IB itself is the problem
NCCL_P2P_LEVEL=NVL        # confirm/force NVLink peer-to-peer within a node rather than falling back to a slower path

FairScale: the historical predecessor worth recognizing, not adopting fresh#

# FairScale's FSDP -- the same core idea native torch FSDP later absorbed
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP_FairScale

model = FSDP_FairScale(model, mixed_precision=True, flatten_parameters=True)
# API deliberately looks almost identical to torch.distributed.fsdp -- that's not a coincidence,
# native FSDP's design directly descends from this implementation

If you encounter this import in a codebase, treat it as a signal the code predates PyTorch 2.0-era native FSDP -- the correct move for new work is almost always porting to torch.distributed.fsdp’s fully_shard, not learning FairScale’s API as a fresh skill.

Megatron-LM / NeMo: purpose-built for transformer pretraining at the largest scale#

# Megatron-LM: tensor + pipeline + sequence parallelism, configured via launch flags, not a Python API you wrap
python pretrain_gpt.py \
    --tensor-model-parallel-size 8 \
    --pipeline-model-parallel-size 16 \
    --sequence-parallel \
    --num-layers 96 --hidden-size 12288 --num-attention-heads 96 \
    --micro-batch-size 1 --global-batch-size 1536 \
    --fp16 \
    --use-flash-attn

# NeMo: the higher-level recipe layer on top of Megatron-LM's kernels -- a config file, not code, drives the run
trainer:
  devices: 8
  num_nodes: 16
model:
  tensor_model_parallel_size: 8
  pipeline_model_parallel_size: 16
  sequence_parallel: true
  micro_batch_size: 1
  global_batch_size: 1536
  optim:
    name: distributed_fused_adam     # a fused, ZeRO-1-style optimizer specific to NeMo/Megatron's own kernels

The practical distinction from DeepSpeed worth remembering: DeepSpeed’s ZeRO stages are primarily a memory-sharding strategy that layers on top of whatever model architecture you hand it, while Megatron-LM’s tensor/pipeline/sequence parallelism is implemented with custom fused kernels specific to the transformer block itself -- which is why frontier pretraining runs (Llama, DeepSeek, and others’ internal infrastructure) commonly build on Megatron-style kernels for the parallelism layer while borrowing ZeRO-style sharding ideas for the optimizer state, rather than treating the two as strictly interchangeable options.

📚 Further reading

• Li et al., 2020 -- PyTorch Distributed: Experiences on Accelerating Data Parallel Training
• Rasley et al., 2020 -- DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters
• Sergeev & Del Balso, 2018 -- Horovod: Fast and Easy Distributed Deep Learning in TensorFlow
• Shoeybi et al., 2019 -- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• NVIDIA NeMo Framework documentation -- docs.nvidia.com/nemo-framework
• Meta AI’s FairScale repository -- github.com/facebookresearch/fairscale
• NVIDIA NCCL documentation -- docs.nvidia.com/deeplearning/nccl

02 · Parallel training types#

❓ Why this topic. “Parallelism” isn’t one technique -- it’s a family, each member splitting a different axis of the problem (batch, weight matrices, layers, sequence position, or expert routing), and real frontier training combines several simultaneously because no single axis has enough room to split a model this large across enough devices alone.

Data, tensor, and pipeline parallelism: the recap#

Data parallelism (every GPU holds a full model copy, different data, gradients synced -- DDP, ch. 1), tensor parallelism (split individual weight matrices column/row-wise across GPUs within a node), and pipeline parallelism (different layers on different GPUs, with the bubble-fraction cost $\frac{p-1}{m}$) are covered with full memory/communication math and worked examples in the LLM lifecycle companion notes, chapter 1 -- this section builds on that rather than repeating it.

Sequence parallelism: splitting the one axis TP and PP don’t touch#

The gap TP and PP leave open. Tensor parallelism splits weight matrices; pipeline parallelism splits layers. Neither splits the sequence length dimension itself -- every GPU in a tensor-parallel group still holds the full activation tensor for every token in a long sequence, which becomes the binding memory constraint specifically at long context lengths (the companion long-context document’s whole chapter 1 problem, now viewed from the training-memory side rather than the inference-compute side).

Two real implementations, different tradeoffs. Ring self-attention (used in Megatron-LM’s sequence parallelism, and the same underlying idea as Ring Attention from the long-context companion document, ch. 1) shards the sequence dimension across GPUs and passes K/V blocks around a ring so every query chunk eventually attends to every key/value chunk, overlapping communication with compute. DeepSpeed-Ulysses (Jacobs et al., 2023) takes a different partition: shard the sequence dimension across GPUs for the non-attention parts of a transformer block (each GPU processes its own sequence chunk independently through the projections), then use an all-to-all communication step specifically around the attention computation to temporarily regroup by attention head instead of sequence position -- since each GPU only needs a subset of heads’ full sequence, not the full sequence for all heads, the communication volume scales with the number of GPUs rather than the sequence length itself, which the Ulysses paper reports as fundamentally more favorable at very long sequence lengths than ring-style communication whose volume grows with sequence length directly.

# DeepSpeed-Ulysses sequence parallelism -- config-level enablement, no attention code rewrite needed
import deepspeed
from deepspeed.sequence.layer import DistributedAttention

class ParallelTransformerBlock(torch.nn.Module):
    def __init__(self, attn, sp_group):
        super().__init__()
        # wraps ordinary attention with the all-to-all head/sequence regroup Ulysses depends on
        self.attn = DistributedAttention(attn, sequence_process_group=sp_group)

    def forward(self, x):
        return self.attn(x)   # sequence-sharded in, sequence-sharded out; regroup-by-head happens internally

Expert parallelism: sharding by which expert, not by layer or matrix#

Why MoE needs its own parallelism axis. A Mixture-of-Experts layer (transformer notes, ch. 5 covers the routing mechanism) has many more total parameters than any single token’s forward pass actually uses -- DeepSeek-V3’s 671B total / 37B active parameters (lifecycle notes, ch. 7) is the extreme real example. Tensor parallelism could split each expert’s weight matrix, but a simpler and more communication-efficient approach for MoE specifically is expert parallelism: place different whole experts on different GPUs, and route each token to whichever GPU holds the expert its router selected -- turning the routing decision into a network communication pattern (an all-to-all: every GPU sends some tokens out to other GPUs’ experts and receives tokens from other GPUs in the same step).

$$\text{all-to-all communication volume per step} \propto (\text{tokens per GPU}) \times (\text{hidden size}) \times (\text{top-}k\text{ experts routed to})$$

The real cost, and why DualPipe (lifecycle notes, ch. 7) exists. This all-to-all is fundamentally cross-node traffic once experts are spread across more GPUs than fit in one node -- exactly the communication DeepSeek-V3’s DualPipe pipeline scheduler was built to overlap with useful compute rather than let it stall the pipeline. A load-balancing problem rides on top of this: if the router sends most tokens to a small subset of experts, those GPUs become both compute- and communication-bottlenecked while others sit idle -- which is precisely why DeepSeek-V3’s own “auxiliary-loss-free load balancing” (lifecycle notes, ch. 7) exists as a training-time mechanism, not just a serving-time convenience.

# Minimal expert-parallel MoE forward pass: route, all-to-all dispatch, compute, all-to-all combine
import torch, torch.distributed as dist

def expert_parallel_forward(x, router, local_expert, ep_group, world_size):
    # x: (tokens, d_model) on this GPU; router picks which GPU's expert each token should go to
    dest_gpu = router(x).argmax(-1)                                  # (tokens,) -- top-1 routing for simplicity
    sorted_idx = dest_gpu.argsort()
    x_sorted = x[sorted_idx]
    counts = torch.bincount(dest_gpu, minlength=world_size)
    send_splits = counts.tolist()

    recv_splits = [None] * world_size
    dist.all_to_all_single(torch.tensor(recv_splits), torch.tensor(send_splits), group=ep_group)  # exchange counts
    x_received = torch.empty(sum(recv_splits), x.size(-1), device=x.device)
    dist.all_to_all_single(x_received, x_sorted, output_split_sizes=recv_splits,
                            input_split_sizes=send_splits, group=ep_group)   # the real token dispatch

    y_local = local_expert(x_received)                                # this GPU's expert processes tokens routed to it

    y_sorted = torch.empty_like(x_sorted)
    dist.all_to_all_single(y_sorted, y_local, output_split_sizes=send_splits,
                            input_split_sizes=recv_splits, group=ep_group)   # send results back to origin GPUs
    y = torch.empty_like(y_sorted)
    y[sorted_idx] = y_sorted
    return y

Hybrid parallelism: 3D and 4D, and why the order of combination matters#

“3D parallelism,” defined plainly, before generalizing. The term refers specifically to running data parallelism × tensor parallelism × pipeline parallelism simultaneously -- three independent axes, each splitting a different dimension of the problem (which data, which weight-matrix slice, which layer), composed into a single training job. It’s called out by name (rather than just “combining parallelism strategies”) because this exact triple is what Megatron-LM and DeepSpeed’s own 3D-parallelism config (chapter 1) both standardize on as the default combination for transformer pretraining beyond single-node scale -- sequence/context parallelism and expert parallelism are later additions on top of this base triple, which is why the field’s real, largest runs are sometimes described as “4D” rather than 3D: the same three axes, plus one more.

Real frontier training never uses one axis alone. The standard combination, in the order it’s applied (outermost to innermost):

$$\text{DP} \times \text{PP} \times \text{TP} \times \text{(SP or EP, if needed)} = \text{total GPU count}$$

Worked example. Llama 3 405B’s reported 4D parallelism (lifecycle notes, ch. 1) across roughly 16,000 H100s combines data parallelism (many independent replicas of the whole TP×PP×SP setup, each seeing different data), pipeline parallelism (splitting layers across GPUs within a replica), tensor parallelism (splitting weight matrices within a pipeline stage, confined to NVLink-connected GPUs within a node per the bandwidth argument in the lifecycle notes’ chapter 8), and context/sequence parallelism (splitting long sequences within a stage). The order matters because of the bandwidth hierarchy from that same chapter: TP’s frequent, small communications go innermost (within-node NVLink), PP’s boundary hand-offs go next (can tolerate cross-node latency better), and DP’s comparatively rare full-model gradient synchronization is the outermost, most bandwidth-tolerant axis -- placing any of these in the wrong position (e.g. TP spanning nodes) reintroduces the exact communication bottleneck the whole hierarchy was designed to avoid.

Gradient accumulation: simulating a bigger batch without more memory#

The problem it solves. A desired global batch size (chosen for training stability, per the gradient-noise-scale reasoning in the lifecycle notes’ chapter 7) may require more activation memory than fits, even after all the parallelism above. Gradient accumulation runs several forward/backward passes on smaller micro-batches, summing (not applying) their gradients, and only calls optimizer.step() after enough micro-batches have accumulated to match the target global batch size -- mathematically identical to one large-batch step, at the cost of doing the forward/backward compute serially rather than in one larger parallel batch.

accumulation_steps = 4          # simulate a 4x larger batch than one micro-batch alone allows
optimizer.zero_grad()
for i, micro_batch in enumerate(dataloader):
    loss = model(micro_batch) / accumulation_steps    # scale down so summed gradients match a true large-batch average
    loss.backward()                                     # gradients accumulate in .grad, not overwritten, across calls
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

⚠️ The detail that silently breaks distributed correctness. Under DDP, every .backward() call triggers a gradient all-reduce by default -- calling it on every micro-batch when only the last one should trigger synchronization wastes communication on every accumulation step. The fix is DDP’s no_sync() context manager (or Fabric’s no_backward_sync, ch. 1), which suppresses the all-reduce on every micro-batch except the final one in the accumulation window:

for i, micro_batch in enumerate(dataloader):
    sync_now = (i + 1) % accumulation_steps == 0
    context = model.no_sync() if not sync_now else contextlib.nullcontext()
    with context:
        loss = model(micro_batch) / accumulation_steps
        loss.backward()                    # all-reduce suppressed except on the final micro-batch of the window
    if sync_now:
        optimizer.step()
        optimizer.zero_grad()

📚 Further reading

• Jacobs et al., 2023 -- DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models
• Liu et al., 2023 -- Ring Attention with Blockwise Transformers (companion long-context document, ch. 1)
• Shoeybi et al., 2019 -- Megatron-LM (tensor/sequence parallelism)
• Lepikhin et al., 2020 -- GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (expert parallelism)

03 · Model sharding#

❓ Why this topic. Parallelism (ch. 2) decides which device computes what; sharding decides which device stores what. They’re related but distinct decisions -- you can be tensor-parallel without sharding the optimizer state, and you can shard the optimizer state (ZeRO-1) without any tensor parallelism at all. This chapter is specifically about the storage side.

Naming it precisely: “sharded data parallelism.” Everything ZeRO/FSDP does (ch. 1) is, in the field’s own terminology, sharded data parallelism specifically -- it’s still data parallelism (every GPU trains on different data, same architecture, results combined) with the extra property that the model’s own state (weights/gradients/optimizer) is also split across that same data-parallel group rather than fully replicated on every member of it. Worth naming explicitly because it clears up a common confusion: sharded data parallelism is not a form of model parallelism (ch. 2) at all -- no GPU is computing a different part of the model’s forward pass than any other, the way tensor or pipeline parallelism does -- it’s data parallelism where the storage happens to be split, which is exactly why it composes cleanly with tensor/pipeline parallelism (chapter 2’s 3D parallelism) rather than competing with it for the same role.

If you’ve worked with database sharding before, the analogy is closer than it looks -- and where it breaks. Database sharding splits rows of a table across servers by some key (user ID, region), so no single server holds the whole table, and a query for a specific row needs to be routed to the right shard. Parameter sharding here is similar in shape -- split a big object (a parameter tensor instead of a table) across devices so no single device holds the whole thing -- but the access pattern is the opposite of a typical database’s: a database is optimized for reading a small slice of the sharded data per query and rarely needing the whole table at once, whereas a forward pass through a sharded model layer needs the entire parameter reassembled (via all-gather, defined in ch. 1) every single time, however briefly. That difference is exactly why parameter sharding’s cost is dominated by communication bandwidth rather than by clever query routing -- there’s no equivalent of a database index that lets you avoid ever touching most of the shards; every shard is needed, every time, just not for very long.

Parameter, gradient, and optimizer-state sharding: the recap#

ZeRO’s three stages -- sharding optimizer state only (stage 1), adding gradients (stage 2), adding parameters themselves (stage 3, equivalent to FSDP) -- are covered with full per-GPU memory arithmetic and a worked 70B-parameter example in the LLM lifecycle companion notes, chapter 8. The one addition worth making here, specific to how FSDP actually shards a parameter tensor rather than just how much memory it saves: FSDP shards each parameter along its first dimension, flattening and evenly dividing it across the data-parallel group, then reconstructing the full tensor via all-gather immediately before it’s needed in a forward or backward pass and freeing it again immediately after -- the parameter physically exists in full on any single GPU only for the brief window it’s actually being multiplied against, not at rest.

# Inspecting FSDP's actual shard: what one GPU holds at rest vs. reconstructed
from torch.distributed.fsdp import fully_shard

model = build_transformer_model(...)
fully_shard(model)
for name, param in model.named_parameters():
    print(name, param.shape, "-- this is the LOCAL SHARD shape, not the full parameter's logical shape")
    # a Linear layer's weight of logical shape (4096, 4096) on an 8-way shard group
    # shows up here as a flat local shard of roughly (4096*4096/8,) -- 1D, evenly divided, not a clean sub-matrix

Activation checkpointing combined with sharded training#

Why this combination needs care, not just stacking two flags. Activation checkpointing (lifecycle notes, ch. 8: trade ~33% more compute for a $\sqrt{L}$ memory reduction) and parameter sharding solve different memory problems -- checkpointing addresses activation memory, sharding addresses weight/gradient/optimizer memory -- and real large-model training needs both simultaneously, but naively wrapping a sharded (FSDP) module with checkpointing can accidentally force an early all-gather of the next layer’s sharded parameters before they’re needed, if the checkpointing boundary and the FSDP sharding boundary don’t line up. The fix is aligning checkpoint segments with FSDP’s own module-wrapping boundaries so recomputation during backward re-triggers the same all-gather/free pattern it would have used anyway, not an extra one.

from torch.distributed.fsdp import fully_shard
from torch.utils.checkpoint import checkpoint

class CheckpointedShardedBlock(torch.nn.Module):
    def __init__(self, block):
        super().__init__()
        fully_shard(block)             # shard this block's parameters specifically
        self.block = block

    def forward(self, x):
        # recomputation during backward re-triggers this block's own all-gather, not a mistimed one elsewhere
        return checkpoint(self.block, x, use_reentrant=False)

Layer-wise partitioning across devices#

Pipeline parallelism’s layer-to-GPU assignment (lifecycle notes, ch. 1) is the most common form of this, but the balance of that assignment is its own real engineering problem: naive equal-layer-count partitioning assumes every layer costs the same compute and memory, which is false the moment embeddings, a final unembedding/LM head, or heterogeneous layer types (attention vs. MoE-FFN, ch. 2) are in the mix -- an unbalanced partition leaves some pipeline stages idle waiting on a slower stage, a different source of bubble than the $\frac{p-1}{m}$ formula alone captures.

def balanced_layer_partition(layer_costs, num_stages):
    # layer_costs: list of estimated FLOPs (or measured wall-clock) per layer, from a profiling pass
    total = sum(layer_costs)
    target_per_stage = total / num_stages
    partitions, current, current_cost = [], [], 0
    for i, cost in enumerate(layer_costs):
        current.append(i)
        current_cost += cost
        if current_cost >= target_per_stage and len(partitions) < num_stages - 1:
            partitions.append(current)
            current, current_cost = [], 0
    partitions.append(current)          # remaining layers go to the last stage
    return partitions   # feed this into pipeline_parallel's stage assignment instead of a naive even split

Shard placement across node boundaries#

The rule from chapter 1’s bandwidth argument, applied to sharding specifically. A sharded parameter’s pieces can, in principle, live on any GPU in the sharding group -- but which GPUs are grouped together for a given shard matters enormously given the NVLink-vs-InfiniBand bandwidth gap (lifecycle notes, ch. 8). FSDP’s process-group construction should place the data-parallel (outer) sharding group so that its per-step all-gather/reduce-scatter traffic, which happens every layer, stays intra-node wherever the model fits that way, or is at minimum arranged so tensor-parallel groups (the highest-frequency communication) never cross a node boundary even when the outer data-parallel/FSDP grouping does.

from torch.distributed.device_mesh import init_device_mesh

# 2 nodes x 8 GPUs: TP confined within a node's 8 GPUs, FSDP/DP sharding spans the 2 nodes
mesh = init_device_mesh("cuda", (2, 8), mesh_dim_names=("dp", "tp"))
tp_mesh = mesh["tp"]     # tensor-parallel group: always the 8 GPUs within one node, per the bandwidth argument
dp_mesh = mesh["dp"]     # data-parallel/FSDP group: spans nodes, tolerates the slower InfiniBand hop

Resharding during checkpoint save/load#

The problem this solves. A checkpoint saved from a training run using 64 GPUs (and therefore sharded 64 ways) needs to be loadable by a different run using, say, 32 GPUs to continue training, or 8 GPUs to fine-tune, or 1 GPU to just inspect the weights -- without resharding support, a checkpoint is locked to the exact GPU count and sharding layout it was saved under, which is a real operational constraint (cluster sizes change, capacity gets reallocated, a training run gets split across two smaller allocations after a large one becomes unavailable).

How it’s actually solved. Rather than saving one sharded file per GPU in a layout tied to that GPU count, modern distributed checkpointing (PyTorch’s torch.distributed.checkpoint, and DeepSpeed’s newer Universal Checkpointing) saves each parameter’s shard with metadata describing its logical position within the full, unsharded tensor -- not just “GPU 7’s piece” but “elements 4096–8191 of this 32768-element parameter.” Loading into a different shard count then becomes a resharding computation: read whichever saved shards logically overlap the new shard boundaries a GPU needs, regardless of how many pieces the tensor was originally split into.

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

# saved under a 64-GPU sharding layout
def save_resharding_checkpoint(model, optimizer, step, path):
    model_state, optim_state = get_state_dict(model, optimizer)
    dcp.save({"model": model_state, "optim": optim_state, "step": step}, checkpoint_id=path)
    # metadata saved alongside each shard describes global tensor shape + this shard's offset within it

# loaded under a DIFFERENT 8-GPU sharding layout -- this is what makes it "resharding," not just "loading"
def load_with_reshard(model, optimizer, path):
    model_state, optim_state = get_state_dict(model, optimizer)   # reflects the NEW 8-way shard layout
    state_dict = {"model": model_state, "optim": optim_state, "step": 0}
    dcp.load(state_dict, checkpoint_id=path)   # dcp computes which of the original 64 shards overlap each of the 8 new ones
    set_state_dict(model, optimizer, model_state_dict=state_dict["model"], optim_state_dict=state_dict["optim"])
    return state_dict["step"]

⚠️ The gotcha this fixes, stated plainly. Before resharding-aware checkpointing existed, changing GPU count mid-project meant either writing custom shard-reassembly scripts or -- the common failure mode -- a checkpoint quietly loading with silently mismatched shard boundaries, corrupting the resumed training run’s weights without necessarily crashing. If a checkpoint format doesn’t document per-shard global offsets, treat any GPU-count change as unsafe until proven otherwise.

📚 Further reading

• Rajbhandari et al., 2020 -- ZeRO (lifecycle notes, ch. 8, for the full memory derivation)
• PyTorch Distributed Checkpoint documentation -- pytorch.org/docs/stable/distributed.checkpoint.html
• DeepSpeed Universal Checkpointing documentation -- deepspeed.ai

04 · Memory and performance#

❓ Why this topic. Getting a model to fit (ch. 2–3) is necessary but not sufficient -- a training run that fits but runs at 20% of achievable throughput is still an expensive way to burn a compute budget. This chapter is about the second question: given a working, correct setup, how much of the hardware’s actual capability is it using.

Mixed precision and gradient checkpointing: the recap#

Both are covered in depth elsewhere in this series -- mixed precision (bf16/fp16 training, loss scaling, numerical stability) in the transformer field notes, chapter 5; activation checkpointing’s memory/compute trade with a worked 126-layer example in the lifecycle notes, chapter 8. What’s worth adding here is how they interact with the parallelism strategies from chapter 2: gradient checkpointing’s ~33% compute overhead is a cost paid per data-parallel replica independently -- it doesn’t change with how many GPUs you’re sharding across, whereas mixed precision’s memory savings compound directly with ZeRO/FSDP sharding (halving bytes-per-parameter before sharding means every shard is also half the size), which is why real configs almost always combine bf16 with FSDP/ZeRO by default rather than treating them as alternatives.

Communication/computation overlap: the single highest-leverage systems optimization#

The idea. Every collective communication op (all-reduce, all-gather, reduce-scatter) is, from the GPU’s perspective, time spent not doing matrix multiplication. If that communication can be issued while the GPU is still busy with an unrelated compute operation -- rather than compute-then-wait-then-communicate-then-wait, sequentially -- the communication’s wall-clock cost partially or fully disappears behind compute that was going to happen anyway.

Where this actually happens in practice. DDP overlaps a given layer’s gradient all-reduce with the backward pass computation of earlier layers (by the time layer 1’s gradient is needed, layer 10’s all-reduce may already be in flight) -- this is on by default and is why DDP’s communication cost is often much smaller in practice than a naive “add up all the all-reduce time” estimate suggests. FSDP does the same for its per-layer all-gather: prefetching the next layer’s parameters while the current layer is still computing.

# Explicit prefetching hints -- FSDP2's forward/backward prefetch controls, the direct lever for this overlap
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

for i, layer in enumerate(model.layers):
    fully_shard(layer, reshard_after_forward=True)      # free this layer's gathered params right after use
    if i + 1 < len(model.layers):
        layer.set_modules_to_forward_prefetch([model.layers[i + 1]])  # start next layer's all-gather early, overlapped

# torch.profiler view of whether overlap is actually happening: look for concurrent NCCL + compute streams
from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    train_step(model, batch)
# in the resulting trace (view in chrome://tracing or Perfetto), overlapping NCCL kernels and compute kernels
# on separate CUDA streams at the same wall-clock timestamp is overlap actually working; back-to-back,
# non-overlapping blocks of "compute" then "nccl:all_gather" then "compute" again is overlap NOT happening
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

Batch size tuning and the throughput/convergence tradeoff#

The tension, stated plainly. A larger batch size increases GPU utilization (bigger matmuls, better hardware efficiency, more tokens processed per unit wall-clock time) but past the gradient noise scale (lifecycle notes, ch. 7) it stops meaningfully improving -- and can start hurting -- how efficiently the model actually learns per token seen, since very large batches average away useful gradient variance along with noise. Real tuning is a search, not a formula: find the largest batch size that (a) fits in memory given the parallelism/sharding setup and (b) doesn’t require a learning-rate increase so large that training destabilizes (the two commonly move together -- bigger batch, proportionally bigger LR, per the standard linear-scaling heuristic -- but that heuristic has a real breaking point).

def find_max_batch_size(model, optimizer, device, start=8, max_memory_frac=0.9):
    batch_size = start
    total_mem = torch.cuda.get_device_properties(device).total_memory
    while True:
        try:
            dummy = torch.randn(batch_size, *input_shape, device=device)
            loss = model(dummy).sum()
            loss.backward()
            optimizer.zero_grad()
            if torch.cuda.max_memory_allocated(device) > total_mem * max_memory_frac:
                return batch_size // 2       # step back from the OOM edge to leave headroom for activation spikes
            batch_size *= 2
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            return batch_size // 2

GPU utilization profiling: MFU, the number that tells you if any of this worked#

Model FLOPs Utilization (MFU) -- actual achieved FLOPs/sec divided by the GPU’s theoretical peak FLOPs/sec -- is the single summary number for “is all of chapters 1–4’s engineering actually paying off,” already introduced in the lifecycle notes’ chapter 1 production table. Computing it directly:

def compute_mfu(model_flops_per_token, tokens_per_second, gpu_peak_flops):
    achieved_flops = model_flops_per_token * tokens_per_second
    return achieved_flops / gpu_peak_flops

# worked example: a 7B dense model, ~2 * 6 * N FLOPs/token (standard forward+backward approximation, N=params)
model_flops_per_token = 2 * 6 * 7e9         # ≈ 8.4e10 FLOPs/token
tokens_per_second = 15_000                   # measured from an actual training step's wall-clock time
gpu_peak_flops = 989e12                      # H100 bf16 peak, per NVIDIA's spec sheet
print(compute_mfu(model_flops_per_token, tokens_per_second, gpu_peak_flops))   # ~0.128, i.e. ~13% MFU -- investigate further

A well-optimized frontier training run commonly reports MFU in the 40–55% range (PaLM’s own published 46.2% is the frequently-cited real number); anything well below that on a first attempt almost always traces back to insufficient communication/compute overlap (above), a suboptimal parallelism topology (ch. 2–3), or a data-loading bottleneck starving the GPU of work rather than a fundamental compute-bound ceiling.

📚 Further reading

• Chowdhery et al., 2022 -- PaLM: Scaling Language Models with Pathways (46.2% MFU)
• PyTorch FSDP2 documentation -- prefetching and overlap controls
• NVIDIA H100 datasheet -- peak FLOPs figures for MFU calculations

05 · Production training concerns#

❓ Why this topic. A training run that’s algorithmically correct and well-parallelized still fails as a production activity if it can’t survive the cluster it runs on for weeks at a time -- the lifecycle notes’ chapter 1 real number (466 interruptions over 54 days on Llama 3’s cluster) is the standing reminder that this isn’t a hypothetical concern.

Job orchestration: Slurm and Kubernetes, the two real defaults#

#!/bin/bash
#SBATCH --job-name=llm-pretrain
#SBATCH --nodes=64
#SBATCH --gpus-per-node=8
#SBATCH --time=72:00:00
#SBATCH --requeue                          # automatically resubmit if preempted or a node fails mid-job

srun torchrun \
    --nnodes=$SLURM_NNODES --nproc_per_node=8 \
    --rdzv_id=$SLURM_JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$(scontrol show hostname $SLURM_NODELIST | head -n1):29500 \
    train.py --resume-from-checkpoint /shared/checkpoints/latest

# Kubernetes, via the Kubeflow PyTorchJob CRD -- the cloud-native equivalent of the Slurm script above
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: llm-pretrain
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      template:
        spec:
          containers:
          - name: pytorch
            image: my-training-image:latest
            command: ["torchrun", "--nnodes=64", "--nproc_per_node=8", "train.py"]
    Worker:
      replicas: 63
      restartPolicy: OnFailure    # k8s-native automatic restart of a failed worker pod

The --requeue/restartPolicy: OnFailure lines are the entire point relative to running python train.py by hand: at 466-interruptions-per-54-days scale, “the orchestrator automatically relaunches a failed job/pod without a human noticing” is a hard requirement, not a convenience.

Checkpointing, resume, and failure recovery, tied together#

import signal, sys

class CheckpointOnSignal:
    """Catches SIGTERM (sent by Slurm/k8s before a preemption or node drain) and forces a checkpoint first."""
    def __init__(self, model, optimizer, step_ref, ckpt_dir):
        self.model, self.optimizer, self.step_ref, self.ckpt_dir = model, optimizer, step_ref, ckpt_dir
        signal.signal(signal.SIGTERM, self._handler)

    def _handler(self, signum, frame):
        save_resharding_checkpoint(self.model, self.optimizer, self.step_ref[0], self.ckpt_dir)  # ch. 3's function
        sys.exit(0)   # exit cleanly so the orchestrator's restart policy relaunches from this exact checkpoint

Combined with --requeue/OnFailure above, this closes the loop: a preemption warning triggers an immediate checkpoint, the job exits, the orchestrator relaunches it, and --resume-from-checkpoint picks up from that just-saved state -- the practical mechanism behind Llama 3’s reported >90% GPU utilization despite the interruption rate (lifecycle notes, ch. 1): most interruptions cost a resume, not a full restart from the last periodic checkpoint.

Experiment tracking: what to log, and the real code#

import wandb

wandb.init(project="llm-pretrain", config={"lr": 3e-4, "batch_size": 4_000_000, "model_size": "7B"})

for step, batch in enumerate(dataloader):
    loss = train_step(model, optimizer, batch)
    if step % 10 == 0:
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        wandb.log({
            "train/loss": loss.item(),
            "train/grad_norm": grad_norm.item(),          # ch. 6: a spike here precedes most loss-spike incidents
            "train/lr": optimizer.param_groups[0]["lr"],
            "system/mfu": compute_mfu(model_flops_per_token, tokens_per_second, gpu_peak_flops),  # ch. 4
        }, step=step)

Real incident data: what actually happens over weeks of training#

Everything in this chapter so far is mechanism. The following numbers are what that mechanism is actually defending against, taken from the small number of frontier runs whose operational logs were made public rather than summarized away in a paper’s methods section:

• OPT-175B (Meta, 2022). Trained on roughly 1,024 NVIDIA A100 80GB GPUs over about 56 calendar days, achieving 147 TFLOPS/GPU. Meta’s own released 114-page operational logbook -- a genuinely unusual level of transparency -- documents 35 manual restarts and roughly 70 automatic restarts due to hardware failures, with over 100 hosts cycled out of the cluster over the run. The team made live algorithmic interventions mid-run when instability appeared: switching the optimizer from AdamW to plain SGD and back, and upgrading the underlying Megatron version mid-training -- decisions a paper’s methods section would normally describe as if they were the plan from day one, when the logbook shows they were real-time responses to instability nobody could have fully planned around in advance.
• Llama 3 405B (Meta, 2024). Across 54 days on roughly 16,384 H100 GPUs, independent analysis of the paper’s own reporting puts unexpected interruptions at 419, roughly one every three hours, with hardware failures responsible for about 78% of them -- GPU failures (including NVLink errors) alone accounting for around 30.1%, and HBM3 memory failures around 17.2%. (A commonly cited related figure from the same paper is 466 total interruptions; the discrepancy is almost certainly a difference in what’s counted as “unexpected” versus total, including planned maintenance -- worth noting as a reminder that even primary-source incident counts depend on the counting methodology, not just the raw event log.) Despite this cadence, automated recovery (ch. 3 and 5’s checkpoint-on-signal plus orchestrator restart policy, working as designed) kept effective training time above 90%.
• ByteDance MegaScale, a production system training a 175B-parameter model across 12,288 GPUs, reports 55.2% MFU sustained over a multi-week run with more than 100 automated failure-recovery events -- a genuinely high MFU by the field’s own standard (compare PaLM’s 46.2%, ch. 4), specifically because the recovery events were handled by automation fast enough not to meaningfully erode sustained throughput.

The throughline across all three: nobody at this scale has a training run with zero incidents. The entire engineering difference between a project that finishes on schedule and one that doesn’t is whether recovery from these incidents is automated and fast (this chapter’s checkpoint-on-signal, orchestrator restart policies, and ch. 6’s fast fault localization) or manual and slow -- OPT-175B’s own account of “someone had to manually intervene each time a machine failed” early on, followed by building automated recovery scripts specifically because manual intervention wasn’t sustainable at that failure cadence, is the concrete before/after case for exactly the automation this chapter describes.

Determinism and reproducibility#

import torch, random, numpy as np

def set_deterministic(seed: int):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)
    torch.use_deterministic_algorithms(True)          # forces deterministic kernel variants where available
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False              # benchmark mode picks fastest kernel per-run, non-deterministically

⚠️ The honest limit. Bit-exact reproducibility across different GPU counts or parallelism configurations is not generally achievable even with every flag above set correctly -- all-reduce’s summation order depends on the communication topology, and floating-point addition is not associative, so the same logical computation genuinely produces slightly different results at a different world size. What determinism settings do guarantee is exact reproducibility on the same hardware/parallelism configuration re-run with the same seed and data order -- useful for debugging (ch. 6) and audit trails (lifecycle notes, ch. 2), not for “this checkpoint will be bit-identical if retrained on a different cluster size.”

📚 Further reading

• Slurm documentation -- slurm.schedmd.com
• Zhang et al. (Meta AI), 2022 -- OPT: Open Pre-trained Transformer Language Models, and the public metaseq OPT-175B chronicles/logbook (github.com/facebookresearch/metaseq)
• Dubey et al., 2024 -- The Llama 3 Herd of Models (interruption and hardware-failure breakdown)
• Jiang et al., 2024 -- MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs (ByteDance)
• Kubeflow PyTorchJob -- kubeflow.org/docs/components/training
• Weights & Biases documentation -- docs.wandb.ai

06 · Debugging and observability#

❓ Why this topic. Everything above this point describes how distributed training is supposed to work. This chapter is about the failure mode that matters most in practice: a distributed job that runs to completion, produces no error, and is silently wrong -- a sharding bug that trains a model on 1/8th of its intended data, a NaN that gets masked by a loss curve that still looks plausible, a checkpoint that loads without complaint but doesn’t actually match what was saved. None of these throw an exception. All of them need to be actively checked for.

The real split between “easy” and “hard” failures, with numbers. ByteDance’s own published account of their production LLM training infrastructure draws a direct, measured distinction between two failure classes. Explicit failures -- a crashed process, a CUDA error, anything that produces an error message or a clear log signature -- are typically detected in around 60 seconds and localized to the faulty machine within 2 to 15 minutes, because the system has something concrete to point at. Implicit failures -- anything that degrades or corrupts training without throwing an error -- made up more than 10% of the incidents they measured, and are categorically harder: their own worked example, a communication hang caused by an underlying CUDA error with no clear log signature, took more than an hour and a half of manual diagnosis to root-cause, because nothing in the system’s own logs pointed at the actual cause. Every check in the rest of this chapter exists specifically to convert an implicit failure into something closer to an explicit one -- a NaN hook, a parity check, or a shard-validation assertion turns “training is silently wrong somewhere” into “training just failed loudly, at this specific line, for this specific reason,” which is the entire difference between a 60-second detection and a 90-minute one.

Verify gradient flow early -- before trusting a single loss number#

The check. After the very first backward pass, every parameter that should receive a gradient must actually have one that’s nonzero -- a silently disconnected module (a wrapping bug, a mistaken .detach(), a frozen parameter that should have been trainable) produces a loss curve that looks completely normal while whole parts of the model never learn anything.

def verify_gradient_flow(model, loss):
    loss.backward()
    missing, zero = [], []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if param.grad is None:
            missing.append(name)                         # disconnected from the loss graph entirely
        elif torch.all(param.grad == 0):
            zero.append(name)                              # connected, but receiving no actual signal
    if missing:
        raise RuntimeError(f"No gradient reached: {missing}")
    if zero:
        print(f"⚠️ Zero gradient (possibly dead, or correctly frozen): {zero}")

Run this once, on the very first training step, on a single GPU before scaling to the full distributed job -- it’s far cheaper to catch a wiring bug at step 0 on 1 GPU than at step 10,000 on 512.

Watch for NaNs/Infs -- with hooks, not just a loss-curve eyeball#

Why waiting for the loss to become NaN is too late. By the time the loss is visibly NaN, the corrupting value has usually already propagated through several layers and possibly several optimizer steps (an Adam second-moment estimate that’s absorbed a NaN doesn’t recover just because a later gradient is clean) -- catching the first NaN activation or gradient, at the layer it appeared in, is what actually makes the incident debuggable.

def register_nan_hooks(model):
    def make_hook(name):
        def hook(module, inp, out):
            out_t = out[0] if isinstance(out, tuple) else out
            if torch.isnan(out_t).any() or torch.isinf(out_t).any():
                raise RuntimeError(f"NaN/Inf detected in forward output of layer: {name}")
        return hook
    for name, module in model.named_modules():
        module.register_forward_hook(make_hook(name))

def register_grad_nan_hooks(model):
    for name, param in model.named_parameters():
        def make_hook(pname):
            def hook(grad):
                if torch.isnan(grad).any() or torch.isinf(grad).any():
                    print(f"⚠️ NaN/Inf gradient at: {pname}")
                return grad
            return hook
        if param.requires_grad:
            param.register_hook(make_hook(name))

These hooks add real overhead (an isnan/isinf check on every layer’s output, every step) -- the standard practice is running them for the first several hundred steps of a new configuration, or behind a debug flag re-enabled specifically when instability is suspected, not left on for an entire multi-week run.

What production frameworks do about this automatically, not just diagnostically. DeepSpeed’s fp16 training doesn’t just let a NaN happen and hope a hook catches it -- its dynamic loss scaler multiplies the loss by a scale factor before backward specifically to keep small gradients representable in fp16’s narrow range, and if an overflow (inf/NaN) is actually detected in the resulting gradients, it skips that step’s optimizer update entirely and lowers the scale factor before retrying, rather than letting a corrupted update reach the weights:

# What this actually looks like from the training loop's perspective under DeepSpeed
for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    overflow = model_engine.optimizer.overflow     # DeepSpeed sets this flag when it detects inf/NaN in gradients
    model_engine.step()                              # internally: skip the actual weight update if overflow==True,
                                                        # and reduce the loss-scale factor for the next attempt
    if overflow:
        print("⚠️ Step skipped due to fp16 overflow -- loss scale reduced, no corrupted update applied")

⚠️ The failure mode this doesn’t fully protect against. If training is inherently unstable -- a learning rate genuinely too high for the current phase of training, not just an occasional fp16 range issue -- the loss scaler keeps reducing its scale factor step after step, chasing a problem that isn’t a precision issue at all, until it hits its own minimum and the run aborts outright with an explicit “loss scale already at minimum” error. That error message is a real, reliable signal that the actual problem is a training-stability issue (learning rate, data, architecture) masquerading as a numerics issue, not a reason to just lower the minimum scale further and hope it goes away.

Single-GPU vs. distributed parity: the check that catches sharding bugs#

The check. Run the exact same small batch, same seed, same initialization, once on a single GPU with no parallelism at all, and once through the full distributed (FSDP/TP/PP) configuration -- the loss and gradient norms should match to within floating-point tolerance. A mismatch here means the parallelism implementation itself -- not the model, not the data -- is the bug.

def check_distributed_parity(model_fn, batch, seed=42, atol=1e-4):
    torch.manual_seed(seed)
    single_gpu_model = model_fn().cuda()
    single_loss = single_gpu_model(batch.cuda())
    single_loss.backward()
    single_grad_norm = torch.nn.utils.clip_grad_norm_(single_gpu_model.parameters(), max_norm=1e9)  # measure, don't clip

    torch.manual_seed(seed)                                  # identical init, identical data order
    distributed_model = wrap_with_fsdp(model_fn())            # your ch. 1/3 setup
    dist_loss = distributed_model(batch)
    dist_loss.backward()
    dist_grad_norm = torch.nn.utils.clip_grad_norm_(distributed_model.parameters(), max_norm=1e9)

    assert torch.allclose(single_loss, dist_loss, atol=atol), f"Loss mismatch: {single_loss} vs {dist_loss}"
    assert torch.allclose(single_grad_norm, dist_grad_norm, atol=atol), "Gradient norm mismatch -- sharding bug likely"

Validate shard parity: does the reconstructed parameter actually match?#

from torch.distributed.fsdp import fully_shard
from torch.distributed._tensor import DTensor

def validate_shard_reconstruction(sharded_model, reference_state_dict):
    for name, param in sharded_model.named_parameters():
        if isinstance(param.data, DTensor):
            full_tensor = param.data.full_tensor()             # explicit all-gather to reconstruct the logical tensor
        else:
            full_tensor = param.data
        ref = reference_state_dict[name]
        if not torch.allclose(full_tensor, ref, atol=1e-5):
            raise RuntimeError(f"Shard reconstruction mismatch at {name} -- resharding or checkpoint-load bug")

Run this specifically after any resharding operation (ch. 3) -- loading a checkpoint saved under one shard count into a different one is exactly where an off-by-one in shard-boundary metadata silently produces a subtly wrong reconstructed parameter.

Measure communication overhead directly, not by inference from wall-clock alone#

import torch.distributed as dist, time

def measure_allreduce_overhead(tensor_size_mb, group, n_trials=20):
    tensor = torch.randn(int(tensor_size_mb * 1e6 / 4), device="cuda")   # fp32 elements
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(n_trials):
        dist.all_reduce(tensor, group=group)
    torch.cuda.synchronize()
    elapsed = (time.perf_counter() - start) / n_trials
    achieved_bandwidth_gb_s = (tensor_size_mb / 1000) / elapsed
    return elapsed, achieved_bandwidth_gb_s
    # compare achieved_bandwidth against NVLink's ~900GB/s (intra-node) or InfiniBand's real NIC spec (inter-node) --
    # a large gap below the hardware's rated bandwidth points at topology misconfiguration (ch. 3), not a hardware fault

Inspect memory fragmentation directly#

print(torch.cuda.memory_summary(device="cuda:0", abbreviated=False))
# key numbers to read from the summary: "reserved" vs "allocated" -- a large, growing gap between them across
# training steps is fragmentation (memory PyTorch's allocator is holding but can't hand out for a new, differently-
# shaped request); torch.cuda.memory_stats()["num_alloc_retries"] rising over time is the same symptom, numerically
torch.cuda.empty_cache()    # releases reserved-but-unallocated memory back to the driver -- a diagnostic, not a fix,
                             # since fragmentation that recurs immediately points at an allocation-pattern problem

Confirm checkpoint integrity -- don’t assume a successful write means a correct one#

import hashlib, torch

def checkpoint_hash(state_dict):
    hasher = hashlib.sha256()
    for key in sorted(state_dict.keys()):
        hasher.update(key.encode())
        hasher.update(state_dict[key].cpu().numpy().tobytes())
    return hasher.hexdigest()

def verify_checkpoint_roundtrip(model, save_fn, load_fn, path):
    original_state = {k: v.clone() for k, v in model.state_dict().items()}
    original_hash = checkpoint_hash(original_state)
    save_fn(model, path)
    load_fn(model, path)                       # load it back into a fresh copy
    reloaded_hash = checkpoint_hash(model.state_dict())
    if original_hash != reloaded_hash:
        raise RuntimeError("Checkpoint round-trip mismatch -- corruption in save, load, or resharding logic")

Run this once for every new checkpointing configuration (a new shard count, a new offload setting, a new storage backend) -- not on every save during a real training run, where the cost of hashing every parameter every step would be prohibitive, but specifically whenever the checkpointing mechanism itself changes.

Track training-speed regressions#

class ThroughputMonitor:
    def __init__(self, baseline_tokens_per_sec, alert_threshold=0.85):
        self.baseline = baseline_tokens_per_sec
        self.threshold = alert_threshold
        self.window = []

    def record(self, tokens_this_step, elapsed_seconds):
        self.window.append(tokens_this_step / elapsed_seconds)
        if len(self.window) >= 50:
            recent_avg = sum(self.window[-50:]) / 50
            if recent_avg < self.baseline * self.threshold:
                self._alert(recent_avg)

    def _alert(self, recent_avg):
        print(f"⚠️ Throughput regression: {recent_avg:.0f} tok/s vs baseline {self.baseline:.0f} tok/s "
              f"({recent_avg/self.baseline:.1%} of expected) -- check for a newly-introduced data-loading "
              f"bottleneck, a degraded network path (ch. 6's all-reduce measurement), or thermal throttling")

A gradual, sustained throughput regression across an otherwise-unremarkable multi-week run is one of the most common real production incidents in large-scale training, and one of the easiest to miss without an explicit monitor like this -- nothing crashes, the loss curve keeps looking reasonable, and the run simply costs more wall-clock time (and money) than it should for weeks before anyone notices the trend.

📚 Further reading

• PyTorch torch.autograd hooks and torch.utils.checkpoint documentation
• PyTorch DTensor / full_tensor() API -- pytorch.org/docs/stable/distributed.tensor.html
• NVIDIA NCCL Debugging Guide -- for interpreting NCCL_DEBUG=INFO output alongside the all-reduce timing check above

07 · Inference-side scaling#

❓ Why this topic, and why this chapter is short. Batch scheduling, KV-cache management, disaggregated serving, and speculative decoding are covered in full depth -- with math, worked examples, and real vLLM/SGLang/TensorRT-LLM code -- in the companion long-context document’s chapters 2–3 and 8. This chapter covers specifically the sharding angle those chapters describe from the serving-engine side: how a model’s weights and its KV cache actually get split across GPUs at inference time, which is the direct continuation of chapters 2–3’s training-time parallelism applied to a very different workload (memory-bandwidth-bound decode instead of compute-bound training).

Tensor-parallel inference: the same idea as chapter 2, a different bottleneck#

Training-time tensor parallelism (ch. 2) splits weight matrices to fit a model that doesn’t fit on one GPU, accepting communication cost as the price. Inference-time tensor parallelism does the same split for the same reason (a 70B+ model’s weights alone may not fit one GPU even without any optimizer state to worry about), but the bottleneck it interacts with is different: every layer’s TP all-reduce now happens on the decode critical path, adding directly to per-token latency rather than being hidden behind a long training step’s other work.

# vLLM: tensor-parallel inference is a single config flag -- the sharding itself is handled internally
from vllm import LLM
llm = LLM(model="meta-llama/Llama-3.1-70B-Instruct", tensor_parallel_size=4)   # weights split across 4 GPUs

KV-cache sharding: splitting the cache along the same axis as the weights#

Why this isn’t a separate decision from TP. When attention’s weight matrices are tensor-parallel-sharded by head (each GPU holding a subset of attention heads, transformer notes ch. 3’s multi-head structure making this a natural split), the KV cache for those heads shards the same way automatically -- GPU $i$ stores K/V only for the heads it owns, meaning the companion document’s whole PagedAttention/RadixAttention memory-management story (ch. 2) runs per TP shard, not once globally: each GPU manages its own slice of every sequence’s cache independently, in lockstep with the others.

# Conceptual shape of what each TP rank actually stores for the KV cache -- this is what vLLM/SGLang do internally
def kv_cache_shard_shape(n_kv_heads_total, tp_size, block_size, d_head, n_blocks):
    n_kv_heads_per_shard = n_kv_heads_total // tp_size    # e.g. 8 total KV heads (GQA), tp_size=4 -> 2 heads/shard
    return (n_blocks, block_size, n_kv_heads_per_shard, d_head)   # each GPU's PagedAttention pool has THIS shape,
                                                                     # not the full n_kv_heads_total -- 1/tp_size the memory

⚠️ The real constraint this creates. GQA (transformer notes, ch. 3) reduces n_kv_heads_total specifically to make this division cleaner -- a model with only 8 KV heads can be TP-sharded 8 ways with exactly 1 head per shard, but TP-sharding it 16 ways is impossible without splitting a single head’s K/V across two GPUs, which most serving engines don’t support cleanly. This is a real, concrete reason a model’s GQA head count caps its viable tensor-parallel degree at inference time, independent of any memory argument -- worth checking before assuming a larger TP size is always available as an option.

Pipeline-parallel inference: the latency-sensitive version of a training-time idea#

Why this exists as a distinct option from TP. Pipeline parallelism at inference splits layers across GPUs (GPU 1 holds layers 1–20, GPU 2 holds layers 21–40, and so on) rather than splitting weight matrices -- the training-time bubble problem (ch. 2’s $\frac{p-1}{m}$ formula) has a direct inference-side analog, but with a crucial difference: a single inference request has no “microbatches” to fill the pipeline with the way a training step does, so a lone request pipelined across many stages spends real wall-clock time with most stages idle, waiting for their turn. This makes pure pipeline-parallel inference a poor fit for low-latency single-request serving, but a reasonable one for high-throughput batch inference (many requests in flight simultaneously naturally fill the pipeline the way training microbatches do) or specifically when a model is too large to fit via tensor parallelism alone within a single node’s NVLink domain and must span nodes, where pipeline parallelism’s less-frequent, boundary-only communication (chapter 2’s argument for why PP tolerates cross-node latency better than TP) makes it the more viable way to spread the model at all.

# TensorRT-LLM / vLLM: combining TP (within a node) and PP (across nodes) for a model too large for TP alone
from vllm import LLM
llm = LLM(
    model="meta-llama/Llama-3.1-405B-Instruct",
    tensor_parallel_size=8,      # within-node split, per GPU's NVLink domain
    pipeline_parallel_size=2,     # across-node split, layers 1-N on one node's TP group, N+1-end on the other's
)

The practical decision rule. Reach for tensor parallelism first, up to the size of one node’s NVLink domain -- it adds latency (an all-reduce per layer) but keeps the whole model responsive per-request. Only add pipeline parallelism on top once the model genuinely doesn’t fit within tensor parallelism’s node-boundary limit, and expect the added pipeline hop to show up as a real, measurable increase in time-to-first-token for a single, non-batched request -- exactly the cost this section’s opening paragraph describes, now paid in production latency rather than training throughput.

Data-parallel serving vs. tensor-parallel serving: the orthogonal axis#

Where TP splits one replica of the model across GPUs to fit it, data-parallel serving runs multiple independent replicas (each possibly itself TP-sharded) to handle more concurrent traffic -- exactly analogous to chapter 2’s DP vs. TP distinction on the training side. The serving-side decision rule: scale replicas (data-parallel) to handle more concurrent requests up to a latency SLO; scale TP degree only when a single replica genuinely doesn’t fit the target latency or memory budget on fewer GPUs, since TP’s communication overhead is a real, per-token cost that pure replication doesn’t pay.

Routing and load balancing across sharded replicas#

# a minimal least-loaded router across N replicas, each itself possibly TP-sharded internally
import itertools

class ReplicaRouter:
    def __init__(self, replica_endpoints):
        self.replicas = {ep: 0 for ep in replica_endpoints}   # endpoint -> current in-flight request count

    def route(self):
        endpoint = min(self.replicas, key=self.replicas.get)   # send to whichever replica has the fewest in-flight requests
        self.replicas[endpoint] += 1
        return endpoint

    def complete(self, endpoint):
        self.replicas[endpoint] -= 1

Real production routers (e.g. the layer in front of a vLLM/SGLang fleet) use this same least-loaded principle but drive it from the queue-depth/Little’s-Law signal from the companion document’s chapter 5, not raw request counts alone -- a request to a replica mid-way through a very long generation counts differently than one that just started.

Model size vs. latency: the tradeoff this whole chapter is in service of#

Every sharding decision above is ultimately in service of one curve: bigger models are more capable but need more GPUs (TP) or more replicas’ worth of memory (DP) to serve at a given latency target, and the actual production decision -- same as chapter 5 of the LLM lifecycle notes’ cost-modeling section -- is where on that curve a given product’s latency SLO and budget land, not “use the biggest model available.”

📚 Further reading

• The companion long-context document, chapters 2, 3, and 8 -- PagedAttention, continuous batching, disaggregated serving, real vLLM/SGLang code
• The LLM lifecycle notes, chapter 5 -- cost modeling, SLOs, autoscaling via queue depth

08 · Learn-in-order path, and the note template#

The path, extended with what each stage should leave you able to debug#

1. Single-GPU training in PyTorch. Before any parallelism at all -- you need a known-correct baseline to run the chapter 6 parity check against later.
2. DDP and gradient accumulation (ch. 1–2). The floor: full model replication, gradient sync, and simulating a larger batch than memory allows. You should leave this stage able to explain no_sync() and why it matters.
3. FSDP or DeepSpeed ZeRO (ch. 1, 3). Parameter/gradient/optimizer sharding. You should leave this stage able to compute a model’s per-GPU memory footprint at a given shard count by hand (lifecycle notes, ch. 8’s formula), not just enable a flag and hope it fits.
4. Tensor/pipeline/sequence/expert parallelism (ch. 2). The axes beyond data parallelism. You should leave this stage able to explain why TP stays within a node and DP/PP can cross nodes, from the bandwidth numbers, not as a memorized rule.
5. Checkpointing and cluster failure recovery (ch. 3, 5). Resharding, signal-triggered checkpointing, orchestrator restart policies. You should leave this stage able to explain what happens, step by step, when a node dies mid-training-run on your specific setup.
6. Memory profiling and kernel efficiency (ch. 4, 6). MFU, communication/compute overlap, fragmentation inspection. You should leave this stage able to look at a training run’s MFU number and say specifically what’s likely leaving it below the achievable ceiling.
7. Inference sharding and serving optimization (ch. 7, and the companion long-context document in full). The same sharding vocabulary, applied to a latency-bound rather than throughput-bound workload.
8. Parallel I/O and data pipelines at scale (ch. 9). The step easiest to skip and most likely to quietly cap your throughput anyway -- a perfectly parallelized, perfectly sharded job still wastes money if the data loader can’t keep up. Leave this stage able to actually measure, not guess, whether a slow run is data-bound.

The note template, extended#

For each framework or technique covered in this chapter, the original structure -- purpose, what gets sharded, parallelism type, memory tradeoffs, communication costs, failure modes, checkpoint behavior, profiling/debugging approach -- is the right one, and applies cleanly to everything above. Three additions worth making explicit, matching the same pattern the LLM lifecycle notes closed on: note which chapter 6 debugging check actually catches this technique’s most common bug (e.g., FSDP’s most common bug is caught by shard-parity validation, not by gradient-flow verification); note the real number, not the textbook default, for whatever this technique’s headline metric is on your actual hardware -- MFU, achieved all-reduce bandwidth versus the hardware’s rated spec, actual checkpoint save/resume wall-clock time -- because, as the LLM lifecycle notes’ chapter 7 makes the case for at length, a hyperparameter or performance figure copied from a paper or a tutorial is a starting guess, not a substitute for what your own cluster, your own model shape, and your own data pipeline actually produce; and, once any of this is running inside an actual organization rather than a research script, note the governance surface -- what compute quota or cost-approval process this technique’s resource footprint has to clear (ch. 10) -- since a technically correct setup that nobody budgeted for is still, in practice, a blocked one.

09 · Parallel I/O and data pipelines at scale#

❓ Why this topic. Chapters 1–4 assume a batch of data simply arrives on each GPU when needed. At real cluster scale, getting the right data to the right rank fast enough not to starve a GPU that costs several dollars an hour to leave idle is its own genuine engineering problem -- a perfectly parallelized, perfectly sharded training job still wastes money if every GPU spends 30% of each step waiting on a slow data loader.

Feeding each rank from different files#

The core requirement: with $N$ data-parallel ranks, each one needs a different, non-overlapping slice of the dataset every epoch -- duplicating data across ranks wastes compute re-processing the same examples multiple times per step for no benefit.

from torch.utils.data.distributed import DistributedSampler

dataset = MyShardedDataset(data_files=glob.glob("/data/shard_*.jsonl"))
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
loader = torch.utils.data.DataLoader(dataset, sampler=sampler, batch_size=32, num_workers=8, pin_memory=True)
# num_workers spins up separate CPU processes to pre-fetch and tokenize the NEXT batch while the GPU
# is still busy with the CURRENT one -- this is the data-side analog of chapter 4's compute/communication overlap

At real scale, don’t shard by sampler alone -- shard by file. DistributedSampler still requires every rank to at least know about the full dataset’s index. At the scale of trillions of tokens spread across many thousands of files, the standard pattern instead assigns entire files to specific ranks up front -- rank $i$ only ever opens files i, i+N, i+2N, ... -- so no rank ever needs a full-dataset index in memory, and no coordination is needed to determine who reads what.

def assign_files_to_rank(all_files, rank, world_size):
    return all_files[rank::world_size]     # rank 0 gets files 0, N, 2N...; rank 1 gets 1, N+1, 2N+1...

my_files = assign_files_to_rank(sorted(glob.glob("/data/shard_*.jsonl")), rank, world_size)

Streaming formats: why frontier-scale data pipelines don’t use plain files at all#

Reading many small files over a shared network filesystem (common in a cluster where storage isn’t local to compute) turns data loading into a filesystem-metadata bottleneck long before it becomes a bandwidth bottleneck -- opening a file has a fixed latency cost that dominates when you’re opening millions of them. WebDataset (and similar sequential-archive formats) solves this by packing many examples into large sequential-read tar-style archives, so a data worker does one long streaming read instead of millions of small ones.

import webdataset as wds

dataset = (
    wds.WebDataset("/data/shard-{000000..012345}.tar", shardshuffle=True)
    .decode()
    .to_tuple("txt", "json")
    .batched(32)
)
loader = wds.WebLoader(dataset, num_workers=8, batch_size=None)   # batching already done above

MPI I/O and parallel filesystems: the HPC-native answer to the same problem#

On genuine HPC clusters (Slurm-managed, with a parallel filesystem like Lustre or GPFS rather than object storage), the equivalent mechanism is MPI I/O -- every rank issues its reads through the MPI library directly against the shared parallel filesystem, which is specifically built to serve many concurrent, non-overlapping read requests efficiently across its own striped storage nodes, rather than relying on a single filesystem server that would become the bottleneck the moment thousands of ranks all request data from it at once. This is less something a training script author codes by hand and more a property of which storage system the cluster provides -- the practical takeaway for a training team is confirming the underlying filesystem is actually built for many-way concurrent access before assuming a straightforward DistributedSampler setup will scale past a few dozen GPUs without becoming I/O-bound.

The failure mode this whole chapter exists to prevent#

# a quick, real diagnostic: is the GPU actually waiting on data, or is data loading keeping up?
import time

def diagnose_data_bottleneck(loader, model, n_steps=20):
    data_wait_time, compute_time = 0.0, 0.0
    it = iter(loader)
    for _ in range(n_steps):
        t0 = time.perf_counter()
        batch = next(it)                    # time spent here is entirely data-loading wait
        t1 = time.perf_counter()
        model(batch)                          # time spent here is entirely GPU compute
        torch.cuda.synchronize()
        t2 = time.perf_counter()
        data_wait_time += t1 - t0
        compute_time += t2 - t1
    print(f"data wait: {data_wait_time:.2f}s, compute: {compute_time:.2f}s, "
          f"data-bound fraction: {data_wait_time/(data_wait_time+compute_time):.1%}")
    # a data-bound fraction above roughly 10-15% is worth investigating -- more num_workers, a streaming
    # format switch (above), or moving data closer to compute, before assuming the parallelism strategy is the issue

This is the same diagnostic instinct as chapter 6’s whole approach to debugging: don’t guess whether data loading is the bottleneck, measure it directly, because a GPU utilization number alone (chapter 4) can’t distinguish “waiting on data” from “waiting on communication” from “genuinely compute-bound” without breaking the wall-clock time down like this.

📚 Further reading

• PyTorch DistributedSampler and DataLoader documentation
• webdataset -- github.com/webdataset/webdataset
• MPI-IO standard (part of the MPI specification) -- for the HPC-native parallel filesystem access pattern

10 · The parallelism taxonomy, and enterprise/industrial context#

❓ Why this topic. Chapters 1–9 covered each technique in isolation. This chapter is a deliberately different cut through the same material -- organized by physical boundary rather than by technique name -- plus the organizational and governance concerns that only appear once this material is being run inside a company rather than a research script.

The taxonomy: organized by boundary, not by technique#

boundary	what lives here	the chapters that cover it
Inside a machine	multi-GPU data/tensor/model parallelism over NVLink -- the highest-bandwidth, lowest-latency boundary, where tensor parallelism’s frequent small communications are confined (ch. 2–3)	ch. 1–4
Across machines, within a cluster	data-parallel replicas and sharded (FSDP/ZeRO) states communicating over InfiniBand/Ethernet -- pipeline parallelism’s less-frequent boundary hand-offs also live here	ch. 1–3, 5
Across clusters / at frontier scale	the regime OPT-175B, Llama 3, and DeepSeek-V3’s real incident data (ch. 5) describe directly -- where failure is a certainty over the run’s duration, not an edge case, and dedicated on-call infrastructure is a requirement, not a nice-to-have	ch. 5–6
Serving vs. training	the same sharding vocabulary (tensor-parallel, KV-cache-sharded) applied to a latency-bound rather than throughput-bound workload -- the same physical boundaries, a different objective function	ch. 7

The reason this reorganization is worth doing explicitly, not just implicitly understood: a question like “should tensor parallelism cross this boundary” has the same answer (no, confine it within a node) at every one of these levels, because the answer is driven by a fixed hardware bandwidth fact, not by which chapter or framework happens to be in scope -- organizing by boundary makes that consistency visible in a way organizing by framework name doesn’t.

Enterprise and industrial context: governance and cost control as real engineering constraints#

Why this belongs in a technical document at all. Everything in chapters 1–9 assumes the only constraints are memory, bandwidth, and compute. Inside an actual organization, two more constraints are just as real: who is allowed to spend how much compute on what, and whether a domain-specific deployment can even use general-purpose techniques as-is.

Cost governance, concretely. The cost-modeling math from the LLM lifecycle companion notes (ch. 5) -- dollars per token, SLO-driven autoscaling -- is the technical half of this; the organizational half is the approval and budgeting process that decides which training runs get greenlit at what cluster size before a single GPU is reserved. A production training platform inside a real company typically enforces this technically, not just procedurally: per-team compute quotas, mandatory cost estimates attached to a job submission (computed from the exact formulas in this document’s chapters 1 and 4 -- parameter count, target token count, expected MFU), and automatic job termination past a budget ceiling, functioning as a hard technical guardrail on top of the human approval process rather than a replacement for it.

Domain-specific adaptation as its own governance question. The general-purpose techniques in this document (FSDP, DeepSpeed, Megatron) are architecture-and-scale concerns; a separate, real question in regulated or highly domain-specific industries (manufacturing, healthcare, finance) is whether a general-purpose LLM approach is even the right starting point versus a domain-structured one. The Industrial Large Knowledge Model (ILKM) framework (Lee & Su, 2023) is a concrete, published example of this distinction made explicit for manufacturing specifically: rather than fine-tuning a general LLM on domain text and hoping structure emerges, it proposes deliberately constructing a structured knowledge library first (human-interpretable and machine-generated industrial data, explicitly organized), then building domain instruction data and a domain-specific model on top of that structure -- a data-centric-first, model-second sequencing, which is a real and different governance decision from “just fine-tune a foundation model on our documents,” worth knowing exists as an alternative even outside manufacturing specifically.

Cross-dataset training and merging: distribution vector merging (DEM). A real, recent (Pappas & Zha, 2024, and its 2026 follow-up “OptiMer”) technique directly relevant to any organization training on multiple distinct data sources or domains: rather than fixing a single data-mixture ratio before an expensive continued-pretraining run and hoping it was the right one (this document’s sibling notes on data mixture in the LLM lifecycle chapter 1), train a separate model checkpoint per data source, extract each one’s distribution vector (the parameter shift that source’s continued training induced, relative to the shared base model), and then search -- after the fact, far more cheaply than a full retrain -- for the optimal weighted combination of those vectors to merge back into a single model. The reported advantage (OptiMer’s own benchmark) is a 15–35x reduction in search cost versus retraining with different mixture ratios directly, precisely because comparing merge weights over already-trained distribution vectors is vastly cheaper than running the full continued-pretraining process again for every candidate ratio. This is a direct, practical answer to a real enterprise scenario: “we have five different domains’ worth of data and don’t know the right mixture” no longer requires guessing the ratio upfront -- it can be decoupled into “train once per domain, then search for the merge afterward.”

📚 Further reading

• Lee & Su, 2023 -- A Unified Industrial Large Knowledge Model Framework in Industry 4.0 and Smart Manufacturing
• Pappas & Zha, 2024 -- DEM: Distribution Edited Model for Training with Mixed Data Distributions
• OptiMer, 2026 -- Optimal Distribution Vector Merging Is Better than Data Mixing for Continual Pre-Training

11 · Practical walkthroughs: from a free notebook to a trillion-parameter model#

❓ Why this chapter. Everything above this point is real, but it’s easy to read a JSON config for ZeRO-3 offload or a 16,384-GPU cluster number and feel like there’s no path from “I want to actually try this” to any of it. There is a path, and it’s continuous -- the exact same code from chapter 1 runs whether it’s on 1 GPU you’re renting for an afternoon or 16,000 GPUs a frontier lab owns. This chapter walks that path one real, costed step at a time, so the gap between “reading about FSDP” and “having run FSDP” is an afternoon and a few dollars, not a mystery.

Step 0 -- completely free: learn the training loop itself, no distributed anything#

Before touching multiple GPUs at all, run a small model’s entire training loop, start to finish, on a single free GPU (Google Colab’s free tier, or Kaggle’s free GPU notebooks both work). Andrej Karpathy’s nanoGPT is the standard choice for this precisely because it’s small enough to train in minutes and simple enough to read every line of.

git clone https://github.com/karpathy/nanoGPT
cd nanoGPT
pip install torch numpy
python data/shakespeare_char/prepare.py         # tiny dataset, tokenized in seconds
python train.py config/train_shakespeare_char.py --device=cuda --compile=False

What you should get out of this step, specifically. Watch the loss number print and go down, on purpose -- that’s the entire idea “training” refers to, made concrete, with a model small enough (a few million parameters) that a single free GPU handles it without any of chapters 1–8’s machinery. Cost: $0. Time: under 15 minutes. Do not skip this step even if you already understand the theory -- actually watching a real loss curve descend on your own screen is a different kind of understanding than reading about one.

Step 1 -- a few dollars: your first real distributed job, two rented GPUs#

Rent two small GPUs from a budget provider (Vast.ai or RunPod, both around $0.30–0.70/hour per RTX 4090, or a similar low-end card, as of mid-2026 pricing) and run the exact DDP code from chapter 1 across both of them, training a small (100M–1B parameter) model.

# on the rented machine, after cloning your training script
pip install torch
torchrun --nproc_per_node=2 train_ddp.py   # the identical script from chapter 1, now actually running on 2 real GPUs

What to actually watch for. Open a second terminal and run nvidia-smi while training runs -- watch both GPUs’ utilization climb together, and watch the loss curve match (in shape, not exact numbers) what step 0 produced, just faster. This is the moment “data parallelism” stops being a diagram and becomes something you watched happen. Realistic cost for a few hours of experimentation: under $5. This is also the right place to deliberately break something -- kill one GPU’s process mid-training and watch the whole job fail, then re-read chapter 5’s checkpoint-and-resume code and add it yourself, so the reason fault tolerance matters is something you caused and fixed, not just read about.

Step 2 -- tens of dollars: fitting a model that doesn’t fit on one GPU#

Rent a single 8-GPU node (a full A100 80GB node from a budget provider runs roughly **$8–16/hour total** depending on provider and card; renting it for a couple of hours to run a real experiment costs well under$50) and actually fine-tune a 7B-parameter open model (not train from scratch -- an individual budget realistically does fine-tuning on top of an already-pretrained model, which is a real, complete, useful skill on its own) using LoRA/QLoRA (the LLM lifecycle companion notes, chapter 2) combined with FSDP or DeepSpeed ZeRO-3 from this document’s chapter 1.

pip install transformers peft deepspeed accelerate
accelerate config     # choose: multi-GPU, DeepSpeed, ZeRO stage 3, 8 GPUs
accelerate launch finetune_7b_lora.py --model_name meta-llama/Llama-2-7b-hf --use_lora --bf16

What you should get out of this step. Before adding --use_lora and ZeRO-3, try loading the plain 7B model without either -- on most 80GB GPUs a full fine-tune of a 7B model without any memory-saving technique genuinely will not fit, and you’ll see a real CUDA out of memory error. That error, seen with your own eyes, is worth more than any paragraph of this document explaining why chapter 1’s frameworks exist -- you just personally hit the wall they’re built to get around. Then add ZeRO-3 (or LoRA, or both) and watch the identical script succeed. This exact experience -- hit the memory wall, apply the fix, watch it work -- is the single most useful hour in this entire learning path, and it costs perhaps $10–30 depending on how much you experiment.

Step 3 -- a few hundred dollars: multi-node, and feeling the network for the first time#

Rent two separate 8-GPU nodes (from a provider offering InfiniBand-connected multi-node clusters -- Lambda Labs and CoreWeave both do, at a real premium over single-node budget pricing, commonly landing in the $150–400 for a several-hour session range for a small multi-node reservation) and run the same fine-tuning job across both nodes using torchrun’s multi-node rendezvous (chapter 1) or Ray Train (also chapter 1) instead of a single-node launch.

# on node 0 (the "master")
torchrun --nnodes=2 --nproc_per_node=8 --node_rank=0 \
    --rdzv_id=job1 --rdzv_backend=c10d --rdzv_endpoint=<node0-ip>:29500 finetune_7b_lora.py
# on node 1, identical except node_rank
torchrun --nnodes=2 --nproc_per_node=8 --node_rank=1 \
    --rdzv_id=job1 --rdzv_backend=c10d --rdzv_endpoint=<node0-ip>:29500 finetune_7b_lora.py

What to watch for. Run the chapter 6 all-reduce bandwidth measurement code, comparing a collective within one node against one that spans both -- this is the first time the NVLink-vs-InfiniBand bandwidth gap from the lifecycle companion notes stops being a number in a table and becomes a number you personally measured on hardware you were, for an afternoon, actually renting.

Step 4 -- the honest wall: what a 70B-to-trillion-parameter run actually requires, and why it’s the same code#

The point of this section is not to talk you out of understanding this -- it’s to be honest about scale. Everything above this point is real, hands-on, and affordable. Training a 70B-parameter model from scratch, or anything in DeepSeek-V3’s 671B-total-parameter class, is not something an individual reasonably does on a personal budget -- but it’s worth being precise about why, using the real numbers already established earlier in this document, rather than leaving it as a vague “it’s expensive”:

what you just did (steps 0–3)	what a real frontier run needs	the actual multiplier
1–16 GPUs, a few hours	OPT-175B: ~1,024 GPUs, ~56 days	roughly 1,000x the GPU-hours
a $10–400 total bill	DeepSeek-V3: 2.788M GPU-hours, ~$5.6M	roughly five to six orders of magnitude more spend
a single rented node or two, booked for an afternoon	Llama 3 405B: ~16,384 H100s, held for 54 continuous days, with a dedicated on-call team responding to ~419 interruptions	an entire organization’s infrastructure team, not a rented reservation

The part worth internalizing precisely. None of the code changes at that scale -- it’s the same torchrun/DeepSpeed/FSDP concepts from chapter 1, the same tensor/pipeline/expert parallelism from chapter 2, the same checkpointing and failure-recovery patterns from chapter 5, just composed across many more GPUs, with a dedicated team on-call for exactly the incidents chapter 5’s real-incident-data section describes. The gap between what you did in steps 0–3 and a frontier run is entirely a gap of money, GPU count, and organizational infrastructure -- not a gap of different knowledge. Having actually run the small version means you understand, concretely rather than abstractly, exactly what a much bigger version of the same job is doing at every step.

📚 Further reading

• Karpathy -- nanoGPT (github.com/karpathy/nanoGPT), the standard starting point for step 0
• Hugging Face -- PEFT + Accelerate + DeepSpeed integration guide, for step 2
• Provider documentation for whichever budget GPU cloud you choose -- pricing and availability change often enough that checking current rates directly, rather than trusting any single cited number for long, is the right habit

12 · How this shows up in interviews#

❓ Why this chapter, and how to read it. This is not a question bank -- it’s notes on the shapes of scenarios that come up when this material gets tested in an ML infrastructure or distributed-systems-for-ML interview, what’s actually being evaluated underneath the specific wording, and how to structure a strong answer out loud. The specific numbers an interviewer uses will vary; the underlying reasoning pattern in each scenario type below won’t.

The “does it fit” sizing scenario#

What it sounds like. “You have a 70B parameter model and 8 GPUs with 80GB each. How would you approach training it?”

What’s actually being tested. Not whether you can recite a framework name -- whether you can do the memory arithmetic from the LLM lifecycle companion notes’ chapter 8 ($16N$ bytes per parameter for full-precision training state) out loud, notice it doesn’t fit, and then reach for a specific tool because the numbers demanded it, rather than naming DeepSpeed first and justifying it after. A strong answer walks through: compute $16 \times 70\text{B} \approx 1.12\text{TB}$ of state, compare against $8 \times 80\text{GB} = 640\text{GB}$ total available, notice it doesn’t fit even fully sharded evenly (and explain why -- ZeRO-3/FSDP alone gives $16N/P$ per GPU, which at $P=8$ is $140$GB per GPU, still over budget once activations are added), and only then propose a specific combination -- likely ZeRO-3 plus activation checkpointing plus possibly CPU offload for optimizer state, precisely because the raw numbers show sharding alone isn’t quite enough at this ratio. The trap to avoid: jumping straight to “use DeepSpeed” without showing the arithmetic that makes it necessary -- an interviewer who’s paying attention will ask “why not just FSDP?” specifically to see if the numbers were ever actually computed or just pattern-matched from a keyword.

The silent-corruption debugging scenario#

What it sounds like. “Training has been running for two weeks across 200 GPUs. The loss looks fine, but eval accuracy has quietly gotten worse over the last three days. Walk me through how you’d investigate.”

What’s actually being tested. Whether “the loss looks fine” is treated as reassuring or as itself suspicious -- a good answer immediately notices that a normal-looking loss curve with degrading downstream quality is close to the textbook symptom of a silent, implicit failure (chapter 6’s whole framing), not a loud one, and reasons through the checklist in that order: first rule out a data pipeline problem (has the data mixture or a specific shard silently changed -- chapter 1 of the lifecycle notes’ data-validation point), then check for a sharding/parity issue introduced by a recent config change (chapter 6’s shard-parity validation), then check whether a recent throughput change correlates with the timing (chapter 6’s throughput-regression monitor, since a quietly failing node sometimes shows up as a speed change before a quality change). The trap to avoid: jumping straight to “maybe the learning rate is too high” -- that’s a real possibility, but it’s the hypothesis you reach for last, after ruling out infrastructure-side silent failures first, because a plausible-sounding algorithmic explanation is exactly the kind of answer that feels satisfying without actually being diagnostic.

The “design the checkpointing strategy” scenario#

What it sounds like. “Design a checkpointing approach for a training job expected to run for a month across a cluster where node failures happen roughly every few hours.”

What’s actually being tested. Whether checkpointing is treated as a single decision (“save every N steps”) or as the multi-part system it actually is: checkpoint frequency (a tradeoff between recompute-on-failure cost and steady-state I/O overhead -- more frequent checkpoints waste less work per failure but cost more disk/network bandwidth continuously), checkpoint format (resharding-aware, chapter 3, so the checkpoint survives a change in GPU count if the cluster allocation shifts), trigger-based checkpointing on a preemption signal (chapter 5’s SIGTERM handler) as a complement to periodic checkpointing rather than a replacement for it, and orchestrator-level automatic restart (chapter 5’s Slurm --requeue/Kubernetes restartPolicy). A strong answer explicitly connects the checkpoint frequency to the stated failure rate in the prompt -- “failures every few hours” should visibly inform a specific frequency choice, not just a generic “checkpoint often” -- which is the detail that separates a memorized answer from a reasoned one.

The “why is utilization low” performance scenario#

What it sounds like. “Your training job is only hitting 25% GPU utilization on paper, despite being correctly parallelized. What would you check, in order?”

What’s actually being tested. Whether the candidate has an actual diagnostic order rather than a list of things that could be wrong in no particular sequence. A strong order, and why: (1) check whether the GPU is actually starved for data first -- a data-loading bottleneck (CPU-bound preprocessing, slow disk I/O) is the single most common cause and the cheapest to rule out with torch.profiler showing idle gaps between compute kernels (chapter 4); (2) check whether communication/compute overlap is actually happening (chapter 4’s prefetch and profiler trace check) -- this is the second most common cause and specifically implicates parallelism configuration, not data; (3) only then suspect a fundamentally compute-bound ceiling (e.g., a genuinely small batch size relative to the hardware, or an inefficient kernel) as the explanation, because it’s the least common of the three in practice and the hardest to “fix” cheaply. The trap to avoid: naming “increase batch size” as a first response -- it’s a real lever, but proposing it before profiling where time is actually going skips the diagnostic step the question is testing for.

The open-ended systems-design scenario#

What it sounds like. “Design the training infrastructure for a team that wants to pretrain a 30B parameter model from scratch.”

What’s actually being tested. Breadth across every chapter of this document and the judgment to sequence them, not depth on any single piece. A strong answer touches, roughly in this order: (1) a scaling-law-informed sizing decision for how much data this compute budget justifies (LLM lifecycle companion notes, ch. 1) before any infra decision at all; (2) a parallelism plan sized to the actual GPU count available, with the reasoning for why that combination (ch. 2 of this document); (3) a sharding strategy chosen by the same memory arithmetic as the first scenario above; (4) an orchestration and checkpointing plan matched to the expected failure rate at that cluster size (ch. 5); (5) an explicit observability plan -- what gets logged, what triggers an alert -- named before the run starts, not improvised after the first incident (ch. 5–6). The trap to avoid: describing only the parallelism/sharding technical choices and never mentioning monitoring, checkpointing, or failure recovery at all -- a design that would work perfectly on hardware that never fails is not a real production design, and every interviewer who’s actually run a training job at scale knows it.

The conceptual-contrast question, and how to answer it well#

What it sounds like. “What’s the difference between data parallelism and tensor parallelism, and when would you use each?” or any of its siblings (ZeRO stage 2 vs. 3, pipeline vs. tensor parallelism, checkpointing vs. sharding).

What separates a good answer from a memorized one. Not defining both terms correctly -- most candidates can do that -- but stating the actual decision rule for choosing between them, grounded in a concrete number: tensor parallelism only when a single layer’s weights don’t fit, confined to a node because of the bandwidth argument (this document’s chapter 3, and the LLM lifecycle notes’ chapter 8); data parallelism whenever the model does fit and the goal is more throughput, not more memory. Naming the bandwidth number itself (NVLink’s roughly 900GB/s versus InfiniBand’s order-of-magnitude-lower effective bandwidth) unprompted, as the actual reason for the node-boundary rule rather than reciting the rule as received wisdom, is consistently the detail that reads as real understanding rather than memorized vocabulary.

---

Zero → Frontier Engineering, log 04. A practical reference, not a tutorial: assumes logs 01–03 and gets denser from there. Named production incidents, papers, and numbers throughout are cited inline per chapter under “further reading.”