Sharded Training: FSDP2, DTensor, Tensor and Pipeline Parallelism

Module 22 counted 16 bytes per parameter for mixed-precision AdamW: 2 for bf16 weights, 2 for bf16 gradients, and 12 for the fp32 master weights plus Adam's two moments. A 7 B model therefore needs 112 GB (104.3 GiB) before any activation, and DDP puts all of it on every rank: more GPUs add throughput but never shrink the per-rank footprint (activations come on top; training considerations).

A large run composes them, one per mesh axis, each mapped to the link its collective tolerates:

ZeRO (Rajbhandari et al., 2020) observes that DDP stores \(N\) identical copies of state while each rank needs only part. Using the paper's 2 + 2 + 12 accounting with world size \(N\):

ZeRO-1/2 floor at \(4P\) and \(2P\) because every rank still holds full bf16 weights; only ZeRO-3 keeps shrinking.

Communication differs too. A ring all-reduce is a reduce-scatter plus an all-gather (module 26), each moving about \(P\) elements per rank. ZeRO-1/2 reduce-scatter the gradients, update the shard each rank owns, and all-gather the updated parameters. That is \(2P\), the same volume as DDP. ZeRO-3 frees the parameters between uses, so it all-gathers them before the forward, all-gathers them again before the backward, and reduce-scatters the gradients. That is \(3P\), 1.5× DDP, which FSDP hides by overlapping each all-gather with the previous layer's compute.

PyTorch's ZeRO-1 is ZeroRedundancyOptimizer (docs, recipe), used next to DDP:

from torch.distributed.optim import ZeroRedundancyOptimizer
model = DDP(model)                      # gradients still all-reduced in full
opt = ZeroRedundancyOptimizer(model.parameters(), optimizer_class=torch.optim.AdamW, lr=3e-4)
# each rank steps the params it owns, then broadcasts them; before saving:
# opt.consolidate_state_dict(to=0)

It partitions whole parameters, not elements: a local two-rank probe with Linear(16, 32) → Linear(32, 4) gave rank 0 the 512-element weight (1,024 Adam state elements) and rank 1 the other three tensors (328). FSDP2 avoids that imbalance by sharding every parameter along dim 0.

FSDP2 is torch.distributed.fsdp.fully_shard (docs, source, tutorial):

fully_shard(module, *, mesh=None, reshard_after_forward=None, shard_placement_fn=None,
            mp_policy=MixedPrecisionPolicy(), offload_policy=OffloadPolicy(), ignored_params=None)

It works in place. Every parameter becomes a Shard(0) DTensor: rank \(r\) keeps chunk \(r\) of torch.chunk(param, world, dim=0), padded internally when rows do not divide evenly. FQNs do not change; the class becomes a generated FSDP<Name> that also inherits FSDPModule. Each call creates one parameter group, one unit of communication:

Two rules follow. Apply fully_shard bottom-up. A call claims every parameter no earlier call has claimed, so shard each Transformer block first, then the root, whose group gets the leftovers: embeddings, final norm, LM head. Sharding only the root makes the whole model one group: one huge all-gather, no overlap, unsharded peak memory. Build the optimizer after sharding. It then sees DTensor shards, and its state is sharded with no extra work.

reshard_after_forward controls whether each group's full parameters survive past the forward:

The root defaults to False because the backward needs its parameters first; freeing and re-gathering them would be wasted work.

Mixed precision. MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, output_dtype=None, cast_forward_inputs=True) keeps sharded parameters and optimizer in fp32, all-gathers bf16 copies (half the bytes) and reduces gradients in fp32. In the lab the shards stayed torch.float32, the output came back bf16 (cast it or set output_dtype), and three steps drifted 0.021 from the fp32 reference versus 1e-7 for fp32 FSDP.

Offload. CPUOffloadPolicy(pin_memory=True) keeps parameter shards, gradients and optimizer state in pinned host memory, and runs the optimizer step on the CPU. Capacity for step time, as in module 22 (not executed here).

Prefetching. FSDP2 overlaps the next group's all-gather with the current group's compute by default; set_modules_to_forward_prefetch([...]) and set_modules_to_backward_prefetch([...]) override the order, trading memory for overlap (GPU-only, not executed here).

Other controls. set_requires_gradient_sync(False) replaces FSDP1's no_sync() (accumulate unsharded gradients, skip the reduce-scatter). set_reshard_after_backward(False) keeps parameters gathered across microbatches. register_fsdp_forward_method(model, "generate") hooks a non-forward method. ignored_params opts parameters out; shard_placement_fn picks another (evenly divisible) shard dim. set_force_sum_reduction_for_comms(True) replaces ReduceOp.AVG with SUM plus FSDP's own division; the lab needs it. Compose per block: activation checkpointing, per-block torch.compile (module 23), fully_shard, then the root (not executed here).

