fsdp2_utils

FSDP2 wrapping of a model’s decoder layers, which needs no Hugging Face package.

The loaders (modelopt.torch.utils.plugins.model_load_utils) wrap with fsdp2_wrap() and export it; it lives here so that a model whose decoder layers are registered with LayerActivationCollector can be wrapped without transformers or huggingface_hub.

Functions

fsdp2_wrap

Auto-detect a HF causal-LM's decoder layers and FSDP2 fully_shard each one.

fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload=False)

Auto-detect a HF causal-LM’s decoder layers and FSDP2 fully_shard each one.

By default (shard_root=True) the root module is wrapped too, so embed/lm_head/norm are sharded instead of replicated per rank; pass shard_root=False to leave the root replicated (only decoder layers sharded). Returns the detected decoder layers so callers can reuse the detection result.

Parameters whose dtype differs from the model’s dominant one are excluded from the wrap (see _off_dtype_params()), since FSDP2 rejects a shard group that mixes dtypes. They stay replicated and are moved onto the shards’ device, which fully_shard does not do for the params it ignores.

Parameters:

cpu_offload (bool)