[OMNIML-6054] Route load_pretrained by the distribution plan

CHANGELOG.rst

@@ -33,10 +33,12 @@ Changelog
 *Misc*
 
 - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs.
+- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_FSDP_SIZE`` variable, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). Until the expert-parallel load lands, an MoE planned expert-parallel, as the default ``MODELOPT_EP_DEGREE=auto`` plans one whose config has a ``base_model_ep_plan``, is refused before it loads: set ``MODELOPT_EP_DEGREE=1`` (with ``MODELOPT_FSDP_SIZE=world`` to FSDP2-shard it) or pass a plan with ``ep_size=1``.
 
 **Backward Breaking Changes**
 
 - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
+- ``modelopt.torch.utils.distributed.fsdp2_wrap`` moved to ``modelopt.torch.utils.plugins.model_load_utils``, with the other loaders. The old import path still works, but calling it now needs the ``hf`` extra.
 - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata.
 
 **Deprecations**

modelopt/torch/utils/distributed.py

@@ -30,7 +30,7 @@
 
 import torch
 import torch.distributed
-from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard
+from torch.distributed.fsdp import FSDPModule
 from torch.distributed.tensor import DTensor
 
 __all__ = [
@@ -734,49 +734,15 @@ def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None:
 
 
 def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False):
-    """Auto-detect a HF causal-LM's decoder layers and FSDP2 ``fully_shard`` each one.
+    """FSDP2-shard a HF causal LM's decoder layers, and the root unless ``shard_root=False``.
 
-    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
-    :func:`_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 loads a model, so it lives with the loaders now:
+    :func:`modelopt.torch.utils.plugins.model_load_utils.fsdp2_wrap`. This keeps the old path, which
+    now needs the ``hf`` extra, as the loaders do.
     """
-    # Lazy import: layerwise_calib imports this module at top level (circular).
-    from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
-
-    decoder_layers = LayerActivationCollector.get_decoder_layers(model)
-    if decoder_layers is None:
-        raise RuntimeError(
-            "Could not auto-detect decoder layers; FSDP2 wrap requires a standard HF causal-LM layout."
-        )
+    from modelopt.torch.utils.plugins.model_load_utils import fsdp2_wrap as wrap
 
-    fsdp_kwargs: dict[str, Any] = {"reshard_after_forward": True}
-    if mp_policy is not None:
-        fsdp_kwargs["mp_policy"] = mp_policy
-    if cpu_offload:
-        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy()
-    ignored_params = _off_dtype_params(model)
-    if ignored_params:
-        fsdp_kwargs["ignored_params"] = ignored_params
-
-    # Snapshot/restore config.architectures: some HF builders mutate it during fully_shard.
-    config = getattr(model, "config", None)
-    architectures = list(getattr(config, "architectures", []) or [])
-    for layer in decoder_layers:
-        fully_shard(layer, **fsdp_kwargs)
-    if shard_root:
-        fully_shard(model, **fsdp_kwargs)
-    if ignored_params:
-        _move_to_fsdp_device(model, ignored_params)
-    if config is not None and architectures:
-        config.architectures = architectures
-
-    return decoder_layers
+    return wrap(model, shard_root=shard_root, mp_policy=mp_policy, cpu_offload=cpu_offload)
 
 
 def broadcast_state_dict(