FSDP2

KempnerForge uses PyTorch’s composable fully_shard() API (FSDP2), not the older FullyShardedDataParallel module. Entry point: apply_fsdp2(model, device_mesh, mp_policy=None, reshard_after_forward=True) in kempnerforge/distributed/parallel.py.

What it shards

Two levels:

1. Per-block — each TransformerBlock in model.layers is wrapped with its own fully_shard() call.

2. Top-level — the whole model is wrapped once more. This covers parameters that live outside the blocks: the token embedding, the final norm, the output head, and (for EP-MoE blocks) the layer norms that don’t get swept up by the sub-module wrap below.

Both wraps share the same MixedPrecisionPolicy and reshard_after_forward setting.

A VLM is wrapped component by component: the blocks, the transformer root, the adapter, and the vision encoder unless it is fully frozen. The encoder’s unit takes the language model’s policy by default. With [vision_encoder] mixed_precision = "bf16" or "fp32" it gets default_mp_policy of that dtype instead, so the encoder computes in that dtype while its parameters, gradients and optimizer state keep the dtype it was loaded in. An encoder FSDP2 does not wrap (frozen, or a run without a dp_shard or dp_replicate dimension, such as one GPU) gets the same from a wrapper around its forward: each forward runs on copies of its parameters cast to that dtype, with its floating-point inputs cast too and autocast off, and autograd returns the gradients to the parameters in their own dtype. Under torch.compile the wrapper runs at a graph break and the encoder’s forward compiles as its own graph, so the fp32 master weights autocast reaches neither its forward nor its backward. Either way its floating-point buffers are cast to the set dtype, and a checkpoint stores them in it, so one GPU and many store, update and compute the encoder alike. Image and video runs build the same encoder, so the setting covers every frame of a clip. It applies to the training build (build_parallel_model); build_vlm_wrapper builds an unsharded model and leaves every dtype to its caller.

HSDP

Both wraps use get_dp_mesh(device_mesh). With dp_replicate > 1 and a dp_shard dimension, that is the 2D mesh mesh["dp_replicate", "dp_shard"], and fully_shard runs hybrid sharding (HSDP):

• Without TP or EP, each parameter is a DTensor placed (Replicate(), Shard(0)): sharded across the ranks of a dp_shard group, which together hold one full copy of the model, and replicated across the dp_shard groups, so the ranks of a dp_replicate group hold the same shard. On a 2 × 4 mesh, ranks 0 to 3 hold one copy and ranks 4 to 7 another, and ranks 0 and 4 hold the same shard.

• In backward, gradients are reduce-scattered within each dp_shard group, then all-reduced within each dp_replicate group, between the ranks that hold the same shard, and averaged over all dp_replicate × dp_shard ranks.

• Each of those ranks reads its own data: get_dp_info gives dp_size = dp_replicate × dp_shard (see Device mesh § Extracting sub-meshes).

• DCP writes each shard once, not once per group, and an HSDP checkpoint loads under a different dp_replicate and dp_shard split (see Resharding).

When dp_shard resolves to 1, init_distributed still keeps a size-1 dp_shard if tp > 1 or ep > 1, so dp_replicate > 1 with TP or EP takes this path too, with one-rank dp_shard groups: the ranks of a dp_replicate group hold the same parameters, and their gradients are all-reduced over that group. A parameter TP has sharded keeps its tp placement as a third entry, for example (Replicate(), Shard(0), Shard(1)), so each rank holds its TP slice rather than a full copy. When dp_shard resolves to more than 1, the mesh has a dp_shard dimension of that size: dp_replicate=2 with tp=2 on 8 GPUs builds a (dp_replicate, dp_shard, tp) = (2, 2, 2) mesh, whose dp_shard groups have two ranks. No shipped config sets dp_replicate > 1.

Mixed precision policy

default_mp_policy(param_dtype=torch.bfloat16)
# returns
MixedPrecisionPolicy(param_dtype=bfloat16, reduce_dtype=float32)

