Distributed Training and Parallelism

Start with ownership and shapes

A training mesh can combine these dimensions. Process groups then have different members and responsibilities. A local rank number is not enough to describe ownership; record named mesh axes and tensor dimensions.

Data parallelism and the global gradient

Ordinary DP keeps a full model replica on every rank. Each replica processes different examples. A synchronized update combines their gradient contributions so all replicas remain consistent.

Unequal valid-token counts require an explicit normalization. If the framework averages gradients across R ranks, a local summed loss can be scaled by R/global_count. Check the framework’s reduction convention before applying this rule.

Gradient accumulation increases work between synchronizations. Suppress intermediate reductions when the framework permits it. Preserve loss scaling across accumulated microbatches. In DDP, wrap intermediate forward and backward passes in ddp.no_sync(); the final pass synchronizes accumulated gradients. A synchronization context around backward alone does not suppress the forward pass’s synchronization setup.

DDP autograd hooks make gradient buckets ready during backward. Bucketed reductions can overlap communication with later backward computation. Small buckets start early but pay more startup costs. Large buckets reduce startup overhead but may leave a large communication tail. Measure the critical path rather than only total bytes.

Matching a single-device update requires matching initial weights, optimizer state and loss normalization. Random operations and floating-point reduction order can prevent bitwise equality even when the mathematical update agrees.

The statistically useful global batch size and the hardware-efficient local batch size are different constraints. Adding replicas shrinks local work if global batch stays fixed. Increasing global batch instead can change optimization. There is no universal “512 GPUs” boundary.

ZeRO and sharded state

DP replicates state that does not need a full copy on every rank at every instant. ZeRO progressively partitions that state. Stage 1 shards optimizer state. Stage 2 also shards gradients. Stage 3 also shards parameters.

This example assumes BF16 model weights and gradients, an FP32 master weight and two FP32 Adam moments. The moments alone require eight bytes. The twelve-byte column includes the master weight. Different optimizers and gradient dtypes change the calculation.

These totals exclude activations, temporary gathers, communication buffers and allocator overhead. A model can fit by the persistent-state formula and still run out of memory during a layer gather.

FSDP is related to full parameter sharding, not an identical implementation of every ZeRO-3 schedule. It gathers a unit’s parameters before use and reduces gradients to their owners. Resharding after forward saves memory but can require another gather for backward.

Compare schedules, not stage numbers alone. A common full-shard step has forward and backward parameter gathers plus a gradient reduce-scatter. Holding parameters longer changes memory and communication. Prefetching can hide transfers while increasing peak live buffers.

Follow one layer across two FSDP ranks → The timeline separates weight gathers, local-batch gradients, ReduceScatter and local optimizer updates.

Count the whole step, not one phase

Let V be the logical bytes of all model parameters. Assume gradients use the same bytes and use a bandwidth-dominated ring on R ranks.

This schedule moves about 50% more bytes than DP. ZeRO-1/2 schedules can instead retain roughly DP’s traffic volume. Different dtypes, retained parameters, accumulation and offload change the comparison. Equal gather-plus-scatter and all-reduce volumes describe one phase; they do not include an extra forward gather.

Tensor parallelism: pair the projections

Consider an MLP XW_up, followed by a pointwise nonlinearity and W_down. Partition W_up by output columns. Every rank computes a different set of intermediate features. The nonlinearity can remain local.

Partition W_down by its matching input rows. Each rank produces a partial contribution to the final hidden vector. Sum those partial outputs with an all-reduce, or reduce-scatter if the next operation wants a sharded result.

For SwiGLU, gate and up projections must share matching feature partitions. For attention, ranks can own different heads and then sum output-projection contributions. GQA’s small KV-head count can limit a naive head partition; MQA has only one KV head. Implementations may replicate KV heads or use another layout.

TP creates frequent layer-level communication. Fast low-latency links often matter more here than for large overlapped DP reductions. Excessive TP can also produce small, inefficient matrix multiplications. Eight-way TP is a common configuration, not a law.

Sequence and context parallelism

SP and CP can both split a sequence dimension. Their names do not make them the same operation.

