Home / Fundamentals / Distributed Training

Distributed TrainingStaffQuestion 5 of 7

Plan the pre-training of a 70B dense model on 1.4T tokens with 2,048 H100s. Choose a parallel layout and estimate the time.

30-second answer

Compute is 6 × 70e9 × 1.4e12 ≈ 5.9e23 FLOPs; at 40% MFU on 2,048 × 989 TFLOP/s that's ~7.3e5 s, about 8.4 days, so budget ~10–11 days with failures. State is ~1.1 TB, so model parallelism is required: TP = 8 within each node, PP = 4 across nodes (20 layers per stage), DP = 64 with ZeRO-1, giving ~9 GB of weights, gradients and optimizer shards per GPU plus activations. With a 4M-token batch (1,024 × 4k sequences), each replica runs 16 micro-batches, a 1F1B bubble of ~16% that interleaving halves. Add sequence parallelism, selective recompute, async checkpoints about every 20 minutes and hot spares.

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

0–5 min Clarify and frame

Clarify: H100 SXM 80 GB, 8 per node with NVLink, 400 Gb/s InfiniBand per GPU; sequence length 4,096; dense decoder (80 layers, h = 8192, GQA); target global batch ~4M tokens. Framing: "Compute budget and time first, then memory, then the layout, then the batch and pipeline schedule, then reliability."

5–15 min Compute and time

FLOPs = 6 · N · D = 6 × 70e9 × 1.4e12 ≈ 5.9e23
Cluster peak = 2,048 × 989e12 ≈ 2.0e18 FLOP/s
At 40% MFU: 8.1e17 FLOP/s  →  5.9e23 / 8.1e17 ≈ 7.3e5 s ≈ 8.4 days

Attention adds about 12 · L · h · s = 12 × 80 × 8192 × 4096 ≈ 3.2e10 FLOPs per token, ~8% on top of 6N = 4.2e11, so ~9 days of pure compute. With failures, restarts and evaluations, plan ~11 days.

Steps: 1.4e12 / 4.2e6 tokens per step ≈ 334k steps, so each step must take ~2.3 s.

15–25 min Memory and layout

State: 70e9 × 16 B ≈ 1.1 TB. Must shard across GPUs.

Scroll sideways to see the whole diagram.

3D parallel layout: TP 8 × PP 4 × DP 64Components: Stage 1 · node (TP 8 · emb, 1–19); Stage 2 · node (TP 8 · layers 20–40); Stage 3 · node (TP 8 · layers 41–61); Stage 4 · node (TP 8 · 62–80 + LM head); Stage 1 × 63 (same weights shard); Stage 2 × 63; Stage 3 × 63; Stage 4 × 63; ZeRO-1 shards (fp32 + Adam ÷ 64). Flow: Stage 1 · node to Stage 2 · node (acts); Stage 2 · node to Stage 3 · node (acts); Stage 3 · node to Stage 4 · node (acts); Stage 1 · node to Stage 1 × 63 (DP reduce-scatter); Stage 2 · node to Stage 2 × 63; Stage 3 · node to Stage 3 × 63; Stage 4 · node to Stage 4 × 63; ZeRO-1 shards to Stage 4 · node (all stages).REPLICA 1 OF 64 · 4 NODES · 32 GPUSREPLICAS 2 … 64 (SAME LAYOUT, DIFFERENT DATA)Stage 1 · nodeTP 8 · emb, 1–19Stage 2 · nodeTP 8 · layers 20–40Stage 3 · nodeTP 8 · layers 41–61Stage 4 · nodeTP 8 · 62–80 + LM headStage 1 × 63same weights shardStage 2 × 63Stage 3 × 63Stage 4 × 63ZeRO-1 shardsfp32 + Adam ÷ 64actsactsactsDP reduce-scatterall stages
3D parallel layout: TP 8 × PP 4 × DP 64. Each node's 8 GPUs form one tensor-parallel group over NVLink and hold one pipeline stage of about 20 layers. Four nodes make one model replica; activations flow stage to stage point-to-point. 64 replicas process different data, and each stage reduce-scatters its gradients with the same stage in every other replica over InfiniBand, with optimizer state sharded by ZeRO-1.

Layout: TP = 8 × PP = 4 × DP = 64 = 2,048.

  • TP = 8 within each node: per-layer all-reduces stay on NVLink. With sequence parallelism, activations are divided by 8.
  • PP = 4 across 4 nodes: 20 layers per stage (the first and last stages hold one or two fewer layers to offset embedding and LM head).
  • DP = 64 across replicas, with ZeRO-1 sharding fp32 optimizer state.

Per GPU:

params per GPU = 70e9 / (8 × 4) ≈ 2.19e9
bf16 weights + grads: 2.19e9 × 4 B ≈ 8.8 GB
fp32 master + Adam, ZeRO-1 over 64: 2.19e9 × 12 / 64 ≈ 0.4 GB

Activations per micro-batch per GPU (s = 4,096, b = 1, TP+SP = 8):

34 × 4096 × 8192 / 8 ≈ 143 MB per layer × 20 layers ≈ 2.9 GB
1F1B keeps up to p = 4 micro-batches in flight on stage 1: ≈ 11.5 GB

Total ~21 GB plus buffers: well within 80 GB, so use selective recomputation only (or none), keeping MFU high. An alternative with FSDP (shard 1.1 TB over 2,048 GPUs = 0.5 GB/GPU) plus TP = 8 also works; PP reduces cross-node traffic for gradients and is well proven at this scale.

25–32 min Batch and pipeline schedule

  • Global batch: 1,024 sequences × 4,096 = 4.2M tokens.
  • Per DP replica: 1,024 / 64 = 16 sequences → m = 16 micro-batches of 1.
  • Bubble: (4 − 1)/(16 + 3) ≈ 16%. Interleaved 1F1B with v = 2 (10 layers per chunk) cuts it to ~3/35 ≈ 9%.
  • Gradient reduce-scatter across DP (64 ranks across nodes) overlaps with the last micro-batches' backward.
  • Batch-size warmup (start smaller) is optional; the LR schedule is WSD so the run can be extended.

Communication check: each TP all-reduce moves b·s·h·2 B = 64 MB per layer, four per layer per micro-batch, over NVLink: fine. PP sends s·h·2 B / 8 = 8 MB (with SP) per micro-batch boundary. DP: each GPU reduces ~4.4 GB of bf16 gradients per step, over ~2 s: comfortably overlapped.

32–37 min Reliability and operations

  • Expect a failure every few hours at 2,048 GPUs (Llama 3 saw one roughly every 3 hours at 16k GPUs; fewer GPUs → less frequent).
  • Checkpoint: ~1.1 TB sharded across ranks, written asynchronously; interval ≈ √(2 × 1 min × MTBF) ≈ 20–40 minutes.
  • Hot-spare nodes, pre-flight NCCL and GPU health checks, automatic restart.
  • Monitor loss, grad norm, per-rank step time, max attention logits; deterministic data order for replays.
  • Validate the whole stack with a short run at full scale (MFU, checkpoint/restore, loss parity with a small reference).

37–40 min Likely follow-ups

  • What if MFU is 30%? ~11 days of compute; profile for exposed communication and bubble.
  • What if we want 32k context later? Add a long-context phase with context parallelism (CP = 4–8), adjusting DP.
  • What about MoE? Add expert parallelism with all-to-all; a different layout.

What interviewers listen for

  • Correct arithmetic, end to end.
  • A layout justified by memory and topology.
  • Micro-batches and bubble calculation.
  • Reliability planning.

Red flags

  • A layout without checking memory.
  • TP across nodes.
  • No allowance for failures.