The unit you wrapped

Wrap a two-layer model in FSDP with a policy that wraps each layer, print the parameters, and the original names are gone. What named_parameters returns is _fsdp_wrapped_module.0._fsdp_wrapped_module._flat_param, one tensor per unit. The FlatParameter docstring at v2.2.2 states the construction plainly: it is comprised of one or more original parameters, flattened and concatenated.

The shard is a contiguous slice of that concatenation, and you can see exactly which slice. A Linear(4, 4) has 16 parameters, which flatten to 16 numbers in row-major order. On a world of two, rank 0's shard holds the first 8, which are rows 0 and 1 of the weight matrix, and rank 1's holds the second 8. Print rank 0's shard next to the original matrix and the numbers line up value for value.

Nothing about that slicing respects the shape of the original parameter, which is the point worth carrying forward. A shard boundary can fall in the middle of a row, in the middle of a tensor, or between two different parameters that happened to be flattened next to each other. The layer's shape only comes back when the unit is gathered.

FSDP.summon_full_params is the context manager that does the gathering on request. Inside it, the wrapped module's weight has shape (4, 4) again on every rank, which is the same all-gather a forward pass would have done, run because you asked rather than because a layer was about to execute.

rank 0: shard of layer0 [-0.0037434101, 0.2682217956, -0.4115225673, -0.3679695129,
                        -0.1925771832, 0.1340786815, -0.0099065900, 0.3964447379]
layer0 full, row 0 and 1: [-0.0037434101, 0.2682217956, -0.4115225673, -0.3679695129]
                          [-0.1925771832, 0.1340786815, -0.0099065900, 0.3964447379]
rank 0: summoned layer0 shape (4, 4)
rank 1: summoned layer0 shape (4, 4)
rank 0: grad shard of layer0 [0.0, 0.0, 0.0, 0.0, 1.2830595970, 1.2830595970, 1.2830595970, 1.2830595970]
rank 1: grad shard of layer0 [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

import functools
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29540"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(4, 4, bias=False), nn.ReLU(),
                          nn.Linear(4, 4, bias=False))
    ref = [p.detach().clone() for p in model.parameters()]
    policy = functools.partial(size_based_auto_wrap_policy, min_num_params=8)
    f = FSDP(model, auto_wrap_policy=policy, device_id=torch.device("cpu"))
    flat = dict(f.named_parameters())
    k0 = "_fsdp_wrapped_module.0._fsdp_wrapped_module._flat_param"
    print(f"rank {rank}: shard of layer0 {flat[k0].detach().tolist()}")
    if rank == 0:
        print("layer0 full, row 0 and 1:", ref[0][0].tolist(), ref[0][1].tolist())
    with FSDP.summon_full_params(f):
        print(f"rank {rank}: summoned layer0 shape "
              f"{tuple(f._fsdp_wrapped_module[0].weight.shape)}")
    f(torch.full((2, 4), 1.0)).sum().backward()
    print(f"rank {rank}: grad shard of layer0 {flat[k0].grad.tolist()}")
    dist.destroy_process_group()

if __name__ == "__main__":
    mp.spawn(worker, args=(2,), nprocs=2, join=True)

An auto_wrap_policy is a predicate FSDP runs over the module tree, and every module it says yes to becomes its own unit with its own flat parameter. Three policies on the same model, two blocks of two Linear(6, 6) layers each, 144 parameters in total, produce three different shapes of job on a world of four.

With no policy at all there is one unit and one flat parameter of 144, so each rank holds 36 numbers and the whole model is gathered at once. ModuleWrapPolicy({Block}) produces three FSDP modules, the root plus one per block, and two flat parameters of 18 per rank. size_based_auto_wrap_policy with a 30-parameter floor wraps each Linear separately: five FSDP modules and four flat parameters of 9.

