Home / Fundamentals / Distributed Training

Distributed TrainingMidQuestion 1 of 7

Estimate the GPU memory needed to train a 7B-parameter model. Where does it go, and how do you reduce it?

30-second answer

Mixed-precision AdamW needs ~16 bytes per parameter: bf16 weights and gradients (4 B) plus fp32 master weights and two Adam moments (12 B), so ~112 GB of state for 7B, more than one 80 GB GPU. Activations add roughly 34·s·b·h bytes per layer: at 4k tokens, h = 4096 and 32 layers that's ~18 GB per sequence. Reduce it by sharding state with ZeRO/FSDP (16/N bytes per parameter with full sharding), activation recomputation (~33% more compute for most of the activation memory), smaller micro-batches with gradient accumulation, and 8-bit optimizer states.

Try answering out loud before reading on. The walkthrough below is paced for a 40-minute interview round.

0–5 min Clarify and frame

Ask: optimizer (AdamW?), precision (bf16 mixed?), sequence length, micro-batch size, GPU type. Framing: "I'll account for model state, then activations, then temporary buffers, give numbers for 7B, and then go through the reduction techniques in order of cost."

5–15 min Core explanation: model state

Item Bytes/param 7B total
bf16 weights 2 14 GB
bf16 gradients 2 14 GB
fp32 master weights 4 28 GB
fp32 Adam m 4 28 GB
fp32 Adam v 4 28 GB
Total 16 112 GB

So even with no activations, a 7B model doesn't fit on one 80 GB GPU for full training with AdamW. (Inference needs only the 14 GB of bf16 weights plus the KV cache.)

15–25 min Deep dive: activations

For each layer, backward needs stored intermediates: inputs to each matmul, normalization inputs, attention outputs, MLP hidden activations. With flash attention (no stored s × s matrix), a standard estimate is:

≈ 34 · s · b · h bytes per layer (bf16)

7B (h = 4096, 32 layers), s = 4096, b = 1:

34 × 4096 × 4096 = 570 MB/layer  → 18.3 GB per sequence

At b = 8 that's ~146 GB, so activations, not weights, often decide the micro-batch size. Without flash attention, the attention scores add 5 · a · s² · b bytes per layer (a = heads), which is huge at long context: for 32 heads at 4k tokens, ~2.7 GB per layer.

Other memory: temporary buffers for communication, the logits (s × V in fp32 can be large: 4096 × 128k × 4 B ≈ 2 GB per sequence; chunked or fused cross-entropy avoids it) and allocator fragmentation.

25–32 min Reducing memory, in order

  1. Shard state with ZeRO/FSDP: 16/N bytes per parameter with full sharding; ~1.75 GB per GPU for 7B on 64 GPUs.
  2. Activation recomputation: keep only layer inputs (2·s·b·h = 34 MB per layer here), costing ~33% more compute; selective recompute gets most of the savings for a few percent.
  3. Smaller micro-batches with gradient accumulation: same effective batch, less activation memory, but smaller matmuls reduce efficiency.
  4. Fused/chunked loss for big vocabularies.
  5. 8-bit optimizer states (−6 bytes/param) or Adafactor.
  6. CPU/NVMe offload of optimizer state: fits large models on few GPUs, but slow.
  7. For fine-tuning: LoRA/QLoRA, which stores optimizer state only for adapters.

32–37 min Worked configuration

7B pre-training on 8 × 80 GB GPUs with FSDP (full sharding within the node): state 112/8 = 14 GB per GPU. With selective recomputation and b = 2 at 4k tokens, activations ≈ 2 × 18 GB × ~0.4 ≈ 15 GB (selective recompute removes the largest attention/MLP intermediates). Total ~30–35 GB plus buffers: comfortable. Throughput, not memory, becomes the constraint.

37–40 min Likely follow-ups

  • Why keep fp32 master weights? bf16 can't represent small updates to large weights; see mixed precision.
  • How does sequence length affect memory? Activations scale linearly with s (with flash attention); without it, quadratically.
  • What's the minimum to fine-tune 70B? QLoRA: 4-bit weights ≈ 35–40 GB plus adapters and activations, so a single 80 GB GPU can work for short sequences.

What interviewers listen for

  • The 16-byte breakdown, fluently.
  • An activation estimate with its dependence on s, b and h.
  • Knowing which technique trades what.

Red flags

  • Estimating training memory as just the weights.
  • Not knowing that activations can dominate.