Efficient Distributed Training: From DP to ZeRO and FlashAttention

Training billion-parameter models on a single GPU is no longer feasible—both memory and compute demands exceed what any single device can provide. The solution is parallelism, but choosing the right parallelism strategy (and combining them) is itself a non-trivial engineering problem. Data parallelism replicates the model but partitions data; tensor parallelism slices individual operators across devices; pipeline parallelism assigns different layers to different GPUs; and sequence parallelism handles the non-tensor-parallel regions. In practice, large-scale training combines all of these into a multi-dimensional parallel strategy.

This post covers the core parallelism methods—Data Parallelism, Megatron-style Tensor Parallelism, GPipe/PipeDream Pipeline Parallelism, and Sequence Parallelism—along with memory optimization techniques including ZeRO, PyTorch FSDP, mixed precision training, activation recomputation, and FlashAttention. For each method, we provide the key formulations, trade-offs, and reference implementations that practitioners need to make informed decisions when scaling training.

Training Parallelism

Data Parallelism

Data Parallelism (DP) replicates the full model on each device and partitions the mini-batch across them. Each device executes forward and backward passes on its data shard independently, then gradients are averaged via all-reduce before the synchronized weight update. DP is simple to implement and scales well when the model fits in a single GPU’s memory.

Model Parallelism (Tensor Parallelism)

Tensor-level Model Parallelism (TP)¹ distributes individual operators across devices, enabling models too large for a single GPU. Megatron-LM pioneered this approach for transformers.

MLP Block

Consider the standard MLP within a transformer:

TP splits $A$ along columns: \(A = [A_1, A_2]\). Since GeLU is element-wise, it can be applied independently to each partition:

Matrix $B$ is then split along rows, allowing direct consumption of the GeLU outputs without inter-device communication until the final all-reduce.

Self-Attention Block

For self-attention:

TP partitions the $Q$, $K$, $V$ projection matrices along columns so that each attention head resides on a single GPU. The output projection is split along rows. This yields an efficient design requiring only two all-reduce operations per transformer layer (one in forward, one in backward).

Non-parallel components: Dropout, layer normalization, and residual connections are replicated (not split) across GPUs. Each GPU maintains a duplicate of the LayerNorm parameters, allowing these operations to proceed locally without communication.

Pipeline Parallelism

Pipeline Parallelism (PP)²³⁴ assigns consecutive groups of layers to different devices, with data flowing through them sequentially as micro-batches.

GPipe

GPipe² splits the mini-batch into $m$ micro-batches that flow through $p$ pipeline stages simultaneously.

Pipeline bubble: The unavoidable idle time at the start and end of each batch. With $p$ stages and $m$ micro-batches:

When $m \gg p$, the bubble overhead becomes negligible.

PipeDream (1F1B)

PipeDream³ introduces the one-forward-one-backward (1F1B) schedule: each device immediately starts the backward pass for a micro-batch as soon as its forward pass completes. This interleaving ensures no GPU sits idle during steady state.

With interleaved scheduling, each device handles $v$ model chunks (multiple non-contiguous stages), reducing per-micro-batch latency by $1/v$ and shrinking the bubble to:

Sequence Parallelism

Sequence Parallelism (SP)⁵ targets the non-tensor-parallel regions of a transformer (LayerNorm, dropout, residual connections), which are independent along the sequence dimension.

The standard non-parallel block:

SP splits $X$ along the sequence dimension: \(X = [X_1^s, X_2^s]\). After LayerNorm (sequence-parallel), an all-gather provides the full $Y$ for the TP region. After the TP GEMM, reduce-scatter returns to sequence-parallel sharding:

SP replaces the all-reduce in TP with all-gather (forward) and reduce-scatter (backward), keeping total communication volume constant while eliminating activation memory redundancy in non-TP regions.

ZeRO: Memory-Efficient Optimizer

Zero Redundancy Optimizer (ZeRO)⁶ eliminates memory redundancy across DP ranks by partitioning (instead of replicating) model states. ZeRO has three progressive stages:

1. \(P_\text{os}\): Partition optimizer states only.
2. \(P_\text{os+g}\): Additionally partition gradients.
3. \(P_\text{os+g+p}\): Additionally partition parameters.

See DeepSpeed ZeRO tutorial for implementation details.

PyTorch FSDP

Fully Sharded Data Parallel (FSDP)⁷ shards model parameters, gradients, and optimizer states across devices. During forward/backward, only one shard is materialized (unsharded) at a time, keeping peak memory proportional to a single shard.

FSDP uses a sharding factor $F$:

• $F = 1$: equivalent to standard DP (all-reduce on gradients).
• $F = W$ (world size): full sharding—each device holds $1/W$ of the model.
• $1 < F < W$: hybrid sharding.

Sharding strategy (flatten-concat-chunk): All parameters within an FSDP unit are flattened into a single contiguous FlatParameter, padded if necessary for divisibility by $F$, then chunked evenly across ranks.

Mixed Precision Training

Mixed precision⁸⁹ executes most operations in FP16 (or BF16) for speed, while maintaining an FP32 master copy of weights for numerical stability.

Training procedure:

1. Maintain FP32 master weights.
2. Each iteration: cast to FP16 → forward → scale loss by $S$ → backward → unscale gradients by $1/S$ → update FP32 master weights.

Dynamic Loss Scaling

Start with a large scaling factor $S$. If no overflow occurs for $N$ iterations, increase $S$. If overflow is detected, skip the update and decrease $S$.

1
2
3
4
5
6
7
8
9
10
11
12
13
 
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
 

AMP Optimization Levels (NVIDIA Apex)

Floating-Point Formats

Memory-Efficient Techniques

CPU Offload

When GPU memory is exhausted, model states (optimizer states, parameters) can be offloaded to CPU memory and fetched back on demand. ZeRO-Offload and ZeRO-Infinity implement this strategy.

Activation Recomputation

Instead of storing all intermediate activations for backpropagation, selective activation recomputation stores only a subset (e.g., activations at layer boundaries) and recomputes the rest during the backward pass. This trades compute for memory, typically with minimal wall-clock overhead.

FlashAttention

FlashAttention¹⁰ reduces HBM (High Bandwidth Memory) accesses by tiling the attention computation into SRAM-sized blocks, avoiding materialization of the full $N \times N$ attention matrix.

FlashAttention v1 tiles K/V in the outer loop and Q in the inner loop; v2 reverses this to reduce SRAM visits. The result is exact attention (not an approximation) with $O(N)$ memory instead of $O(N^2)$.

References