In the common TP-associated SP layout, each rank keeps some token rows with their full hidden width. Token-wise normalization and dropout run locally. An all-gather restores token rows before a projection that needs them. A reduce-scatter can restore sequence ownership after the paired projection.

SP reduces replicated activation storage around those boundaries. A token-wise activation of shape $[B,S,H]$ becomes $[B,\frac{S}{p},H]$ on each of p ranks. That activation uses about 1/p of the replicated storage when S divides evenly. This is not a 1/p reduction of all training memory.

An all-gather plus reduce-scatter has roughly one all-reduce’s ring byte volume for the same logical tensor; startup latency and overlap can still differ.

SP does not, by itself, solve attention across a partitioned long context. CP assigns different query positions to ranks while preserving the context each query can access.

A ring-attention CP implementation keeps query blocks on their owners and circulates key/value blocks. Rank i sends to rank i+1 and receives from rank i−1, with wraparound.

Asynchronous sends can overlap independent attention work; insufficient local compute or shared-link contention can still expose the transfer.

Moving KV blocks solves ownership. Combining their results must still preserve one softmax over all permitted keys. A block with more probability mass must contribute more, even when blocks have equal sizes.

For one query row, $s_j=\frac{q^\top k_j}{\sqrt{d_{\rm head}}}$ is the scaled score for visible key j. The query and key head width is d_head; v_j is the value vector. The running maximum m sets the exponential scale. The scalar ℓ sums scaled exponential weights; the vector u sums those weights times values. The incoming block has maximum m_block. Primes mark the state after adding that block.

If a larger score arrives, the old sums use the wrong scale. Multiply both old sums by the same factor before adding the new block. The represented output stays unchanged by that rescaling.

In the example, the shared maximum is ln(3). Block B contributes weight $e^{0-\ln3}=\frac13$. The accumulated denominator is 4/3, the numerator is 8, and the output remains 6. Processing B first gives the same result in exact arithmetic. Finite-precision summation can differ slightly with block order. See FlashAttention, Appendix B.3 ↗ for the blockwise update and Ring Attention, Section 3 ↗ for distributed block ownership.

Initialize an empty row with ℓ=0 and u=0. Skip fully masked blocks; use a guarded first-valid-block update rather than evaluating exp(−∞−(−∞)). A row with no valid keys needs an explicit output policy, not division by zero.

Causal attention only sees permitted earlier positions. Early and late query chunks have different work. Zigzag assignments can pair them to balance ranks while retaining original positions and visibility.

Why pair early and late queries?

Consider sixteen query positions, numbered 1–16, including each position’s own key. Query q has q valid keys. Eight ranks each own two queries.

Contiguous ownership ranges from 3 to 31 valid pairs per rank. Zigzag ownership gives every rank 17. Both assignments perform the same 136 valid pairs. Real kernels operate on blocks, with padding, triangular tiles and transfer costs; equal pair counts do not guarantee equal times. Query positions never change, so the causal mask remains intact.

Local-window layers exchange only the reachable neighborhood. Recurrent attention can require an associative state scan instead of circulating all KV.

CP’s benefit depends on communication volume, kernel overlap and sequence layout. It does not make arbitrary long contexts cheap. Attention: masks, recurrence and online softmax →

Pipeline parallelism and idle time

PP assigns groups of layers to different stages. A microbatch sends activations forward and activation gradients backward. Model state becomes smaller per stage, but the stages are dependent.

A pipeline must fill and drain. For an ideal equal-cost model with p stages and m microbatches, distinguish idle overhead relative to useful work from idle fraction of elapsed capacity.

GPipe’s all-forward/all-backward schedule retains many activations. A 1F1B schedule alternates forward and backward work after warmup, often reducing live activation storage. It does not remove every bubble.

Interleaved schedules give each physical stage multiple virtual layer chunks. With v virtual chunks per physical stage, ideal balanced interleaving can reduce the non-interleaved bubble by roughly 1/v. This assumes a compatible microbatch schedule and ignores transfer overhead. Smaller chunks send messages more often; 1/v is not a general measured speedup. Zero-bubble schedules can defer weight-gradient computation after the urgently needed activation gradient; this trades dependencies, memory and scheduling complexity.

What crosses a stage boundary?