• param_dtype — the dtype FSDP2 casts parameters to during all-gather. bfloat16 is the default. If mixed_precision = "fp8", the param dtype stays bfloat16 and FP8 is applied as a separate layer on top (see FP8); FSDP2 never sees FP8 types directly.

• reduce_dtype = float32 — gradient reductions happen in fp32 regardless of param dtype. This is the standard “bf16 compute, fp32 reduce” policy and the source of FSDP2’s numerical stability advantage over bare bf16 training.

Pass a custom MixedPrecisionPolicy as the mp_policy kwarg if you need a different combination — most of the time the default is what you want.

fp32 master weights

By default the model itself is cast to param_dtype, so with mixed_precision = "bf16" the sharded parameters, their gradients and the optimizer state (Muon’s momentum included) are all bf16, and the param_dtype cast above has nothing to do. At small learning rates many bf16 weight updates round to nothing.

[train]
mixed_precision = "bf16"
fp32_master_weights = true

keeps the model in fp32 (TrainConfig.master_dtype). FSDP2 casts each parameter shard to param_dtype before the all-gather, so compute and all-gather traffic stay bf16 and the forward matches a bf16 copy of the weights bit for bit. Gradients are reduce-scattered in fp32 and stay fp32 on the shards, and the optimizer steps fp32 parameters with fp32 state. Parameter, gradient and optimizer-state memory double. A checkpoint saved in either mode loads into the other: DCP copies each saved tensor into the live one, casting it.

FSDP2 is what does the cast, and there is no FSDP2 wrap on a single GPU or on a mesh with only a pp dimension. There, the training and in-loop eval forward passes run under torch.autocast(param_dtype) instead (TrainingSession.autocast()), which casts per op: matmuls and attention run in bf16 while the residual stream and norms stay fp32, and a vision encoder’s matmuls run in bf16 too unless [vision_encoder] mixed_precision sets its dtype, which keeps the encoder out of the autocast. Those numerics differ slightly from the FSDP2 path, and the pipeline then passes fp32 activations between stages. scripts/eval.py builds its model in param_dtype either way.

Limits of the autocast path:

• On a pp-only mesh the whole schedule.step() runs under autocast, backward included, because the schedule interleaves the two. PyTorch recommends leaving autocast before backward. It trains (one training-step test covers it), but its numerics and memory have not been compared with the FSDP2 path.

• An MoE model is rejected on a single GPU: its expert inputs would stay fp32, which turns off grouped GEMM.

fp32_master_weights with mixed_precision = "fp32" is rejected, since the weights are already fp32.

reshard_after_forward

apply_fsdp2(model, device_mesh, reshard_after_forward=True)

Value	Behavior	When to use
True (default)	All-gather on forward, reshard after, all-gather again on backward	Maximum memory savings; every forward pass reshards
False	All-gather once, keep gathered across the step	Pipeline parallel: 1F1B schedule sends many microbatches through each stage; resharding between them would be wasted work
int	Rate-limit concurrency: keep at most N blocks gathered at once	Rarely used; middle-ground for memory tuning

training/entry.py calls apply_fsdp2(model, device_mesh, mp_policy=mp_policy) on all paths without overriding reshard_after_forward, so both dense and PP runs default to True. The docstring on pipeline_parallel.py recommends False for PP (to amortize the all-gather over microbatches), but that’s not currently wired through the training script — pass it manually if you need it.

EP-MoE: per-sub-module wrapping

The MoE+EP path needs a wrapping pattern that the dense path does not:

if _has_ep_moe(layer):
    fully_shard(layer.attention, ...)
    fully_shard(layer.mlp, ...)          # MoE MLP wrapped separately
else:
    fully_shard(layer, ...)               # whole block wrapped together

Why: under EP, the MoE MLP’s backward pass issues two all-to-all calls (unpermute + backward dispatch). If FSDP2’s reduce-scatter for the whole block fires between them, different ranks end up in different communication phases — deadlock.

Wrapping attention and mlp separately means attention’s reduce- scatter fires after attention’s backward, and MLP’s reduce-scatter fires after both all-to-alls complete. All dp_shard peers share the same EP rank and so reach each phase together.

