Distributed Training: FSDP, DeepSpeed, and Megatron

Under one common mixed-precision Adam recipe, a 70-billion-parameter model needs 1.12 TB for weights, gradients, and optimizer state. That exceeds an 80 GB GPU before activations enter the budget. Distributing the state lets us train a model that would otherwise be too large.

Replicating the full model can increase throughput, but it cannot solve that storage problem. Sharding state reduces each device’s persistent allocation. Splitting the computation adds other choices about which tensors cross the network.

Prerequisites

This post assumes you've read:

• Gradients and Learning for how forward, loss, backward, and optimizer step fit together.
• GPUs and Floating-Point Formats so "fp16" and "fp32 master weights" mean something.
• Training View vs. Inference View for the work and memory involved in a forward pass over a full sequence.

Why one GPU isn't enough

For one common mixed-precision Adam accounting scheme, a model with $P$ parameters stores:

• Weights in fp16: $2P$ bytes.
• Gradients in fp16: $2P$ bytes.
• Optimizer states: an fp32 master copy of the weights ($4P$), first moment ($4P$), and second moment ($4P$). So $12P$ bytes.
• Activations for the backward pass. Its size varies with batch size, sequence length, architecture, and checkpointing.

That's roughly $16P$ bytes of model state plus activations. For 70B parameters: $16 \times 70 \times 10^9 = 1.12$ TB just for the model state. Even with aggressive activation checkpointing, you're at a terabyte-class problem.

An 80 GB H100 cannot hold the 1.12 TB of model state in this example. Even a 7B model needs 112 GB under these assumptions, before activations. Other precisions, optimizers, CPU offload, or parameter-efficient fine-tuning change the budget; the claim here concerns full training with the stated on-device state.

The optimizer state in this accounting is $12P$ bytes, six times the 16-bit weights. Sharding only that state across $N$ ranks reduces persistent storage from $16P$ to $4P + 12P/N$. The remaining weights and gradients still occupy memory.

Data parallelism: replicate the model, split the batch

The simplest distributed training you can do: put a complete copy of the model on every GPU, give each GPU a different chunk of the batch, and have them synchronize gradients at the end of each step. This is data parallelism, and in PyTorch it's DistributedDataParallel (DDP).

Each of $N$ GPUs has an identical replica of the model weights. At each step, the global batch gets split into $N$ micro-batches, one per GPU. Forward and backward run independently on each GPU, and at the end each GPU has its own gradient tensor. Then all $N$ gradient tensors get averaged across the cluster via an all-reduce operation, so every GPU ends up with the same averaged gradient. Each GPU runs the optimizer step locally. Because they started from the same weights and just applied the same averaged gradient, they apply the same averaged update, subject to implementation and floating-point details.

DDP can scale throughput well when each rank has enough computation to amortize communication. Implementations bucket gradients and overlap all-reduces with backward computation. Communication is not free, and scaling depends on workload and interconnect.

Each GPU still holds the full weights, gradients, and optimizer state. DDP splits the batch across devices, so it cannot by itself make an oversized model fit.

The ZeRO insight: you don't need a full copy on every GPU

After gradient synchronization, DDP replicas maintain equivalent model and optimizer state. ZeRO reduces that replication by assigning portions of the state to different ranks. Communication supplies the parameters or reduced gradients when each rank needs them.

The wrinkle is timing. You can shard weights, but at some point each GPU needs the full weight matrix to do a forward pass. You can shard gradients, but the optimizer step needs the gradient for each parameter on the GPU that owns that parameter. ZeRO figures out the exact communication pattern that makes each stage work.

The three ZeRO stages shard progressively more state, with different collective communication patterns.

• Stage 1 shards optimizer states. Each rank updates its assigned parameters; communication makes the updated weights available to the other ranks.
• Stage 2 also shards gradients, using reduce-scatter so each owner receives its reduced gradient shard. At four ranks, the ideal persistent state changes from $7P$ bytes in Stage 1 to $5.5P$, not a twofold reduction.
• Stage 3 also shards parameters. All-gathers temporarily materialize each wrapped unit for computation; persistent state approaches $16P/N$ bytes under this accounting.

Stage 3 on 64 ranks gives an ideal 64-fold reduction in persistent model state relative to DDP. Temporary gathers, activations, and communication remain part of the budget.

PyTorch FSDP and DeepSpeed ZeRO-3 implement related fully sharded data-parallel strategies. Their APIs, buffer lifetimes, and scheduling options differ. FSDP2 is PyTorch’s newer per-parameter implementation; the sketch below shows the older wrapper API explicitly.

FSDP's rhythm: all-gather, compute, reshard

For a 4-layer transformer sharded across 4 GPUs, consider a forward pass without prefetching. Each GPU holds 1/4 of every layer's weights in permanent storage. When the forward pass starts, layer 1's weights are reassembled via all-gather: every GPU temporarily holds the full layer-1 weight tensor. The matmul runs. Then the gathered weights get dropped, leaving only the sharded 1/4 behind. Next, layer 2 goes through the same gather-compute-reshard cycle. And so on.

