Training and Inference Systems › Distributed Training and Sequence Parallelism

Distributed training is the collection of techniques that allow one logical model to be trained using the memory and compute of many accelerators. It exists for two reasons that should be kept separate in your head, because they call for different remedies and are often confused with one another:

• Memory capacity: Parameters, gradients, optimizer state, activations and the input batch must fit. Two fp32 Adam moment tensors for a billion parameters occupy 8 GB decimal, excluding step counters and any master-weight copy. That is not the entire training footprint. Dtypes, optimizer variants and sharding determine what each rank actually stores.
• Throughput and wall-clock: Even when everything fits, dividing work can reduce completion time. A hypothetical day-long run that finishes in one hour across thirty-two workers has a 24-fold speedup, not ideal 32-fold scaling. Throughput affects experiment iteration and cost; under a fixed deadline or budget it can also determine whether a run is feasible.

The distinction matters because the two problems have different fixes. You should identify which tensors cause the memory limit. Tensor and pipeline parallelism partition model computation. Sharded data parallelism can instead partition optimizer states, gradients and parameters while retaining data-parallel work assignment, as described in ZeRO. Activation checkpointing trades recomputation for activation memory. Adding fully replicated data-parallel workers does not reduce the replicated parameter or optimizer footprint.

Trajectory length creates another possible memory limit. As an illustrative input, a five-minute video at 30 Hz contains 9,000 frames; a simulated trajectory could instead contain 10,000 transitions. These are examples, not typical training-window sizes. Partitioning only the batch still leaves each assigned sequence on its device. For suitable architectures, context partitioning distributes token positions and their attention work, while some frameworks use sequence parallelism for selected activation storage alongside tensor parallelism. Neither term licenses ignoring temporal dependencies. We will make those meanings explicit below.

A simple baseline is fully replicated data parallelism: every device holds a full copy of the model, receives its assigned minibatch, computes gradients locally, and then synchronizes gradients before the optimizer step. Mathematically, if worker $i$ has local loss $\mathcal{L}_i(\theta)$ computed from its own shard of the global batch, the goal is that every worker takes the same step using the global gradient

$$g = \frac{1}{N} \sum_{i=1}^{N} \nabla_\theta \mathcal{L}_i(\theta)$$

where:

• $g$: the global gradient estimate, the average of the per-worker gradients
• $N$: the number of workers sharing the global batch
• $\mathcal{L}_i(\theta)$: the local loss computed by worker $i$ from its own shard of the global batch
• $\nabla_\theta \mathcal{L}_i(\theta)$: the gradient of worker $i$'s local loss with respect to the parameters $\theta$
• $\theta$: the model parameters, identical on every worker before the optimizer step

The gradient exchange operation is called all-reduce: every rank contributes a local gradient and receives the reduced result. In ordinary replicated data parallelism, matching initial parameters, optimizer state, hyperparameters and update order let ranks apply the same synchronized update. Gradient synchronization does not synchronize arbitrary mutable model state or repair differing optimizer state. PyTorch DDP organizes synchronization into gradient buckets; it is not necessarily one collective per whole step.

For a loss that is the mean of independent per-example terms, equal-size shards justify the displayed unweighted mean. Unequal shard sizes require weights proportional to the number of contributing examples or valid tokens, with matching loss normalization. This is a mathematical equivalence, not a promise of bitwise equality: floating-point reduction order changes rounding, and batch-dependent operations such as local BatchNorm can change the computation itself. Learning-rate and batch-size choices still require validation rather than following automatically from distribution.

Two details decide efficiency. First, a billion fp32 gradient entries occupy 4 GB decimal locally; this is a payload size, not the bytes sent over the network. For the ideal ring algorithm below, each rank transmits $2(N-1)G/N$ bytes for gradient payload $G$. Other algorithms and compression change that count. Second, bucketing can overlap ready gradient reductions with remaining backward computation. Whether overlap hides much communication depends on bucket readiness, compute duration and network contention; a large model is not enough to guarantee it. Profile the particular encoder and dynamics workload.

