Scaling Distributed Training: Foundations and First Principles

1. High-Level Overview: The Three Fundamental Challenges

Every technique in large-scale model training addresses one or more of three orthogonal resource constraints. These constraints form the scaling trilemma of distributed deep learning.

1.1 Memory Usage — The Hard Constraint

Memory is a binary gate: if the aggregate memory required by a single training step exceeds the available GPU high-bandwidth memory (HBM), training cannot proceed at all — there is no graceful degradation, only an Out-Of-Memory (OOM) crash.

During a single training step, four principal memory occupants coexist on the GPU:

Additionally, there are minor but non-negligible constant overheads:

• CUDA context/kernels: typically $1$–$2$ GB upon first CUDA tensor allocation.
• Fragmentation & temporary buffers: memory that exists but cannot be utilized due to allocator fragmentation or short-lived intermediate results.

1.2 Compute Efficiency — Maximizing Arithmetic Throughput

Modern accelerators (e.g., NVIDIA H100 SXM5) deliver peak throughput of approximately $989$ TFLOPS in BF16 Tensor Core operations. However, achieved throughput is invariably lower due to:

• Memory-bound operations: kernels whose execution time is dominated by data movement (loads/stores) rather than arithmetic.
• Kernel launch overhead: CPU-side latency to dispatch GPU kernels.
• Pipeline bubbles: idle time when certain stages of computation must wait for others.

The goal is to maximize the compute utilization ratio:

Every second the GPU spends waiting — for data transfers from CPU to GPU, for memory allocation, or for synchronization barriers — directly reduces this ratio.

1.3 Communication Overhead — Minimizing GPU Idle Time

In multi-GPU and multi-node settings, GPUs must exchange data (gradients, activations, parameters). Communication occurs over interconnects with heterogeneous bandwidths:

Communication overhead keeps GPUs idle (not performing useful arithmetic). Two primary strategies mitigate this:

1. Overlap communication with computation: launch asynchronous collective operations (e.g., AllReduce) concurrently with backward-pass gradient computation on independent layers.
2. Topology-aware placement: assign communication-heavy operations to the fast intra-node links and minimize the volume of data traversing slow inter-node links.

1.4 The Scaling Trilemma: Trading Off Resources

These three resources — memory, compute, and communication — are not independent. Optimizing one often comes at the cost of another. Two canonical examples:

Formally, if we denote total training step wall-clock time as $T$, we can decompose it as:

The objective of efficient scaling is:

where $M_{\text{peak}}$is the peak memory during a training step and$M_{\text{GPU}}$ is the physical HBM capacity.

2. First Steps: Training on One GPU

2.1 The Three Phases of a Training Step

A single training step on one GPU consists of three sequential phases:

Phase 1: Forward Pass

The input batch $X \in \mathbb{R}^{bs \times seq \times d_{\text{input}}}$is propagated through$L$successive layers of the model. For a generic layer$\ell$with parameters$W_\ell$, the forward computation is:

where $a^{(0)} = X$is the input embedding and$a^{(\ell)}$is the activation (hidden state) output of layer$\ell$. The final output $\hat{y} = a^{(L)}$ is used to compute the scalar loss:

Critical point: Every intermediate activation $a^{(\ell)}$ must be stored in GPU memory because it is needed for the backward pass (see Phase 2).

Phase 2: Backward Pass (Backpropagation)

Using the chain rule, gradients of the loss with respect to each layer’s parameters are computed in reverse order ($\ell = L, L-1, \ldots, 1$):

The upstream gradient $\frac{\partial \mathcal{L}}{\partial a^{(\ell)}}$ is propagated backward through:

As each layer $\ell$completes its gradient computation, the stored activation$a^{(\ell)}$is freed from memory, while the gradient tensor$\nabla_{W_\ell}\mathcal{L}$ is allocated.

Phase 3: Optimizer Step

The optimizer uses the accumulated gradients to update all parameters. For the Adam optimizer, the update rule for each parameter tensor $\theta$at step$t$ is:

where $g_t = \nabla_\theta \mathcal{L}$, $m_t$is the first moment (momentum),$v_t$is the second moment (variance),$\eta$is the learning rate, and$\epsilon$is a numerical stability constant (typically$10^{-8}$).

