Distributed Training

112 GB for model states alone — that exceeds a single A100's 80 GB even before accounting for activations. Without ZeRO, you can't even load the optimizer state on one GPU.

Stage 1 alone cuts memory from 112 GB to 29.3 GB. The 4 bytes/param for full weights + gradients is the floor — you can't avoid keeping the local copy of parameters and the full gradient for the backward pass.

Stage 2 halves the Stage 1 cost by also sharding gradients. Now 15.5 GB fits comfortably on an 80 GB A100, leaving ~65 GB for activations. This is the default in most FSDP deployments.

1.75 GB! But there's a catch: Stage 3 must gather parameters on-the-fly before each forward/backward computation, then discard them. This adds AllGather communication at every layer. The trade-off is clear: memory scales perfectly (1/N) but communication overhead increases.

At 100 GB/s, that's ~420 ms of pure communication per step. Compare to ZeRO-1/2 which only does one AllReduce of gradients (~28 GB). ZeRO-3 trades 1.5x more communication for 8-17x less memory. Worth it when the model literally can't fit otherwise.

ZeRO-3 with 8 GPUs gives 16/8 = 2 bytes/param × 30B = 60 GB for model states. Add 20 GB for activations = 80 GB — exactly the A100 limit. In practice you'd want some headroom, so you might add activation checkpointing to reduce the 20 GB, or use 16 GPUs with ZeRO-2.

64 micro-batches per optimizer step. That means 64 forward+backward passes between each weight update. The AllReduce only happens once at the end, so communication is 1/64th of what it would be without accumulation.

3.15 seconds saved per optimizer step! Note however that without accumulation you'd also get 64 optimizer updates instead of 1 — the training dynamics differ. Gradient accumulation gives you a specific batch size, not just speed. Large batch training converges differently from small batch training.

Communication is only 50/7731 = 0.6% of the step time. Gradient accumulation makes the training almost perfectly compute-bound. Without it, communication would be 3200/(64×120 + 3200) = 29% of total time.

javascript
function effectiveBatch(microBatch, accumSteps, dpDegree) {
  return microBatch * accumSteps * dpDegree;
}
function requiredAccumSteps(targetBatch, microBatch, dpDegree) {
  return Math.ceil(targetBatch / (microBatch * dpDegree));
}

With micro=8: 2048/(8×8) = 32 accumulation steps. If each step takes T ms, total = 32T.

With micro=16 + checkpoint: 2048/(16×8) = 16 steps. Each step processes 2x tokens but takes ~2×1.33 = 2.66x time (doubled data, 33% recompute overhead). Total = 16 × 2.66T = 42.6T.

32T < 42.6T. Larger micro-batch with checkpointing is ~33% slower. Avoid activation checkpointing when you have enough accumulation steps to hit your batch size. Only use it when you literally can't fit the micro-batch you need.

An 8x increase in batch size gives an 8x increase in LR. This is the rule used in the original ResNet large-batch paper (Goyal et al.). It works well up to batch sizes around 8K-32K for vision and ~1M-4M tokens for LLMs.

Square root scaling is more conservative — 2.83x instead of 8x. It's derived from the noise-signal ratio analysis and tends to work better at very large batch sizes (>32K). Most modern LLM training runs use something between linear and sqrt.

8000/100000 = 8% of training, which falls in the 5-10% sweet spot. Warmup is critical at large batch sizes because the early gradient estimates are very noisy (the model is randomly initialized), and a large LR amplifies that noise. Ramping up slowly lets the loss landscape stabilize before taking big steps.

Below B_crit, doubling the batch nearly doubles training speed (fewer noisy steps needed). Above B_crit, the gradient is already clean, so more samples per batch don't help — you just burn more compute per step without reaching the minimum faster. For GPT-3 scale models, B_crit ≈ 1-2M tokens. This is why Llama 3 uses batch rampup: start at 4M tokens/batch, scale to 16M as the loss landscape becomes more predictable.

The cosine schedule should decay over the actual 50K steps. The key quantity is per-token LR exposure: at step t, each token in the batch sees LR(t). With 2x batch and 2x LR, each token gets twice the update magnitude. But you have half the steps, so the integrated LR exposure (area under the LR curve) stays roughly the same. The cosine must decay over the actual step count to hit the minimum LR at the right time.

~66 ms for a single AllReduce. For reference, a forward+backward pass of a 10B model on 64 GPUs with decent micro-batch takes 200-400 ms. So communication is 15-30% of the step — significant but overlappable.

97.5% overlap means only 1.5 ms of the 60 ms AllReduce is visible as overhead. The effective communication cost is 1.5 ms on top of a 200 ms backward — less than 1% overhead. This is why bucketed, overlapped AllReduce is so effective: even though total data transfer is huge, almost all of it hides behind the backward pass.

350 ms for the cross-node step alone — 5x slower than the full NVLink AllReduce. The intra-node steps add maybe 10 ms total. This is why modern clusters use fat InfiniBand fabrics (400 Gb/s = 50 GB/s per port, 8 ports per node = 400 GB/s aggregate) and why tensor parallelism is always kept within a node.

