FSDP — Training Models That Don't Fit on One GPU

When the model is too big for one GPU

DDP replicates the model on every GPU. That works until the model itself doesn't fit. FullyShardedDataParallel (FSDP) shards the model parameters, gradients, AND optimizer states across GPUs. Each GPU only holds 1/N of each, gathering the slices it needs for forward and backward then releasing them. This is what makes it possible to train a 70B-parameter model on a cluster of 24GB GPUs.

Sharding strategies

FULL_SHARD — shard parameters, gradients, optimizer states. Most memory-efficient. Default choice for the largest models.
SHARD_GRAD_OP — shard gradients and optimizer state only; replicate parameters. Faster than FULL_SHARD because no parameter all-gather, but uses more memory.
NO_SHARD — equivalent to DDP. Useful for unit tests that share the same training script.

FSDP vs FSDP2

FSDP1 wraps modules and is the legacy API; FSDP2 (the fully_shard() per-parameter approach) is the new in-development API with a cleaner interface and better composability with torch.compile. For new code in 2026, prefer FSDP2 if you're targeting bleeding-edge PyTorch — but FSDP1 is still production-stable and the documentation is more mature.

What stays the same

Inside the training loop, FSDP looks identical to DDP. You wrap the model and proceed as usual. The complexity is in setup (auto-wrap policies for nested modules, mixed-precision settings) — not in the loop.

Code

FSDP1 — the wrap-the-model API·python

import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

dist.init_process_group("nccl")
rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(rank)

model = MyHugeModel().cuda(rank)
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=rank,
)

# The training loop is the SAME as DDP
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for epoch in range(num_epochs):
    for x, y in loader:
        optimizer.zero_grad()
        loss = criterion(model(x.cuda(rank)), y.cuda(rank))
        loss.backward()
        optimizer.step()

Auto-wrap — let FSDP handle nested modules·python

import torch
import functools
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
import torch.nn as nn

# Wrap any nn.Module submodule with > 100M params automatically
auto_wrap_policy = functools.partial(
    size_based_auto_wrap_policy,
    min_num_params=100_000_000,
)

model = FSDP(
    MyHugeModel(),
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=rank,
)

# For Transformer-shaped models, prefer transformer_auto_wrap_policy
# which wraps each block — much better for FSDP performance

FSDP2 — the modern per-parameter API·python

import torch.nn as nn
from torch.distributed._composable.fsdp import fully_shard

# FSDP2 is per-parameter and per-block, not whole-model
class TransformerStack(nn.Module):
    def __init__(self):
        super().__init__()
        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(24)])
    def forward(self, x):
        for b in self.blocks:
            x = b(x)
        return x

model = TransformerStack()

# Apply FSDP2 to each block — finer control than FSDP1's wrap policy
for layer in model.blocks:
    fully_shard(layer)
fully_shard(model)

# FSDP2 advantages:
#   - Cleaner interaction with torch.compile
#   - Per-parameter sharding (mix sharded / non-sharded freely)
#   - More composable with other parallelism strategies (TP, PP)

Exercise

Take a model that fits comfortably on a single GPU. Run it through FSDP1 with FULL_SHARD on a multi-GPU box. Verify the per-GPU memory usage is roughly 1/N of single-GPU usage — that's the sharding actually working. Don't expect a speedup vs DDP at this scale; FSDP wins when the model wouldn't fit otherwise.