The FSDP paper (Zhao et al., 2023) describes FSDP1's FlatParameter. The FSDP2 tutorial's migration guide and the FSDP notes give the differences:

A DeviceMesh is an n-D array of ranks with named axes, and it has one process group per axis slice (DTensor docs, source):

mesh = init_device_mesh("cuda", (dp, tp), mesh_dim_names=("dp", "tp"))
tp_mesh, dp_mesh = mesh["tp"], mesh["dp"]

A DTensor is a local torch.Tensor plus (mesh, placements), one placement per mesh dimension; the logical tensor is what one device would hold. A (4, 2) tensor on a 2-rank mesh:

Partial is what a matmul with a sharded contraction dimension produces: a full-shape partial sum per rank (Shard(1) @ Shard(0) → Partial(), fence 3). Ops run locally until a placement change, redistribute, forces communication; each transition is one collective:

On CPU, gloo warned it "does not support alltoall yet, falling back with allgather + chunk": correct, not bandwidth-optimal. distribute_tensor(t, mesh, placements, src_data_rank=0) scatters rank 0's copy; src_data_rank=None keeps each rank's own data. DTensor.from_local wraps existing shards without communication; full_tensor() gathers the whole tensor; to_local() returns this rank's piece. distribute_module(module, mesh, partition_fn, input_fn, output_fn) shards a module with your function; experimental local_map runs a plain-tensor function on local shards. Mixing a DTensor and a plain tensor in one op raises "got mixed torch.Tensor and DTensor".

Megatron-LM (Shoeybi et al., 2019) splits the MLP \(Y = \mathrm{GELU}(XA)\,B\), where \(X\) is (B, T, d_model), \(A\) is (d_model, d_ff) and \(B\) is (d_ff, d_model). Split \(A\) by columns, \(A = [A_1, A_2]\). GELU is elementwise, so \(\mathrm{GELU}(XA) = [\mathrm{GELU}(XA_1), \mathrm{GELU}(XA_2)]\), and each rank computes its half of the hidden layer without communicating. Split \(B\) by rows, \(B = [B_1; B_2]\):

Each rank holds a full-shape Partial; one all-reduce completes the forward. Splitting \(A\) by rows would put the sum inside the nonlinearity, \(\mathrm{GELU}(X_1A_1 + X_2A_2) \neq \mathrm{GELU}(X_1A_1) + \mathrm{GELU}(X_2A_2)\). Fence 2 measures that error at 0.2. In backward the input gradient sums over column shards, so one all-reduce runs at the MLP input: Megatron's \(f\) (identity forward, all-reduce backward) and \(g\) (the reverse). Attention is the same: Q, K, V column-parallel (each rank owns \(H/\text{tp}\) whole heads), output projection row-parallel. That is two all-reduces per layer in forward and two in backward, on the critical path, which is why TP stays inside a node.

PyTorch names each submodule's role (docs, source, tutorial):

parallelize_module(mlp, tp_mesh, {"w1": ColwiseParallel(), "w2": RowwiseParallel()})

nn.Linear stores (out_features, in_features) and computes \(XW^\top\), so "column-parallel" shards dim 0 of the stored weight, the output features. The lab asserts local shapes (4·d/2, d) for w1 and (d, 4·d/2) for w2. By default ColwiseParallel maps Replicate() input to Shard(-1) output, and RowwiseParallel maps Shard(-1) input to Replicate() via the all-reduce. On two gloo ranks, outputs and weight gradients matched the unsharded MLP to 6e-8 and 2.4e-7.

Sequence parallelism (Korthikanti et al., 2022) covers the regions TP leaves replicated: LayerNorm, dropout and residual adds hold full (B, T, d_model) activations on every TP rank. SequenceParallel() shards them along T and replaces each all-reduce with a reduce-scatter plus an all-gather: same bytes, activation memory in those regions divided by the TP degree. PrepareModuleInput and PrepareModuleOutput declare layout conversions at module boundaries. loss_parallel() computes cross-entropy on vocab-sharded logits (ColwiseParallel(output_layouts=Shard(-1), use_local_output=False) on the LM head) without materialising (B, T, V) on any rank. For serving, see the inference course's TP chapter.