Communication exceeds compute! The effective MFU drops from 50% to 50% × 242/(242+350) = 20%. This is why you never do pure data parallelism across slow inter-node links for large models. Instead, combine TP within the node (high-BW NVLink) + DP across nodes, or use gradient accumulation to amortize the communication cost.

javascript
function commComputeRatio(params_B, N, bw_gbs, flops_per_step_T, gpu_tflops) {
  const S = params_B * 2;  // FP16 grads in GB
  const commTime = 2 * S * (N - 1) / (N * bw_gbs);
  const computeTime = flops_per_step_T / gpu_tflops;
  return commTime / computeTime;
}

The bubble represents the warmup (first p-1 forward-only passes) and drain (last p-1 backward-only passes). During these phases, at least one stage is idle. With 12 microbatches, the steady-state (where all stages are busy) covers only 9/12 = 75% of the time.

Compare to GPipe which stashes all m=12 micro-batches: 12 × 2 = 24 GB. 1F1B uses 4x less activation memory because it starts backward passes as soon as possible, freeing stashed activations early. This is why 1F1B is universally preferred over GPipe in practice.

4x reduction in bubble — from 25% to 6.25%. The cost is 4x more point-to-point sends between stages (each virtual stage boundary requires a send/recv). With fast NVLink this is usually worth it. Megatron-LM's interleaved schedule uses v=number_of_layers / (PP × layers_per_chunk).

180 GB of point-to-point traffic per step. At 600 GB/s NVLink, the raw transfer time is 300 ms. But these transfers are pipelined and overlapped with computation, so the actual overhead is much smaller. The key constraint is that adjacent virtual stages must be on different GPUs connected by fast links.

Line 5 is the bug. The 1F1B bubble fraction formula is (p-1)/m for standard and (p-1)/(m×v) for interleaved — both have the total in the denominator alone, NOT (p-1) / ((p-1) + m×v). The GPipe formula uses (p-1) / (p-1+m), but the comment says "1F1B." Line 5 should be return bubbleSteps / totalSteps;

m=140 gives exactly 5%, so you need m ≥ 141 for strictly less than 5%. In practice with PP=8, you'd use interleaved stages (v=2 gives bubble < 5% at m=70, v=4 at m=35) rather than cramming 140+ micro-batches into a step.

Each GPU stores 8192 × 1024 = 8.39M parameters instead of 67.1M. For attention, this naturally maps to splitting heads: with 64 heads and TP=8, each GPU handles 8 heads.

56 MB per AllReduce per GPU. With 2 AllReduces per layer and 80 layers, total TP communication = 160 × 56 = 8960 MB ≈ 8.75 GB per GPU per step (forward only, double for backward). This is why TP must use NVLink — at 600 GB/s, each AllReduce takes 56 MB / 600 GB/s ≈ 0.09 ms. Over InfiniBand (100 GB/s), it would be 0.56 ms per AllReduce × 160 = 90 ms overhead.

TP=8 splits parameters perfectly — each GPU holds 1/8 of every layer's weights. Unlike ZeRO-3, there's no need to AllGather parameters before computation because each GPU already has its shard of the weight matrix and can compute its portion of the output directly. The trade-off is the 2 AllReduces per layer during both forward and backward.

With TP=8 and 64 Q heads: 64/8 = 8 Q heads per GPU. With 8 KV heads: 8/8 = 1 KV head per GPU. This divides cleanly because the number of KV heads (8) is divisible by TP (8). Each GPU has 8 Q heads that all share 1 KV head — the GQA ratio of 8:1 is preserved locally. If TP were 16, you'd need to replicate KV heads (16/8 = 2 copies each), or use sequence parallelism instead.

~10% overhead with NVLink is acceptable. But TP AllReduces are on the critical path — they cannot be overlapped because the next layer's computation depends on the AllReduce result. This is why TP is capped at 8 GPUs (one node) in most training configs. Going to TP=16 across nodes with 100 GB/s InfiniBand would make each AllReduce 6x slower, pushing overhead to ~58%.

98 seconds = 1.6 minutes per checkpoint. With distributed checkpointing (each GPU writes its own shard in parallel via ZeRO), the effective write speed can be N × 10 GB/s, reducing time to 98/N seconds. With 64 GPUs: 1.5 seconds!

Checkpoint every ~68 minutes. If each step takes 10 seconds, that's every ~410 steps. On a crash, you lose on average half the interval = 34 minutes of work. The checkpoint overhead is 98 seconds per 68.5 minutes = 2.4% — a small price for crash recovery.

~$600 per crash. With MTBF = 24 hours over a 30-day training run, that's ~30 crashes × $618 = ~$18,500 in wasted compute. This is why Meta and Google invest heavily in faster checkpointing, fault-tolerant training (elastic training that can continue with fewer GPUs), and hardware reliability.

