Training at Scale

Training is where AI hits your cluster hardest. One job can hold hundreds or thousands of GPUs for weeks. It pushes HBM capacity, NVLink and the scale-out fabric, the parallel file system and your on-call rotation all at once. This chapter is for the people who run those jobs and the people who support them. It covers what happens inside a training step, how to account for memory before you ask for an allocation, the parallelism strategies and how they map onto real hardware, the frameworks worth knowing in 2026, and what large labs have published about keeping giant jobs alive. The hardware itself is covered in Accelerators and Interconnects. Inference is its own chapter: Inference at Scale.

Anatomy of a Training Step

Every synchronous data-parallel training job runs the same loop. Knowing its phases is the first step in figuring out why a job is slow.

1. Load a batch. Read pre-tokenized samples, move them to the GPU, and in model-parallel setups send them to the right ranks.
2. Forward pass. Compute activations layer by layer and keep what the backward pass will need.
3. Loss. For language models this is cross-entropy over the vocabulary. The output projection and softmax can be a large share of memory.
4. Backward pass. Compute gradients. It costs about twice the forward pass, which is where the "6" in the 6N FLOPs-per-token rule comes from: 2N forward plus 4N backward.[5]
5. Gradient synchronization. All-reduce (DDP) or reduce-scatter (FSDP/ZeRO) across data-parallel ranks. Good implementations overlap this with the backward pass. DDP groups gradients into buckets and starts an asynchronous all-reduce as soon as a bucket is ready.[18]
6. Optimizer step. Update weights with AdamW or a similar optimizer. This step is memory-bound and touches every optimizer state byte.
7. Periodic work. Logging, evaluation and checkpointing. Done badly, these stall every GPU in the job.

Where the Time Goes

Wall-clock time per step breaks down into four buckets. The first is useful math: the big matrix multiplies on tensor cores. The second is memory-bound kernels such as normalization, activations, optimizer updates and attention without fused kernels. The third is exposed communication, meaning collectives the framework could not hide behind compute. The fourth is idle time: pipeline bubbles, data-loader stalls, host-side launch overhead and waiting on a slow rank. The standard efficiency metric is Model FLOPs Utilization (MFU): achieved model FLOPs per second divided by the hardware's peak. Estimate it as 6 x N x (tokens per second) / (GPU count x peak FLOP/s per GPU). For calibration, Meta reports 38-43% BF16 MFU (380-430 TFLOPs per GPU) for Llama 3 405B pre-training on 8,192 to 16,384 H100s.[7] A well-tuned research job in the 30-50% range is doing fine. If you are at 10%, something is exposed: communication, data loading or tiny per-GPU batches.

Field note: Before tuning parallelism, profile one step with the PyTorch profiler or Nsight Systems and look for gaps on the GPU compute stream. Gaps next to NCCL kernels point to exposed communication. Gaps at the start of the step point to the data loader. Gaps everywhere with small kernels point to host overhead, which torch.compile or CUDA graphs often fix. See Observability for tooling.

Memory Accounting for Training

Most "CUDA out of memory" tickets come from people who never did this arithmetic. Training memory has two parts. Model states are weights, gradients and optimizer states, and they scale with parameter count. Residual states are activations, temporary buffers and fragmentation, and they scale with batch size, sequence length and model width.

The 16 Bytes per Parameter Rule

The ZeRO paper gives the standard breakdown for mixed-precision training with Adam. Call the parameter count Psi. You keep 16-bit parameters (2 Psi bytes) and 16-bit gradients (2 Psi bytes). The optimizer holds an FP32 master copy of the parameters, FP32 momentum and FP32 variance: 4 Psi bytes each, 12 Psi bytes in total. That adds up to 16 bytes per parameter before a single activation is stored. The paper's own example is a 1.5B-parameter GPT-2, which needs at least 24 GB for model states even though the 16-bit weights alone are only 3 GB.[1]

Worked Example: 7B and 70B Models

The table below applies the ZeRO formulas to 7B and 70B models: plain DP, then stage 1 (P_os, shard optimizer states: 4 Psi + 12 Psi/N_d), stage 2 (P_os+g, also shard gradients: 2 Psi + 14 Psi/N_d) and stage 3 (P_os+g+p, shard everything: 16 Psi/N_d), where N_d is the data-parallel degree.[1] Values are model states only, in GB (10^9 bytes), per GPU.

Note: computed from the ZeRO paper's formulas with K = 12 for mixed-precision Adam.[1] Activations, buffers and fragmentation are extra.

What this means in practice: a 7B model does not fit for full training on a single 80 GB GPU (the H100s Meta used have 80 GB of HBM3[7]), no matter how small the batch. Shard it across one 8-GPU node with ZeRO-3/FSDP and model states drop to 14 GB per GPU, leaving plenty of room for activations. A 70B model at 140 GB per GPU does not fit on an 8-GPU node of 80 GB parts even with full sharding. You need more nodes, or tensor/pipeline parallelism, or both. The ZeRO paper's Table 1 shows the same pattern: a 7.5B model drops from 120 GB per device with plain DP to 1.88 GB with all three stages at N_d = 64.[1]

Activation Memory

