From 2dc637d61cff1cdde12e4dffcfba1eb75d9e4e11 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 07:07:52 +0000
Subject: [PATCH 1/9] Move fsdp2_wrap to the model loaders, keeping its old
 path

The FSDP2 load in model_load_utils wraps its model with fsdp2_wrap, so the
wrap moves there, beside the load that calls it. It reaches the dtype helpers
it uses from distributed, and imports the layer detection where it uses it:
importing the quantization package registers every mode and plugin, which only
a wrap needs. modelopt.torch.utils.distributed.fsdp2_wrap forwards to it, so
callers of the old path keep working, given the hf extra the loaders need.

test_fsdp2_wrap_keeps_its_old_path checks the old path forwards every argument
and the result.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst                                 |  1 +
 modelopt/torch/utils/distributed.py           | 46 ++---------------
 .../torch/utils/plugins/model_load_utils.py   | 50 ++++++++++++++++++-
 tests/unit/torch/utils/test_distributed.py    | 15 ++++++
 4 files changed, 69 insertions(+), 43 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 89fb2a8b8d7..a0ecf9c2cf6 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -51,6 +51,7 @@ Changelog
 - Bump minimum required PyTorch version to 2.10.
 - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra.
 - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later.
+- ``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.
 - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm.
 - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly.
diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 26892ef112d..1e032865858 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/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__ = [
@@ -776,49 +776,13 @@ 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.
+    """Forward to :func:`~modelopt.torch.utils.plugins.model_load_utils.fsdp2_wrap`.
 
-    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 moved there with the FSDP2 load that wraps with it, so this old path needs the ``hf`` extra.
     """
-    # 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(
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index b9cb961317b..b6b5b6d1875 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -31,7 +31,7 @@
     from collections.abc import Iterable
 from torch.distributed.checkpoint import HuggingFaceStorageReader
 from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
-from torch.distributed.fsdp import FSDPModule
+from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard
 from torch.distributed.tensor import DTensor
 from transformers import (
     AutoConfig,
@@ -49,9 +49,10 @@
 from modelopt.torch.utils import env
 from modelopt.torch.utils.distributed import (
     DistributionPlan,
+    _move_to_fsdp_device,
+    _off_dtype_params,
     barrier,
     broadcast_state_dict,
-    fsdp2_wrap,
     is_initialized,
 )
 from modelopt.torch.utils.logging import print_rank_0, warn_rank_0
@@ -504,6 +505,51 @@ def load_pretrained(
     return model
 
 
+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.
+
+    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:`~modelopt.torch.utils.distributed._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.
+    """
+    # Importing the quantization package registers every mode and plugin; only a wrap needs it.
+    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."
+        )
+
+    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
+
+
 def _map_ckpt_keys_to_model(model: "torch.nn.Module", ckpt_keys, model_keys) -> dict:
     """Map each checkpoint key to the model ``state_dict`` key it feeds.
 
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index 25b3bc12268..df1d0066f20 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -192,6 +192,21 @@ def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(cpu_dist_workers)
     cpu_dist_workers.run(_no_plan_on_several_ranks_job)
 
 
+def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
+    """fsdp2_wrap lives with the loaders that wrap with it; callers of the old path still reach it."""
+    model_load_utils = pytest.importorskip("modelopt.torch.utils.plugins.model_load_utils")
+    calls = []
+    monkeypatch.setattr(
+        model_load_utils,
+        "fsdp2_wrap",
+        lambda model, **kwargs: calls.append((model, kwargs)) or "layers",
+    )
+    model = nn.Linear(2, 2)
+
+    assert dist.fsdp2_wrap(model, shard_root=False, cpu_offload=True) == "layers"
+    assert calls == [(model, {"shard_root": False, "mp_policy": None, "cpu_offload": True})]
+
+
 # ------------------------------------------------------------------------------------ collectives
 
 

From 480c3405ec07d7d86f39192122ce4becfe154fa3 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 07:22:37 +0000
Subject: [PATCH 2/9] Name the hf extra when fsdp2_wrap's old path cannot reach
 the loaders

modelopt.torch.utils.distributed.fsdp2_wrap forwards to model_load_utils,
which imports transformers and accelerate at module scope, so without the hf
extra the old path failed on the missing transformers module, with nothing
saying it moved or which extra it needs. It now raises an ImportError naming
the move and the extra. test_fsdp2_wrap_names_the_extra_its_old_path_needs
fails without it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py        |  9 +++++++--
 tests/unit/torch/utils/test_distributed.py | 11 +++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 1e032865858..c39f3a12d59 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -780,8 +780,13 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False
 
     It moved there with the FSDP2 load that wraps with it, so this old path needs the ``hf`` extra.
     """
-    from modelopt.torch.utils.plugins.model_load_utils import fsdp2_wrap as wrap
-
+    try:
+        from modelopt.torch.utils.plugins.model_load_utils import fsdp2_wrap as wrap
+    except ImportError as error:
+        raise ImportError(
+            "fsdp2_wrap moved to modelopt.torch.utils.plugins.model_load_utils, which needs the hf "
+            "extra (pip install 'nvidia-modelopt[hf]') to import."
+        ) from error
     return wrap(model, shard_root=shard_root, mp_policy=mp_policy, cpu_offload=cpu_offload)
 
 
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index df1d0066f20..c315671cc83 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -18,6 +18,7 @@
 
 import os
 import pickle
+import sys
 from contextlib import nullcontext
 from datetime import timedelta
 
@@ -207,6 +208,16 @@ def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
     assert calls == [(model, {"shard_root": False, "mp_policy": None, "cpu_offload": True})]
 
 
+def test_fsdp2_wrap_names_the_extra_its_old_path_needs(monkeypatch):
+    """Without the hf extra, the old path says what it needs rather than which module is missing."""
+    monkeypatch.setitem(sys.modules, "modelopt.torch.utils.plugins.model_load_utils", None)
+
+    with pytest.raises(
+        ImportError, match=r"needs the hf extra \(pip install 'nvidia-modelopt\[hf\]'\)"
+    ):
+        dist.fsdp2_wrap(nn.Linear(2, 2))
+
+
 # ------------------------------------------------------------------------------------ collectives
 
 

From e6db81b1699d8fb398717788dd2f05a13d3503f8 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 07:23:45 +0000
Subject: [PATCH 3/9] Move fsdp2_wrap's dtype helpers with it

_off_dtype_params and _move_to_fsdp_device have no caller but fsdp2_wrap, so
they move beside it into model_load_utils rather than stay private names in
distributed that another module reaches for; they take model_load_utils's own
warn_rank_0, which distributed had to import lazily. Their tests move into
test_model_load_utils.py unchanged, and fsdp2_wrap's docstring names the helper
in its own module.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py           | 62 ------------------
 .../torch/utils/plugins/model_load_utils.py   | 64 +++++++++++++++++--
 tests/unit/torch/utils/test_distributed.py    | 42 ------------
 .../unit/torch/utils/test_model_load_utils.py | 45 +++++++++++++
 4 files changed, 104 insertions(+), 109 deletions(-)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index c39f3a12d59..046c9f994a5 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -713,68 +713,6 @@ def is_fsdp2_model(model) -> bool:
     return any(isinstance(m, FSDPModule) for m in model.modules())
 
 
-def _off_dtype_params(model) -> set[torch.nn.Parameter]:
-    """Params whose dtype differs from the model's dominant (by element count) param dtype.
-
-    FSDP2 needs one dtype per shard group, but HF models routinely keep a few params in fp32 for
-    stability (e.g. MoE router gates). Pass these to ``fully_shard(ignored_params=...)``.
-
-    TODO: Drop this and shard the off-dtype params once a stable PyTorch release includes FSDP2
-    mixed-precision parameter dtype support (already on nightly).
-    """
-    numel_by_dtype: dict[torch.dtype, int] = {}
-    for param in model.parameters():
-        numel_by_dtype[param.dtype] = numel_by_dtype.get(param.dtype, 0) + param.numel()
-    if len(numel_by_dtype) <= 1:
-        return set()
-
-    # Lazy import: logging imports this module at top level (circular).
-    from modelopt.torch.utils.logging import warn_rank_0
-
-    dominant = max(numel_by_dtype, key=lambda d: numel_by_dtype[d])
-    off_dtype = {n: p for n, p in model.named_parameters() if p.dtype != dominant}
-    off_numel = sum(numel_by_dtype[d] for d in numel_by_dtype if d != dominant)
-    names = sorted(off_dtype)
-    warn_rank_0(
-        f"Model has mixed parameter dtypes {set(numel_by_dtype)}; FSDP2 needs one dtype per shard "
-        f"group, so {len(names)} non-{dominant} parameter(s) "
-        f"({100 * off_numel / sum(numel_by_dtype.values()):.2f}% of elements) will stay replicated "
-        f"rather than sharded: {names[:3]}{' ...' if len(names) > 3 else ''}"
-    )
-    return set(off_dtype.values())
-
-
-def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None:
-    """Move ``params`` onto the device FSDP2 computes on for ``model``.
-
-    ``fully_shard`` only moves the params it manages, so ignored ones would be stranded on
-    whatever device the caller built the model on. Meta params are left alone for deferred init.
-
-    The device comes from a sharded param's mesh rather than its local shard: under
-    ``cpu_offload`` the shard rests on CPU while compute still happens on the accelerator.
-    """
-    # Lazy import: logging imports this module at top level (circular).
-    from modelopt.torch.utils.logging import warn_rank_0
-
-    mesh = next((p.device_mesh for p in model.parameters() if isinstance(p, DTensor)), None)
-    if mesh is None:
-        warn_rank_0(
-            f"FSDP2 sharded no parameter of {type(model).__name__}, so the compute device for "
-            f"{len(params)} unsharded off-dtype parameter(s) cannot be determined; leaving them "
-            "where they are. Move them to the compute device or the forward will fail."
-        )
-        return
-
-    device = (
-        torch.device("cpu")
-        if mesh.device_type == "cpu"
-        else torch.device(mesh.device_type, getattr(torch, mesh.device_type).current_device())
-    )
-    for param in params:
-        if not param.is_meta and param.device != device:
-            param.data = param.data.to(device)
-
-
 def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False):
     """Forward to :func:`~modelopt.torch.utils.plugins.model_load_utils.fsdp2_wrap`.
 
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index b6b5b6d1875..2b86e8902dd 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -49,8 +49,6 @@
 from modelopt.torch.utils import env
 from modelopt.torch.utils.distributed import (
     DistributionPlan,
-    _move_to_fsdp_device,
-    _off_dtype_params,
     barrier,
     broadcast_state_dict,
     is_initialized,
@@ -505,6 +503,62 @@ def load_pretrained(
     return model
 
 
+def _off_dtype_params(model) -> set[torch.nn.Parameter]:
+    """Params whose dtype differs from the model's dominant (by element count) param dtype.
+
+    FSDP2 needs one dtype per shard group, but HF models routinely keep a few params in fp32 for
+    stability (e.g. MoE router gates). Pass these to ``fully_shard(ignored_params=...)``.
+
+    TODO: Drop this and shard the off-dtype params once a stable PyTorch release includes FSDP2
+    mixed-precision parameter dtype support (already on nightly).
+    """
+    numel_by_dtype: dict[torch.dtype, int] = {}
+    for param in model.parameters():
+        numel_by_dtype[param.dtype] = numel_by_dtype.get(param.dtype, 0) + param.numel()
+    if len(numel_by_dtype) <= 1:
+        return set()
+
+    dominant = max(numel_by_dtype, key=lambda d: numel_by_dtype[d])
+    off_dtype = {n: p for n, p in model.named_parameters() if p.dtype != dominant}
+    off_numel = sum(numel_by_dtype[d] for d in numel_by_dtype if d != dominant)
+    names = sorted(off_dtype)
+    warn_rank_0(
+        f"Model has mixed parameter dtypes {set(numel_by_dtype)}; FSDP2 needs one dtype per shard "
+        f"group, so {len(names)} non-{dominant} parameter(s) "
+        f"({100 * off_numel / sum(numel_by_dtype.values()):.2f}% of elements) will stay replicated "
+        f"rather than sharded: {names[:3]}{' ...' if len(names) > 3 else ''}"
+    )
+    return set(off_dtype.values())
+
+
+def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None:
+    """Move ``params`` onto the device FSDP2 computes on for ``model``.
+
+    ``fully_shard`` only moves the params it manages, so ignored ones would be stranded on
+    whatever device the caller built the model on. Meta params are left alone for deferred init.
+
+    The device comes from a sharded param's mesh rather than its local shard: under
+    ``cpu_offload`` the shard rests on CPU while compute still happens on the accelerator.
+    """
+    mesh = next((p.device_mesh for p in model.parameters() if isinstance(p, DTensor)), None)
+    if mesh is None:
+        warn_rank_0(
+            f"FSDP2 sharded no parameter of {type(model).__name__}, so the compute device for "
+            f"{len(params)} unsharded off-dtype parameter(s) cannot be determined; leaving them "
+            "where they are. Move them to the compute device or the forward will fail."
+        )
+        return
+
+    device = (
+        torch.device("cpu")
+        if mesh.device_type == "cpu"
+        else torch.device(mesh.device_type, getattr(torch, mesh.device_type).current_device())
+    )
+    for param in params:
+        if not param.is_meta and param.device != device:
+            param.data = param.data.to(device)
+
+
 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.
 
@@ -514,9 +568,9 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False
     detection result.
 
     Parameters whose dtype differs from the model's dominant one are excluded from the wrap (see
-    :func:`~modelopt.torch.utils.distributed._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.
+    :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.
     """
     # Importing the quantization package registers every mode and plugin; only a wrap needs it.
     from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index c315671cc83..23d1e1fa605 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -36,7 +36,6 @@
     DistributedState,
     DistributionPlan,
     _deserialize,
-    _off_dtype_params,
     _serialize,
     all_reduce,
     broadcast_from,
@@ -46,47 +45,6 @@
 )
 
 
-def _model(*sizes_and_dtypes) -> nn.Module:
-    model = nn.Module()
-    for i, (numel, dtype) in enumerate(sizes_and_dtypes):
-        model.register_parameter(f"p{i}", nn.Parameter(torch.zeros(numel, dtype=dtype)))
-    return model
-
-
-def test_off_dtype_params_uniform_is_empty(recwarn):
-    model = _model((8, torch.bfloat16), (4, torch.bfloat16))
-    assert _off_dtype_params(model) == set()
-    assert not recwarn.list
-
-
-def test_off_dtype_params_returns_only_the_minority():
-    # An fp32 MoE router gate next to bf16 weights, as Nemotron-3-Nano ships it.
-    model = _model((100, torch.bfloat16), (5, torch.float32))
-    with pytest.warns(UserWarning, match="mixed parameter dtypes"):
-        off = _off_dtype_params(model)
-    assert off == {model.p1}
-
-
-@pytest.mark.parametrize(
-    ("params", "expected"),
-    [
-        # Dominant dtype is by element count, not parameter count: three small fp32 params
-        # lose to one large bf16 param, and vice versa.
-        ([(100, torch.bfloat16), (5, torch.float32), (5, torch.float32), (5, torch.float32)], "p0"),
-        (
-            [(100, torch.float32), (5, torch.bfloat16), (5, torch.bfloat16), (5, torch.bfloat16)],
-            "p0",
-        ),
-    ],
-)
-def test_off_dtype_params_dominant_is_by_numel(params, expected):
-    model = _model(*params)
-    with pytest.warns(UserWarning, match="mixed parameter dtypes"):
-        off = _off_dtype_params(model)
-    kept = {n for n, p in model.named_parameters() if p not in off}
-    assert kept == {expected}
-
-
 class _Gadget:
     def __reduce__(self):
         return (os.system, ("exit 1",))
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index ca9e71b97f5..432d15940b2 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -64,6 +64,7 @@
     _auto_class,
     _conversion_plan,
     _convert_keys,
+    _off_dtype_params,
     _resolve_target,
     build_meta_causal_lm,
     record_unplaced_source_keys,
@@ -188,6 +189,50 @@ def test_meta_qwen_vlm_preserves_architecture_and_vision(model_type):
     assert all(parameter.is_meta for parameter in model.parameters())
 
 
+# ----------------------------------------------------------------------- fsdp2_wrap's dtype helpers
+
+
+def _model(*sizes_and_dtypes) -> torch.nn.Module:
+    model = torch.nn.Module()
+    for i, (numel, dtype) in enumerate(sizes_and_dtypes):
+        model.register_parameter(f"p{i}", torch.nn.Parameter(torch.zeros(numel, dtype=dtype)))
+    return model
+
+
+def test_off_dtype_params_uniform_is_empty(recwarn):
+    model = _model((8, torch.bfloat16), (4, torch.bfloat16))
+    assert _off_dtype_params(model) == set()
+    assert not recwarn.list
+
+
+def test_off_dtype_params_returns_only_the_minority():
+    # An fp32 MoE router gate next to bf16 weights, as Nemotron-3-Nano ships it.
+    model = _model((100, torch.bfloat16), (5, torch.float32))
+    with pytest.warns(UserWarning, match="mixed parameter dtypes"):
+        off = _off_dtype_params(model)
+    assert off == {model.p1}
+
+
+@pytest.mark.parametrize(
+    ("params", "expected"),
+    [
+        # Dominant dtype is by element count, not parameter count: three small fp32 params
+        # lose to one large bf16 param, and vice versa.
+        ([(100, torch.bfloat16), (5, torch.float32), (5, torch.float32), (5, torch.float32)], "p0"),
+        (
+            [(100, torch.float32), (5, torch.bfloat16), (5, torch.bfloat16), (5, torch.bfloat16)],
+            "p0",
+        ),
+    ],
+)
+def test_off_dtype_params_dominant_is_by_numel(params, expected):
+    model = _model(*params)
+    with pytest.warns(UserWarning, match="mixed parameter dtypes"):
+        off = _off_dtype_params(model)
+    kept = {n for n, p in model.named_parameters() if p not in off}
+    assert kept == {expected}
+
+
 # -------------------------------------------------------------- deciding how a model is distributed
 
 

From 5fa18febd6af46cdfa066e68d68164b68342d637 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 07:36:58 +0000
Subject: [PATCH 4/9] Name the hf extra only when transformers or accelerate is
 missing

The old path's import error named the hf extra for any ImportError from
model_load_utils, so one from inside it, or from a dependency that moved,
sent a user who has the extra to install it again. It now names the extra
only when transformers or accelerate is the missing module, and lets any
other import error through as it is; the test has a row for each.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py        |  4 +++-
 tests/unit/torch/utils/test_distributed.py | 20 ++++++++++++++------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 046c9f994a5..aa80a8f7304 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -720,7 +720,9 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False
     """
     try:
         from modelopt.torch.utils.plugins.model_load_utils import fsdp2_wrap as wrap
-    except ImportError as error:
+    except ModuleNotFoundError as error:
+        if (error.name or "").split(".")[0] not in ("transformers", "accelerate"):
+            raise
         raise ImportError(
             "fsdp2_wrap moved to modelopt.torch.utils.plugins.model_load_utils, which needs the hf "
             "extra (pip install 'nvidia-modelopt[hf]') to import."
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index 23d1e1fa605..a2c2ba4f891 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -166,13 +166,21 @@ def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
     assert calls == [(model, {"shard_root": False, "mp_policy": None, "cpu_offload": True})]
 
 
-def test_fsdp2_wrap_names_the_extra_its_old_path_needs(monkeypatch):
-    """Without the hf extra, the old path says what it needs rather than which module is missing."""
-    monkeypatch.setitem(sys.modules, "modelopt.torch.utils.plugins.model_load_utils", None)
+@pytest.mark.parametrize(
+    ("missing", "message"),
+    [
+        ("transformers", r"needs the hf extra \(pip install 'nvidia-modelopt\[hf\]'\)"),
+        ("modelopt.torch.utils.plugins.model_load_utils", "halted; None in sys.modules"),
+    ],
+    ids=["hf-extra", "another-import-error"],
+)
+def test_fsdp2_wrap_names_the_extra_its_old_path_needs(monkeypatch, missing, message):
+    """Without the hf extra, the old path says what it needs rather than which module is missing;
+    any other import error reaches the caller as it is."""
+    monkeypatch.delitem(sys.modules, "modelopt.torch.utils.plugins.model_load_utils", raising=False)
+    monkeypatch.setitem(sys.modules, missing, None)
 
-    with pytest.raises(
-        ImportError, match=r"needs the hf extra \(pip install 'nvidia-modelopt\[hf\]'\)"
-    ):
+    with pytest.raises(ImportError, match=message):
         dist.fsdp2_wrap(nn.Linear(2, 2))
 
 

From 447cfb50a2777ccaf7c99f8541bf7967600eff92 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 07:47:59 +0000
Subject: [PATCH 5/9] Name the hf extra when huggingface_hub is missing too

model_load_utils imports huggingface_hub at module scope, before transformers, so an
install without it fails on that import first.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py        | 3 ++-
 tests/unit/torch/utils/test_distributed.py | 8 ++++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index aa80a8f7304..9fa3d44a0ee 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -721,7 +721,8 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False
     try:
         from modelopt.torch.utils.plugins.model_load_utils import fsdp2_wrap as wrap
     except ModuleNotFoundError as error:
-        if (error.name or "").split(".")[0] not in ("transformers", "accelerate"):
+        missing = (error.name or "").split(".")[0]
+        if missing not in ("transformers", "accelerate", "huggingface_hub"):
             raise
         raise ImportError(
             "fsdp2_wrap moved to modelopt.torch.utils.plugins.model_load_utils, which needs the hf "
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index a2c2ba4f891..1ee1494d63f 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -166,13 +166,17 @@ def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
     assert calls == [(model, {"shard_root": False, "mp_policy": None, "cpu_offload": True})]
 
 
+_HF_EXTRA = r"needs the hf extra \(pip install 'nvidia-modelopt\[hf\]'\)"
+
+
 @pytest.mark.parametrize(
     ("missing", "message"),
     [
-        ("transformers", r"needs the hf extra \(pip install 'nvidia-modelopt\[hf\]'\)"),
+        ("transformers", _HF_EXTRA),
+        ("huggingface_hub", _HF_EXTRA),
         ("modelopt.torch.utils.plugins.model_load_utils", "halted; None in sys.modules"),
     ],
-    ids=["hf-extra", "another-import-error"],
+    ids=["transformers", "huggingface_hub", "another-import-error"],
 )
 def test_fsdp2_wrap_names_the_extra_its_old_path_needs(monkeypatch, missing, message):
     """Without the hf extra, the old path says what it needs rather than which module is missing;

From 7f16c2c216e7a77b0db940ecd49c910397999973 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 08:01:14 +0000
Subject: [PATCH 6/9] Name the hf extra only when a whole package is missing

A module missing from an installed transformers is a version mismatch, which the extra's
install message would hide.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 modelopt/torch/utils/distributed.py        | 4 ++--
 tests/unit/torch/utils/test_distributed.py | 6 ++++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 9fa3d44a0ee..9fd804e5f1d 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -721,8 +721,8 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False
     try:
         from modelopt.torch.utils.plugins.model_load_utils import fsdp2_wrap as wrap
     except ModuleNotFoundError as error:
-        missing = (error.name or "").split(".")[0]
-        if missing not in ("transformers", "accelerate", "huggingface_hub"):
+        # A module missing from an installed package is a version mismatch, not a missing extra.
+        if error.name not in ("transformers", "accelerate", "huggingface_hub"):
             raise
         raise ImportError(
             "fsdp2_wrap moved to modelopt.torch.utils.plugins.model_load_utils, which needs the hf "
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index 1ee1494d63f..459e51202f6 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -174,13 +174,15 @@ def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
     [
         ("transformers", _HF_EXTRA),
         ("huggingface_hub", _HF_EXTRA),
+        ("transformers.conversion_mapping", "halted; None in sys.modules"),
         ("modelopt.torch.utils.plugins.model_load_utils", "halted; None in sys.modules"),
     ],
-    ids=["transformers", "huggingface_hub", "another-import-error"],
+    ids=["transformers", "huggingface_hub", "transformers-submodule", "another-import-error"],
 )
 def test_fsdp2_wrap_names_the_extra_its_old_path_needs(monkeypatch, missing, message):
     """Without the hf extra, the old path says what it needs rather than which module is missing;
-    any other import error reaches the caller as it is."""
+    any other import error, such as a module an installed transformers lacks, reaches the caller as
+    it is."""
     monkeypatch.delitem(sys.modules, "modelopt.torch.utils.plugins.model_load_utils", raising=False)
     monkeypatch.setitem(sys.modules, missing, None)
 

From 368a533d5fd30a611f24a0550fb974163e031295 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 08:17:33 +0000
Subject: [PATCH 7/9] Skip the transformers-submodule row of the old-path test
 without transformers

Without transformers, the loaders' import of it fails first, so the old path rightly names
the missing extra; the partial-install unit lane runs this file so.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/utils/test_distributed.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index 459e51202f6..87f12d926b8 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -16,6 +16,7 @@
 """Tests for ``modelopt.torch.utils.distributed``: its pure functions, the distribution plan and a
 model's record of it, the collectives that fail together, and the process group's lifetime."""
 
+import importlib.util
 import os
 import pickle
 import sys
@@ -174,7 +175,14 @@ def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
     [
         ("transformers", _HF_EXTRA),
         ("huggingface_hub", _HF_EXTRA),
-        ("transformers.conversion_mapping", "halted; None in sys.modules"),
+        pytest.param(
+            "transformers.conversion_mapping",
+            "halted; None in sys.modules",
+            # Without transformers, its own import fails first: the extra is what is missing.
+            marks=pytest.mark.skipif(
+                importlib.util.find_spec("transformers") is None, reason="needs transformers"
+            ),
+        ),
         ("modelopt.torch.utils.plugins.model_load_utils", "halted; None in sys.modules"),
     ],
     ids=["transformers", "huggingface_hub", "transformers-submodule", "another-import-error"],

From 9a55cde4e39b0f45b7c68f89a335849c2fa48e06 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 17:46:56 +0000
Subject: [PATCH 8/9] Keep fsdp2_wrap free of the Hugging Face packages, in a
 module of its own

In the loaders' module, fsdp2_wrap needed transformers and huggingface_hub to import, and
so did its old path, which wrapped without them before: a model whose decoder layers are
registered with LayerActivationCollector needs neither. It now lives with its dtype
helpers in modelopt.torch.utils.fsdp2_utils, which imports no Hugging Face package; the
loaders import it from there, and the old path forwards to it. The changelog no longer
lists the move as breaking.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 CHANGELOG.rst                                 |   1 -
 modelopt/torch/utils/distributed.py           |  16 +--
 modelopt/torch/utils/fsdp2_utils.py           | 132 ++++++++++++++++++
 .../torch/utils/plugins/model_load_utils.py   | 104 +-------------
 tests/unit/torch/utils/test_distributed.py    |  80 ++++++-----
 tests/unit/torch/utils/test_fsdp2_utils.py    |  63 +++++++++
 .../unit/torch/utils/test_model_load_utils.py |  45 ------
 7 files changed, 247 insertions(+), 194 deletions(-)
 create mode 100644 modelopt/torch/utils/fsdp2_utils.py
 create mode 100644 tests/unit/torch/utils/test_fsdp2_utils.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index a0ecf9c2cf6..89fb2a8b8d7 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -51,7 +51,6 @@ Changelog
 - Bump minimum required PyTorch version to 2.10.
 - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra.
 - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later.
-- ``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.
 - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm.
 - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly.
diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py
index 9fd804e5f1d..e82257a3132 100644
--- a/modelopt/torch/utils/distributed.py
+++ b/modelopt/torch/utils/distributed.py
@@ -714,20 +714,10 @@ def is_fsdp2_model(model) -> bool:
 
 
 def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False):
-    """Forward to :func:`~modelopt.torch.utils.plugins.model_load_utils.fsdp2_wrap`.
+    """Forward to :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap`, where it moved."""
+    # Lazy import: fsdp2_utils imports modelopt.torch.utils.logging, which imports this module.
+    from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap as wrap
 
-    It moved there with the FSDP2 load that wraps with it, so this old path needs the ``hf`` extra.
-    """
-    try:
-        from modelopt.torch.utils.plugins.model_load_utils import fsdp2_wrap as wrap
-    except ModuleNotFoundError as error:
-        # A module missing from an installed package is a version mismatch, not a missing extra.
-        if error.name not in ("transformers", "accelerate", "huggingface_hub"):
-            raise
-        raise ImportError(
-            "fsdp2_wrap moved to modelopt.torch.utils.plugins.model_load_utils, which needs the hf "
-            "extra (pip install 'nvidia-modelopt[hf]') to import."
-        ) from error
     return wrap(model, shard_root=shard_root, mp_policy=mp_policy, cpu_offload=cpu_offload)
 
 
diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py
new file mode 100644
index 00000000000..14db7266b5a
--- /dev/null
+++ b/modelopt/torch/utils/fsdp2_utils.py
@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""FSDP2 wrapping of a model's decoder layers, which needs no Hugging Face package.
+
+The loaders (:mod:`modelopt.torch.utils.plugins.model_load_utils`) wrap with :func:`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``.
+"""
+
+from typing import Any
+
+import torch
+from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard
+from torch.distributed.tensor import DTensor
+
+from modelopt.torch.utils.logging import warn_rank_0
+
+__all__ = ["fsdp2_wrap"]
+
+
+def _off_dtype_params(model) -> set[torch.nn.Parameter]:
+    """Params whose dtype differs from the model's dominant (by element count) param dtype.
+
+    FSDP2 needs one dtype per shard group, but HF models routinely keep a few params in fp32 for
+    stability (e.g. MoE router gates). Pass these to ``fully_shard(ignored_params=...)``.
+
+    TODO: Drop this and shard the off-dtype params once a stable PyTorch release includes FSDP2
+    mixed-precision parameter dtype support (already on nightly).
+    """
+    numel_by_dtype: dict[torch.dtype, int] = {}
+    for param in model.parameters():
+        numel_by_dtype[param.dtype] = numel_by_dtype.get(param.dtype, 0) + param.numel()
+    if len(numel_by_dtype) <= 1:
+        return set()
+
+    dominant = max(numel_by_dtype, key=lambda d: numel_by_dtype[d])
+    off_dtype = {n: p for n, p in model.named_parameters() if p.dtype != dominant}
+    off_numel = sum(numel_by_dtype[d] for d in numel_by_dtype if d != dominant)
+    names = sorted(off_dtype)
+    warn_rank_0(
+        f"Model has mixed parameter dtypes {set(numel_by_dtype)}; FSDP2 needs one dtype per shard "
+        f"group, so {len(names)} non-{dominant} parameter(s) "
+        f"({100 * off_numel / sum(numel_by_dtype.values()):.2f}% of elements) will stay replicated "
+        f"rather than sharded: {names[:3]}{' ...' if len(names) > 3 else ''}"
+    )
+    return set(off_dtype.values())
+
+
+def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None:
+    """Move ``params`` onto the device FSDP2 computes on for ``model``.
+
+    ``fully_shard`` only moves the params it manages, so ignored ones would be stranded on
+    whatever device the caller built the model on. Meta params are left alone for deferred init.
+
+    The device comes from a sharded param's mesh rather than its local shard: under
+    ``cpu_offload`` the shard rests on CPU while compute still happens on the accelerator.
+    """
+    mesh = next((p.device_mesh for p in model.parameters() if isinstance(p, DTensor)), None)
+    if mesh is None:
+        warn_rank_0(
+            f"FSDP2 sharded no parameter of {type(model).__name__}, so the compute device for "
+            f"{len(params)} unsharded off-dtype parameter(s) cannot be determined; leaving them "
+            "where they are. Move them to the compute device or the forward will fail."
+        )
+        return
+
+    device = (
+        torch.device("cpu")
+        if mesh.device_type == "cpu"
+        else torch.device(mesh.device_type, getattr(torch, mesh.device_type).current_device())
+    )
+    for param in params:
+        if not param.is_meta and param.device != device:
+            param.data = param.data.to(device)
+
+
+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.
+
+    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.
+    """
+    # Importing the quantization package registers every mode and plugin; only a wrap needs it.
+    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."
+        )
+
+    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
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 2b86e8902dd..1967708c248 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -31,7 +31,7 @@
     from collections.abc import Iterable
 from torch.distributed.checkpoint import HuggingFaceStorageReader
 from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
-from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard
+from torch.distributed.fsdp import FSDPModule
 from torch.distributed.tensor import DTensor
 from transformers import (
     AutoConfig,
@@ -53,6 +53,7 @@
     broadcast_state_dict,
     is_initialized,
 )
+from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap
 from modelopt.torch.utils.logging import print_rank_0, warn_rank_0
 from modelopt.torch.utils.moe import (
     _choose_ep_degree,
@@ -503,107 +504,6 @@ def load_pretrained(
     return model
 
 
-def _off_dtype_params(model) -> set[torch.nn.Parameter]:
-    """Params whose dtype differs from the model's dominant (by element count) param dtype.
-
-    FSDP2 needs one dtype per shard group, but HF models routinely keep a few params in fp32 for
-    stability (e.g. MoE router gates). Pass these to ``fully_shard(ignored_params=...)``.
-
-    TODO: Drop this and shard the off-dtype params once a stable PyTorch release includes FSDP2
-    mixed-precision parameter dtype support (already on nightly).
-    """
-    numel_by_dtype: dict[torch.dtype, int] = {}
-    for param in model.parameters():
-        numel_by_dtype[param.dtype] = numel_by_dtype.get(param.dtype, 0) + param.numel()
-    if len(numel_by_dtype) <= 1:
-        return set()
-
-    dominant = max(numel_by_dtype, key=lambda d: numel_by_dtype[d])
-    off_dtype = {n: p for n, p in model.named_parameters() if p.dtype != dominant}
-    off_numel = sum(numel_by_dtype[d] for d in numel_by_dtype if d != dominant)
-    names = sorted(off_dtype)
-    warn_rank_0(
-        f"Model has mixed parameter dtypes {set(numel_by_dtype)}; FSDP2 needs one dtype per shard "
-        f"group, so {len(names)} non-{dominant} parameter(s) "
-        f"({100 * off_numel / sum(numel_by_dtype.values()):.2f}% of elements) will stay replicated "
-        f"rather than sharded: {names[:3]}{' ...' if len(names) > 3 else ''}"
-    )
-    return set(off_dtype.values())
-
-
-def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None:
-    """Move ``params`` onto the device FSDP2 computes on for ``model``.
-
-    ``fully_shard`` only moves the params it manages, so ignored ones would be stranded on
-    whatever device the caller built the model on. Meta params are left alone for deferred init.
-
-    The device comes from a sharded param's mesh rather than its local shard: under
-    ``cpu_offload`` the shard rests on CPU while compute still happens on the accelerator.
-    """
-    mesh = next((p.device_mesh for p in model.parameters() if isinstance(p, DTensor)), None)
-    if mesh is None:
-        warn_rank_0(
-            f"FSDP2 sharded no parameter of {type(model).__name__}, so the compute device for "
-            f"{len(params)} unsharded off-dtype parameter(s) cannot be determined; leaving them "
-            "where they are. Move them to the compute device or the forward will fail."
-        )
-        return
-
-    device = (
-        torch.device("cpu")
-        if mesh.device_type == "cpu"
-        else torch.device(mesh.device_type, getattr(torch, mesh.device_type).current_device())
-    )
-    for param in params:
-        if not param.is_meta and param.device != device:
-            param.data = param.data.to(device)
-
-
-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.
-
-    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.
-    """
-    # Importing the quantization package registers every mode and plugin; only a wrap needs it.
-    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."
-        )
-
-    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
-
-
 def _map_ckpt_keys_to_model(model: "torch.nn.Module", ckpt_keys, model_keys) -> dict:
     """Map each checkpoint key to the model ``state_dict`` key it feeds.
 
diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index 87f12d926b8..013b9f57bfd 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -16,12 +16,13 @@
 """Tests for ``modelopt.torch.utils.distributed``: its pure functions, the distribution plan and a
 model's record of it, the collectives that fail together, and the process group's lifetime."""
 
-import importlib.util
 import os
 import pickle
 import sys
 from contextlib import nullcontext
 from datetime import timedelta
+from functools import partial
+from unittest import mock
 
 import pytest
 import torch
@@ -29,7 +30,7 @@
 import torch.nn as nn
 from _test_utils.torch.distributed.utils import bounded_collectives
 from torch.distributed.device_mesh import init_device_mesh
-from torch.distributed.fsdp import fully_shard
+from torch.distributed.fsdp import FSDPModule, fully_shard
 
 import modelopt.torch.utils.distributed as dist
 from modelopt.torch.nas.hparams.container import DepthHparam
@@ -153,12 +154,10 @@ def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(cpu_dist_workers)
 
 
 def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
-    """fsdp2_wrap lives with the loaders that wrap with it; callers of the old path still reach it."""
-    model_load_utils = pytest.importorskip("modelopt.torch.utils.plugins.model_load_utils")
+    """fsdp2_wrap moved to fsdp2_utils; callers of the old path still reach it."""
     calls = []
     monkeypatch.setattr(
-        model_load_utils,
-        "fsdp2_wrap",
+        "modelopt.torch.utils.fsdp2_utils.fsdp2_wrap",
         lambda model, **kwargs: calls.append((model, kwargs)) or "layers",
     )
     model = nn.Linear(2, 2)
@@ -167,35 +166,50 @@ def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
     assert calls == [(model, {"shard_root": False, "mp_policy": None, "cpu_offload": True})]
 
 
-_HF_EXTRA = r"needs the hf extra \(pip install 'nvidia-modelopt\[hf\]'\)"
+class _Decoder(nn.Module):
+    """Decoder layers of plain PyTorch, which no Hugging Face plugin knows."""
 
+    def __init__(self):
+        super().__init__()
+        self.layers = nn.ModuleList([nn.Linear(4, 4), nn.Linear(4, 4)])
 
-@pytest.mark.parametrize(
-    ("missing", "message"),
-    [
-        ("transformers", _HF_EXTRA),
-        ("huggingface_hub", _HF_EXTRA),
-        pytest.param(
-            "transformers.conversion_mapping",
-            "halted; None in sys.modules",
-            # Without transformers, its own import fails first: the extra is what is missing.
-            marks=pytest.mark.skipif(
-                importlib.util.find_spec("transformers") is None, reason="needs transformers"
-            ),
-        ),
-        ("modelopt.torch.utils.plugins.model_load_utils", "halted; None in sys.modules"),
-    ],
-    ids=["transformers", "huggingface_hub", "transformers-submodule", "another-import-error"],
-)
-def test_fsdp2_wrap_names_the_extra_its_old_path_needs(monkeypatch, missing, message):
-    """Without the hf extra, the old path says what it needs rather than which module is missing;
-    any other import error, such as a module an installed transformers lacks, reaches the caller as
-    it is."""
-    monkeypatch.delitem(sys.modules, "modelopt.torch.utils.plugins.model_load_utils", raising=False)
-    monkeypatch.setitem(sys.modules, missing, None)
-
-    with pytest.raises(ImportError, match=message):
-        dist.fsdp2_wrap(nn.Linear(2, 2))
+
+# As an install without the hf extra has them: importing any of them fails.
+_HUGGING_FACE = ("transformers", "huggingface_hub", "accelerate")
+
+
+def _wraps_without_hugging_face_job(rank, size):
+    # Imported while the Hugging Face packages are there, as the pooled worker keeps them.
+    from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
+    from modelopt.torch.utils import fsdp2_utils
+
+    model, mesh = _Decoder(), init_device_mesh("cpu", (size,))
+    hidden = (*_HUGGING_FACE, "modelopt.torch.utils.plugins.model_load_utils")
+    saved = {name: sys.modules.pop(name) for name in hidden if name in sys.modules}
+    sys.modules.update(dict.fromkeys(_HUGGING_FACE))
+    try:
+        with (
+            mock.patch.object(LayerActivationCollector, "_decoder_layer_support", []),
+            # A CPU mesh, whatever accelerator the host has.
+            mock.patch.object(fsdp2_utils, "fully_shard", partial(fully_shard, mesh=mesh)),
+        ):
+            LayerActivationCollector.register_decoder_layer_support(
+                lambda model: isinstance(model, _Decoder), lambda model: model.layers
+            )
+            layers = dist.fsdp2_wrap(model, shard_root=False)
+    finally:
+        for name in hidden:
+            sys.modules.pop(name, None)
+        sys.modules.update(saved)
+    assert layers is model.layers
+    assert all(isinstance(layer, FSDPModule) for layer in layers)
+    assert not isinstance(model, FSDPModule)
+
+
+def test_fsdp2_wrap_needs_no_hugging_face(cpu_dist_workers):
+    """The old path wraps a model whose decoder layers are registered, with no Hugging Face package
+    to import and the loaders unimportable, as it did before fsdp2_wrap moved."""
+    cpu_dist_workers.run(_wraps_without_hugging_face_job)
 
 
 # ------------------------------------------------------------------------------------ collectives
diff --git a/tests/unit/torch/utils/test_fsdp2_utils.py b/tests/unit/torch/utils/test_fsdp2_utils.py
new file mode 100644
index 00000000000..b1bcadc933a
--- /dev/null
+++ b/tests/unit/torch/utils/test_fsdp2_utils.py
@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ``modelopt.torch.utils.fsdp2_utils``: which parameters stay out of FSDP2's shards."""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from modelopt.torch.utils.fsdp2_utils import _off_dtype_params
+
+
+def _model(*sizes_and_dtypes) -> nn.Module:
+    model = nn.Module()
+    for i, (numel, dtype) in enumerate(sizes_and_dtypes):
+        model.register_parameter(f"p{i}", nn.Parameter(torch.zeros(numel, dtype=dtype)))
+    return model
+
+
+def test_off_dtype_params_uniform_is_empty(recwarn):
+    model = _model((8, torch.bfloat16), (4, torch.bfloat16))
+    assert _off_dtype_params(model) == set()
+    assert not recwarn.list
+
+
+def test_off_dtype_params_returns_only_the_minority():
+    # An fp32 MoE router gate next to bf16 weights, as Nemotron-3-Nano ships it.
+    model = _model((100, torch.bfloat16), (5, torch.float32))
+    with pytest.warns(UserWarning, match="mixed parameter dtypes"):
+        off = _off_dtype_params(model)
+    assert off == {model.p1}
+
+
+@pytest.mark.parametrize(
+    ("params", "expected"),
+    [
+        # Dominant dtype is by element count, not parameter count: three small fp32 params
+        # lose to one large bf16 param, and vice versa.
+        ([(100, torch.bfloat16), (5, torch.float32), (5, torch.float32), (5, torch.float32)], "p0"),
+        (
+            [(100, torch.float32), (5, torch.bfloat16), (5, torch.bfloat16), (5, torch.bfloat16)],
+            "p0",
+        ),
+    ],
+)
+def test_off_dtype_params_dominant_is_by_numel(params, expected):
+    model = _model(*params)
+    with pytest.warns(UserWarning, match="mixed parameter dtypes"):
+        off = _off_dtype_params(model)
+    kept = {n for n, p in model.named_parameters() if p not in off}
+    assert kept == {expected}
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 432d15940b2..ca9e71b97f5 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -64,7 +64,6 @@
     _auto_class,
     _conversion_plan,
     _convert_keys,
-    _off_dtype_params,
     _resolve_target,
     build_meta_causal_lm,
     record_unplaced_source_keys,
@@ -189,50 +188,6 @@ def test_meta_qwen_vlm_preserves_architecture_and_vision(model_type):
     assert all(parameter.is_meta for parameter in model.parameters())
 
 
-# ----------------------------------------------------------------------- fsdp2_wrap's dtype helpers
-
-
-def _model(*sizes_and_dtypes) -> torch.nn.Module:
-    model = torch.nn.Module()
-    for i, (numel, dtype) in enumerate(sizes_and_dtypes):
-        model.register_parameter(f"p{i}", torch.nn.Parameter(torch.zeros(numel, dtype=dtype)))
-    return model
-
-
-def test_off_dtype_params_uniform_is_empty(recwarn):
-    model = _model((8, torch.bfloat16), (4, torch.bfloat16))
-    assert _off_dtype_params(model) == set()
-    assert not recwarn.list
-
-
-def test_off_dtype_params_returns_only_the_minority():
-    # An fp32 MoE router gate next to bf16 weights, as Nemotron-3-Nano ships it.
-    model = _model((100, torch.bfloat16), (5, torch.float32))
-    with pytest.warns(UserWarning, match="mixed parameter dtypes"):
-        off = _off_dtype_params(model)
-    assert off == {model.p1}
-
-
-@pytest.mark.parametrize(
-    ("params", "expected"),
-    [
-        # Dominant dtype is by element count, not parameter count: three small fp32 params
-        # lose to one large bf16 param, and vice versa.
-        ([(100, torch.bfloat16), (5, torch.float32), (5, torch.float32), (5, torch.float32)], "p0"),
-        (
-            [(100, torch.float32), (5, torch.bfloat16), (5, torch.bfloat16), (5, torch.bfloat16)],
-            "p0",
-        ),
-    ],
-)
-def test_off_dtype_params_dominant_is_by_numel(params, expected):
-    model = _model(*params)
-    with pytest.warns(UserWarning, match="mixed parameter dtypes"):
-        off = _off_dtype_params(model)
-    kept = {n for n, p in model.named_parameters() if p not in off}
-    assert kept == {expected}
-
-
 # -------------------------------------------------------------- deciding how a model is distributed
 
 

From 30ab95bdb5398693262940e8a2e76e27d0b9bec1 Mon Sep 17 00:00:00 2001
From: Shengliang Xu <shengliangx@nvidia.com>
Date: Fri, 30 Oct 2026 21:03:40 +0000
Subject: [PATCH 9/9] Drop the old path's mock-only forwarding test

The two-rank test wraps a registered decoder for real through the old path, and the GPU
tests wrap through it with cpu_offload, so a check that the forwarder hands its arguments
to a stand-in adds nothing.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
---
 tests/unit/torch/utils/test_distributed.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py
index 013b9f57bfd..f292843c89a 100644
--- a/tests/unit/torch/utils/test_distributed.py
+++ b/tests/unit/torch/utils/test_distributed.py
@@ -153,19 +153,6 @@ def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(cpu_dist_workers)
     cpu_dist_workers.run(_no_plan_on_several_ranks_job)
 
 
-def test_fsdp2_wrap_keeps_its_old_path(monkeypatch):
-    """fsdp2_wrap moved to fsdp2_utils; callers of the old path still reach it."""
-    calls = []
-    monkeypatch.setattr(
-        "modelopt.torch.utils.fsdp2_utils.fsdp2_wrap",
-        lambda model, **kwargs: calls.append((model, kwargs)) or "layers",
-    )
-    model = nn.Linear(2, 2)
-
-    assert dist.fsdp2_wrap(model, shard_root=False, cpu_offload=True) == "layers"
-    assert calls == [(model, {"shard_root": False, "mp_policy": None, "cpu_offload": True})]
-
-
 class _Decoder(nn.Module):
     """Decoder layers of plain PyTorch, which no Hugging Face plugin knows."""