javascript
function optimalCheckpointInterval(mtbf_hours, ckpt_seconds) {
  const C = ckpt_seconds / 3600;  // hours
  const interval = Math.sqrt(2 * C * mtbf_hours); // hours
  const interval_min = interval * 60;
  // Waste per crash: avg lost = interval/2, plus ckpt overhead per interval
  const crashes = 1;  // per MTBF period
  const lost = interval / 2;  // hours lost per crash
  const overhead = C * (mtbf_hours / interval); // total ckpt time in mtbf
  const waste = (lost + overhead) / mtbf_hours * 100;
  return [interval_min, waste];
}

Anything smaller than 6.1 × 10^-5 underflows to zero (or subnormal with reduced precision). Typical gradient magnitudes for a well-trained deep network are 10^-4 to 10^-7. So without loss scaling, the bottom end of the gradient distribution is silently lost.

S = 16384 = 2¹⁴. Each overflow halves S. Doubling requires W=2000 consecutive good steps, so it never triggers here. In practice, frequent overflows early in training cause S to drop quickly until it finds a stable range. A healthy training run sees S stabilize after a few hundred steps.

Surprise: mixed precision doesn't save optimizer memory! The savings come from activation memory (FP16 activations are 2x smaller) and communication (FP16 gradients are 2x smaller to AllReduce). The FP32 master copy + optimizer states still cost 12 bytes/param either way. This is why ZeRO is orthogonal to mixed precision — you need both.

Activations during forward pass tend to be well-behaved (bounded by activation functions like GELU, LayerNorm), so they benefit from more precision bits (3 mantissa = 8 representable fractions per exponent). Gradients can have extreme outliers (especially during loss spikes or early training), so they need wider dynamic range (5 exponent bits = range up to 57344). The precision loss from only 2 mantissa bits is acceptable for gradients because the optimizer averages over many steps.

172 GB saved — that's more than 2 full A100-80GB GPUs worth of memory. This is a major advantage of FP8: even though the optimizer still uses FP32, the activation memory (which dominates at long sequences) is halved compared to FP16.

javascript
function dynamicLossScale(steps, initScale, growInterval) {
  let scale = initScale;
  let goodCount = 0;
  for (const overflow of steps) {
    if (overflow) {
      scale = scale / 2;
      goodCount = 0;
    } else {
      goodCount++;
      if (goodCount >= growInterval) {
        scale = scale * 2;
        goodCount = 0;
      }
    }
  }
  return scale;
}

In practice, Meta assigns the first and last stages fewer layers because they also handle the embedding layer and LM head respectively. The embedding (V×d = 128256×16384 ≈ 2.1B params) is non-trivial at this scale and must be accounted for in the memory budget of stage 0.

7.4 GB for weights alone, leaving ~72 GB on an 80 GB H100 for optimizer states, activations, and communication buffers. With ZeRO-1 across DP=128, optimizer states add 12/128 = 0.094 bytes/param × 3.69B = 0.35 GB. The main memory pressure comes from activations.

~439K petaFLOP-days. For comparison, GPT-4 is estimated at ~100K PF-days. Llama 3.1 405B is roughly 4x the compute of GPT-4. At Meta's scale, this represents tens of millions of dollars in electricity and GPU-hours.

~68 days of continuous training. Meta reported ~54 days for pre-training (implying ~48% MFU). With crashes and restarts, the actual wall-clock time was ~77 days. At 16,384 GPUs × 77 days × 24 hours × ~$2/GPU-hr, the GPU cost alone is ~$60 million.

Still terrible! This is why the 405B training is so challenging. Meta reports using a customized pipeline schedule and overlapping communication with computation to bring the effective bubble down to ~18-20%. They also experimented with larger accumulation steps. The raw 1F1B formula shows why PP=16 is aggressive — you need many microbatches or clever scheduling to overcome the bubble.

Wait, let me recalculate. The ring formula gives total data per GPU = 2S(N-1)/N. With S=7.38 GB and N=128: 2 × 7.38 × 127/128 = 14.63 GB per GPU total. But this is split across 2(N-1) steps. At 400 Gb/s = 50 GB/s InfiniBand, time = 14.63/50 ≈ 293 ms. This is overlapped with backward compute via bucketed AllReduce.

From innermost to outermost: TP → PP → CP → DP.

TP needs the highest bandwidth (AllReduce on every layer's forward/backward — on the critical path). PP needs fast point-to-point (activation transfers between stages — pipelined). CP (context parallelism) uses ring attention which requires fast ring bandwidth but is less frequent. DP uses AllReduce of gradients which can be fully overlapped with backward — so it tolerates the slowest interconnect.

~$1.06 million in electricity alone. This is just the GPU power — add ~40% for cooling, networking, storage, and other infrastructure (PUE ≈ 1.4), bringing total energy cost to ~$1.5M. The GPU rental cost ($2/GPU-hr × 16384 × 77 × 24 ≈ $60.5M) dwarfs the electricity. This is why MFU optimization matters so much — every 1% improvement in MFU saves ~$600K at this scale.