A stage sends activations forward, then receives their gradients during backward. It does not send every parameter for each microbatch.

Sequence or tensor sharding changes the boundary’s local shape. Check whether adjacent stages agree on ownership; a hidden gather or redistribution can erase expected savings.

Full-shard parameter gathers can repeat across pipeline microbatches. Keeping parameters live avoids some gathers but consumes memory. Choose the PP and sharding schedule together.

Equal layer counts do not imply equal stage times. Embeddings, output projection, loss and expert imbalance matter. The slowest stage constrains throughput. Increasing microbatch count at fixed global batch also makes each microbatch smaller, which can hurt matrix efficiency and overlap.

PP is especially useful for capacity and aggregate throughput. A single dependent request does not become faster merely because it crosses more stages. In inference, each stage still owns the KV cache for its layers.

Expert parallelism

EP assigns experts to ranks. The router selects experts for each token, so token representations must move to their owners. Experts compute locally. Outputs then return to the token’s original owner for combination.

The forward pass therefore commonly contains a dispatch all-to-all and a return all-to-all. Token counts vary by destination. The slowest loaded rank can determine the layer’s completion time even if average FLOPs look low.

Load-balancing objectives, capacity policies, dropless routing and redundant expert copies change both learning and traffic. Redundancy needs a logical parameter owner or correctly synchronized replicas; copying experts without an update policy changes the model.

Grouped GEMMs improve execution of many small expert batches. Dispatch/computation overlap helps only when resources and buffer lifetimes permit it. Top-k routing multiplies token-expert assignments, not necessarily every tensor byte in the same way.

Do not use a universal E/k performance multiplier. Expert width, token count, load distribution and network placement determine the actual compute-to-communication ratio. Mixture of Experts: routing, balancing and stability →

Infer communication from matrix layouts

Write C[I,K] = A[I,J]B[J,K]. J is contracted. I and K survive in the output. A sharding decision can split an output dimension or split the sum itself.

Collective Communication animates each pattern with explicit senders and receivers. All-gather removes a partition by collecting shards. All-reduce sums partial contributions and replicates the result. Reduce-scatter sums contributions and leaves different result shards on different ranks. All-to-all redistributes ownership, often switching which tensor dimension is partitioned.

The backward operation must respect the transpose of the forward linear communication. An all-gather’s adjoint is a summing reduce-scatter; a summing reduce-scatter’s adjoint is an all-gather. If the implementation averages instead, account for that scaling explicitly.

A shape is not an ownership description

Use a subscript for a named mesh axis. The global shape stays unchanged; the local shape depends on which dimensions that axis partitions.

A batching dimension is retained in both inputs and the output. For example, each batch in C[b,I,K] = A[b,I,J]B[b,J,K] has its own product; b is not contracted. Two operands can have correct local dimensions yet incompatible ownership. In the table above, sharding I and K over the same axis leaves only diagonal rank-pair tiles, not the full product.

For a large ring collective on R ranks and a logical tensor of V bytes, each rank transfers approximately (R−1)V/R for all-gather or reduce-scatter. An all-reduce uses roughly twice that volume. Latency and topology still affect elapsed time. The Scaling Book’s ideal bandwidth-bound bidirectional ring gives about one quarter of the all-gather time for an all-to-all layout exchange. This is not a universal collective law. Its bandwidth convention differs from the unidirectional ring volume estimate above.

# Contract-axis sharding, two ranks, matching slices.
partial = A_local @ B_local  # [I, K], not the final product
dist.all_reduce(partial, op=dist.ReduceOp.SUM)
# Every rank now has C. Use reduce_scatter if C should stay sharded.

The animation above performs this exact calculation. Named dimensions and local shapes often reveal a missing collective before a distributed run starts.

Map communication onto the machine

Separate a tightly connected scale-up domain from a slower scale-out network. A scale-up domain is not always identical to one host. NIC placement, rails and oversubscription determine which rank pairs can communicate efficiently.

Frequent TP collectives often belong on fast local links. EP benefits from bisection bandwidth because many destinations are active together. DP/FSDP can sometimes tolerate scale-out links through larger work units and overlap. PP often sends smaller activation payloads across stages. These are starting hypotheses, not mandatory mappings.

