Worked Examples — infra.bea.ai

This chapter walks through four real-world parallelism planning exercises, applying every formula from the previous chapters. Each example follows the same six-step process: fit check → TP → PP → DP → communication → memory verification.

Example 1: 8B Dense on 8× H100 (Single Node)

Setup: Llama-3.1 8B, $h=4096$ , $L=32$ , GQA-8, $d_{ff}=14336$ . Training on 1 DGX H100 (8 GPUs, 80 GB each, NVLink 4.0).

Step 1: Fit Check

\text{Static memory} = 16 \times 8 \times 10^9 = 128 \text{ GB}

128 GB > 80 GB → doesn’t fit on one GPU. But with FSDP across 8 GPUs:

\frac{128 \text{ GB}}{8} = 16 \text{ GB per GPU}

✅ Easily fits. No TP/PP needed.

Step 2-4: Configuration

Parameter	Value	Reasoning
TP	1	Model fits with FSDP, no TP needed
PP	1	Not needed
DP (FSDP)	8	All GPUs for data parallelism

Step 5: Communication

FSDP AllGather + ReduceScatter per step:

T_{\text{FSDP comm}} = \frac{2 \times 8B \times 2 \text{ bytes}}{900 \text{ GB/s (NVLink)}} = 35.6 \text{ ms}

All communication is intra-node (NVLink) — very fast. Backward takes ~100 ms → mostly overlapped.

Step 6: Memory Verification

Component	Value
Parameters (FSDP sharded)	16 GB / 8 = 2 GB
Optimizer (sharded)	96 GB / 8 = 12 GB
Gradients (sharded)	16 GB / 8 = 2 GB
Activations (b_μ=4, full ckpt)	~4 GB
Buffers	3 GB
Total	23 GB ✅

Can increase micro-batch to 8-16 and still have headroom. Global batch via gradient accumulation.

Expected MFU

Compute per step: $6 \times 8B \times b_{global} \times s$ FLOPs (forward + backward)
H100 peak BF16: 989 TFLOPS × 8 = 7,912 TFLOPS
Target MFU: 45-55% for FSDP on a single node

Example 2: 70B Dense on 64× H100 (8 Nodes)

Setup: Llama-2 70B, $h=8192$ , $L=80$ , GQA-8, $d_{ff}=28672$ . Training on 8 DGX H100 nodes, IB NDR (400 GB/s per node).

Step 1: Fit Check

\text{Weights (BF16)} = 2 \times 70B = 140 \text{ GB} > 80 \text{ GB}

Weights alone exceed one GPU → need TP or PP.

Step 2: TP Degree

\text{Min TP} = \lceil 140 / (80 \times 0.7) \rceil = \lceil 2.5 \rceil = 4

But TP=8 (all NVLink GPUs) gives better compute/comm ratio:

TP=4: 35 GB params per GPU → fits, but TP AllReduce uses only half the NVLink bandwidth
TP=8: 17.5 GB per GPU → fits with more headroom, uses full NVLink

Choose TP=8 — maximize NVLink utilization.

Step 3-4: PP and DP

PP=1 (model fits with TP=8)
DP = 64 / 8 = 8 (one DP rank per node)

Step 5: Communication

Communication	Calculation	Time
TP AllReduce (per layer)	$2 \times \frac{2 \times 64\text{MB}}{900\text{GB/s}}$	0.28 ms
TP total (80 layers)	$80 \times 0.28$	22.4 ms
DP AllReduce	$\frac{2 \times 17.5\text{GB}}{400\text{GB/s}}$	87.5 ms
Backward compute	(estimated)	~500 ms

DP AllReduce (87.5 ms) fully overlaps with backward (500 ms). TP comm (22.4 ms) is exposed but small vs. total compute.

Communication efficiency: $\frac{500}{500 + 22.4} \approx 96\%$ ✅

Step 6: Memory Verification

Component	Value
Parameters (TP=8)	17.5 GB
Gradients (TP=8)	17.5 GB
Optimizer (TP=8, DP=8)	840 GB / 64 = 13.1 GB (sharded per DP)
Activations (b_μ=2, full ckpt)	~3 GB
Buffers	3 GB
Total	54.1 GB ✅

Note: Without FSDP sharding the optimizer, optimizer memory would be 840/8 = 105 GB → OOM. Must use at least ZeRO-1 (optimizer sharding) with TP=8 and DP=8.

Example 3: 405B Dense on 256× H100 (32 Nodes)

Setup: Llama-3.1 405B, $h=16384$ , $L=126$ , GQA with 8 KV heads. Training on 32 DGX H100 nodes.

