FSDP: sharding a model across GPUs

Setup#

pip install torch

Be warned before you plan around this lesson: FSDP is a GPU technique. The sharding and the forward pass below do run on CPU with the gloo backend, and that is how the outputs were captured. A full training step does not — reduce_scatter, which FSDP needs for gradients, is not implemented in gloo. Everything past the forward pass here needs real CUDA GPUs.

The modern API is fully_shard, sometimes called FSDP2. It replaced the older FullyShardedDataParallel wrapper class, which is still present and still works.

One import note, since this API moved recently. from torch.distributed.fsdp import fully_shard is the public path from PyTorch 2.6 onward. On 2.5 the same function lives at torch.distributed._composable.fsdp, and that is the path used to capture the output below. If the public import raises ImportError on your machine, you are on 2.5 or older.

Watching a model get cut up#

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.device_mesh import DeviceMesh
from torch.distributed.fsdp import fully_shard

dist.init_process_group(backend="gloo")
rank, world = dist.get_rank(), dist.get_world_size()
mesh = DeviceMesh.from_group(dist.group.WORLD, "cpu")   # "cuda" on GPUs

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 64))
total = sum(p.numel() for p in model.parameters())

for layer in model:                      # shard each block, then the whole model
    if isinstance(layer, nn.Linear):
        fully_shard(layer, mesh=mesh)
fully_shard(model, mesh=mesh)

w = model[0].weight                      # now a DTensor: a global view over shards
print(f"rank {rank}: 0.weight logical shape {tuple(w.shape)}, "
      f"my shard {tuple(w.to_local().shape)}, placement {w.placements}")
kept = sum(p.to_local().numel() for p in model.parameters())
print(f"rank {rank}: {total} parameters in the model, {kept} stored on this rank")

out = model(torch.randn(2, 64))          # gathers each block just in time
print(f"rank {rank}: forward output {tuple(out.shape)}")

dist.destroy_process_group()

torchrun --nproc_per_node=2 fsdp_shard.py

rank 0: 0.weight logical shape (64, 64), my shard (32, 64), placement (Shard(dim=0),)
rank 0: 8320 parameters in the model, 4160 stored on this rank
rank 0: forward output (2, 64)
rank 1: 0.weight logical shape (64, 64), my shard (32, 64), placement (Shard(dim=0),)
rank 1: 8320 parameters in the model, 4160 stored on this rank
rank 1: forward output (2, 64)

Lines are grouped by rank for reading. Three things to take from that output.

The weight still says it is 64x64. model[0].weight is now a DTensor — a tensor that knows it is a view over pieces living on several ranks. Your code, and print(model), still see the logical shape. w.to_local() reveals the real 32x64 slice this rank owns.

placement=Shard(dim=0) says the cut runs along the first dimension. Rows 0–31 on rank 0, rows 32–63 on rank 1.

8320 parameters, 4160 per rank. Exactly half each, with two ranks. That halving is the entire point.

The forward pass still returns the right shape, because FSDP gathered each block back to full size the instant it was needed, then released it.

Caveat, stated plainly: on this CPU run, calling .backward() on that output raises an autograd in-place-modification error. FSDP's backward path is written against CUDA collectives. Treat the CPU version as a way to see the sharding, not as a way to train.

A real GPU training loop#

This block needs at least two CUDA GPUs, so no output is shown for it. Presenting invented multi-GPU numbers would be worse than showing none.

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard

local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")
mesh = init_device_mesh("cuda", (dist.get_world_size(),))

model = nn.Sequential(*[nn.Linear(1024, 1024) for _ in range(8)]).cuda()

for block in model:                       # shard per block, not per parameter
    fully_shard(block, mesh=mesh)
fully_shard(model, mesh=mesh)

opt = torch.optim.AdamW(model.parameters(), lr=1e-4)   # AFTER sharding
x = torch.randn(8, 1024, device="cuda")

for step in range(3):
    opt.zero_grad(set_to_none=True)
    model(x).sum().backward()
    opt.step()
    if dist.get_rank() == 0:
        print(f"step {step} done")

dist.destroy_process_group()

Two lines carry most of the meaning. Shard per block, then the root: each fully_shard call creates one gather-and-release unit, so its size controls the memory/communication trade. Shard nothing and you are back to DDP; shard every single parameter and you drown in tiny collectives. One transformer block per call is the standard choice.

Build the optimizer after sharding. model.parameters() must return the sharded DTensor objects, or the optimizer allocates state for full-size parameters and you lose the memory saving you came for.

What FSDP actually saves#

For a model with P parameters trained in float32 with Adam, the per-GPU fixed cost is roughly:

At 8 ranks, a 1-billion-parameter model drops from about 16 GB of fixed state per GPU to about 2 GB. Activations are not sharded by this — they scale with batch size, and gradient checkpointing is the lever for those.

Communication rises by roughly 50% against DDP: an all-gather in the forward pass, an all-gather in the backward pass, and a reduce-scatter for gradients, instead of DDP's single all-reduce.

Common mistakes#

Building the optimizer before fully_shard. No error, no warning, no memory saved.

Sharding every parameter individually. Thousands of tiny collectives, each paying full network latency. Group at block level.

Saving with plain torch.save(model.state_dict()). Each rank writes its own shard, so you get R incomplete files. Use torch.distributed.checkpoint — see saving from one rank.

Reaching for FSDP when DDP fits. FSDP is strictly slower per step. If the model, gradients and optimizer state fit on one card, DDP is the right answer. Try mixed precision and checkpointing before you reach for sharding.

Calling .to_local() and forgetting it. Once you unwrap a DTensor, the result is an ordinary tensor holding one rank's slice. Arithmetic on it is silently per-rank.

Try it yourself#

Run fsdp_shard.py with --nproc_per_node=4. The stored-parameter count should fall to a quarter. Then remove the per-layer fully_shard calls, keep the root one, and see how the reported shard changes.

What to learn next#

• Pipeline and tensor parallelism — for when a single layer does not fit.
• GPU memory anatomy — where the 16 bytes per parameter comes from.
• Gradient checkpointing — the other half of the memory problem.