Korthikanti et al. derive per-layer activation memory for a standard transformer in 16-bit precision: sbh(34 + 5as/h) bytes, where s is sequence length, b is micro-batch size, h is hidden size and a is the number of attention heads.[4] Take an illustrative 7B-class shape (h = 4096, 32 layers, 32 heads) at s = 4096 and b = 1. The formula gives about 3.25 GB per layer and 104 GB for the full stack. Most of that comes from the 5as/h term, which is the attention score matrix that grows quadratically with sequence length. Fused attention kernels such as FlashAttention do not store that large intermediate matrix for the backward pass,[37] which removes most of the quadratic term. Without that term, the formula gives 34sbh per layer (about 18 GB over 32 layers in this example). From there you have three levers:

• Full activation recomputation keeps only each layer's input (2sbhL total, about 1.1 GB here) and recomputes everything else during backward. The authors measured a 30-40% execution-time overhead.[4]
• Selective recomputation recomputes only the cheap-to-recompute, memory-heavy parts of attention. Combined with sequence parallelism it cut activation memory 5x and cut recomputation overhead by over 90% in their experiments.[4]
• Tensor plus sequence parallelism divides the 34sbh term by the TP degree t.[4]

Parallelism Strategies

Every parallelism strategy cuts some dimension of the problem across GPUs: the batch, the model states, the individual weight matrices, the layers, the sequence, or the experts. Each cut has its own communication pattern and its own bandwidth needs.

Distributed Data Parallel (DDP)

Each rank holds a full replica of the model and processes its own slice of the global batch. Gradients are averaged with all-reduce, bucketed and overlapped with the backward pass.[18] DDP is simple, robust and fast, and it only works if the full 16-bytes-per-parameter model state plus activations fits on one GPU. For the many research models under about a billion parameters, DDP is still the right answer.

ZeRO and FSDP

ZeRO removes the redundancy in data parallelism by partitioning model states across DP ranks. Stages 1 and 2 keep DDP's communication volume. Stage 3 (parameter partitioning) makes memory drop linearly with DP degree for a modest 50% increase in communication volume.[1] PyTorch's implementation of this idea is Fully Sharded Data Parallel. The current version, FSDP2 (torch.distributed.fsdp.fully_shard), shards each parameter along dim-0 as a DTensor. It all-gathers parameters before each module's forward and backward and reduce-scatters gradients afterward.[13] The PyTorch tutorial describes FSDP as DDP's all-reduce split into a reduce-scatter and an all-gather, and marks FSDP1 as deprecated.[14]

Several FSDP2 details matter in practice:[13]

• Apply fully_shard bottom-up, once per transformer block and then on the root. Each call creates one communication group. Sharding only the root gives one huge blocking all-gather with no overlap, which the docs call "almost never what you want."
• reshard_after_forward=True is ZeRO-3 behavior. False keeps the gathered parameters until backward, trading memory for one less all-gather, much like ZeRO-2. An integer reshards to a smaller group, for example the GPUs within a node. Llama 3 used the "do not reshard after forward" choice to avoid the extra all-gather in backward.[7]
• A 2D device mesh gives HSDP (hybrid sharding): parameters are sharded along one mesh dimension and replicated along the other. Typically you shard within a node or rack and replicate across them.
• MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) computes in BF16 and reduces gradients in FP32. Because FSDP keeps the sharded FP32 parameters anyway, it needs no extra FP32 master copy.
• CPUOffloadPolicy moves parameters, gradients and optimizer states to host memory and runs the optimizer step on CPU. It is slow, but sometimes it is the only way to fit a model onto a small allocation.

Tensor Parallelism

Megatron-LM's tensor parallelism splits each transformer block's weight matrices across GPUs. The first MLP matrix is split by columns and the second by rows, so the block needs only one all-reduce in the forward pass and one in the backward pass. Attention heads are split the same way.[2] Every layer adds a latency-sensitive collective to the critical path, so TP only performs well on very fast links. Narayanan et al.'s first takeaway is that tensor parallelism should generally be used up to degree g on g-GPU servers, with pipeline parallelism scaling beyond the server.[3] In PyTorch, TP is expressed with parallelize_module and styles such as ColwiseParallel, RowwiseParallel and SequenceParallel, all built on DTensor.[40]

Sequence and Context Parallelism

Two related ideas share this name. Sequence parallelism (Korthikanti et al.) is a companion to TP. It shards the parts TP leaves replicated, such as LayerNorm and dropout, along the sequence dimension, so activation memory divides cleanly by the TP degree.[4] Megatron Core recommends always enabling it with TP.[19] Context parallelism splits the whole sequence, attention included, across GPUs. It exchanges key/value blocks between ranks, in the style of Ring Attention, which overlaps KV communication with blockwise attention compute and supports sequences up to device-count times longer.[36] CP is what makes long-context training feasible. Llama 3's long-context stage ran at a 131,072-token sequence length with CP = 16 on 16,384 GPUs.[7] Megatron Core suggests CP for sequences of 8K tokens and up.[19]

Pipeline Parallelism

Pipeline parallelism assigns contiguous groups of layers to different GPUs, called stages, and streams micro-batches through them. GPipe introduced batch-splitting pipelining for general sequential networks.[38] The cost is the pipeline bubble: stages sit idle while the pipeline fills and drains. With p stages and m micro-batches per batch, the bubble fraction is (p - 1)/m, so you need m much larger than p.[3] Worked example: at p = 16 and m = 32 the bubble is 15/32, about 47% of ideal compute time. At m = 128 it falls to about 12%.