The backward pass can regather parameters and reduce-scatter gradients to their owning ranks. Peak memory includes persistent shards, gathered units, activations, communication buffers, and allocator overhead. Prefetching or retaining parameters can keep more than one unit materialized.

A simple fully sharded schedule uses parameter all-gathers in forward and backward plus a gradient reduce-scatter. For equal-sized, equal-dtype payloads, this moves roughly three units of data per rank, compared with two for a ring all-reduce in DDP, each scaled by $(N-1)/N$. Actual byte counts depend on parameter and reduction dtypes, resharding, and the collective algorithm. Overlap can hide some of that time.

Minimal math: where the memory actually goes

For a model with $P$ parameters, trained in mixed precision with Adam, on $N$ GPUs:

For 70B parameters on 64 ranks, the ideal persistent state is 17.5 GB per GPU. This subtotal excludes activations, temporary gathers, and workspace.

For 405B parameters on an illustrative 128 ranks, persistent state is about 50.6 GB per rank. Whether either model fits depends on the remaining allocations.

Worked example

The model. 70 billion parameters. Mixed precision with Adam. That's $16 \times 70 \times 10^9 = 1.12$ TB of model state, plus activations.

The hypothetical hardware: 64 GPUs with 80 GB each, arranged in eight nodes. Total device memory is 5.12 TB. The local and cross-node links affect how much collective traffic overlaps with computation.

Strategy. FSDP across all 64 GPUs.

Per-GPU budget under FSDP.

• Weights: $2 \times 70 \times 10^9 / 64 \approx 2.2$ GB
• Gradients: $2 \times 70 \times 10^9 / 64 \approx 2.2$ GB
• Optimizer states: $12 \times 70 \times 10^9 / 64 \approx 13.1$ GB
• Model state subtotal: ~17.5 GB

About 62.5 GB remains per GPU before activations, temporary gathers, and workspace. Measure those allocations for the chosen sequence length and micro-batch; parameter-state arithmetic alone cannot establish a batch-size limit.

Assume, for a rough estimate, 25 PFLOP/s sustained across the cluster. The dense-transformer approximation $C \approx 6PD$ gives $5.88\times10^{23}$ FLOPs for 70B parameters and 1.4T tokens. Dividing by $25\times10^{15}$ FLOP/s gives about 272 days, before interruptions or costs omitted by the approximation. This is an assumed throughput, not a measured FSDP result.

When FSDP isn't enough: Megatron and tensor parallelism

More ranks reduce persistent state, but the largest gathered unit and the available compute per rank can still limit a fully sharded run.

One limit is a large gathered unit. A hypothetical $16{,}384 \times 65{,}536$ weight matrix occupies about 2.15 GB at two bytes per value. Moving 2 GB through a 100 GB/s link takes at least 20 ms under a simple payload/bandwidth estimate. Whether that limits throughput depends on collective traffic and how much transfer overlaps with compute.

Another limit is too little computation per rank. For a fixed global batch, increasing the data-parallel degree reduces local work. Small local matrices may use the GPU less efficiently and leave less compute to overlap communication.

Tensor and pipeline parallelism split computation as well as storage. They introduce their own communication and scheduling costs.

Tensor parallelism (TP) can split a weight matrix along its output dimension. Each GPU computes different output coordinates; an all-gather can concatenate those slices. If the next layer consumes them in a row-parallel matmul, each rank instead produces a partial sum of the next output, which a reduction combines. All-reduce adds partial sums; it does not concatenate output slices.

Tensor-parallel groups are often placed on fast local links because they communicate within layers. The right group size depends on tensor shapes, device memory, and topology; TP=8 is an example rather than a rule.

Pipeline parallelism (PP) splits the layers of the model across different GPU groups. Stage 1 holds layers 1-8, stage 2 holds layers 9-16, and so on. A forward pass walks through the stages: micro-batch 1 goes through stage 1, then stage 2, while stage 1 picks up micro-batch 2. This is a wavefront pattern. The pipeline fills, runs steady-state, then drains.

Pipeline fill and drain create idle stage-slots, called the bubble. More micro-batches amortize that cost. At the same stage and micro-batch counts, non-interleaved 1F1B primarily reduces saved activations compared with doing all forwards before backwards; it does not inherently shrink the ideal bubble. Interleaving can reduce the bubble at the cost of more communication. Stage balance and memory determine which schedule is useful.

Combining parallelism axes

Large training runs often combine these strategies. A three-axis example uses:

• Data parallelism (or ZeRO/FSDP within data parallelism) across one axis. Each data-parallel replica sees a different micro-batch.
• Tensor parallelism within each replica, usually across the 8 GPUs in a node, using NVLink.
• Pipeline parallelism across nodes, with each stage holding a contiguous range of layers.

In this three-axis example, a rank has coordinates (dp_rank, tp_rank, pp_stage). Real systems can add more axes: the Llama 3 report describes context parallelism as well, making its training scheme four-dimensional.

