Advanced Distributed Training

What is advanced distributed training?

Advanced distributed training coordinates multiple accelerators so one model-training job can exceed the memory or throughput of a single device. It partitions the batch, model state, layers, sequence, or experts, then uses communication collectives to keep those partitions consistent.

In plain language, more GPUs do not automatically mean faster training. Every device must receive useful work, fit its part of the model, exchange the right tensors, and agree on the next optimizer step. A slow input shard, misplaced tensor-parallel group, missing collective, or incomplete checkpoint can stall or invalidate the whole run.

The core invariant is one coherent global step: every participating rank must derive its update from the intended batch and synchronization contract, or the job must fail and recover without publishing a partial step.

This lesson assumes familiarity with model training and data pipelines. It focuses on choosing a partition, estimating memory, diagnosing communication, and recovering a long-running job.

Trace one distributed optimizer step

1
Input
Assign deterministic work
Partition examples or tokens by data-parallel rank. Every sample should have an intentional owner, including the final uneven batch.
2
Forward + backward
Compute local tensors
Run the local model partition and produce activations, losses, and gradients for the rank's assigned work.
3
Collectives
Exchange required state
All-reduce gradients, all-gather parameters, reduce-scatter shards, or send activations between pipeline stages according to the chosen layout.
4
State
Update and commit progress
Apply a logically consistent optimizer step, advance the data position, and periodically publish a recoverable checkpoint.

The end-to-end step time is constrained by the slowest required path:

Input delivery: one rank waits for data decoding, remote reads, or an imbalanced shard.
Compute: one stage has more work, slower kernels, or a different tensor shape.
Communication: collectives cross a slower link or cannot overlap enough with backward compute.
Recovery storage: synchronous checkpoint writes pause training or saturate the storage path.

Synchronous workers meet at communication boundaries. A single straggler therefore raises tail step time for the group even when average device utilization looks acceptable.

Choose the dimension that removes the real constraint

Partition the batch

Data parallelism

Replicate the model, assign different examples to each rank, and synchronize gradients. It is the simplest option when one complete training replica fits on every device.

Partition model state

Sharded data parallelism

Shard optimizer state, gradients, and eventually parameters across data-parallel ranks. FSDP and ZeRO reduce persistent memory at the cost of all-gather and reduce-scatter traffic.

Partition each layer

Tensor parallelism

Split matrix operations across a tightly connected group. It helps when an individual layer is too large, but introduces communication within the layer's critical path.

Partition model depth

Pipeline parallelism

Place groups of layers on stages and stream micro-batches through them. It reduces per-stage state but creates bubbles when the pipeline fills, drains, or becomes imbalanced.

Additional dimensions solve narrower shapes:

Context parallelism partitions long sequences when activation memory grows with context length.
Expert parallelism distributes mixture-of-experts modules and routes tokens to the owning experts.
Hybrid or 3D parallelism composes data, tensor, and pipeline groups so each expensive communication pattern stays on an appropriate link.

Parallelism degrees are a topology contract. For a dense hybrid job, world size = data-parallel degree x tensor-parallel degree x pipeline-parallel degree x context-parallel degree. The factors must match the launched ranks, and the most communication-heavy groups should stay on the fastest links.

Fit the training state before estimating speed

Use the planner to change model size, accelerator memory, reserve for activations, GPU count, and state-sharding stage. It intentionally estimates capacity, not tokens per second: throughput requires measured kernel, input, and collective performance from the target hardware.

The planner uses a transparent mixed-precision Adam proxy:

2 bytes per parameter for low-precision parameters;
2 bytes per parameter for gradients;
4 bytes per parameter for an FP32 master copy;
8 bytes per parameter for the two FP32 Adam moments.

That is 16 GB of persistent model state per billion parameters before sharding. The selected strategy divides only the states its contract owns, then adds a user-controlled reserve for activations and temporary collective buffers. Real allocations vary with optimizer, precision, tied parameters, allocator fragmentation, activation checkpointing, framework internals, and maximum live all-gathers, so confirm the result with a measured peak-memory profile.

Reproduce the planner's model-state estimate

Read communication as part of the algorithm

The communication primitive determines what every rank owns before and after an operation.