• 1F1B (one-forward-one-backward, from PipeDream) limits how many micro-batches' activations are in flight. The first stage still has to hold activations for p micro-batches, which works out to L layers' worth regardless of p.[4]
• Interleaved 1F1B gives each GPU v non-contiguous "virtual" stages. This cuts the bubble by a factor of v and raises communication by the same factor.[3]
• Zero-bubble schedules split the backward pass into input-gradient and weight-gradient computations and reorder them. The authors report up to 23% higher throughput than 1F1B at similar memory, and 31% with the memory constraint relaxed.[10] DeepSeek-V3's DualPipe overlaps compute and communication across forward and backward chunks.[9]

Pipelines also have load-balance problems. Meta found that the first stage uses more memory (embeddings plus warm-up micro-batches) and the last stage becomes the latency bottleneck (output projection plus loss). They also found that standard schedules forced the per-GPU batch size to be divisible by the stage count, so they modified the schedule to allow arbitrary micro-batch counts.[7] PyTorch ships torch.distributed.pipelining, still labeled alpha, with GPipe, 1F1B, Interleaved 1F1B and Looped BFS schedules.[39]

Expert Parallelism for MoE

A Mixture-of-Experts model replaces dense MLPs with many expert MLPs and a router that sends each token to a few of them. DeepSeek-V3, for example, has 671B total parameters but activates only 37B per token.[9] Expert parallelism puts different experts on different GPUs and moves tokens to their experts and back with all-to-all collectives, which are the most network-hostile pattern in training. DeepSeek-V3 used 16-way PP, 64-way EP spanning 8 nodes, and ZeRO-1 DP on 2,048 H800 GPUs (8 per node on NVLink/NVSwitch, InfiniBand between nodes), with custom cross-node all-to-all kernels.[9] In Megatron Core, EP is set with --expert-model-parallel-size and an all-to-all token dispatcher. If you combine EP with TP, you must also enable sequence parallelism.[19][43] Router load balancing and token-drop policy affect both model quality and step time. Put expert-imbalance metrics in your logs.

Multi-Dimensional Parallelism

Real large-model jobs combine all of the above. The GPU count is the product of the degrees. Megatron Core writes it as TP x PP x CP x EP x DP. Its example trains a Llama-3 70B-shaped model on 64 GPUs with TP = 4, PP = 4, CP = 2 and DP = 2.[19] Llama 3 405B used "4D parallelism", with GPUs grouped in the order [TP, CP, PP, DP] from innermost to outermost and DP implemented as FSDP.[7] Its published configurations:

Source: Llama 3 paper, Table 4. Global batch was 16M tokens in all three.[7]

TorchTitan measured the payoff of combining techniques on H100s against optimized baselines: 65.08% speedup for Llama 3.1 8B on 128 GPUs (1D), 12.59% for 70B on 256 GPUs (2D), and 30% for 405B on 512 GPUs (3D).[8]

Mapping Parallelism onto Hardware

The rule is simple: the more often a parallel dimension communicates, the closer together its GPUs must be. Llama 3's [TP, CP, PP, DP] ordering puts it into practice. The innermost dimension (TP, a collective every layer) is placed on adjacent ranks, which on a standard launcher means GPUs in the same server. Meta's servers had eight GPUs connected by NVLink, and TP was 8.[7]

• TP (and usually SP) inside the NVLink domain. On an 8-GPU HGX-style node, TP is at most 8. Rack-scale NVLink systems widen the domain. NVIDIA's GB200 NVL72 connects 72 Blackwell GPUs in one NVLink domain.[42] That leaves room for larger TP or EP groups without touching the scale-out network. See Interconnects.
• CP and EP next. Both are bandwidth-heavy. Keep them in the NVLink domain when you can, and otherwise on the best rail-aligned scale-out links.
• PP across nodes. Point-to-point activation transfers are small and tolerate slower links. That is why PP is the standard way to scale depth across servers.[3][39]
• DP/FSDP outermost. Gradient reduction happens once per step (or once per layer group under FSDP) and overlaps with compute. HSDP shards within a fast domain and replicates across slower ones.

Rank ordering comes from the launcher. torchrun places the LOCAL_RANK processes of one node on consecutive ranks,[17] so building your DeviceMesh with the TP dimension last (fastest-varying) keeps TP groups inside a node. Slurm topology-aware placement (see Scheduling) keeps DP and PP groups under as few switch tiers as possible.

Field note: The most common facilitation mistake is a user setting TP = 16 on 8-GPU nodes because "the model is big." The job runs, but every layer now crosses the scale-out NIC, and throughput collapses. Keep TP within one NVLink domain. Add PP or FSDP to go bigger.

Frameworks

PyTorch Native

PyTorch 2.14 is the current release as of this writing. Its distributed highlights include a new nccl2 backend and fault tolerance as a first-class c10d concept, with in-place process-group reconfiguration and a Flight Recorder that works for any backend.[35] The building blocks are:

• DeviceMesh and DTensor. init_device_mesh("cuda", (dp, tp), mesh_dim_names=("dp", "tp")) describes the N-dimensional layout of GPUs.[46] DTensor records how each tensor is placed on that mesh (Shard(dim), Replicate(), Partial()). FSDP2, TP and distributed checkpointing all build on it.
• DDP and FSDP2, described above.
• torch.compile. By default it compiles through the Inductor backend.[47] With DDP, wrap in DDP first and compile second, so the compiler can insert graph breaks at bucket boundaries and keep all-reduce overlapped with backward.[18]
• torch.distributed.pipelining for PP (alpha) and torch.distributed.checkpoint (DCP) for sharded checkpoints.

TorchTitan

TorchTitan is the PyTorch team's clean-room reference for large-scale training. It is the best place to see how FSDP2, TP, PP, CP, activation checkpointing, torch.compile, Float8, async DCP checkpointing and Flight Recorder fit together in about as little code as possible. Recent features include MXFP8 training on Blackwell, TorchFT integration, checkpointable data loading and BF16 optimizer states. AMD maintains an optimized fork for its GPUs.[27] It is also the companion code to the TorchTitan paper.[8]

Megatron-LM, Megatron Core and NeMo

Megatron Core is NVIDIA's composable library of GPU-optimized transformer building blocks. It supports TP, PP, DP, EP and CP and FP16/BF16/FP8/FP4 precision. Megatron-LM is the reference training scripts built on it. Megatron Bridge handles Hugging Face to Megatron checkpoint conversion in both directions.[20] Its parallelism guide documents the flags you will see in user job scripts: --tensor-model-parallel-size, --sequence-parallel, --pipeline-model-parallel-size, --num-layers-per-virtual-pipeline-stage, --context-parallel-size, --expert-model-parallel-size, --use-distributed-optimizer, and the overlap switches --overlap-grad-reduce, --overlap-param-gather and --tp-comm-overlap.[19] For LLM work, the NVIDIA NeMo Framework container now defers to Megatron Bridge's documentation, starting with the 26.02 container.[34] Megatron Bridge 0.5.0 was released in June 2026.[33]

DeepSpeed

DeepSpeed made ZeRO practical. Its main selling point is still that you turn on ZeRO stages through the JSON config without changing model code, and that ZeRO-Infinity can offload to CPU and NVMe.[21] The config enforces the batch identity every facilitator should know: train_batch_size = train_micro_batch_size_per_gpu x gradient_accumulation_steps x number of GPUs.[22] The project remains active. Its 2026 news includes Muon optimizer support and SDMA-offloaded collectives for ZeRO-3 on AMD GPUs.[49]

JAX and XLA

JAX takes the compiler route. You annotate how arrays are sharded over a device mesh, and XLA partitions the computation and inserts collectives. shard_map gives explicit per-shard control when you need it, and the two approaches compose.[26] Multi-node JAX runs one process per host (or per GPU), each calling jax.distributed.initialize(). XLA uses high-speed links such as TPU ICI or NVLink where they exist and host networking elsewhere.[25] JAX is the native path on Google TPUs (see Cloud) and is fully workable on GPU clusters.

Hugging Face Accelerate

Accelerate wraps an existing PyTorch loop with an Accelerator object and a few changed lines, then dispatches to DDP, FSDP or DeepSpeed based on the launch config.[28] Setting fsdp_version=2 selects FSDP2. The default is still 1 for backward compatibility, so check what your users are actually running.[29] Accelerate is the lowest-effort path for fine-tuning jobs built on the Hugging Face stack.

Field note: For a campus cluster, standardize on two paths and document them well. PyTorch-native FSDP2 (or TorchTitan) covers people writing their own models. Megatron Bridge/NeMo or Accelerate covers people adapting existing checkpoints. Supporting all seven frameworks equally is how a facilitation team burns out.

Precision, Activation Checkpointing and Gradient Accumulation

Mixed Precision: BF16 as the Baseline

BF16 compute with FP32 master weights and FP32 gradient reduction is the default recipe on current hardware. FP16 needs dynamic loss scaling to keep gradients from overflowing or underflowing.[24] BF16 training does not require loss scaling.[22] Precision formats are covered in more depth in The Convergence.

FP8 Training and Transformer Engine

NVIDIA Transformer Engine (TE) provides FP8 on Hopper, Ada and Blackwell GPUs, plus MXFP8 and NVFP4 on Blackwell. It ships drop-in modules (te.Linear, full transformer layers) that track the scaling factors FP8 needs internally.[23] Two FP8 encodings are used. E4M3 (range +/-448) is typically used for forward activations and weights. E5M2 (range +/-57344) trades precision for dynamic range and suits gradients.[24] A minimal TE usage pattern:

import torch
import transformer_engine.pytorch as te
from transformer_engine.common import recipe

model = te.Linear(768, 3072, bias=True)
inp = torch.randn(2048, 768, device="cuda")
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)

with te.autocast(enabled=True, recipe=fp8_recipe):
    out = model(inp)
out.sum().backward()

Adapted from the Transformer Engine documentation example.[23]

FP8 now works at frontier scale. DeepSeek-V3 was pre-trained with an FP8 mixed-precision framework that uses fine-grained tile- and block-wise quantization. In validation runs of about 1 trillion tokens, its relative loss error against a BF16 baseline stayed below 0.25%.[9] Megatron Core and TorchTitan both expose FP8 and MXFP8 recipes.[20][27] Treat FP8 as a throughput optimization to validate against a BF16 loss curve, not as a default to switch on blindly.