More pipeline stages can increase bubbles, larger tensor groups can increase communication, and more data replicas change the relationship between local and global batch size.

Implementation sketch

For the code block, the legacy FSDP1 wrapper illustrates the communication boundaries with one wrapper around a model and one training step. The sketch assumes an application-provided transformer; it is not a complete runnable trainer. Current PyTorch documentation introduces FSDP2 via fully_shard.

import os
import torch
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools
 
# --- Init process group: one process per GPU ---
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
 
# --- Construct the application model on CPU on each rank ---
# (In practice you'd use meta-tensors to avoid ever materializing the full model)
from my_transformer import TransformerDecoder, TransformerBlock
model = TransformerDecoder(vocab=50257, d_model=4096, n_layers=32, n_heads=32)
 
# --- Wrap with FSDP ---
# auto_wrap_policy tells FSDP to shard each TransformerBlock separately,
# so the all-gather scope is one block at a time, not the whole model.
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},
)
 
mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)
 
model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # == ZeRO-3
    mixed_precision=mp_policy,
    device_id=local_rank,
)
 
# --- Optimizer: sees only the local shards, not the full params ---
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
 
# --- One training step ---
def train_step(batch):
    # FSDP auto-gathers each block's params at forward time,
    # then reshards as soon as compute is done.
    batch = {k: v.to(local_rank) for k, v in batch.items()}
    logits = model(batch["input_ids"])
    loss = torch.nn.functional.cross_entropy(
        logits[:, :-1].reshape(-1, logits.size(-1)),
        batch["labels"][:, 1:].reshape(-1),  # next-token targets
    )
    # Backward pass: gradients get reduce-scattered across ranks,
    # so each rank ends up with only its shard of the gradient.
    loss.backward()
    optimizer.step()       # updates only the local param shard
    optimizer.zero_grad()
    return loss.item()

The wrapper arranges parameter gathers and gradient reductions around module execution. Data must still be partitioned across ranks, targets must be shifted for causal language modeling, and distributed initialization and checkpointing need their own handling.

auto_wrap_policy selects the units that FSDP wraps. FSDP shards one "unit" at a time. If the unit is the whole model, you get a single gigantic all-gather at every forward pass and no overlap with compute. If the unit is one transformer block, you get many smaller gathers that can be overlapped with the compute of neighboring blocks. Picking the right granularity is the main knob you turn when tuning FSDP for throughput.

Misconceptions

"FSDP is model parallelism." Not quite. FSDP is still data parallelism in the sense that every rank runs the full computation on a different micro-batch; every GPU does a forward pass on every layer. What FSDP shards is the storage of the parameters, not the compute. Tensor parallelism is model parallelism proper: each GPU computes only a slice of each operation, and the results have to be combined. They compose, but they aren't the same thing.

"Adding more GPUs always helps." More ranks can reduce persistent state or increase throughput, but communication, local workload, and global batch size limit scaling. There is no universal rank count or batch size where scaling stops helping.

"ZeRO-3 and FSDP are unrelated." Both implement fully sharded data-parallel training, but their APIs, buffer lifetimes, and scheduling options differ. Choose and measure the implementation for the intended workload rather than assuming identical behavior.

What's next

The next post covers supervised fine-tuning, how you take a pretrained base model and teach it to follow instructions. Its examples demonstrate the responses we want the model to produce.

Additional reading (and watching)

• Korthikanti, V., et al. (2022). Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198. Good source for the per-GPU memory accounting, including activations under selective checkpointing.

• Rajbhandari, S., et al. (2019). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054. The foundational paper that established the three stages and the per-GPU memory math this whole post is built on.

• Li, S., et al. (2020). PyTorch Distributed: Experiences on Accelerating Data Parallel Training. VLDB 2020. The DDP paper — documents the gradient bucketing and overlap tricks that make naive data parallelism fast in PyTorch.

• NVIDIA. (2020). NCCL: Optimized primitives for collective multi-GPU communication. Documentation. A collective-communication library used in many GPU training stacks.

• Zhao, Y., et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. VLDB 2023. The PyTorch team's FSDP paper, covering the design choices, auto-wrap policy, and how FSDP composes with other parallelism axes.

• Touvron, H., et al. (2023). LLaMA 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288. Training recipe for Llama 2, which used FSDP-style sharding at 70B scale.

• Meta AI. (2024). The Llama 3 Herd of Models. arXiv:2407.21783. The engineering appendix of this paper is a detailed source for what frontier-scale distributed training actually looks like in practice, including its four-axis parallelism recipe, failure rates, and bandwidth accounting.

• Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053. The paper that established tensor parallelism as a practical technique; the Narayanan et al. (2021) follow-up Efficient Large-Scale Language Model Training on GPU Clusters adds pipeline parallelism and the 3D recipe.

• Huang, Y., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. NeurIPS 2019. The original pipeline-parallelism paper, with the classic bubble-time analysis.

• PyTorch. Getting Started with FSDP2. Current per-parameter API and parameter/gradient communication lifecycle.