Chapter 14: Distributed Model Parallelism

Parallel mesh design

Construct rank coordinates, process groups, memory timelines, collective tables, and physical placement for one feasible combined mesh

Before you begin

Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.

All prior sections' ownership, memory, shape, and communication rules
Chapter 13 physical topology and route measurement

A cluster contains eight servers, and every server contains eight GPUs connected by a fast local GPU fabric. A target model needs eight-way tensor partitioning to make each layer fit and eight replicas to reach the desired global throughput. The resulting 64 processes need two different communication neighborhoods: frequent tensor collectives among the eight GPUs inside each server and gradient or sharding collectives among corresponding GPUs across servers.

Assign every process a coordinate (data index, tensor index). Rank (3, 5), for example, belongs to tensor group [(3, 0)…(3, 7)] and data group [(0, 5)…(7, 5)]. A device mesh is the coordinate system and the associated process groups. The coordinates describe logical ownership; a separate placement map relates each coordinate to a physical GPU, CPU locality, and network interface.

Print the process-group membership and local tensor shapes at startup. Capture one representative block with PyTorch Profiler or Nsight Systems. The trace should show tensor collectives at layer cadence inside each server and cross-server sharding communication at parameter-group boundaries. Pairwise topology tests and NCCL logs should confirm that each logical group uses the intended links and network rails.

The mesh degrees must multiply to 64, but arithmetic is only the first check. A candidate must also fit peak transient memory, preserve efficient local matrix shapes, place frequent communication on suitable links, maintain the global-batch contract, and support checkpoint restore. A written per-rank memory timeline and collective table make those constraints inspectable before a full training run.

List the constraints before choosing degrees

Calculate the minimum partitioning required for parameters, optimizer state, gradients, activations, and workspaces to fit with headroom. Identify compute-time or context constraints that require additional partitioning even when memory fits. Record global-batch and quality requirements.

A mesh axis should solve a stated constraint. Data parallelism adds replicas for throughput. Full sharding reduces replicated state. Tensor and pipeline axes divide model computation. Context and expert axes divide specialized workload structure.

Define rank coordinates and process groups

Represent each rank by coordinates such as (pipeline, tensor, data). Ranks that differ only along the tensor coordinate form a tensor group. Ranks that differ only along data form a gradient or sharding group. Each group has its own collective sequence.

The product of degrees must equal world size, but group construction must also be consistent. Write concrete rank lists for a small example before scaling. This catches accidental group overlap, inconsistent order, and fixed-rank assumptions.

Calculate local work

For every candidate, derive local matrix and attention shapes. Tensor degree divides selected feature dimensions. Pipeline degree changes the number and type of layers per stage. Data degree changes local batch under a fixed global batch. Context degree changes local sequence work.

Reject a candidate that produces inefficient local kernels unless another constraint requires it and measured end-to-end value remains positive. More ranks can reduce arithmetic per rank faster than fixed launch and synchronization cost.

Calculate memory timelines

List persistent state by axis: which parameters are replicated, tensor-sharded, FSDP-sharded, or assigned to a pipeline stage? Then add transient full parameters, activations, pipeline microbatches, collective outputs, and workspaces on a timeline.

Parallel methods interact. Pipeline stages change which FSDP groups live on a rank. Tensor sharding reduces parameter and activation dimensions. Sequence parallelism changes activation ownership. Do not sum isolated method estimates.

Build a collective table

For each process group, record operation, payload, frequency per layer or step, dependency depth, first-ready time, consumption deadline, and numerical format. Convert payload and rounds into a path-specific alpha–beta estimate.

Mark simultaneous collectives from different axes. Tensor all-reduces, FSDP all-gathers, data reductions, pipeline sends, and expert all-to-alls can share device and network resources. The sum on each physical edge must fit the scheduled interval.

Map logical groups to physical domains

Place high-frequency, latency-sensitive groups in the smallest fast domain with enough ranks. Tensor parallelism usually remains within a node because tensor ranks communicate inside each layer. The within-node placement rule is conditional: a system with a different interconnect hierarchy or a degree larger than local capacity requires another mapping.

Align data or sharding groups across corresponding local positions when they communicate across nodes. Respect GPU-to-NIC affinity and rails. A logical group that alternates nodes and sockets can force traffic through weak paths even when group size is correct.

Schedule the critical path and reliability

Combine local compute, readiness, collective duration, pipeline cadence, and overlap into a rank timeline. Use the maximum stage and rank times. Include checkpoint pauses or interference when the objective is long-duration throughput.

More ranks increase the job's exposure to interruption and can enlarge restart cost. Record failure domains, checkpoint portability, spare-resource requirements, and whether elastic restart can reconstruct the same or a different mesh.

Validate the largest assumptions first