Activation Checkpointing

Activation checkpointing (recomputation) trades compute for memory, as quantified above. In PyTorch, use torch.utils.checkpoint.checkpoint(fn, *args, use_reentrant=False), usually wrapped around each transformer block.[45] TorchTitan supports both full and per-op selective checkpointing.[27] Start with no checkpointing, then selective, then full, and stop at the first setting that fits.

Gradient Accumulation

Gradient accumulation runs several micro-batches before each optimizer step, so the global batch can be larger than memory allows. Skip gradient communication on every micro-batch except the last. FSDP2 does this with model.set_requires_gradient_sync(False), the equivalent of FSDP1's no_sync.[13] Remember that under PP, micro-batches are also what fill the pipeline, so accumulation and bubble size are linked.[3]

Scaling Laws and Compute Budgets

Kaplan and Chinchilla

Kaplan et al. (2020) showed that language-model loss follows power laws in model size, dataset size and compute across more than seven orders of magnitude. They concluded that compute-optimal training uses very large models trained on relatively modest data and stops well before convergence.[5] Hoffmann et al. (2022, "Chinchilla") trained over 400 models from 70M to over 16B parameters on 5 to 500B tokens. They found instead that parameters and tokens should be scaled equally: every doubling of model size should come with a doubling of training tokens. Their 70B Chinchilla model, trained on 1.4 trillion tokens (20 tokens per parameter) with the same compute as the 280B Gopher, outperformed Gopher.[6] Modern labs fit their own scaling laws on their own data. Meta did this to choose the 405B size for its 3.8 x 10^25 FLOP budget.[7] Compute-optimal is not always inference-optimal, either. Smaller models trained on more tokens cost less to serve, as the Chinchilla paper itself points out.[6]

The 6ND Estimate

Training compute is approximately C = 6 x N x D FLOPs, where N is non-embedding parameters and D is training tokens.[5] This formula is how you turn a research idea into an allocation request.

• Sanity check against a published run: Llama 3 405B on 15.6T tokens gives 6 x 405 x 10^9 x 15.6 x 10^12 = 3.79 x 10^25 FLOPs, which matches the 3.8 x 10^25 Meta reports.[7] At Meta's reported 400 TFLOP/s per GPU on 16,384 GPUs, that is about 67 days of pure compute, before counting any interruptions.
• Chinchilla: 6 x 70 x 10^9 x 1.4 x 10^12 = 5.9 x 10^23 FLOPs.[6]
• A research-scale proposal: a 1B-parameter model at the Chinchilla ratio (20B tokens) needs 6 x 10^9 x 20 x 10^9 = 1.2 x 10^20 FLOPs. At a sustained 400 TFLOP/s per GPU, that is 1.2 x 10^20 / 4 x 10^14 = 3 x 10^5 GPU-seconds, about 83 GPU-hours. At half that throughput, which is realistic for small models with short sequences, budget about 170 GPU-hours plus headroom for restarts and evaluation.

For MoE models, use the active parameter count for N. DeepSeek-V3's 37B active parameters over 14.8T tokens gives about 3.3 x 10^24 FLOPs by this rule.[9]

Fine-Tuning versus Pretraining

Most university GPU hours go to fine-tuning, evaluation and small-model pretraining, not frontier-scale pretraining. Know which one a user needs before sizing their request.

• Full fine-tuning has the same memory profile as pretraining (16 bytes per parameter plus activations) but uses far fewer tokens. A 7B full fine-tune is an FSDP job on one 8-GPU node.
• LoRA freezes the base weights and trains small low-rank update matrices. On GPT-3 175B it reduced trainable parameters by 10,000 times and GPU memory by 3 times compared with Adam fine-tuning, matched or beat full fine-tuning quality, and added no inference latency.[11] Arithmetic: a frozen BF16 7B base is 14 GB, and the adapter's optimizer states are negligible.
• QLoRA backpropagates through a frozen 4-bit base model into LoRA adapters, using the 4-bit NormalFloat (NF4) type, double quantization and paged optimizers. It fine-tuned a 65B model on a single 48 GB GPU while preserving 16-bit fine-tuning task performance.[12]
• Hugging Face PEFT implements LoRA and related methods and integrates with Transformers and Accelerate.[30] FSDP2's per-parameter sharding handles mixed frozen and trainable parameters, and even NF4 all-gathers for QLoRA, without the workarounds FSDP1 needed.[14]

Field note: When someone asks for "8 H100s for a week to fine-tune a 70B model," ask whether QLoRA or LoRA on a smaller allocation would answer their research question. Often it will, and the queue will thank you. Save full-parameter runs for work that actually needs them.

Checkpointing, Fault Tolerance and Data at Scale

What the Llama 3 Run Taught Us

Meta's report is the best public data on failure rates at scale. During a 54-day snapshot of Llama 3 405B pre-training on up to 16K H100s, there were 466 job interruptions: 47 planned (firmware upgrades, configuration or dataset changes) and 419 unexpected. About 78% of the unexpected ones were confirmed or suspected hardware issues, and GPU issues alone were 58.7%. Only three needed significant manual intervention. Automation handled the rest, and the job still achieved higher than 90% effective training time.[7]