FSDP × TP uses a 2-D mesh (dp, tp): apply TP with mesh["tp"], then fully_shard(..., mesh=mesh["dp"]). A colwise weight then carries two placements, (Shard(0), Shard(0)). HSDP passes a 2-D mesh to fully_shard itself. Parameters become (Replicate(), Shard(0)): sharded within a group (usually a node) and replicated across groups, so all-gathers stay on fast links and only a gradient all-reduce crosses nodes. torchtitan is the reference implementation of FSDP2 + TP (including async TP) + PP + CP with meta init, activation checkpointing, compile, DCP and low-precision training (reviewed from its README; not executed here).

Bandwidth numbers: collectives and topology.

Pipeline parallelism cuts the layer stack into \(p\) stages, one per rank; in the naive single-machine model-parallel baseline only one stage is busy at a time. GPipe (Huang et al., 2019) splits the batch into \(m\) microbatches. With one time slot per microbatch-stage forward, the last stage finishes at slot \(m+p-1\) and each stage computes in \(m\) of those slots. The backward mirrors this, so the idle fraction, the bubble, is

With \(p = m = 4\) the bubble is 3/7 ≈ 43 %, and with \(m = 32\) it is 8.6 % (fence 5).

Each row has \(2(m+p-1) = 14\) slots, and every stage sits idle for \(2(p-1) = 6\) of them.

1F1B caps activation memory, not the bubble, so you can raise \(m\), which does shrink it. Interleaving (Narayanan et al., 2021) gives each rank \(v\) non-adjacent layer chunks.

The API (docs, source, tutorial) offers a tracing front end, pipeline(module, mb_args=(x_mb,), split_spec={"layers.2": SplitPoint.BEGINNING}), which cuts an exported graph. The manual front end, used in the lab, wraps a submodule you slice yourself:

stage = PipelineStage(layers[rank*k:(rank+1)*k], stage_index=rank, num_stages=world, device=dev)
sched = ScheduleGPipe(stage, n_microbatches=4, loss_fn=F.mse_loss)
if rank == 0:            sched.step(x)                      # first stage feeds input
elif rank == world - 1:  sched.step(target=y, losses=out)   # last stage computes loss
else:                    sched.step()

The schedule splits x into microbatches, runs the backward, and with the default scale_grads=True divides gradients by n_microbatches. In the lab the last stage's four losses matched the single-process ones exactly. Serving-side pipelining: inference course.

Context parallelism shards T for every activation; attention, the only op mixing positions, needs every rank's K and V. Ring Attention (Liu et al., 2023) passes K/V shards around a ring and merges the partial results with the online-softmax rescaling from hardware-aware attention. In 2.8, torch.distributed.tensor.experimental.context_parallel(mesh, buffers=..., buffer_seq_dims=...) swaps scaled_dot_product_attention for ring attention. Every sequence-indexed buffer must be listed (tutorial); omitting RoPE's freq_cis breaks rotary embeddings. Experimental, not executed here.

Expert parallelism places an MoE layer's experts (mixture of experts) on different ranks: one all_to_all dispatches routed tokens to their experts' ranks and a second combines the outputs. Monarch (tutorial) is a prototype single-controller actor framework that drives a whole process mesh from one script instead of SPMD.

torch.save(model.state_dict()) on an FSDP2 model saves this rank's DTensor shards: from rank 0 alone it loses the rest; from every rank it is tied to one layout. Distributed Checkpoint (docs, recipe) fixes both:

model_sd, optim_sd = get_state_dict(model, opt)      # FQN-keyed DTensor state
dcp.save({"model": model_sd, "optim": optim_sd}, checkpoint_id=ckpt_dir)
state = dict(zip(("model", "optim"), get_state_dict(model, opt)))  # target layout, any world size
dcp.load(state, checkpoint_id=ckpt_dir)              # in place
set_state_dict(model, opt, model_state_dict=state["model"], optim_state_dict=state["optim"])

Each rank writes its own shard files; a coordinator writes .metadata, mapping every FQN's global shape and chunks to file offsets. On load each rank reads only the chunks overlapping its target placement, so resharding (2 ranks to 1, FSDP to FSDP+TP) is planning, not conversion. get_state_dict keys state by normalised FQN, so checkpoints survive rewrapping. dcp.async_save stages the state in CPU memory and writes it from a background thread, returning a Future (async recipe). format_utils.dcp_to_torch_save(dir, file) consolidates a checkpoint into one torch.save file for export. Single-process checkpoints: module 16; capstone use: module 37.

Fence 1 recomputes the ZeRO table and the 1.5× ratio: arithmetic, not a framework measurement.

P = 7e9  # bytes per param: 2 bf16 weights + 2 bf16 grads + 12 fp32 master/m/v

