diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 26892ef112d..e82257a3132 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__ = [ @@ -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( 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 b9cb961317b..1967708c248 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/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, diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py index 25b3bc12268..f292843c89a 100644 --- a/tests/unit/torch/utils/test_distributed.py +++ b/tests/unit/torch/utils/test_distributed.py @@ -18,8 +18,11 @@ 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 @@ -27,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 @@ -35,7 +38,6 @@ 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 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}