Source: Llama 3 paper, Table 5; counts and percentages as published.[7]

The same paper records two effects that matter for facility planning. Throughput varied 1-2% with time of day because higher midday temperatures affected GPU DVFS. Synchronized power swings across tens of thousands of GPUs, for example when every GPU waits on a checkpoint or collective, caused data-center-level fluctuations on the order of tens of megawatts.[7] See Facilities.

Distributed and Asynchronous Checkpoints

Never gather a large model onto rank 0 to save it. PyTorch Distributed Checkpoint (DCP) has every rank write its own shards in parallel, producing at least one file per rank. It also reshards at load time, so you can save on one topology and resume on another.[15] dcp.async_save first copies state to CPU staging memory and then writes from a background thread, so training continues. This costs host RAM equal to the per-rank checkpoint size, and the recipe recommends at most one async save in flight.[16] TorchTitan measured a 5-15x reduction in checkpoint overhead with async DCP compared with synchronous DCP for Llama 3.1 8B.[8] Size storage for the bursts. Meta's Tectonic fabric for Llama 3 offered 240 PB with 2 TB/s sustained and 7 TB/s peak throughput, and bursty checkpoint writes of 1 MB to 4 GB per GPU were still a named challenge.[7] Storage design is covered in Memory and Storage.

Elastic Training and Restarts

torchrun supports fault-tolerant (--max-restarts) and elastic (--nnodes=MIN:MAX) launches through the c10d rendezvous backend. On any failure or membership change it kills all surviving workers and restarts them. That makes your checkpoint frequency your loss budget, and RANK is not stable across restarts.[17] Newer tools go finer-grained. TorchFT provides per-step fault tolerance: it heartbeats replicas, reconfigures process groups, and recovers live from healthy peers for fault-tolerant DDP and HSDP.[31] The NVIDIA Resiliency Extension (NVRx), marked experimental, adds hung-rank detection, in-job restart without reallocating Slurm nodes, async and local checkpointing, and straggler detection.[32]

Stragglers and Hangs

In synchronous training, the slowest GPU sets the pace for the whole job. Meta relied heavily on PyTorch's NCCL Flight Recorder to diagnose hangs and built tools to find the few process groups with suspicious, slow communication.[7] TorchTitan integrates Flight Recorder,[8] and PyTorch 2.14 extends it beyond NCCL.[35] On your own cluster, run node health checks (GPU diagnostics, NCCL all-reduce tests, ECC and XID counters) in the Slurm prolog, drain nodes that fail, and keep a per-node history. The same few nodes cause most of the trouble.

Field note: Work out the checkpoint interval instead of guessing. If a job loses on average half an interval of work per failure, and checkpoints cost a few seconds with async DCP, checkpointing every 15-30 minutes is cheap insurance on a multi-node job. Test the resume path in week one, not after the first node failure.

Data Pipelines

A stalled data loader wastes the whole allocation just as surely as a slow network. Tokenize and pack sequences offline into large shard files rather than millions of small files, which are what kill parallel file systems. Give each DP rank a disjoint shard stream. Make loader state part of the checkpoint so a resumed job does not repeat or skip data. TorchTitan's loader is checkpointable for this reason.[27] The JAX documentation sets out the trade-off clearly. Distributed loading, where each process reads only the shards its devices need, is more efficient but more complex. Having every process load the full global data is simpler but spends network bandwidth and can block the training loop.[48] Stage datasets to node-local NVMe when they fit, and keep the parallel file system for checkpoints and shared corpora (see Memory and Storage).

Launching Multi-Node FSDP2 on Slurm

A Minimal FSDP2 Training Skeleton

The skeleton below uses the APIs as documented for PyTorch 2.14: init_device_mesh, bottom-up fully_shard with a MixedPrecisionPolicy, an optimizer built after sharding, and DCP save/load through a Stateful wrapper using get_state_dict/set_state_dict.[13][14][15][16] We ran it on PyTorch 2.14.0 (CPU build, gloo backend, 2 processes under torchrun, with a reduced model size) to confirm that training, async checkpointing and resume all work. The CUDA path uses the same calls but was not run here. Async saves use a separate gloo process group so background checkpoint collectives do not contend with training collectives. In our CPU test, async saves on the default group hung. The model is a toy. Swap in your own model and a real data loader.

# train_fsdp2.py - minimal FSDP2 skeleton (PyTorch 2.14 APIs)
# Launch: torchrun --nnodes=N --nproc-per-node=G train_fsdp2.py
import os
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.distributed.checkpoint.stateful import Stateful

CKPT_DIR = os.environ.get("CKPT_DIR", "checkpoints")


