Distributed Training Deep Dive

Distributed Training Deep Dive

N.1.1 The Memory Math: Why You Need Parallelism at All

Before choosing a strategy you need to know whether your job fits. For an LLM in mixed precision (bf16 weights, fp32 optimizer state), Adam needs roughly 20 bytes per parameter: 2 for the bf16 weight, 4 for the fp32 master, 4 for the gradient, 8 for the two Adam moments. A 7B model therefore needs ~140 GB of static optimizer-state memory before any activation; an 80 GB H100 cannot hold it without sharding. Activations are the second variable: a 7B Llama at seqlen 4096, batch 4 per GPU adds ~30-40 GB on top. Activation checkpointing (recompute on backward) trades ~30% extra FLOPs for a 4-8x activation-memory reduction and is on by default in modern scripts.

N.1.2 Distributed Data Parallel (DDP): The Simplest Case

DDP is the default when the model fits on a single GPU and you want to scale throughput by processing more data per step. Each rank holds a full replica, processes a different mini-batch, and synchronizes gradients via an all-reduce after the backward pass. The PyTorch implementation in torch.distributed uses NCCL on NVIDIA hardware and is the baseline against which every other strategy is measured. The script reads its rank from environment variables and wraps the model in DistributedDataParallel:

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import AutoModelForCausalLM, AutoTokenizer

def main():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16
    ).to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
    # ... training loop with DistributedSampler ...
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

Launch with torchrun --nproc-per-node=8 train.py on a single node, or torchrun --nnodes=4 --nproc-per-node=8 --rdzv-backend=c10d --rdzv-endpoint=$MASTER:29500 train.py for multi-node. The torchrun docs describe elastic restart semantics; SageMaker and Azure ML wrap this with rendezvous services so you do not maintain the cluster yourself.

N.1.3 Fully Sharded Data Parallel (FSDP)

When the model no longer fits on one GPU, the cheapest upgrade is FSDP (FSDP2, the rewritten API, stable in PyTorch 2.4 (2024)). FSDP shards weights, gradients, and optimizer states across ranks; each rank gathers parameters for the current layer, computes, then frees them. Two strategies: FULL_SHARD (lowest memory) and HYBRID_SHARD (full shard within a node, replicate across nodes), the default for bandwidth-limited fabrics:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, MixedPrecision
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
import functools, torch

mp = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)
wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={LlamaDecoderLayer},
)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.HYBRID_SHARD,  # FULL_SHARD for single node
    auto_wrap_policy=wrap_policy,
    mixed_precision=mp,
    device_id=torch.cuda.current_device(),
    use_orig_params=True,
)

The auto_wrap_policy matters more than people realise: wrap at the transformer-block level for clean compute-communication overlap. Wrap too coarsely (whole model) and you defeat sharding; wrap too finely (every nn.Linear) and you flood NCCL with tiny collectives. The Hugging Face Accelerate FSDP guide (updated 2025) documents the wrap policies for the major model families.

N.1.4 DeepSpeed ZeRO: Three Stages and Two Offloads

ZeRO (Zero Redundancy Optimizer) from Microsoft predates FSDP and remains the standard when you need offload to CPU or NVMe. ZeRO-1 shards optimizer states; ZeRO-2 also shards gradients; ZeRO-3 also shards parameters (equivalent to FSDP FULL_SHARD). ZeRO-Infinity spills optimizer states and parameters to NVMe, letting a single node train a 100B+ model at the cost of slower steps. The DeepSpeed Ulysses work (March 2024) added sequence parallelism for long-context training, enabling million-token-context fine-tuning on commodity clusters. DeepSpeed remains the project of choice when you need to mix ZeRO with pipeline parallelism in the same run.

N.1.5 Pipeline and Tensor Parallelism: When Sharding is Not Enough

Beyond roughly 30B parameters, sharding alone runs out of room. Pipeline parallelism cuts the model into stages along its depth and streams micro-batches through them; the GPipe scheme and its 1F1B refinement are built into PyTorch as torch.distributed.pipelining (stable in PyTorch 2.5, late 2024). Tensor parallelism splits individual matrices across GPUs so one GeMM is computed by several ranks in parallel; the Megatron-LM patterns (NVIDIA's 2024 Megatron-Core release exposes the kernels as a library) split QKV and MLP projections across the tensor-parallel group. Tensor parallelism is communication-heavy and stays inside a node where NVLink is fast; pipeline parallelism is communication-light and spans nodes.

Frontier training composes all three axes into 3D parallelism (DP across replica groups, PP across nodes, TP within a node), used by NVIDIA, Meta (Llama-3 added a 4D context-parallel variant), and DeepSeek. Megatron-LM and NeMo are the production stacks; Hugging Face's nanotron (2024) is the readable reference implementation.

N.1.6 Decision Table: Which Strategy Fits

The table maps training scenarios to a recommended strategy. The fits-on-one-GPU column assumes 80 GB H100; halve for A100-40GB, quarter for a 24 GB consumer card.

N.1.7 Production Realities and 2024-2026 Tooling

Two trends matter today. FSDP has eaten most of DeepSpeed's mind-share for jobs that fit without offload; PyTorch 2.4's FSDP2 closed the configurability gap and the Hugging Face stack defaults to FSDP. DeepSpeed remains essential for NVMe offload, MoE expert parallelism, or composition with pipeline parallelism. Composable launchers (Accelerate, PyTorch Lightning) let you switch between DDP, FSDP, and DeepSpeed via config flags. On hardware, the 2025 GB200 NVL72 treats 72 GPUs as one coherent memory pool, collapsing the tensor-vs-pipeline distinction for sub-500B models. Track MosaicML LLM Foundry and Databricks' training stack as living references.