Choose a Parallelism Plan for Training a 70B Model on 64 GPUs
You must fine-tune (continued-pretrain) a 70B dense transformer (L = 80, h = 8192, sequence length 4096) on 64 H100 80 GB GPUs: 8 nodes of 8 GPUs, NVLink inside each node, about 50 GB/s of InfiniBand per GPU between nodes. Global batch: 4,194,304 tokens per step. AdamW, bf16 mixed precision, FlashAttention.
- Show that DDP is impossible and compute the static state per GPU if everything is sharded across all 64 GPUs.
- Estimate the step time at 40 percent MFU and the FSDP communication per GPU per step. Is communication likely to be hidden?
- Evaluate a 2D alternative: tensor parallelism 8 inside each node with FSDP/DP 8 across nodes. What are its communication costs and where do they run?
- Recommend a plan, say what you would measure first, and state the conditions under which you would change it (smaller global batch, slower network, 512 GPUs).
1. Memory. Static state is 16 \times 70 \times 10^9 = 1{,}120 GB, so DDP (full replica per GPU) is impossible against 80 GB. Fully sharded over 64 GPUs: 1120 / 64 = 17.5 GB per GPU, leaving about 60 GB for activations and transient gathered layers, which with activation checkpointing is ample (verified arithmetic).
2. Step time and FSDP communication. FLOPs per step = 6 \times 70\text{e}9 \times 4{,}194{,}304 = 1.76 \times 10^{18}. Peak of 64 H100s = 64 \times 989\text{e}12 = 6.33 \times 10^{16} FLOPs/s (spec value, from memory). At 40 percent MFU the step takes 69.6 s (27.8 s at 100 percent). FSDP sends 3 \times (63/64) \times 2 \times 70 GB = 413 GB per GPU per step; at 50 GB/s that is 8.3 s, about 12 percent of the step, and because layer i+1's weights are prefetched during layer i's compute it can be mostly hidden. The collectives cross nodes, so this is the topology-sensitive part, but with this large a global batch the ratio is favourable: the same plan with a 1M-token batch has a 17.4 s step and the 8.3 s of communication becomes about 48 percent of it, hard to hide.
3. TP=8 x DP=8. Tensor parallelism inside a node: each GPU holds 8.75B parameters' worth of weights; FSDP across the 8 data-parallel replicas moves 3 \times (7/8) \times 2 \times 8.75 GB = 45.9 GB per GPU per step over InfiniBand, 0.92 s at 50 GB/s: nine times less cross-node traffic. The cost moves onto NVLink and onto the critical path: each data-parallel replica processes 4{,}194{,}304/8 = 524{,}288 tokens, and every layer needs 4 all-reduces (2 forward, 2 backward) of 524{,}288 \times 8192 \times 2 B = 8.6 GB; a GPU sends 2 \times 7/8 \times 8.6 = 15.0 GB per all-reduce, so 4 \times 80 \times 15.0 = 4.8 TB per GPU per step. At 450 GB/s per direction (H100 NVLink, spec value from memory) that is 10.7 s, 15 percent of the 69.6 s step, and unlike FSDP all-gathers it is mostly not overlappable because the next operation needs the all-reduced result. So TP trades 8.3 s of (hideable) cross-node traffic for about 10.7 s of (exposed) NVLink traffic plus the engineering cost; it is not a win here.
4. Recommendation. Start with FSDP full sharding across all 64 GPUs (or HSDP: shard within the node, replicate across 8 nodes) with per-block wrapping, FlashAttention, activation checkpointing sized to the memory headroom, and micro-batch as large as fits. Why: the global batch is large, the static state fits easily, communication is predicted to be about 12 percent and hideable, and it uses fewest parallel dimensions. Measure first: MFU from tokens/s, then a profiler trace for exposed NCCL time and per-rank step time spread. Change the plan when: (a) the global batch shrinks (say 1M tokens), making all-gather time comparable to compute: use HSDP, or TP=2 to 4 inside the node so fewer GPUs participate in each cross-node gather; (b) the network is slower (25 GB/s or Ethernet): HSDP or TP+PP so only gradients, or only small pipeline activations, cross nodes; (c) 512 GPUs: a full-shard group of 512 makes every all-gather span the whole cluster and the batch per GPU shrink, so go 3D (TP=8 in node, pipeline across nodes with many micro-batches to keep the bubble (p-1)/(m+p-1) under about 10 percent, DP/ZeRO-1 outside). Principle: the chattiest axis on the fastest link, the fewest axes that fit.
Share this question