Decision Process

Given a model and a cluster, how do you choose the parallelism configuration? This chapter provides the systematic process — not guesswork, but a deterministic algorithm that any AI infrastructure engineer can follow. Then we show how the chosen configuration maps to PyTorch process groups.

---

1. The Inputs

You need three things:

Input	What you need	Example
Model	Parameter count, architecture, MoE config	70B dense, GQA-8, 80 layers, d=8192
Cluster	GPU count, memory, NVLink/IB bandwidth	32× H100 SXM (80 GB), NVLink 4.0, IB NDR
Target	Throughput, latency, or training time	Train on 2T tokens in 30 days

---

2. The Six-Step Algorithm

Step 1: Can the Model Fit on One GPU?

Training memory per parameter (BF16 mixed precision + AdamW):

Component	Precision	Bytes per param	Purpose
Parameters	BF16	2Ψ	Forward/backward working copy
Gradients	BF16	2Ψ	Backward output
FP32 master weights	FP32	4Ψ	Optimizer’s high-precision copy
Adam first moment (m)	FP32	4Ψ	Exponential moving average of gradients
Adam second moment (v)	FP32	4Ψ	EMA of squared gradients
Static total		16Ψ	Without activations

Key insight from optimizer accounting: The optimizer states alone (12Ψ) consume 75% of static training memory. This is exactly why ZeRO/FSDP shards optimizer states first — it’s the highest-ROI optimization.

Adding activation memory (depends on batch size $b$, sequence length $s$, hidden dimension $h$, layers $L$):

$$M_{\text{total}} = 16\Psi + \underbrace{k \cdot b \cdot s \cdot h \cdot L}_{\text{activations}}$$

where $k$ depends on whether activation checkpointing is used (roughly $k \approx 34$ without checkpointing, $\approx 2$ with full checkpointing for a Transformer layer).

Model	Static (16Ψ)	On 80 GB GPU?
7B	112 GB	❌ (needs FSDP or TP)
13B	208 GB	❌
70B	1,120 GB	❌

Surprise: Even a 7B model’s training state (112 GB) doesn’t fit on an 80 GB GPU! But with FSDP (ZeRO-3), the optimizer states and parameters are sharded across all GPUs. With 8 GPUs: 112/8 = 14 GB static per GPU — easily fits.

If sharded training fits on available GPUs: FSDP + DP is sufficient. Skip to Step 4. If model weights alone (2Ψ in BF16) exceed one GPU: Need TP and/or PP.

Step 2: Determine Minimum TP Degree

$$\text{TP degree} = \left\lceil \frac{N \times 2}{M_{\text{GPU}} \times 0.7} \right\rceil \qquad \text{(70\% of GPU memory for weights, rest for activations)}$$

Round up to the nearest power of 2 that divides the attention head count.

Model	Weights (BF16)	Min TP (80 GB GPU)	Typical TP
7B	14 GB	1 (fits)	1
13B	26 GB	1 (fits)	1-2
70B	140 GB	3 → TP=4	4-8
405B	810 GB	15 → TP=8 + PP	8 + PP=4

Constraint: TP ≤ GPUs per node. TP uses NVLink; crossing node boundaries is too slow.

For TP communication patterns (column/row splitting, AllReduce per layer), see Tensor Parallelism.

Step 3: Determine PP Degree (If Needed)

If TP alone doesn’t make the model fit:

$$\text{PP degree} = \left\lceil \frac{N \times 2}{\text{TP} \times M_{\text{GPU}} \times 0.7} \right\rceil$$

Model	With TP=8	Memory per GPU	PP needed?
70B	140/8 = 17.5 GB	17.5 GB	No
405B	810/8 = 101 GB > 80 GB	Overflow	Yes: PP=2 → 50.6 GB

For PP scheduling (1F1B, interleaved), see Pipeline Parallelism.

Step 4: Assign Remaining GPUs to DP

$$\text{DP degree} = \frac{\text{Total GPUs}}{\text{TP} \times \text{PP}}$$

For 70B on 256 GPUs: TP=8, PP=1 → DP = 256/8 = 32. For 405B on 256 GPUs: TP=8, PP=2 → DP = 256/16 = 16.

Step 5: Verify Communication Budget

Check that communication time ≤ compute time:

Strategy	Communication per step	Formula
TP AllReduce	Per layer (very frequent)	$\frac{2 \times B \times S \times d \times 2}{B_{\text{NVLink}}}$
DP AllReduce	Per training step (overlapped)	$\frac{2 \times \Psi \times 2}{B_{\text{IB}}}$
PP P2P	Per micro-batch	$\frac{B_\mu \times S \times d \times 2}{B_{\text{IB}}}$

If any check fails → increase TP (if NVLink available), reduce DP, or add PP.

For detailed communication cost formulas, see Communication Overhead Estimation.

Step 6: Add EP for MoE (If Applicable)

$$\text{EP degree} = \frac{\text{num\_experts}}{\text{experts\_per\_GPU}}$$

EP reduces DP degree: $\text{DP} = \frac{\text{Total GPUs}}{\text{TP} \times \text{PP} \times \text{EP}}$

---

3. From Configuration to Process Groups

Once you’ve determined (TP, PP, DP), you need to create sub-communication groups in PyTorch so each parallelism dimension has its own independent communication domain.

Why Sub-Groups?

The default process group contains all ranks. But TP AllReduce should only happen within 8 GPUs, not all 256. Each parallelism dimension needs its own group — like separate rehearsal rooms where different ensembles practice independently.

Example: TP=2, DP=4 on 8 GPUs

import torch.distributed as dist

# After init_process_group(backend="nccl")
rank = dist.get_rank()
tp_size, dp_size = 2, 4

# TP groups: ranks that do tensor-parallel together
# {0,1}, {2,3}, {4,5}, {6,7}
tp_group_ranks = [
    list(range(i, i + tp_size))
    for i in range(0, tp_size * dp_size, tp_size)
]

tp_group = None
for ranks in tp_group_ranks:
    group = dist.new_group(ranks=ranks)  # ALL ranks must call this!
    if rank in ranks:
        tp_group = group

# DP groups: same TP-position across TP groups
# {0,2,4,6}, {1,3,5,7}
dp_group_ranks = [
    list(range(i, tp_size * dp_size, tp_size))
    for i in range(tp_size)
]

dp_group = None
for ranks in dp_group_ranks:
    group = dist.new_group(ranks=ranks)
    if rank in ranks:
        dp_group = group

# Usage:
# dist.all_reduce(tensor, group=tp_group)   # TP AllReduce within {0,1}
# dist.all_reduce(grads, group=dp_group)     # DP AllReduce within {0,2,4,6}

Critical: new_group() must be called by all processes even if they don’t belong to that group. If any process skips the call, the others will deadlock waiting.

3D Parallelism: Three Dimensions of Groups

For TP=8, PP=4, DP=8 on 256 GPUs:

Rank layout (TP varies fastest, then PP, then DP):
  DP group 0: Ranks   0-31   (TP=8 × PP=4)
  DP group 1: Ranks  32-63
  ...
  DP group 7: Ranks 224-255

Within each DP group:
  PP stage 0: Ranks 0-7   (TP group)
  PP stage 1: Ranks 8-15
  PP stage 2: Ranks 16-23
  PP stage 3: Ranks 24-31

Three communication domains:
  TP group: {0,1,2,3,4,5,6,7}     — NVLink AllReduce per layer
  PP group: {0,8,16,24}           — IB P2P per micro-batch
  DP group: {0,32,64,...,224}     — IB AllReduce per step

Megatron-LM’s mpu.initialize_model_parallel() sets all this up automatically given (tp_size, pp_size, dp_size).

---

4. Output: The Configuration

Model	GPUs	TP	PP	DP	EP	CP
7B training	8	1	1	8	—	1
70B training	64	8	1	8	—	1
405B training	256	8	4	8	—	1
DeepSeek-V3 (671B MoE)	2048	4	—	8	64	1
70B, 128K context	64	8	1	4	—	2

---

Self-Check

• Walk through the 6-step algorithm for a 13B model on 16 GPUs
• Calculate static training memory using the 16Ψ formula
• Explain why optimizer states dominate training memory (12 out of 16 bytes/param)
• Determine when PP is needed on top of TP
• Create TP and DP sub-groups with dist.new_group() for a given configuration
• Explain why new_group() must be called by all processes

---

References

• Megatron-LM: Training Multi-Billion Parameter Language Models (Shoeybi et al., 2019)
• Efficient Large-Scale Language Model Training (Narayanan et al., 2021)
• PyTorch Distributed: Process Groups