After the optimizer step, the gradient buffers are zeroed and the cycle repeats.

2.2 Batch Size: Definitions, Impact, and Practical Ranges

Definition

The batch size ($bs$) is the number of independent input samples processed in a single forward–backward pass before the optimizer updates the parameters.

In the LLM pretraining community, batch sizes are reported in tokens ($bst$) to decouple from the choice of sequence length ($seq$):

Impact on Convergence

OpenAI’s seminal study on large batch training [1] demonstrated the existence of a critical batch size $B_{\text{crit}}$ for each model and dataset, below which doubling the batch size nearly halves the number of required optimization steps, and above which diminishing returns set in rapidly.

Key practical insight: The sensitivity of final model performance to the exact batch size is low around the optimum — i.e., batch size can be varied within a broad range near $B_{\text{crit}}$ without significant degradation.

Real-World Examples

The sweet spot for contemporary LLM pretraining is typically:

3. Memory Usage in Transformers: Detailed Breakdown

3.1 The Four Principal Memory Occupants

Items 1–3 are static with respect to batch size and sequence length — they depend only on the model architecture (parameter count $N$). Item 4 is dynamic and is the dominant memory consumer for large batch sizes and long sequences.

3.2 Numerical Precision Formats

The memory footprint of any tensor is:

4. Memory for Weights, Gradients, and Optimizer States

4.1 Parameter Count of a Transformer LLM

For a decoder-only transformer with:

• $h$: hidden dimension
• $v$: vocabulary size
• $L$: number of layers
• No fixed positional embeddings (e.g., using RoPE)

The total parameter count is:

Breakdown of the per-layer term $12h^2 + 13h$:

The term $h \cdot v$accounts for the token embedding matrix, and the final$2h$ accounts for the final LayerNorm (gain and bias).

Scaling observation: As $h$grows, the dominant term is$12Lh^2$, which grows quadratically in the hidden dimension. The linear terms ($13Lh$, $hv$, $2h$) become negligible for very large models.

4.2 Full Precision (FP32) Training Memory

In pure FP32 training, every value — parameters, gradients, optimizer states — is stored in 32-bit floating point ($4$ bytes):

Total (FP32):

4.3 Mixed Precision (BF16 + FP32 Master Weights) Training Memory

In standard mixed precision training:

• Forward/backward computations are done in BF16 ($2$ bytes per value).
• A master copy of weights is maintained in FP32 for numerical stability during the optimizer update.
• Optimizer states (Adam $m_t$, $v_t$) are stored in FP32.

Total (mixed precision, BF16 gradients):

Total (mixed precision, FP32 gradient accumulation):

If gradients are accumulated in FP32 (for stability with small gradient values in BF16), an additional $4N$ bytes is required:

Key insight: Mixed precision training does not reduce the total memory for weights + gradients + optimizer states compared to full FP32 training (both are $16N$ bytes without FP32 gradient accumulation). The advantages of mixed precision are:

1. Faster computation: BF16 Tensor Core operations are $2\times$ faster than FP32.
2. Reduced activation memory: activations stored in BF16 use half the memory of FP32.

4.4 Practical Memory Requirements Table

For reference, an NVIDIA H100 SXM5 has 80 GB HBM3. At $7$B parameters, the weights + gradients + optimizer states alone ($112$–$140$ GB) already exceed a single GPU’s capacity — before accounting for activations.

5. Memory for Activations

5.1 Why Activations Must Be Stored

During the backward pass, computing $\frac{\partial \mathcal{L}}{\partial W_\ell}$requires the input activation$a^{(\ell-1)}$to layer$\ell$:

For a linear layer $a^{(\ell)} = W_\ell \, a^{(\ell-1)}$, this becomes:

Without the stored activation $a^{(\ell-1)}$, this gradient cannot be computed. Hence, all intermediate activations must be retained in memory from the forward pass until they are consumed in the backward pass.

5.2 Activation Memory Formula

For a transformer model in mixed precision (BF16 activations), the total activation memory is:

where:

Derivation Sketch (per layer, per sample)

Within each transformer layer, the following intermediate tensors must be stored for backpropagation:

Summing all per-layer contributions for a single sample and then scaling by $L$layers and$bs$ samples yields the formula above. The full accounting is detailed in the NVIDIA recomputation paper [4].

