Sharding a model with FSDP

FSDP is the API for the sharding described in the previous section: each device owns a slice of the parameters, gradients and optimiser state, and the pieces are gathered back for the moment they are used. The mental model that makes it behave predictably is that sharding happens per wrapped unit — you choose the units, and that choice sets both the memory saving and the communication cost. This section covers choosing them, the mixed-precision settings, combining it with activation checkpointing, saving a model whose weights are spread across eight machines, and what to look at when it is slower than you expected.

Sharding per block#

Wrap the repeating unit of the model — a transformer block — rather than the whole thing. Wrapping only the root gathers every parameter at once, which defeats the purpose, since peak memory then includes the entire model. Wrapping each block means only one block's parameters are gathered at a time, and the next block's gather can overlap with the current block's compute.

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

for block in model.transformer_blocks:   # the repeating unit, innermost first
    fully_shard(block)
fully_shard(model)                       # then the root, for anything left over

# Peak memory now holds one block's worth of gathered parameters, not the model.

The one tuning knob worth knowing is reshard_after_forward. Left on, parameters are freed after the forward pass and gathered again for the backward, which is the lowest-memory setting and costs one extra all-gather. Turned off, they stay resident between the two, which is faster and uses more memory. The usual configuration is to reshard everywhere except the last block, which is gathered again almost immediately.

Mixed precision settings#

from torch.distributed.fsdp import MixedPrecisionPolicy

policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,     # what the forward and backward compute in
    reduce_dtype=torch.float32,     # what gradients are reduced in
)
for block in model.transformer_blocks:
    fully_shard(block, mp_policy=policy)

That split is the answer to a common interview question about why bf16 is safe for training when fp16 needs loss scaling. bf16 has the same exponent range as fp32, so it does not underflow the way fp16 does and needs no scaler; what it has instead is fewer mantissa bits, so the place to be careful is accumulation — which is exactly what reduce_dtype controls.

Activation checkpointing alongside#

Sharding and checkpointing solve different problems and are almost always used together: sharding handles everything that scales with parameter count, checkpointing handles activations, which scale with batch and sequence length. For the 7B-shaped model in the previous section that was 16 GiB of activations reduced to 1 GiB, for 33 per cent more compute.

from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper, CheckpointImpl,
)

for i, block in enumerate(model.transformer_blocks):
    if i % 2 == 0:                       # every block, or every second one
        model.transformer_blocks[i] = checkpoint_wrapper(
            block, checkpoint_impl=CheckpointImpl.NO_REENTRANT,
        )
# Apply checkpointing BEFORE fully_shard, so FSDP shards the wrapped module.

Checkpoints for a sharded model#

This is where people lose a training run. Calling state_dict on a sharded model gives you that rank's shard, so saving it from rank zero saves an eighth of the model, and the failure is only discovered at load time. There are two correct approaches and the right one depends on what the checkpoint is for.

# For resuming: every rank writes its own shard, in parallel. Fast, and the
# only sane option at scale -- no rank ever holds the whole model.
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

model_sd, optim_sd = get_state_dict(model, optimizer)
dcp.save({"model": model_sd, "optim": optim_sd}, checkpoint_id="step_1000")

# Resuming: load into the same shape, then push back into the live objects.
state = {"model": model_sd, "optim": optim_sd}
dcp.load(state, checkpoint_id="step_1000")
set_state_dict(model, optimizer, model_state_dict=state["model"],
               optim_state_dict=state["optim"])

For the other purpose — a single file to hand to inference — you need the full unsharded weights gathered onto one rank, which requires that rank to have enough memory for the whole model. Do that once at the end of training rather than every few hundred steps, and keep it separate from the resume checkpoints, which are the ones that must be fast.

When it is slow#

FSDP can run correctly and slowly, and the symptom is always the same — the memory fits, the loss falls, and the throughput is a fraction of what the hardware should give. The causes are a short list, in the order worth checking.

symptom / cause                          what to do
wrapped only the root                    wrap each repeating block
wrapped too finely, every Linear         wrap blocks, not layers
stage 3 on a slow interconnect           try stage 2, or accumulate gradients
reshard_after_forward on everywhere      leave the last block resident
gathering full state_dict every save     use dcp.save for resume checkpoints
CPU offload left on                      turn it off unless you truly need it

Measure it the way the previous section measures DDP: samples per second on one device against samples per second on all of them. Sharding that has taken throughput below what a smaller batch on plain data parallelism would have given is a configuration to abandon, not to tune further — and knowing that boundary is most of what the FSDP question in an interview is really asking.