Stochastic latent variables add a gradient-estimation question, not a requirement that data-parallel ranks sample identical latents. Independent random streams allow independent draws across ranks; after gradient reduction, parameters can still receive the same update. If the sampling distribution depends on parameters, differentiating a sampled value is not automatically an unbiased estimator. Valid pathwise or score-function estimators need their respective assumptions; straight-through estimators may be biased. Stochastic computation graphs make these dependencies explicit. Sharing seeds can correlate estimates and may increase or decrease variance depending on covariance; it does not generally buy lower variance. Replacing a sampled latent by its mean generally changes a nonlinear stochastic objective, although special cases such as affine costs can agree. Treat mean substitution as a modeling choice rather than a synchronization fix.

When layer or model storage exceeds one device's capacity, model-parallel schemes are possible remedies alongside state sharding or offload. Two useful families are:

• Tensor parallelism (TP): partition operations inside a layer. For a column-partitioned matrix $W=[W_1,\dots,W_k]$ in $Y=XW$, rank $r$ computes the output-feature shard $Y_r=XW_r$. Reconstructing $Y=[Y_1,\dots,Y_k]$ requires concatenation, implemented by an all-gather when every rank needs the full output. A compatible next layer can consume the shards without gathering them first. In a row-partitioned multiply, each rank instead computes a contribution to the same output entries, and the contributions must be summed, for example with an all-reduce or a reduce-scatter. Megatron's linear-layer implementation makes this distinction explicit. Sharding reduces local weight storage, but layer-internal communication can become latency-critical. Its collective count and overlap depend on shard layouts, scheduling and kernels, not simply the number of layers. Fast links are valuable where communication is exposed; profiling determines the useful rank count.
• Pipeline parallelism (PP): consecutive layer blocks are assigned to stages, and independent microbatches flow through them. One stage can process a later microbatch while the next processes an earlier one. Filling and draining leave idle stage slots. For an ideal balanced forward pipeline with $P$ stages, $M$ microbatches, equal-duration stage work, no communication overhead and no overlap between separate batches, each stage does $M$ active slots within $M+P-1$ elapsed slots. Its idle fraction is

$$\text{bubble fraction} = \frac{P-1}{M+P-1}$$

where:

• $P$: the number of pipeline stages (each stage is a contiguous block of layers on its own device)
• $M$: the number of microbatches into which each global batch is split
• $P - 1$: idle slots for an individual stage over the complete forward schedule; the split between filling and draining depends on the stage's position

At fixed $P$, increasing $M$ makes this ideal idle fraction approach zero. It does not eliminate real communication, imbalance or memory costs. Training adds backward work and may use a different schedule, so apply the formula only to a schedule for which its slot accounting holds. The conveyor-belt intuition is useful: more independent items amortize startup and teardown, provided stations can process them at compatible rates.

For a large video encoder, tensor parallelism may address a layer-size limit; deep stacks can provide pipeline partitions. Rollout training still needs a workload trace before choosing either. A temporal transition is not necessarily an optimizer update or synchronization boundary: a learner can form one loss from an entire batch of imagined trajectories. Independent environments or particles can supply microbatches, but causal dependencies inside each trajectory remain. When utilization is poor, measure stage imbalance, communication, memory and schedule bubbles. Larger independent batches are one possible remedy, not a guaranteed cure or a requirement of all imagination-based training.

Long-context training can exceed activation memory even when parameters fit. Partitioning known token positions distributes some of that storage and computation. The relevant length is the actual training context, not the total duration of the source video or a hypothetical imagination horizon. Tokens, frames and environment steps are different units, and there is no family-wide rule that world models use longer contexts than language models. The key question is which operations can be distributed without changing the computation.

The semantics differ by architecture. For a recurrent state-space model with hidden state $h_t$, the recurrence

$$h_t = f(h_{t-1}, x_t)$$

where:

• $h_t$: the hidden state at time $t$
• $h_{t-1}$: the hidden state at the previous time step
• $x_t$: the input at time $t$
• $f$: the deterministic state-transition function (the recurrent core)

has a causal dependency: a shard needs the preceding shard's final state. For an arbitrary nonlinear recurrent core, computing each chunk from zero and combining its final states cannot reconstruct the true trajectory. Parallel scans require a compact associative representation of transition composition. For example, input-conditioned affine maps $h_t=A_t h_{t-1}+b_t$ compose as $(A_2A_1,A_2b_1+b_2)$; whether that representation is efficient depends on matrix structure and whether coefficients can be computed without the previous state. This does not establish an efficient scan for a nonlinear Dreamer-style recurrence. Such a model can batch independent episodes or particles while preserving time order within each trajectory; time partitioning needs actual boundary states, recomputation or an explicitly approximate training scheme.

