Distributed Training and Inference

Distribution goals and execution dependencies

Distribution serves three distinct goals: fitting required state, finishing fixed work sooner, and processing more independent work. Communication and waiting can offset saved computation. A capacity improvement therefore need not improve request latency. [1]

A tensor is a multidimensional array. Parameters are learned values; activations are intermediate results. Forward computation produces predictions, backward computation produces gradients—the derivatives used for updates—and an optimizer changes parameters using those gradients and retained optimizer state. Distribution must preserve that chain. Training objectives and device execution supply the underlying context. [6]

Prefill processes a known prompt; autoregressive decoding selects successive tokens conditioned on earlier selections. Processing a selected token produces new internal representations; the token itself is not a key or value vector. Distributing stages preserves these dependencies rather than removing them. Prefill and iterative decoding explains the local execution. [7]

State ownership and per-device memory

A replica is a complete copy; sharding assigns separate portions to different owners. Distributed memory requires communication to access remote values: aggregate capacity is not automatically a shared address space. [1]

Activation recomputation, also called activation checkpointing, regenerates intermediates during backward computation instead of retaining all of them. It exchanges computation for memory; it does not save a durable restart point. Offloading selected intermediates to CPU memory adds another transfer-and-prefetch tradeoff. [13][14]

In serving, profile non-cache demand before assigning remaining memory to request state. Weights fitting alone does not establish useful concurrency. KV state and the resident memory budget explains the cache; distributed placement adds ownership and transfer constraints. [15]

Workers, groups and collective communication

A worker is an executing process. Its rank identifies it within a participating group; a communicator supplies that group's communication context. A collective involves the group under a shared operation contract. NCCL implements such communication for participating GPU devices. [16]

NCCL participants must supply matching counts and datatypes. Violations can hang, crash or corrupt data. Collective definitions specify the ownership changes. [17]

MPI's nonblocking contract separates initiation from local buffer-safe completion. Participants initiate collectives in matching order; local completion need not mean every peer finished. GPU frameworks require their own completion rules. A point-to-point transfer instead connects a particular sender and receiver. [18][19]

Physical topology and communication cost

A node is a machine; an interconnect carries data between devices or machines. Latency measures startup delay, while bandwidth measures bytes transferred per time. Shared interfaces and competing traffic can reduce effective bandwidth. [1]

Collective algorithms determine communication rounds and paths. Rings circulate pieces through participants and can use bandwidth efficiently, but startup costs grow with group size. Trees reduce communication depth. A hierarchy can reduce-scatter locally, all-reduce between corresponding shard owners across nodes, then all-gather locally. NVIDIA's historical NCCL report illustrates these choices; no algorithm is universally fastest. [20]

Frequent small exchanges can dominate decoding even when aggregate memory increases. Alex Cheema describes this capacity-versus-speed problem in a local tensor-parallel cluster. Advertised link bandwidth alone cannot predict the accumulated synchronization delay. [21]

Data-parallel updates and loss normalization

Synchronous data-parallel workers start from matching parameters, compute local gradients, aggregate them, then update. Matching optimizer state and synchronized gradients preserve matching replicas. DistributedDataParallel, or DDP, performs gradient synchronization; it does not make arbitrary input weighting correct. [2]

Two tokens with mean loss 1 and six with mean loss 3 yield token mean 2.5, not worker mean 2. DDP averaging requires compensation. [22]

# Equal iteration counts across ranks; no implicit loss scaling.
window = next_local_window()
N = sum_all_ranks(count_scored_tokens(window))
optimizer.zero_grad()
if N > 0:
    for i, batch in enumerate(window):
        with synchronize_gradients(i == len(window) - 1):
            loss_sum = model.summed_token_loss(batch)
            (world_size * loss_sum / N).backward()
    optimizer.step()

Use actual counts for partial windows. Framework accumulation scaling needs separate compensation. Sampling and loss contribution explains the objective. [22]