Critical scaling behavior:

• Linear in $bs$ (batch size)
• Linear in $L$ (depth)
• Quadratic in $seq$(due to the$n_{\text{heads}} \cdot seq^2 / h$ attention score term)
• The attention score memory $\propto seq^2$ dominates for long sequences

This means:

For short sequences, the $34 \cdot seq \cdot h$term dominates and memory is approximately linear in$seq$. For long sequences, the quadratic term takes over and activation memory explodes.

6. Activation Recomputation (Gradient Checkpointing / Rematerialization)

6.1 Core Idea

Activation recomputation trades compute for memory: instead of storing all intermediate activations during the forward pass, we discard most of them and recompute them on-the-fly during the backward pass from a small set of saved checkpoints.

Without recomputation, memory for activations is $\mathcal{O}(L)$(all$L$layers’ activations stored). With full recomputation, if we checkpoint only every$k$-th layer, activation memory becomes $\mathcal{O}(L/k)$, but we pay an additional forward-pass compute cost to recompute the discarded activations.

6.2 Strategies

Strategy 1: Full Recomputation

• What is stored: only the activation at the boundary of each transformer layer (i.e., $a^{(\ell)}$for$\ell = 0, 1, \ldots, L$, but none of the intermediate tensors within each layer).
• Recomputation cost: during the backward pass, for each layer $\ell$, the entire forward pass through layer $\ell$ must be re-executed to reconstruct internal activations (attention scores, FFN intermediates, etc.).
• Net effect: approximately one additional full forward pass is performed during each backward pass.
• Compute overhead: typically $+30\%$to$+40\%$ wall-clock time increase per training step.
• Memory savings: maximal — all intra-layer activations are freed.

Strategy 2: Selective Recomputation

The key observation from the NVIDIA paper [4] is that not all activations are equally costly to store or cheap to recompute:

Selective strategy: discard and recompute only the attention-related activations (which are large but cheap to recompute), while retaining the FFN activations (which are expensive to recompute).

For GPT-3 (175B):

• Activation memory reduction: $\sim 70\%$
• Compute overhead: only $\sim 2.7\%$

This is a dramatically better trade-off than full recomputation.

Example — DeepSeek-V3: uses Multi-Head Latent Attention (MLA), which compresses the $Q$, $K$, $V$ representations into a low-rank latent space. This reduces the attention activation memory even further, making selective checkpointing even more effective:

where $d_c$ is the compressed latent dimension, drastically reducing the stored activation size for key-value pairs.

6.3 FlashAttention and Native Recomputation

FlashAttention (Dao et al.) is a hardware-aware exact attention algorithm that:

1. Computes attention in tiled blocks that fit in GPU SRAM (on-chip memory), avoiding materialization of the full $seq \times seq$ attention matrix in HBM.
2. Natively integrates selective recomputation: during the backward pass, attention scores are recomputed from $Q$, $K$, $V$ rather than loaded from HBM.

Since FlashAttention is now the default in most training frameworks, practitioners using it are already benefiting from selective recomputation without explicit gradient checkpointing configuration for the attention layers.

6.4 Hardware FLOPS Utilization (HFU) vs. Model FLOPS Utilization (MFU)

Recomputation adds “extra” floating-point operations that are not part of the theoretical minimum computation for a forward + backward pass. This creates an important distinction in efficiency metrics:

Hardware FLOPS Utilization (HFU):

This measures how effectively the hardware arithmetic units are utilized, including recomputation. A high HFU means the GPU is kept busy, but some of that work is “redundant.”

Model FLOPS Utilization (MFU):

This measures how much useful (non-redundant) computation the hardware performs per unit time. MFU is the better metric for comparing different hardware or different training configurations, because it rewards setups that can avoid recomputation (e.g., by having more memory available).

For an ideal training setup with no recomputation: $\text{HFU} = \text{MFU}$.

For a setup with full recomputation: $\text{HFU} > \text{MFU}$ (because the denominator is the same but HFU’s numerator includes the extra forward pass).

7. Gradient Accumulation

7.1 Core Mechanism

Gradient accumulation decouples the global batch size (which determines the gradient quality and convergence behavior) from the micro-batch size (which determines the activation memory per forward–backward pass).