Hierarchical collectives can exploit node-local bandwidth before crossing the fabric. Network-assisted reduction systems such as SHARP can accelerate supported reductions, but it does not accelerate arbitrary all-to-all routing automatically.

One accelerator link, summed scale-up links per accelerator, summed NIC bandwidth per node and achieved application bandwidth are four different measurements. A collective can achieve less than the advertised hardware rate. Benchmark the real group sizes and payloads.

Two operations can overlap in the schedule yet compete for the same link. CPU offload, checkpoint I/O, DP reductions and EP transfers all create contention. Measure node egress, fabric bisection and the slowest rank, not only average accelerator utilization.

How much work can hide communication?

Communication becomes harder to hide when local compute shrinks. Use consistent units: C in FLOP/s, W in byte/s, and communication volume in bytes.

The parameter count cancels in this simplified DP threshold. Local work still has to be large enough relative to network arithmetic intensity. Attention at long context, sparse MoE compute and different precision can invalidate the dense approximation.

FSDP moves weights. TP moves activations. Their costs respond differently to a mixed mesh. Let $N=XY$ devices, with X the FSDP degree and Y the TP degree. In a simplified forward MLP model, the two traffic terms vary approximately as $X$ and $\frac1X$ after substituting $Y=\frac NX$.

The expression is a continuous candidate, not a deployable mesh. Round to valid hardware divisors and benchmark nearby configurations. The derivation assumes an MLP-dominated forward pass and a particular topology model. It does not predict every GPU cluster or long-context model.

Why mixing the axes helps

At the balanced split, doubling global token count improves this compute/communication ratio by only a square-root factor. Pure DP’s simplified ratio grows linearly with local token count. Mixing axes can support smaller local batches, but does not remove the communication limit.

For TP alone, a common rough limit is $Y<\frac{FW}{C}$, with byte and FLOP constants understood. Too many tensor shards create less local math without removing every activation collective. Preserve those constants when comparing dtypes or hardware.

Collective-GEMM overlap works tile by tile. It requires available compute, communication resources and enough local tile work. Double buffering also requires enough memory and correct buffer lifetimes. Merely launching transfers asynchronously does not prove that their cost disappears.

Architecture changes the distributed system

Model architecture changes the best distributed layout. These reported configurations are examples, not prescriptions. Follow the linked recipe sections for their conditions and evidence.

These examples explain why architecture and distributed execution need joint design. Recurrent attention, cross-layer state sharing and memory tables create different dependencies from an ordinary independent stack of dense blocks. A mesh that worked for one model is not a plug-in recipe for the next.

Complete source reports: deepseek v4 1 flash → · gemma 4 technical report → · glm 5 agentic engineering → · kimi k3 open frontier intelligence → · mai thinking 1 building a hill climbing machine → · mimo v2 6 scaling rl towards self improvement → · the llama 3 herd of models →.

Reproduction and debugging

1. Ownership. Record global/local shapes, mesh axes, replicated values and logical parameter owners.
2. Exactness. Compare a tiny distributed step against an unsharded reference, including masks and global loss normalization.
3. Updates. Verify tied/shared parameters, optimizer states, gradient sums and accumulation scaling.
4. Memory. Measure persistent state and peak temporary buffers separately.
5. Scheduling. Check dependencies, buffer lifetimes and slowest-stage/rank timing.
6. Communication. Measure startup, achieved bandwidth, overlap and critical-link contention.
7. Recovery. Verify checkpoint ownership, resumability and mesh reconfiguration.

Start with two ranks and tiny matrices. Many errors are layout or normalization mistakes rather than kernel failures. Inspect the compiler/runtime trace for hidden reshards between regions. Start the mesh design from a statistically justified global batch, not the device count. Increase scale only after the numerical reference and ownership model agree.

Sources and connected notes

Scaling Book: sharded matrices ↗ · Scaling Book: training ↗ · Scaling Book: GPUs ↗ · Scaling Book: transformer costs ↗

ZeRO ↗ · Megatron-LM ↗ · Reducing activation recomputation ↗ · GPipe ↗ · PipeDream ↗ · PyTorch DDP: synchronization ↗