Distributed Machine Learning Systems

Distributed training is not “run the same program on more machines.” It is a state-coordination problem. Workers compute from different data, but training is meaningful only if their updates produce one defined sequence of model and optimizer states.

The 2016 TensorFlow paper made one dataflow graph span CPUs, GPUs, and machines, with mutable tensors enabling parameter-server training. That architecture is historically important. Current TensorFlow 2 executes eagerly by default, traces selected Python functions into graphs with tf.function, and exposes distributed execution through tf.distribute strategies. The enduring ideas are placement, communication, synchronization, failure, and state—not one version's API surface.

Three ways to divide training #

Data parallelism replicates the model and gives each worker a different batch shard. Workers compute gradients locally and combine them before or while updating one logical model. It is the default when a full model fits on each accelerator.

Model parallelism divides model state or operations across devices. It is required when the model, an embedding table, or one operation does not fit on one device. Every partition boundary may become communication on the forward and backward passes.

Pipeline parallelism places consecutive model stages on different devices and overlaps microbatches. It reduces idle time compared with moving one full batch through stages serially, but introduces pipeline bubbles and rules about which weight version each microbatch observes.

These forms compose. A large language model may use data-parallel replica groups, tensor-parallel layers inside each replica, and pipeline stages across groups. Correctness then depends on multiple nested synchronization boundaries.

Synchronous replicated data parallelism #

With $P$ workers, each reads a local batch of size $b$, computes gradient $g_i$, and participates in a collective reduction:

$$g = \frac{1}{P}\sum_{i=1}^{P} g_i.$$

Every replica applies the same optimizer update from the same starting state, so replicas remain equal. The effective batch size is usually $Pb$ unless gradients are accumulated across more steps. Changing $P$ therefore changes optimization semantics unless learning rate, normalization, and schedule are considered.

A ring all-reduce moves approximately

$$2\frac{P-1}{P}S$$

bytes through each worker for a gradient payload of size $S$, excluding latency, protocol overhead, topology mismatch, and retries. This scales the aggregate transfer better than sending every gradient through one coordinator, but a ring step advances at the pace of its slowest participant.

Collectives can overlap gradient communication with backpropagation: reduce a layer's gradients while earlier layers are still computing. Small tensors are often fused because launching a collective per tensor makes latency dominate. The right bucket size balances earlier overlap against fewer, larger transfers.

Parameter servers and sharded state #

A parameter-server design assigns variables to one or more stateful server tasks. Workers read parameters, compute gradients, and send updates to the server that owns each variable. Sharding distributes storage and update load; placement near the consumers reduces network cost.

This design is attractive for enormous sparse embeddings because one worker may touch only a small subset of rows. Replicating and all-reducing the full table would communicate mostly zeros. A sparse pull/push transfers selected indices and values instead.

The server path can be synchronous or asynchronous:

• synchronous workers contribute to a step barrier, preserving a clear global step but inheriting the slowest worker's latency;
• asynchronous workers apply updates without a global barrier, improving utilization while allowing gradients computed from older parameter versions to modify newer state.

If a worker reads version $t$ and its gradient is applied at $t+s$, then $s$ is its staleness. The gradient points downhill for the old model, not necessarily the current one. Small, bounded staleness may train well; unbounded staleness can destabilize optimization and makes exact reproduction difficult.

The straggler problem #

For a synchronous step,

$$T_{step} \approx \max_i(T_{compute,i} + T_{communicate,i}) + T_{update}.$$

One slow input shard, host, accelerator, network link, or compilation event holds every participant at the barrier. Average worker time is therefore a weak signal; tail latency and barrier wait reveal the scaling limit.

The 2016 TensorFlow paper demonstrated backup workers: request more worker results than the update requires and proceed after a quorum, discarding late duplicates. This can remove random tail latency, but it spends extra compute and changes which examples contribute to a step. It does not solve systematic data skew, a slow shared server, or a network bottleneck.

Other controls include deterministic sharding, balanced input preprocessing, bounded batches, health-based worker replacement, gradient accumulation, and elastic membership protocols that explicitly redefine batch and optimizer state when the worker count changes.

Placement is part of the algorithm #

A graph edge crossing a device boundary is a transfer. A model can have enough aggregate FLOPS and still slow down because placement creates a cut with more bytes per second than the interconnect carries.

Useful placement questions are:

• Which device owns each variable and optimizer slot?
• Which operations consume it, and how many bytes cross that boundary?
• Can an operation move to the data instead of moving the tensor?
• Can communication overlap independent computation?
• Does sparse structure survive serialization, or become a dense payload?
• Does the physical topology match the collective algorithm's assumed links?

Accelerators increase the importance of these questions because faster local compute leaves less time to hide transfer. A placement that was adequate on CPUs can become communication-bound on GPUs.

Eager execution and traced graphs #

Eager execution evaluates operations immediately, which makes Python control flow and debugging direct. tf.function traces compatible Python code into a callable graph that the runtime can optimize, serialize, and place. Tracing is not the same as “execute Python later”: Python side effects normally happen during tracing, and new input signatures can create new concrete graphs.

Distributed execution benefits from graphs because the runtime can see a larger dataflow region, schedule independent operations, and reduce per-op coordination. But correctness still belongs to explicit variable, random-state, input, and update semantics. A graph does not make an impure function deterministic or an external side effect retry-safe.

Checkpoints define recoverable training state #

A useful checkpoint contains more than weights:

• model variables;
• optimizer slots such as momentum and adaptive moments;
• the global step and learning-rate schedule state;
• random-number generator state when reproducibility requires it;
• input position or an explicit rule for safe replay;
• mixed-precision loss-scaling and other controller state.

Saving these values from incompatible logical times produces a checkpoint that loads but does not represent any real training step. Synchronous training has a natural barrier at which to snapshot. Asynchronous training needs a protocol that coordinates variable shards and in-flight updates or accepts a documented fuzzy boundary.

Retry semantics also cross the input boundary. After worker failure, replaying a batch is usually safer than skipping it, but an exactly reproducible sequence requires deterministic data order, seed ownership, augmentation, worker membership, collective order, and kernels. “Same checkpoint” alone is not bitwise reproducibility.

Diagnosing distributed training #

Profile a complete step timeline before changing topology. A low accelerator utilization number cannot distinguish input starvation, host launch overhead, collective wait, parameter-server queueing, or checkpoint I/O.

Choosing the coordination model #

Use replicated synchronous data parallelism when the model fits per device, gradients are reasonably dense, and predictable optimizer semantics matter. Use parameter-server or sharded-state designs when state is too large to replicate, access is sparse, or asynchronous progress is worth controlled staleness. Add model or pipeline parallelism when a single replica no longer fits, after measuring the communication cuts that division creates.

The real scaling result is time to a target model quality under a named failure and cost envelope. Device count, examples per second, and graph size are intermediate measurements. A distributed design is successful only when its coordination semantics remain correct and the removed compute time exceeds the communication, synchronization, and recovery it introduces.