[OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils

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__ = [
@@ -713,112 +713,12 @@ 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):
-    """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.
-    """
-    # 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."
-        )
+    """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
 
-    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(

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

modelopt/torch/utils/plugins/model_load_utils.py

@@ -51,9 +51,9 @@
     DistributionPlan,
     barrier,
     broadcast_state_dict,
-    fsdp2_wrap,
     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,

tests/unit/torch/utils/test_distributed.py

@@ -18,24 +18,26 @@
 
 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
 import torch.distributed
 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
 from modelopt.torch.utils.distributed import (
     DistributedState,
     DistributionPlan,
     _deserialize,
-    _off_dtype_params,
     _serialize,
     all_reduce,
     broadcast_from,
@@ -45,47 +47,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",))
@@ -192,6 +153,52 @@ 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)
 
 
+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)])
+
+
+# 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