Equal partition lengths do not ensure exactly-once coverage. PyTorch's DistributedSampler can pad with repeated indices or drop a tail. Shared shuffle seeds and set_epoch control coordinated reshuffling. Check sample identities, not just counts; a rank cannot simply stop participating while peers continue ordinary collectives. [23][18]

Sharded training state and temporary materialization

Fully sharded data parallelism temporarily gathers parameters for computation. Layer-level grouping limits materialization; deeper prefetch can overlap communication with computation, but increases residency. Workers still execute model operations on different inputs. [8]

Partitioning persistent state does not eliminate activation pressure. Long-context training can still exhaust memory after parameter sharding. Smaller activation buffers or recomputation address a different allocation category; adding more parameter shards alone may leave the limiting peak unchanged. [10]

Tensor partitions and result reconstruction

For X=[1,2] and W=[[1,−1],[2,1]], contracted partitions produce [1,−1] and [4,2]. Their sum is [5,1]. Applying ReLU, which replaces negative values with zero, before summing produces [5,2] instead: the reconstruction boundary matters. [3]

Megatron pairs a column-partitioned first matrix with local elementwise activation and a row-partitioned second matrix. Intermediate features stay local; second-matrix partial outputs reduce. Backward likewise combines contributions to shared inputs. Its original MLP uses one forward and one backward all-reduce, not a universal count. [3]

Smaller local work must repay repeated communication. In ordinary decoding, the next prediction waits for the current model execution, exposing collective latency along each token's path. [21]

Pipeline schedules and bubbles

Pipeline stages exchange forward activations and backward gradients; microbatches overlap. Fill-drain runs all forwards, then all backwards, retaining activations until used. Synchronous one-forward-one-backward (1F1B) warms up with forwards, alternates forward/backward, then drains remaining backwards. Earlier backward completion releases activations, bounding outstanding microbatches by pipeline depth. Both retain dependency-induced idle intervals—pipeline bubbles—and update parameters only after all backwards, preserving consistent versions. [19]

Unequal prompt lengths can unbalance stages. Increasing microbatch count cannot remove a persistently slow stage or transfer bottleneck. During inference, independent requests can overlap across stages, while each autoregressive request still depends on its preceding token. [24][7]

Composed parallelism and device placement

A device mesh organizes logical groups; it does not certify physical links. This eight-device assignment uses two data replicas, two pipeline stages and two tensor workers per stage. [25]

Tensor groups are {0,1}, {2,3}, {4,5}, {6,7}; pipeline groups are {0,2}, {1,3}, {4,6}, {5,7}; data groups are {0,4}, {1,5}, {2,6}, {3,7}. Each varies one coordinate while fixing the others. [25]

Corresponding parameter slices repeat across data replicas unless further state sharding partitions them. A separate replicate-by-shard layout can shard within each host and replicate across hosts. State ownership follows the selected group, not the total device count. [9][25]

First reject layouts whose peak state cannot fit. Then compare local computation, collective frequency and stage balance on the actual links. Keeping frequent exchanges within a node is a useful candidate, not a rule. In serving, fewer devices per executable instance may leave room for more independent replicas. [20][5]

Overlap, readiness and the critical path

The critical path is the dependency chain determining completion. DDP groups gradients into buckets and starts reduction when a bucket is ready, preserving collective order. Independent backward work can continue before the optimizer consumes synchronized gradients. [2]

Small buckets invite more startups; large buckets can become ready too late. Bucket contents that poorly match gradient readiness delay launch. Measure exposed waiting rather than treating every communication interval as added elapsed time. [2][20]

A straggler is a participant that finishes later than its peers. Variable generation lengths make synchronous sampling wait for the longest response while completed requests leave capacity idle. Removing that barrier changes coordination semantics; asynchronous reinforcement learning additionally introduces policy-staleness constraints. [26][27][28]

Aligned worker traces should separate readiness, transfer and dependent execution. A late collective completion can reflect a late producer, not merely a slow network. Timeline-led diagnosis and span boundaries explain how to preserve those distinctions. [2][29]

Serving replicas, routing and request state

