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_root_unsharded`	Gather the FSDP2 root's own parameters while a submodule of `model` is called directly.
`fsdp2_wrap`	Auto-detect a HF causal-LM's decoder layers and FSDP2 `fully_shard` each one.

fsdp2_root_unsharded(model)#

Gather the FSDP2 root’s own parameters while a submodule of model is called directly.

Calling a submodule skips the root’s pre-forward, which makes model the FSDP2 root and unshards the parameters it holds outside its own units (e.g. the embeddings). No-op for a model that is not an FSDP2 root.

Parameters:: model (Module)
Return type:: Iterator[None]

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.

Its callers load causal LMs and text-only VLMs; the expert-parallel dense load, which also takes VLMs, picks its blocks through _fsdp2_wrap_modules(), keeping a vision tower replicated.

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). With the root sharded, auxiliary modules beside the decoder stack (a vision tower, a projector, an MTP head; see _auxiliary_modules()) get groups of their own. 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.

It records no distribution plan, and refuses a model that, or any of whose submodules, records one, as a model load_pretrained returned does: wrap one loaded with from_pretrained.

Parameters:: cpu_offload (bool)