Notice the count mismatch in the middle row, because it explains the structure. Three FSDP modules, two flat parameters. The root unit owns no parameters of its own once its children have claimed them all, so it gathers nothing and issues no collective. Counting isinstance(m, FSDP) and counting flat parameters answer two different questions.

The tradeoff runs in one direction. One large unit means one all-gather and a peak that holds the entire model unsharded. Many small units mean many all-gathers and a peak that holds one unit unsharded. Wrapping per transformer block is the usual middle, which is what ModuleWrapPolicy was built to express.

rank 0 | no policy: fsdp units 1 flat params [(36,)]
rank 0 | ModuleWrapPolicy({Block}): fsdp units 3 flat params [(18,), (18,)]
rank 0 | size_based(min=30): fsdp units 5 flat params [(9,), (9,), (9,), (9,)]

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(6, 6, bias=False)
        self.b = nn.Linear(6, 6, bias=False)
    def forward(self, x):
        return self.b(torch.relu(self.a(x)))

def build():
    torch.manual_seed(0)
    return nn.Sequential(Block(), Block())

def report(tag, f, rank):
    units = sum(1 for m in f.modules() if isinstance(m, FSDP))
    shapes = [tuple(p.shape) for _, p in f.named_parameters()]
    print(f"rank {rank} | {tag}: fsdp units {units} flat params {shapes}")

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29552"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    cpu = torch.device("cpu")
    report("no policy", FSDP(build(), device_id=cpu), rank)
    report("ModuleWrapPolicy({Block})",
           FSDP(build(), auto_wrap_policy=ModuleWrapPolicy({Block}), device_id=cpu), rank)
    pol = functools.partial(size_based_auto_wrap_policy, min_num_params=30)
    report("size_based(min=30)", FSDP(build(), auto_wrap_policy=pol, device_id=cpu), rank)
    dist.destroy_process_group()

if __name__ == "__main__":
    mp.spawn(worker, args=(4,), nprocs=4, join=True)

# rank 0 | no policy: fsdp units 1 flat params [(36,)]
# rank 0 | ModuleWrapPolicy({Block}): fsdp units 3 flat params [(18,), (18,)]
# rank 0 | size_based(min=30): fsdp units 5 flat params [(9,), (9,), (9,), (9,)]
# (ranks 1, 2 and 3 printed the same three lines)

The collectives FSDP issues are countable without a profiler. Wrap dist.all_gather, dist.all_gather_into_tensor and dist.reduce_scatter_tensor in counters before the forward, run one step, and read the counts off. On three blocks wrapped one unit each, the forward issues three all-gathers and the backward issues three more plus three reduce-scatters.

Each of those numbers has an obvious owner. Three all-gathers in the forward is one per unit, gathering that unit's parameters just before the layer runs and freeing them again after. Three more in the backward is the same gathering repeated, because the parameters were discarded after the forward and the backward needs them again. Three reduce-scatters is the gradient path: each unit's gradient is reduced across ranks and scattered so every rank keeps only its own shard.

That last one is the mechanical difference from DDP in a single word. DDP all-reduces, so every rank ends up with the whole gradient. FSDP reduce-scatters, so every rank ends up with the slice matching the parameters it owns, which is exactly what its shard of the optimizer state needs.

One honest note on the printout: the gathers went through the list-form dist.all_gather here rather than all_gather_into_tensor, which is the path this build took on gloo. The count is the fact being taught, not the choice of entry point.

seen = Counter()
for n in ("all_gather_into_tensor", "reduce_scatter_tensor", "all_gather", "all_reduce"):
    fn = getattr(dist, n)
    def wrapped(*a, _n=n, _fn=fn, **k):
        seen[_n] += 1
        return _fn(*a, **k)
    setattr(dist, n, wrapped)

model = nn.Sequential(Block(), Block(), Block())
f = FSDP(model, auto_wrap_policy=ModuleWrapPolicy({Block}),
         device_id=torch.device("cpu"))
