AAI Agentic Handbook
Question 01 · ML infrastructure & serving

Explain data, tensor and pipeline parallelism and FSDP at a whiteboard level. How would you lay out training of a 70B model on 1,024 H100s?

Reviewed Oct 2026
How to spend the time
  1. 5 min — define the four strategies with one picture each
  2. 10 min — memory per GPU and communication per strategy
  3. 10 min — lay out 70B on 1,024 GPUs against topology
  4. 10 min — FLOPs, wall clock, MFU and failure handling
  5. 5 min — ownership and framework choice
01Concept
  • Data parallelism (DDP): every GPU holds a full model copy, processes a different micro-batch, and all-reduces gradients. Simple, but mixed-precision Adam needs ~16 bytes per parameter (BF16 weights and grads, FP32 master copy and two Adam moments) → 70B × 16 = 1.12 TB per copy.
  • FSDP / ZeRO-3: shard parameters, gradients and optimizer state across the data-parallel ranks; all-gather each layer's weights just in time, reduce-scatter gradients. Memory divides by N; communication is ~1.5× DDP.
  • Tensor parallelism (Megatron): split each matmul across GPUs, with an all-reduce per layer — keep it inside the NVLink domain (8 GPUs on HGX H100, 72 on GB200 NVL72).
  • Pipeline parallelism: split layers into stages across nodes; micro-batches fill the pipe; bubble ≈ (p−1)/(m+p−1) for p stages and m micro-batches.
  • Add context parallelism for long sequences and expert parallelism for MoE.
02Practical use case

70B on 1,024 H100s (128 nodes), 8k sequence: TP=8 inside each node, PP=4 across nodes (20 layers per stage), DP=32 with ZeRO-1 (optimizer sharded). Global batch ~4M tokens = 512 sequences × 8k; with 16 micro-batches per pipeline the bubble is 3/19 ≈ 16%, so use an interleaved schedule to get to ~5–8%; selective activation checkpointing. The simpler alternative on PyTorch FSDP2 + DTensor (torchtitan): TP=8 + HSDP (shard within a node group, replicate across groups).

03Staff framing
  • Topology decides the layout. NVLink gives ~900 GB/s per H100; a 400 Gb/s InfiniBand NIC gives ~50 GB/s per GPU — an 18× gap. TP stays inside a node; PP and DP cross nodes.
  • Target 35–45% BF16 MFU on H100; the Llama 3 report gives 38–43% for its 405B run.
  • Failure modes: stragglers (one slow GPU slows every rank — watch per-rank step time); NCCL hangs (watchdog timeouts); loss spikes (skip the batch, lower LR, rewind a checkpoint).
  • Org scope: a training-infra team owns frameworks and launch templates; researchers own hyperparameters. Use Megatron-LM/NeMo, torchtitan or DeepSpeed — never write parallelism from scratch. For fine-tuning up to 70B, plain FSDP on 1–4 nodes is the default; 3D/4D parallelism is for pretraining scale.
04Estimate

Assumptions: 70B dense; 15T training tokens (far past Chinchilla-optimal ~1.4T, as modern runs over-train for cheaper inference — see scaling laws); training FLOPs ≈ 6·N·D; H100 BF16 dense peak 989 TFLOPS; 40% MFU; illustrative $2.50/H100-hour.

  • Compute: 6 × 70e9 × 15e12 ≈ 6.3e24 FLOPs.
  • Per GPU: 989 × 0.4 ≈ 396 TFLOPS → 6.3e24 / 3.96e14 ≈ 1.6e10 GPU-seconds ≈ 4.4M GPU-hours.
  • Wall clock: 1,024 GPUs ≈ 4,300 h ≈ 180 days; 8,192 GPUs ≈ 22 days. So 1,024 GPUs suits a 70B fine-tune or a ~1–2T-token continued pretrain, not a 15T pretrain.
  • $: 4.4M × $2.50 ≈ $11M for the main run, before restarts (~5–10%) and ablations (often 1–2× the main run).
  • State per GPU (64-way data-parallel group): DDP 1,120 GB (impossible); ZeRO-1 ≈ 293 GB; ZeRO-2 ≈ 155 GB; FSDP/ZeRO-3 ≈ 17.5 GB, leaving ~60 GB for activations.
DDP
1,120 GB per GPU
ZeRO-1
293 GB per GPU
ZeRO-2
155.2 GB per GPU
FSDP / ZeRO-3
17.5 GB per GPU
WeightsGradsOptimizer
70B, mixed-precision Adam, 64-way sharding; activations excluded. An H100 has 80 GB.

What the answer sounds like at each level

Mid

Knows DDP copies the model and syncs gradients, and that big models need sharding.

Senior

Explains FSDP, TP and PP with their communication patterns and memory per GPU, and picks TP inside NVLink.

Staff

Derives memory and FLOPs, lays out 3D parallelism against the network topology, sets an MFU target, and decides what the infra team owns versus researchers.

Follow-ups the interviewer will push

1Now it's 405B.

~6.5 TB of training state: TP=8, PP=16, ZeRO-1 DP and context parallelism for 128k sequences. Meta's Llama 3 405B used four-dimensional parallelism (TP, CP, PP, DP) on up to 16k H100s.

2Fine-tune a 70B with LoRA on one node.

Freeze the BF16 base (140 GB) and shard it with FSDP (17.5 GB per GPU), or use QLoRA 4-bit on fewer GPUs. LoRA optimizer state is tiny; activations dominate, so use gradient checkpointing.

3MFU fell from 40% to 28% going from 512 to 1,024 GPUs.

Communication: DP all-reduce now crosses more spine hops. Increase gradient-bucket size and overlap with compute, use topology-aware placement, raise micro-batches per step, and hunt stragglers.

Red flags interviewers listen for

  • Puts tensor parallelism across nodes over InfiniBand.
  • Cannot estimate bytes per parameter for training state.
  • Confuses data parallelism with pipeline parallelism.
  • Proposes building a custom parallelism framework.

How it's scored

DimensionWhat strong looks like
ConceptsAccurate memory and communication for DDP, FSDP, TP and PP, including the pipeline bubble.
Estimation6·N·D FLOPs, MFU, GPU-hours and memory per GPU worked out.
LayoutMaps strategies to NVLink vs InfiniBand and justifies degrees.
PragmatismUses existing frameworks; knows FSDP alone suffices for fine-tuning.

Sources

  1. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism — Shoeybi et al., 2019paper
  2. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — Rajbhandari et al., 2019paper
  3. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel — Zhao et al., 2023paper
  4. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM — Narayanan et al., 2021paper
  5. The Llama 3 Herd of Models — Llama Team, Meta, 2024paper