Test local GEMM shapes without the complete model. Test representative process groups and message sizes. Run a few layers with the intended overlap and memory schedule. Shape, communication, and memory experiments eliminate bad candidates before an expensive full-scale run.

Then scale in stages and compare traces with the predicted schedule. Revise parameters rather than adding unexplained correction factors. Preserve one rejected mesh and the measurements that rejected it.

Worked example: Sketching a 64-GPU mesh

Eight nodes each contain eight GPUs with a fast intra-node GPU fabric and two local NICs. A 64-GPU model requires tensor partitioning for layer capacity and sharded data parallelism for optimizer and gradient state.

Define coordinates (data=8, tensor=8). Each node contains one tensor group of eight GPUs. Each data or shard group contains the same local tensor coordinate from all eight nodes.
Calculate tensor-local shapes. For a projection with global output width N = 8,192 and tensor degree T = 8, each rank receives N/T = 1,024 output features. Benchmark these shapes and compare with T = 4, which produces 2,048-wide output shards.
Calculate persistent and transient memory for tensor degree eight with eight-way full sharding. Include gathered FSDP groups, tensor-partitioned activations, communication buffers, and operating headroom.
Create the collective table. Tensor collectives occur inside every layer on local fabric. FSDP all-gathers and reduce-scatters cross nodes for corresponding tensor coordinates. Assign the eight groups across the two NIC rails according to GPU affinity.
Schedule one representative block. Check whether local tensor collectives and cross-node prefetch overlap or serialize through shared device resources. Measure peak memory during the overlap.
Evaluate the tensor-degree-four alternative. The alternative needs another axis or more replicated capacity, but its local GEMMs can be more efficient. Compare complete step-time and memory models rather than selecting the maximum local degree automatically.
Run two nodes, then four, then eight. At each scale, compare message curves, rank arrival skew, local kernels, and checkpoint behavior with the prediction.

Conclusion: The candidate mesh is supported only if the measured peak stays below the memory budget with headroom, the 1,024-wide local matrix shapes meet the declared kernel-efficiency threshold, the printed process groups use the intended local fabric and NIC rails, and the measured step timeline agrees with the predicted collective order and critical path. The memory, kernel, topology, and timeline observations replace an undefined claim that the mesh is simply good.

Common errors

Choosing axis degrees only because they multiply to world size. Capacity, local shapes, communication, and topology must all agree.
Mapping rank numbers without hardware coordinates. Logical adjacency does not imply GPU, socket, NIC, or fabric locality.
Modeling each parallel axis independently. Memory lifetimes, collectives, and resource contention interact in the executed schedule.
Accepting a mesh because it fits memory. A configuration can fit while narrow kernels, exposed collectives, or recovery cost make it slower.

Reference summary

A parallel mesh assigns each rank coordinates along selected axes such as data, shard, tensor, pipeline, context, and expert parallelism. Each axis defines state ownership, local compute shapes, collective groups, and communication frequency. The product of axis degrees equals the participating world size, but that arithmetic does not determine whether the mesh fits or performs well.

Capacity: calculate persistent and transient state for every rank coordinate and leave headroom for workspaces and communication buffers.
Compute: derive local GEMM and attention shapes and confirm that the selected degree preserves efficient kernels.
Bandwidth: calculate bytes and frequency for every process group, then place frequent and latency-sensitive groups on suitable links.
Critical path: schedule readiness, collective duration, overlap, pipeline bubbles, and the slowest-rank effect instead of adding all durations.
Global semantics: record global batch, sequence partition, routing behavior, reduction format, and quality requirements.
Reliability: account for the number of participating ranks, failure domains, checkpoint layout, and replacement requirements.
Operation: prefer a configuration that can be tested, diagnosed, checkpointed, and recovered within the required completion objective.

Knowledge check

Why should tensor parallelism usually remain inside a node before data parallelism does? State the assumptions behind the rule.

Show the answer guide

Tensor-parallel collectives often occur inside every transformer layer and lie directly between dependent matrix operations, so their frequency and latency sensitivity favor the fastest small physical domain.
Data-parallel or sharding communication often occurs at bucket or parameter-group boundaries and can sometimes use larger messages with more overlap, making cross-node placement less harmful.
The rule assumes faster intra-node links, enough local GPUs for the tensor degree, compatible local matrix shapes, and no alternative fabric hierarchy; topology and measured critical paths can reverse it.

Practice

Compare two parallel meshes

Design two feasible meshes for the same declared model and world size. Hold global batch and quality semantics fixed, then map every process group to physical GPUs and links.

Evidence to keep: Rank coordinates, group lists, local shapes, peak memory timelines, collective bytes and readiness, topology mapping, predicted step schedules, and one measured slice that rejects or supports each candidate.

Validate a proposed parallel mesh with distributed measurements. Compare the measured critical path with the predicted computation and communication times.

References

Sources and further reading

The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.