The procedure for a single optimizer step with $G$ gradient accumulation steps:

7.2 Batch Size Relationship

where:

In tokens:

7.3 Memory Trade-Off

Activation memory reduction factor: $\frac{gbs}{mbs} = G$

Drawback: Gradient accumulation requires $G$sequential forward–backward passes per optimizer step. The wall-clock time per optimizer step increases approximately linearly with$G$:

Subtle memory note: Gradient accumulation requires a persistent gradient buffer of size $m_{\text{grad}}$that persists across all$G$ micro-batch passes. Without gradient accumulation, gradients can be computed and immediately consumed during the backward pass, enabling slightly lower peak memory through operator fusion. With accumulation, gradients from previous micro-batches must persist, creating a small memory overhead.

7.4 The Parallelism Opportunity

The $G$micro-batch forward–backward passes are independent computations (different input samples, no inter-dependencies). This independence is precisely what enables Data Parallelism: distribute the$G$micro-batches across$G$ GPUs, perform forward–backward passes simultaneously, and synchronize gradients via an AllReduce collective before the optimizer step. This is covered in the next section of the book.

8. Profiling GPU Compute and Communication

8.1 The PyTorch Profiler

PyTorch provides a built-in profiler that traces CPU and GPU activity at the kernel level. The API:

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/profile'),
    with_stack=True
) as prof:
    for step in range(steps):
        train_step()
        prof.step()

Parameters explained:

8.2 Anatomy of a Profiler Trace

A profiler trace visualized in TensorBoard or Chrome’s chrome://tracing reveals multiple concurrent timelines:

Key patterns to identify:

1. Sequential compute and communication: If a gradient AllReduce operation appears after the backward pass completes (rather than overlapping with it), there is unnecessary GPU idle time. The fix is to launch AllReduce asynchronously as soon as each layer’s gradients are ready.

2. Idle GPU time (gaps between kernels): May indicate:

   • CPU bottleneck (kernel launch overhead, data preprocessing)
   • Synchronization barriers (torch.cuda.synchronize(), dist.barrier())
   • Memory allocation stalls (fragmentation forcing the CUDA caching allocator to search or defragment)

3. CUDA Syncs and CPU↔GPU data transfers: cudaMemcpy (host-to-device or device-to-host) operations appear as blocking periods. These should be minimized or overlapped with computation using pinned memory and non-blocking transfers.

4. First step anomaly: The first training step exhibits a markedly different profile due to:

   • CUDA context initialization: loading CUDA kernels, JIT-compiling fused operations.
   • PyTorch caching allocator warm-up: the allocator performs trial allocations to build a free-block cache (see Zach DeVito’s blog on the PyTorch CUDA caching allocator). Subsequent steps reuse these cached blocks, making allocation nearly free.
   • Optimizer state initialization: Adam’s $m_t$and$v_t$ tensors are allocated after the first backward pass and persist for all subsequent steps.

Practical consequence: A model that fits in memory during step 1 may OOM at step 2, because the optimizer states ($8N$ bytes for Adam in FP32) are allocated only after the first backward pass, permanently increasing the memory baseline.

9. Summary: The Landscape Before Multi-GPU Scaling

These single-GPU foundations — memory anatomy, activation recomputation, and gradient accumulation — form the building blocks upon which all multi-GPU parallelism strategies (Data Parallelism, Tensor Parallelism, Pipeline Parallelism, ZeRO, etc.) are constructed.

Corrected Production-Accurate Memory Analysis for Transformer Training

0. The Fundamental Error in the Naïve “$16N$ Bytes” Estimate

The statement “$16N$ bytes for FP32 training” commits a critical analytical error: it treats all four memory occupants — parameters, gradients, optimizer states, and activations — as if they coexist simultaneously at their full sizes throughout the entire training step. In reality:

1. Gradients do not exist before the backward pass begins. They are allocated layer-by-layer during backpropagation.
2. Activations do not vanish the instant the backward pass starts. They are freed progressively as each layer’s backward computation consumes them.
3. The true peak memory occurs at a specific transient moment — the start of the backward pass — when gradients begin accumulating while nearly all activations are still resident.