Step 1: Fit Check

\text{Weights (BF16)} = 810 \text{ GB}

TP=8 gives 101 GB per GPU → still exceeds 80 GB. Need PP.

Step 2-4: Configuration

\text{TP=8, PP=} \lceil \frac{810}{8 \times 80 \times 0.7} \rceil = \lceil 1.82 \rceil = 2

PP=2: 810 / 16 = 50.6 GB per GPU → fits.

TP=8 (intra-node NVLink)
PP=2 (across 2 nodes, IB)
DP = 256 / (8 × 2) = 16

Step 5: Communication

Communication	Time	Exposed?
TP AllReduce (126 layers, NVLink)	~35 ms	Yes
PP P2P (b_μ=1, IB)	$\frac{128\text{MB}}{50\text{GB/s}}$ = 2.56 ms	Partially (bubble)
PP bubble (PP=2, m=32)	1/32 = 3.1%	Yes
DP AllReduce (IB, 50.6 GB)	~253 ms	Overlapped with backward

Backward compute: ~1200 ms → DP AllReduce fully overlapped. Total exposed communication: ~38 ms out of ~1238 ms → 3% overhead ✅

Step 6: Memory Verification

Component	Value
Parameters (TP=8, PP=2)	50.6 GB
Gradients (TP=8, PP=2)	50.6 GB
Wait — that’s already over 80 GB!	❌ Need FSDP

Correction: Must use FSDP (ZeRO-3) within the DP dimension:

Component	Value
Parameters (TP=8, PP=2, FSDP gathered on demand)	~6.3 GB resident
Gradients (sharded)	50.6 / 16 = 3.16 GB
Optimizer (sharded)	303.75 / 16 = 18.98 GB
Activations (63 layers, full ckpt, b_μ=1)	~3 GB
Temp parameter gather (one layer at a time)	~0.8 GB
Buffers	3 GB
Total	~35 GB ✅

Lesson: At 405B scale, FSDP is mandatory even with TP+PP. Raw TP+PP leaves 50.6 GB of parameters + 50.6 GB of gradients per GPU — doesn’t fit.

Example 4: DeepSeek-V3 (671B MoE) on 2048× H800

Setup: 671B MoE (37B active per token), 256 experts, top-2, $h=7168$ , $d_{ff}=18432$ per expert, 61 layers. 2048 H800 GPUs (NVLink 4.0, IB).

Configuration (from their paper)

Parameter	Value	Reasoning
TP	1	MoE reduces per-expert size enough to skip TP
EP	64	256 experts / 4 experts per GPU = 64
PP	16	61 layers spread across 16 stages
DP	2048 / (1 × 16 × 1) = 128	But EP=64 means EP groups of 64 handle expert routing

Why TP=1? Each expert has only 18432 × 7168 × 2 = 264M params (0.53 GB in BF16). No need to split across GPUs. Instead, 4 experts per GPU × 64 EP ranks = 256 experts.

Communication Pattern

Per MoE layer:
  1. All-to-All dispatch: send each token's activation to its expert's GPU
  2. Expert compute
  3. All-to-All combine: send results back

Cross-node EP communication: the #1 bottleneck for MoE training

DeepSeek mitigates this with:

Auxiliary-loss-free load balancing → even expert distribution → predictable All-to-All traffic
FP8 dispatch → 2× less data to send
Overlapping EP comm with non-MoE compute

Memory per GPU

Shared params: ~11B → 22 GB (BF16), sharded across PP
Expert params: 4 experts × 264M = 1.06B → 2.12 GB per GPU
Optimizer: sharded across DP
Total: ~25 GB per GPU → comfortable fit on 80 GB

Summary Table

For the decision algorithm behind these configurations, see Decision Process. For memory formulas, see Memory Budget Verification.

Model	GPUs	TP	PP	DP	EP	FSDP	Comm overhead	MFU
8B	8	1	1	8	—	ZeRO-3	<5%	45-55%
70B	64	8	1	8	—	ZeRO-1	~4%	40-50%
405B	256	8	2	16	—	ZeRO-3	~3%	35-45%
671B MoE	2048	1	16	128	64	ZeRO-1	~8%	33% (reported)

Self-Check

Plan parallelism for a 13B model on 32 GPUs (4 nodes × 8 GPUs)
Identify when FSDP becomes mandatory (hint: optimizer memory)
Explain why DeepSeek-V3 uses TP=1 while Llama-405B uses TP=8
Calculate the pipeline bubble for PP=4 with 16 micro-batches
Verify memory fits for each example’s configuration