A serving replica is a complete executable model instance, possibly spread across a tensor-and-pipeline group. Routing chooses an instance; communication inside that instance executes the request. One shard cannot independently serve the whole model. KV capacity still limits concurrent state after weights fit; a logged concurrency estimate is not measured sustainable throughput. [5]

Request affinity favors workers holding reusable state. Cache-aware routing balances that benefit against projected prefill and decode load: a colder but less busy destination can finish sooner. Cache creation and release events inform locality estimates. Destination selection and movement of KV state remain separate responsibilities. [30]

Prefix reuse saves processing for compatible repeated token prefixes, not semantically similar text or future output tokens. Cache compatibility and local batching remain relevant inside every distributed instance. [31]

Admission control bounds accepted work against available resources. Request count alone hides differences in computational cost. Preserve explicit rejection behavior when capacity is exhausted; queueing and sustainable capacity develops this boundary. Operational ownership belongs in AI Platform Engineering. [32]

Prefill–decode separation and handoff

Separate prefill and decode pools can reduce scheduling interference and support different resource allocations. Short prompts offer less interference to remove; excessive prefill supply queues behind decode, while excess decode capacity waits idle. The operating point and request-length distribution determine whether separation helps. [33][34]

DistServe transfers prompt KV state and the first generated token to decoding workers. Transfer demand depends on prompt length and arrival rate. [24]

The network can erase phase-specific advantages. In a heterogeneous local demonstration, cache traffic accidentally used Wi-Fi, leaving transfer exposed before decode. The intended Ethernet connection was a setup choice, not a universal sufficient bandwidth threshold. [35]

A usable handoff needs compatible model execution, cache layout, positions and destination storage. Matching registered service names does not establish those properties. Dynamo's documented route returns prefill or handoff failure after dispatch rather than automatically falling back to aggregated execution. Local scheduling supplies the colocated alternative. [36]

Correctness across partition boundaries

A bounded reference comparison starts with cloned parameters and optimizer state, identical examples and masks, and one intended update. Compare intermediate quantities as well as final parameters. A matching scalar loss can conceal a different gradient or input allocation. [22][2]

Floating-point operations are not associative, so reordered reductions can differ numerically. Choose tolerances from numerical type, value scale and a trusted reference; no universal threshold excuses arbitrary error. Compare fixed-context inference outputs separately from sampled text. Numerical correctness explains the testing discipline. [37]

Useful scaling and measurement boundaries

Strong scaling increases devices for fixed total work. Weak scaling increases total work while holding work per device fixed. These answer different performance questions. [1]

Training throughput does not determine time to useful quality when batch size or learning behavior changes. MLPerf Training measures wall-clock time to a specified quality target and repeats runs. Retain that distinction when comparing your own corpus, objective and training recipe. [39]

Serving goodput counts completed requests meeting declared latency criteria per benchmark second. Report offered load, errors and latency distributions separately; timely output can still be incorrect. Define first-content and continuation timing precisely. Serving experiments explains workload controls; cost and performance engineering extends the economic comparison. [40]

Busy devices need not perform useful arithmetic efficiently. Krea used tensor-core utilization as a diagnostic proxy beyond GPU busy time. Neither replaces end-to-end completion or quality measurements. [41]

Failure propagation and coordinated shutdown

A failure domain contains components exposed to a common failure. Missing participants can prevent collective completion and disable a cooperating model group. Independent replicas can continue only if their required devices and links remain usable. NCCL documents asynchronous communication errors and communicator abort/recreation; recreating communication does not restore lost application state. [16][42]

Separate slow progress from process or device loss, broken communication, exhausted memory and numerical corruption. An elapsed timeout identifies missing progress, not its cause. Krea's experience shows why device counters alone are insufficient: fabric errors exposed faults despite apparently healthy GPUs. Failure localization requires competing explanations and per-worker evidence. [43][16][37]

Recovery policy is implementation-specific. Under torchrun's documented elastic policy, worker failure stops and restarts the worker group within its restart limit. Membership changes also form a new group with new ranks and world size. Surviving processes are not a substitute for checkpointed progress. [44]