The corrected analysis requires tracing memory occupancy as a function of the training phase, not collapsing it to a single static number.

1. Notation and Assumptions

1.1 Symbols

1.2 Precision Assumption (FP32 Baseline)

All tensors — parameters, gradients, optimizer states — stored in FP32 ($4$bytes per scalar). Adam optimizer stores two auxiliary states per parameter: first moment$m_t$and second moment$v_t$.

1.3 Generalized Parameter Count

The formula $N = hv + L(12h^2 + 13h) + 2h$is architecture-specific (standard GPT-style decoder with$4h$ FFN intermediate dimension, bias terms in all projections, two LayerNorms per layer, and no tied embeddings). A generalized parameter count for an arbitrary transformer should account for:

where each layer $\ell$ contributes:

with $d_{\text{head}} = h / n_{\text{heads}}$and$d_{\text{ff}}$being the FFN intermediate dimension (often$4h$, but $\frac{8h}{3}$in SwiGLU-based models like Llama). Whether biases exist, whether embeddings are tied to the output head, whether GQA (grouped-query attention) is used — all change$N$. The specific formula must be derived per architecture; no single formula is universal.

For the remainder of this analysis, we treat $N$ as a known constant for a given model and focus on the memory dynamics during training, which are architecture-independent in structure.

2. Fixed Memory: Parameters and Optimizer States

These components depend only on $N$ and are independent of batch size, sequence length, and training phase (once initialized):

Critical observation: The gradient tensor’s $4N$bytes is not always resident. Before the backward pass, no gradient memory is allocated (or it is zero-filled from the previous step’s clearing). It builds up during the backward pass. This temporal behavior is what makes the simple “$16N$” summary misleading.

3. Phase-by-Phase Memory Analysis

3.1 Phase 0: Before Forward Pass (Steady-State Baseline)

At this point in the training loop, the previous step’s optimizer update has completed, gradients have been zeroed, and activations from the previous step have been fully freed.

Note on Step 0 vs. Step 1+: At the very first step, the optimizer states $m_0$and$v_0$do not yet exist (they are initialized to zero tensors only after the first backward pass). Therefore, memory before the first forward pass is only$4N$(parameters alone). After the first optimizer step, the persistent baseline jumps to$12N$. This explains why a model that fits in GPU memory at step 0 can OOM at step 1.

3.2 Phase 1: During the Forward Pass

As the forward pass proceeds through layers $\ell = 1, 2, \ldots, L$, each layer produces intermediate activations $a^{(\ell)}$ that must be retained for the backward pass. The activation memory monotonically increases as more layers are processed.

After completing the forward pass through layer $k$(i.e., layers$1$through$k$ are done):

3.3 Phase 2: End of Forward Pass (All Activations Stored)

When the forward pass is complete ($k = L$), all layer activations are simultaneously resident:

At this point, the loss $\mathcal{L}$ is computed, and the backward pass is about to begin. No gradients have been allocated yet.

For the activation memory formula in mixed precision (BF16):

For FP32, all activation bytes double (replace the constant $34$and coefficient$5$ with their FP32 equivalents derived from the same accounting with 4-byte storage per element instead of 2-byte).

3.4 Phase 3: During the Backward Pass — The True Peak

The backward pass proceeds in reverse layer order: $\ell = L, L-1, \ldots, 1$. At each backward step $k$(processing layer$k$ in reverse):

1. The activation $a^{(k-1)}$(input to layer$k$) is read to compute $\frac{\partial \mathcal{L}}{\partial W_k}$.
2. The gradient $\nabla_{W_k} \mathcal{L}$is written to the gradient buffer ($4N_k$bytes, where$N_k$is the parameter count of layer$k$).
3. After layer $k$’s backward computation completes, activation $a^{(k)}$ is freed.

Define the remaining activation memory when the backward pass has just reached layer $k$(i.e., layers$L, L-1, \ldots, k+1$have completed their backward, and layer$k$ is about to begin):

And the accumulated gradient memory at this point (gradients for layers $L, L-1, \ldots, k+1$have been computed, and layer$k$’s gradient is about to be computed):

(Approximate, assuming equal parameter distribution across layers.)

The total memory at backward step $k$:

3.5 Identifying the True Peak: Start of Backward ($k = L$)

