From b26c26d02d499dc38c0ec24af938f0a7d7e6aee9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 18:45:06 +0000 Subject: [PATCH 01/41] Route load_pretrained by the distribution plan load_pretrained becomes the one loader for a model ModelOpt will optimize and export, whole or distributed. A new distribution argument picks the load: "sequential" (from_pretrained as before), "fsdp2" (parallel_load_and_prepare_fsdp2), or a DistributionPlan -- FSDP2 over the world, or a whole copy on every rank's own GPU. A plan that shards the routed experts raises: no load here does that yet. Left at None, the active ModelOptSession decides: a single process loads whole, a distributed session decides the plan with the new resolve_distribution_plan, through decide_distribution_plan with the MODELOPT_EP_DEGREE and MODELOPT_FSDP_SIZE knobs (modelopt.torch.utils.env). The plan a load runs is kept as the session's synced plan, the same on every rank. Every load makes a Hub ID local and records the unplaced source keys. The FSDP2 load cannot run the model type's loading hook, so it refuses a checkpoint the hook would build modules for (a Nemotron-H MTP head), and it refuses from_pretrained arguments it would otherwise ignore, a quantized checkpoint, whose tensors it would read as stored, and a partial fsdp_size, which it cannot shard. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + .../torch/utils/plugins/model_load_utils.py | 226 +++++++++- .../test_load_pretrained_distribution.py | 99 +++++ .../test_load_pretrained_distribution.py | 419 ++++++++++++++++++ 4 files changed, 730 insertions(+), 15 deletions(-) create mode 100644 tests/gpu/torch/utils/test_load_pretrained_distribution.py create mode 100644 tests/unit/torch/utils/test_load_pretrained_distribution.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e943e511c9a..3ca53274a89 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -148,6 +148,7 @@ Changelog - Add a ``kd_loss_alpha`` option to HuggingFace ``KDTrainer`` to blend cross-entropy and knowledge distillation losses during training. - A tracked ``examples/hf_ptq/hf_ptq.py`` run now writes ``.experiment.json`` into ``--export_path`` and uploads the same file with the run, so a checkpoint on disk names the experiment and MLflow run id that produced it. The pointer is written only once the export completes, and an export that is not tracked removes one it would otherwise inherit from a reused ``--export_path`` or from a quantized source checkpoint. - ``export_hf_checkpoint`` and the vLLM fake-quant export now copy the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged. They read the source from local disk only: get a local copy of a Hugging Face Hub model with ``modelopt.torch.export.ensure_local_checkpoint`` and load it from there. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Inside a distributed ``modelopt.torch.utils.distributed.modelopt_session()`` it FSDP2-shards the model or gives every rank a whole copy, as the model and the ``MODELOPT_FSDP_SIZE`` variable decide; its new ``distribution`` argument picks the load explicitly. - Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with Qwen-Image training, distributed-checkpoint export, and PDD-2/4/8 inference. AutoModel remains an unmodified pinned runtime dependency. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c6078aded80..d6f915f5dcf 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""HuggingFace model loading: :func:`load_pretrained`, and the plan to distribute a model by.""" +"""HuggingFace model loading: :func:`load_pretrained`, by a plan to distribute the model.""" import logging import os @@ -400,17 +400,190 @@ def record_unplaced_source_keys( return keys +_FSDP2 = "FSDP2" +# What each sharded load takes besides the checkpoint. These loads build the model class and place +# it themselves, so from_pretrained's own arguments (device_map, max_memory, ...) do not reach them. +_SHARDED_LOAD_OPTIONS = { + _FSDP2: frozenset( + { + "trust_remote_code", + "attn_implementation", + "hf_config", + "cpu_offload", + "mp_policy", + "broadcast_chunk_size", + } + ), +} + + +def _distributed_session(): + """The active ``ModelOptSession`` when it runs distributed, else ``None``.""" + from modelopt.torch.utils.distributed import get_active_mo_session + + session = get_active_mo_session(required=False) + return session if session.initialized and session.distributed else None + + +def _decide_distribution(ckpt_path: str, trust_remote_code: bool, ep_size, fsdp_size): + """:func:`resolve_distribution_plan` for a local ``ckpt_path``.""" + from modelopt.torch.utils import env + from modelopt.torch.utils.moe import decide_distribution_plan + + session = _distributed_session() + if session is None: + return None, None + hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code) + plan = decide_distribution_plan( + hf_config, + session.world_size, + ep_size=env.ep_degree() if ep_size is None else ep_size, + fsdp_size=env.fsdp_size() if fsdp_size is None else fsdp_size, + ) + if session.is_main: + print(f"Distribution: {plan!r}", flush=True) + return plan, hf_config + + +def resolve_distribution_plan( + ckpt_path: str, + *, + trust_remote_code: bool = False, + ep_size: "int | str | None" = None, + fsdp_size: "int | str | None" = None, +) -> tuple[DistributionPlan | None, PretrainedConfig | None]: + """Decide how :func:`load_pretrained` distributes ``ckpt_path`` in this session. + + Returns ``(plan, hf_config)``: a :class:`~modelopt.torch.utils.distributed.DistributionPlan` + and the config it was decided from, which a sharded load takes back as + ``load_pretrained(ckpt_path, distribution=plan, hf_config=hf_config)`` rather than read it + again. Outside a distributed :class:`~modelopt.torch.utils.distributed.ModelOptSession`: + ``(None, None)``, the model loads whole. ``load_pretrained`` decides the same way on its own; + deciding first is for a caller that loads an unsharded model differently. + + :func:`~modelopt.torch.utils.moe.decide_distribution_plan` decides, once on rank 0, from + ``ep_size`` and ``fsdp_size``; ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` + (:mod:`modelopt.torch.utils.env`), which default to ``"auto"``. Collective: every rank must + call it. + """ + if _distributed_session() is None: + return None, None + _, ckpt_path = ensure_local_checkpoint(ckpt_path) + return _decide_distribution(ckpt_path, trust_remote_code, ep_size, fsdp_size) + + +def _route(distribution: DistributionPlan | str) -> str: + """Which load ``distribution`` asks for: a sharded one, ``"replicated"``, or ``"sequential"``.""" + if isinstance(distribution, str): + if distribution in ("sequential", "fsdp2"): + return _FSDP2 if distribution == "fsdp2" else "sequential" + raise ValueError( + f"distribution={distribution!r}: expected None, 'sequential', 'fsdp2' or a " + "DistributionPlan." + ) + plan, session = distribution, _distributed_session() + world_size = session.world_size if session is not None else 1 + if plan.world_size != world_size: + raise ValueError(f"{plan!r} is for {plan.world_size} ranks, but this run has {world_size}.") + if plan.is_expert_parallel: + raise NotImplementedError( + f"{plan!r} shards the routed experts, which load_pretrained has no load for. Plan " + "ep_size=1 (MODELOPT_EP_DEGREE=1) to give every rank the whole model, or shard it " + "with FSDP2." + ) + if plan.is_fsdp_dense: + if plan.fsdp_size != plan.world_size: + raise NotImplementedError( + f"{plan!r}: without expert parallelism the FSDP2 load shards over the whole world, " + f"so fsdp_size must be {plan.world_size}." + ) + return _FSDP2 + return "replicated" if session is not None else "sequential" + + +def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) -> nn.Module: + """The sharded load ``route`` names, refusing what it cannot honor before it starts.""" + from modelopt.torch.models.hf import checkpoint_has_mtp + from modelopt.torch.utils import distributed as dist + + if auto_class not in (None, AutoModelForCausalLM): + raise NotImplementedError( + f"The {route} load builds its model class itself; auto_class={auto_class.__name__} " + "does not apply to it." + ) + kwargs = dict(kwargs) + if kwargs.get("dtype", "auto") == "auto": + kwargs.pop("dtype", None) # A sharded load keeps the checkpoint's dtype. + unknown = sorted(set(kwargs) - _SHARDED_LOAD_OPTIONS[route]) + if unknown: + raise TypeError( + f"{unknown} do not apply to the {route} load, which places the model itself; it " + f"takes {sorted(_SHARDED_LOAD_OPTIONS[route])}." + ) + if checkpoint_has_mtp(model_type, ckpt_path): + # prepare_model_for_loading builds the head on a whole model only; carried over, the + # head would be exported unquantized. + raise NotImplementedError( + f"The {route} load cannot build the MTP head this {model_type} checkpoint carries; " + "load it unsharded." + ) + trust_remote_code = kwargs.pop("trust_remote_code", False) + hf_config = kwargs.pop("hf_config", None) + if not dist.is_initialized(): + raise RuntimeError("The FSDP2 load needs a torch.distributed process group (torchrun).") + if hf_config is None: + hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code) + if getattr(hf_config, "quantization_config", None) is not None: + # It reads the tensors as stored: quantized weights would load as if they were not. + raise NotImplementedError( + f"The FSDP2 load cannot dequantize this {model_type} checkpoint; load it unsharded." + ) + return parallel_load_and_prepare_fsdp2( + ckpt_path, + torch.device("cuda", dist.local_rank()), + dist.rank(), + dist.size(), + trust_remote_code=trust_remote_code, + hf_config=hf_config, + **kwargs, + ) + + def load_pretrained( ckpt_path: str, auto_class: Any = None, *, model_type: str | None = None, + distribution: DistributionPlan | str | None = None, **kwargs, ) -> nn.Module: - """Load a Hugging Face model for ModelOpt to optimize and export. + """Load a Hugging Face model for ModelOpt to optimize and export, whole or distributed. + + ``distribution`` decides how the model is placed: + + * ``None`` (default): the active :class:`~modelopt.torch.utils.distributed.ModelOptSession` + decides. Without one, or in a single-process one, the model loads whole; in a distributed one, + as :func:`resolve_distribution_plan` decides. + * ``"sequential"``: ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` + defaults to ``AutoModelForCausalLM``), placed as ``kwargs`` say. Each rank that calls it gets + the whole model. + * ``"fsdp2"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its + own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, ``mp_policy`` and + ``broadcast_chunk_size`` ``kwargs`` may set). + * A :class:`~modelopt.torch.utils.distributed.DistributionPlan`, as + :func:`resolve_distribution_plan` returns it: + + - ``fsdp_size > 1``: as ``"fsdp2"``; + - otherwise as ``"sequential"`` on every rank, ``device_map`` defaulting to the rank's GPU. + + A plan with ``ep_size > 1`` raises: no load here shards the routed experts. The plan is kept + as ``session.plan``, the same on every rank. + + A sharded load builds the model class and places it itself. Besides its own options it takes + ``trust_remote_code``, ``attn_implementation`` and ``hf_config`` (the config to load with), and + refuses any other keyword rather than ignore it. - ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to - ``AutoModelForCausalLM``), plus what ModelOpt needs around the load: + Whatever the placement: * A Hub model ID is made local first (:func:`ensure_local_checkpoint`, which takes the ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` ``kwargs``): the export reads @@ -422,7 +595,8 @@ def load_pretrained( * The model type's loading hook runs (:func:`~modelopt.torch.models.hf.prepare_model_for_loading`), e.g. to keep the MTP head a Nemotron-H checkpoint carries. ``model_type`` selects it; when omitted, the type of the - ``config`` passed for the load, else of the checkpoint's config. + config passed for the load, else of the checkpoint's config. A sharded load cannot run it, so + it refuses a checkpoint the hook would build modules for. * The checkpoint tensors the model has no parameter for are recorded (:func:`record_unplaced_source_keys`), so the export carries them over unchanged. The model is returned alone: ``output_loading_info=True`` is refused. So is a ``variant``: the record @@ -443,7 +617,6 @@ def load_pretrained( # Importing modelopt.torch.models registers every model spec; only callers of this need them. from modelopt.torch.models.hf import prepare_model_for_loading - auto_class = auto_class or AutoModelForCausalLM hub_options = { k: kwargs.pop(k) for k in ("token", "cache_dir", "local_files_only") if k in kwargs } @@ -453,21 +626,44 @@ def load_pretrained( if subfolder: ckpt_path = os.path.join(ckpt_path, subfolder) trust_remote_code = kwargs.get("trust_remote_code", False) + if distribution is None: + plan, hf_config = _decide_distribution(ckpt_path, trust_remote_code, None, None) + distribution = "sequential" if plan is None else plan + if plan is not None and plan.is_fsdp_dense: + kwargs.setdefault("hf_config", hf_config) + route = _route(distribution) + session = _distributed_session() + if session is not None and isinstance(distribution, DistributionPlan): + session.plan = distribution # What the model is loaded with, the same on every rank. if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. - model_type = getattr(kwargs.get("config"), "model_type", None) or ( + config = kwargs.get("hf_config") or kwargs.get("config") + model_type = getattr(config, "model_type", None) or ( AutoConfig.from_pretrained( ckpt_path, trust_remote_code=trust_remote_code, **hub_options ).model_type ) - with prepare_model_for_loading( - model_type, - ckpt_path, - trust_remote_code, - model_class=None if auto_class in (AutoModel, AutoModelForCausalLM) else auto_class, - ): - model = auto_class.from_pretrained(ckpt_path, **hub_options, **kwargs) - unplaced = record_unplaced_source_keys(model, ckpt_path) + if route in _SHARDED_LOAD_OPTIONS: + model = _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) + # A sharded load finds the unplaced keys while it places the weights. + keys = model._modelopt_unplaced_source_keys + else: + if "hf_config" in kwargs: + raise TypeError("hf_config applies to a sharded load only.") + if route == "replicated": + # Each rank takes the whole model onto its own GPU. Left to "auto", every rank's + # accelerate could pick the same idle GPU. + kwargs.setdefault("device_map", {"": str(session.device)}) + auto_class = auto_class or AutoModelForCausalLM + with prepare_model_for_loading( + model_type, + ckpt_path, + trust_remote_code, + model_class=None if auto_class in (AutoModel, AutoModelForCausalLM) else auto_class, + ): + model = auto_class.from_pretrained(ckpt_path, **hub_options, **kwargs) + keys = None + unplaced = record_unplaced_source_keys(model, ckpt_path, keys) if unplaced: print_rank_0( f"{len(unplaced)} checkpoint tensor(s) the model has no parameter for " diff --git a/tests/gpu/torch/utils/test_load_pretrained_distribution.py b/tests/gpu/torch/utils/test_load_pretrained_distribution.py new file mode 100644 index 00000000000..3abb9d9dd44 --- /dev/null +++ b/tests/gpu/torch/utils/test_load_pretrained_distribution.py @@ -0,0 +1,99 @@ +# 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. + +"""``load_pretrained`` in a distributed session, on real GPUs. + +Left to decide, it FSDP2-shards a dense model when asked to, and otherwise gives every rank the whole +model on its own GPU. Whichever load runs, every rank ends up with the checkpoint's weights and the +load's record of what it could not place. +""" + +import os + +import pytest +import torch + +transformers = pytest.importorskip("transformers") + +from torch.distributed.tensor import DTensor +from transformers import AutoModelForCausalLM, LlamaConfig + +from modelopt.torch.utils.distributed import get_active_mo_session +from modelopt.torch.utils.plugins.model_load_utils import load_pretrained + +_TINY = { + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 128, + "max_position_embeddings": 64, +} + + +def _save(config, ckpt_dir): + torch.manual_seed(0) + AutoModelForCausalLM.from_config(config).to(torch.bfloat16).save_pretrained(ckpt_dir) + + +def _full(tensor): + return tensor.full_tensor() if isinstance(tensor, DTensor) else tensor + + +def _load_and_check(rank, size, ckpt_dir, knobs, expected): + """Enter a distributed session as torchrun would, load with the plan left to decide, check it.""" + device = torch.device(f"cuda:{rank}") + session = get_active_mo_session(required=False) + session.initialize( + distributed=True, + rank=rank, + world_size=size, + local_rank=rank, + device=device, + is_main=rank == 0, + ) + os.environ.update(knobs) + try: + model = load_pretrained(ckpt_dir) + whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16) + + params = dict(model.named_parameters()) + whole_params = dict(whole.named_parameters()) + assert params.keys() == whole_params.keys() + sharded = {name for name, param in params.items() if isinstance(param, DTensor)} + if expected == "replicated": + assert not sharded + assert all(param.device == device for param in params.values()) + else: + assert {name for name in params if ".layers." in name} <= sharded + for name, param in params.items(): # the same order on every rank: full_tensor gathers + torch.testing.assert_close(_full(param).cpu(), whole_params[name]) + assert model._modelopt_unplaced_source_keys == [] + assert model._modelopt_source_checkpoint == ckpt_dir + finally: + for knob in knobs: + del os.environ[knob] + session._reset() + + +def test_a_dense_model_loads_whole_on_every_rank(need_2_gpus, dist_workers, tmp_path): + _save(LlamaConfig(**_TINY), tmp_path) + dist_workers.run(_load_and_check, str(tmp_path), {}, "replicated") + + +def test_a_dense_model_asked_to_shard_loads_fsdp2_sharded(need_2_gpus, dist_workers, tmp_path): + _save(LlamaConfig(**_TINY), tmp_path) + dist_workers.run(_load_and_check, str(tmp_path), {"MODELOPT_FSDP_SIZE": "world"}, "fsdp2") diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py new file mode 100644 index 00000000000..f3981762752 --- /dev/null +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -0,0 +1,419 @@ +# 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. + +"""How ``load_pretrained`` decides a model's distribution, and which load that routes it to. + +Nothing here needs a process group or a GPU: the active ``ModelOptSession`` is initialized the way a +torchrun launch leaves it, the configs are tiny real transformers configs served from memory, and the +loaders are replaced by stand-ins that record how they were called. +""" + +import inspect + +import pytest +import torch + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version +from transformers import AutoModel, LlamaConfig, Qwen3MoeConfig + +import modelopt.torch.models.hf as models_hf +import modelopt.torch.utils.distributed as dist +import modelopt.torch.utils.moe as moe +from modelopt.torch.utils import env +from modelopt.torch.utils.distributed import DistributionPlan, get_active_mo_session +from modelopt.torch.utils.plugins import model_load_utils +from modelopt.torch.utils.plugins.model_load_utils import load_pretrained, resolve_distribution_plan + +# Small enough that the dense part fits any GPU, so "auto" dense sharding always resolves to +# replication whether or not the machine running the tests has one. +_TINY = { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "vocab_size": 32, +} + + +def _qwen3_moe(num_experts): + """Qwen3-MoE, whose base_model_ep_plan ships with transformers itself.""" + return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY) + + +@pytest.fixture(autouse=True) +def _no_distribution_knobs(monkeypatch): + """A shell that exports the knobs would otherwise steer every decision below.""" + monkeypatch.delenv("MODELOPT_EP_DEGREE", raising=False) + monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False) + + +@pytest.fixture +def session(monkeypatch): + """Enter the active session as rank ``rank`` of a ``world_size``-rank launch would. + + Every rank acts as its group's first, deciding for itself: these tests are about what is + decided, and test_distributed_session.py about how the ranks agree on it. + """ + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) + active = get_active_mo_session(required=False) + + def enter(world_size, rank=0): + active.initialize( + distributed=world_size > 1, + rank=rank, + world_size=world_size, + local_rank=rank, + device=torch.device(f"cuda:{rank}"), + is_main=rank == 0, + ) + return active + + yield enter + active._reset() + + +@pytest.fixture +def serve(monkeypatch): + """Answer every config read with ``config`` rather than a checkpoint on disk.""" + + def serve(config): + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **k: config) + return config + + return serve + + +class _Loaded(torch.nn.Module): + """What a sharded loader returns: a model whose load recorded what it could not place.""" + + def __init__(self): + super().__init__() + self._modelopt_unplaced_source_keys = [] + + +@pytest.fixture +def loaders(monkeypatch): + """Stand-ins for the sharded loaders, each recording its call by parameter name.""" + calls = {} + + def recording(name, loader): + def stub(*args, **kwargs): + calls[name] = inspect.signature(loader).bind(*args, **kwargs).arguments + return _Loaded() + + return stub + + monkeypatch.setattr( + model_load_utils, + "parallel_load_and_prepare_fsdp2", + recording("fsdp2", model_load_utils.parallel_load_and_prepare_fsdp2), + ) + return calls + + +@pytest.fixture +def auto_class(): + """An ``auto_class`` whose from_pretrained records its arguments.""" + + class Auto: + calls: list = [] + + @classmethod + def from_pretrained(cls, ckpt_path, **kwargs): + cls.calls.append(kwargs) + return torch.nn.Linear(1, 1) + + return Auto + + +@pytest.fixture +def process_group(monkeypatch): + """Rank 1 of a 2-rank process group, as the FSDP2 load reads it.""" + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "rank", lambda group=None: 1) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(dist, "local_rank", lambda: 1) + + +# ----------------------------------------------------------------------------------- env knobs + + +@pytest.mark.parametrize( + ("knob", "name"), + [(env.ep_degree, "MODELOPT_EP_DEGREE"), (env.fsdp_size, "MODELOPT_FSDP_SIZE")], +) +def test_distribution_knobs_default_to_auto(monkeypatch, knob, name): + """Unset and exported-but-empty both mean "let the plan decide".""" + assert knob() == "auto" + monkeypatch.setenv(name, "") + assert knob() == "auto" + + +@pytest.mark.parametrize( + ("knob", "name", "other"), + [ + (env.ep_degree, "MODELOPT_EP_DEGREE", env.fsdp_size), + (env.fsdp_size, "MODELOPT_FSDP_SIZE", env.ep_degree), + ], +) +def test_distribution_knobs_are_read_when_called(monkeypatch, knob, name, other): + """A value snapshotted at import would miss what the launcher sets later; and each knob reads + only its own variable.""" + monkeypatch.setenv(name, "2") + assert knob() == "2" + monkeypatch.setenv(name, "8") + assert knob() == "8" + assert other() == "auto" + + +# -------------------------------------------------------------------------------- the decision + + +@pytest.mark.parametrize("world_size", [None, 1], ids=["no_session", "single_process_session"]) +def test_outside_a_distributed_session_nothing_is_distributed(session, serve, tmp_path, world_size): + """With no ranks to shard over, even a MoE loads whole.""" + if world_size is not None: + session(world_size) + serve(_qwen3_moe(num_experts=8)) + + assert resolve_distribution_plan(str(tmp_path)) == (None, None) + + +@pytest.mark.skipif( + Version(transformers.__version__) < Version("5.7"), + reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", +) +@pytest.mark.parametrize("rank", [0, 1]) +def test_an_ep_capable_moe_is_planned_expert_parallel(session, serve, tmp_path, capsys, rank): + """Auto picks the largest EP degree dividing both the world and the expert count, hands back + the config it decided from, and only rank 0 reports the decision.""" + session(4, rank) + serve(_qwen3_moe(num_experts=8)) + + plan, hf_config = resolve_distribution_plan(str(tmp_path)) + + assert (plan.world_size, plan.ep_size, plan.edp) == (4, 4, 1) + assert not plan.is_fsdp_dense + assert moe.is_expert_parallel_capable(hf_config) + assert ("Distribution:" in capsys.readouterr().out) == (rank == 0) + + +def test_a_dense_model_is_replicated(session, serve, tmp_path): + """A dense model that fits one GPU has nothing to shard: every rank holds it whole.""" + session(4) + serve(LlamaConfig(**_TINY)) + + plan, _ = resolve_distribution_plan(str(tmp_path)) + + assert (plan.ep_size, plan.fsdp_size) == (1, 1) + + +def test_the_env_knobs_override_the_automatic_plan(session, serve, monkeypatch, tmp_path): + """The decision honours both knobs over its automatic choice.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + session(4) + serve(_qwen3_moe(num_experts=8)) + + plan, _ = resolve_distribution_plan(str(tmp_path)) + + assert (plan.ep_size, plan.edp, plan.fsdp_size) == (2, 2, 4) + + +def test_arguments_override_the_env_knobs(session, serve, monkeypatch, tmp_path): + monkeypatch.setenv("MODELOPT_EP_DEGREE", "4") + session(4) + serve(_qwen3_moe(num_experts=8)) + + plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=2, fsdp_size=1) + + assert (plan.ep_size, plan.fsdp_size) == (2, 1) + + +# ----------------------------------------------------------------------------------- the loads + + +def test_without_a_session_the_model_loads_whole_as_asked(serve, loaders, auto_class, tmp_path): + """No session, no plan: from_pretrained places the model exactly as the caller says.""" + serve(LlamaConfig(**_TINY)) + + model = load_pretrained(str(tmp_path), auto_class, device_map="auto", dtype="auto") + + assert auto_class.calls == [{"device_map": "auto", "dtype": "auto"}] + assert model._modelopt_unplaced_source_keys == [] + assert loaders == {} + + +@pytest.mark.parametrize( + ("device_map", "placed"), + [(None, {"": "cuda:1"}), ("cpu", "cpu")], + ids=["defaults_to_the_rank_gpu", "caller_placement_wins"], +) +def test_a_replicated_plan_loads_the_whole_model_on_every_rank( + session, serve, loaders, auto_class, tmp_path, device_map, placed +): + """Each rank takes the whole model onto its own GPU; "auto" would let every rank's accelerate + pick the same idle GPU. The plan the load decided is kept on the session.""" + active = session(4, rank=1) + serve(LlamaConfig(**_TINY)) + kwargs = {} if device_map is None else {"device_map": device_map} + + load_pretrained(str(tmp_path), auto_class, **kwargs) + + assert auto_class.calls == [{"device_map": placed}] + assert loaders == {} + assert (active.plan.ep_size, active.plan.fsdp_size) == (1, 1) + + +@pytest.mark.parametrize("distribution", ["fsdp2", DistributionPlan(2, fsdp_size=2)]) +def test_fsdp2_loads_through_the_parallel_fsdp2_loader( + session, serve, loaders, process_group, tmp_path, distribution +): + """Named or planned, FSDP2 shards over the process group it reads its rank and GPU from.""" + session(2, rank=1) + served = serve(LlamaConfig(**_TINY)) + + load_pretrained(str(tmp_path), distribution=distribution, cpu_offload=True, dtype="auto") + + call = loaders.pop("fsdp2") + assert (call["device"], call["rank"], call["world_size"]) == (torch.device("cuda:1"), 1, 2) + assert call["cpu_offload"] is True + assert call["hf_config"] is served + assert loaders == {} + + +def test_a_load_it_decides_reads_the_config_once( + session, loaders, process_group, monkeypatch, tmp_path +): + """The config the plan was decided from is the one the FSDP2 load gets, not a second read.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + session(2, rank=1) + config, reads = LlamaConfig(**_TINY), [] + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a) or config + ) + + load_pretrained(str(tmp_path)) + + assert len(reads) == 1 + assert loaders["fsdp2"]["hf_config"] is config + + +def test_partial_dense_sharding_without_expert_parallelism_is_refused(session, serve, tmp_path): + """The FSDP2 load shards over the whole world; it cannot honour a smaller dense degree.""" + session(4) + serve(LlamaConfig(**_TINY)) + + with pytest.raises(NotImplementedError, match="fsdp_size must be 4"): + load_pretrained(str(tmp_path), distribution=DistributionPlan(4, fsdp_size=2)) + + +def test_an_expert_parallel_plan_is_refused(session, serve, loaders, auto_class, tmp_path): + """No load here shards the routed experts: an EP plan fails before anything loads, and is not + kept on the session.""" + active = session(2) + serve(_qwen3_moe(num_experts=8)) + + with pytest.raises(NotImplementedError, match="shards the routed experts"): + load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) + assert active.plan is None + assert loaders == {} + assert auto_class.calls == [] + + +def test_the_plan_a_load_runs_is_kept_on_the_session( + session, serve, loaders, process_group, tmp_path +): + """Handed a plan, the load records it, as it does one it decides itself.""" + active = session(2) + serve(LlamaConfig(**_TINY)) + + load_pretrained(str(tmp_path), distribution=DistributionPlan(2, fsdp_size=2)) + + assert active.plan.fsdp_size == 2 + + +def test_a_plan_for_another_world_is_refused(session, serve, tmp_path): + session(2) + serve(LlamaConfig(**_TINY)) + + with pytest.raises(ValueError, match="is for 4 ranks, but this run has 2"): + load_pretrained(str(tmp_path), distribution=DistributionPlan(4, fsdp_size=4)) + + +@pytest.mark.parametrize( + ("change", "error", "match"), + [ + ({"kwargs": {"device_map": "auto"}}, TypeError, r"\['device_map'\] do not apply"), + ({"auto_class": AutoModel}, NotImplementedError, "builds its model class itself"), + ({"mtp": True}, NotImplementedError, "cannot build the MTP head"), + ], + ids=["placement", "model_class", "mtp_head"], +) +def test_a_sharded_load_refuses_what_it_cannot_honor( + session, serve, loaders, process_group, monkeypatch, tmp_path, change, error, match +): + """Placement, a model class, a head only a whole model builds: none reaches a sharded load, + which fails before it starts rather than load something other than what was asked.""" + session(2) + serve(LlamaConfig(**_TINY)) + if change.get("mtp"): + monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) + + with pytest.raises(error, match=match): + load_pretrained( + str(tmp_path), + change.get("auto_class"), + distribution="fsdp2", + **change.get("kwargs", {}), + ) + assert loaders == {} + + +def test_the_fsdp2_load_refuses_a_quantized_checkpoint(serve, loaders, process_group, tmp_path): + """It reads the tensors as stored, so quantized weights would load as if they were not.""" + config = LlamaConfig(**_TINY) + config.quantization_config = {"quant_method": "fp8"} + serve(config) + + with pytest.raises(NotImplementedError, match="cannot dequantize"): + load_pretrained(str(tmp_path), distribution="fsdp2") + assert loaders == {} + + +def test_the_fsdp2_load_needs_a_process_group(serve, loaders, tmp_path): + serve(LlamaConfig(**_TINY)) + + with pytest.raises(RuntimeError, match=r"needs a torch\.distributed process group"): + load_pretrained(str(tmp_path), distribution="fsdp2") + + +def test_an_unknown_distribution_is_refused(serve, tmp_path): + serve(LlamaConfig(**_TINY)) + + with pytest.raises(ValueError, match="expected None, 'sequential', 'fsdp2'"): + load_pretrained(str(tmp_path), distribution="ep") + + +def test_a_whole_model_load_refuses_a_sharded_loads_config(serve, auto_class, tmp_path): + """from_pretrained would take hf_config for a model argument and fail far from here.""" + config = serve(LlamaConfig(**_TINY)) + + with pytest.raises(TypeError, match="hf_config applies to a sharded load only"): + load_pretrained(str(tmp_path), auto_class, hf_config=config) From 3755dc95e13dc06f9d93865f820336dbc8113e74 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 23:07:12 +0000 Subject: [PATCH 02/41] Load the plan rank 0 was handed on every rank load_pretrained kept a plan the caller passed as session.plan, which holds rank 0's, but routed and loaded with this rank's own: ranks handed different plans built mismatched meshes and hung while the session said they agreed. Every rank now loads the plan the session holds, once its own plan passed the checks, so a refused plan still leaves the session as it was. A plan the session decides itself is already held, and is no longer sent a second time. The test that rank 0 alone works out the unplaced keys answers the plan the load now decides first. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 10 +++++--- .../test_load_pretrained_distribution.py | 24 +++++++++++++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d6f915f5dcf..a58618300ff 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -626,15 +626,19 @@ def load_pretrained( if subfolder: ckpt_path = os.path.join(ckpt_path, subfolder) trust_remote_code = kwargs.get("trust_remote_code", False) + session = _distributed_session() if distribution is None: plan, hf_config = _decide_distribution(ckpt_path, trust_remote_code, None, None) distribution = "sequential" if plan is None else plan if plan is not None and plan.is_fsdp_dense: kwargs.setdefault("hf_config", hf_config) + elif session is not None and isinstance(distribution, DistributionPlan): + _route(distribution) # A plan this rank refuses leaves the session as it was. + # Every rank loads the plan rank 0 was handed, as it would one the session decided: ranks + # loading different plans would build mismatched meshes and hang. + session.plan = distribution + distribution = session.plan route = _route(distribution) - session = _distributed_session() - if session is not None and isinstance(distribution, DistributionPlan): - session.plan = distribution # What the model is loaded with, the same on every rank. if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. config = kwargs.get("hf_config") or kwargs.get("config") diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index f3981762752..f51f9c7fc1c 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -337,6 +337,30 @@ def test_an_expert_parallel_plan_is_refused(session, serve, loaders, auto_class, assert auto_class.calls == [] +def test_every_rank_loads_the_plan_rank0_was_handed( + session, serve, loaders, process_group, auto_class, monkeypatch, tmp_path +): + """Ranks loading different plans would build mismatched meshes and hang: rank 1, handed a plan + of its own, loads the one rank 0 was handed.""" + active = session(2, rank=1) + serve(LlamaConfig(**_TINY)) + rank0_plan = DistributionPlan(2, fsdp_size=2) + + def broadcast_object_list(message, group=None, group_src=None): + label = message[0][0] + message[0] = (label, rank0_plan if label == "plan" else [], None) + + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 1) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) + monkeypatch.setattr(model_load_utils, "AutoModelForCausalLM", auto_class) + + load_pretrained(str(tmp_path), distribution=DistributionPlan(2)) + + assert auto_class.calls == [], "rank 1 loaded the whole model, as its own plan says" + assert "fsdp2" in loaders + assert active.plan is rank0_plan + + def test_the_plan_a_load_runs_is_kept_on_the_session( session, serve, loaders, process_group, tmp_path ): From 78613c557ba47c440a30e9d717d8f87c546c8dbb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 02:38:28 +0000 Subject: [PATCH 03/41] Take a plan the caller hands load_pretrained through the plan handle Every rank loads the plan rank 0 was handed, now set with session.plan.decide(lambda: handed) rather than an assignment, and read back with session.plan.value. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 4 ++-- .../unit/torch/utils/test_load_pretrained_distribution.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index a58618300ff..6310e44f755 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -636,8 +636,8 @@ def load_pretrained( _route(distribution) # A plan this rank refuses leaves the session as it was. # Every rank loads the plan rank 0 was handed, as it would one the session decided: ranks # loading different plans would build mismatched meshes and hang. - session.plan = distribution - distribution = session.plan + handed = distribution + distribution = session.plan.decide(lambda: handed) route = _route(distribution) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index f51f9c7fc1c..c3f2dc1316a 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -278,7 +278,7 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( assert auto_class.calls == [{"device_map": placed}] assert loaders == {} - assert (active.plan.ep_size, active.plan.fsdp_size) == (1, 1) + assert (active.plan.value.ep_size, active.plan.value.fsdp_size) == (1, 1) @pytest.mark.parametrize("distribution", ["fsdp2", DistributionPlan(2, fsdp_size=2)]) @@ -332,7 +332,7 @@ def test_an_expert_parallel_plan_is_refused(session, serve, loaders, auto_class, with pytest.raises(NotImplementedError, match="shards the routed experts"): load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) - assert active.plan is None + assert active.plan.value is None assert loaders == {} assert auto_class.calls == [] @@ -358,7 +358,7 @@ def broadcast_object_list(message, group=None, group_src=None): assert auto_class.calls == [], "rank 1 loaded the whole model, as its own plan says" assert "fsdp2" in loaders - assert active.plan is rank0_plan + assert repr(active.plan.value) == repr(rank0_plan) def test_the_plan_a_load_runs_is_kept_on_the_session( @@ -370,7 +370,7 @@ def test_the_plan_a_load_runs_is_kept_on_the_session( load_pretrained(str(tmp_path), distribution=DistributionPlan(2, fsdp_size=2)) - assert active.plan.fsdp_size == 2 + assert active.plan.value.fsdp_size == 2 def test_a_plan_for_another_world_is_refused(session, serve, tmp_path): From 64faf38ae39cbbe74b00b3b68ca3ba401c98f85d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 03:08:45 +0000 Subject: [PATCH 04/41] Hand every rank rank 0's plan through the renamed broadcast_from Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 6310e44f755..41df34b6203 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -637,7 +637,7 @@ def load_pretrained( # Every rank loads the plan rank 0 was handed, as it would one the session decided: ranks # loading different plans would build mismatched meshes and hang. handed = distribution - distribution = session.plan.decide(lambda: handed) + distribution = session.plan.broadcast_from(lambda: handed) route = _route(distribution) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. From f94b9188686d05752b38718a8a1be3cc0ad58d14 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 17:21:09 +0000 Subject: [PATCH 05/41] Route load_pretrained by the live process group instead of the session Rank and world size come from the process group, and a plan handed to every rank goes out from rank 0 through broadcast_from. Rank 0 now checks that plan before it sends, so a plan it refuses raises on every rank instead of leaving the others in the broadcast. Nothing keeps the plan after the load. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/utils/plugins/model_load_utils.py | 64 +++++----- .../test_load_pretrained_distribution.py | 15 +-- .../test_load_pretrained_distribution.py | 109 +++++++----------- 4 files changed, 78 insertions(+), 112 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3ca53274a89..0362518f83c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -148,7 +148,7 @@ Changelog - Add a ``kd_loss_alpha`` option to HuggingFace ``KDTrainer`` to blend cross-entropy and knowledge distillation losses during training. - A tracked ``examples/hf_ptq/hf_ptq.py`` run now writes ``.experiment.json`` into ``--export_path`` and uploads the same file with the run, so a checkpoint on disk names the experiment and MLflow run id that produced it. The pointer is written only once the export completes, and an export that is not tracked removes one it would otherwise inherit from a reused ``--export_path`` or from a quantized source checkpoint. - ``export_hf_checkpoint`` and the vLLM fake-quant export now copy the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged. They read the source from local disk only: get a local copy of a Hugging Face Hub model with ``modelopt.torch.export.ensure_local_checkpoint`` and load it from there. -- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Inside a distributed ``modelopt.torch.utils.distributed.modelopt_session()`` it FSDP2-shards the model or gives every rank a whole copy, as the model and the ``MODELOPT_FSDP_SIZE`` variable decide; its new ``distribution`` argument picks the load explicitly. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, as the model and the ``MODELOPT_FSDP_SIZE`` variable decide; its new ``distribution`` argument picks the load explicitly. - Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with Qwen-Image training, distributed-checkpoint export, and PDD-2/4/8 inference. AutoModel remains an unmodified pinned runtime dependency. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 41df34b6203..b61a17815f3 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -417,30 +417,22 @@ def record_unplaced_source_keys( } -def _distributed_session(): - """The active ``ModelOptSession`` when it runs distributed, else ``None``.""" - from modelopt.torch.utils.distributed import get_active_mo_session - - session = get_active_mo_session(required=False) - return session if session.initialized and session.distributed else None - - def _decide_distribution(ckpt_path: str, trust_remote_code: bool, ep_size, fsdp_size): """:func:`resolve_distribution_plan` for a local ``ckpt_path``.""" + from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env from modelopt.torch.utils.moe import decide_distribution_plan - session = _distributed_session() - if session is None: + if not dist.is_distributed_run(): return None, None hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code) plan = decide_distribution_plan( hf_config, - session.world_size, + dist.size(), ep_size=env.ep_degree() if ep_size is None else ep_size, fsdp_size=env.fsdp_size() if fsdp_size is None else fsdp_size, ) - if session.is_main: + if dist.is_master(): print(f"Distribution: {plan!r}", flush=True) return plan, hf_config @@ -452,21 +444,23 @@ def resolve_distribution_plan( ep_size: "int | str | None" = None, fsdp_size: "int | str | None" = None, ) -> tuple[DistributionPlan | None, PretrainedConfig | None]: - """Decide how :func:`load_pretrained` distributes ``ckpt_path`` in this session. + """Decide how :func:`load_pretrained` distributes ``ckpt_path`` in this run. Returns ``(plan, hf_config)``: a :class:`~modelopt.torch.utils.distributed.DistributionPlan` and the config it was decided from, which a sharded load takes back as ``load_pretrained(ckpt_path, distribution=plan, hf_config=hf_config)`` rather than read it - again. Outside a distributed :class:`~modelopt.torch.utils.distributed.ModelOptSession`: - ``(None, None)``, the model loads whole. ``load_pretrained`` decides the same way on its own; - deciding first is for a caller that loads an unsharded model differently. + again. Outside a distributed run: ``(None, None)``, the model loads whole. + ``load_pretrained`` decides the same way on its own; deciding first is for a caller that loads + an unsharded model differently. :func:`~modelopt.torch.utils.moe.decide_distribution_plan` decides, once on rank 0, from ``ep_size`` and ``fsdp_size``; ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` (:mod:`modelopt.torch.utils.env`), which default to ``"auto"``. Collective: every rank must call it. """ - if _distributed_session() is None: + from modelopt.torch.utils import distributed as dist + + if not dist.is_distributed_run(): return None, None _, ckpt_path = ensure_local_checkpoint(ckpt_path) return _decide_distribution(ckpt_path, trust_remote_code, ep_size, fsdp_size) @@ -481,8 +475,10 @@ def _route(distribution: DistributionPlan | str) -> str: f"distribution={distribution!r}: expected None, 'sequential', 'fsdp2' or a " "DistributionPlan." ) - plan, session = distribution, _distributed_session() - world_size = session.world_size if session is not None else 1 + from modelopt.torch.utils import distributed as dist + + plan, distributed = distribution, dist.is_distributed_run() + world_size = dist.size() if distributed else 1 if plan.world_size != world_size: raise ValueError(f"{plan!r} is for {plan.world_size} ranks, but this run has {world_size}.") if plan.is_expert_parallel: @@ -498,7 +494,7 @@ def _route(distribution: DistributionPlan | str) -> str: f"so fsdp_size must be {plan.world_size}." ) return _FSDP2 - return "replicated" if session is not None else "sequential" + return "replicated" if distributed else "sequential" def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) -> nn.Module: @@ -561,9 +557,8 @@ def load_pretrained( ``distribution`` decides how the model is placed: - * ``None`` (default): the active :class:`~modelopt.torch.utils.distributed.ModelOptSession` - decides. Without one, or in a single-process one, the model loads whole; in a distributed one, - as :func:`resolve_distribution_plan` decides. + * ``None`` (default): a single process loads the model whole; a distributed run loads it as + :func:`resolve_distribution_plan` decides. * ``"sequential"``: ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to ``AutoModelForCausalLM``), placed as ``kwargs`` say. Each rank that calls it gets the whole model. @@ -576,8 +571,8 @@ def load_pretrained( - ``fsdp_size > 1``: as ``"fsdp2"``; - otherwise as ``"sequential"`` on every rank, ``device_map`` defaulting to the rank's GPU. - A plan with ``ep_size > 1`` raises: no load here shards the routed experts. The plan is kept - as ``session.plan``, the same on every rank. + A plan with ``ep_size > 1`` raises: no load here shards the routed experts. On a distributed + run every rank loads the plan rank 0 was handed. A sharded load builds the model class and places it itself. Besides its own options it takes ``trust_remote_code``, ``attn_implementation`` and ``hf_config`` (the config to load with), and @@ -616,6 +611,7 @@ def load_pretrained( ) # Importing modelopt.torch.models registers every model spec; only callers of this need them. from modelopt.torch.models.hf import prepare_model_for_loading + from modelopt.torch.utils import distributed as dist hub_options = { k: kwargs.pop(k) for k in ("token", "cache_dir", "local_files_only") if k in kwargs @@ -626,18 +622,22 @@ def load_pretrained( if subfolder: ckpt_path = os.path.join(ckpt_path, subfolder) trust_remote_code = kwargs.get("trust_remote_code", False) - session = _distributed_session() if distribution is None: plan, hf_config = _decide_distribution(ckpt_path, trust_remote_code, None, None) distribution = "sequential" if plan is None else plan if plan is not None and plan.is_fsdp_dense: kwargs.setdefault("hf_config", hf_config) - elif session is not None and isinstance(distribution, DistributionPlan): - _route(distribution) # A plan this rank refuses leaves the session as it was. - # Every rank loads the plan rank 0 was handed, as it would one the session decided: ranks - # loading different plans would build mismatched meshes and hang. + elif dist.is_distributed_run() and isinstance(distribution, DistributionPlan): + # Every rank loads the plan rank 0 was handed, as it would one it decided: ranks loading + # different plans would build mismatched meshes and hang. Rank 0 checks it before it sends, + # so a plan it refuses raises on every rank instead of leaving them in the broadcast. handed = distribution - distribution = session.plan.broadcast_from(lambda: handed) + + def checked() -> DistributionPlan: + _route(handed) + return handed + + distribution = dist.broadcast_from("plan", checked) route = _route(distribution) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. @@ -657,7 +657,7 @@ def load_pretrained( if route == "replicated": # Each rank takes the whole model onto its own GPU. Left to "auto", every rank's # accelerate could pick the same idle GPU. - kwargs.setdefault("device_map", {"": str(session.device)}) + kwargs.setdefault("device_map", {"": f"cuda:{dist.local_rank()}"}) auto_class = auto_class or AutoModelForCausalLM with prepare_model_for_loading( model_type, diff --git a/tests/gpu/torch/utils/test_load_pretrained_distribution.py b/tests/gpu/torch/utils/test_load_pretrained_distribution.py index 3abb9d9dd44..9f3e58c3a9a 100644 --- a/tests/gpu/torch/utils/test_load_pretrained_distribution.py +++ b/tests/gpu/torch/utils/test_load_pretrained_distribution.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""``load_pretrained`` in a distributed session, on real GPUs. +"""``load_pretrained`` in a distributed run, on real GPUs. Left to decide, it FSDP2-shards a dense model when asked to, and otherwise gives every rank the whole model on its own GPU. Whichever load runs, every rank ends up with the checkpoint's weights and the @@ -30,7 +30,6 @@ from torch.distributed.tensor import DTensor from transformers import AutoModelForCausalLM, LlamaConfig -from modelopt.torch.utils.distributed import get_active_mo_session from modelopt.torch.utils.plugins.model_load_utils import load_pretrained _TINY = { @@ -54,17 +53,8 @@ def _full(tensor): def _load_and_check(rank, size, ckpt_dir, knobs, expected): - """Enter a distributed session as torchrun would, load with the plan left to decide, check it.""" + """Load on every rank of the run with the plan left to decide, and check what each rank holds.""" device = torch.device(f"cuda:{rank}") - session = get_active_mo_session(required=False) - session.initialize( - distributed=True, - rank=rank, - world_size=size, - local_rank=rank, - device=device, - is_main=rank == 0, - ) os.environ.update(knobs) try: model = load_pretrained(ckpt_dir) @@ -86,7 +76,6 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected): finally: for knob in knobs: del os.environ[knob] - session._reset() def test_a_dense_model_loads_whole_on_every_rank(need_2_gpus, dist_workers, tmp_path): diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index c3f2dc1316a..fb01956b9a3 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -15,8 +15,8 @@ """How ``load_pretrained`` decides a model's distribution, and which load that routes it to. -Nothing here needs a process group or a GPU: the active ``ModelOptSession`` is initialized the way a -torchrun launch leaves it, the configs are tiny real transformers configs served from memory, and the +Nothing here needs a process group or a GPU: the process group's rank and size are answered the way +a torchrun launch would, the configs are tiny real transformers configs served from memory, and the loaders are replaced by stand-ins that record how they were called. """ @@ -34,7 +34,7 @@ import modelopt.torch.utils.distributed as dist import modelopt.torch.utils.moe as moe from modelopt.torch.utils import env -from modelopt.torch.utils.distributed import DistributionPlan, get_active_mo_session +from modelopt.torch.utils.distributed import DistributionPlan from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import load_pretrained, resolve_distribution_plan @@ -63,29 +63,22 @@ def _no_distribution_knobs(monkeypatch): @pytest.fixture -def session(monkeypatch): - """Enter the active session as rank ``rank`` of a ``world_size``-rank launch would. +def run_as(monkeypatch): + """Run as rank ``rank`` of a ``world_size``-rank launch. Every rank acts as its group's first, deciding for itself: these tests are about what is - decided, and test_distributed_session.py about how the ranks agree on it. + decided, and test_distributed_run.py about how the ranks agree on it. """ monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) - active = get_active_mo_session(required=False) def enter(world_size, rank=0): - active.initialize( - distributed=world_size > 1, - rank=rank, - world_size=world_size, - local_rank=rank, - device=torch.device(f"cuda:{rank}"), - is_main=rank == 0, - ) - return active + monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1) + monkeypatch.setattr(dist, "size", lambda group=None: world_size) + monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) + monkeypatch.setattr(dist, "local_rank", lambda: rank) - yield enter - active._reset() + return enter @pytest.fixture @@ -185,11 +178,11 @@ def test_distribution_knobs_are_read_when_called(monkeypatch, knob, name, other) # -------------------------------------------------------------------------------- the decision -@pytest.mark.parametrize("world_size", [None, 1], ids=["no_session", "single_process_session"]) -def test_outside_a_distributed_session_nothing_is_distributed(session, serve, tmp_path, world_size): +@pytest.mark.parametrize("world_size", [None, 1], ids=["no_process_group", "one_rank"]) +def test_outside_a_distributed_run_nothing_is_distributed(run_as, serve, tmp_path, world_size): """With no ranks to shard over, even a MoE loads whole.""" if world_size is not None: - session(world_size) + run_as(world_size) serve(_qwen3_moe(num_experts=8)) assert resolve_distribution_plan(str(tmp_path)) == (None, None) @@ -200,10 +193,10 @@ def test_outside_a_distributed_session_nothing_is_distributed(session, serve, tm reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", ) @pytest.mark.parametrize("rank", [0, 1]) -def test_an_ep_capable_moe_is_planned_expert_parallel(session, serve, tmp_path, capsys, rank): +def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, capsys, rank): """Auto picks the largest EP degree dividing both the world and the expert count, hands back the config it decided from, and only rank 0 reports the decision.""" - session(4, rank) + run_as(4, rank) serve(_qwen3_moe(num_experts=8)) plan, hf_config = resolve_distribution_plan(str(tmp_path)) @@ -214,9 +207,9 @@ def test_an_ep_capable_moe_is_planned_expert_parallel(session, serve, tmp_path, assert ("Distribution:" in capsys.readouterr().out) == (rank == 0) -def test_a_dense_model_is_replicated(session, serve, tmp_path): +def test_a_dense_model_is_replicated(run_as, serve, tmp_path): """A dense model that fits one GPU has nothing to shard: every rank holds it whole.""" - session(4) + run_as(4) serve(LlamaConfig(**_TINY)) plan, _ = resolve_distribution_plan(str(tmp_path)) @@ -224,11 +217,11 @@ def test_a_dense_model_is_replicated(session, serve, tmp_path): assert (plan.ep_size, plan.fsdp_size) == (1, 1) -def test_the_env_knobs_override_the_automatic_plan(session, serve, monkeypatch, tmp_path): +def test_the_env_knobs_override_the_automatic_plan(run_as, serve, monkeypatch, tmp_path): """The decision honours both knobs over its automatic choice.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") - session(4) + run_as(4) serve(_qwen3_moe(num_experts=8)) plan, _ = resolve_distribution_plan(str(tmp_path)) @@ -236,9 +229,9 @@ def test_the_env_knobs_override_the_automatic_plan(session, serve, monkeypatch, assert (plan.ep_size, plan.edp, plan.fsdp_size) == (2, 2, 4) -def test_arguments_override_the_env_knobs(session, serve, monkeypatch, tmp_path): +def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path): monkeypatch.setenv("MODELOPT_EP_DEGREE", "4") - session(4) + run_as(4) serve(_qwen3_moe(num_experts=8)) plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=2, fsdp_size=1) @@ -249,8 +242,8 @@ def test_arguments_override_the_env_knobs(session, serve, monkeypatch, tmp_path) # ----------------------------------------------------------------------------------- the loads -def test_without_a_session_the_model_loads_whole_as_asked(serve, loaders, auto_class, tmp_path): - """No session, no plan: from_pretrained places the model exactly as the caller says.""" +def test_a_single_process_loads_the_model_whole_as_asked(serve, loaders, auto_class, tmp_path): + """No process group, no plan: from_pretrained places the model exactly as the caller says.""" serve(LlamaConfig(**_TINY)) model = load_pretrained(str(tmp_path), auto_class, device_map="auto", dtype="auto") @@ -266,11 +259,11 @@ def test_without_a_session_the_model_loads_whole_as_asked(serve, loaders, auto_c ids=["defaults_to_the_rank_gpu", "caller_placement_wins"], ) def test_a_replicated_plan_loads_the_whole_model_on_every_rank( - session, serve, loaders, auto_class, tmp_path, device_map, placed + run_as, serve, loaders, auto_class, tmp_path, device_map, placed ): """Each rank takes the whole model onto its own GPU; "auto" would let every rank's accelerate - pick the same idle GPU. The plan the load decided is kept on the session.""" - active = session(4, rank=1) + pick the same idle GPU.""" + run_as(4, rank=1) serve(LlamaConfig(**_TINY)) kwargs = {} if device_map is None else {"device_map": device_map} @@ -278,15 +271,14 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( assert auto_class.calls == [{"device_map": placed}] assert loaders == {} - assert (active.plan.value.ep_size, active.plan.value.fsdp_size) == (1, 1) @pytest.mark.parametrize("distribution", ["fsdp2", DistributionPlan(2, fsdp_size=2)]) def test_fsdp2_loads_through_the_parallel_fsdp2_loader( - session, serve, loaders, process_group, tmp_path, distribution + run_as, serve, loaders, process_group, tmp_path, distribution ): """Named or planned, FSDP2 shards over the process group it reads its rank and GPU from.""" - session(2, rank=1) + run_as(2, rank=1) served = serve(LlamaConfig(**_TINY)) load_pretrained(str(tmp_path), distribution=distribution, cpu_offload=True, dtype="auto") @@ -299,11 +291,11 @@ def test_fsdp2_loads_through_the_parallel_fsdp2_loader( def test_a_load_it_decides_reads_the_config_once( - session, loaders, process_group, monkeypatch, tmp_path + run_as, loaders, process_group, monkeypatch, tmp_path ): """The config the plan was decided from is the one the FSDP2 load gets, not a second read.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") - session(2, rank=1) + run_as(2, rank=1) config, reads = LlamaConfig(**_TINY), [] monkeypatch.setattr( transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a) or config @@ -315,34 +307,32 @@ def test_a_load_it_decides_reads_the_config_once( assert loaders["fsdp2"]["hf_config"] is config -def test_partial_dense_sharding_without_expert_parallelism_is_refused(session, serve, tmp_path): +def test_partial_dense_sharding_without_expert_parallelism_is_refused(run_as, serve, tmp_path): """The FSDP2 load shards over the whole world; it cannot honour a smaller dense degree.""" - session(4) + run_as(4) serve(LlamaConfig(**_TINY)) with pytest.raises(NotImplementedError, match="fsdp_size must be 4"): load_pretrained(str(tmp_path), distribution=DistributionPlan(4, fsdp_size=2)) -def test_an_expert_parallel_plan_is_refused(session, serve, loaders, auto_class, tmp_path): - """No load here shards the routed experts: an EP plan fails before anything loads, and is not - kept on the session.""" - active = session(2) +def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, tmp_path): + """No load here shards the routed experts: an EP plan fails before anything loads.""" + run_as(2) serve(_qwen3_moe(num_experts=8)) with pytest.raises(NotImplementedError, match="shards the routed experts"): load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) - assert active.plan.value is None assert loaders == {} assert auto_class.calls == [] def test_every_rank_loads_the_plan_rank0_was_handed( - session, serve, loaders, process_group, auto_class, monkeypatch, tmp_path + run_as, serve, loaders, process_group, auto_class, monkeypatch, tmp_path ): """Ranks loading different plans would build mismatched meshes and hang: rank 1, handed a plan of its own, loads the one rank 0 was handed.""" - active = session(2, rank=1) + run_as(2, rank=1) serve(LlamaConfig(**_TINY)) rank0_plan = DistributionPlan(2, fsdp_size=2) @@ -357,24 +347,11 @@ def broadcast_object_list(message, group=None, group_src=None): load_pretrained(str(tmp_path), distribution=DistributionPlan(2)) assert auto_class.calls == [], "rank 1 loaded the whole model, as its own plan says" - assert "fsdp2" in loaders - assert repr(active.plan.value) == repr(rank0_plan) - - -def test_the_plan_a_load_runs_is_kept_on_the_session( - session, serve, loaders, process_group, tmp_path -): - """Handed a plan, the load records it, as it does one it decides itself.""" - active = session(2) - serve(LlamaConfig(**_TINY)) - - load_pretrained(str(tmp_path), distribution=DistributionPlan(2, fsdp_size=2)) - - assert active.plan.value.fsdp_size == 2 + assert "fsdp2" in loaders, "rank 0's plan shards the model" -def test_a_plan_for_another_world_is_refused(session, serve, tmp_path): - session(2) +def test_a_plan_for_another_world_is_refused(run_as, serve, tmp_path): + run_as(2) serve(LlamaConfig(**_TINY)) with pytest.raises(ValueError, match="is for 4 ranks, but this run has 2"): @@ -391,11 +368,11 @@ def test_a_plan_for_another_world_is_refused(session, serve, tmp_path): ids=["placement", "model_class", "mtp_head"], ) def test_a_sharded_load_refuses_what_it_cannot_honor( - session, serve, loaders, process_group, monkeypatch, tmp_path, change, error, match + run_as, serve, loaders, process_group, monkeypatch, tmp_path, change, error, match ): """Placement, a model class, a head only a whole model builds: none reaches a sharded load, which fails before it starts rather than load something other than what was asked.""" - session(2) + run_as(2) serve(LlamaConfig(**_TINY)) if change.get("mtp"): monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) From d6bd85deabc13c917dd965e5bc628fdc5c39b97d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 19:48:15 +0000 Subject: [PATCH 06/41] Load by a distribution plan only, and record it on the loaded model load_pretrained's distribution is now a DistributionPlan or None, which decides through resolve_distribution_plan; the 'sequential' and 'fsdp2' strings that picked a load without a plan are gone. A single process loads by a single process's plan, so every load has one. Whatever the route, the plan is recorded on the model and every submodule, where calibration and export read it back instead of inspecting the model. parallel_load_and_prepare_fsdp2 records its own plan as well, for a caller that uses it directly. The decision's copy here goes too: it lives in resolve_distribution_plan now. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/utils/plugins/model_load_utils.py | 125 ++++-------- .../test_load_pretrained_distribution.py | 3 + .../test_load_pretrained_distribution.py | 181 ++++-------------- 4 files changed, 77 insertions(+), 234 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0362518f83c..95d0972cf22 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -148,7 +148,7 @@ Changelog - Add a ``kd_loss_alpha`` option to HuggingFace ``KDTrainer`` to blend cross-entropy and knowledge distillation losses during training. - A tracked ``examples/hf_ptq/hf_ptq.py`` run now writes ``.experiment.json`` into ``--export_path`` and uploads the same file with the run, so a checkpoint on disk names the experiment and MLflow run id that produced it. The pointer is written only once the export completes, and an export that is not tracked removes one it would otherwise inherit from a reused ``--export_path`` or from a quantized source checkpoint. - ``export_hf_checkpoint`` and the vLLM fake-quant export now copy the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged. They read the source from local disk only: get a local copy of a Hugging Face Hub model with ``modelopt.torch.export.ensure_local_checkpoint`` and load it from there. -- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, as the model and the ``MODELOPT_FSDP_SIZE`` variable decide; its new ``distribution`` argument picks the load explicitly. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_FSDP_SIZE`` variable, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). - Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with Qwen-Image training, distributed-checkpoint export, and PDD-2/4/8 inference. AutoModel remains an unmodified pinned runtime dependency. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index b61a17815f3..cb757e606d0 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -44,6 +44,7 @@ broadcast_state_dict, fsdp2_wrap, is_initialized, + set_distribution_plan, ) from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -417,84 +418,21 @@ def record_unplaced_source_keys( } -def _decide_distribution(ckpt_path: str, trust_remote_code: bool, ep_size, fsdp_size): - """:func:`resolve_distribution_plan` for a local ``ckpt_path``.""" +def _route(plan: DistributionPlan) -> str | None: + """The sharded load ``plan`` takes, or ``None`` when every rank loads the whole model.""" from modelopt.torch.utils import distributed as dist - from modelopt.torch.utils import env - from modelopt.torch.utils.moe import decide_distribution_plan - - if not dist.is_distributed_run(): - return None, None - hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code) - plan = decide_distribution_plan( - hf_config, - dist.size(), - ep_size=env.ep_degree() if ep_size is None else ep_size, - fsdp_size=env.fsdp_size() if fsdp_size is None else fsdp_size, - ) - if dist.is_master(): - print(f"Distribution: {plan!r}", flush=True) - return plan, hf_config - - -def resolve_distribution_plan( - ckpt_path: str, - *, - trust_remote_code: bool = False, - ep_size: "int | str | None" = None, - fsdp_size: "int | str | None" = None, -) -> tuple[DistributionPlan | None, PretrainedConfig | None]: - """Decide how :func:`load_pretrained` distributes ``ckpt_path`` in this run. - - Returns ``(plan, hf_config)``: a :class:`~modelopt.torch.utils.distributed.DistributionPlan` - and the config it was decided from, which a sharded load takes back as - ``load_pretrained(ckpt_path, distribution=plan, hf_config=hf_config)`` rather than read it - again. Outside a distributed run: ``(None, None)``, the model loads whole. - ``load_pretrained`` decides the same way on its own; deciding first is for a caller that loads - an unsharded model differently. - - :func:`~modelopt.torch.utils.moe.decide_distribution_plan` decides, once on rank 0, from - ``ep_size`` and ``fsdp_size``; ``None`` reads ``MODELOPT_EP_DEGREE`` and ``MODELOPT_FSDP_SIZE`` - (:mod:`modelopt.torch.utils.env`), which default to ``"auto"``. Collective: every rank must - call it. - """ - from modelopt.torch.utils import distributed as dist - - if not dist.is_distributed_run(): - return None, None - _, ckpt_path = ensure_local_checkpoint(ckpt_path) - return _decide_distribution(ckpt_path, trust_remote_code, ep_size, fsdp_size) - -def _route(distribution: DistributionPlan | str) -> str: - """Which load ``distribution`` asks for: a sharded one, ``"replicated"``, or ``"sequential"``.""" - if isinstance(distribution, str): - if distribution in ("sequential", "fsdp2"): - return _FSDP2 if distribution == "fsdp2" else "sequential" + if plan.world_size != dist.size(): raise ValueError( - f"distribution={distribution!r}: expected None, 'sequential', 'fsdp2' or a " - "DistributionPlan." + f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." ) - from modelopt.torch.utils import distributed as dist - - plan, distributed = distribution, dist.is_distributed_run() - world_size = dist.size() if distributed else 1 - if plan.world_size != world_size: - raise ValueError(f"{plan!r} is for {plan.world_size} ranks, but this run has {world_size}.") if plan.is_expert_parallel: raise NotImplementedError( f"{plan!r} shards the routed experts, which load_pretrained has no load for. Plan " "ep_size=1 (MODELOPT_EP_DEGREE=1) to give every rank the whole model, or shard it " "with FSDP2." ) - if plan.is_fsdp_dense: - if plan.fsdp_size != plan.world_size: - raise NotImplementedError( - f"{plan!r}: without expert parallelism the FSDP2 load shards over the whole world, " - f"so fsdp_size must be {plan.world_size}." - ) - return _FSDP2 - return "replicated" if distributed else "sequential" + return _FSDP2 if plan.is_fsdp_dense else None def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) -> nn.Module: @@ -550,29 +488,25 @@ def load_pretrained( auto_class: Any = None, *, model_type: str | None = None, - distribution: DistributionPlan | str | None = None, + distribution: DistributionPlan | None = None, **kwargs, ) -> nn.Module: """Load a Hugging Face model for ModelOpt to optimize and export, whole or distributed. - ``distribution`` decides how the model is placed: + ``distribution`` is the :class:`~modelopt.torch.utils.distributed.DistributionPlan` to spread + the model over the run's ranks by, :func:`resolve_distribution_plan`'s by default. By the plan: - * ``None`` (default): a single process loads the model whole; a distributed run loads it as - :func:`resolve_distribution_plan` decides. - * ``"sequential"``: ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` - defaults to ``AutoModelForCausalLM``), placed as ``kwargs`` say. Each rank that calls it gets - the whole model. - * ``"fsdp2"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its + * ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, ``mp_policy`` and ``broadcast_chunk_size`` ``kwargs`` may set). - * A :class:`~modelopt.torch.utils.distributed.DistributionPlan`, as - :func:`resolve_distribution_plan` returns it: - - - ``fsdp_size > 1``: as ``"fsdp2"``; - - otherwise as ``"sequential"`` on every rank, ``device_map`` defaulting to the rank's GPU. + * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to + ``AutoModelForCausalLM``), placed as ``kwargs`` say. Every rank of a distributed run gets the + whole model, ``device_map`` defaulting to the rank's GPU. - A plan with ``ep_size > 1`` raises: no load here shards the routed experts. On a distributed - run every rank loads the plan rank 0 was handed. + A plan with ``ep_size > 1`` raises: no load here shards the routed experts. On a distributed run + every rank loads the plan rank 0 was handed. The plan is recorded on the model + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), where calibration and export + read it back. A sharded load builds the model class and places it itself. Besides its own options it takes ``trust_remote_code``, ``attn_implementation`` and ``hf_config`` (the config to load with), and @@ -622,12 +556,12 @@ def load_pretrained( if subfolder: ckpt_path = os.path.join(ckpt_path, subfolder) trust_remote_code = kwargs.get("trust_remote_code", False) + hf_config = kwargs.get("hf_config") if distribution is None: - plan, hf_config = _decide_distribution(ckpt_path, trust_remote_code, None, None) - distribution = "sequential" if plan is None else plan - if plan is not None and plan.is_fsdp_dense: + plan, hf_config = resolve_distribution_plan(ckpt_path, trust_remote_code=trust_remote_code) + if _route(plan) is not None: kwargs.setdefault("hf_config", hf_config) - elif dist.is_distributed_run() and isinstance(distribution, DistributionPlan): + elif isinstance(distribution, DistributionPlan): # Every rank loads the plan rank 0 was handed, as it would one it decided: ranks loading # different plans would build mismatched meshes and hang. Rank 0 checks it before it sends, # so a plan it refuses raises on every rank instead of leaving them in the broadcast. @@ -637,24 +571,26 @@ def checked() -> DistributionPlan: _route(handed) return handed - distribution = dist.broadcast_from("plan", checked) - route = _route(distribution) + plan = dist.broadcast_from("plan", checked) + else: + raise TypeError(f"distribution={distribution!r}: expected a DistributionPlan or None.") + route = _route(plan) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. - config = kwargs.get("hf_config") or kwargs.get("config") + config = hf_config or kwargs.get("config") model_type = getattr(config, "model_type", None) or ( AutoConfig.from_pretrained( ckpt_path, trust_remote_code=trust_remote_code, **hub_options ).model_type ) - if route in _SHARDED_LOAD_OPTIONS: + if route is not None: model = _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: if "hf_config" in kwargs: raise TypeError("hf_config applies to a sharded load only.") - if route == "replicated": + if plan.world_size > 1: # Each rank takes the whole model onto its own GPU. Left to "auto", every rank's # accelerate could pick the same idle GPU. kwargs.setdefault("device_map", {"": f"cuda:{dist.local_rank()}"}) @@ -667,6 +603,7 @@ def checked() -> DistributionPlan: ): model = auto_class.from_pretrained(ckpt_path, **hub_options, **kwargs) keys = None + set_distribution_plan(model, plan) unplaced = record_unplaced_source_keys(model, ckpt_path, keys) if unplaced: print_rank_0( @@ -1002,6 +939,9 @@ def parallel_load_and_prepare_fsdp2( ``broadcast_chunk_size`` sets how many of a source's owned layers are broadcast per collective: a smaller value lowers the peak transient GPU memory at the cost of more collectives (default 8; pass ``None`` to broadcast all of a source's layers at once). + + The model records how it is distributed + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) weight_map = indexed_weight_map(resolved_path) @@ -1093,6 +1033,7 @@ def parallel_load_and_prepare_fsdp2( _promote_non_dtensor_to_gpu(model, device) if hasattr(model, "tie_weights"): model.tie_weights() + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) return model diff --git a/tests/gpu/torch/utils/test_load_pretrained_distribution.py b/tests/gpu/torch/utils/test_load_pretrained_distribution.py index 9f3e58c3a9a..7f9eb7a24f3 100644 --- a/tests/gpu/torch/utils/test_load_pretrained_distribution.py +++ b/tests/gpu/torch/utils/test_load_pretrained_distribution.py @@ -30,6 +30,7 @@ from torch.distributed.tensor import DTensor from transformers import AutoModelForCausalLM, LlamaConfig +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins.model_load_utils import load_pretrained _TINY = { @@ -73,6 +74,8 @@ def _load_and_check(rank, size, ckpt_dir, knobs, expected): torch.testing.assert_close(_full(param).cpu(), whole_params[name]) assert model._modelopt_unplaced_source_keys == [] assert model._modelopt_source_checkpoint == ckpt_dir + dense = "replicate" if expected == "replicated" else "fsdp" + assert distribution_plan(model) == DistributionPlan(size, dense=dense) finally: for knob in knobs: del os.environ[knob] diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index fb01956b9a3..d55dbfc60c6 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""How ``load_pretrained`` decides a model's distribution, and which load that routes it to. +"""Which load ``load_pretrained`` routes a distribution plan to, and the plan it records. Nothing here needs a process group or a GPU: the process group's rank and size are answered the way a torchrun launch would, the configs are tiny real transformers configs served from memory, and the @@ -27,19 +27,14 @@ transformers = pytest.importorskip("transformers") -from packaging.version import Version from transformers import AutoModel, LlamaConfig, Qwen3MoeConfig import modelopt.torch.models.hf as models_hf import modelopt.torch.utils.distributed as dist -import modelopt.torch.utils.moe as moe -from modelopt.torch.utils import env -from modelopt.torch.utils.distributed import DistributionPlan +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils -from modelopt.torch.utils.plugins.model_load_utils import load_pretrained, resolve_distribution_plan +from modelopt.torch.utils.plugins.model_load_utils import load_pretrained -# Small enough that the dense part fits any GPU, so "auto" dense sharding always resolves to -# replication whether or not the machine running the tests has one. _TINY = { "hidden_size": 16, "intermediate_size": 32, @@ -50,11 +45,6 @@ } -def _qwen3_moe(num_experts): - """Qwen3-MoE, whose base_model_ep_plan ships with transformers itself.""" - return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY) - - @pytest.fixture(autouse=True) def _no_distribution_knobs(monkeypatch): """A shell that exports the knobs would otherwise steer every decision below.""" @@ -66,8 +56,8 @@ def _no_distribution_knobs(monkeypatch): def run_as(monkeypatch): """Run as rank ``rank`` of a ``world_size``-rank launch. - Every rank acts as its group's first, deciding for itself: these tests are about what is - decided, and test_distributed_run.py about how the ranks agree on it. + Every rank acts as its group's first, deciding for itself: these tests are about how a plan + loads, and test_resolve_distribution_plan.py about how the ranks agree on it. """ monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) @@ -144,112 +134,22 @@ def process_group(monkeypatch): monkeypatch.setattr(dist, "local_rank", lambda: 1) -# ----------------------------------------------------------------------------------- env knobs - - -@pytest.mark.parametrize( - ("knob", "name"), - [(env.ep_degree, "MODELOPT_EP_DEGREE"), (env.fsdp_size, "MODELOPT_FSDP_SIZE")], -) -def test_distribution_knobs_default_to_auto(monkeypatch, knob, name): - """Unset and exported-but-empty both mean "let the plan decide".""" - assert knob() == "auto" - monkeypatch.setenv(name, "") - assert knob() == "auto" - - -@pytest.mark.parametrize( - ("knob", "name", "other"), - [ - (env.ep_degree, "MODELOPT_EP_DEGREE", env.fsdp_size), - (env.fsdp_size, "MODELOPT_FSDP_SIZE", env.ep_degree), - ], -) -def test_distribution_knobs_are_read_when_called(monkeypatch, knob, name, other): - """A value snapshotted at import would miss what the launcher sets later; and each knob reads - only its own variable.""" - monkeypatch.setenv(name, "2") - assert knob() == "2" - monkeypatch.setenv(name, "8") - assert knob() == "8" - assert other() == "auto" - - -# -------------------------------------------------------------------------------- the decision - - -@pytest.mark.parametrize("world_size", [None, 1], ids=["no_process_group", "one_rank"]) -def test_outside_a_distributed_run_nothing_is_distributed(run_as, serve, tmp_path, world_size): - """With no ranks to shard over, even a MoE loads whole.""" - if world_size is not None: - run_as(world_size) - serve(_qwen3_moe(num_experts=8)) - - assert resolve_distribution_plan(str(tmp_path)) == (None, None) - - -@pytest.mark.skipif( - Version(transformers.__version__) < Version("5.7"), - reason="qwen3_moe's base_model_ep_plan ships in transformers 5.7", -) -@pytest.mark.parametrize("rank", [0, 1]) -def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, capsys, rank): - """Auto picks the largest EP degree dividing both the world and the expert count, hands back - the config it decided from, and only rank 0 reports the decision.""" - run_as(4, rank) - serve(_qwen3_moe(num_experts=8)) - - plan, hf_config = resolve_distribution_plan(str(tmp_path)) - - assert (plan.world_size, plan.ep_size, plan.edp) == (4, 4, 1) - assert not plan.is_fsdp_dense - assert moe.is_expert_parallel_capable(hf_config) - assert ("Distribution:" in capsys.readouterr().out) == (rank == 0) - - -def test_a_dense_model_is_replicated(run_as, serve, tmp_path): - """A dense model that fits one GPU has nothing to shard: every rank holds it whole.""" - run_as(4) - serve(LlamaConfig(**_TINY)) - - plan, _ = resolve_distribution_plan(str(tmp_path)) - - assert (plan.ep_size, plan.fsdp_size) == (1, 1) - - -def test_the_env_knobs_override_the_automatic_plan(run_as, serve, monkeypatch, tmp_path): - """The decision honours both knobs over its automatic choice.""" - monkeypatch.setenv("MODELOPT_EP_DEGREE", "2") - monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") - run_as(4) - serve(_qwen3_moe(num_experts=8)) - - plan, _ = resolve_distribution_plan(str(tmp_path)) - - assert (plan.ep_size, plan.edp, plan.fsdp_size) == (2, 2, 4) - - -def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path): - monkeypatch.setenv("MODELOPT_EP_DEGREE", "4") - run_as(4) - serve(_qwen3_moe(num_experts=8)) - - plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=2, fsdp_size=1) - - assert (plan.ep_size, plan.fsdp_size) == (2, 1) - - # ----------------------------------------------------------------------------------- the loads -def test_a_single_process_loads_the_model_whole_as_asked(serve, loaders, auto_class, tmp_path): - """No process group, no plan: from_pretrained places the model exactly as the caller says.""" +@pytest.mark.parametrize("kwargs", [{"device_map": "auto", "dtype": "auto"}, {}]) +def test_a_single_process_loads_the_model_whole_as_asked( + serve, loaders, auto_class, tmp_path, kwargs +): + """No process group, a single process's plan: from_pretrained places the model exactly as the + caller says, and with nothing said, as it does by default.""" serve(LlamaConfig(**_TINY)) - model = load_pretrained(str(tmp_path), auto_class, device_map="auto", dtype="auto") + model = load_pretrained(str(tmp_path), auto_class, **kwargs) - assert auto_class.calls == [{"device_map": "auto", "dtype": "auto"}] + assert auto_class.calls == [kwargs] assert model._modelopt_unplaced_source_keys == [] + assert distribution_plan(model) == DistributionPlan() assert loaders == {} @@ -267,26 +167,28 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( serve(LlamaConfig(**_TINY)) kwargs = {} if device_map is None else {"device_map": device_map} - load_pretrained(str(tmp_path), auto_class, **kwargs) + model = load_pretrained(str(tmp_path), auto_class, **kwargs) assert auto_class.calls == [{"device_map": placed}] + assert distribution_plan(model) == DistributionPlan(4) assert loaders == {} -@pytest.mark.parametrize("distribution", ["fsdp2", DistributionPlan(2, fsdp_size=2)]) -def test_fsdp2_loads_through_the_parallel_fsdp2_loader( - run_as, serve, loaders, process_group, tmp_path, distribution +def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( + run_as, serve, loaders, process_group, tmp_path ): - """Named or planned, FSDP2 shards over the process group it reads its rank and GPU from.""" + """FSDP2 shards over the process group it reads its rank and GPU from.""" run_as(2, rank=1) served = serve(LlamaConfig(**_TINY)) + plan = DistributionPlan(2, dense="fsdp") - load_pretrained(str(tmp_path), distribution=distribution, cpu_offload=True, dtype="auto") + model = load_pretrained(str(tmp_path), distribution=plan, cpu_offload=True, dtype="auto") call = loaders.pop("fsdp2") assert (call["device"], call["rank"], call["world_size"]) == (torch.device("cuda:1"), 1, 2) assert call["cpu_offload"] is True assert call["hf_config"] is served + assert distribution_plan(model) == plan assert loaders == {} @@ -301,25 +203,17 @@ def test_a_load_it_decides_reads_the_config_once( transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a) or config ) - load_pretrained(str(tmp_path)) + model = load_pretrained(str(tmp_path)) assert len(reads) == 1 assert loaders["fsdp2"]["hf_config"] is config - - -def test_partial_dense_sharding_without_expert_parallelism_is_refused(run_as, serve, tmp_path): - """The FSDP2 load shards over the whole world; it cannot honour a smaller dense degree.""" - run_as(4) - serve(LlamaConfig(**_TINY)) - - with pytest.raises(NotImplementedError, match="fsdp_size must be 4"): - load_pretrained(str(tmp_path), distribution=DistributionPlan(4, fsdp_size=2)) + assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, tmp_path): """No load here shards the routed experts: an EP plan fails before anything loads.""" run_as(2) - serve(_qwen3_moe(num_experts=8)) + serve(Qwen3MoeConfig(num_experts=8, moe_intermediate_size=8, **_TINY)) with pytest.raises(NotImplementedError, match="shards the routed experts"): load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) @@ -334,7 +228,7 @@ def test_every_rank_loads_the_plan_rank0_was_handed( of its own, loads the one rank 0 was handed.""" run_as(2, rank=1) serve(LlamaConfig(**_TINY)) - rank0_plan = DistributionPlan(2, fsdp_size=2) + rank0_plan = DistributionPlan(2, dense="fsdp") def broadcast_object_list(message, group=None, group_src=None): label = message[0][0] @@ -344,10 +238,11 @@ def broadcast_object_list(message, group=None, group_src=None): monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) monkeypatch.setattr(model_load_utils, "AutoModelForCausalLM", auto_class) - load_pretrained(str(tmp_path), distribution=DistributionPlan(2)) + model = load_pretrained(str(tmp_path), distribution=DistributionPlan(2)) assert auto_class.calls == [], "rank 1 loaded the whole model, as its own plan says" assert "fsdp2" in loaders, "rank 0's plan shards the model" + assert distribution_plan(model) == rank0_plan def test_a_plan_for_another_world_is_refused(run_as, serve, tmp_path): @@ -355,7 +250,7 @@ def test_a_plan_for_another_world_is_refused(run_as, serve, tmp_path): serve(LlamaConfig(**_TINY)) with pytest.raises(ValueError, match="is for 4 ranks, but this run has 2"): - load_pretrained(str(tmp_path), distribution=DistributionPlan(4, fsdp_size=4)) + load_pretrained(str(tmp_path), distribution=DistributionPlan(4, dense="fsdp")) @pytest.mark.parametrize( @@ -381,35 +276,39 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( load_pretrained( str(tmp_path), change.get("auto_class"), - distribution="fsdp2", + distribution=DistributionPlan(2, dense="fsdp"), **change.get("kwargs", {}), ) assert loaders == {} -def test_the_fsdp2_load_refuses_a_quantized_checkpoint(serve, loaders, process_group, tmp_path): +def test_the_fsdp2_load_refuses_a_quantized_checkpoint( + run_as, serve, loaders, process_group, tmp_path +): """It reads the tensors as stored, so quantized weights would load as if they were not.""" + run_as(2) config = LlamaConfig(**_TINY) config.quantization_config = {"quant_method": "fp8"} serve(config) with pytest.raises(NotImplementedError, match="cannot dequantize"): - load_pretrained(str(tmp_path), distribution="fsdp2") + load_pretrained(str(tmp_path), distribution=DistributionPlan(2, dense="fsdp")) assert loaders == {} def test_the_fsdp2_load_needs_a_process_group(serve, loaders, tmp_path): + """FSDP2 on one rank still shards through a process group, of one.""" serve(LlamaConfig(**_TINY)) with pytest.raises(RuntimeError, match=r"needs a torch\.distributed process group"): - load_pretrained(str(tmp_path), distribution="fsdp2") + load_pretrained(str(tmp_path), distribution=DistributionPlan(dense="fsdp")) -def test_an_unknown_distribution_is_refused(serve, tmp_path): +def test_anything_but_a_plan_is_refused(serve, tmp_path): serve(LlamaConfig(**_TINY)) - with pytest.raises(ValueError, match="expected None, 'sequential', 'fsdp2'"): - load_pretrained(str(tmp_path), distribution="ep") + with pytest.raises(TypeError, match="expected a DistributionPlan or None"): + load_pretrained(str(tmp_path), distribution="fsdp2") def test_a_whole_model_load_refuses_a_sharded_loads_config(serve, auto_class, tmp_path): From 7bc64618f366a8fc46eecd2842e2ecb0f35ecad6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 20:26:10 +0000 Subject: [PATCH 07/41] Record the plan where FSDP2 wraps a model, in fsdp2_wrap fsdp2_wrap FSDP2-shards a model over the process group, so it records that plan on the model itself. A caller that wraps a model with it -- the FSDP2 load among them, which no longer records its own -- exports and calibrates it by the plan without a further call. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 5 ++++- .../torch/utils/plugins/model_load_utils.py | 5 ++--- tests/gpu/torch/utils/test_distributed.py | 3 +++ tests/unit/torch/utils/test_distributed_run.py | 17 +++++++++++++++++ 4 files changed, 26 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 7add05e1be7..94b1452f504 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -767,6 +767,9 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False :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. + + The model records that it is FSDP2-sharded over the process group + (:func:`set_distribution_plan`). """ # Lazy import: layerwise_calib imports this module at top level (circular). from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -797,7 +800,7 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False _move_to_fsdp_device(model, ignored_params) if config is not None and architectures: config.architectures = architectures - + set_distribution_plan(model, DistributionPlan(size(), dense="fsdp")) return decoder_layers diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index cb757e606d0..398b11e704c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -940,8 +940,8 @@ def parallel_load_and_prepare_fsdp2( a smaller value lowers the peak transient GPU memory at the cost of more collectives (default 8; pass ``None`` to broadcast all of a source's layers at once). - The model records how it is distributed - (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). + The model records how it is distributed, as :func:`~modelopt.torch.utils.distributed.fsdp2_wrap` + does. """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) weight_map = indexed_weight_map(resolved_path) @@ -1033,7 +1033,6 @@ def parallel_load_and_prepare_fsdp2( _promote_non_dtensor_to_gpu(model, device) if hasattr(model, "tie_weights"): model.tie_weights() - set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) return model diff --git a/tests/gpu/torch/utils/test_distributed.py b/tests/gpu/torch/utils/test_distributed.py index d1581e1ca03..a1470bad932 100644 --- a/tests/gpu/torch/utils/test_distributed.py +++ b/tests/gpu/torch/utils/test_distributed.py @@ -29,8 +29,10 @@ from modelopt.torch.nas.traced_hp import TracedHp from modelopt.torch.utils.distributed import ( + DistributionPlan, allgather, broadcast, + distribution_plan, fsdp2_wrap, gather_bytes, gather_state_dicts, @@ -82,6 +84,7 @@ def _test_fsdp2_wrap_mixed_dtypes(rank, size): } fsdp2_wrap(model) + assert distribution_plan(model) == DistributionPlan(size, dense="fsdp") # Raised "FSDP expects uniform original parameter dtype" before the ignored-param fix. input_ids = torch.randint(0, VOCAB_SIZE, (1, 8), device=device) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index 0d39403af24..ac56fb54e15 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -163,6 +163,23 @@ 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_records_that_the_model_is_sharded_over_the_process_group(monkeypatch): + """Whatever wraps the model, the export and calibration read the plan rather than the wrap.""" + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + model = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)) + monkeypatch.setattr( + LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) + ) + monkeypatch.setattr(dist, "fully_shard", lambda module, **kwargs: None) + monkeypatch.setattr(dist, "size", lambda group=None: 4) + + dist.fsdp2_wrap(model) + + plan = DistributionPlan(4, dense="fsdp") + assert all(distribution_plan(module) == plan for module in model.modules()) + + # ------------------------------------------------------------------------------------ collectives From 1672ee49277630e935b95b1472db1269806b85bf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 20:55:15 +0000 Subject: [PATCH 08/41] Give get_model's tests a module for load_pretrained to record its plan on load_pretrained records the plan on every module of the model it returns, so the stand-in model these tests have from_pretrained return is a torch module now. Signed-off-by: Shengliang Xu --- tests/examples/hf_ptq/test_example_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index b68e9e5bff9..7e294c0ed0b 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -230,7 +230,7 @@ def test_get_model_uses_expected_dtype_kwarg( torch_dtype=torch.bfloat16, ) - class FakeModel: + class FakeModel(torch.nn.Module): # load_pretrained records its plan on every module def eval(self): calls["eval"] = True @@ -315,7 +315,7 @@ def test_get_model_device_map_for_diffusion_gemma( torch_dtype=torch.bfloat16, ) - class FakeModel: + class FakeModel(torch.nn.Module): # load_pretrained records its plan on every module def eval(self): calls["eval"] = True @@ -444,7 +444,7 @@ def _load_config(*args, **kwargs): raise builtin_config_error return hf_config - class FakeModel: + class FakeModel(torch.nn.Module): # load_pretrained records its plan on every module def eval(self): return None From aba6de6b017eb536516227559c17fc3fc2384d21 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 22:58:38 +0000 Subject: [PATCH 09/41] Set up a single process's FSDP2 group in the load that needs it FSDP2 shards through a process group even on one rank, and the FSDP2 load refused to start without one, leaving every caller to set it up first. load_pretrained sets up the group of one itself, after anything that could refuse the load, so a refused load leaves none behind; the run (distributed_run) destroys it with the run. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 10 ++++-- .../test_load_pretrained_distribution.py | 34 +++++++++++++++---- 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 398b11e704c..27ab4a16730 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -463,8 +463,6 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) -> nn.Module ) trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) - if not dist.is_initialized(): - raise RuntimeError("The FSDP2 load needs a torch.distributed process group (torchrun).") if hf_config is None: hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code) if getattr(hf_config, "quantization_config", None) is not None: @@ -472,6 +470,10 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) -> nn.Module raise NotImplementedError( f"The FSDP2 load cannot dequantize this {model_type} checkpoint; load it unsharded." ) + if not dist.is_initialized(): + # FSDP2 shards through a process group even on one rank; a plan for more ranks than one + # was refused without a group. A run (distributed_run) destroys the group on exit. + dist.setup() return parallel_load_and_prepare_fsdp2( ckpt_path, torch.device("cuda", dist.local_rank()), @@ -498,7 +500,9 @@ def load_pretrained( * ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, ``mp_policy`` and - ``broadcast_chunk_size`` ``kwargs`` may set). + ``broadcast_chunk_size`` ``kwargs`` may set). A single process has no group to shard through, + so the load sets up one of one rank, which + :func:`~modelopt.torch.utils.distributed.distributed_run` destroys with the run. * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to ``AutoModelForCausalLM``), placed as ``kwargs`` say. Every rank of a distributed run gets the whole model, ``device_map`` defaulting to the rank's GPU. diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index d55dbfc60c6..94dabafe286 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -296,12 +296,34 @@ def test_the_fsdp2_load_refuses_a_quantized_checkpoint( assert loaders == {} -def test_the_fsdp2_load_needs_a_process_group(serve, loaders, tmp_path): - """FSDP2 on one rank still shards through a process group, of one.""" - serve(LlamaConfig(**_TINY)) - - with pytest.raises(RuntimeError, match=r"needs a torch\.distributed process group"): - load_pretrained(str(tmp_path), distribution=DistributionPlan(dense="fsdp")) +@pytest.mark.parametrize("quantized", [False, True]) +def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( + monkeypatch, serve, loaders, tmp_path, quantized +): + """FSDP2 on one rank still shards through a process group, of one, which the load sets up once + nothing can refuse it: a refused load leaves no group behind.""" + config = serve(LlamaConfig(**_TINY)) + if quantized: + config.quantization_config = {"quant_method": "fp8"} + live = [] + monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) + monkeypatch.setattr(dist, "setup", lambda timeout=None: live.append(0)) + # What a group of one answers once it is up. + monkeypatch.setattr(dist, "rank", lambda group=None: 0) + monkeypatch.setattr(dist, "size", lambda group=None: 1) + monkeypatch.setattr(dist, "local_rank", lambda: 0) + plan = DistributionPlan(dense="fsdp") + + if quantized: + with pytest.raises(NotImplementedError, match="cannot dequantize"): + load_pretrained(str(tmp_path), distribution=plan) + assert live == [] + return + load_pretrained(str(tmp_path), distribution=plan) + + assert live == [0] + call = loaders.pop("fsdp2") + assert (call["rank"], call["world_size"]) == (0, 1) def test_anything_but_a_plan_is_refused(serve, tmp_path): From 5eab3fa439755deae414b9c087de4e7a3ca1d8ab Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:21:50 +0000 Subject: [PATCH 10/41] Move fsdp2_wrap to the model loaders, keeping its old path fsdp2_wrap shards a model as part of loading it, the one load that calls it lives in model_load_utils, and so does every other loader. It moves there; the FSDP2 helpers it uses stay in distributed.py, and modelopt.torch.utils.distributed.fsdp2_wrap forwards to it for callers of the old path. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 50 ++-------------- .../torch/utils/plugins/model_load_utils.py | 58 +++++++++++++++++-- .../unit/torch/utils/test_distributed_run.py | 22 ++++--- .../test_load_pretrained_distribution.py | 20 +++++++ 4 files changed, 90 insertions(+), 60 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 94b1452f504..8f96373f887 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__ = [ @@ -756,52 +756,14 @@ def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None: def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False): - """Auto-detect a HF causal-LM's decoder layers and FSDP2 ``fully_shard`` each one. + """FSDP2-shard a HF causal LM's decoder layers, and the root unless ``shard_root=False``. - By default (``shard_root=True``) the root module is wrapped too, so embed/lm_head/norm are - sharded instead of replicated per rank; pass ``shard_root=False`` to leave the root replicated - (only decoder layers sharded). Returns the detected decoder layers so callers can reuse the - detection result. - - Parameters whose dtype differs from the model's dominant one are excluded from the wrap (see - :func:`_off_dtype_params`), since FSDP2 rejects a shard group that mixes dtypes. They stay - replicated and are moved onto the shards' device, which ``fully_shard`` does not do for the - params it ignores. - - The model records that it is FSDP2-sharded over the process group - (:func:`set_distribution_plan`). + It loads a model, so it lives with the loaders now: + :func:`modelopt.torch.utils.plugins.model_load_utils.fsdp2_wrap`. This keeps the old path. """ - # 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 - set_distribution_plan(model, DistributionPlan(size(), dense="fsdp")) - 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 27ab4a16730..713d54a76ca 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, AutoModel, AutoModelForCausalLM, PretrainedConfig from transformers.conversion_mapping import get_model_conversion_mapping @@ -40,11 +40,13 @@ 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, set_distribution_plan, + size, ) from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -617,6 +619,55 @@ def checked() -> DistributionPlan: 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. + + The model records that it is FSDP2-sharded over the process group + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). + """ + # 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." + ) + + 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 + set_distribution_plan(model, DistributionPlan(size(), dense="fsdp")) + 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. @@ -944,8 +995,7 @@ def parallel_load_and_prepare_fsdp2( a smaller value lowers the peak transient GPU memory at the cost of more collectives (default 8; pass ``None`` to broadcast all of a source's layers at once). - The model records how it is distributed, as :func:`~modelopt.torch.utils.distributed.fsdp2_wrap` - does. + The model records how it is distributed, as :func:`fsdp2_wrap` does. """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) weight_map = indexed_weight_map(resolved_path) diff --git a/tests/unit/torch/utils/test_distributed_run.py b/tests/unit/torch/utils/test_distributed_run.py index ac56fb54e15..4778f9c15fd 100644 --- a/tests/unit/torch/utils/test_distributed_run.py +++ b/tests/unit/torch/utils/test_distributed_run.py @@ -163,21 +163,19 @@ 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_records_that_the_model_is_sharded_over_the_process_group(monkeypatch): - """Whatever wraps the model, the export and calibration read the plan rather than the wrap.""" - from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector - - model = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)) +def test_fsdp2_wrap_keeps_its_old_path(monkeypatch): + """fsdp2_wrap loads a model, so it lives with the loaders; callers of the old path still reach it.""" + model_load_utils = pytest.importorskip("modelopt.torch.utils.plugins.model_load_utils") + calls = [] monkeypatch.setattr( - LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) + model_load_utils, + "fsdp2_wrap", + lambda model, **kwargs: calls.append((model, kwargs)) or "layers", ) - monkeypatch.setattr(dist, "fully_shard", lambda module, **kwargs: None) - monkeypatch.setattr(dist, "size", lambda group=None: 4) - - dist.fsdp2_wrap(model) + model = nn.Linear(2, 2) - plan = DistributionPlan(4, dense="fsdp") - assert all(distribution_plan(module) == plan for module in model.modules()) + 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 diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 94dabafe286..2e775005768 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -339,3 +339,23 @@ def test_a_whole_model_load_refuses_a_sharded_loads_config(serve, auto_class, tm with pytest.raises(TypeError, match="hf_config applies to a sharded load only"): load_pretrained(str(tmp_path), auto_class, hf_config=config) + + +# --------------------------------------------------------------------------------- fsdp2_wrap + + +def test_fsdp2_wrap_records_that_the_model_is_sharded_over_the_process_group(monkeypatch): + """Whatever wraps the model, the export and calibration read the plan rather than the wrap.""" + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + model = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + monkeypatch.setattr( + LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) + ) + monkeypatch.setattr(model_load_utils, "fully_shard", lambda module, **kwargs: None) + monkeypatch.setattr(model_load_utils, "size", lambda group=None: 4) + + model_load_utils.fsdp2_wrap(model) + + plan = DistributionPlan(4, dense="fsdp") + assert all(distribution_plan(module) == plan for module in model.modules()) From fb57614779638b548d3bcb2ba7a4e55abf1881b2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:43:44 +0000 Subject: [PATCH 11/41] Record a load's plan in load_pretrained alone fsdp2_wrap recorded a plan it derived from the process group, and load_pretrained then recorded the one it loaded by over it: two writers, one silently replacing the other. load_pretrained is the one that knows the plan, so it alone records it, once, at the end of every load. fsdp2_wrap records nothing, as on main; the tests that wrap or load a model through it directly record the plan themselves, as any caller distributing a model itself must. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 10 +++++----- tests/gpu/torch/export/test_export_carry_over.py | 10 +++++++++- tests/gpu/torch/utils/test_distributed.py | 3 --- tests/gpu/torch/utils/test_model_load_utils.py | 8 +++++++- .../torch/utils/test_load_pretrained_distribution.py | 11 ++++++----- 5 files changed, 27 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 713d54a76ca..8a8549ad5a8 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -46,7 +46,6 @@ broadcast_state_dict, is_initialized, set_distribution_plan, - size, ) from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -632,8 +631,10 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False that mixes dtypes. They stay replicated and are moved onto the shards' device, which ``fully_shard`` does not do for the params it ignores. - The model records that it is FSDP2-sharded over the process group - (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). + It records no distribution plan: :func:`load_pretrained` records the plan it loads by, and a + caller that wraps a model itself records one + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) before calibrating or + exporting it. """ # Lazy import: layerwise_calib imports this module at top level (circular). from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -664,7 +665,6 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False _move_to_fsdp_device(model, ignored_params) if config is not None and architectures: config.architectures = architectures - set_distribution_plan(model, DistributionPlan(size(), dense="fsdp")) return decoder_layers @@ -995,7 +995,7 @@ def parallel_load_and_prepare_fsdp2( a smaller value lowers the peak transient GPU memory at the cost of more collectives (default 8; pass ``None`` to broadcast all of a source's layers at once). - The model records how it is distributed, as :func:`fsdp2_wrap` does. + It records no distribution plan, as :func:`fsdp2_wrap` does not; :func:`load_pretrained` does. """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) weight_map = indexed_weight_map(resolved_path) diff --git a/tests/gpu/torch/export/test_export_carry_over.py b/tests/gpu/torch/export/test_export_carry_over.py index fa528f4cdf9..34a54d23f5e 100644 --- a/tests/gpu/torch/export/test_export_carry_over.py +++ b/tests/gpu/torch/export/test_export_carry_over.py @@ -40,7 +40,12 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.unified_export_hf import export_hf_checkpoint from modelopt.torch.quantization.utils import patch_fsdp_mp_dtypes -from modelopt.torch.utils.distributed import fsdp2_wrap, is_fsdp2_model +from modelopt.torch.utils.distributed import ( + DistributionPlan, + fsdp2_wrap, + is_fsdp2_model, + set_distribution_plan, +) # Small enough to force the tiny model across several shards, so the multi-file layout -- and the # weight index that makes it loadable -- are exercised rather than collapsing to one model.safetensors. @@ -108,6 +113,8 @@ def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwar model.eval() fsdp2_wrap(model) + # Wrapped by hand, the model records how it is distributed, as load_pretrained would. + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) # The FSDP2 export path is selected by exactly this predicate, so assert it here: without it # the export would silently fall back to the rank-0 gather and the test would pass while # never touching the code under test. @@ -201,6 +208,7 @@ def _ptq_language_model_and_export(rank, size, *, src_dir, export_dir): ) model.eval() fsdp2_wrap(model) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) assert is_fsdp2_model(model), "fsdp2_wrap did not shard the VLM" torch.distributed.barrier() diff --git a/tests/gpu/torch/utils/test_distributed.py b/tests/gpu/torch/utils/test_distributed.py index a1470bad932..d1581e1ca03 100644 --- a/tests/gpu/torch/utils/test_distributed.py +++ b/tests/gpu/torch/utils/test_distributed.py @@ -29,10 +29,8 @@ from modelopt.torch.nas.traced_hp import TracedHp from modelopt.torch.utils.distributed import ( - DistributionPlan, allgather, broadcast, - distribution_plan, fsdp2_wrap, gather_bytes, gather_state_dicts, @@ -84,7 +82,6 @@ def _test_fsdp2_wrap_mixed_dtypes(rank, size): } fsdp2_wrap(model) - assert distribution_plan(model) == DistributionPlan(size, dense="fsdp") # Raised "FSDP expects uniform original parameter dtype" before the ignored-param fix. input_ids = torch.randint(0, VOCAB_SIZE, (1, 8), device=device) diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index 3296ccbfc86..05b524fb19a 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -26,7 +26,11 @@ from torch.distributed.tensor import DTensor from modelopt.torch.export.unified_export_hf import export_hf_checkpoint -from modelopt.torch.utils.distributed import broadcast_state_dict +from modelopt.torch.utils.distributed import ( + DistributionPlan, + broadcast_state_dict, + set_distribution_plan, +) from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 VOCAB_SIZE = 64 @@ -61,6 +65,7 @@ def _test_parallel_load_and_export(rank, size, ckpt_dir, export_dir, cpu_offload size, cpu_offload=cpu_offload, ) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) # Decoder layers AND root params (embed/lm_head) are sharded DTensors under shard_root=True. decoder_params = list(model.model.layers[0].parameters()) @@ -106,6 +111,7 @@ def _test_carry_over_unplaced_weights(rank, size, ckpt_dir, export_dir, orphan_p """Load, export, and require the weights the model could not place to survive.""" device = torch.device(f"cuda:{rank}") model = parallel_load_and_prepare_fsdp2(ckpt_dir, device, rank, size) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) unplaced = getattr(model, "_modelopt_unplaced_source_keys", None) assert unplaced, ( diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 2e775005768..5a3be45bdc8 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -344,8 +344,9 @@ def test_a_whole_model_load_refuses_a_sharded_loads_config(serve, auto_class, tm # --------------------------------------------------------------------------------- fsdp2_wrap -def test_fsdp2_wrap_records_that_the_model_is_sharded_over_the_process_group(monkeypatch): - """Whatever wraps the model, the export and calibration read the plan rather than the wrap.""" +def test_fsdp2_wrap_leaves_the_plan_to_the_load(monkeypatch): + """load_pretrained records the plan once, at the end of the load; the wrap it calls records none, + so a caller that wraps a model itself must record one before anything reads it.""" from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector model = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) @@ -353,9 +354,9 @@ def test_fsdp2_wrap_records_that_the_model_is_sharded_over_the_process_group(mon LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) ) monkeypatch.setattr(model_load_utils, "fully_shard", lambda module, **kwargs: None) - monkeypatch.setattr(model_load_utils, "size", lambda group=None: 4) + monkeypatch.setattr(dist, "size", lambda group=None: 4) model_load_utils.fsdp2_wrap(model) - plan = DistributionPlan(4, dense="fsdp") - assert all(distribution_plan(module) == plan for module in model.modules()) + with pytest.raises(RuntimeError, match="no distribution plan recorded"): + distribution_plan(model) From db6f15ef8f8d073efe1786196b64208130a42d91 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 17:30:20 +0000 Subject: [PATCH 12/41] File the distributed load_pretrained entry under 0.49.0 in the changelog main opened 0.49.0 for what merges after 0.48.0's cut, which this PR does. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 95d0972cf22..838b58cfdfb 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -35,6 +35,7 @@ Changelog *Misc* - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_FSDP_SIZE`` variable, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). **Backward Breaking Changes** @@ -148,7 +149,6 @@ Changelog - Add a ``kd_loss_alpha`` option to HuggingFace ``KDTrainer`` to blend cross-entropy and knowledge distillation losses during training. - A tracked ``examples/hf_ptq/hf_ptq.py`` run now writes ``.experiment.json`` into ``--export_path`` and uploads the same file with the run, so a checkpoint on disk names the experiment and MLflow run id that produced it. The pointer is written only once the export completes, and an export that is not tracked removes one it would otherwise inherit from a reused ``--export_path`` or from a quantized source checkpoint. - ``export_hf_checkpoint`` and the vLLM fake-quant export now copy the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged. They read the source from local disk only: get a local copy of a Hugging Face Hub model with ``modelopt.torch.export.ensure_local_checkpoint`` and load it from there. -- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_FSDP_SIZE`` variable, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). - Add Parallel Decoding Distillation (PDD) to ``modelopt.torch.fastgen`` with Qwen-Image training, distributed-checkpoint export, and PDD-2/4/8 inference. AutoModel remains an unmodified pinned runtime dependency. From ead7dc3735615ae7445ac3489c15f9832bb61899 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 02:04:28 +0000 Subject: [PATCH 13/41] Answer the checkpoint-read gather in the load-routing tests' launches resolve_distribution_plan reads the checkpoint through the failure-carrying all_reduce, which gathers from every rank of a launch. The run_as fixture answers that gather as it answers the plan's broadcast: every rank read the checkpoint as this one did. Signed-off-by: Shengliang Xu --- .../torch/utils/test_load_pretrained_distribution.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 5a3be45bdc8..aa98fd1638f 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -52,19 +52,27 @@ def _no_distribution_knobs(monkeypatch): monkeypatch.delenv("MODELOPT_FSDP_SIZE", raising=False) +def _gathered_alike(gathered, mine, group=None): + """Answer a gather as if every rank handed it what this rank did.""" + gathered[:] = [mine] * len(gathered) + + @pytest.fixture def run_as(monkeypatch): """Run as rank ``rank`` of a ``world_size``-rank launch. - Every rank acts as its group's first, deciding for itself: these tests are about how a plan - loads, and test_resolve_distribution_plan.py about how the ranks agree on it. + Every rank acts as its group's first, deciding for itself, and reads the checkpoint as this one + does: these tests are about how a plan loads, and test_resolve_distribution_plan.py about how + the ranks agree on it. """ monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) + monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike) def enter(world_size, rank=0): monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1) monkeypatch.setattr(dist, "size", lambda group=None: world_size) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size) monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) monkeypatch.setattr(dist, "local_rank", lambda: rank) From 627b8799f78dae54c341da63c3997a2ab72d406c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:14:15 +0000 Subject: [PATCH 14/41] Test on two real ranks that every rank loads the plan rank 0 was handed The test played rank 1 against a stand-in broadcast that answered with rank 0's plan, and checked which stub loader ran. It now runs on the session's two gloo workers: rank 0 is handed a replicated plan and rank 1 an FSDP2 one, both load a real checkpoint onto the CPU, and both must load it whole and record rank 0's plan. A rank that loaded its own plan would refuse the placement an FSDP2 load cannot honor, so a regression fails at once instead of waiting in a collective. Signed-off-by: Shengliang Xu --- .../test_load_pretrained_distribution.py | 40 ++++++++----------- 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index aa98fd1638f..4e67dab0e5b 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -15,19 +15,21 @@ """Which load ``load_pretrained`` routes a distribution plan to, and the plan it records. -Nothing here needs a process group or a GPU: the process group's rank and size are answered the way -a torchrun launch would, the configs are tiny real transformers configs served from memory, and the -loaders are replaced by stand-ins that record how they were called. +Nothing here needs a GPU, and only one test a process group, two gloo ranks: elsewhere the process +group's rank and size are answered the way a torchrun launch would, the configs are tiny real +transformers configs served from memory, and the loaders are replaced by stand-ins that record how +they were called. """ import inspect import pytest import torch +from _test_utils.torch.distributed.utils import bounded_collectives transformers = pytest.importorskip("transformers") -from transformers import AutoModel, LlamaConfig, Qwen3MoeConfig +from transformers import AutoModel, LlamaConfig, LlamaForCausalLM, Qwen3MoeConfig import modelopt.torch.models.hf as models_hf import modelopt.torch.utils.distributed as dist @@ -229,28 +231,18 @@ def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, assert auto_class.calls == [] -def test_every_rank_loads_the_plan_rank0_was_handed( - run_as, serve, loaders, process_group, auto_class, monkeypatch, tmp_path -): - """Ranks loading different plans would build mismatched meshes and hang: rank 1, handed a plan - of its own, loads the one rank 0 was handed.""" - run_as(2, rank=1) - serve(LlamaConfig(**_TINY)) - rank0_plan = DistributionPlan(2, dense="fsdp") - - def broadcast_object_list(message, group=None, group_src=None): - label = message[0][0] - message[0] = (label, rank0_plan if label == "plan" else [], None) - - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 1) - monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) - monkeypatch.setattr(model_load_utils, "AutoModelForCausalLM", auto_class) +def _every_rank_loads_rank0s_plan_job(rank, size, ckpt): + handed = DistributionPlan(size) if rank == 0 else DistributionPlan(size, dense="fsdp") + with bounded_collectives(): + model = load_pretrained(ckpt, distribution=handed, device_map="cpu") + assert distribution_plan(model) == DistributionPlan(size), rank - model = load_pretrained(str(tmp_path), distribution=DistributionPlan(2)) - assert auto_class.calls == [], "rank 1 loaded the whole model, as its own plan says" - assert "fsdp2" in loaders, "rank 0's plan shards the model" - assert distribution_plan(model) == rank0_plan +def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path): + """Over a real broadcast. Ranks loading different plans would build mismatched meshes and hang: + rank 1, handed an FSDP2 plan of its own, loads the whole model by the plan rank 0 was handed.""" + LlamaForCausalLM(LlamaConfig(**_TINY)).save_pretrained(tmp_path) + cpu_dist_workers.run(_every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path)) def test_a_plan_for_another_world_is_refused(run_as, serve, tmp_path): From 3fd4e6deb1947f70d3b4aacd309e8103ad566cc0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 18:01:12 +0000 Subject: [PATCH 15/41] Trim the load routing's tests Fold near-duplicates into one test per behavior: - test_a_load_refuses_what_it_cannot_load takes the another-world, quantized-checkpoint, not-a-plan and hf-config-on-a-whole-load refusals; each case still fails if its refusal goes. - test_a_single_process_loads_the_model_whole_as_asked goes, and its plan check moves to test_unplaced_source_keys.py's real single-process load. That load fails if a single process is placed on the rank GPU, loses the caller's kwargs, records another plan or reaches the FSDP2 loader. Signed-off-by: Shengliang Xu --- .../test_load_pretrained_distribution.py | 91 +++++++++---------- .../torch/utils/test_unplaced_source_keys.py | 2 + 2 files changed, 46 insertions(+), 47 deletions(-) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 4e67dab0e5b..14adfa9b9dc 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -147,22 +147,6 @@ def process_group(monkeypatch): # ----------------------------------------------------------------------------------- the loads -@pytest.mark.parametrize("kwargs", [{"device_map": "auto", "dtype": "auto"}, {}]) -def test_a_single_process_loads_the_model_whole_as_asked( - serve, loaders, auto_class, tmp_path, kwargs -): - """No process group, a single process's plan: from_pretrained places the model exactly as the - caller says, and with nothing said, as it does by default.""" - serve(LlamaConfig(**_TINY)) - - model = load_pretrained(str(tmp_path), auto_class, **kwargs) - - assert auto_class.calls == [kwargs] - assert model._modelopt_unplaced_source_keys == [] - assert distribution_plan(model) == DistributionPlan() - assert loaders == {} - - @pytest.mark.parametrize( ("device_map", "placed"), [(None, {"": "cuda:1"}), ("cpu", "cpu")], @@ -245,14 +229,6 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path): cpu_dist_workers.run(_every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path)) -def test_a_plan_for_another_world_is_refused(run_as, serve, tmp_path): - run_as(2) - serve(LlamaConfig(**_TINY)) - - with pytest.raises(ValueError, match="is for 4 ranks, but this run has 2"): - load_pretrained(str(tmp_path), distribution=DistributionPlan(4, dense="fsdp")) - - @pytest.mark.parametrize( ("change", "error", "match"), [ @@ -282,18 +258,54 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( assert loaders == {} -def test_the_fsdp2_load_refuses_a_quantized_checkpoint( - run_as, serve, loaders, process_group, tmp_path +@pytest.mark.parametrize( + ("distribution", "quantized", "whole_model_kwargs", "error", "match"), + [ + ( + DistributionPlan(4, dense="fsdp"), + False, + False, + ValueError, + "is for 4 ranks, but this run has 2", + ), + (DistributionPlan(2, dense="fsdp"), True, False, NotImplementedError, "cannot dequantize"), + ("fsdp2", False, False, TypeError, "expected a DistributionPlan or None"), + ( + DistributionPlan(2), + False, + True, + TypeError, + "hf_config applies to a sharded load only", + ), + ], + ids=["another-world", "quantized-checkpoint", "not-a-plan", "hf-config-on-a-whole-load"], +) +def test_a_load_refuses_what_it_cannot_load( + run_as, + serve, + loaders, + auto_class, + process_group, + tmp_path, + distribution, + quantized, + whole_model_kwargs, + error, + match, ): - """It reads the tensors as stored, so quantized weights would load as if they were not.""" + """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would + read as stored; anything but a plan; and the sharded loads' config, which from_pretrained would + take for a model argument and fail far from here.""" run_as(2) - config = LlamaConfig(**_TINY) - config.quantization_config = {"quant_method": "fp8"} - serve(config) + config = serve(LlamaConfig(**_TINY)) + if quantized: + config.quantization_config = {"quant_method": "fp8"} + whole = {"auto_class": auto_class, "hf_config": config} if whole_model_kwargs else {} - with pytest.raises(NotImplementedError, match="cannot dequantize"): - load_pretrained(str(tmp_path), distribution=DistributionPlan(2, dense="fsdp")) + with pytest.raises(error, match=match): + load_pretrained(str(tmp_path), distribution=distribution, **whole) assert loaders == {} + assert auto_class.calls == [] @pytest.mark.parametrize("quantized", [False, True]) @@ -326,21 +338,6 @@ def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( assert (call["rank"], call["world_size"]) == (0, 1) -def test_anything_but_a_plan_is_refused(serve, tmp_path): - serve(LlamaConfig(**_TINY)) - - with pytest.raises(TypeError, match="expected a DistributionPlan or None"): - load_pretrained(str(tmp_path), distribution="fsdp2") - - -def test_a_whole_model_load_refuses_a_sharded_loads_config(serve, auto_class, tmp_path): - """from_pretrained would take hf_config for a model argument and fail far from here.""" - config = serve(LlamaConfig(**_TINY)) - - with pytest.raises(TypeError, match="hf_config applies to a sharded load only"): - load_pretrained(str(tmp_path), auto_class, hf_config=config) - - # --------------------------------------------------------------------------------- fsdp2_wrap diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 31920ea6799..8632c924772 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -41,6 +41,7 @@ Mxfp4Config, ) +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( load_pretrained, @@ -105,6 +106,7 @@ def test_records_the_stray_tensors_and_the_checkpoint(tmp_path, extra, subfolder assert model._modelopt_unplaced_source_keys == sorted(extra), "sorted, for a stable export" assert model._modelopt_source_checkpoint == str(checkpoint) assert model.dtype == torch.bfloat16, "from_pretrained kwargs must pass through" + assert distribution_plan(model) == DistributionPlan(), "a single process holds it whole" class _Recorder: From 4edb4d69129ed0d8a761ae6bf20921b61dc0a73f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:52:50 +0000 Subject: [PATCH 16/41] Import LayerActivationCollector at the load-routing tests' module scope test_fsdp2_wrap_leaves_the_plan_to_the_load imported it inside the test, though it is neither optional nor circular. At module scope it sits after the file's importorskip("transformers"), so the transformers-free lane still skips the module first. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_load_pretrained_distribution.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 14adfa9b9dc..007d4c18a11 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -33,6 +33,7 @@ import modelopt.torch.models.hf as models_hf import modelopt.torch.utils.distributed as dist +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import load_pretrained @@ -344,8 +345,6 @@ def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( def test_fsdp2_wrap_leaves_the_plan_to_the_load(monkeypatch): """load_pretrained records the plan once, at the end of the load; the wrap it calls records none, so a caller that wraps a model itself must record one before anything reads it.""" - from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector - model = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) monkeypatch.setattr( LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) From 088f445536d11041b69fc1e147cf9c25d62d0c73 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 00:44:17 +0000 Subject: [PATCH 17/41] Pass the Hub options on to the plan's checkpoint read load_pretrained now resolves the distribution plan before it loads, and resolve_distribution_plan resolved the checkpoint and read its config without token, cache_dir or local_files_only. That read reaches the Hub for custom code an auto_map names in another repo, as the loads after it do. resolve_distribution_plan takes the three options, and load_pretrained passes on the ones it was given. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 8a8549ad5a8..8981d399b58 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -563,7 +563,9 @@ def load_pretrained( trust_remote_code = kwargs.get("trust_remote_code", False) hf_config = kwargs.get("hf_config") if distribution is None: - plan, hf_config = resolve_distribution_plan(ckpt_path, trust_remote_code=trust_remote_code) + plan, hf_config = resolve_distribution_plan( + ckpt_path, trust_remote_code=trust_remote_code, **hub_options + ) if _route(plan) is not None: kwargs.setdefault("hf_config", hf_config) elif isinstance(distribution, DistributionPlan): From 37e8f90b181b82427c43926f02b4c809ac3f2c45 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 01:09:31 +0000 Subject: [PATCH 18/41] Test that the plan's checkpoint read takes the Hub options The load_pretrained test cannot see whether the plan's read resolves the checkpoint with token, cache_dir and local_files_only: load_pretrained has made the checkpoint local by then. Drive resolve_distribution_plan with a Hub ID and check that the resolution and the config read both get them. (Test from the epdp port of this change.) Signed-off-by: Shengliang Xu --- .../utils/test_resolve_distribution_plan.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 14c60b7b359..f3c9b66bcca 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -358,6 +358,30 @@ def test_a_single_process_holds_the_whole_model( ) +def test_the_hub_options_reach_the_plans_checkpoint_read(monkeypatch, tmp_path): + """The plan's read resolves a Hub ID and loads its config with the options it was given: the + config reaches the Hub for the custom code its ``auto_map`` names in another repo.""" + options = {"token": "secret", "cache_dir": str(tmp_path), "local_files_only": True} + calls = [] + monkeypatch.setattr( + model_load_utils, + "ensure_local_checkpoint", + lambda path, **kwargs: calls.append(("resolve", path, kwargs)) or (path, str(tmp_path)), + ) + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + lambda path, **kwargs: calls.append(("config", path, kwargs)) or _qwen3_moe(), + ) + + resolve_distribution_plan("org/model", **options) + + assert calls == [ + ("resolve", "org/model", options), + ("config", str(tmp_path), {"trust_remote_code": False, **options}), + ] + + @needs_qwen3_moe_ep_plan @pytest.mark.parametrize("rank", [0, 1]) def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, capsys, rank): From 27239b2fd59640f4dbb5c2101d927ee6fed1850c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 01:43:29 +0000 Subject: [PATCH 19/41] Read a sharded load's config with the Hub options too load_pretrained keeps token, cache_dir and local_files_only out of the keywords it routes, and a sharded load given a plan but no config read the config without them. That read reaches the Hub for custom code an auto_map names in another repo, as the whole-model load's do. _load_sharded now takes the Hub options for its config read. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 8 +++++--- .../utils/test_load_pretrained_distribution.py | 15 +++++++++++---- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 8981d399b58..d3075fa3d2c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -436,7 +436,7 @@ def _route(plan: DistributionPlan) -> str | None: return _FSDP2 if plan.is_fsdp_dense else None -def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) -> nn.Module: +def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) -> nn.Module: """The sharded load ``route`` names, refusing what it cannot honor before it starts.""" from modelopt.torch.models.hf import checkpoint_has_mtp from modelopt.torch.utils import distributed as dist @@ -465,7 +465,9 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) -> nn.Module trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) if hf_config is None: - hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code) + hf_config = AutoConfig.from_pretrained( + ckpt_path, trust_remote_code=trust_remote_code, **hub_options + ) if getattr(hf_config, "quantization_config", None) is not None: # It reads the tensors as stored: quantized weights would load as if they were not. raise NotImplementedError( @@ -591,7 +593,7 @@ def checked() -> DistributionPlan: ).model_type ) if route is not None: - model = _load_sharded(route, ckpt_path, auto_class, model_type, kwargs) + model = _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 007d4c18a11..60680c59f6d 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -170,19 +170,26 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( - run_as, serve, loaders, process_group, tmp_path + run_as, loaders, process_group, monkeypatch, tmp_path ): - """FSDP2 shards over the process group it reads its rank and GPU from.""" + """FSDP2 shards over the process group it reads its rank and GPU from, by a config read with + the Hub options the load was given (custom code an ``auto_map`` names in another repo).""" run_as(2, rank=1) - served = serve(LlamaConfig(**_TINY)) + served, reads = LlamaConfig(**_TINY), [] + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(k) or served + ) plan = DistributionPlan(2, dense="fsdp") - model = load_pretrained(str(tmp_path), distribution=plan, cpu_offload=True, dtype="auto") + model = load_pretrained( + str(tmp_path), distribution=plan, cpu_offload=True, dtype="auto", local_files_only=True + ) call = loaders.pop("fsdp2") assert (call["device"], call["rank"], call["world_size"]) == (torch.device("cuda:1"), 1, 2) assert call["cpu_offload"] is True assert call["hf_config"] is served + assert reads and all(read.get("local_files_only") is True for read in reads) assert distribution_plan(model) == plan assert loaders == {} From 3e09c8ccaeeacfe5d2d68b491643ae21c6d829fe Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 00:45:23 +0000 Subject: [PATCH 20/41] Decide the plan from a config passed for the load load_pretrained honored a config passed for the load only in the whole-model from_pretrained: the plan was still decided from the checkpoint's config.json, so a directory without one failed before loading, and a sharded load refused config as an option it does not take. resolve_distribution_plan takes the caller's config and decides from a copy of it (deciding sets fields on it) without reading the checkpoint's; load_pretrained passes it on, and a sharded load builds from it as its hf_config. Passing both config and hf_config is refused. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 28 ++++++++++++---- .../test_load_pretrained_distribution.py | 32 +++++++++++++++++-- 2 files changed, 51 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d3075fa3d2c..215ddfdbd74 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -15,6 +15,7 @@ """HuggingFace model loading: :func:`load_pretrained`, by a plan to distribute the model.""" +import copy import logging import os import re @@ -517,7 +518,8 @@ def load_pretrained( A sharded load builds the model class and places it itself. Besides its own options it takes ``trust_remote_code``, ``attn_implementation`` and ``hf_config`` (the config to load with), and - refuses any other keyword rather than ignore it. + refuses any other keyword rather than ignore it. A ``config`` passed for the load is the one the + plan is decided from, and a sharded load takes it as its ``hf_config``. Whatever the placement: @@ -563,10 +565,14 @@ def load_pretrained( if subfolder: ckpt_path = os.path.join(ckpt_path, subfolder) trust_remote_code = kwargs.get("trust_remote_code", False) + config = kwargs.get("config") + config = config if isinstance(config, PretrainedConfig) else None + if config is not None and "hf_config" in kwargs: + raise TypeError("config and hf_config both name the config to load with; pass one.") hf_config = kwargs.get("hf_config") if distribution is None: plan, hf_config = resolve_distribution_plan( - ckpt_path, trust_remote_code=trust_remote_code, **hub_options + ckpt_path, config=config, trust_remote_code=trust_remote_code, **hub_options ) if _route(plan) is not None: kwargs.setdefault("hf_config", hf_config) @@ -584,10 +590,13 @@ def checked() -> DistributionPlan: else: raise TypeError(f"distribution={distribution!r}: expected a DistributionPlan or None.") route = _route(plan) + if route is not None and config is not None: + # The sharded load builds the model from hf_config: the plan's copy of this one, or it. + del kwargs["config"] + kwargs.setdefault("hf_config", config) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. - config = hf_config or kwargs.get("config") - model_type = getattr(config, "model_type", None) or ( + model_type = getattr(hf_config or config, "model_type", None) or ( AutoConfig.from_pretrained( ckpt_path, trust_remote_code=trust_remote_code, **hub_options ).model_type @@ -1257,6 +1266,7 @@ def _decide_distribution( def resolve_distribution_plan( ckpt_path: str, *, + config: PretrainedConfig | None = None, trust_remote_code: bool = False, ep_size: int | str | None = None, dense: str | None = None, @@ -1296,6 +1306,9 @@ def resolve_distribution_plan( checkpoint as :func:`load_pretrained` takes them, so the plan is decided for what the load reads: all but ``subfolder`` resolve a Hub model ID (:func:`ensure_local_checkpoint`), and the config read takes all but ``revision``, for custom code its ``auto_map`` names in another Hub repo. + + ``config``, the config the load will use when the caller has one, is decided from instead of + the checkpoint's, which is then not read: each rank takes a copy, as deciding may add to it. """ from modelopt.torch.utils import distributed as dist @@ -1313,14 +1326,17 @@ def checkpoint_read() -> bool: ) if subfolder: local_path = os.path.join(local_path, subfolder) - config = AutoConfig.from_pretrained( + if config is not None: + read.append((local_path, copy.deepcopy(config))) + return True + hf_config = AutoConfig.from_pretrained( local_path, trust_remote_code=trust_remote_code, token=token, cache_dir=cache_dir, local_files_only=local_files_only, ) - read.append((local_path, config)) + read.append((local_path, hf_config)) return True dist.all_reduce("checkpoint read", checkpoint_read, all) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 60680c59f6d..0d17fdf7248 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -212,6 +212,26 @@ def test_a_load_it_decides_reads_the_config_once( assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") +def test_a_config_passed_for_the_load_is_the_one_it_decides_from( + run_as, loaders, process_group, monkeypatch, tmp_path +): + """No config is read: the plan is decided from the one passed for the load, which the FSDP2 + load then builds from. Deciding adds to that config, so it works on a copy of it.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(2, rank=1) + reads: list = [] + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a)) + config = LlamaConfig(**{**_TINY, "num_hidden_layers": 2}) + as_passed = config.to_dict() + + model = load_pretrained(str(tmp_path), config=config) + + assert reads == [] + assert loaders["fsdp2"]["hf_config"].num_hidden_layers == 2 + assert config.to_dict() == as_passed + assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") + + def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, tmp_path): """No load here shards the routed experts: an EP plan fails before anything loads.""" run_as(2) @@ -243,14 +263,20 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path): ({"kwargs": {"device_map": "auto"}}, TypeError, r"\['device_map'\] do not apply"), ({"auto_class": AutoModel}, NotImplementedError, "builds its model class itself"), ({"mtp": True}, NotImplementedError, "cannot build the MTP head"), + ( + {"kwargs": {"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}}, + TypeError, + "config and hf_config both name", + ), ], - ids=["placement", "model_class", "mtp_head"], + ids=["placement", "model_class", "mtp_head", "two_configs"], ) def test_a_sharded_load_refuses_what_it_cannot_honor( run_as, serve, loaders, process_group, monkeypatch, tmp_path, change, error, match ): - """Placement, a model class, a head only a whole model builds: none reaches a sharded load, - which fails before it starts rather than load something other than what was asked.""" + """Placement, a model class, a head only a whole model builds, two configs to load with: none + reaches a sharded load, which fails before it starts rather than load something other than + what was asked.""" run_as(2) serve(LlamaConfig(**_TINY)) if change.get("mtp"): From 3641c47686e68dd141dc8bb548d1ff29cb9337b6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 01:11:18 +0000 Subject: [PATCH 21/41] Test a passed config with a plan handed to the load too With a plan handed in, nothing is decided, so the model type comes from the passed config itself rather than from the plan's copy of it. The test only covered a plan load_pretrained decides, and a load that ignored the passed config there (model type from hf_config alone) still passed it. Signed-off-by: Shengliang Xu --- .../torch/utils/test_load_pretrained_distribution.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 0d17fdf7248..f68a1e5e096 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -212,19 +212,22 @@ def test_a_load_it_decides_reads_the_config_once( assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") +@pytest.mark.parametrize("handed", [False, True], ids=["decided", "handed"]) def test_a_config_passed_for_the_load_is_the_one_it_decides_from( - run_as, loaders, process_group, monkeypatch, tmp_path + run_as, loaders, process_group, monkeypatch, tmp_path, handed ): """No config is read: the plan is decided from the one passed for the load, which the FSDP2 - load then builds from. Deciding adds to that config, so it works on a copy of it.""" + load then builds from, as it does with a plan handed in. Deciding adds to that config, so it + works on a copy of it.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") run_as(2, rank=1) reads: list = [] monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a)) config = LlamaConfig(**{**_TINY, "num_hidden_layers": 2}) as_passed = config.to_dict() + plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} - model = load_pretrained(str(tmp_path), config=config) + model = load_pretrained(str(tmp_path), config=config, **plan) assert reads == [] assert loaders["fsdp2"]["hf_config"].num_hidden_layers == 2 From 99629af06774f5db7616b9112c43eb2ed089f3cc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:09:59 +0000 Subject: [PATCH 22/41] Refuse hf_config without the plan it was resolved with A sharded load builds from hf_config as given. Without a plan handed in, load_pretrained decided one from the checkpoint's own config and then loaded the caller's hf_config, so the model loaded was not the one the plan was decided for, nor one the model spec's checks had seen. hf_config is the config resolve_distribution_plan returns with its plan; without that plan it is now refused, pointing at config= for a config the plan should be decided from. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 15 ++++++++++++--- .../utils/test_load_pretrained_distribution.py | 14 ++++++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 215ddfdbd74..07ccf96a7e7 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -517,9 +517,11 @@ def load_pretrained( read it back. A sharded load builds the model class and places it itself. Besides its own options it takes - ``trust_remote_code``, ``attn_implementation`` and ``hf_config`` (the config to load with), and - refuses any other keyword rather than ignore it. A ``config`` passed for the load is the one the - plan is decided from, and a sharded load takes it as its ``hf_config``. + ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other keyword + rather than ignore it. ``hf_config`` is the config :func:`resolve_distribution_plan` returned with + the plan handed in as ``distribution``: the load builds from it as given, so without that plan + it is refused. A ``config`` passed for the load is the one the plan is decided from, and a + sharded load takes it as its ``hf_config``. Whatever the placement: @@ -569,6 +571,13 @@ def load_pretrained( config = config if isinstance(config, PretrainedConfig) else None if config is not None and "hf_config" in kwargs: raise TypeError("config and hf_config both name the config to load with; pass one.") + if distribution is None and "hf_config" in kwargs: + # The load builds from hf_config as given, so the plan must be the one it came with, not + # one decided here from another config. + raise TypeError( + "hf_config is the config resolve_distribution_plan returned with a plan: pass that plan " + "as distribution=, or pass config= to have the plan decided from your config." + ) hf_config = kwargs.get("hf_config") if distribution is None: plan, hf_config = resolve_distribution_plan( diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index f68a1e5e096..afdc0d8e6f2 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -235,6 +235,20 @@ def test_a_config_passed_for_the_load_is_the_one_it_decides_from( assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") +def test_hf_config_is_refused_without_the_plan_it_came_with( + run_as, serve, loaders, process_group, monkeypatch, tmp_path +): + """A sharded load builds from hf_config as given, so the plan must be the one it came with: a + plan decided from the checkpoint's own config would load a model it was not decided for.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(2, rank=1) + serve(LlamaConfig(**_TINY)) + + with pytest.raises(TypeError, match="returned with a plan: pass that plan"): + load_pretrained(str(tmp_path), hf_config=LlamaConfig(**{**_TINY, "num_hidden_layers": 2})) + assert loaders == {} + + def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, tmp_path): """No load here shards the routed experts: an EP plan fails before anything loads.""" run_as(2) From 80a2ee33b0e2dc3ba1e53f91b8d4680a76d911f5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:10:04 +0000 Subject: [PATCH 23/41] Say in load_pretrained that the FSDP2 load refuses a quantized checkpoint Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 07ccf96a7e7..31b7b706125 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -506,7 +506,8 @@ def load_pretrained( own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, ``mp_policy`` and ``broadcast_chunk_size`` ``kwargs`` may set). A single process has no group to shard through, so the load sets up one of one rank, which - :func:`~modelopt.torch.utils.distributed.distributed_run` destroys with the run. + :func:`~modelopt.torch.utils.distributed.distributed_run` destroys with the run. It reads the + checkpoint's tensors as stored, so it refuses a quantized checkpoint. * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to ``AutoModelForCausalLM``), placed as ``kwargs`` say. Every rank of a distributed run gets the whole model, ``device_map`` defaulting to the rank's GPU. From 3abad9859edbb24545e7d70eb47abcb48a6cc54f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:02:01 +0000 Subject: [PATCH 24/41] Tell a launcher-less FSDP2 load to use torchrun, and copy a handed config A single process's FSDP2 load sets up its group of one with the launcher's environment, so a process started with plain `python` failed inside init_process_group ("environment variable RANK expected"); it is now refused with the command that works, `torchrun --nproc_per_node=1`. The FSDP2 load built from the config passed with a handed plan as is, and building writes to it; it now builds from a copy, as the decided path and from_pretrained do. Refusing the placement options a sharded load takes over now also names how to place the model yourself. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 27 ++++++++++++++----- .../test_load_pretrained_distribution.py | 15 +++++++++++ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 31b7b706125..2039d5e74c4 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -452,9 +452,15 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) kwargs.pop("dtype", None) # A sharded load keeps the checkpoint's dtype. unknown = sorted(set(kwargs) - _SHARDED_LOAD_OPTIONS[route]) if unknown: + whole = ( + " To place the model yourself, have every rank load it whole: MODELOPT_FSDP_SIZE=1, or " + f"distribution=DistributionPlan({dist.size()}, dense='replicate')." + if route == _FSDP2 + else "" + ) raise TypeError( f"{unknown} do not apply to the {route} load, which places the model itself; it " - f"takes {sorted(_SHARDED_LOAD_OPTIONS[route])}." + f"takes {sorted(_SHARDED_LOAD_OPTIONS[route])}.{whole}" ) if checkpoint_has_mtp(model_type, ckpt_path): # prepare_model_for_loading builds the head on a whole model only; carried over, the @@ -477,6 +483,12 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) if not dist.is_initialized(): # FSDP2 shards through a process group even on one rank; a plan for more ranks than one # was refused without a group. A run (distributed_run) destroys the group on exit. + if "RANK" not in os.environ: + raise RuntimeError( + f"The {route} load shards through a process group, even on one rank, which a " + "process started without a launcher cannot set up; launch it with " + "`torchrun --nproc_per_node=1`." + ) dist.setup() return parallel_load_and_prepare_fsdp2( ckpt_path, @@ -504,9 +516,10 @@ def load_pretrained( * ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, ``mp_policy`` and - ``broadcast_chunk_size`` ``kwargs`` may set). A single process has no group to shard through, - so the load sets up one of one rank, which - :func:`~modelopt.torch.utils.distributed.distributed_run` destroys with the run. It reads the + ``broadcast_chunk_size`` ``kwargs`` may set). A single process launched with + ``torchrun --nproc_per_node=1`` has no group to shard through, so the load sets up one of one + rank, which :func:`~modelopt.torch.utils.distributed.distributed_run` destroys with the run; a + process started without a launcher cannot set one up and is refused. It reads the checkpoint's tensors as stored, so it refuses a quantized checkpoint. * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to ``AutoModelForCausalLM``), placed as ``kwargs`` say. Every rank of a distributed run gets the @@ -601,9 +614,11 @@ def checked() -> DistributionPlan: raise TypeError(f"distribution={distribution!r}: expected a DistributionPlan or None.") route = _route(plan) if route is not None and config is not None: - # The sharded load builds the model from hf_config: the plan's copy of this one, or it. + # The sharded load builds the model from hf_config: the plan's copy of this one, or a copy + # of it, since building writes to the config. del kwargs["config"] - kwargs.setdefault("hf_config", config) + if "hf_config" not in kwargs: + kwargs["hf_config"] = copy.deepcopy(config) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. model_type = getattr(hf_config or config, "model_type", None) or ( diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index afdc0d8e6f2..46b9df7c5d2 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -231,6 +231,7 @@ def test_a_config_passed_for_the_load_is_the_one_it_decides_from( assert reads == [] assert loaders["fsdp2"]["hf_config"].num_hidden_layers == 2 + assert loaders["fsdp2"]["hf_config"] is not config # building it writes to it assert config.to_dict() == as_passed assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") @@ -368,6 +369,7 @@ def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( config = serve(LlamaConfig(**_TINY)) if quantized: config.quantization_config = {"quant_method": "fp8"} + monkeypatch.setenv("RANK", "0") # torchrun --nproc_per_node=1 live = [] monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) monkeypatch.setattr(dist, "setup", lambda timeout=None: live.append(0)) @@ -389,6 +391,19 @@ def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( assert (call["rank"], call["world_size"]) == (0, 1) +def test_a_single_process_fsdp2_load_needs_a_launcher(monkeypatch, serve, loaders, tmp_path): + """Started without one, the process cannot set its group of one up; it is told to launch with + torchrun rather than fail inside init_process_group.""" + serve(LlamaConfig(**_TINY)) + monkeypatch.delenv("RANK", raising=False) + monkeypatch.setattr(dist, "is_initialized", lambda: False) + monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set a group up")) + + with pytest.raises(RuntimeError, match="torchrun --nproc_per_node=1"): + load_pretrained(str(tmp_path), distribution=DistributionPlan(dense="fsdp")) + assert loaders == {} + + # --------------------------------------------------------------------------------- fsdp2_wrap From b3008f8472ed8b8bd5922352169e185fcbf6d873 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:06:57 +0000 Subject: [PATCH 25/41] Refuse an FSDP2 load in a multi-rank launch whose group is not up, and name dtype A process launched as one of several ranks has no group until distributed_run sets it up, so dist.size() reads 1 and a one-rank FSDP2 plan passes the routing check. The load then set up the launch's whole group and sharded over every rank while recording a one-rank plan. It now refuses, telling the caller to set the launch's group up first, and the plan-size refusal says the same when that is why the sizes differ. A sharded load keeps the checkpoint's dtype; an explicit dtype is refused saying so, and that FSDP2's mp_policy sets the compute dtype, instead of reading like an unknown keyword. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 20 +++++++++++++++- .../test_load_pretrained_distribution.py | 23 +++++++++++++++---- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 2039d5e74c4..7394391da2e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -425,8 +425,14 @@ def _route(plan: DistributionPlan) -> str | None: from modelopt.torch.utils import distributed as dist if plan.world_size != dist.size(): + unset = ( + f" This process is one of a {os.environ['WORLD_SIZE']}-rank launch whose process " + "group is not set up: set it up first, with distributed_run()." + if dist.is_distributed_launch() and not dist.is_initialized() + else "" + ) raise ValueError( - f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." + f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}.{unset}" ) if plan.is_expert_parallel: raise NotImplementedError( @@ -450,6 +456,12 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) kwargs = dict(kwargs) if kwargs.get("dtype", "auto") == "auto": kwargs.pop("dtype", None) # A sharded load keeps the checkpoint's dtype. + if "dtype" in kwargs: + raise TypeError( + f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, so " + f"dtype={kwargs['dtype']!r} does not apply to it; FSDP2's mp_policy sets the dtype " + "it computes in." + ) unknown = sorted(set(kwargs) - _SHARDED_LOAD_OPTIONS[route]) if unknown: whole = ( @@ -489,6 +501,12 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) "process started without a launcher cannot set up; launch it with " "`torchrun --nproc_per_node=1`." ) + if dist.is_distributed_launch(): + # Its group would be the launch's, over ranks this plan was not decided for. + raise RuntimeError( + f"The {route} load needs this {os.environ['WORLD_SIZE']}-rank launch's process " + "group, which is not set up yet; set it up first, with distributed_run()." + ) dist.setup() return parallel_load_and_prepare_fsdp2( ckpt_path, diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 46b9df7c5d2..a7671e3b33d 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -279,6 +279,7 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path): ("change", "error", "match"), [ ({"kwargs": {"device_map": "auto"}}, TypeError, r"\['device_map'\] do not apply"), + ({"kwargs": {"dtype": torch.bfloat16}}, TypeError, "keeps their dtype"), ({"auto_class": AutoModel}, NotImplementedError, "builds its model class itself"), ({"mtp": True}, NotImplementedError, "cannot build the MTP head"), ( @@ -287,7 +288,7 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path): "config and hf_config both name", ), ], - ids=["placement", "model_class", "mtp_head", "two_configs"], + ids=["placement", "dtype", "model_class", "mtp_head", "two_configs"], ) def test_a_sharded_load_refuses_what_it_cannot_honor( run_as, serve, loaders, process_group, monkeypatch, tmp_path, change, error, match @@ -391,15 +392,27 @@ def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( assert (call["rank"], call["world_size"]) == (0, 1) -def test_a_single_process_fsdp2_load_needs_a_launcher(monkeypatch, serve, loaders, tmp_path): +@pytest.mark.parametrize( + ("launch", "match"), + [(None, "torchrun --nproc_per_node=1"), ("8", "8-rank launch's process group")], + ids=["no-launcher", "launch-without-its-group"], +) +def test_a_single_process_fsdp2_load_needs_a_launcher( + monkeypatch, serve, loaders, tmp_path, launch, match +): """Started without one, the process cannot set its group of one up; it is told to launch with - torchrun rather than fail inside init_process_group.""" + torchrun rather than fail inside init_process_group. One of a multi-rank launch whose group is + not up yet would set up the launch's group instead, over ranks its plan was not decided for.""" serve(LlamaConfig(**_TINY)) - monkeypatch.delenv("RANK", raising=False) + if launch is None: + monkeypatch.delenv("RANK", raising=False) + else: + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", launch) monkeypatch.setattr(dist, "is_initialized", lambda: False) monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set a group up")) - with pytest.raises(RuntimeError, match="torchrun --nproc_per_node=1"): + with pytest.raises(RuntimeError, match=match): load_pretrained(str(tmp_path), distribution=DistributionPlan(dense="fsdp")) assert loaders == {} From 68c668d35eece3765f830fe9d40bf89978920ba9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:46:19 +0000 Subject: [PATCH 26/41] Drop the mocked Hub-options test the plan's cached-revision test covers [5/21]'s resolve_distribution_plan now takes every Hub option load_pretrained does, and a test reads a two-revision cache offline through them. The mocked test pinned the exact keywords of the old three-option read, so it failed once the read also passed revision, and covered nothing the cached test does not. Signed-off-by: Shengliang Xu --- .../utils/test_resolve_distribution_plan.py | 24 ------------------- 1 file changed, 24 deletions(-) diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index f3c9b66bcca..14c60b7b359 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -358,30 +358,6 @@ def test_a_single_process_holds_the_whole_model( ) -def test_the_hub_options_reach_the_plans_checkpoint_read(monkeypatch, tmp_path): - """The plan's read resolves a Hub ID and loads its config with the options it was given: the - config reaches the Hub for the custom code its ``auto_map`` names in another repo.""" - options = {"token": "secret", "cache_dir": str(tmp_path), "local_files_only": True} - calls = [] - monkeypatch.setattr( - model_load_utils, - "ensure_local_checkpoint", - lambda path, **kwargs: calls.append(("resolve", path, kwargs)) or (path, str(tmp_path)), - ) - monkeypatch.setattr( - transformers.AutoConfig, - "from_pretrained", - lambda path, **kwargs: calls.append(("config", path, kwargs)) or _qwen3_moe(), - ) - - resolve_distribution_plan("org/model", **options) - - assert calls == [ - ("resolve", "org/model", options), - ("config", str(tmp_path), {"trust_remote_code": False, **options}), - ] - - @needs_qwen3_moe_ep_plan @pytest.mark.parametrize("rank", [0, 1]) def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, capsys, rank): From 278d43ed8d3e70fb7919aeced257024954cbc2f2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:25:33 +0000 Subject: [PATCH 27/41] Take rank 0's distribution on every rank; load a caller's class whole load_pretrained branched on each rank's own distribution argument: a handed plan went straight to a broadcast, none to the decision's collectives, and a non-plan raised before any collective. Ranks passing different kinds, as distribution=plan if rank == 0 else None does, waited in each other's collectives. Every rank now takes rank 0's distribution, a plan or none, through one broadcast before it branches, and rank 0 refuses a non-plan inside it, which fails every rank. Deciding a plan built the model on the meta device with a class of the decision's choosing, so a multi-rank load of a model only the caller's class builds, such as an encoder-decoder loaded through AutoModel, failed there. The sharded loads refuse any class but AutoModelForCausalLM, so resolve_distribution_plan now takes load_pretrained's auto_class: with another class, and no knob or argument asking to shard, every rank loads the whole model, planned without building it. fsdp2_wrap imports LayerActivationCollector with the module's other imports, as nothing imported from here leads back to this module. load_pretrained's docstring says the FSDP2 load builds its model without the Hub options, and the old fsdp2_wrap path says it needs the hf extra. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 3 +- .../torch/utils/plugins/model_load_utils.py | 84 ++++++++++++------- .../test_load_pretrained_distribution.py | 58 +++++++++++-- 3 files changed, 104 insertions(+), 41 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 8f96373f887..bdcacba93de 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -759,7 +759,8 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False """FSDP2-shard a HF causal LM's decoder layers, and the root unless ``shard_root=False``. It loads a model, so it lives with the loaders now: - :func:`modelopt.torch.utils.plugins.model_load_utils.fsdp2_wrap`. This keeps the old path. + :func:`modelopt.torch.utils.plugins.model_load_utils.fsdp2_wrap`. This keeps the old path, which + now needs the ``hf`` extra, as the loaders do. """ from modelopt.torch.utils.plugins.model_load_utils import fsdp2_wrap as wrap diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 7394391da2e..65eaf226e19 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -38,6 +38,7 @@ from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils import env from modelopt.torch.utils.distributed import ( DistributionPlan, @@ -530,7 +531,8 @@ def load_pretrained( """Load a Hugging Face model for ModelOpt to optimize and export, whole or distributed. ``distribution`` is the :class:`~modelopt.torch.utils.distributed.DistributionPlan` to spread - the model over the run's ranks by, :func:`resolve_distribution_plan`'s by default. By the plan: + the model over the run's ranks by. Every rank takes rank 0's, or, when rank 0 passes none, + :func:`resolve_distribution_plan`'s. By the plan: * ``dense="fsdp"``: a causal LM FSDP2-sharded over the whole process group, each rank reading its own layers (:func:`parallel_load_and_prepare_fsdp2`, whose ``cpu_offload``, ``mp_policy`` and @@ -560,10 +562,11 @@ def load_pretrained( * A Hub model ID is made local first (:func:`ensure_local_checkpoint`, which takes the ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` ``kwargs``): the export reads the source checkpoint from disk. Under ``torch.distributed`` every rank must call this. The - config and model loads still take all but ``revision``: an ``auto_map`` entry can name - custom code in another Hub repo, which Transformers fetches with them. A ``subfolder`` is - part of where the checkpoint is: the config, the loading hook and the record below read it, - as the load does. + config read and the whole-model load still take all but ``revision``: an ``auto_map`` entry + can name custom code in another Hub repo, which Transformers fetches with them. A sharded + load builds its model from the config with ``from_config``, which takes none, so such code + comes with the environment's Hub settings. A ``subfolder`` is part of where the checkpoint + is: the config, the loading hook and the record below read it, as the load does. * The model type's loading hook runs (:func:`~modelopt.torch.models.hf.prepare_model_for_loading`), e.g. to keep the MTP head a Nemotron-H checkpoint carries. ``model_type`` selects it; when omitted, the type of the @@ -603,33 +606,38 @@ def load_pretrained( config = config if isinstance(config, PretrainedConfig) else None if config is not None and "hf_config" in kwargs: raise TypeError("config and hf_config both name the config to load with; pass one.") - if distribution is None and "hf_config" in kwargs: - # The load builds from hf_config as given, so the plan must be the one it came with, not - # one decided here from another config. - raise TypeError( - "hf_config is the config resolve_distribution_plan returned with a plan: pass that plan " - "as distribution=, or pass config= to have the plan decided from your config." - ) + + def handed() -> DistributionPlan | None: + if distribution is None: + return None + if not isinstance(distribution, DistributionPlan): + raise TypeError(f"distribution={distribution!r}: expected a DistributionPlan or None.") + _route(distribution) + return distribution + + # Every rank takes rank 0's distribution, a plan or none, so every rank runs the same load and + # its collectives: a rank going by its own could wait in collectives the others never join, and + # ranks loading different plans would build mismatched meshes. Rank 0 checks it before it sends, + # so one it refuses raises on every rank instead of leaving them in the broadcast. + plan = dist.broadcast_from("distribution", handed) hf_config = kwargs.get("hf_config") - if distribution is None: + if plan is None: + if "hf_config" in kwargs: + # The load builds from hf_config as given, so the plan must be the one it came with, + # not one decided here from another config. + raise TypeError( + "hf_config is the config resolve_distribution_plan returned with a plan: pass that " + "plan as distribution=, or pass config= to have the plan decided from your config." + ) plan, hf_config = resolve_distribution_plan( - ckpt_path, config=config, trust_remote_code=trust_remote_code, **hub_options + ckpt_path, + config=config, + auto_class=auto_class, + trust_remote_code=trust_remote_code, + **hub_options, ) if _route(plan) is not None: kwargs.setdefault("hf_config", hf_config) - elif isinstance(distribution, DistributionPlan): - # Every rank loads the plan rank 0 was handed, as it would one it decided: ranks loading - # different plans would build mismatched meshes and hang. Rank 0 checks it before it sends, - # so a plan it refuses raises on every rank instead of leaving them in the broadcast. - handed = distribution - - def checked() -> DistributionPlan: - _route(handed) - return handed - - plan = dist.broadcast_from("plan", checked) - else: - raise TypeError(f"distribution={distribution!r}: expected a DistributionPlan or None.") route = _route(plan) if route is not None and config is not None: # The sharded load builds the model from hf_config: the plan's copy of this one, or a copy @@ -692,9 +700,6 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) before calibrating or exporting it. """ - # 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( @@ -1229,7 +1234,7 @@ def _ep_degree(ep_size) -> int | str: def _decide_distribution( - ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool + ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool, auto_class ) -> DistributionPlan: """The plan :func:`resolve_distribution_plan` returns, decided by this rank alone.""" from modelopt.torch.utils.moe import ( @@ -1252,6 +1257,15 @@ def _decide_distribution( ) ep_size = _ep_degree(ep_size) dense = _dense_from_env(world_size) if dense is None else dense + if ( + auto_class not in (None, AutoModelForCausalLM) + and ep_size in ("auto", 1) + and dense in ("auto", "replicate") + ): + # The sharded loads build their own model class and refuse another, so with nothing asking + # to shard, a load with this one gives every rank the whole model: there is nothing to + # measure, nor a model to build for it. + return DistributionPlan(world_size) try: model = build_meta_causal_lm( ckpt_path, @@ -1310,6 +1324,7 @@ def resolve_distribution_plan( ckpt_path: str, *, config: PretrainedConfig | None = None, + auto_class: Any = None, trust_remote_code: bool = False, ep_size: int | str | None = None, dense: str | None = None, @@ -1350,6 +1365,11 @@ def resolve_distribution_plan( all but ``subfolder`` resolve a Hub model ID (:func:`ensure_local_checkpoint`), and the config read takes all but ``revision``, for custom code its ``auto_map`` names in another Hub repo. + ``auto_class`` is the class the load builds the model with, :func:`load_pretrained`'s. The + sharded loads build their own and refuse any other than ``AutoModelForCausalLM``, so with + another, and no knob or argument asking to shard, every rank loads the whole model, planned + without building it. + ``config``, the config the load will use when the caller has one, is decided from instead of the checkpoint's, which is then not read: each rank takes a copy, as deciding may add to it. """ @@ -1388,7 +1408,7 @@ def checkpoint_read() -> bool: plan = dist.broadcast_from( "plan", lambda: _decide_distribution( - ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code + ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code, auto_class ), ) if plan != DistributionPlan(): diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index a7671e3b33d..d0a321df350 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -250,6 +250,31 @@ def test_hf_config_is_refused_without_the_plan_it_came_with( assert loaders == {} +@pytest.mark.parametrize("fsdp_size", [None, "world"], ids=["no-knob", "knob-shards"]) +def test_a_model_class_of_the_callers_own_loads_whole( + run_as, serve, loaders, auto_class, monkeypatch, tmp_path, fsdp_size +): + """The sharded loads build their own model class and refuse another, so with nothing asking to + shard, a class of the caller's own loads whole on every rank, planned without building the + model on the meta device, which the decision's own pick of class may not build (an + encoder-decoder loaded through AutoModel). A knob asking to shard still decides, and the + sharded load refuses the class.""" + run_as(2) + serve(LlamaConfig(**_TINY)) + if fsdp_size is None: + monkeypatch.setattr( + model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") + ) + model = load_pretrained(str(tmp_path), auto_class) + assert distribution_plan(model) == DistributionPlan(2) + assert len(auto_class.calls) == 1 + else: + monkeypatch.setenv("MODELOPT_FSDP_SIZE", fsdp_size) + with pytest.raises(NotImplementedError, match="builds its model class itself"): + load_pretrained(str(tmp_path), auto_class) + assert loaders == {} + + def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, tmp_path): """No load here shards the routed experts: an EP plan fails before anything loads.""" run_as(2) @@ -261,18 +286,35 @@ def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, assert auto_class.calls == [] -def _every_rank_loads_rank0s_plan_job(rank, size, ckpt): - handed = DistributionPlan(size) if rank == 0 else DistributionPlan(size, dense="fsdp") +def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed): + model, refused = None, None with bounded_collectives(): - model = load_pretrained(ckpt, distribution=handed, device_map="cpu") - assert distribution_plan(model) == DistributionPlan(size), rank + try: # not pytest.raises: its failure would end the pooled worker instead of the job + model = load_pretrained(ckpt, distribution=handed[rank], device_map="cpu") + except Exception as error: + refused = str(error) + if isinstance(handed[0], DistributionPlan): + assert refused is None and distribution_plan(model) == handed[0], (rank, refused) + else: + # Rank 0's refusal, on every rank, rather than a rank left waiting in the broadcast. + assert "expected a DistributionPlan or None" in str(refused), (rank, refused) -def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path): - """Over a real broadcast. Ranks loading different plans would build mismatched meshes and hang: - rank 1, handed an FSDP2 plan of its own, loads the whole model by the plan rank 0 was handed.""" +@pytest.mark.parametrize( + "handed", + [ + (DistributionPlan(2), DistributionPlan(2, dense="fsdp")), + (DistributionPlan(2), None), + ("fsdp2", DistributionPlan(2)), + ], + ids=["another-plan", "none", "not-a-plan"], +) +def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, handed): + """Over a real broadcast, whatever rank 1 was handed. Ranks loading different plans would build + mismatched meshes, and one handed none would decide a plan through collectives the others never + join: both would hang. A distribution rank 0 refuses fails every rank.""" LlamaForCausalLM(LlamaConfig(**_TINY)).save_pretrained(tmp_path) - cpu_dist_workers.run(_every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path)) + cpu_dist_workers.run(_every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path), handed=handed) @pytest.mark.parametrize( From ec43eddbf71790acb095a0a929899c587341d2e7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 02:07:23 +0000 Subject: [PATCH 28/41] Decide from a copy of the config; read an FSDP2 load's config once Deciding a plan builds the model on the meta device, on rank 0 alone, from the config it then returns for a sharded load to build from, and building writes to that config: its dtype and its attention and experts implementations. Rank 0's FSDP2 load so built from a config the other ranks' loads did not. The decision now builds from a copy, so every rank returns the config as it read it. With a plan handed in and no config, load_pretrained read the checkpoint's config for the model type, and the FSDP2 load read it again to build from. It is now read once and handed to the load, so the quantized-checkpoint refusal checks the config the model type came from. The changelog entry now says that, until the expert-parallel load lands, an MoE planned expert-parallel, as the default MODELOPT_EP_DEGREE=auto plans one whose config has a base_model_ep_plan, is refused before it loads, and how to load it meanwhile. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- modelopt/torch/utils/plugins/model_load_utils.py | 11 ++++++++++- .../torch/utils/test_load_pretrained_distribution.py | 10 +++++----- .../torch/utils/test_resolve_distribution_plan.py | 11 +++++++++++ 4 files changed, 27 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 838b58cfdfb..ad20411376e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -35,7 +35,7 @@ Changelog *Misc* - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. -- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_FSDP_SIZE`` variable, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now also loads a model distributed. Under ``torchrun`` it FSDP2-shards the model or gives every rank a whole copy, by the plan the new ``resolve_distribution_plan`` decides from the model and the ``MODELOPT_FSDP_SIZE`` variable, or by a plan passed as its new ``distribution`` argument. The loaded model records the plan (``modelopt.torch.utils.distributed.distribution_plan``). Until the expert-parallel load lands, an MoE planned expert-parallel, as the default ``MODELOPT_EP_DEGREE=auto`` plans one whose config has a ``base_model_ep_plan``, is refused before it loads: set ``MODELOPT_EP_DEGREE=1`` (with ``MODELOPT_FSDP_SIZE=world`` to FSDP2-shard it) or pass a plan with ``ep_size=1``. **Backward Breaking Changes** diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 65eaf226e19..7a46dc6004e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -647,6 +647,12 @@ def handed() -> DistributionPlan | None: kwargs["hf_config"] = copy.deepcopy(config) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. + if route == _FSDP2 and hf_config is None and config is None: + # The FSDP2 load builds from the checkpoint's config too: read it once, for both. + hf_config = AutoConfig.from_pretrained( + ckpt_path, trust_remote_code=trust_remote_code, **hub_options + ) + kwargs["hf_config"] = hf_config model_type = getattr(hf_config or config, "model_type", None) or ( AutoConfig.from_pretrained( ckpt_path, trust_remote_code=trust_remote_code, **hub_options @@ -1266,6 +1272,9 @@ def _decide_distribution( # to shard, a load with this one gives every rank the whole model: there is nothing to # measure, nor a model to build for it. return DistributionPlan(world_size) + # Building writes to the config it builds from, and only rank 0 decides: build from a copy, so + # the config returned with the plan, which a sharded load builds from, is alike on every rank. + hf_config = copy.deepcopy(hf_config) try: model = build_meta_causal_lm( ckpt_path, @@ -1371,7 +1380,7 @@ def resolve_distribution_plan( without building it. ``config``, the config the load will use when the caller has one, is decided from instead of - the checkpoint's, which is then not read: each rank takes a copy, as deciding may add to it. + the checkpoint's, which is then not read: each rank takes a copy, as building writes to it. """ from modelopt.torch.utils import distributed as dist diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index d0a321df350..620729b4ea2 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -172,8 +172,8 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( run_as, loaders, process_group, monkeypatch, tmp_path ): - """FSDP2 shards over the process group it reads its rank and GPU from, by a config read with - the Hub options the load was given (custom code an ``auto_map`` names in another repo).""" + """FSDP2 shards over the process group it reads its rank and GPU from, by a config read once, + with the Hub options the load was given (custom code an ``auto_map`` names in another repo).""" run_as(2, rank=1) served, reads = LlamaConfig(**_TINY), [] monkeypatch.setattr( @@ -189,7 +189,7 @@ def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( assert (call["device"], call["rank"], call["world_size"]) == (torch.device("cuda:1"), 1, 2) assert call["cpu_offload"] is True assert call["hf_config"] is served - assert reads and all(read.get("local_files_only") is True for read in reads) + assert [read.get("local_files_only") for read in reads] == [True] assert distribution_plan(model) == plan assert loaders == {} @@ -217,8 +217,8 @@ def test_a_config_passed_for_the_load_is_the_one_it_decides_from( run_as, loaders, process_group, monkeypatch, tmp_path, handed ): """No config is read: the plan is decided from the one passed for the load, which the FSDP2 - load then builds from, as it does with a plan handed in. Deciding adds to that config, so it - works on a copy of it.""" + load then builds from, as it does with a plan handed in. Building writes to that config, so the + load works on a copy of it.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") run_as(2, rank=1) reads: list = [] diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 14c60b7b359..164e5f94f46 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -665,6 +665,17 @@ def remote_from_config(config, **kwargs): assert ProbeConfig not in mapping +def test_the_config_returned_with_the_plan_is_as_every_rank_read_it(run_as, serve, tmp_path): + """Rank 0 alone decides, building the model, which writes to the config it builds from; a + sharded load builds from the config returned, so rank 0's must be the one the others read.""" + run_as(2) + serve(LlamaConfig(**_TINY)) + + _, hf_config = resolve_distribution_plan(str(tmp_path)) + + assert vars(hf_config) == vars(LlamaConfig(**_TINY)) + + def test_the_hub_options_locate_the_checkpoint_the_load_reads(tmp_path): """``revision``, ``cache_dir``, ``local_files_only`` and ``subfolder`` find the checkpoint as ``load_pretrained`` takes them, so the plan and its config are those of what the load reads.""" From 207989e6f17e2fe357ae6e0dfc2202cceed53a37 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 03:54:44 +0000 Subject: [PATCH 29/41] Refuse an FSDP2 load of a VLM before it builds one as a causal LM The FSDP2 load builds its model through AutoModelForCausalLM, and refused only a model class the caller passed, not a config that needs another. A VLM's config reached it whenever its plan FSDP2-shards it: auto picks that for a VLM too big for one GPU, and MODELOPT_FSDP_SIZE=world asks for it. AutoModelForCausalLM then refused the config deep in the loader (Qwen2.5-VL), or built the language model alone (Mllama; Qwen3.5 on Transformers 5.18), whose vision weights the export would carry over unquantized. _load_sharded now refuses a config that _auto_class does not build through AutoModelForCausalLM, before it sets up a group or reads anything, and says how to load the model whole on every rank: its model class as auto_class, with a replicated plan. resolve_distribution_plan's docstring says to hand the config it returns to load_pretrained only with a plan that shards the model: a whole-model load reads the checkpoint's config itself, and refuses hf_config. The load refusal test covers a VLM checkpoint; it fails without the refusal. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 12 +++++- .../test_load_pretrained_distribution.py | 42 ++++++++++++++----- 2 files changed, 42 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 7a46dc6004e..4478c52c2f6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -493,6 +493,14 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) raise NotImplementedError( f"The FSDP2 load cannot dequantize this {model_type} checkpoint; load it unsharded." ) + if _auto_class(hf_config, trust_remote_code=trust_remote_code) is not AutoModelForCausalLM: + # It builds through AutoModelForCausalLM, which refuses most VLMs' configs and, for some, + # builds the language model alone, whose other towers the export would carry unquantized. + raise NotImplementedError( + f"The FSDP2 load builds a causal LM, which this {model_type} checkpoint is not. To load " + "it whole on every rank, pass its model class as auto_class, with MODELOPT_FSDP_SIZE=1 " + f"or distribution=DistributionPlan({dist.size()}, dense='replicate')." + ) if not dist.is_initialized(): # FSDP2 shards through a process group even on one rank; a plan for more ranks than one # was refused without a group. A run (distributed_run) destroys the group on exit. @@ -1349,7 +1357,9 @@ def resolve_distribution_plan( records it on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), and every later step reads it back from there (:func:`~modelopt.torch.utils.distributed.distribution_plan`). Returns ``(plan, hf_config)``, - ``hf_config`` being the config the plan was decided from, for a sharded load to build from. + ``hf_config`` being the config the plan was decided from, for a sharded load to build from: hand + it to :func:`load_pretrained` only with a plan that shards the model, as a whole-model load reads + the checkpoint's config itself, and refuses ``hf_config``. Rank 0 decides, from the model built on the meta device: from its parameters rather than estimates of them, and from its routed experts as an expert-parallel load finds them. Every rank diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 620729b4ea2..4d3b29f2ed8 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -354,26 +354,45 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( @pytest.mark.parametrize( - ("distribution", "quantized", "whole_model_kwargs", "error", "match"), + ("distribution", "checkpoint", "whole_model_kwargs", "error", "match"), [ ( DistributionPlan(4, dense="fsdp"), - False, + "causal-lm", False, ValueError, "is for 4 ranks, but this run has 2", ), - (DistributionPlan(2, dense="fsdp"), True, False, NotImplementedError, "cannot dequantize"), - ("fsdp2", False, False, TypeError, "expected a DistributionPlan or None"), ( - DistributionPlan(2), + DistributionPlan(2, dense="fsdp"), + "quantized", + False, + NotImplementedError, + "cannot dequantize", + ), + ( + DistributionPlan(2, dense="fsdp"), + "vlm", False, + NotImplementedError, + "builds a causal LM, which this llava checkpoint is not", + ), + ("fsdp2", "causal-lm", False, TypeError, "expected a DistributionPlan or None"), + ( + DistributionPlan(2), + "causal-lm", True, TypeError, "hf_config applies to a sharded load only", ), ], - ids=["another-world", "quantized-checkpoint", "not-a-plan", "hf-config-on-a-whole-load"], + ids=[ + "another-world", + "quantized-checkpoint", + "vlm-checkpoint", + "not-a-plan", + "hf-config-on-a-whole-load", + ], ) def test_a_load_refuses_what_it_cannot_load( run_as, @@ -383,17 +402,18 @@ def test_a_load_refuses_what_it_cannot_load( process_group, tmp_path, distribution, - quantized, + checkpoint, whole_model_kwargs, error, match, ): """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would - read as stored; anything but a plan; and the sharded loads' config, which from_pretrained would - take for a model argument and fail far from here.""" + read as stored; a VLM, which FSDP2 would build through AutoModelForCausalLM, refused or as its + language model alone; anything but a plan; and the sharded loads' config, which from_pretrained + would take for a model argument and fail far from here.""" run_as(2) - config = serve(LlamaConfig(**_TINY)) - if quantized: + config = serve(transformers.LlavaConfig() if checkpoint == "vlm" else LlamaConfig(**_TINY)) + if checkpoint == "quantized": config.quantization_config = {"quant_method": "fp8"} whole = {"auto_class": auto_class, "hf_config": config} if whole_model_kwargs else {} From 4a98e5f4cacc84173591db278f1d31b18ff1d329 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:43 +0000 Subject: [PATCH 30/41] Refuse a load on a rank of a launch whose process group is not set up Without its group every collective runs on that rank alone, so each rank of the launch planned and loaded the model as a single process's, the whole model on every rank, and recorded a one-rank plan on a model the whole launch holds. load_pretrained now refuses it first, on every route; the two route-specific hints it makes unreachable go. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 28 ++++++------- .../test_load_pretrained_distribution.py | 41 +++++++++++-------- 2 files changed, 37 insertions(+), 32 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 4478c52c2f6..078e757f2ef 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -426,15 +426,7 @@ def _route(plan: DistributionPlan) -> str | None: from modelopt.torch.utils import distributed as dist if plan.world_size != dist.size(): - unset = ( - f" This process is one of a {os.environ['WORLD_SIZE']}-rank launch whose process " - "group is not set up: set it up first, with distributed_run()." - if dist.is_distributed_launch() and not dist.is_initialized() - else "" - ) - raise ValueError( - f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}.{unset}" - ) + raise ValueError(f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}.") if plan.is_expert_parallel: raise NotImplementedError( f"{plan!r} shards the routed experts, which load_pretrained has no load for. Plan " @@ -510,12 +502,6 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) "process started without a launcher cannot set up; launch it with " "`torchrun --nproc_per_node=1`." ) - if dist.is_distributed_launch(): - # Its group would be the launch's, over ranks this plan was not decided for. - raise RuntimeError( - f"The {route} load needs this {os.environ['WORLD_SIZE']}-rank launch's process " - "group, which is not set up yet; set it up first, with distributed_run()." - ) dist.setup() return parallel_load_and_prepare_fsdp2( ckpt_path, @@ -554,7 +540,9 @@ def load_pretrained( whole model, ``device_map`` defaulting to the rank's GPU. A plan with ``ep_size > 1`` raises: no load here shards the routed experts. On a distributed run - every rank loads the plan rank 0 was handed. The plan is recorded on the model + every rank loads the plan rank 0 was handed. A rank of a multi-rank launch whose process group + is not set up is refused on every route: set it up first, with + :func:`~modelopt.torch.utils.distributed.distributed_run`. The plan is recorded on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), where calibration and export read it back. @@ -601,6 +589,14 @@ def load_pretrained( from modelopt.torch.models.hf import prepare_model_for_loading from modelopt.torch.utils import distributed as dist + if dist.is_distributed_launch() and not dist.is_initialized(): + # Without its group every collective below runs on this rank alone: each rank would plan + # and load the model as a single process's, and record a one-rank plan on a model the whole + # launch holds. + raise RuntimeError( + f"This process is one of a {os.environ['WORLD_SIZE']}-rank launch whose process group " + "is not set up; set it up first, with distributed_run()." + ) hub_options = { k: kwargs.pop(k) for k in ("token", "cache_dir", "local_files_only") if k in kwargs } diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 4d3b29f2ed8..c4b6e970e00 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -454,31 +454,40 @@ def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( assert (call["rank"], call["world_size"]) == (0, 1) -@pytest.mark.parametrize( - ("launch", "match"), - [(None, "torchrun --nproc_per_node=1"), ("8", "8-rank launch's process group")], - ids=["no-launcher", "launch-without-its-group"], -) -def test_a_single_process_fsdp2_load_needs_a_launcher( - monkeypatch, serve, loaders, tmp_path, launch, match -): +def test_a_single_process_fsdp2_load_needs_a_launcher(monkeypatch, serve, loaders, tmp_path): """Started without one, the process cannot set its group of one up; it is told to launch with - torchrun rather than fail inside init_process_group. One of a multi-rank launch whose group is - not up yet would set up the launch's group instead, over ranks its plan was not decided for.""" + torchrun rather than fail inside init_process_group.""" serve(LlamaConfig(**_TINY)) - if launch is None: - monkeypatch.delenv("RANK", raising=False) - else: - monkeypatch.setenv("RANK", "0") - monkeypatch.setenv("WORLD_SIZE", launch) + monkeypatch.delenv("RANK", raising=False) monkeypatch.setattr(dist, "is_initialized", lambda: False) monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set a group up")) - with pytest.raises(RuntimeError, match=match): + with pytest.raises(RuntimeError, match="torchrun --nproc_per_node=1"): load_pretrained(str(tmp_path), distribution=DistributionPlan(dense="fsdp")) assert loaders == {} +@pytest.mark.parametrize("distribution", [None, DistributionPlan(dense="fsdp")], ids=["decided", "fsdp"]) +def test_a_rank_of_a_launch_without_its_group_is_refused( + monkeypatch, serve, loaders, auto_class, tmp_path, distribution +): + """Without its group every collective runs on this rank alone, so each rank of the launch would + plan and load the model as a single process's, the whole model on every rank, and record a + one-rank plan on a model the whole launch holds.""" + serve(LlamaConfig(**_TINY)) + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "8") + monkeypatch.setattr(dist, "is_initialized", lambda: False) + monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set a group up")) + + with pytest.raises(RuntimeError, match="8-rank launch whose process group is not set up"): + load_pretrained( + str(tmp_path), auto_class if distribution is None else None, distribution=distribution + ) + assert auto_class.calls == [] + assert loaders == {} + + # --------------------------------------------------------------------------------- fsdp2_wrap From c8b7f4c5cd1554bb2314a8a01d6b90ab132b02d6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:50 +0000 Subject: [PATCH 31/41] Refuse a caller's model class before building the model a sharded plan needs With a knob asking to shard, the decision built the model on the meta device through its own pick of class, which fails for a model the caller loads through another class (an encoder-decoder through AutoModel), before the sharded load could refuse the class clearly. Every sharded load refuses it, so the decision does, before building. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 9 +++++++++ .../utils/test_load_pretrained_distribution.py | 13 ++++++------- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 078e757f2ef..ccdf78ff873 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1276,6 +1276,15 @@ def _decide_distribution( # to shard, a load with this one gives every rank the whole model: there is nothing to # measure, nor a model to build for it. return DistributionPlan(world_size) + if auto_class not in (None, AutoModelForCausalLM): + # Something asks to shard, and every sharded load refuses this class: refuse it here, rather + # than build the model of the decision's own pick of class, which may not build at all. + raise NotImplementedError( + f"A sharded load builds its model class itself, so auto_class={auto_class.__name__} " + f"does not apply to it, and ep_size={ep_size!r}, dense={dense!r} ask to shard. To give " + "every rank the whole model, plan ep_size=1 and dense='replicate' " + "(MODELOPT_EP_DEGREE=1, MODELOPT_FSDP_SIZE=1)." + ) # Building writes to the config it builds from, and only rank 0 decides: build from a copy, so # the config returned with the plan, which a sharded load builds from, is alike on every rank. hf_config = copy.deepcopy(hf_config) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index c4b6e970e00..476e7af8d91 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -255,16 +255,15 @@ def test_a_model_class_of_the_callers_own_loads_whole( run_as, serve, loaders, auto_class, monkeypatch, tmp_path, fsdp_size ): """The sharded loads build their own model class and refuse another, so with nothing asking to - shard, a class of the caller's own loads whole on every rank, planned without building the - model on the meta device, which the decision's own pick of class may not build (an - encoder-decoder loaded through AutoModel). A knob asking to shard still decides, and the - sharded load refuses the class.""" + shard, a class of the caller's own loads whole on every rank, and with a knob asking to shard + it is refused. Neither builds the model on the meta device, which the decision's own pick of + class may not build (an encoder-decoder loaded through AutoModel).""" run_as(2) serve(LlamaConfig(**_TINY)) + monkeypatch.setattr( + model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") + ) if fsdp_size is None: - monkeypatch.setattr( - model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") - ) model = load_pretrained(str(tmp_path), auto_class) assert distribution_plan(model) == DistributionPlan(2) assert len(auto_class.calls) == 1 From d267cb8dfceae9a527688820b9b7e8fadea9552e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:50 +0000 Subject: [PATCH 32/41] Note in the changelog that fsdp2_wrap moved with the loaders Calling it through its old path now needs the hf extra, as the loaders do. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ad20411376e..1fa2fb0d93e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -41,6 +41,7 @@ Changelog - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 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. **Deprecations** From 3f028c3c48ae8e26223c7badeac886326fd1cb37 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:15:21 +0000 Subject: [PATCH 33/41] Wrap two lines of the launch refusal as ruff formats them Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 4 +++- tests/unit/torch/utils/test_load_pretrained_distribution.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ccdf78ff873..ad0edb2302e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -426,7 +426,9 @@ def _route(plan: DistributionPlan) -> str | None: from modelopt.torch.utils import distributed as dist if plan.world_size != dist.size(): - raise ValueError(f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}.") + raise ValueError( + f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." + ) if plan.is_expert_parallel: raise NotImplementedError( f"{plan!r} shards the routed experts, which load_pretrained has no load for. Plan " diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 476e7af8d91..c22b17ce316 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -466,7 +466,9 @@ def test_a_single_process_fsdp2_load_needs_a_launcher(monkeypatch, serve, loader assert loaders == {} -@pytest.mark.parametrize("distribution", [None, DistributionPlan(dense="fsdp")], ids=["decided", "fsdp"]) +@pytest.mark.parametrize( + "distribution", [None, DistributionPlan(dense="fsdp")], ids=["decided", "fsdp"] +) def test_a_rank_of_a_launch_without_its_group_is_refused( monkeypatch, serve, loaders, auto_class, tmp_path, distribution ): From 34ff861421df2668aec2cfd28249e443d04273f2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 17:58:06 +0000 Subject: [PATCH 34/41] Take dtype=None and torch_dtype in a sharded load as from_pretrained does from_pretrained reads dtype=None as "auto", and torch_dtype as the older name of dtype. The sharded load refused dtype=None, which keeps the checkpoint's dtype as "auto" does, and gave torch_dtype, which hf_ptq's example_utils sets for DeciLM checkpoints, the placement refusal instead of the dtype one. It now checks both names: None and "auto" keep the checkpoint's dtype, any other value is refused for the dtype it asks for. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 17 ++++++++-------- .../test_load_pretrained_distribution.py | 20 +++++++++++++++---- 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ad0edb2302e..807d00dadc5 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -449,14 +449,15 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) "does not apply to it." ) kwargs = dict(kwargs) - if kwargs.get("dtype", "auto") == "auto": - kwargs.pop("dtype", None) # A sharded load keeps the checkpoint's dtype. - if "dtype" in kwargs: - raise TypeError( - f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, so " - f"dtype={kwargs['dtype']!r} does not apply to it; FSDP2's mp_policy sets the dtype " - "it computes in." - ) + # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. + for name in ("dtype", "torch_dtype"): + value = kwargs.pop(name, None) + if value not in (None, "auto"): # A sharded load keeps the checkpoint's dtype. + raise TypeError( + f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, " + f"so {name}={value!r} does not apply to it; FSDP2's mp_policy sets the dtype it " + "computes in." + ) unknown = sorted(set(kwargs) - _SHARDED_LOAD_OPTIONS[route]) if unknown: whole = ( diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index c22b17ce316..4547578b8d1 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -169,11 +169,18 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( assert loaders == {} +@pytest.mark.parametrize( + "dtype", + [{"dtype": "auto"}, {"dtype": None}, {"torch_dtype": "auto"}], + ids=["auto", "none", "torch_dtype"], +) def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( - run_as, loaders, process_group, monkeypatch, tmp_path + run_as, loaders, process_group, monkeypatch, tmp_path, dtype ): """FSDP2 shards over the process group it reads its rank and GPU from, by a config read once, - with the Hub options the load was given (custom code an ``auto_map`` names in another repo).""" + with the Hub options the load was given (custom code an ``auto_map`` names in another repo). + It keeps the checkpoint's dtype, as from_pretrained does for "auto" and for None, under either + name of the argument.""" run_as(2, rank=1) served, reads = LlamaConfig(**_TINY), [] monkeypatch.setattr( @@ -182,7 +189,7 @@ def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( plan = DistributionPlan(2, dense="fsdp") model = load_pretrained( - str(tmp_path), distribution=plan, cpu_offload=True, dtype="auto", local_files_only=True + str(tmp_path), distribution=plan, cpu_offload=True, local_files_only=True, **dtype ) call = loaders.pop("fsdp2") @@ -321,6 +328,11 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, [ ({"kwargs": {"device_map": "auto"}}, TypeError, r"\['device_map'\] do not apply"), ({"kwargs": {"dtype": torch.bfloat16}}, TypeError, "keeps their dtype"), + ( + {"kwargs": {"torch_dtype": torch.bfloat16}}, + TypeError, + r"keeps their dtype, so torch_dtype=torch\.bfloat16", + ), ({"auto_class": AutoModel}, NotImplementedError, "builds its model class itself"), ({"mtp": True}, NotImplementedError, "cannot build the MTP head"), ( @@ -329,7 +341,7 @@ def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, "config and hf_config both name", ), ], - ids=["placement", "dtype", "model_class", "mtp_head", "two_configs"], + ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs"], ) def test_a_sharded_load_refuses_what_it_cannot_honor( run_as, serve, loaders, process_group, monkeypatch, tmp_path, change, error, match From 11d2e1c30cb6fce3a3fb206cb2a1f2f740d93122 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:03:42 +0000 Subject: [PATCH 35/41] Import the layer detection where fsdp2_wrap uses it Importing it at module level loaded the whole quantization package, which registers every mode and plugin, into every import of the loaders. Only a wrap needs it. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 807d00dadc5..ec700a4d3d9 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -38,7 +38,6 @@ from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key -from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils import env from modelopt.torch.utils.distributed import ( DistributionPlan, @@ -713,6 +712,9 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) before calibrating or exporting it. """ + # 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( From 4a885a83c9f22ebe809142e8d9a812a876ddedae Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:06:41 +0000 Subject: [PATCH 36/41] Check a handed plan before the checkpoint is fetched Making the checkpoint local can mean a download of the whole model, and a plan the load refuses failed only after it. Rank 0 now checks and broadcasts the plan first, and the config arguments are checked before the fetch too. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 35 ++++++++++--------- .../test_load_pretrained_distribution.py | 11 ++++++ 2 files changed, 29 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ec700a4d3d9..276108d533d 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -599,15 +599,6 @@ def load_pretrained( f"This process is one of a {os.environ['WORLD_SIZE']}-rank launch whose process group " "is not set up; set it up first, with distributed_run()." ) - hub_options = { - k: kwargs.pop(k) for k in ("token", "cache_dir", "local_files_only") if k in kwargs - } - revision = {"revision": kwargs.pop("revision")} if "revision" in kwargs else {} - _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) - subfolder = kwargs.pop("subfolder", None) - if subfolder: - ckpt_path = os.path.join(ckpt_path, subfolder) - trust_remote_code = kwargs.get("trust_remote_code", False) config = kwargs.get("config") config = config if isinstance(config, PretrainedConfig) else None if config is not None and "hf_config" in kwargs: @@ -624,17 +615,27 @@ def handed() -> DistributionPlan | None: # Every rank takes rank 0's distribution, a plan or none, so every rank runs the same load and # its collectives: a rank going by its own could wait in collectives the others never join, and # ranks loading different plans would build mismatched meshes. Rank 0 checks it before it sends, - # so one it refuses raises on every rank instead of leaving them in the broadcast. + # so one it refuses raises on every rank instead of leaving them in the broadcast, and before + # the checkpoint is fetched. plan = dist.broadcast_from("distribution", handed) + if plan is None and "hf_config" in kwargs: + # The load builds from hf_config as given, so the plan must be the one it came with, not + # one decided here from another config. + raise TypeError( + "hf_config is the config resolve_distribution_plan returned with a plan: pass that " + "plan as distribution=, or pass config= to have the plan decided from your config." + ) + hub_options = { + k: kwargs.pop(k) for k in ("token", "cache_dir", "local_files_only") if k in kwargs + } + revision = {"revision": kwargs.pop("revision")} if "revision" in kwargs else {} + _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) + subfolder = kwargs.pop("subfolder", None) + if subfolder: + ckpt_path = os.path.join(ckpt_path, subfolder) + trust_remote_code = kwargs.get("trust_remote_code", False) hf_config = kwargs.get("hf_config") if plan is None: - if "hf_config" in kwargs: - # The load builds from hf_config as given, so the plan must be the one it came with, - # not one decided here from another config. - raise TypeError( - "hf_config is the config resolve_distribution_plan returned with a plan: pass that " - "plan as distribution=, or pass config= to have the plan decided from your config." - ) plan, hf_config = resolve_distribution_plan( ckpt_path, config=config, diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 4547578b8d1..3c27a8da360 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -434,6 +434,17 @@ def test_a_load_refuses_what_it_cannot_load( assert auto_class.calls == [] +def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched(run_as, monkeypatch, tmp_path): + """Making the checkpoint local can mean downloading it: a plan the load refuses fails first.""" + run_as(2) + monkeypatch.setattr( + model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") + ) + + with pytest.raises(ValueError, match="is for 4 ranks"): + load_pretrained(str(tmp_path), distribution=DistributionPlan(4)) + + @pytest.mark.parametrize("quantized", [False, True]) def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( monkeypatch, serve, loaders, tmp_path, quantized From 3fd79f8b2d80f847052f9226edc579c461aca0bf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:09:09 +0000 Subject: [PATCH 37/41] Load whole on each process of a launch that sets no group up The launch guard refused every route whenever WORLD_SIZE named several ranks and no process group was set up, so hf_ptq launched with torchrun but without --use_fsdp2, where nothing sets a group up, failed where it used to load. Each such process now loads as a single process does, the whole model with no plan to agree on, as before plans existed. Only what needs the group is refused there: a plan for several ranks, before anything is fetched, and a sharded load, which would otherwise set up a group of all the launch's ranks for a one-rank plan. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 32 ++++++++---- .../test_load_pretrained_distribution.py | 51 +++++++++++++------ 2 files changed, 58 insertions(+), 25 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 276108d533d..ad7b1785d27 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -498,6 +498,12 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) if not dist.is_initialized(): # FSDP2 shards through a process group even on one rank; a plan for more ranks than one # was refused without a group. A run (distributed_run) destroys the group on exit. + if dist.is_distributed_launch(): + # One rank of a launch of several cannot set the group up for itself. + raise RuntimeError( + f"The {route} load shards through this {os.environ['WORLD_SIZE']}-rank launch's " + "process group, which is not set up; load inside distributed_run(), which sets it up." + ) if "RANK" not in os.environ: raise RuntimeError( f"The {route} load shards through a process group, even on one rank, which a " @@ -542,8 +548,10 @@ def load_pretrained( whole model, ``device_map`` defaulting to the rank's GPU. A plan with ``ep_size > 1`` raises: no load here shards the routed experts. On a distributed run - every rank loads the plan rank 0 was handed. A rank of a multi-rank launch whose process group - is not set up is refused on every route: set it up first, with + every rank loads the plan rank 0 was handed. A launch whose processes set no process group up (a + launcher that only exports ``WORLD_SIZE``) has each process load as a single process does, the + whole model with no plan to agree on; a plan for several ranks and a sharded load need the group, + so there they are refused: set it up first, with :func:`~modelopt.torch.utils.distributed.distributed_run`. The plan is recorded on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`), where calibration and export read it back. @@ -591,13 +599,15 @@ def load_pretrained( from modelopt.torch.models.hf import prepare_model_for_loading from modelopt.torch.utils import distributed as dist - if dist.is_distributed_launch() and not dist.is_initialized(): - # Without its group every collective below runs on this rank alone: each rank would plan - # and load the model as a single process's, and record a one-rank plan on a model the whole - # launch holds. + # A launcher can start several processes that set no process group up (hf_ptq without + # --use_fsdp2, say): each then loads as a single process does, the whole model with no plan to + # agree on, as before plans existed. A plan for several ranks needs the group. + alone = dist.is_distributed_launch() and not dist.is_initialized() + if alone and isinstance(distribution, DistributionPlan) and distribution.world_size > 1: raise RuntimeError( - f"This process is one of a {os.environ['WORLD_SIZE']}-rank launch whose process group " - "is not set up; set it up first, with distributed_run()." + f"{distribution!r} is for {distribution.world_size} ranks, whose process group this " + f"{os.environ['WORLD_SIZE']}-rank launch has not set up; load inside distributed_run(), " + "which sets it up." ) config = kwargs.get("config") config = config if isinstance(config, PretrainedConfig) else None @@ -617,7 +627,7 @@ def handed() -> DistributionPlan | None: # ranks loading different plans would build mismatched meshes. Rank 0 checks it before it sends, # so one it refuses raises on every rank instead of leaving them in the broadcast, and before # the checkpoint is fetched. - plan = dist.broadcast_from("distribution", handed) + plan = handed() if alone else dist.broadcast_from("distribution", handed) if plan is None and "hf_config" in kwargs: # The load builds from hf_config as given, so the plan must be the one it came with, not # one decided here from another config. @@ -635,7 +645,9 @@ def handed() -> DistributionPlan | None: ckpt_path = os.path.join(ckpt_path, subfolder) trust_remote_code = kwargs.get("trust_remote_code", False) hf_config = kwargs.get("hf_config") - if plan is None: + if plan is None and alone: + plan = DistributionPlan() + elif plan is None: plan, hf_config = resolve_distribution_plan( ckpt_path, config=config, diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 3c27a8da360..f6d1fc99d24 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -489,26 +489,47 @@ def test_a_single_process_fsdp2_load_needs_a_launcher(monkeypatch, serve, loader assert loaders == {} -@pytest.mark.parametrize( - "distribution", [None, DistributionPlan(dense="fsdp")], ids=["decided", "fsdp"] -) -def test_a_rank_of_a_launch_without_its_group_is_refused( - monkeypatch, serve, loaders, auto_class, tmp_path, distribution -): - """Without its group every collective runs on this rank alone, so each rank of the launch would - plan and load the model as a single process's, the whole model on every rank, and record a - one-rank plan on a model the whole launch holds.""" - serve(LlamaConfig(**_TINY)) +@pytest.fixture +def launch_without_its_group(monkeypatch): + """Rank 0 of an 8-rank launch whose processes set no process group up.""" monkeypatch.setenv("RANK", "0") monkeypatch.setenv("WORLD_SIZE", "8") monkeypatch.setattr(dist, "is_initialized", lambda: False) monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set a group up")) - with pytest.raises(RuntimeError, match="8-rank launch whose process group is not set up"): - load_pretrained( - str(tmp_path), auto_class if distribution is None else None, distribution=distribution - ) - assert auto_class.calls == [] + +def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( + launch_without_its_group, serve, loaders, auto_class, tmp_path +): + """A launcher can start several processes that never set a group up (hf_ptq without + --use_fsdp2): each loads the whole model as a single process does, with no collective to wait + in and no plan to agree on.""" + serve(LlamaConfig(**_TINY)) + + model = load_pretrained(str(tmp_path), auto_class) + + assert auto_class.calls == [{}] + assert distribution_plan(model) == DistributionPlan() + assert loaders == {} + + +@pytest.mark.parametrize( + ("distribution", "match"), + [ + (DistributionPlan(8), "is for 8 ranks, whose process group this 8-rank launch has not"), + (DistributionPlan(dense="fsdp"), "shards through this 8-rank launch's process group"), + ], + ids=["several-ranks", "sharded"], +) +def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( + launch_without_its_group, serve, loaders, tmp_path, distribution, match +): + """A plan for several ranks and a sharded load need the launch's group, which one rank cannot + set up for itself: refused before anything loads.""" + serve(LlamaConfig(**_TINY)) + + with pytest.raises(RuntimeError, match=match): + load_pretrained(str(tmp_path), distribution=distribution) assert loaders == {} From 6aad91365b62ac0b984faaf319ffac25a65b10cc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:12:38 +0000 Subject: [PATCH 38/41] Say when an automatic layout is what refuses a placement With neither a plan nor MODELOPT_FSDP_SIZE, rank 0's GPU memory decides whether the dense part is sharded, and the sharded load then refuses device_map, quantization_config or a dtype as if the caller had asked for it, so the same call loads on one GPU and fails on a smaller one with no word of why. The refusal now says the layout was decided from rank 0's memory and how to keep the load whole, and the docstring says which calls depend on it. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 29 ++++++++++++++++--- .../test_load_pretrained_distribution.py | 15 ++++++++++ 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ad7b1785d27..f0756de4a20 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -420,6 +420,13 @@ def record_unplaced_source_keys( } +def _whole_load_options(kwargs) -> list[str]: + """The load ``kwargs`` that place a whole model, which the FSDP2 load refuses.""" + names = set(kwargs) - _SHARDED_LOAD_OPTIONS[_FSDP2] - {"config", "dtype", "torch_dtype"} + dtypes = {name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")} + return sorted(names | dtypes) + + def _route(plan: DistributionPlan) -> str | None: """The sharded load ``plan`` takes, or ``None`` when every rank loads the whole model.""" from modelopt.torch.utils import distributed as dist @@ -558,10 +565,13 @@ def load_pretrained( A sharded load builds the model class and places it itself. Besides its own options it takes ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other keyword - rather than ignore it. ``hf_config`` is the config :func:`resolve_distribution_plan` returned with - the plan handed in as ``distribution``: the load builds from it as given, so without that plan - it is refused. A ``config`` passed for the load is the one the plan is decided from, and a - sharded load takes it as its ``hf_config``. + rather than ignore it. With neither a plan nor ``MODELOPT_FSDP_SIZE``, rank 0's GPU memory + decides whether the dense part is sharded, so a call that places the model (``device_map``, + ``quantization_config``, a ``dtype``) loads only where it fits one GPU: pass ``distribution=`` + or set ``MODELOPT_FSDP_SIZE=1`` to keep it whole. ``hf_config`` is the config + :func:`resolve_distribution_plan` returned with the plan handed in as ``distribution``: the load + builds from it as given, so without that plan it is refused. A ``config`` passed for the load is + the one the plan is decided from, and a sharded load takes it as its ``hf_config``. Whatever the placement: @@ -645,6 +655,7 @@ def handed() -> DistributionPlan | None: ckpt_path = os.path.join(ckpt_path, subfolder) trust_remote_code = kwargs.get("trust_remote_code", False) hf_config = kwargs.get("hf_config") + decided = plan is None and not alone if plan is None and alone: plan = DistributionPlan() elif plan is None: @@ -658,6 +669,16 @@ def handed() -> DistributionPlan | None: if _route(plan) is not None: kwargs.setdefault("hf_config", hf_config) route = _route(plan) + placed = _whole_load_options(kwargs) + if decided and route == _FSDP2 and env.fsdp_size().lower() == "auto" and placed: + # Rank 0's GPU memory chose to shard: say so, rather than refuse the options below as if + # the caller had asked for the sharded load. + raise TypeError( + f"{placed} place a whole model, but no plan or MODELOPT_FSDP_SIZE asked for a layout, " + f"and by rank 0's GPU memory the model does not fit one GPU, so {plan!r} shards it. " + "To load it whole on every rank and place it yourself, set MODELOPT_FSDP_SIZE=1 or " + f"pass distribution=DistributionPlan({plan.world_size}, dense='replicate')." + ) if route is not None and config is not None: # The sharded load builds the model from hf_config: the plan's copy of this one, or a copy # of it, since building writes to the config. diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index f6d1fc99d24..2c58dc0513f 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -434,6 +434,21 @@ def test_a_load_refuses_what_it_cannot_load( assert auto_class.calls == [] +def test_an_automatic_layout_says_so_when_it_refuses_a_placement( + run_as, serve, loaders, monkeypatch, tmp_path +): + """With neither a plan nor MODELOPT_FSDP_SIZE, rank 0's GPU memory decides to shard a model too + big for it, and the sharded load refuses options that place a whole model. The refusal says the + layout was decided rather than asked for, and how to keep the load whole.""" + run_as(2) + serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: 1) + + with pytest.raises(TypeError, match=r"\['device_map'\] place a whole model, but no plan"): + load_pretrained(str(tmp_path), device_map="auto") + assert loaders == {} + + def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched(run_as, monkeypatch, tmp_path): """Making the checkpoint local can mean downloading it: a plan the load refuses fails first.""" run_as(2) From d4ff352a83eea059457fcf92bfc9e62c623e63e5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:13:20 +0000 Subject: [PATCH 39/41] Name the route in the two sharded-load refusals that said FSDP2 Every other refusal of the sharded load names the route it is for; these two named FSDP2 literally. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f0756de4a20..632cd547534 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -492,13 +492,13 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) if getattr(hf_config, "quantization_config", None) is not None: # It reads the tensors as stored: quantized weights would load as if they were not. raise NotImplementedError( - f"The FSDP2 load cannot dequantize this {model_type} checkpoint; load it unsharded." + f"The {route} load cannot dequantize this {model_type} checkpoint; load it unsharded." ) if _auto_class(hf_config, trust_remote_code=trust_remote_code) is not AutoModelForCausalLM: # It builds through AutoModelForCausalLM, which refuses most VLMs' configs and, for some, # builds the language model alone, whose other towers the export would carry unquantized. raise NotImplementedError( - f"The FSDP2 load builds a causal LM, which this {model_type} checkpoint is not. To load " + f"The {route} load builds a causal LM, which this {model_type} checkpoint is not. To load " "it whole on every rank, pass its model class as auto_class, with MODELOPT_FSDP_SIZE=1 " f"or distribution=DistributionPlan({dist.size()}, dense='replicate')." ) From 385aa424b6c2cce4e0e0e064465778b0d5bb146e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:49:35 +0000 Subject: [PATCH 40/41] Refuse knobs asking to shard a launch that sets no group up A launch whose processes set no process group up has each process load the whole model, but MODELOPT_FSDP_SIZE naming the world or an MODELOPT_EP_DEGREE above 1 asks to shard it over the launch's ranks: loading whole there ignored the knob without a word, eight whole copies where it asked for one sharded. Such knobs are now refused before anything is fetched, with the way to set the group up; the default, auto, keeps the whole load. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 12 ++++++++++ .../test_load_pretrained_distribution.py | 23 +++++++++++++------ 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 632cd547534..054573107fe 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -619,6 +619,18 @@ def load_pretrained( f"{os.environ['WORLD_SIZE']}-rank launch has not set up; load inside distributed_run(), " "which sets it up." ) + if alone and distribution is None: + # The knobs ask to shard over the launch, rather than leave each process the whole model. + ep, world = env.ep_degree(), int(os.environ["WORLD_SIZE"]) + asked = [f"MODELOPT_EP_DEGREE={ep}"] if ep.isdigit() and int(ep) > 1 else [] + if _dense_from_env(world) == "fsdp": + asked.append(f"MODELOPT_FSDP_SIZE={env.fsdp_size()}") + if asked: + raise RuntimeError( + f"{' and '.join(asked)} ask to shard the model over this " + f"{os.environ['WORLD_SIZE']}-rank launch, whose process group is not set up; load " + "inside distributed_run(), which sets it up." + ) config = kwargs.get("config") config = config if isinstance(config, PretrainedConfig) else None if config is not None and "hf_config" in kwargs: diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 2c58dc0513f..0a1d66c63ce 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -529,19 +529,28 @@ def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( @pytest.mark.parametrize( - ("distribution", "match"), + ("distribution", "knobs", "match"), [ - (DistributionPlan(8), "is for 8 ranks, whose process group this 8-rank launch has not"), - (DistributionPlan(dense="fsdp"), "shards through this 8-rank launch's process group"), + (DistributionPlan(8), {}, "is for 8 ranks, whose process group this 8-rank launch has not"), + ( + DistributionPlan(dense="fsdp"), + {}, + "shards through this 8-rank launch's process group", + ), + (None, {"MODELOPT_FSDP_SIZE": "world"}, "MODELOPT_FSDP_SIZE=world ask to shard"), + (None, {"MODELOPT_EP_DEGREE": "4"}, "MODELOPT_EP_DEGREE=4 ask to shard"), ], - ids=["several-ranks", "sharded"], + ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob"], ) def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( - launch_without_its_group, serve, loaders, tmp_path, distribution, match + launch_without_its_group, serve, loaders, monkeypatch, tmp_path, distribution, knobs, match ): - """A plan for several ranks and a sharded load need the launch's group, which one rank cannot - set up for itself: refused before anything loads.""" + """A plan for several ranks, a sharded load, and knobs asking to shard over the launch need its + group, which one rank cannot set up for itself: refused before anything loads, rather than + leave each process the whole model the knobs did not ask for.""" serve(LlamaConfig(**_TINY)) + for name, value in knobs.items(): + monkeypatch.setenv(name, value) with pytest.raises(RuntimeError, match=match): load_pretrained(str(tmp_path), distribution=distribution) From 6d87a6a31996f604afcdc0de424a174849dfc478 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:49:48 +0000 Subject: [PATCH 41/41] Name the plan that works in the refusal of a groupless sharded load The one-rank FSDP2 plan that reaches it can never load on a launch of several ranks, so loading inside distributed_run() as advised only met a second refusal for the plan's world size. The refusal now names the plan for the launch. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 9 ++++++--- .../torch/utils/test_load_pretrained_distribution.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 054573107fe..2962ae1c774 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -506,10 +506,13 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) # FSDP2 shards through a process group even on one rank; a plan for more ranks than one # was refused without a group. A run (distributed_run) destroys the group on exit. if dist.is_distributed_launch(): - # One rank of a launch of several cannot set the group up for itself. + # One rank of a launch of several cannot set the group up for itself, and the one-rank + # plan that reaches here is no plan for the launch. + world = os.environ["WORLD_SIZE"] raise RuntimeError( - f"The {route} load shards through this {os.environ['WORLD_SIZE']}-rank launch's " - "process group, which is not set up; load inside distributed_run(), which sets it up." + f"The {route} load shards through this {world}-rank launch's process group, which " + "is not set up; load inside distributed_run(), which sets it up, with a plan for " + f"the launch: DistributionPlan({world}, dense='fsdp')." ) if "RANK" not in os.environ: raise RuntimeError( diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 0a1d66c63ce..775c637b4b5 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -535,7 +535,7 @@ def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( ( DistributionPlan(dense="fsdp"), {}, - "shards through this 8-rank launch's process group", + r"shards through this 8-rank launch's .* DistributionPlan\(8, dense='fsdp'\)", ), (None, {"MODELOPT_FSDP_SIZE": "world"}, "MODELOPT_FSDP_SIZE=world ask to shard"), (None, {"MODELOPT_EP_DEGREE": "4"}, "MODELOPT_EP_DEGREE=4 ask to shard"),