For a transformer world model, context parallelism partitions token positions. In Ring Attention, queries stay on their owning device while key/value blocks circulate. Blockwise attention combines partial results using online softmax statistics rather than materializing the full attention matrix. Causal masks prevent attending to future positions. Dense attention still has quadratic total arithmetic, while token storage is distributed and temporary score memory is bounded by the chosen block size. Long-context attention over already available inputs does not remove autoregressive dependencies when generating new trajectory tokens. NVIDIA also distinguishes sequence parallelism from context parallelism: the former in Megatron shards selected activations alongside tensor parallelism, while the latter partitions the input context and attention work.

For an ordinary sequential diffusion or flow sampler, each solver step consumes the previous step's state. Assigning successive dependent steps to different devices therefore requires transfers at their boundaries, not necessarily a global collective. One exact baseline keeps solver steps ordered while distributing within-step spatial-temporal token work. Independent samples can also be batched. Alternative parallel or approximate solvers require their own convergence and quality analysis; the sequential dependency alone does not prove that every other distributed schedule is inefficient.

Let's compare the communication patterns to build intuition:

• Replicated data parallel: gradient-bucket collectives per synchronized backward pass, with payload proportional to gradient entries; overlap depends on the workload.
• Tensor parallel: layer-internal gathers, reductions or scatters, determined by partition layout. Count each collective's payload and rounds over the actual forward/backward schedule; some dependencies expose latency while others can overlap work.
• Pipeline parallel: point-to-point activation transfers at stage boundaries for each microbatch, with gradient transfers in backward. For $P$ stages and $M$ microbatches, the basic forward schedule has $M(P-1)$ boundary transfers; their byte counts depend on the boundary tensors. More microbatches create overlap opportunities but do not guarantee hidden traffic.
• Ring context parallel attention: an ideal unpruned forward sweep over $N$ ranks can visit all KV blocks with $N-1$ sends per rank, without a final return exchange. Each block has size proportional to $S d_{\mathrm{KV}}/N$, so this schedule transmits per-rank KV entries proportional to $(N-1)S d_{\mathrm{KV}}/N$, linear in total tokens $S$, not quadratic. Here $d_{\mathrm{KV}}$ is the combined key/value width. Count the actual implementation's exchanges: the paper's Appendix A forward loop also exchanges after the last computation, making $N$ sends. Backward communication and causal pruning need separate accounting.

These patterns trade bytes, synchronization frequency and dependency placement. There is no universal ranking in which data parallelism always moves the most bytes or tensor parallelism always moves less. Model width, context length, microbatch size and rank count determine the comparison. Ring attention's quadratic arithmetic and linear per-rank forward KV traffic are different quantities; chunking bounds working memory rather than converting quadratic network traffic into linear traffic.

Large-model systems can combine several parallelism axes. A combination of data, tensor and pipeline parallelism is often termed 3D parallelism; an additional context or other partition introduces another axis, but these names do not define one mandatory world-model recipe. The mixture depends on interconnect, tensor shapes, sequence length and the measured schedule. A useful heuristic is to put the most exposed communication on the fastest available links. Data-parallel reductions can span slower links when enough backward work or gradient accumulation amortizes their cost; short synchronized updates can be latency-sensitive too.

To build intuition about gradient synchronization, consider an idealized ring all-reduce. A reduce-scatter phase and an all-gather phase each take $N-1$ rounds, with $G/N$ bytes sent per rank in each round. This gives the following communication-only model for $N>1$:

$$T(N) = 2(N-1)\alpha + \frac{2(N-1)}{N}\beta G$$

where:

• $T(N)$: the time to complete the all-reduce collective with $N$ workers
• $\alpha$: latency per communication round in seconds
• $\beta$: the per-byte transfer cost of the interconnect
• $G$: gradient bytes held by each rank before the collective
• $N$: worker count; a one-worker job has no synchronization, so $T(1)=0$

This derivation counts transmitted bytes and rounds, not reduction arithmetic, contention, or compute/communication overlap. Ring algorithms are analyzed by Patarasuk and Yuan (2009). The values below are hypothetical, not measurements of a GPU cluster. Holding global work fixed, we assume computation takes $3/N$ seconds per worker and add the unoverlapped synchronization time. The byte-transfer term approaches $2\beta G$ as workers increase, while the round-latency term grows.