TopolearnFSDP is PyTorch's native way to train multi-billion-parameter models.
DDP keeps a full copy of parameters, gradients and optimizer state on every GPU. For a model with P parameters trained with AdamW in mixed precision, that is roughly 16P bytes per GPU, too much once P reaches billions. Fully Sharded Data Parallel splits all three across ranks (the idea known as ZeRO stage 3).
How it runs: each rank stores only its shard. Before a sharded unit's forward, FSDP all-gathers its full parameters, computes, then frees them. Backward all-gathers again and reduce-scatters gradients so each rank receives only its shard's gradients. The optimizer then updates local shards.
In current PyTorch, use FSDP2: torch.distributed.fsdp.fully_shard. You apply it bottom-up: to each transformer block (each becomes one all-gather unit), then to the root model. Parameters become DTensors, sharded along dimension 0. Create the optimizer after sharding.
Practical points:
MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)gathers weights in bf16 and reduces gradients in fp32.- Build huge models on the
metadevice, shard, thenmodel.to_empty(device="cuda")and initialise, so no rank ever materialises the full model. - Save with
torch.distributed.checkpoint, which writes sharded checkpoints in parallel and can reshard on load. - Combine with activation checkpointing; for even larger scale, add tensor or pipeline parallelism.
Example
# Launch: torchrun --nproc_per_node=8 train_fsdp.py
import torch
import torch.distributed as dist
from torch import nn
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
dist.init_process_group("nccl")
torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())
model = nn.Sequential(*[nn.TransformerEncoderLayer(1024, 16, batch_first=True)
for _ in range(12)]).cuda()
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
for block in model:
fully_shard(block, mp_policy=mp) # each block = one all-gather unit
fully_shard(model, mp_policy=mp) # root wraps whatever is left
opt = torch.optim.AdamW(model.parameters(), lr=1e-4) # create after sharding
x = torch.randn(4, 128, 1024, device="cuda")
loss = model(x).float().pow(2).mean()
loss.backward() # gradients reduce-scattered to shards
opt.step()
print(type(next(model.parameters())).__name__) # DTensor: each rank holds 1/8
dist.destroy_process_group()Remember
- FSDP shards params, grads and optimizer state (ZeRO-3)
- All-gather before compute, reduce-scatter after backward
- Apply
fully_shardper block, then to the root; build the optimizer afterwards
Watch out for
- Wrapping only the root module, which all-gathers the entire model at once and saves no peak memory.