def per_rank_bytes(stage, n):
    params, grads, optim = 2 * P, 2 * P, 12 * P
    optim /= n if stage >= 1 else 1
    grads /= n if stage >= 2 else 1
    params /= n if stage >= 3 else 1
    return params + grads + optim

for n in (1, 8, 64):
    row = [per_rank_bytes(s, n) for s in range(4)]
    assert all(a >= b for a, b in zip(row, row[1:])), row
    assert row[1] == 4 * P + 12 * P / n and row[2] == 2 * P + 14 * P / n
assert per_rank_bytes(3, 8) == 16 * P / 8 and per_rank_bytes(0, 64) == 112e9
ddp, zero12, zero3 = 2 * P, P + P, 3 * P  # AR = RS + AG; ZeRO-3 adds a 2nd all-gather
assert zero12 == ddp and zero3 / ddp == 1.5
print({n: [round(per_rank_bytes(s, n) / 2**30, 1) for s in range(4)] for n in (1, 8, 64)})

Fence 2 simulates two TP ranks in one process: it proves the column/row algebra and that the row-first split fails, but exercises no parallelize_module or collective (the lab does).

import torch
import torch.nn.functional as F

torch.manual_seed(0)
d, ranks = 8, 2
up = torch.nn.Linear(d, 4 * d, bias=False)    # weight (4d, d) = (out, in)
down = torch.nn.Linear(4 * d, d, bias=False)  # weight (d, 4d)
X = torch.randn(3, 5, d)                      # (B, T, d_model)
ref = down(F.gelu(up(X)))

A_shards = up.weight.chunk(ranks, dim=0)      # column-parallel: OUT features, (2d, d)
B_shards = down.weight.chunk(ranks, dim=1)    # row-parallel: IN features, (d, 2d)
partials = [F.gelu(X @ A.T) @ B.T for A, B in zip(A_shards, B_shards)]
assert partials[0].shape == ref.shape         # each rank: a full-shape addend
torch.testing.assert_close(sum(partials), ref, atol=1e-5, rtol=1e-5)  # = the all-reduce

A_rows = up.weight.chunk(ranks, dim=1)        # wrong: split IN features first
wrong = sum(F.gelu(x @ A.T) for x, A in zip(X.chunk(ranks, -1), A_rows)) @ down.weight.T
assert not torch.allclose(wrong, ref, atol=1e-3)
print("colwise->rowwise exact; rowwise-first error", round((wrong - ref).abs().max().item(), 3))

Fence 3 uses a one-rank mesh, so every collective is trivial; it proves the API contract (placements, redistribute, full_tensor, the mixed-operand error, Shard(1) @ Shard(0) → Partial), not data movement.

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor

dist.init_process_group("gloo", store=dist.HashStore(), rank=0, world_size=1)
try:
    torch.manual_seed(0)
    mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp",))
    t = torch.arange(8.0).reshape(4, 2)
    dt = distribute_tensor(t, mesh, [Shard(0)])
    assert isinstance(dt, DTensor) and dt.placements == (Shard(dim=0),)
    rep = dt.redistribute(mesh, [Replicate()])
    assert torch.equal(rep.to_local(), t) and torch.equal(dt.full_tensor(), t)
    try:
        dt @ torch.ones(2, 3)
        raise AssertionError("mixing Tensor and DTensor should fail")
    except RuntimeError as err:
        assert "mixed torch.Tensor and DTensor" in str(err)
    out = dt @ distribute_tensor(torch.ones(2, 3), mesh, [Replicate()])
    assert out.placements == (Shard(dim=0),)
    a = distribute_tensor(torch.randn(4, 6), mesh, [Shard(1)])
    b = distribute_tensor(torch.randn(6, 3), mesh, [Shard(0)])
    assert (a @ b).placements == (Partial(),)  # sharded contraction -> partial sums
    print(dt.placements, rep.placements, out.placements, (a @ b).placements)
finally:
    dist.destroy_process_group()

Fence 4 applies fully_shard at world size 1; it proves the API contract (DTensor parameters, unchanged FQNs, the FSDP<Name> class, sharded optimizer state, identical AdamW math) but cannot catch a gradient-scaling bug, which is what the two-rank lab found.

import copy
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import FSDPModule, fully_shard
from torch.distributed.tensor import DTensor