out = f(torch.ones(2, 6))
print("after forward :", dict(seen))
out.sum().backward()
print("after backward:", dict(seen))
print("fsdp units    :", sum(1 for m in f.modules() if isinstance(m, FSDP)))

# after forward : {"all_gather": 3}
# after backward: {"all_gather": 6, "reduce_scatter_tensor": 3}
# fsdp units    : 4

Swap sharding_strategy=ShardingStrategy.SHARD_GRAD_OP into the same script and the backward all-gathers disappear: three in the forward, three in total, three reduce-scatters unchanged. The ShardingStrategy docstring at v2.2.2 predicts exactly that, saying SHARD_GRAD_OP unshards before the forward, does not reshard after the forward, and only reshards after the backward.

So the dial is a memory-against-bandwidth choice, and the measured counts put a number on it. FULL_SHARD gathers twice per unit per step and holds one unit unsharded at a time. SHARD_GRAD_OP gathers once per unit per step and holds the whole model unsharded between forward and backward. NO_SHARD replicates and all-reduces, which is DDP's deal wearing FSDP's interface.

The two hybrid strategies are the ones a single machine cannot show you. HYBRID_SHARD applies FULL_SHARD inside a node and replicates across nodes, so the expensive collectives stay on the fast intra-node links. That is a claim about a network this course has no second node to measure, and it stays a citation.

A flat parameter has to divide evenly by the world size, and models do not cooperate. Wrap a Linear(3, 3) with 9 parameters on a world of four and each rank reports a shard of 3, which is 12 numbers spread over four ranks. Three of those numbers are padding that exists only so the division works.

The FlatParameter docstring names both sizes for this reason, _unpadded_unsharded_size and _padded_unsharded_size, the second being the first with right-hand-side padding for divisibility by the world size. On a small model that overhead is a third of the tensor. On a real one it is a rounding error, which is why nobody notices until they wrap something tiny and the arithmetic stops matching.

This is the same discipline the sharding lessons on the jax path apply from the other side of the fence: a shard is a number you can compute before you run anything, and if the number you computed does not match the number the framework reports, one of your assumptions about the layout is wrong.

odd = nn.Linear(3, 3, bias=False)          # 9 parameters, world of 4
fo = FSDP(odd, device_id=torch.device("cpu"))
p = next(fo.parameters())
print(f"rank {rank} | 9 params over world {world}: shard {tuple(p.shape)}")

# rank 0 | 9 params over world 4: shard (3,)
# rank 1 | 9 params over world 4: shard (3,)
# rank 2 | 9 params over world 4: shard (3,)
# rank 3 | 9 params over world 4: shard (3,)
# 4 ranks x 3 = 12, so 3 of the 12 are padding

FSDP on a CPU-only build needs one argument that a GPU job never passes. Construct it without device_id and _init_device_handle falls through to torch.device("cuda", torch.cuda.current_device()), which on this build raises AssertionError: Torch not compiled with CUDA enabled from inside the FSDP constructor. Passing device_id=torch.device("cpu") is what made every run in this lesson possible, and it is the first thing to try when FSDP refuses to build on a machine with no GPU.

Three features stay cited rather than measured. Mixed precision through MixedPrecision is a real memory and bandwidth win on hardware with fast reduced-precision arithmetic, and measuring it on a CPU would produce a number that means nothing. cpu_offload moves shards to host memory between uses, which only has a point when the shards were somewhere else. limit_all_gathers throttles prefetching against the allocator, and the allocator it is protecting is CUDA's.

The other absence is a version boundary rather than a hardware one. Everything in this lesson is the FSDP that ships in 2.2.2, wrapper classes and flat parameters. The rewrite that arrived later, built on DTensor with per-parameter sharding and a fully_shard function instead of a wrapper class, is the next lesson's machinery applied to this lesson's problem. Read it in the current docs and hold the two apart by version, because the printouts here are 2.2.2 and will not match.