Distributed Training Planner

When planning a distributed training run for a large language model, use this framework to determine the parallelism strategy, memory budget, communication overhead, and expected throughput.

Input Requirements

Provide:

• Model size (parameters in billions)
• Target training tokens (in trillions)
• Available GPUs (type: A100/H100/H200, count, interconnect: NVLink/InfiniBand)
• GPU memory (80GB for A100/H100, 141GB for H200)
• Nodes (GPUs per node, number of nodes)
• Budget constraints (max cost in dollars, max wall-clock time)

Step 1: Memory Budget

Calculate per-GPU memory for each component:

Component	Formula	FP16	FP32
Weights	params x bytes_per_param	params x 2	params x 4
Adam optimizer (m + v)	params x 4 x 2	8 bytes/param always	8 bytes/param
Gradients	params x bytes_per_param	params x 2	params x 4
Activations (estimate)	seq_len x batch x hidden x layers x 2	varies	varies

If total exceeds GPU memory, sharding is required. Try in order:

1. ZeRO-1 (shard optimizer only) -- cheapest communication
2. ZeRO-2 (+ gradients) -- moderate communication
3. FSDP/ZeRO-3 (+ weights) -- highest communication but maximum memory savings
4. Add activation checkpointing if activations still too large
5. Add tensor parallelism if a single layer does not fit on one GPU

Step 2: Parallelism Strategy

Decision Tree

1. Does one layer fit on one GPU?

   • No: You need tensor parallelism. Set TP = 2, 4, or 8 (within a node).
   • Yes: Skip tensor parallelism.

2. Does the full model (with sharding) fit on GPUs within one node?

   • No: You need pipeline parallelism. Set PP = number of nodes / groups.
   • Yes: Skip pipeline parallelism.

3. How many remaining GPUs for data parallelism?

   • DP = total_gpus / (TP x PP)

4. What sharding level within the data parallel group?

   • Start with FSDP (ZeRO-3). Reduce to ZeRO-2 or ZeRO-1 if communication is bottleneck.

Typical Configurations

Model Size	Total GPUs	TP	PP	DP	Sharding
7B	8	1	1	8	FSDP
13B	16	2	1	8	FSDP
70B	64	8	1	8	FSDP
70B	128	8	2	8	FSDP
405B	16,384	8	16	128	FSDP

Step 3: Communication Analysis

Estimate communication volume per training step:

• Data parallel (all-reduce): 2 x gradient_size x (N-1)/N per step
• FSDP (all-gather + reduce-scatter): ~3 x weight_size x (N-1)/N per step (higher than DP)
• Tensor parallel (all-reduce per layer): 2 x activation_size x num_layers per step (needs NVLink)
• Pipeline parallel (point-to-point): activation_size per stage boundary (minimal)

If communication time exceeds 20% of compute time, the strategy is communication-bound. Solutions:

• Gradient accumulation (reduce all-reduce frequency)
• Overlap communication with computation (FSDP does this by default)
• Increase micro-batch size (better compute-to-communication ratio)
• Switch to a less communication-heavy sharding stage

Step 4: Throughput and Cost Estimate

FLOPS per training step:

• Forward: ~2 x params x tokens_per_batch
• Backward: ~4 x params x tokens_per_batch (2x forward)
• Total: ~6 x params x tokens_per_batch

Training time:

• total_flops = 6 x params x total_tokens
• time_seconds = total_flops / (num_gpus x gpu_tflops x 1e12 x utilization)
• Typical utilization: 35-45% (accounting for communication, pipeline bubbles, memory overhead)

Cost:

• total_gpu_hours = num_gpus x time_seconds / 3600
• cost = total_gpu_hours x cost_per_gpu_hour

Step 5: Validation Checklist

Before launching:

1. Per-GPU memory fits within hardware limit (with 10% headroom)
2. Effective batch size matches target (per_gpu_batch x DP x gradient_accumulation_steps)
3. Communication-to-compute ratio is below 20%
4. Pipeline bubble fraction is below 15% (enough micro-batches)
5. Learning rate is scaled for the effective batch size
6. Checkpointing frequency accounts for failure probability (save every 1-2 hours for large runs)
7. Gradient clipping is set (typically 1.0 for large models)
8. Warmup steps are proportional to total steps (typically 0.1-1% of total)

Red Flags

• TP > 8: Tensor parallelism across nodes (over InfiniBand) is almost always slower than pipeline parallelism
• Pipeline stages > 32: Bubble overhead becomes significant even with many micro-batches
• Effective batch size > 10M tokens: Diminishing returns; may harm convergence
• Utilization below 30%: Communication-bound -- re-evaluate parallelism strategy
• No activation checkpointing above 13B: You will run out of memory during the backward pass
• No gradient accumulation with small per-GPU batch: Gradient noise increases; accumulate to effective batch of 256+ samples