This fix is why the MoE benchmark calls out “per-sub-module FSDP wrapping” as the load-bearing optimization.

Activation checkpointing

apply_ac(model, mode) runs before apply_fsdp2 — the FSDP2 wrap sees the already-checkpointed block as its unit.

Three modes from ActivationCheckpointing:

act_ckpt	What gets rematerialized
"none"	Nothing; fastest; uses most memory
"full"	Every block, each child of layers and cross_attention_layers whatever its class (TransformerBlock, MoTBlock, MoMaBlock, CrossAttentionBlock); the entire block is recomputed on backward
"selective"	Only the attention module within each block (Attention, MoTAttention or CrossAttention); MLP activations stay materialized

apply_ac logs the classes it wrapped and returns the number of modules it wrapped. Under "selective", a block that holds none of those attention classes stays unwrapped, and apply_ac warns with the block’s class; "full" covers it. A block that already holds a CheckpointWrapper is left as it is, so a second call wraps nothing.

Wrapping adds a _checkpoint_wrapped_module. segment to the names named_modules() and named_parameters() give whatever is inside, for example transformer.layers.0._checkpoint_wrapped_module.attn.q_proj.image.weight. state_dict() and DCP checkpoints drop it, so a checkpoint loads under any mode. Freeze patterns match the named_parameters() names: the default module aliases end in * and still match, but a raw pattern that spells a path inside a wrapped module matches nothing: under full a path inside a block, such as transformer.layers.0.attn.*, and under selective a path inside an attention module, such as transformer.layers.0.attn.q_proj.*. That was already so for TransformerBlock and Attention; it now applies to the MoT, MoMa and cross-attention blocks and to MoTAttention and CrossAttention too.

All three use CheckpointImpl.NO_REENTRANT — the modern torch.utils.checkpoint implementation that composes cleanly with FSDP2 and torch.compile.

"selective" is the right default when memory is tight but you still want MLP’s full forward — attention is both the biggest activation consumer (O(seq_len^2)) and cheap to recompute (SDPA fused kernel). "full" is the fallback when even selective doesn’t fit.

Application order inside build_parallel_model

Non-PP path:

model = build_model(...)                   # on device, not meta
apply_expert_parallel(model, mesh)          # MoE only — partitions experts
apply_float8(model)                         # optional — converts nn.Linear
apply_ac(model, mode)                       # optional — wraps blocks / attention
apply_fsdp2(model, mesh, mp_policy)         # shards remaining params

TP-enabled path:

with torch.device("meta"): model = build_model(...)
apply_tensor_parallel(model, mesh)          # shards Linears with DTensor
apply_expert_parallel(model, mesh)
apply_float8(model)
apply_ac(model, mode)
apply_fsdp2(model, mesh, mp_policy)
model.to_empty(device=device)               # materialize on GPU
model.init_weights_and_freqs()              # init RoPE + weights
model.to(dtype=param_dtype)                 # bf16, or fp32 with fp32_master_weights

The meta-device init is required when TP is active: parameters must already be DTensors before they’re materialized, so FSDP2 can wrap them correctly. For non-TP runs, meta-device init is skipped — the model is built directly on the target device.

See the full breakdown in Parallelism order.

Fused optimizer compatibility

When TP or EP are active, apply_fsdp2 always runs with a dp_shard dimension in the mesh — even if that dimension has size 1. Why: the fused AdamW kernel requires every parameter to be a DTensor (not a mix of DTensor and plain tensor). If TP has already wrapped some Linears as DTensors, FSDP2 has to upgrade the rest — which requires an (even trivial) dp_shard axis to wrap along.

This is handled automatically in init_distributed — see Device mesh § Special case.

See also

• Device mesh — how the mesh that apply_fsdp2 consumes is built.

• Tensor parallelism — what apply_tensor_parallel does before FSDP2 runs.

• Expert parallelism § EP-MoE + FSDP2 — the deadlock case that per-sub-module wrapping avoids.

• Parallelism order — the full sequence of passes with reasoning.

• Gradient utilities § clip_grad_norm_ — the DTensor-aware clipping FSDP2-sharded models need.