Gang scheduling admits cooperating workers together. It addresses resource availability before execution, not restoration after failure. Krea separated queued-workload priority from execution priority and encountered stale capacity quotas as nodes changed. [45]

Coherent checkpoints and training restart

A distributed restart needs one coherent training state, not unrelated shard files. A checkpoint manifest identifies saved tensor pieces. Distributed Checkpoint coordinates planning, rank writes and coordinator completion; loading writes into destination-allocated storage and can reshard tensors. Training continuity defines the resumable state beyond weights. [46]

Asynchronous saving separates staging into CPU memory from storage upload. Protect source parameters until staging completes, including before the next optimizer mutation; other mutable training state needs equivalent protection. Upload completion is a later boundary. Await it before another save or final cleanup. Background staging also consumes CPU capacity and memory. [47]

In the inspected fresh-directory filesystem path, synchronized shard writes precede final metadata publication. The reader requires .metadata. Final metadata is not independent shard-integrity verification. Overwriting an existing checkpoint lacks an atomic old-or-new guarantee, and the path lacks parent-directory synchronization needed to establish stronger crash-durability claims. [48]

Reported rank-write failures prevent coordinator finish. A vanished participant is outside that ordinary gathered-exception path. [49]

Resharding changes storage ownership, not automatically the intended next update. Preserve effective batch size, loss weighting and input progress when worker count changes. Validate restoration through the next update, not merely successful tensor loading. Compare checkpoint overhead, lost work and actual restart time together. [46][44][50]

Serving restoration and interrupted requests

A deployment can require whole-group replacement after losing a model shard. Its re-entry policy should explicitly require model loading, communication initialization, warmup, executable checks and routing registration. These are conditions to implement and test, not one universal Dynamo recovery sequence; configured grouping and controller policy determine replacement scope. [51][53]

Model artifacts, reconstructible KV state and client-observed output have different recovery boundaries. A service can fail an interrupted request or replay retained context. Planned shutdown first stops admission and can drain existing work. Replacing capacity does not by itself complete either policy for requests already underway. [42][38]

The inspected migrator records returned token IDs before forwarding, adjusts remaining generation budget, and preserves request identity. That retained history is not a client acknowledgement cursor. [54]

Replay support is bounded by tracking length and attempt limits. Multi-choice and structured-output requests need additional continuation state and are unsupported by the documented migration path. Replaying tokens does not advance a newly initialized constraint machine. Generation constraints explains why text history alone is insufficient. [38]

Cancellation propagates through linked request contexts, but engines must observe it and stop computation. A received-signal metric does not confirm resource reclamation. Graceful stop does not retract already-streamed results; kill behavior depends on engine support. Streaming delivery separates generation, buffering and client observation. [55]

A worker-loss test should record routing removal, restored executable capacity, interrupted-request outcome and client-visible output separately. Inject cancellation during recovery and observe actual engine termination. A passing minimal canary establishes executable reachability, not representative-load performance or seamless continuation. [53][55][42]

Open questions

1. 01

   Adaptive placement must respond to changing request lengths without spending its savings on cache movement. Progress requires workload-shift experiments that include transfer, queueing and latency compliance, rather than a fixed optimal worker ratio. [30]

2. 02

   Checkpoint publication still needs backend-specific integrity and power-loss guarantees. Metadata visibility alone cannot establish them; progress would include corrupted-shard, interrupted-write and concurrent-reader tests with explicit accepted restart states. [48]

3. 03

   Elastic restart must preserve the intended update after group membership changes. Tensor resharding is easier than preserving sample progress and weighting; a decisive test compares the next restored update with an uninterrupted reference using identical examples. [46][23]

4. 04

   Streaming recovery needs a delivery contract across frontend loss and reconnection. Backend replay lacks a client acknowledgement boundary; progress requires token-offset reconciliation plus cancellation tests that verify actual engine cleanup. [54][55]