The peak memory occurs at the very beginning of the backward pass, when:

• All activations are still resident: $A_{\text{rem}}(1) \approx A_{\text{tot}}$
• Gradients are beginning to be allocated: $G_{\text{acc}}$ is initially small but the gradient buffer for the full model is typically pre-allocated

In most frameworks (PyTorch, JAX), the full gradient buffer of size $4N$ is allocated at once when .backward() is called, not incrementally. Therefore, at the first instant of the backward pass:

This is the worst-case peak memory — the absolute maximum GPU memory required during training.

This is strictly greater than $16N$:

3.6 Memory Decrease During Backward Pass

As the backward pass progresses from layer $L$down to layer$1$:

• $A_{\text{rem}}(k)$ decreases (activations freed after each layer’s backward)
• $G_{\text{acc}}(k)$ increases (gradients accumulate) — but since the gradient buffer is pre-allocated, this doesn’t actually change total memory

Net effect: memory decreases monotonically during the backward pass, from the peak of $16N + A_{\text{tot}}$toward$16N$ (all activations freed, all gradients computed).

3.7 Phase 4: Optimizer Step

After the backward pass completes:

The optimizer reads $\nabla_W \mathcal{L}$, updates $m_t$, $v_t$, and $W$ in-place. After the update, gradients are zeroed (or freed and re-allocated next step).

3.8 Phase 5: After Optimizer Step (Return to Baseline)

Parameters ($4N$) + Adam states ($8N$) persist. Gradient buffers are zeroed/freed. The cycle repeats.

4. Complete Phase-Accurate Memory Table (FP32 + Adam)

5. The System-Level Invariant

In words: The peak training memory occurs at the start of the backward pass and equals the full parameter + gradient + optimizer state memory ($16N$) plus all unreleased activations ($A_{\text{tot}}$).

This is an exact invariant (modulo constant CUDA context overhead and allocator fragmentation) that holds for:

• Any transformer architecture (encoder, decoder, encoder-decoder)
• Any depth $L$, hidden dimension $h$, batch size $bs$, sequence length $seq$
• FP32 training with the Adam optimizer

For mixed precision (BF16 + FP32 master weights, BF16 gradients):

For mixed precision with FP32 gradient accumulation:

6. Corrected Practical Memory Requirements

6.1 Activation Memory Estimates

Using the activation formula for mixed precision:

Example for Llama-2 7B ($L=32$, $h=4096$, $n_{\text{heads}}=32$, $bs=1$, $seq=4096$):

6.2 Corrected Total Peak Memory Table (Mixed Precision, BF16 Gradients, $bs=1$, $seq=4096$)

The naïve “$16N$” table dramatically underestimates the true memory requirement by ignoring $A_{\text{tot}}$, which can be comparable to or larger than the parameter-related memory, especially for long sequences and large batch sizes.

7. Why This Correction Matters in Practice

7.1 OOM Debugging

When a training run crashes with OOM, practitioners often check only whether “$16N$ fits in GPU memory.” This is insufficient. The true check must be:

where $M_{\text{CUDA context}} \approx 1\text{–}2$ GB.

7.2 Batch Size / Sequence Length Selection

Since $A_{\text{tot}} \propto bs \cdot seq$(linearly) and$A_{\text{tot}} \propto seq^2$(quadratically, through the attention term), the maximum feasible$bs$and$seq$are not determined by$16N$alone but by the residual memory after subtracting the persistent$12N$ baseline:

7.3 Impact on Activation Recomputation Decisions

With the corrected peak $16N + A_{\text{tot}}$, the value of activation recomputation becomes quantifiable. Selective recomputation reduces $A_{\text{tot}}$by$\sim 70\%$:

Full recomputation reduces $A_{\text{tot}}$ to near zero (only layer-boundary checkpoints remain):

The memory saved is:

This saving can convert an infeasible training configuration into a feasible one.

8. Summary: The Corrected System-Level Memory Invariant

The steady-state baseline (between steps) is:

The peak-to-baseline ratio:

For large batch sizes or long sequences where $A_{\text{tot}} \gg N$, the peak can be many times the baseline, making the naïve $16N$ estimate not just slightly wrong but qualitatively misleading about whether a training configuration will fit in memory.