Same result on every rank

All-reduce

Reduce corresponding values from all ranks and return the reduced result to every rank. DDP commonly uses it to synchronize gradient buckets.

Reduced shard per rank

Reduce-scatter

Reduce values across ranks while leaving each rank with only its assigned shard. Fully sharded training uses it for gradients.

Reconstruct from shards

All-gather

Collect each rank's shard so participants can temporarily reconstruct the parameters required for a module's computation.

Neighboring stages

Point-to-point send

Move activations and gradients between pipeline stages. Stage imbalance and too few micro-batches leave visible idle bubbles.

Diagnose a scaling regression in this order:

Measure one rank and the full group. Compare input time, forward time, backward time, collective time, and checkpoint time by rank.
Find the maximum, not only the mean. One slow rank determines a synchronous boundary.
Separate bandwidth from latency. Large collectives stress bytes per second; many small collectives stress launch and network latency.
Inspect placement. Keep tensor-parallel and sharding groups on high-bandwidth links before sending their traffic across nodes.
Change one variable. Test bucket size, micro-batch count, activation checkpointing, or group layout against the same measured workload.

All ranks in an NCCL collective must call compatible operations in the same order. If one rank skips or reorders a collective, peers can wait indefinitely rather than producing a partial answer.

Design recovery as a distributed state transition

A checkpoint is useful only if the training process can reconstruct a valid continuation. Switch scenarios, inspect the active path, and predict what the scheduler may safely do.

A resumable training checkpoint normally preserves the state that influences the next update:

model parameters or recoverable parameter shards;
optimizer moments and parameter-group configuration;
learning-rate scheduler and gradient-scaler state;
completed global step, epoch, and data-sampler position;
random-number-generator state when exact continuation matters;
parallelism layout, software version, and manifest metadata needed to interpret the shards.

Publish checkpoints atomically: write immutable shards, validate expected files and metadata, then commit a small manifest that makes the checkpoint discoverable. Keep at least one previously validated checkpoint until the new one has been restored successfully. PyTorch Distributed Checkpoint can save and load from multiple ranks and supports load-time resharding, but the application still owns its complete training-state contract.

Verify interruption and resume against an uninterrupted run

Review the design with explicit invariants

Build these habits

Start with the least complex strategy that fits, then add dimensions only for a measured memory or throughput constraint.
Keep communication-heavy groups inside a node or fabric island when the hardware topology allows it.
Calculate global batch = micro-batch per rank x gradient accumulation steps x data-parallel degree; tensor and pipeline ranks collaborate on the same examples and do not multiply the batch.
Track useful throughput, step-time percentiles, per-rank idle time, collective duration, input stalls, memory high-water mark, and checkpoint age.
Exercise rank loss, storage interruption, and restore under the intended world size before a long or expensive run.
Revalidate convergence when changing global batch, optimizer, precision, loss scaling, or communication compression.

Avoid these traps

Do not choose a topology from GPU count alone; model shape, link hierarchy, batch constraints, and failure domains determine the layout.
Do not report theoretical FLOPS as training throughput without measuring useful tokens or examples completed.
Do not let a new checkpoint replace the last good checkpoint before its manifest and restore path are validated.
Do not treat elastic worker replacement as exact recovery unless data ownership, optimizer state, and collective membership are reconstructed coherently.
Do not hide stragglers behind cluster averages; inspect distributions and the slowest rank at every synchronization boundary.

Use current primary references

PyTorch DDP design note: gradient buckets, all-reduce ordering, and synchronization behavior.
PyTorch FSDP2 documentation: parameter all-gathers, gradient reduce-scatters, and sharded state ownership.
PyTorch Distributed Checkpoint: multi-rank save/load and load-time resharding.
DeepSpeed configuration reference: the current ZeRO stage definitions and offload controls.
NVIDIA Megatron Core parallelism guide: data, tensor, pipeline, context, expert, and fully sharded parallelism.
NVIDIA NCCL collective operations: collective participation and result semantics.

Framework APIs and supported strategy combinations change. Recheck the versioned documentation for the exact PyTorch, DeepSpeed, Megatron Core, accelerator, and communication-library versions used by the training environment.

No quiz questions available

Could not load questions file