AI Training

As models grow larger and datasets expand, training on a single accelerator becomes impractical. Distributed training splits the workload across multiple devices, potentially across multiple machines, to reduce training time from weeks to hours. Understanding distributed training is essential for any engineer working with modern large-scale ML systems.

The fundamental challenge of distributed training is maintaining the mathematical equivalence of optimization while dividing work across devices. Communication overhead between devices, synchronization strategies, and fault tolerance all add complexity that does not exist in single-device training.

Failure Modes in Distributed Systems

Distributed training introduces new failure modes that single-device training does not have. Network partitions, device failures, stragglers (slow devices), and communication deadlocks can all disrupt training. Production distributed training systems must handle these failures gracefully, typically through checkpointing and automatic restart mechanisms.

• Network partitions — Devices lose connectivity, causing communication operations to hang or timeout.
• Device failures — A GPU may crash or produce incorrect results due to hardware faults.
• Stragglers — One slow device forces all others to wait during synchronization barriers.
• Communication deadlocks — Incorrect ordering of collective operations blocks all progress.
• Checkpoint corruption — Failed writes during checkpointing can corrupt the saved state.

Choosing a Distribution Strategy

Data parallelism is the most common distributed training strategy, where each device holds a complete copy of the model and processes a different subset of the training data. After computing local gradients, devices synchronize by aggregating (typically averaging) their gradients before updating model parameters.

Synchronous Data Parallelism

Synchronous data parallelism uses an all-reduce operation to aggregate gradients from all devices before any device updates its parameters. This maintains mathematical equivalence with single-device training (up to floating-point precision) and is the default approach in most frameworks. The all-reduce operation can be implemented using ring all-reduce, which achieves optimal bandwidth utilization.

Asynchronous Data Parallelism

Asynchronous data parallelism allows devices to update parameters independently without waiting for all gradients to arrive. This eliminates the synchronization bottleneck but introduces gradient staleness, where devices may use slightly outdated parameters. Asynchronous methods converge faster in wall-clock time for some problems but can be less stable and harder to reproduce.

Scaling: Batch Size and Learning Rate

Scaling data parallelism effectively requires careful attention to the relationship between batch size and learning rate. The linear scaling rule suggests increasing the learning rate proportionally to the number of devices, but this breaks down at very large batch sizes.

When a model is too large to fit in a single device's memory, model parallelism distributes different parts of the model across devices. Tensor parallelism splits individual layers across devices, while pipeline parallelism assigns different layers to different devices. Both approaches are essential for training the largest foundation models.

Tensor Parallelism

Tensor parallelism partitions weight matrices across devices and uses collective communication to combine partial results. For Transformer models, the attention and feed-forward layers can be split across devices along specific dimensions. Megatron-LM popularized efficient tensor parallelism strategies for large language models.

Pipeline Parallelism

Pipeline parallelism assigns sequential groups of layers to different devices, creating a pipeline where micro-batches flow through devices in sequence. GPipe and PipeDream introduced techniques to reduce the "pipeline bubble" (idle time when devices wait for data), including micro-batch interleaving and asynchronous scheduling.

Hybrid (3D) Parallelism

In practice, the largest models use a combination of all three parallelism strategies. A typical setup might use tensor parallelism within a node (where NVLink provides high bandwidth), pipeline parallelism across nodes, and data parallelism across groups of nodes.

Mixed precision training uses lower-precision number formats (typically FP16 or BF16) for most computations while maintaining FP32 for critical operations. This approach can nearly double training throughput on modern GPUs with tensor cores while using significantly less memory.

The Mixed Precision Recipe

The key technique is maintaining a FP32 master copy of weights while performing forward and backward passes in half precision. Gradients are computed in FP16/BF16 and then cast to FP32 for the weight update. This preserves the precision of the update step, which is most sensitive to numerical errors.

1. Maintain FP32 master weights (the "source of truth" for parameters).
2. Cast weights to FP16/BF16 for the forward pass.
3. Compute loss and backward pass in FP16/BF16.
4. Cast gradients back to FP32.
5. Update FP32 master weights using the FP32 gradients.

Loss Scaling for FP16

Loss scaling is essential for FP16 training to prevent gradient underflow. Small gradient values that are representable in FP32 may round to zero in FP16. Loss scaling multiplies the loss by a large factor before backpropagation (scaling up gradients) and then divides the gradients after the backward pass. Dynamic loss scaling automatically adjusts the scale factor during training.

BFloat16: The Simpler Alternative

BFloat16 (BF16), developed for Google's TPUs and now supported on NVIDIA Ampere and newer GPUs, uses the same exponent range as FP32 but with reduced mantissa precision. This eliminates the need for loss scaling because BF16 can represent the same range of values as FP32.

class="tok-comment"># PyTorch automatic mixed precision (AMP)
scaler = torch.cuda.amp.GradScaler()  class="tok-comment"># Only needed for FP16

for batch in dataloader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
        output = model(batch)
        loss = loss_fn(output, targets)
    class="tok-comment"># With BF16, no scaler needed:
    loss.backward()
    optimizer.step()

Efficient gradient communication is often the bottleneck in distributed training. For a model with billions of parameters, each gradient synchronization step must transfer gigabytes of data across the network. Optimizing this communication is essential for achieving good scaling efficiency.

Gradient Compression

Gradient compression reduces communication volume by transmitting approximate gradients. Techniques include gradient quantization (reducing precision of transmitted gradients), sparsification (transmitting only the largest gradients), and error feedback (accumulating compression error for future correction). These methods can reduce communication by 10-100x with minimal impact on convergence.

Gradient Accumulation

Gradient accumulation is a practical technique for simulating larger batch sizes without additional devices. Instead of synchronizing gradients after every micro-batch, gradients are accumulated locally over multiple micro-batches before synchronization. This reduces communication frequency and allows training with effectively larger batch sizes on limited hardware.

class="tok-comment"># Gradient accumulation: simulate batch_size * accum_steps
accum_steps = class="tok-number">4
for i, batch in enumerate(dataloader):
    loss = model(batch) / accum_steps  class="tok-comment"># Scale loss
    loss.backward()                     class="tok-comment"># Accumulate gradients
    if (i + class="tok-number">1) % accum_steps == class="tok-number">0:
        optimizer.step()                class="tok-comment"># Update after N steps
        optimizer.zero_grad()           class="tok-comment"># Reset gradients

Overlapping Communication and Computation

Overlapping computation with communication is critical for hiding communication latency. By starting gradient communication for earlier layers while later layers are still computing, the communication overhead can be partially or fully masked. Modern frameworks implement this through bucket-based gradient reduction, where gradients are grouped into buckets that are communicated as soon as all gradients in the bucket are available.