Finding the Best Training Configuration for Distributed Large Model Training

1. Problem Statement and Decision Framework

Given a model with $P$total parameters, a target global batch size$B_{\text{global}}$, a cluster of $N_{\text{nodes}}$nodes each containing$G$GPUs (total GPU count$N_{\text{GPU}} = N_{\text{nodes}} \times G$), and per-GPU memory capacity $M_{\text{GPU}}$, the objective is to determine the optimal 5D parallelism configuration:

where the configuration vector is:

The search space $\mathcal{S}$ is constrained by the fundamental parallelism decomposition constraint:

and the global batch size equation:

where $s$ is the sequence length in tokens.

Model FLOPs Utilization (MFU) is the primary throughput metric:

where $F_{\text{model}}$is the forward-pass FLOPs per token (approximately$2P$for dense transformers),$T_{\text{step}}$is the wall-clock time per training step, and$\Phi_{\text{peak}}$is the peak FLOPs/s per GPU (e.g.,$\approx 989\;\text{TFLOP/s}$ for H100 SXM in BF16).

2. Step 1 — Fitting a Training Step in Memory

2.1 Memory Accounting

The total per-GPU memory requirement decomposes as:

Here $h$is the hidden dimension and$L_{\text{local}} = L / D_{\text{PP}}$is the number of transformer layers assigned to each pipeline stage, with$L$ total layers.

2.2 ZeRO Optimization Stages and Their Memory Impact

ZeRO (Zero Redundancy Optimizer) partitions optimizer state, gradients, and optionally parameters across $D_{\text{DP}}$ data-parallel ranks:

The memory feasibility constraint is therefore:

2.3 GPU-Rich Case — Decision Heuristics

The decision tree is determined by parameter count $P$and available GPU count$N_{\text{GPU}}$:

Case A: $P < 10\text{B}$ (Small-to-Medium Models)

A single parallelism dimension typically suffices:

• Option 1: Tensor Parallelism with $D_{\text{TP}} \leq 8$ (intra-node).

  Per-GPU parameter memory:

  $$M_{\text{params}}^{\text{TP}} = \frac{2P}{D_{\text{TP}}} \;\text{bytes}$$

• Option 2: ZeRO-3 with Data Parallelism across 8 GPUs plus full activation recomputation.

  Per-GPU total model state memory:

  $$M_{\text{model}}^{\text{Z3}} = \frac{16P}{D_{\text{DP}}} \;\text{bytes}$$

For a 7B parameter model on 8 GPUs with ZeRO-3:

This easily fits within the 80 GB of an H100, leaving ample room for activations.

Case B: $10\text{B} \leq P \leq 100\text{B}$ (Large Models, Multi-Node)

Multiple parallelism dimensions must be composed. The number of GPUs required exceeds one node ($> 8$ GPUs), introducing inter-node communication as a critical consideration. Viable configurations include:

Communication volume analysis dictates the preferred choice. For ZeRO-3, the per-step communication volume per GPU scales as:

The factor of 2 accounts for one all-gather (forward) and one reduce-scatter (backward) per training step.

Case C: $N_{\text{GPU}} \geq 512$ (Large-Scale Clusters)

At this scale, pure ZeRO-3 becomes communication-bound because the all-gather and reduce-scatter collectives span hundreds of GPUs across many nodes. The communication time for a ring all-gather over $D_{\text{DP}}$ ranks is:

where $\beta_{\text{inter}}$is the inter-node bandwidth,$\alpha$is per-message latency, and$M_{\text{shard}} = 2P / (D_{\text{DP}} \cdot D_{\text{TP}} \cdot D_{\text{PP}})$.

As $D_{\text{DP}} \to \infty$, the latency term $(D_{\text{DP}} - 1)\alpha$dominates, making pure DP inefficient. The solution is to limit$D_{\text{DP}}$ by introducing TP and PP:

Case D: $N_{\text{GPU}} \geq 1024$ (Frontier Scale)

The recommended configuration is a 3D parallelism combination:

with ZeRO Stage 2 (sharding optimizer states and gradients but not parameters, thus avoiding the additional all-gather for parameters during forward pass).

Special Considerations

Context Parallelism (CP) for very long sequences ($s \gg 4096$): Activation memory scales linearly with sequence length. When $s$ is large, activations dominate:

Context parallelism partitions the sequence dimension across $D_{\text{CP}}$ GPUs:

CP is placed across nodes since its communication (ring attention all-to-all for KV exchange) is less bandwidth-intensive than TP.

Expert Parallelism (EP) for Mixture-of-Experts (MoE): If the model has $E$experts per MoE layer, each expert has parameters$P_{\text{expert}}$. Expert parallelism distributes experts across $D_{\text{EP}}$ GPUs:

EP is placed across nodes because the all-to-all communication pattern (dispatching tokens to experts) can tolerate higher latency compared to the tight synchronization required by TP.