class Block(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.norm = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, num_heads=8, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        h = self.norm(x)
        x = x + self.attn(h, h, h, need_weights=False)[0]
        return x + self.mlp(self.norm(x))


class TinyLM(nn.Module):
    def __init__(self, vocab=32000, d=1024, n_layers=8):
        super().__init__()
        self.embed = nn.Embedding(vocab, d)
        self.layers = nn.ModuleList(Block(d) for _ in range(n_layers))
        self.head = nn.Linear(d, vocab, bias=False)

    def forward(self, tokens):
        x = self.embed(tokens)
        for layer in self.layers:
            x = layer(x)
        return self.head(x)


class AppState(Stateful):
    """Lets DCP save/load model + optimizer as sharded (DTensor) state."""

    def __init__(self, model, optim):
        self.model, self.optim = model, optim

    def state_dict(self):
        msd, osd = get_state_dict(self.model, self.optim)
        return {"model": msd, "optim": osd}

    def load_state_dict(self, sd):
        set_state_dict(self.model, self.optim,
                       model_state_dict=sd["model"], optim_state_dict=sd["optim"])


def main():
    use_cuda = torch.cuda.is_available()
    local_rank = int(os.environ["LOCAL_RANK"])
    if use_cuda:
        torch.cuda.set_device(local_rank)
    device_type = "cuda" if use_cuda else "cpu"
    world = int(os.environ["WORLD_SIZE"])
    mesh = init_device_mesh(device_type, (world,), mesh_dim_names=("dp",))

    model = TinyLM()
    mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
    for layer in model.layers:          # shard bottom-up: one comm group per block
        fully_shard(layer, mesh=mesh, mp_policy=mp)
    fully_shard(model, mesh=mesh, mp_policy=mp)   # root: embed + head

    optim = torch.optim.AdamW(model.parameters(), lr=3e-4)  # build AFTER fully_shard

    step0 = 0
    if os.path.isdir(CKPT_DIR) and os.listdir(CKPT_DIR):
        latest = max(os.listdir(CKPT_DIR), key=lambda s: int(s.split("_")[-1]))
        dcp.load({"app": AppState(model, optim)},
                 checkpoint_id=os.path.join(CKPT_DIR, latest))
        step0 = int(latest.split("_")[-1]) + 1

    ckpt_pg = dist.new_group(backend="gloo")  # separate group for background checkpoint I/O
    pending = None
    steps = int(os.environ.get("STEPS", "20"))
    for step in range(step0, steps):
        tokens = torch.randint(0, 32000, (4, 512), device=device_type)  # stand-in for a real loader
        logits = model(tokens[:, :-1])
        loss = nn.functional.cross_entropy(
            logits.flatten(0, 1).float(), tokens[:, 1:].flatten())
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optim.step()
        optim.zero_grad(set_to_none=True)

        if step % 10 == 0 or step == steps - 1:
            if pending is not None:
                pending.result()        # keep at most one async save in flight
            pending = dcp.async_save({"app": AppState(model, optim)},
                                     checkpoint_id=os.path.join(CKPT_DIR, f"step_{step}"),
                                     process_group=ckpt_pg)
            if dist.get_rank() == 0:
                print(f"step {step} loss {loss.item():.4f}", flush=True)

    if pending is not None:
        pending.result()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

The skeleton leaves out several things a production job needs. It picks the latest directory by name without checking that the write completed. A real job should write a marker file after pending.result() returns and resume only from marked checkpoints. It also has no learning-rate schedule, real data loader or loader-state checkpointing. For those, start from TorchTitan rather than growing this file.

Slurm Launch Script

Use one Slurm task per node. That task runs torchrun, which spawns one worker per GPU and sets RANK, LOCAL_RANK and WORLD_SIZE. Rendezvous uses the c10d backend on the first node of the allocation.[17] Values marked as examples are site-specific.

#!/bin/bash
#SBATCH --job-name=fsdp2-train
#SBATCH --nodes=4                    # example: 4 nodes
#SBATCH --ntasks-per-node=1          # one torchrun launcher per node
#SBATCH --gpus-per-node=8            # example: 8-GPU nodes
#SBATCH --cpus-per-task=64           # example: size to your nodes
#SBATCH --exclusive
#SBATCH --time=24:00:00
#SBATCH --output=%x-%j.out

set -euo pipefail

# Environment (example: adjust to your site's modules or container)
source /path/to/venv/bin/activate

export OMP_NUM_THREADS=8
export CKPT_DIR=/scratch/$USER/ckpt/${SLURM_JOB_NAME}   # stable across requeues
export NCCL_DEBUG=WARN
# Site-specific NCCL settings (fabric interfaces, HCAs) belong here;
# see the Interconnects chapter and your center's documentation.

HEAD_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
GPUS_PER_NODE=8

srun --kill-on-bad-exit=1 torchrun \
    --nnodes="$SLURM_NNODES" \
    --nproc-per-node="$GPUS_PER_NODE" \
    --rdzv-id="$SLURM_JOB_ID" \
    --rdzv-backend=c10d \
    --rdzv-endpoint="$HEAD_NODE:29500" \
    --max-restarts=3 \
    train_fsdp2.py

Notes on the script. --kill-on-bad-exit=1 makes Slurm tear down the whole step when any task fails, instead of leaving the other nodes blocked in a collective until the NCCL watchdog fires.[50] --max-restarts lets torchrun restart its workers after a crash, and the script resumes from the latest DCP checkpoint.[17] NCCL_DEBUG=WARN is a reasonable production setting. Leave the settings listed under "Debugging" in the NCCL docs out of production scripts.[44] For jobs longer than your wall-clock limit, sbatch --signal can warn the job a set number of seconds before its end time, so it can checkpoint before it is killed.[41] Keep CKPT_DIR keyed to something stable, such as the job name, so a follow-on job finds it. Queue policy for long training runs is covered in Scheduling, and building the containers in Systems Software.

Key Takeaways

• Do the memory math first. Mixed-precision Adam costs 16 bytes per parameter before activations, so a 7B full fine-tune needs sharding across several GPUs and a 70B model needs more than one 8-GPU node of 80 GB parts.
• Use the 6ND rule for every allocation request and MFU (6N x tokens/s over peak) to judge efficiency. Published frontier runs sit around 38-43% MFU.
• Start simple. Use DDP if the model fits, then FSDP2/ZeRO, then add TP inside the NVLink domain, PP across nodes, CP for long sequences and EP for MoE.
• Match communication frequency to link speed. TP and EP stay on NVLink, PP crosses nodes, and DP/FSDP is outermost.
• BF16 is the baseline. FP8 through Transformer Engine or TorchTitan is production-proven but should be validated against a BF16 loss curve.
• Failures are routine at scale. Llama 3 saw 419 unexpected interruptions in 54 days, most of them hardware. Use sharded async checkpoints, automated restarts and node health checks.
• Most research users need fine-tuning, not pretraining. LoRA and QLoRA cut memory by large factors, so steer users toward them before granting large full-parameter allocations.

Sources

1. Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (arXiv / SC20, 2019-2020)
2. Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (arXiv, 2019)
3. Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (arXiv / SC21, 2021)
4. Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models (arXiv, 2022)
5. Kaplan et al., Scaling Laws for Neural Language Models (arXiv, 2020)
6. Hoffmann et al., Training Compute-Optimal Large Language Models (arXiv, 2022)
7. Llama Team, Meta, The Llama 3 Herd of Models (arXiv, 2024)
8. Liang, Liu et al., TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training (arXiv / ICLR 2025)
9. DeepSeek-AI, DeepSeek-V3 Technical Report (arXiv, 2024-2025)
10. Qi et al., Zero Bubble Pipeline Parallelism (arXiv, 2023)
11. Hu et al., LoRA: Low-Rank Adaptation of Large Language Models (arXiv, 2021)
12. Dettmers et al., QLoRA: Efficient Finetuning of Quantized LLMs (arXiv, 2023)
13. torch.distributed.fsdp.fully_shard (PyTorch 2.14 documentation, 2026)
14. Getting Started with Fully Sharded Data Parallel (FSDP2) (PyTorch Tutorials, 2025)
15. Distributed Checkpoint - torch.distributed.checkpoint (PyTorch 2.14 documentation, 2026)
16. Asynchronous Saving with Distributed Checkpoint (DCP) (PyTorch Recipes, 2025)
17. torchrun (Elastic Launch) (PyTorch 2.14 documentation, 2026)
18. Distributed Data Parallel design notes (PyTorch 2.14 documentation, 2026)
19. Megatron Core Parallelism Strategies Guide (NVIDIA, 2026)
20. NVIDIA Megatron-LM and Megatron Core repository README (NVIDIA / GitHub, 2026)
21. Zero Redundancy Optimizer tutorial (DeepSpeed, updated 2026)
22. DeepSpeed Configuration JSON (DeepSpeed, 2026)
23. Transformer Engine documentation (NVIDIA, 2026)
24. Using FP8 and FP4 with Transformer Engine, TE 2.18.0 (NVIDIA, 2026)
25. Introduction to multi-controller JAX (JAX documentation, 2026)
26. Manual parallelism with shard_map (JAX documentation, 2026)
27. pytorch/torchtitan repository README (PyTorch / GitHub, 2026)
28. Accelerate documentation (Hugging Face, 2026)
29. FSDP1 vs FSDP2 (Hugging Face Accelerate, 2026)
30. PEFT documentation (Hugging Face, 2026)
31. torchft: Easy Per Step Fault Tolerance for PyTorch (Meta PyTorch / GitHub, 2026)
32. NVIDIA Resiliency Extension (NVIDIA / GitHub, 2026)
33. NeMo Megatron Bridge documentation (NVIDIA, 2026)
34. NVIDIA NeMo Framework Overview (NVIDIA, 2026)
35. PyTorch 2.14.0 Release Notes (PyTorch / GitHub, 2026)
36. Liu, Zaharia, Abbeel, Ring Attention with Blockwise Transformers for Near-Infinite Context (arXiv, 2023)
37. Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv, 2022)
38. Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (arXiv, 2018)
39. Pipeline Parallelism - torch.distributed.pipelining (PyTorch 2.14 documentation, 2026)
40. Tensor Parallelism - torch.distributed.tensor.parallel (PyTorch 2.14 documentation, 2026)
41. sbatch manual page (SchedMD, 2026)
42. NVIDIA GB200 NVL72 product page (NVIDIA, 2026)
43. Megatron Core Mixture of Experts package (NVIDIA, 2026)
44. NCCL Environment Variables (NVIDIA, 2026)
45. torch.utils.checkpoint (PyTorch 2.14 documentation, 2026)
46. Distributed communication package - torch.distributed, incl. init_device_mesh (PyTorch 2.14 documentation, 2026)
47. torch.compile (PyTorch 2.14 documentation, 2026)
48. Distributed data loading (JAX documentation, 2026)
49. deepspeedai/DeepSpeed repository README (DeepSpeed / GitHub, 2026)
50. srun manual page (SchedMD, 2026)