dist.init_process_group("gloo", store=dist.HashStore(), rank=0, world_size=1)
try:
    torch.manual_seed(0)
    plain = nn.Sequential(nn.Linear(8, 16), nn.GELU(), nn.Linear(16, 4))
    sharded = copy.deepcopy(plain)
    names = [n for n, _ in plain.named_parameters()]
    mesh = init_device_mesh("cpu", (1,))
    for layer in (sharded[0], sharded[2]):    # bottom-up: leaves, then root
        fully_shard(layer, mesh=mesh)
    fully_shard(sharded, mesh=mesh)
    assert isinstance(sharded, FSDPModule) and type(sharded).__name__ == "FSDPSequential"
    assert [n for n, _ in sharded.named_parameters()] == names
    assert all(isinstance(p, DTensor) for p in sharded.parameters())
    opts = [torch.optim.AdamW(m.parameters(), lr=1e-2) for m in (plain, sharded)]
    x, y = torch.randn(6, 8), torch.randn(6, 4)
    for model, opt in zip((plain, sharded), opts):
        nn.functional.mse_loss(model(x), y).backward()
        opt.step()
    assert isinstance(opts[1].state[next(sharded.parameters())]["exp_avg"], DTensor)
    for p, q in zip(plain.parameters(), sharded.parameters()):
        torch.testing.assert_close(q.full_tensor(), p.detach(), atol=1e-6, rtol=0)
    print(type(sharded).__name__, "params after one AdamW step match")
finally:
    dist.destroy_process_group()

Fence 5 checks the bubble formula with exact fractions, not any schedule's real timing.

from fractions import Fraction

def bubble(p, m):  # idle p-1 of m+p-1 slots, in forward and again in backward
    return Fraction(p - 1, m + p - 1)

assert bubble(4, 4) == Fraction(3, 7)
fracs = [bubble(4, m) for m in range(1, 33)]
assert all(a > b for a, b in zip(fracs, fracs[1:]))  # more microbatches, smaller bubble
assert bubble(4, 1) == Fraction(3, 4) and bubble(4, 32) < Fraction(1, 10)
print({m: round(float(bubble(4, m)), 3) for m in (1, 4, 8, 16, 32)})

./code/27_fsdp_tp_gloo.py spawns two gloo processes (file:// rendezvous) and compares each experiment with a same-seed single-process reference. Tested with Python 3.11 and PyTorch 2.8.0 on CPU.

python content/courses/pytorch/code/27_fsdp_tp_gloo.py --world-size 2
python site/test_pytorch_27_fsdp_tp_gloo.py

• FSDP2. Four MLP-residual blocks (d_model = 32) plus a head, sharded per block then root; three AdamW steps on half of an 8-row batch per rank; plus a bf16 MixedPrecisionPolicy run.
• TP and redistribute. ColwiseParallel/RowwiseParallel on an MLP (outputs and weight gradients); six placement transitions on a (4, 2) tensor.
• PP. One PipelineStage per rank under ScheduleGPipe(n_microbatches=4).
• DCP. Save model + optimizer; reload into a differently initialised sharded model, then into a plain model (2 shards → 1) and via dcp_to_torch_save.

One local PyTorch 2.8 CPU run on an Apple-silicon laptop produced:

The ratio exceeds 0.5 because the head's (1, 32) weight and (1,) bias cannot split: rank 1 holds an empty (0, 32) shard.

The bug Adam hid. The first version of this lab reported fsdp_max_diff = 1.6e-4, plausible as reduction-order noise, yet the gradients were exactly 2× the reference. FSDP2 reduce-scatters fp32 gradients with ReduceOp.AVG. In 2.8, gloo's reduce_scatter_tensor_coalesced (which _reduce_scatter_base, and so reduce_scatter_tensor, also calls) builds AllreduceOptions with that op and then calls allreduce(inp) without them, so AVG becomes SUM (ProcessGroupGloo.cpp, v2.8.0). A direct probe confirmed it: reduce_scatter_tensor(op=AVG) summed while all_reduce(op=AVG) averaged. Adam is scale-invariant except through eps and weight decay, so doubled gradients moved the parameters by only 1.6e-4 (SGD would have shown a 2× step). The workaround is set_force_sum_reduction_for_comms(True) on every FSDP module, so FSDP reduces with SUM and divides by the world size itself. gloo_avg_grad_ratio keeps the bug visible, and the test accepts 2.0 or a fixed 1.0. NCCL supports AVG natively (not executed here). Lesson: compare a sharded run's gradients with a reference, not just Adam-trained parameters.

Scope: sharding semantics on gloo/CPU only; no memory savings, overlap, prefetching, HSDP or NCCL. On gloo reduce_scatter_tensor is a full all-reduce plus a chunk, so even the communication volume is not FSDP's. The test runs the script via subprocess (120 s timeout) and asserts FSDP < 1e-5, bf16 < 5e-2, ratio in [0.45, 0.55], TP and PP < 1e-5, all redistribute checks, and DCP < 1e-7.