2.4 GPU-Poor Case — Memory Reduction Techniques

When $M_{\text{total}} > M_{\text{GPU}}$ and additional GPUs are unavailable, two primary strategies reduce memory:

Full Activation Recomputation: Instead of storing all intermediate activations during the forward pass, discard them and recompute during the backward pass. This eliminates $M_{\text{act}}$ at the cost of one additional forward pass:

Since $F_{\text{bwd}} \approx 2 F_{\text{fwd}}$:

This is approximately a 33% increase in compute time for a near-complete elimination of activation memory.

Gradient Accumulation: Process the global batch as $N_{\text{GAS}}$ sequential micro-batches, accumulating gradients before the optimizer step:

Increasing $N_{\text{GAS}}$allows using a smaller$b_{\text{mbs}}$, reducing peak activation memory:

3. Step 2 — Achieving the Target Global Batch Size

After Step 1 establishes a memory-feasible configuration with some initial values of $D_{\text{DP}}$, $N_{\text{GAS}}$, and $b_{\text{mbs}}$, the current effective global batch size is:

The target is $B_{\text{global}} = B_{\text{target}}$(e.g.,$1\text{M tokens}$).

3.1 Increasing $B_{\text{current}}$to$B_{\text{target}}$

If $B_{\text{current}} < B_{\text{target}}$, increase the batch size via:

The scaling factor required is:

3.2 Decreasing $B_{\text{current}}$to$B_{\text{target}}$

If $B_{\text{current}} > B_{\text{target}}$ (too many data-parallel replicas), reallocate GPUs from DP to other parallelism dimensions:

where $r$is the reallocation factor, preserving$N_{\text{GPU}} = D_{\text{DP}} \times D_{\text{TP}} \times D_{\text{PP}} \times D_{\text{CP}} \times D_{\text{EP}}$.

4. Step 3 — Optimizing Training Throughput

With memory feasibility and correct batch size established, the final objective is maximizing MFU. The per-step training time decomposes as:

where:

• $T_{\text{compute}}$: time for forward + backward + optimizer computation
• $T_{\text{comm}}$: total communication time across all parallelism dimensions
• $T_{\text{overlap}}$: time where communication is hidden behind computation
• $T_{\text{idle}}$: idle time due to pipeline bubbles, synchronization barriers, load imbalance

4.1 Communication Time Analysis per Parallelism Dimension

Tensor Parallelism communication per layer (2 all-reduce operations in forward, 2 in backward):

where $M_{\text{tensor}} \propto b_{\text{mbs}} \times s \times h$is the activation tensor size, and$\beta_{\text{intra}}$is the intra-node bandwidth (NVLink:$\sim 900\;\text{GB/s}$ bidirectional on H100).

Data Parallelism (ZeRO-2) communication per step (one all-reduce of gradients, decomposed as reduce-scatter + all-gather of optimizer states):

where $M_{\text{grad}} = 2P / (D_{\text{TP}} \cdot D_{\text{PP}})$and$\beta_{\text{eff}}$ is the effective cross-node bandwidth.

Pipeline Parallelism introduces bubble overhead rather than bandwidth-limited communication. For a 1F1B schedule with $D_{\text{PP}}$stages and$m$ micro-batches:

The pipeline bubble fraction is:

where $m = D_{\text{DP}} \times N_{\text{GAS}}$ is the total number of micro-batches in the pipeline.

4.2 Throughput Optimization Heuristics (Ordered by Priority)

Priority 1 — Maximize TP within a node:

Since TP communication occurs over NVLink ($\beta_{\text{intra}} \gg \beta_{\text{inter}}$), this minimizes communication latency. However, TP introduces synchronization points per layer, so there exists a diminishing-returns threshold where computation per GPU becomes too small relative to communication.

The compute-to-communication ratio for TP must satisfy:

When $\rho_{\text{TP}}$drops below$\sim 1$, TP becomes communication-bound.

Priority 2 — Scale DP with ZeRO-3 while maintaining $B_{\text{target}}$:

If $D_{\text{DP}} \times N_{\text{GAS}} \times b_{\text{mbs}} \times s = B_{\text{target}}$can be maintained, increasing$D_{\text{DP}}$distributes computation while keeping the batch size constant (by reducing$N_{\text{GAS}}$). This is beneficial because:

Fewer gradient accumulation steps means fewer sequential micro-batch forward-backward passes, reducing $T_{\text{step}}$.

Priority 3 — Transition to PP when DP communication saturates:

When DP communication time $T_{\text{comm}}^{\text{DP}}$can no longer be overlapped with backward computation (i.e.,$T_{\text{comm}}^{\text{DP}} > T_{\text{compute}}^{\text{bwd}}$), introduce pipeline parallelism to reduce $D_{\text{DP}}$:

This trades DP communication overhead for pipeline bubble overhead. The transition is beneficial when:

Priority 4 — Tune micro-batch size $b_{\text{mbs}}$:

The micro-batch size affects multiple performance dimensions simultaneously:

The optimal $b_{\text{mbs}}^{*}$ balances these competing effects and must be found empirically.

5. Benchmarking Thousands of Configurations

5.1 Search Space Enumeration

For a given model size $P$and cluster size$N_{\text{GPU}}$, the total number of valid configurations is:

where $\mathcal{B}$is the set of candidate micro-batch sizes,$\mathcal{G}$is the set of valid gradient accumulation steps, and the indicator function$\mathbb{1}[\cdot]$ filters configurations that satisfy memory constraints and batch size targets.

Even after pruning infeasible configurations, $|\mathcal{S}|$ remains in the thousands across all model sizes and cluster sizes.

5.2 Benchmark Setup and Experimental Conditions

The benchmarks referenced in the content were conducted with:

For each $(P, N_{\text{GPU}})$pair, every valid configuration$\mathcal{C} \in \mathcal{S}$ was benchmarked, and the MFU was recorded.

5.3 Heatmap Analysis and Key Insights

The heatmap visualization plots the optimal configuration $\mathcal{C}^{*}$and its corresponding MFU for each combination of model size$P$and node count$N_{\text{nodes}}$.

Insight 1: Efficiency Decreases with Increasing Node Count (Especially for Small Models)

For a fixed model size $P$, increasing $N_{\text{nodes}}$(and hence$N_{\text{GPU}}$) reduces MFU. The root cause is the arithmetic intensity drop:

When $P$is small,$F_{\text{model}} = 2P$is small, so the numerator shrinks as$N_{\text{GPU}}$increases while$V_{\text{comm}}$does not decrease proportionally. For small models, even increasing$B_{\text{global}}$to compensate is impossible when constrained to$1\text{M tokens}$.

The scaling efficiency can be quantified as:

For small models, $\eta_{\text{scale}}$ decays sharply because the compute-to-communication ratio falls below the threshold for efficient overlap.

Insight 2: Large Models Face Memory Walls on Small Clusters

For large $P$(e.g.,$80\text{B}$) on few nodes (e.g., 4 nodes = 32 GPUs):

This is the absolute minimum (ZeRO-3 with no activations). With activations, the memory exceeds 80 GB, forcing aggressive recomputation and small $b_{\text{mbs}}$, which in turn leads to:

• Poor GPU utilization (small matrix multiplications)
• Large pipeline bubble fractions (many stages, few micro-batches)
• Net result: low MFU despite the model fitting

Insight 3: Implementation Quality Dominates Configuration Choice

The relative performance ranking between TP and PP is not fixed — it depends critically on implementation quality:

Specifically:

• TP with naive synchronous all-reduce: Exposes full communication latency, making TP slower than PP.
• TP with asynchronous communication-computation overlap: Hides communication behind computation via CUDA streams, potentially making TP faster.
• PP with optimized interleaved scheduling (e.g., interleaved 1F1B, zero-bubble schedules): Reduces $f_{\text{bubble}}$from$\frac{D_{\text{PP}}-1}{m}$ toward zero.

The lesson is that $\text{MFU}(\mathcal{C})$ is not a pure function of the parallelism configuration — it is also a function of the implementation:

6. Practical Engineering Challenges in Large-Scale Benchmarking

6.1 Infrastructure Failure Modes

Running thousands of distributed configurations exposes failure modes invisible at small scale:

6.2 CUDA Memory Allocator Behavior

The CUDA caching allocator (used by PyTorch) maintains pools of allocated memory blocks. Fragmentation occurs when freed blocks cannot be coalesced:

This means the theoretical memory computation $M_{\text{total}}$ underestimates actual GPU memory consumption. In practice:

where $\epsilon_{\text{frag}} \in [0.05, 0.20]$ depending on allocation patterns.

6.3 NCCL Communication Overhead

NCCL (NVIDIA Collective Communications Library) implements collectives using GPU Streaming Multiprocessors (SMs). This creates a critical hidden cost:

When communication is overlapped with computation, NCCL kernels consume SMs that would otherwise perform matrix multiplications, resulting in:

This SM contention is the fundamental reason why the assumption “computation and communication can be efficiently overlapped without throughput impact” is violated in practice. The actual throughput during overlap is:

where $N_{\text{SM}}^{\text{NCCL}}$is the number of SMs used by NCCL and$N_{\text{SM}}^{\text{total}}$ is the total SM count (e.g., 132 on H100).

7. Summary — Complete Decision Algorithm

The complete configuration search algorithm can be formalized as:

The entire process is iterative and empirical: theoretical analysis narrows the search space $\mathcal{S}$, but the final optimal $\mathcal{C}^{*}$ must be discovered through systematic benchmarking on the target hardware, as implementation quality, network topology, and GPU-level resource contention introduce performance variations that analytical models cannot fully capture.