From 3d752188b65d5a2cd5cb1975854ffaba5ce917c6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 18:45:06 +0000 Subject: [PATCH 01/99] 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 89fb2a8b8d7..613fef67b6b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -177,6 +177,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 1967708c248..77fa00d86ff 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 @@ -428,17 +428,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 @@ -450,7 +623,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 @@ -471,7 +645,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 } @@ -481,21 +654,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 8a4a18b5f69d0dfcab3f86d36c98afa8fddc248a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 13 Oct 2026 23:07:12 +0000 Subject: [PATCH 02/99] 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 77fa00d86ff..2cd5a4b83a1 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -654,15 +654,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 9bf3418671ec1cdb39018d79ee683ca426c45d68 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 02:38:28 +0000 Subject: [PATCH 03/99] 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 2cd5a4b83a1..e6d3d6dd580 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -664,8 +664,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 76440d953bc3f7f06837de75bbf19312d538d004 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 03:08:45 +0000 Subject: [PATCH 04/99] 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 e6d3d6dd580..5ec6fa290a0 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -665,7 +665,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 05b015d6931afc7d7162563ee66f29bd8429a64a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 17:21:09 +0000 Subject: [PATCH 05/99] 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 613fef67b6b..ca67e778c26 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -177,7 +177,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 5ec6fa290a0..9d135588584 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -445,30 +445,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 @@ -480,21 +472,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) @@ -509,8 +503,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: @@ -526,7 +522,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: @@ -589,9 +585,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. @@ -604,8 +599,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 @@ -644,6 +639,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 @@ -654,18 +650,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. @@ -685,7 +685,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 5412517b32eed54e7bae15c0b0ddb4754e8795a7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 19:48:15 +0000 Subject: [PATCH 06/99] 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 ca67e778c26..afda1bba5f2 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -177,7 +177,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 9d135588584..7076ad4371a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -52,6 +52,7 @@ barrier, broadcast_state_dict, is_initialized, + set_distribution_plan, ) from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 @@ -445,84 +446,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: @@ -578,29 +516,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 @@ -650,12 +584,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. @@ -665,24 +599,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()}"}) @@ -695,6 +631,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( @@ -1030,6 +967,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) @@ -1121,6 +1061,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 b00273fb2bb2e74758c01d8df491ebe41e7baa89 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 20:26:10 +0000 Subject: [PATCH 07/99] 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/plugins/model_load_utils.py | 5 ++--- tests/gpu/torch/utils/test_distributed.py | 3 +++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 7076ad4371a..0009d5134e0 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -968,8 +968,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) @@ -1061,7 +1061,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) From c226e40a83bcb6d3f83845bb399e65e64d735097 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 20:55:15 +0000 Subject: [PATCH 08/99] 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 392b3d2807e..01e0634c93a 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -284,7 +284,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 @@ -369,7 +369,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 @@ -498,7 +498,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 f70010938cb8a2e4a52f769810f512bfb9d884c8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 22:58:38 +0000 Subject: [PATCH 09/99] 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 0009d5134e0..29cae17c5c0 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -491,8 +491,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: @@ -500,6 +498,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()), @@ -526,7 +528,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 cb26e7ec4cf6d32a13172281ad6713cbaf6bb690 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:21:50 +0000 Subject: [PATCH 10/99] 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 --- .../torch/utils/plugins/model_load_utils.py | 4 ++-- .../test_load_pretrained_distribution.py | 20 +++++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 29cae17c5c0..32cb2f75c5f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -53,6 +53,7 @@ broadcast_state_dict, is_initialized, set_distribution_plan, + size, ) from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 @@ -972,8 +973,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_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 23d6a2f767f7a975888a09b9acd8fb31f5461a13 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:43:44 +0000 Subject: [PATCH 11/99] 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/fsdp2_utils.py | 5 +++++ modelopt/torch/utils/plugins/model_load_utils.py | 3 +-- 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 ++++++----- 6 files changed, 28 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index 14db7266b5a..059d29abdc4 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -99,6 +99,11 @@ 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. + + 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. """ # Importing the quantization package registers every mode and plugin; only a wrap needs it. from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 32cb2f75c5f..6c450ed0e2b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -53,7 +53,6 @@ broadcast_state_dict, is_initialized, set_distribution_plan, - size, ) from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 @@ -973,7 +972,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 71241620222bb060e24e6d6374b41fb02df55705 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 17:30:20 +0000 Subject: [PATCH 12/99] 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 afda1bba5f2..9ab13a9c65a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,6 +43,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** @@ -177,7 +178,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 7211eb6e258cbc8f28e7019ce3658b074e44cf5c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 02:04:28 +0000 Subject: [PATCH 13/99] 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 363f931f62057b1fddd6de98857204928acff8ef Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:14:15 +0000 Subject: [PATCH 14/99] 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 88bf8c51f15151ba66e4361354006b0448d43b35 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 18:01:12 +0000 Subject: [PATCH 15/99] 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 a8f31f331a7c096cd5c4facbe8896eca052b0b78 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:52:50 +0000 Subject: [PATCH 16/99] 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 859a32556f397da28ab46d24bbb79b76e72e7b51 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 00:44:17 +0000 Subject: [PATCH 17/99] 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 6c450ed0e2b..02042f36ff5 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -590,7 +590,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 ca1aefce311e79eab163f5066b8761ff28b41115 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 01:09:31 +0000 Subject: [PATCH 18/99] 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 --- .../unit/torch/utils/test_model_load_utils.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index ca9e71b97f5..b0dba2e0089 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -393,6 +393,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 91a2b2554566dd91e8e4efcd9e9a3bd7fbc3f6b1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 01:43:29 +0000 Subject: [PATCH 19/99] 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 02042f36ff5..025ad0fdabb 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -463,7 +463,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 @@ -492,7 +492,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( @@ -618,7 +620,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 0e29bf0ea2d95fc7e24c4a13329b98df65b708de Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 00:45:23 +0000 Subject: [PATCH 20/99] 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 025ad0fdabb..ca5bea85a88 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 @@ -544,7 +545,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: @@ -590,10 +592,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) @@ -611,10 +617,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 @@ -1244,6 +1253,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, @@ -1287,6 +1297,9 @@ def resolve_distribution_plan( returns no path, so a caller that loads the checkpoint resolves it first, with :func:`ensure_local_checkpoint`, and passes the local directory: the plan is then decided for the snapshot the load reads. + + ``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. """ # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather # than leave the others waiting in a broadcast below for a rank that already raised. @@ -1302,14 +1315,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 99d850f79daff963cbb9aefc11adea055cb236df Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 01:11:18 +0000 Subject: [PATCH 21/99] 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 2c71259eb79d0e906c2b1f92803a4b7a52fa5bd1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:09:59 +0000 Subject: [PATCH 22/99] 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 ca5bea85a88..a975553c18b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -544,9 +544,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: @@ -596,6 +598,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 4007acf5cd0c553ca6e1982370381ab9e1c00a56 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:10:04 +0000 Subject: [PATCH 23/99] 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 a975553c18b..27d226d7a78 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -533,7 +533,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 e1bc5a9ad8c6e8cfbbffdfb484340caecc2c05ff Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:02:01 +0000 Subject: [PATCH 24/99] 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 27d226d7a78..fbbfeb17424 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -479,9 +479,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 @@ -504,6 +510,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, @@ -531,9 +543,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 @@ -628,9 +641,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 0ee033192948542040d48c78e5682d3ff322a6a9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:06:57 +0000 Subject: [PATCH 25/99] 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 fbbfeb17424..8fdb49450d2 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -452,8 +452,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( @@ -477,6 +483,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 = ( @@ -516,6 +528,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 5b018fdb87bba0a916367e8a007c691b799b6599 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:46:19 +0000 Subject: [PATCH 26/99] 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 --- .../unit/torch/utils/test_model_load_utils.py | 24 ------------------- 1 file changed, 24 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b0dba2e0089..ca9e71b97f5 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -393,30 +393,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 3c37f8371ce3eec0a3964ee25f3b0257732fa474 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:25:33 +0000 Subject: [PATCH 27/99] 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 --- .../torch/utils/plugins/model_load_utils.py | 81 ++++++++++++------- .../test_load_pretrained_distribution.py | 58 +++++++++++-- 2 files changed, 102 insertions(+), 37 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 8fdb49450d2..170170111ff 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -46,6 +46,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 distributed as dist from modelopt.torch.utils import env from modelopt.torch.utils.distributed import ( @@ -557,7 +558,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 @@ -587,10 +589,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 @@ -630,33 +633,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 @@ -1213,7 +1221,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.""" # What the caller set the EP degree with, for a refusal to name it. @@ -1229,6 +1237,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, @@ -1297,6 +1314,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, @@ -1341,6 +1359,11 @@ def resolve_distribution_plan( :func:`ensure_local_checkpoint`, and passes the local directory: the plan is then decided for the snapshot the load reads. + ``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. """ @@ -1377,7 +1400,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 1bebb908829309dcac89f916ef091000f4a9f25f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 02:07:23 +0000 Subject: [PATCH 28/99] 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 +++++----- tests/unit/torch/utils/test_model_load_utils.py | 11 +++++++++++ 4 files changed, 27 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9ab13a9c65a..150c8e918b3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +43,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 170170111ff..03e159a3c90 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -674,6 +674,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 @@ -1246,6 +1252,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, @@ -1365,7 +1374,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. """ # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather # than leave the others waiting in a broadcast below for a rank that already raised. 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_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index ca9e71b97f5..5330fd0cfad 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -735,6 +735,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 be069c4b49cd925b5c7fc9753be06faa2058ca89 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 03:54:44 +0000 Subject: [PATCH 29/99] 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 03e159a3c90..08140c7f199 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -520,6 +520,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. @@ -1340,7 +1348,9 @@ def resolve_distribution_plan( every later step reads it back from there (:func:`~modelopt.torch.utils.distributed.distribution_plan`). Returns ``(plan, hf_config)``, ``hf_config`` being this rank's own copy of the config (rank 0 decides the plan from its copy), - for a sharded load to build 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 6f04f0ba22443ba3e416e17b79efc6e8e531ef90 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:43 +0000 Subject: [PATCH 30/99] 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 08140c7f199..0dfb0291c9c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -453,15 +453,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 " @@ -537,12 +529,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, @@ -581,7 +567,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. @@ -628,6 +616,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 0d3cf2cb7f6800ffb4c0fb272abc43602b3b1942 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:50 +0000 Subject: [PATCH 31/99] 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 0dfb0291c9c..137bc5fcb88 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1256,6 +1256,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 d4f7d45173c0c86ddf406008ba223ad5ba20d90c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:15:21 +0000 Subject: [PATCH 32/99] 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 137bc5fcb88..2a73ac5cc52 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -453,7 +453,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 c13cd3d27225c698e803b46413c7b32062ce756f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 17:58:06 +0000 Subject: [PATCH 33/99] 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 2a73ac5cc52..3ac3bf93a24 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -476,14 +476,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 f657b30b1cff8188217c4e02054e86553add84f0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:03:42 +0000 Subject: [PATCH 34/99] 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 | 1 - 1 file changed, 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3ac3bf93a24..d5ce4cd1292 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -46,7 +46,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 distributed as dist from modelopt.torch.utils import env from modelopt.torch.utils.distributed import ( From b11a56ba631561ac09ea7b0b8a951975a4b96dd3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:06:41 +0000 Subject: [PATCH 35/99] 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 d5ce4cd1292..2ea105a316f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -626,15 +626,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: @@ -651,17 +642,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 d4bbf654b5505f75551224cabf4035419fc05468 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:09:09 +0000 Subject: [PATCH 36/99] 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 2ea105a316f..b1e2522e93b 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -525,6 +525,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 " @@ -569,8 +575,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. @@ -618,13 +626,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 @@ -644,7 +654,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. @@ -662,7 +672,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 63ebfde445534e0203a03660d097a26ffe165afb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:12:38 +0000 Subject: [PATCH 37/99] 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 b1e2522e93b..c9e8d52ef06 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -447,6 +447,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 @@ -585,10 +592,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: @@ -672,6 +682,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: @@ -685,6 +696,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 7c0bd7b8495368a8593a049d1bb3e72516a41d68 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:13:20 +0000 Subject: [PATCH 38/99] 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 c9e8d52ef06..d73988ece35 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -519,13 +519,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 f01b7524e55dc0354d37d59e36fa80a4ad73c227 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:49:35 +0000 Subject: [PATCH 39/99] 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 d73988ece35..ccff33da076 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -646,6 +646,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 c136624e087c3ad7a39fee072ab7a39b25e4aedd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:49:48 +0000 Subject: [PATCH 40/99] 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 ccff33da076..138c27869f3 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -533,10 +533,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"), From 8f900fc9289248b584cf3266ca0476319d851053 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:05:16 +0000 Subject: [PATCH 41/99] Read the EP knob as the plan decision does in a launch without its group The launch-without-group check matched MODELOPT_EP_DEGREE with str.isdigit(), so a value that is not a degree ("eigth") was dropped there, while the plan decision refuses it with a ValueError once a group exists. It now reads the knob through _ep_degree, as the decision does. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 5 ++- .../test_load_pretrained_distribution.py | 37 +++++++++++++++---- 2 files changed, 32 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 138c27869f3..bbc296c662f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -651,8 +651,9 @@ def load_pretrained( ) 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 [] + # Read as the plan decision reads them, so a value that is not one raises here too. + ep, world = _ep_degree(None), int(os.environ["WORLD_SIZE"]) + asked = [f"MODELOPT_EP_DEGREE={ep}"] if isinstance(ep, int) and ep > 1 else [] if _dense_from_env(world) == "fsdp": asked.append(f"MODELOPT_FSDP_SIZE={env.fsdp_size()}") if asked: diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 775c637b4b5..ee2013b5535 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -529,30 +529,51 @@ def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( @pytest.mark.parametrize( - ("distribution", "knobs", "match"), + ("distribution", "knobs", "error", "match"), [ - (DistributionPlan(8), {}, "is for 8 ranks, whose process group this 8-rank launch has not"), + ( + DistributionPlan(8), + {}, + RuntimeError, + "is for 8 ranks, whose process group this 8-rank launch has not", + ), ( DistributionPlan(dense="fsdp"), {}, + RuntimeError, 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"), + ( + None, + {"MODELOPT_FSDP_SIZE": "world"}, + RuntimeError, + "MODELOPT_FSDP_SIZE=world ask to shard", + ), + (None, {"MODELOPT_EP_DEGREE": "4"}, RuntimeError, "MODELOPT_EP_DEGREE=4 ask to shard"), + (None, {"MODELOPT_EP_DEGREE": "eigth"}, ValueError, "MODELOPT_EP_DEGREE='eigth'"), ], - ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob"], + ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob", "ep-knob-not-a-degree"], ) def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( - launch_without_its_group, serve, loaders, monkeypatch, tmp_path, distribution, knobs, match + launch_without_its_group, + serve, + loaders, + monkeypatch, + tmp_path, + distribution, + knobs, + error, + match, ): """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.""" + leave each process the whole model the knobs did not ask for. A knob that is not one is refused + as the plan decision refuses it, not dropped.""" serve(LlamaConfig(**_TINY)) for name, value in knobs.items(): monkeypatch.setenv(name, value) - with pytest.raises(RuntimeError, match=match): + with pytest.raises(error, match=match): load_pretrained(str(tmp_path), distribution=distribution) assert loaders == {} From 6558b26aaaee45eec422129ec5b0a8334cb801a0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:05:33 +0000 Subject: [PATCH 42/99] Say the automatic layout weighs the model in the checkpoint's dtype With no plan, the dense layout is decided from the meta model built in the checkpoint's own dtype, not one the load asks for, so a load that widens it should be handed its plan. 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 bbc296c662f..8c6aa224a37 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -598,7 +598,9 @@ def load_pretrained( 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 + or set ``MODELOPT_FSDP_SIZE=1`` to keep it whole. The model is weighed in the checkpoint's own + dtype (its config's, else bfloat16), not one the load asks for: hand a load that widens it, + ``dtype=torch.float32`` say, its plan. ``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``. From 8c2fab573d8648fc83fc2fb684aea98c4e95c706 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:09:11 +0000 Subject: [PATCH 43/99] Parse the EP knob in the groupless check itself #2746 makes _ep_degree resolve an automatic degree from the model's spec, which this check runs before: called here, its later signature would leave the check without the model type it needs. The check now parses the knob itself, with the plan decision's spellings and message, so a value that is not a degree ("eigth", "-1") is refused and a padded one (" 4") counts as the degree it is. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 8c6aa224a37..78766eb8d10 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -652,10 +652,12 @@ def load_pretrained( "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. - # Read as the plan decision reads them, so a value that is not one raises here too. - ep, world = _ep_degree(None), int(os.environ["WORLD_SIZE"]) - asked = [f"MODELOPT_EP_DEGREE={ep}"] if isinstance(ep, int) and ep > 1 else [] + # The knobs ask to shard over the launch, rather than leave each process the whole model. A + # value that is not an EP degree is refused, as the plan decision refuses it. + ep, world = env.ep_degree().strip(), int(os.environ["WORLD_SIZE"]) + if ep.lower() not in ("auto", "0") and not ep.isdigit(): + raise ValueError(f"MODELOPT_EP_DEGREE={ep!r}: expected auto, 0 (auto) or an EP degree.") + 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: From cc0ae8296a2bff32e4bb4a64fe3863126575f9be Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:53:56 +0000 Subject: [PATCH 44/99] Refuse a groupless launch's sharded load before fetching the checkpoint A one-rank plan for a sharded load, handed in a multi-rank launch whose group is not set up, was refused in _load_sharded, after the checkpoint was made local and its config read: under torchrun with a Hub ID, the whole snapshot downloaded first. The refusal now sits with the other groupless checks, which run before anything is fetched; the test checks that nothing is. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 22 +++++++++---------- .../test_load_pretrained_distribution.py | 5 ++++- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 78766eb8d10..9119aef586c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -530,17 +530,8 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) 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. - if dist.is_distributed_launch(): - # 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 {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')." - ) + # FSDP2 shards through a process group even on one rank; a launch without its group was + # refused before the load. 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 " @@ -651,6 +642,15 @@ 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 isinstance(distribution, DistributionPlan) and _route(distribution) is not None: + # One rank of a launch of several cannot set the group up for itself, and a one-rank plan + # for a sharded load is no plan for the launch. + launch = os.environ["WORLD_SIZE"] + raise RuntimeError( + f"The {_route(distribution)} load shards through this {launch}-rank launch's process " + "group, which is not set up; load inside distributed_run(), which sets it up, with a " + f"plan for the launch: DistributionPlan({launch}, dense='fsdp')." + ) if alone and distribution is None: # The knobs ask to shard over the launch, rather than leave each process the whole model. A # value that is not an EP degree is refused, as the plan decision refuses it. diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index ee2013b5535..cd0bd663959 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -566,12 +566,15 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( match, ): """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 + group, which one rank cannot set up for itself: refused before anything is fetched, rather than leave each process the whole model the knobs did not ask for. A knob that is not one is refused as the plan decision refuses it, not dropped.""" serve(LlamaConfig(**_TINY)) for name, value in knobs.items(): monkeypatch.setenv(name, value) + monkeypatch.setattr( + model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") + ) with pytest.raises(error, match=match): load_pretrained(str(tmp_path), distribution=distribution) From be5cde4a5f12842d20a9bf1e092cc881084a39a0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 22:07:11 +0000 Subject: [PATCH 45/99] Point a model without a plan at load_pretrained again, now that it records one load_pretrained records the plan it loads by from this PR on, so the advice for a model with no plan names it beside set_distribution_plan. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index e82257a3132..7febd2fe36a 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -355,10 +355,10 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non """Record that ``model`` is distributed as ``plan``, on it and on every submodule it has now. Code handed a submodule, such as a VLM's language model, then reads the same plan as code handed - the whole model; a module added afterwards records none. Code that distributes a model records - its plan here before calibrating or exporting the model. The plan says how the model is placed - over the run; the :class:`ParallelState` a quantized module holds names the process groups its - collectives run on. + the whole model; a module added afterwards records none. ``load_pretrained`` records the plan it + loads by once the model is built; code that distributes a model itself records it here before + calibrating or exporting the model. The plan says how the model is placed over the run; the + :class:`ParallelState` a quantized module holds names the process groups its collectives run on. A model is distributed once, and its plan recorded once: recording a plan on a model that, or any of whose submodules, already records one raises, the same plan included, and leaves the @@ -403,9 +403,10 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: if sharded or size() > 1: how = "is FSDP2-sharded" if sharded else f"is held in a run of {size()} ranks" raise RuntimeError( - f"{type(model).__name__} {how}, but has no distribution plan recorded. Record how it " - "is distributed with set_distribution_plan(); a module added to a model after its plan " - "was recorded has none, so read the plan off that model." + f"{type(model).__name__} {how}, but has no distribution plan recorded. Load it with " + "load_pretrained, or record how it is distributed with set_distribution_plan(); a " + "module added to a model after its plan was recorded has none, so read the plan off " + "that model." ) return DistributionPlan() From 19b98a3ddb57153a92e950d94ff040f873674cdd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 23:04:48 +0000 Subject: [PATCH 46/99] Refuse hf_config with a whole-model plan before the checkpoint is fetched Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 13 +++++------ .../test_load_pretrained_distribution.py | 23 ++++++++++++++----- 2 files changed, 23 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 9119aef586c..d3fdcf39a38 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -685,12 +685,13 @@ def handed() -> DistributionPlan | None: # so one it refuses raises on every rank instead of leaving them in the broadcast, and before # the checkpoint is fetched. 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. + if "hf_config" in kwargs and (plan is None or _route(plan) is None): + # Only a sharded load builds from hf_config, as given, so the plan must be the one it came + # with, not one decided here from another config nor one that loads the model whole. 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 is the config resolve_distribution_plan returned with a plan that shards " + "the model: 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 @@ -750,8 +751,6 @@ def handed() -> DistributionPlan | None: # 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 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. diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index cd0bd663959..ed3064364f4 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -252,7 +252,7 @@ def test_hf_config_is_refused_without_the_plan_it_came_with( run_as(2, rank=1) serve(LlamaConfig(**_TINY)) - with pytest.raises(TypeError, match="returned with a plan: pass that plan"): + with pytest.raises(TypeError, match="returned with a plan that shards the model: pass that"): load_pretrained(str(tmp_path), hf_config=LlamaConfig(**{**_TINY, "num_hidden_layers": 2})) assert loaders == {} @@ -394,7 +394,7 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( "causal-lm", True, TypeError, - "hf_config applies to a sharded load only", + "returned with a plan that shards the model", ), ], ids=[ @@ -449,15 +449,26 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( 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.""" +@pytest.mark.parametrize( + ("distribution", "options", "error", "match"), + [ + (DistributionPlan(4), {}, ValueError, "is for 4 ranks"), + (DistributionPlan(2), {"hf_config": LlamaConfig(**_TINY)}, TypeError, "plan that shards"), + ], + ids=["another-world", "hf-config-on-a-whole-load"], +) +def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched( + run_as, monkeypatch, tmp_path, distribution, options, error, match +): + """Making the checkpoint local can mean downloading it: a plan the load refuses fails first, as + does hf_config with a plan that loads the model whole.""" 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)) + with pytest.raises(error, match=match): + load_pretrained(str(tmp_path), distribution=distribution, **options) @pytest.mark.parametrize("quantized", [False, True]) From 8b75662840374252c135c62e4fad924d3fcca6b2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:08:01 +0000 Subject: [PATCH 47/99] Refuse a handed sharding plan's unusable load options before the checkpoint is fetched The arguments alone decide whether the FSDP2 load refuses auto_class, a dtype or an option that places a whole model, but those refusals ran inside the load, after ensure_local_checkpoint had downloaded the checkpoint. With the plan handed in, the load is known up front, so the same checks now run before the fetch, as hf_config's does; a decided plan still meets them in the load. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 32 +++++++++++++++---- .../test_load_pretrained_distribution.py | 11 +++++-- 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d3fdcf39a38..461e2e92de0 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -471,9 +471,8 @@ 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, 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 +def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: + """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide.""" from modelopt.torch.utils import distributed as dist if auto_class not in (None, AutoModelForCausalLM): @@ -481,17 +480,16 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) f"The {route} load builds its model class itself; auto_class={auto_class.__name__} " "does not apply to it." ) - kwargs = dict(kwargs) # 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) + value = kwargs.get(name) 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]) + unknown = sorted(set(kwargs) - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) if unknown: whole = ( " To place the model yourself, have every rank load it whole: MODELOPT_FSDP_SIZE=1, or " @@ -503,6 +501,15 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) f"{unknown} do not apply to the {route} load, which places the model itself; it " f"takes {sorted(_SHARDED_LOAD_OPTIONS[route])}.{whole}" ) + + +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 + + _refuse_sharded_load_options(route, auto_class, kwargs) + kwargs = {k: v for k, v in kwargs.items() if k not in ("dtype", "torch_dtype")} 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. @@ -697,6 +704,19 @@ def handed() -> DistributionPlan | None: 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 {} + if plan is not None and _route(plan) is not None: + # A handed plan names its load before anything is fetched, and the arguments alone decide + # what that load refuses. A config passed for it becomes its hf_config; subfolder is where + # the checkpoint is. + _refuse_sharded_load_options( + _route(plan), + auto_class, + { + k: v + for k, v in kwargs.items() + if k != "subfolder" and not (k == "config" and config is not None) + }, + ) _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) subfolder = kwargs.pop("subfolder", None) if subfolder: diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index ed3064364f4..3b553cc06ef 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -454,14 +454,21 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( [ (DistributionPlan(4), {}, ValueError, "is for 4 ranks"), (DistributionPlan(2), {"hf_config": LlamaConfig(**_TINY)}, TypeError, "plan that shards"), + ( + DistributionPlan(2, dense="fsdp"), + {"device_map": "auto"}, + TypeError, + r"\['device_map'\] do not apply to the FSDP2 load", + ), ], - ids=["another-world", "hf-config-on-a-whole-load"], + ids=["another-world", "hf-config-on-a-whole-load", "whole-load-option-on-a-sharded-plan"], ) def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched( run_as, monkeypatch, tmp_path, distribution, options, error, match ): """Making the checkpoint local can mean downloading it: a plan the load refuses fails first, as - does hf_config with a plan that loads the model whole.""" + do hf_config with a plan that loads the model whole and an option a sharded plan's load + refuses.""" run_as(2) monkeypatch.setattr( model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") From d420344ed915d2623daa3bea4dd58ef0c8bcf1bd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 08:03:20 +0000 Subject: [PATCH 48/99] Drop the load routing's local imports of the distributed helpers The module imports modelopt.torch.utils.distributed at the top now, so the four function-local imports of it bound the same module again. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 461e2e92de0..50b45dbdadd 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -456,8 +456,6 @@ def _whole_load_options(kwargs) -> list[str]: 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 - if plan.world_size != dist.size(): raise ValueError( f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." @@ -473,8 +471,6 @@ def _route(plan: DistributionPlan) -> str | None: def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide.""" - 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__} " @@ -506,7 +502,6 @@ def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: 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 _refuse_sharded_load_options(route, auto_class, kwargs) kwargs = {k: v for k, v in kwargs.items() if k not in ("dtype", "torch_dtype")} @@ -637,7 +632,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 - from modelopt.torch.utils import distributed as dist # 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 From ee93eb8cba1928beede6ebcace966585210346ae Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 08:03:29 +0000 Subject: [PATCH 49/99] Pin the auto_class and dtype refusals before the checkpoint is fetched The handed-plan test pinned only the unknown-option arm of the pre-fetch refusal. The auto_class and dtype rows each fail with their arm left out of the pre-fetch call, where the refusal would come only after the download. Signed-off-by: Shengliang Xu --- .../test_load_pretrained_distribution.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 3b553cc06ef..3c8ed95b137 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -460,8 +460,26 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( TypeError, r"\['device_map'\] do not apply to the FSDP2 load", ), + ( + DistributionPlan(2, dense="fsdp"), + {"auto_class": AutoModel}, + NotImplementedError, + "auto_class=AutoModel does not apply", + ), + ( + DistributionPlan(2, dense="fsdp"), + {"dtype": torch.float32}, + TypeError, + "keeps their dtype", + ), + ], + ids=[ + "another-world", + "hf-config-on-a-whole-load", + "whole-load-option-on-a-sharded-plan", + "auto-class-on-a-sharded-plan", + "dtype-on-a-sharded-plan", ], - ids=["another-world", "hf-config-on-a-whole-load", "whole-load-option-on-a-sharded-plan"], ) def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched( run_as, monkeypatch, tmp_path, distribution, options, error, match From 3dee145dbcd3c497619c6c7b590877a0b9487f98 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:03:57 +0000 Subject: [PATCH 50/99] Send the load-routing tests' plan broadcast from the rank deciding it main's broadcast_from now picks its sending rank by modelopt's dist.rank() rather than torch.distributed.get_rank(). The process_group fixture sets dist.rank() to 1 for the FSDP2 load, so a test acting as its group's first rank no longer sent the plan, and decoded nothing. run_as now runs the real broadcast with this rank as the sender. Signed-off-by: Shengliang Xu --- .../torch/utils/test_load_pretrained_distribution.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 3c8ed95b137..39f4e88a8a0 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -71,6 +71,15 @@ def run_as(monkeypatch): 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) + broadcast_from = dist.broadcast_from + + def broadcast_as_the_first(label, fn, src=0): + # broadcast_from sends from dist.rank(), which process_group sets for the FSDP2 load. + with monkeypatch.context() as sending: + sending.setattr(dist, "rank", lambda group=None: src) + return broadcast_from(label, fn, src) + + monkeypatch.setattr(dist, "broadcast_from", broadcast_as_the_first) def enter(world_size, rank=0): monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1) From e1a579a76e1a4cbe43b2e6583dd3b6ce31bf1839 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:23:24 +0000 Subject: [PATCH 51/99] Move the load-routing tests into the files of the module they cover load_pretrained's routing tests move from test_load_pretrained_distribution.py into tests/unit/torch/utils/test_model_load_utils.py and its GPU tests into tests/gpu/torch/utils/test_model_load_utils.py, as CONTRIBUTING asks of a test file. They share the plan decision's fixtures there; run_as now also sends a broadcast as its source rank and answers the rank's local GPU, which only the load's tests need. Test bodies are unchanged. Signed-off-by: Shengliang Xu --- .../test_load_pretrained_distribution.py | 91 --- .../gpu/torch/utils/test_model_load_utils.py | 75 +- .../test_load_pretrained_distribution.py | 645 ------------------ .../unit/torch/utils/test_model_load_utils.py | 575 +++++++++++++++- 4 files changed, 641 insertions(+), 745 deletions(-) delete mode 100644 tests/gpu/torch/utils/test_load_pretrained_distribution.py delete mode 100644 tests/unit/torch/utils/test_load_pretrained_distribution.py diff --git a/tests/gpu/torch/utils/test_load_pretrained_distribution.py b/tests/gpu/torch/utils/test_load_pretrained_distribution.py deleted file mode 100644 index 7f9eb7a24f3..00000000000 --- a/tests/gpu/torch/utils/test_load_pretrained_distribution.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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 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 -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 DistributionPlan, distribution_plan -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): - """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}") - 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 - dense = "replicate" if expected == "replicated" else "fsdp" - assert distribution_plan(model) == DistributionPlan(size, dense=dense) - finally: - for knob in knobs: - del os.environ[knob] - - -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/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index 05b524fb19a..54b96461485 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GPU/distributed tests for the FSDP2 load path and its helpers.""" +"""GPU/distributed tests for the loads: ``load_pretrained`` by its plan, the FSDP2 load, and their +helpers.""" import json import os @@ -24,14 +25,19 @@ from _test_utils.torch.transformers_models import create_tiny_llama_dir from safetensors.torch import load_file from torch.distributed.tensor import DTensor +from transformers import AutoModelForCausalLM, LlamaConfig from modelopt.torch.export.unified_export_hf import export_hf_checkpoint from modelopt.torch.utils.distributed import ( DistributionPlan, broadcast_state_dict, + distribution_plan, set_distribution_plan, ) -from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 +from modelopt.torch.utils.plugins.model_load_utils import ( + load_pretrained, + parallel_load_and_prepare_fsdp2, +) VOCAB_SIZE = 64 @@ -177,3 +183,68 @@ def test_carry_over_unplaced_weights(dist_workers, tmp_path): orphan_prefix=f"model.layers.{orphan_idx}.", ) ) + + +# ------------------------------------------------------------------ load_pretrained, left to decide + +_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): + """Load on every rank of the run with the plan left to decide, and check what each rank holds. + + Left to decide, ``load_pretrained`` 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. + """ + device = torch.device(f"cuda:{rank}") + 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 + dense = "replicate" if expected == "replicated" else "fsdp" + assert distribution_plan(model) == DistributionPlan(size, dense=dense) + finally: + for knob in knobs: + del os.environ[knob] + + +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 deleted file mode 100644 index 39f4e88a8a0..00000000000 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ /dev/null @@ -1,645 +0,0 @@ -# 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. - -"""Which load ``load_pretrained`` routes a distribution plan to, and the plan it records. - -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, LlamaForCausalLM, Qwen3MoeConfig - -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 - -_TINY = { - "hidden_size": 16, - "intermediate_size": 32, - "num_hidden_layers": 1, - "num_attention_heads": 2, - "num_key_value_heads": 1, - "vocab_size": 32, -} - - -@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) - - -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, 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) - broadcast_from = dist.broadcast_from - - def broadcast_as_the_first(label, fn, src=0): - # broadcast_from sends from dist.rank(), which process_group sets for the FSDP2 load. - with monkeypatch.context() as sending: - sending.setattr(dist, "rank", lambda group=None: src) - return broadcast_from(label, fn, src) - - monkeypatch.setattr(dist, "broadcast_from", broadcast_as_the_first) - - 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) - - return enter - - -@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) - - -# ----------------------------------------------------------------------------------- the loads - - -@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( - 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.""" - run_as(4, rank=1) - serve(LlamaConfig(**_TINY)) - kwargs = {} if device_map is None else {"device_map": device_map} - - 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( - "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, 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). - 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( - 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, local_files_only=True, **dtype - ) - - 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 [read.get("local_files_only") for read in reads] == [True] - assert distribution_plan(model) == plan - assert loaders == {} - - -def test_a_load_it_decides_reads_the_config_once( - 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") - run_as(2, rank=1) - config, reads = LlamaConfig(**_TINY), [] - monkeypatch.setattr( - transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a) or config - ) - - model = load_pretrained(str(tmp_path)) - - assert len(reads) == 1 - assert loaders["fsdp2"]["hf_config"] is config - 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, 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. 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 = [] - 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, **plan) - - 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") - - -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 that shards the model: pass that"): - load_pretrained(str(tmp_path), hf_config=LlamaConfig(**{**_TINY, "num_hidden_layers": 2})) - 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, 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: - 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) - 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)) - assert loaders == {} - assert auto_class.calls == [] - - -def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed): - model, refused = None, None - with bounded_collectives(): - 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) - - -@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), handed=handed) - - -@pytest.mark.parametrize( - ("change", "error", "match"), - [ - ({"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"), - ( - {"kwargs": {"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}}, - TypeError, - "config and hf_config both name", - ), - ], - 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 -): - """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"): - 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=DistributionPlan(2, dense="fsdp"), - **change.get("kwargs", {}), - ) - assert loaders == {} - - -@pytest.mark.parametrize( - ("distribution", "checkpoint", "whole_model_kwargs", "error", "match"), - [ - ( - DistributionPlan(4, dense="fsdp"), - "causal-lm", - False, - ValueError, - "is for 4 ranks, but this run has 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, - "returned with a plan that shards the model", - ), - ], - 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, - serve, - loaders, - auto_class, - process_group, - tmp_path, - distribution, - checkpoint, - whole_model_kwargs, - error, - match, -): - """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would - 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(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 {} - - with pytest.raises(error, match=match): - load_pretrained(str(tmp_path), distribution=distribution, **whole) - assert loaders == {} - 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 == {} - - -@pytest.mark.parametrize( - ("distribution", "options", "error", "match"), - [ - (DistributionPlan(4), {}, ValueError, "is for 4 ranks"), - (DistributionPlan(2), {"hf_config": LlamaConfig(**_TINY)}, TypeError, "plan that shards"), - ( - DistributionPlan(2, dense="fsdp"), - {"device_map": "auto"}, - TypeError, - r"\['device_map'\] do not apply to the FSDP2 load", - ), - ( - DistributionPlan(2, dense="fsdp"), - {"auto_class": AutoModel}, - NotImplementedError, - "auto_class=AutoModel does not apply", - ), - ( - DistributionPlan(2, dense="fsdp"), - {"dtype": torch.float32}, - TypeError, - "keeps their dtype", - ), - ], - ids=[ - "another-world", - "hf-config-on-a-whole-load", - "whole-load-option-on-a-sharded-plan", - "auto-class-on-a-sharded-plan", - "dtype-on-a-sharded-plan", - ], -) -def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched( - run_as, monkeypatch, tmp_path, distribution, options, error, match -): - """Making the checkpoint local can mean downloading it: a plan the load refuses fails first, as - do hf_config with a plan that loads the model whole and an option a sharded plan's load - refuses.""" - run_as(2) - monkeypatch.setattr( - model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") - ) - - with pytest.raises(error, match=match): - load_pretrained(str(tmp_path), distribution=distribution, **options) - - -@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"} - 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)) - # 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_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 == {} - - -@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")) - - -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", "knobs", "error", "match"), - [ - ( - DistributionPlan(8), - {}, - RuntimeError, - "is for 8 ranks, whose process group this 8-rank launch has not", - ), - ( - DistributionPlan(dense="fsdp"), - {}, - RuntimeError, - r"shards through this 8-rank launch's .* DistributionPlan\(8, dense='fsdp'\)", - ), - ( - None, - {"MODELOPT_FSDP_SIZE": "world"}, - RuntimeError, - "MODELOPT_FSDP_SIZE=world ask to shard", - ), - (None, {"MODELOPT_EP_DEGREE": "4"}, RuntimeError, "MODELOPT_EP_DEGREE=4 ask to shard"), - (None, {"MODELOPT_EP_DEGREE": "eigth"}, ValueError, "MODELOPT_EP_DEGREE='eigth'"), - ], - ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob", "ep-knob-not-a-degree"], -) -def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( - launch_without_its_group, - serve, - loaders, - monkeypatch, - tmp_path, - distribution, - knobs, - error, - match, -): - """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 is fetched, rather than - leave each process the whole model the knobs did not ask for. A knob that is not one is refused - as the plan decision refuses it, not dropped.""" - serve(LlamaConfig(**_TINY)) - for name, value in knobs.items(): - monkeypatch.setenv(name, value) - monkeypatch.setattr( - model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") - ) - - with pytest.raises(error, match=match): - load_pretrained(str(tmp_path), distribution=distribution) - assert loaders == {} - - -# --------------------------------------------------------------------------------- fsdp2_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.""" - 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(dist, "size", lambda group=None: 4) - - model_load_utils.fsdp2_wrap(model) - - with pytest.raises(RuntimeError, match="no distribution plan recorded"): - distribution_plan(model) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 5330fd0cfad..cd0653667cc 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -15,12 +15,14 @@ """Tests for ``modelopt.torch.utils.plugins.model_load_utils``. -Nothing here needs a GPU. The tests of ``resolve_distribution_plan`` answer the process group's rank -and size as a torchrun launch would, and serve tiny real transformers configs from memory, so the -decision builds their real models on the meta device; only the last ones use a process group, two -gloo ranks. +Nothing here needs a GPU. The tests of ``resolve_distribution_plan`` and of the loads +``load_pretrained`` routes a plan to answer the process group's rank and size as a torchrun launch +would, and serve tiny real transformers configs from memory: the decision builds their real models on +the meta device, and stand-ins for the sharded loaders record how they were called. Only the last +tests of each part use a process group, two gloo ranks. """ +import inspect import math import os from contextlib import contextmanager @@ -49,6 +51,7 @@ AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq, LlamaConfig, + LlamaForCausalLM, Phi4MultimodalConfig, Qwen2_5_VLConfig, Qwen3MoeConfig, @@ -56,9 +59,10 @@ WhisperConfig, ) +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 +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 ( _auto_class, @@ -66,6 +70,7 @@ _convert_keys, _resolve_target, build_meta_causal_lm, + load_pretrained, record_unplaced_source_keys, resolve_distribution_plan, unplaced_source_keys, @@ -236,17 +241,27 @@ def _gathered_alike(gathered, mine, group=None): def run_as(monkeypatch): """Run as rank ``rank`` of a ``world_size``-rank launch, deciding for itself. - Every rank reads the checkpoint as this one does. How the ranks agree on one plan is tested at - the end of this file. + Every rank acts as its group's first, and reads the checkpoint as this one does. How the ranks + agree on one plan, and on one load, is tested on two gloo ranks at the end of each part. """ 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) + broadcast_from = dist.broadcast_from + + def broadcast_as_the_first(label, fn, src=0): + # broadcast_from sends from dist.rank(), which process_group sets for the FSDP2 load. + with monkeypatch.context() as sending: + sending.setattr(dist, "rank", lambda group=None: src) + return broadcast_from(label, fn, src) + + monkeypatch.setattr(dist, "broadcast_from", broadcast_as_the_first) 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(dist, "is_master", lambda group=None: rank == 0) + monkeypatch.setattr(dist, "local_rank", lambda: rank) monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size) return enter @@ -822,3 +837,549 @@ def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers """Before the plan's broadcast, which rank 0 would otherwise wait in for a rank that raised.""" LlamaConfig(**_TINY).save_pretrained(tmp_path) cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path)) + + +# ------------------------------------------------------------------------------ loading by the plan + + +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) + + +# ----------------------------------------------------------------------------------- the loads + + +@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( + 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.""" + run_as(4, rank=1) + serve(LlamaConfig(**_TINY)) + kwargs = {} if device_map is None else {"device_map": device_map} + + 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( + "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, 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). + 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( + 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, local_files_only=True, **dtype + ) + + 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 [read.get("local_files_only") for read in reads] == [True] + assert distribution_plan(model) == plan + assert loaders == {} + + +def test_a_load_it_decides_reads_the_config_once( + 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") + run_as(2, rank=1) + config, reads = LlamaConfig(**_TINY), [] + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a) or config + ) + + model = load_pretrained(str(tmp_path)) + + assert len(reads) == 1 + assert loaders["fsdp2"]["hf_config"] is config + 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, 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. 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 = [] + 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, **plan) + + 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") + + +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 that shards the model: pass that"): + load_pretrained(str(tmp_path), hf_config=LlamaConfig(**{**_TINY, "num_hidden_layers": 2})) + 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, 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: + 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) + 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)) + assert loaders == {} + assert auto_class.calls == [] + + +def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed): + model, refused = None, None + with bounded_collectives(): + 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) + + +@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), handed=handed) + + +@pytest.mark.parametrize( + ("change", "error", "match"), + [ + ({"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"), + ( + {"kwargs": {"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}}, + TypeError, + "config and hf_config both name", + ), + ], + 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 +): + """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"): + 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=DistributionPlan(2, dense="fsdp"), + **change.get("kwargs", {}), + ) + assert loaders == {} + + +@pytest.mark.parametrize( + ("distribution", "checkpoint", "whole_model_kwargs", "error", "match"), + [ + ( + DistributionPlan(4, dense="fsdp"), + "causal-lm", + False, + ValueError, + "is for 4 ranks, but this run has 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, + "returned with a plan that shards the model", + ), + ], + 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, + serve, + loaders, + auto_class, + process_group, + tmp_path, + distribution, + checkpoint, + whole_model_kwargs, + error, + match, +): + """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would + 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(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 {} + + with pytest.raises(error, match=match): + load_pretrained(str(tmp_path), distribution=distribution, **whole) + assert loaders == {} + 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 == {} + + +@pytest.mark.parametrize( + ("distribution", "options", "error", "match"), + [ + (DistributionPlan(4), {}, ValueError, "is for 4 ranks"), + (DistributionPlan(2), {"hf_config": LlamaConfig(**_TINY)}, TypeError, "plan that shards"), + ( + DistributionPlan(2, dense="fsdp"), + {"device_map": "auto"}, + TypeError, + r"\['device_map'\] do not apply to the FSDP2 load", + ), + ( + DistributionPlan(2, dense="fsdp"), + {"auto_class": AutoModel}, + NotImplementedError, + "auto_class=AutoModel does not apply", + ), + ( + DistributionPlan(2, dense="fsdp"), + {"dtype": torch.float32}, + TypeError, + "keeps their dtype", + ), + ], + ids=[ + "another-world", + "hf-config-on-a-whole-load", + "whole-load-option-on-a-sharded-plan", + "auto-class-on-a-sharded-plan", + "dtype-on-a-sharded-plan", + ], +) +def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched( + run_as, monkeypatch, tmp_path, distribution, options, error, match +): + """Making the checkpoint local can mean downloading it: a plan the load refuses fails first, as + do hf_config with a plan that loads the model whole and an option a sharded plan's load + refuses.""" + run_as(2) + monkeypatch.setattr( + model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") + ) + + with pytest.raises(error, match=match): + load_pretrained(str(tmp_path), distribution=distribution, **options) + + +@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"} + 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)) + # 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_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 == {} + + +@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")) + + +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", "knobs", "error", "match"), + [ + ( + DistributionPlan(8), + {}, + RuntimeError, + "is for 8 ranks, whose process group this 8-rank launch has not", + ), + ( + DistributionPlan(dense="fsdp"), + {}, + RuntimeError, + r"shards through this 8-rank launch's .* DistributionPlan\(8, dense='fsdp'\)", + ), + ( + None, + {"MODELOPT_FSDP_SIZE": "world"}, + RuntimeError, + "MODELOPT_FSDP_SIZE=world ask to shard", + ), + (None, {"MODELOPT_EP_DEGREE": "4"}, RuntimeError, "MODELOPT_EP_DEGREE=4 ask to shard"), + (None, {"MODELOPT_EP_DEGREE": "eigth"}, ValueError, "MODELOPT_EP_DEGREE='eigth'"), + ], + ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob", "ep-knob-not-a-degree"], +) +def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( + launch_without_its_group, + serve, + loaders, + monkeypatch, + tmp_path, + distribution, + knobs, + error, + match, +): + """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 is fetched, rather than + leave each process the whole model the knobs did not ask for. A knob that is not one is refused + as the plan decision refuses it, not dropped.""" + serve(LlamaConfig(**_TINY)) + for name, value in knobs.items(): + monkeypatch.setenv(name, value) + monkeypatch.setattr( + model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") + ) + + with pytest.raises(error, match=match): + load_pretrained(str(tmp_path), distribution=distribution) + assert loaders == {} + + +# --------------------------------------------------------------------------------- fsdp2_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.""" + 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(dist, "size", lambda group=None: 4) + + model_load_utils.fsdp2_wrap(model) + + with pytest.raises(RuntimeError, match="no distribution plan recorded"): + distribution_plan(model) From d8a0e941f4b5ef12c044908782738bec5b502641 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:24:05 +0000 Subject: [PATCH 52/99] Refuse a load on every rank when any rank refuses it Each rank checks its own arguments and reads the config itself, so a rank could refuse the load alone: one handed another dtype, both config and hf_config, or a quantized hf_config raised while its peers went on into the checkpoint's fetch or the FSDP2 load's collectives and waited there. The argument checks now run after the plan's broadcast, and they and the checks after the fetch (the config reads and the sharded load's refusals) go through the failure-carrying all_reduce: a refusal on any rank raises on every rank before the next collective. _load_sharded keeps the group setup and the load; its checks move to _sharded_load_kwargs, which runs inside that all_reduce. A test runs the load on two real gloo ranks, rank 1 alone refused before the fetch or before the FSDP2 load, with the fetch and the load standing in as collectives on a bounded group: without the change rank 0 fails on a dropped connection instead of rank 1's refusal. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 197 +++++++++++------- .../unit/torch/utils/test_model_load_utils.py | 56 +++++ 2 files changed, 176 insertions(+), 77 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 50b45dbdadd..60acd5bfb93 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -20,6 +20,7 @@ import os import re from collections import defaultdict +from contextlib import contextmanager from itertools import chain from typing import TYPE_CHECKING, Any @@ -469,6 +470,32 @@ def _route(plan: DistributionPlan) -> str | None: return _FSDP2 if plan.is_fsdp_dense else None +@contextmanager +def _refused_together(label: str, alone: bool): + """Raise on every rank if the block raises on any, through :func:`dist.all_reduce`. + + A rank refusing a load alone would leave the others waiting in the load's next collective (the + checkpoint's fetch, the sharded load). A process of a launch that set no group up (``alone``) has + no peers to tell. + """ + error = None + try: + yield + # A SystemExit too: the others would wait for this rank all the same. + except BaseException as e: + error = e + if alone: + if error is not None: + raise error + return + + def refuse() -> None: + if error is not None: + raise error + + dist.all_reduce(label, refuse, combine=lambda _: None) + + def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide.""" if auto_class not in (None, AutoModelForCausalLM): @@ -499,8 +526,8 @@ def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: ) -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.""" +def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_options) -> dict: + """What the sharded load ``route`` names takes, refusing what it cannot honor.""" from modelopt.torch.models.hf import checkpoint_has_mtp _refuse_sharded_load_options(route, auto_class, kwargs) @@ -531,24 +558,27 @@ def _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) "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(): + if not dist.is_initialized() and "RANK" not in os.environ: # FSDP2 shards through a process group even on one rank; a launch without its group was - # refused before the load. 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`." - ) + # refused before the load. + 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`." + ) + return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} + + +def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: + """The FSDP2 load, by the arguments :func:`_sharded_load_kwargs` checked.""" + if not dist.is_initialized(): + # A single process's group of one; a run (distributed_run) destroys it on exit. dist.setup() 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, + **load_kwargs, ) @@ -578,7 +608,8 @@ 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 launch whose processes set no process group up (a + every rank loads the plan rank 0 was handed, and a refusal on any rank, of its own arguments or + of the config it reads, raises on every rank. 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 @@ -618,18 +649,6 @@ def load_pretrained( is returned alone: ``output_loading_info=True`` is refused. So is a ``variant``: the record and the export read the checkpoint's default weight files. """ - if kwargs.pop("output_loading_info", False): - raise ValueError( - "load_pretrained returns the model alone and does not take output_loading_info=True; " - "the checkpoint tensors the model has no parameter for are recorded on it instead " - "(see record_unplaced_source_keys)." - ) - if kwargs.get("variant") is not None: - raise ValueError( - f"load_pretrained does not take variant={kwargs['variant']!r}: the checkpoint tensors " - "the model has no parameter for are recorded from, and carried into the export from, " - "the checkpoint's default weight files." - ) # Importing modelopt.torch.models registers every model spec; only callers of this need them. from modelopt.torch.models.hf import prepare_model_for_loading @@ -669,8 +688,6 @@ def load_pretrained( ) 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.") def handed() -> DistributionPlan | None: if distribution is None: @@ -686,31 +703,48 @@ def handed() -> DistributionPlan | None: # so one it refuses raises on every rank instead of leaving them in the broadcast, and before # the checkpoint is fetched. plan = handed() if alone else dist.broadcast_from("distribution", handed) - if "hf_config" in kwargs and (plan is None or _route(plan) is None): - # Only a sharded load builds from hf_config, as given, so the plan must be the one it came - # with, not one decided here from another config nor one that loads the model whole. - raise TypeError( - "hf_config is the config resolve_distribution_plan returned with a plan that shards " - "the model: 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 {} - if plan is not None and _route(plan) is not None: - # A handed plan names its load before anything is fetched, and the arguments alone decide - # what that load refuses. A config passed for it becomes its hf_config; subfolder is where - # the checkpoint is. - _refuse_sharded_load_options( - _route(plan), - auto_class, - { - k: v - for k, v in kwargs.items() - if k != "subfolder" and not (k == "config" and config is not None) - }, - ) + # Each rank's arguments are its own, so each rank checks its own, and a refusal on one fails + # every rank before the checkpoint is fetched. + with _refused_together("load arguments", alone): + if kwargs.pop("output_loading_info", False): + raise ValueError( + "load_pretrained returns the model alone and does not take " + "output_loading_info=True; the checkpoint tensors the model has no parameter for " + "are recorded on it instead (see record_unplaced_source_keys)." + ) + if kwargs.get("variant") is not None: + raise ValueError( + f"load_pretrained does not take variant={kwargs['variant']!r}: the checkpoint " + "tensors the model has no parameter for are recorded from, and carried into the " + "export from, the checkpoint's default weight files." + ) + 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 "hf_config" in kwargs and (plan is None or _route(plan) is None): + # Only a sharded load builds from hf_config, as given, so the plan must be the one it + # came with, not one decided here from another config nor one that loads it whole. + raise TypeError( + "hf_config is the config resolve_distribution_plan returned with a plan that " + "shards the model: pass that plan as distribution=, or pass config= to have the " + "plan decided from your config." + ) + if plan is not None and _route(plan) is not None: + # A handed plan names its load before anything is fetched, and the arguments alone + # decide what that load refuses. A config passed for it becomes its hf_config; + # subfolder is where the checkpoint is. + _refuse_sharded_load_options( + _route(plan), + auto_class, + { + k: v + for k, v in kwargs.items() + if k != "subfolder" and not (k == "config" and config is not None) + }, + ) _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) subfolder = kwargs.pop("subfolder", None) if subfolder: @@ -731,37 +765,46 @@ 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. - del kwargs["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. - 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 + # Each rank reads the checkpoint's config itself, so a refusal of what it reads, or of its own + # arguments, fails every rank before the load's collectives. + with _refused_together("load checks", alone): + 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 " + f"layout, 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 pass " + f"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. + del kwargs["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. + 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 + ).model_type ) - 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 - ).model_type + load_kwargs = ( + {} + if route is None + else _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_options) ) if route is not None: - model = _load_sharded(route, ckpt_path, auto_class, model_type, kwargs, hub_options) + model = _load_sharded(ckpt_path, load_kwargs) # 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_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index cd0653667cc..fd54d1c573f 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1072,6 +1072,62 @@ 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), handed=handed) +def _fetched_in_step(ckpt_path, **hub_options): + """The checkpoint's fetch, as the collective it is: a rank not in it leaves the others waiting.""" + torch.distributed.barrier() + return None, ckpt_path + + +def _loaded_in_step(*args, **kwargs): + """The FSDP2 load, as the collectives it runs.""" + torch.distributed.barrier() + return _Loaded() + + +def _one_rank_refused_job(rank, size, ckpt, refused_on_1, match): + refused = None + with ( + bounded_collectives(), + mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), + mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", _loaded_in_step), + ): + try: # not pytest.raises: its failure would end the pooled worker instead of the job + load_pretrained( + ckpt, + distribution=DistributionPlan(2, dense="fsdp"), + **(refused_on_1 if rank == 1 else {}), + ) + except Exception as error: + refused = str(error) + # Rank 1's refusal on both ranks, rather than rank 0 failing at the group's timeout. + assert refused is not None and match in refused, (rank, refused) + assert rank == 1 or "rank 1 failed" in refused, refused + + +@pytest.mark.parametrize( + ("refused_on_1", "match"), + [ + ({"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}, "both name"), + ({"dtype": torch.float32}, "keeps their dtype"), + ( + {"hf_config": LlamaConfig(**_TINY, quantization_config={"quant_method": "fp8"})}, + "cannot dequantize", + ), + ], + ids=["two-configs", "dtype", "quantized-config"], +) +def test_a_load_refused_on_one_rank_fails_every_rank( + cpu_dist_workers, tmp_path, refused_on_1, match +): + """Rank 1 alone is refused, by its arguments before the checkpoint is fetched or by the config + it was handed before the FSDP2 load: every rank raises the refusal, rather than leave rank 0 + waiting in the fetch's or the load's collectives, which stand-ins on a bounded group run here.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run( + _one_rank_refused_job, ckpt=str(tmp_path), refused_on_1=refused_on_1, match=match + ) + + @pytest.mark.parametrize( ("change", "error", "match"), [ From c2b95bcfd7e34eb1c536ae139cdb3ab349b5e60a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:24:20 +0000 Subject: [PATCH 53/99] Keep load_pretrained's changelog entry to two sentences The entry keeps what a user sees, the distributed load and the expert-parallel refusal until that load lands; how the plan is decided and recorded is in load_pretrained's and resolve_distribution_plan's docstrings. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 150c8e918b3..189df7305a2 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +43,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``). 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``. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now loads a model distributed under ``torchrun``, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument. Until the expert-parallel load lands, an MoE planned expert-parallel is refused; set ``MODELOPT_EP_DEGREE=1`` to load it. **Backward Breaking Changes** From 86336d1451dce035da3a53ae895db448ed1ea905 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:53:30 +0000 Subject: [PATCH 54/99] Fail every rank when one cannot see the fetched checkpoint ensure_local_checkpoint ends on each rank's own check that it can see the checkpoint rank 0 fetched, after the fetch's collectives: a rank that cannot raised alone, and its peers went on into the load's next collective and waited there, as nothing ends that wait when the caller owns the group or catches the error. The fetch now runs inside the same failure-carrying all_reduce as the load's other refusals. The two-rank visibility test also loads by a handed plan, its error caught on both ranks: without the change rank 0 fails on a dropped connection instead of rank 1's refusal. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 15 +++++++++------ .../unit/torch/utils/test_model_load_utils.py | 18 ++++++++++++------ 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 60acd5bfb93..a53bde36239 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -608,11 +608,11 @@ 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, and a refusal on any rank, of its own arguments or - of the config it reads, raises on every rank. 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 + every rank loads the plan rank 0 was handed, and a refusal on any rank, of its own arguments, of a + checkpoint it cannot see or of the config it reads, raises on every rank. 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. @@ -745,7 +745,10 @@ def handed() -> DistributionPlan | None: if k != "subfolder" and not (k == "config" and config is not None) }, ) - _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) + # Rank 0 fetches the checkpoint, and each rank then checks it can see it: one that cannot fails + # every rank, rather than leave the others in the load's next collective. + with _refused_together("checkpoint fetch", alone): + _, 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) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index fd54d1c573f..15984865db3 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -817,26 +817,32 @@ def test_every_rank_takes_the_plan_rank0_decided(cpu_dist_workers, tmp_path): ) -def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt): +def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt, load): isdir = os.path.isdir with ( bounded_collectives(), mock.patch("os.path.isdir", lambda path: (rank == 0 or path != ckpt) and isdir(path)), ): try: - resolve_distribution_plan(ckpt) + if load: + load_pretrained(ckpt, distribution=DistributionPlan(size), device_map="cpu") + else: + resolve_distribution_plan(ckpt) except RuntimeError as error: raised = str(error) else: raised = "no error" - # Rank 1's own error, carried to rank 0, rather than rank 0 timing out in the plan's broadcast. + # Rank 1's own error, carried to rank 0, rather than rank 0 timing out in the next collective. assert "which this rank cannot see" in raised, raised -def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers, tmp_path): - """Before the plan's broadcast, which rank 0 would otherwise wait in for a rank that raised.""" +@pytest.mark.parametrize("load", [False, True], ids=["plan", "load"]) +def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers, tmp_path, load): + """Deciding the plan, or loading by a handed one: rank 0 would otherwise wait for the rank that + raised in the next collective, the plan's broadcast or the load's checks. The group is the + caller's and the error is caught, so nothing else ends rank 0's wait.""" LlamaConfig(**_TINY).save_pretrained(tmp_path) - cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path)) + cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path), load=load) # ------------------------------------------------------------------------------ loading by the plan From 878063215cd7664dbc4647c7b9d4cf4cfff64412 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:25:21 +0000 Subject: [PATCH 55/99] Decide the plan from a config path passed for the load from_pretrained takes config= as a config or as a path to one, and a caller that keeps the config apart from the weights, such as a weights-only checkpoint loaded through its model class with model_type given, loaded before load_pretrained decided a plan. The decision kept config= only as a config object, so it read the checkpoint's own config.json, which such a checkpoint does not have, and the load failed before it started. A config path is now read as the load reads it, inside the fetch's failure-carrying all_reduce, and the plan, the loading hook and a sharded load use that config. test_a_config_passed_for_the_load_needs_none_on_disk also loads a weights-only checkpoint through LlamaForCausalLM with config= naming a separate directory, with and without model_type; both fail without the change. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 14 +++++++-- .../torch/utils/test_unplaced_source_keys.py | 30 ++++++++++++++----- 2 files changed, 33 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index a53bde36239..d36b6a91b02 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -626,8 +626,9 @@ def load_pretrained( dtype (its config's, else bfloat16), not one the load asks for: hand a load that widens it, ``dtype=torch.float32`` say, its plan. ``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``. + builds from it as given, so without that plan it is refused. A ``config`` passed for the load, or + a path to one as ``from_pretrained`` takes, is the one the plan is decided from, and a sharded + load takes it as its ``hf_config``. Whatever the placement: @@ -686,8 +687,9 @@ def load_pretrained( f"{os.environ['WORLD_SIZE']}-rank launch, whose process group is not set up; load " "inside distributed_run(), which sets it up." ) + # The config to load with, or, as from_pretrained also takes it, a path to read it from. config = kwargs.get("config") - config = config if isinstance(config, PretrainedConfig) else None + config = config if isinstance(config, (PretrainedConfig, str, os.PathLike)) else None def handed() -> DistributionPlan | None: if distribution is None: @@ -749,6 +751,12 @@ def handed() -> DistributionPlan | None: # every rank, rather than leave the others in the load's next collective. with _refused_together("checkpoint fetch", alone): _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) + if config is not None and not isinstance(config, PretrainedConfig): + # Read as the load reads it, so the plan and the loading hook follow the config the + # model is built from, which need not be in the checkpoint's directory. + config = AutoConfig.from_pretrained( + config, trust_remote_code=kwargs.get("trust_remote_code", False), **hub_options + ) subfolder = kwargs.pop("subfolder", None) if subfolder: ckpt_path = os.path.join(ckpt_path, subfolder) diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 8632c924772..4004ccb1641 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -194,15 +194,29 @@ def test_an_option_it_cannot_honor_is_refused_before_anything_loads(option, refu load_pretrained("no/such/checkpoint", **option) -def test_a_config_passed_for_the_load_needs_none_on_disk(tmp_path): - """The load uses the config it is passed, as ``from_pretrained`` does, and so does the choice of - its loading hook: a directory without a config.json loads.""" +@pytest.mark.parametrize( + ("as_path", "model_type"), + [(False, None), (True, None), (True, "llama")], + ids=["config", "config-path", "config-path-and-model-type"], +) +def test_a_config_passed_for_the_load_needs_none_on_disk(tmp_path, as_path, model_type): + """The load uses the config it is passed, or reads it from the directory it is passed, as + ``from_pretrained`` does, and so do its plan and the choice of its loading hook: a checkpoint + directory without a config.json loads, through a model class of the caller's own too.""" config = _tiny_llama_config() - _save(tmp_path, config, {**LlamaForCausalLM(config).state_dict(), STRAY: torch.ones(32)}) - (tmp_path / "config.json").unlink() - (tmp_path / "generation_config.json").write_text("{}") # from_pretrained reads it instead - - model = _load(tmp_path, config=config) + weights = tmp_path / "weights" + weights.mkdir() + _save(weights, config, {**LlamaForCausalLM(config).state_dict(), STRAY: torch.ones(32)}) + if as_path: + (tmp_path / "config").mkdir() + (weights / "config.json").rename(tmp_path / "config" / "config.json") + passed = {"auto_class": LlamaForCausalLM, "config": str(tmp_path / "config")} + else: + (weights / "config.json").unlink() + passed = {"config": config} + (weights / "generation_config.json").write_text("{}") # from_pretrained reads it instead + + model = _load(weights, model_type=model_type, **passed) assert model._modelopt_unplaced_source_keys == [STRAY] From 4cd6518f8cd0107ab1e93fbbbaa4a36cc6f3324b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:37:07 +0000 Subject: [PATCH 56/99] Read a config path through the model class's own config class from_pretrained reads config= given as a path through its model class's config_class, while the plan's read went through AutoConfig, which needs the config's model_type in its mapping: a model class of the caller's own whose config class AutoConfig does not know loaded before, and was refused by that read. The read now uses the model class's config_class, and the auto classes, which have none, still read through AutoConfig. test_a_config_passed_for_the_load_needs_none_on_disk has a row for such a model class (own-config-class-path); it fails with a KeyError without the change. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 6 ++-- .../torch/utils/test_unplaced_source_keys.py | 34 +++++++++++++++---- 2 files changed, 31 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d36b6a91b02..69227bb60ce 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -753,8 +753,10 @@ def handed() -> DistributionPlan | None: _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) if config is not None and not isinstance(config, PretrainedConfig): # Read as the load reads it, so the plan and the loading hook follow the config the - # model is built from, which need not be in the checkpoint's directory. - config = AutoConfig.from_pretrained( + # model is built from, which need not be in the checkpoint's directory. A model class + # of the caller's own reads it through its own config class, as from_pretrained does. + reader = getattr(auto_class, "config_class", None) or AutoConfig + config = reader.from_pretrained( config, trust_remote_code=kwargs.get("trust_remote_code", False), **hub_options ) subfolder = kwargs.pop("subfolder", None) diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 4004ccb1641..d269cc47b59 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -194,23 +194,43 @@ def test_an_option_it_cannot_honor_is_refused_before_anything_loads(option, refu load_pretrained("no/such/checkpoint", **option) +class _OwnConfig(LlamaConfig): + """A config class of the caller's own, which AutoConfig has no mapping for.""" + + model_type = "modelopt_test_own_llama" + + +class _OwnModel(LlamaForCausalLM): + config_class = _OwnConfig + + @pytest.mark.parametrize( - ("as_path", "model_type"), - [(False, None), (True, None), (True, "llama")], - ids=["config", "config-path", "config-path-and-model-type"], + ("as_path", "model_type", "model_class"), + [ + (False, None, None), + (True, None, LlamaForCausalLM), + (True, "llama", LlamaForCausalLM), + (True, None, _OwnModel), + ], + ids=["config", "config-path", "config-path-and-model-type", "own-config-class-path"], ) -def test_a_config_passed_for_the_load_needs_none_on_disk(tmp_path, as_path, model_type): +def test_a_config_passed_for_the_load_needs_none_on_disk( + tmp_path, as_path, model_type, model_class +): """The load uses the config it is passed, or reads it from the directory it is passed, as - ``from_pretrained`` does, and so do its plan and the choice of its loading hook: a checkpoint - directory without a config.json loads, through a model class of the caller's own too.""" + ``from_pretrained`` does, through the model class's own config class, and so do its plan and the + choice of its loading hook: a checkpoint directory without a config.json loads, through a model + class of the caller's own too.""" config = _tiny_llama_config() + if model_class is _OwnModel: + config = _OwnConfig(**config.to_diff_dict()) weights = tmp_path / "weights" weights.mkdir() _save(weights, config, {**LlamaForCausalLM(config).state_dict(), STRAY: torch.ones(32)}) if as_path: (tmp_path / "config").mkdir() (weights / "config.json").rename(tmp_path / "config" / "config.json") - passed = {"auto_class": LlamaForCausalLM, "config": str(tmp_path / "config")} + passed = {"auto_class": model_class, "config": str(tmp_path / "config")} else: (weights / "config.json").unlink() passed = {"config": config} From 92290af04532a6343ffe3aefd36c7322453b14a8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:49:31 +0000 Subject: [PATCH 57/99] Decide nothing for a process alone, and read no config for it A process without a group, or in a group of one, decided a plan through resolve_distribution_plan, which reads the checkpoint's config with AutoConfig before it finds a single rank holds the whole model. A model class of the caller's own, with its model_type given, read no config before plans existed and loaded; one whose config class AutoConfig has no mapping for was refused by that read. Such a process now takes DistributionPlan(), what the decision rules for one rank whatever the model, without reading anything. A config path read through a model class's own config class no longer gets trust_remote_code, an AutoConfig option the class only warns about. test_a_model_class_of_the_callers_own_reads_its_config_itself loads such a class with its model_type given and the config in the checkpoint; it fails without the change. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 17 ++++++++++++----- .../torch/utils/test_unplaced_source_keys.py | 13 +++++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 69227bb60ce..08c96c22dfa 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -755,17 +755,24 @@ def handed() -> DistributionPlan | None: # Read as the load reads it, so the plan and the loading hook follow the config the # model is built from, which need not be in the checkpoint's directory. A model class # of the caller's own reads it through its own config class, as from_pretrained does. - reader = getattr(auto_class, "config_class", None) or AutoConfig - config = reader.from_pretrained( - config, trust_remote_code=kwargs.get("trust_remote_code", False), **hub_options + reader = getattr(auto_class, "config_class", None) + config = ( + reader.from_pretrained(config, **hub_options) + if reader is not None + else AutoConfig.from_pretrained( + config, trust_remote_code=kwargs.get("trust_remote_code", False), **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") - decided = plan is None and not alone - if plan is None and alone: + # A process alone, without a group or in a group of one, holds the whole model, as the decision + # rules for one rank whatever the model: nothing to decide, nor a config to read for it, which a + # model class of the caller's own reads through its own config class. + decided = plan is None and dist.size() > 1 + if plan is None and not decided: plan = DistributionPlan() elif plan is None: plan, hf_config = resolve_distribution_plan( diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index d269cc47b59..4ddbffeedec 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -241,6 +241,19 @@ class of the caller's own too.""" assert model._modelopt_unplaced_source_keys == [STRAY] +def test_a_model_class_of_the_callers_own_reads_its_config_itself(tmp_path): + """Its model type given, a single process reads no config for a model class of the caller's + own: the class reads the checkpoint's through its own config class, which AutoConfig need not + know, and the plan of a single process is the same whatever the model.""" + config = _OwnConfig(**_tiny_llama_config().to_diff_dict()) + _save(tmp_path, config, {**LlamaForCausalLM(config).state_dict(), STRAY: torch.ones(32)}) + + model = _load(tmp_path, _OwnModel, model_type=_OwnConfig.model_type) + + assert model._modelopt_unplaced_source_keys == [STRAY] + assert distribution_plan(model) == DistributionPlan() + + def test_an_ignore_rule_does_not_hide_an_unplaced_tensor(tmp_path): """Qwen3-Next ignores ``^mtp.*``, DeepSeek-V3 and GLM their MTP layers. From e100c5850c9ef9530884c753bca3c390ff1e29b1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 02:03:20 +0000 Subject: [PATCH 58/99] Read the checkpoint's config for the plan as the model class does With several ranks the plan decision read the checkpoint's config through AutoConfig, which refuses a config whose class it has no mapping for, before it found that a model class of the caller's own loads whole without needing it: such a class, with its model_type given, loaded before plans existed and was refused by that read. The decision now reads the config as the load does, through the model class's own config_class, or AutoConfig for an auto class, in one helper the config-path read uses too. The single-process branch reads as one condition instead of a double negative. test_a_model_class_of_the_callers_own_loads_whole has a two-rank row with such a config on disk and no stubbed config read; it fails without the change. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 38 ++++++++++--------- .../unit/torch/utils/test_model_load_utils.py | 24 ++++++++++-- 2 files changed, 41 insertions(+), 21 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 08c96c22dfa..968dc09fce9 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -496,6 +496,17 @@ def refuse() -> None: dist.all_reduce(label, refuse, combine=lambda _: None) +def _read_config(path, auto_class, trust_remote_code: bool = False, **hub_options): + """``path``'s config, read as ``auto_class.from_pretrained`` reads it. + + A model class reads it through its own config class, an auto class, which has none, through + ``AutoConfig``. + """ + if (reader := getattr(auto_class, "config_class", None)) is not None: + return reader.from_pretrained(path, **hub_options) + return AutoConfig.from_pretrained(path, trust_remote_code=trust_remote_code, **hub_options) + + def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide.""" if auto_class not in (None, AutoModelForCausalLM): @@ -753,28 +764,16 @@ def handed() -> DistributionPlan | None: _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) if config is not None and not isinstance(config, PretrainedConfig): # Read as the load reads it, so the plan and the loading hook follow the config the - # model is built from, which need not be in the checkpoint's directory. A model class - # of the caller's own reads it through its own config class, as from_pretrained does. - reader = getattr(auto_class, "config_class", None) - config = ( - reader.from_pretrained(config, **hub_options) - if reader is not None - else AutoConfig.from_pretrained( - config, trust_remote_code=kwargs.get("trust_remote_code", False), **hub_options - ) - ) + # model is built from, which need not be in the checkpoint's directory. + trust = kwargs.get("trust_remote_code", False) + config = _read_config(config, auto_class, trust, **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") - # A process alone, without a group or in a group of one, holds the whole model, as the decision - # rules for one rank whatever the model: nothing to decide, nor a config to read for it, which a - # model class of the caller's own reads through its own config class. decided = plan is None and dist.size() > 1 - if plan is None and not decided: - plan = DistributionPlan() - elif plan is None: + if decided: plan, hf_config = resolve_distribution_plan( ckpt_path, config=config, @@ -784,6 +783,10 @@ def handed() -> DistributionPlan | None: ) if _route(plan) is not None: kwargs.setdefault("hf_config", hf_config) + elif plan is None: + # A process alone, without a group or in a group of one, holds the whole model, as the + # decision rules for one rank whatever the model: nothing to decide, nor a config to read. + plan = DistributionPlan() route = _route(plan) # Each rank reads the checkpoint's config itself, so a refusal of what it reads, or of its own # arguments, fails every rank before the load's collectives. @@ -1540,8 +1543,9 @@ def checkpoint_read() -> bool: if config is not None: read.append((local_path, copy.deepcopy(config))) return True - hf_config = AutoConfig.from_pretrained( + hf_config = _read_config( local_path, + auto_class, trust_remote_code=trust_remote_code, token=token, cache_dir=cache_dir, diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 15984865db3..f24fb2e30d0 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1012,16 +1012,32 @@ 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"]) +class _OwnConfig(LlamaConfig): + """A config class of the caller's own, which AutoConfig has no mapping for.""" + + model_type = "modelopt_test_own_llama" + + +@pytest.mark.parametrize( + ("fsdp_size", "own_config"), + [(None, False), ("world", False), (None, True)], + ids=["no-knob", "knob-shards", "own-config-class"], +) def test_a_model_class_of_the_callers_own_loads_whole( - run_as, serve, loaders, auto_class, monkeypatch, tmp_path, fsdp_size + run_as, serve, loaders, auto_class, monkeypatch, tmp_path, fsdp_size, own_config ): """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, 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).""" + class may not build (an encoder-decoder loaded through AutoModel). The decision reads the + checkpoint's config as the class does, through its own config class, which AutoConfig need not + know.""" run_as(2) - serve(LlamaConfig(**_TINY)) + if own_config: + _OwnConfig(**_TINY).save_pretrained(tmp_path) + auto_class.config_class = _OwnConfig + else: + serve(LlamaConfig(**_TINY)) monkeypatch.setattr( model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") ) From 3ba6b062fa1ea56d9f5cda0801122880a3f9b55a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 02:04:26 +0000 Subject: [PATCH 59/99] Keep the config reader's docstring and the one-rank comment to a line each Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 968dc09fce9..e62c8fc1f3c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -497,11 +497,7 @@ def refuse() -> None: def _read_config(path, auto_class, trust_remote_code: bool = False, **hub_options): - """``path``'s config, read as ``auto_class.from_pretrained`` reads it. - - A model class reads it through its own config class, an auto class, which has none, through - ``AutoConfig``. - """ + """``path``'s config as ``auto_class.from_pretrained`` reads it: config_class, or AutoConfig.""" if (reader := getattr(auto_class, "config_class", None)) is not None: return reader.from_pretrained(path, **hub_options) return AutoConfig.from_pretrained(path, trust_remote_code=trust_remote_code, **hub_options) @@ -784,8 +780,7 @@ def handed() -> DistributionPlan | None: if _route(plan) is not None: kwargs.setdefault("hf_config", hf_config) elif plan is None: - # A process alone, without a group or in a group of one, holds the whole model, as the - # decision rules for one rank whatever the model: nothing to decide, nor a config to read. + # One rank holds the whole model, whatever the model: there is nothing to decide or read. plan = DistributionPlan() route = _route(plan) # Each rank reads the checkpoint's config itself, so a refusal of what it reads, or of its own From c73831710ad6b574a3e37acafa26793d0ce613ee Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 02:21:26 +0000 Subject: [PATCH 60/99] Refuse FSDP2-only options on a whole load, and parse the EP knob once cpu_offload, mp_policy and broadcast_chunk_size apply only to the FSDP2 load, which refuses what it cannot honor, but on a plan that loads the model whole they reached from_pretrained, which fails on them in the model's constructor or drops them: with MODELOPT_FSDP_SIZE=auto the same call takes either route. They are now refused on the whole route, with the other load checks. A launch without its group parsed MODELOPT_EP_DEGREE by rules of its own, and the two parses disagreed: a padded auto, +2 and a degree below 1 were taken by one and refused by the other. That launch now reads the knob through the plan decision's parse, which takes surrounding spaces and refuses a degree below 1. New rows fail without the change: an FSDP2-only option on a replicated plan, a padded auto, and a degree of -2 from the knob or the argument. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 25 +++++++--- .../unit/torch/utils/test_model_load_utils.py | 50 ++++++++++++------- 2 files changed, 49 insertions(+), 26 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index e62c8fc1f3c..ad790d076f2 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -681,11 +681,9 @@ def load_pretrained( ) if alone and distribution is None: # The knobs ask to shard over the launch, rather than leave each process the whole model. A - # value that is not an EP degree is refused, as the plan decision refuses it. - ep, world = env.ep_degree().strip(), int(os.environ["WORLD_SIZE"]) - if ep.lower() not in ("auto", "0") and not ep.isdigit(): - raise ValueError(f"MODELOPT_EP_DEGREE={ep!r}: expected auto, 0 (auto) or an EP degree.") - asked = [f"MODELOPT_EP_DEGREE={ep}"] if ep.isdigit() and int(ep) > 1 else [] + # value that is not an EP degree is refused, by the plan decision's own parse. + ep, world = _ep_degree(None), int(os.environ["WORLD_SIZE"]) + asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if ep != "auto" and ep > 1 else [] if _dense_from_env(world) == "fsdp": asked.append(f"MODELOPT_FSDP_SIZE={env.fsdp_size()}") if asked: @@ -786,6 +784,14 @@ def handed() -> DistributionPlan | None: # Each rank reads the checkpoint's config itself, so a refusal of what it reads, or of its own # arguments, fails every rank before the load's collectives. with _refused_together("load checks", alone): + # Options only the FSDP2 load takes would reach from_pretrained, which has no use for them. + shared = {"trust_remote_code", "attn_implementation", "hf_config"} # hf_config: see above + own = _SHARDED_LOAD_OPTIONS[_FSDP2] - shared + if route is None and (sharded := sorted(set(kwargs) & own)): + raise TypeError( + f"{sharded} only apply to the FSDP2 load, but {plan!r} gives every rank the whole " + "model; pass a plan with dense='fsdp' to shard it." + ) 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 @@ -1351,13 +1357,16 @@ def _dense_from_env(world_size: int) -> str: def _ep_degree(ep_size) -> int | str: """``ep_size`` as an EP degree or ``"auto"``; ``None`` reads ``MODELOPT_EP_DEGREE``.""" value = env.ep_degree() if ep_size is None else ep_size - if str(value).lower() in ("auto", "0"): + if str(value).strip().lower() in ("auto", "0"): return "auto" try: - return int(value) + degree = int(value) except ValueError: + degree = 0 + if degree < 1: name = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" - raise ValueError(f"{name}={value!r}: expected auto, 0 (auto) or an EP degree.") from None + raise ValueError(f"{name}={value!r}: expected auto, 0 (auto) or an EP degree.") + return degree def _decide_distribution( diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index f24fb2e30d0..89e2100d1d0 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -501,31 +501,35 @@ def test_an_ep_degree_that_does_not_divide_the_world_names_where_it_came_from( resolve_distribution_plan(str(tmp_path), **arguments) +@pytest.mark.parametrize("value", ["world", "-2"], ids=["word", "below-1"]) @pytest.mark.parametrize( ("world_size", "source"), [(4, "MODELOPT_EP_DEGREE"), (1, "ep_size")], ids=["knob", "argument-in-a-single-process"], ) def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from( - run_as, serve, monkeypatch, tmp_path, world_size, source + run_as, serve, monkeypatch, tmp_path, world_size, source, value ): - """A single process takes its own argument as given, but parses it as the decision does.""" + """A single process takes its own argument as given, but parses it as the decision does, and a + degree below 1 is no degree either.""" run_as(world_size) serve(LlamaConfig(**_TINY)) arguments = {} if source == "ep_size": - arguments["ep_size"] = "world" + arguments["ep_size"] = value else: - monkeypatch.setenv(source, "world") + monkeypatch.setenv(source, value) - with pytest.raises(ValueError, match=f"{source}='world'"): + with pytest.raises(ValueError, match=f"{source}='{value}'"): resolve_distribution_plan(str(tmp_path), **arguments) @needs_qwen3_moe_ep_plan -@pytest.mark.parametrize("value", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"]) +@pytest.mark.parametrize( + "value", ["0", "", "AUTO", " auto "], ids=["zero", "exported-empty", "capitals", "padded"] +) def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value): - """And so is one exported empty, and auto in any case: the plan decides.""" + """And so is one exported empty, and auto in any case or padded: the plan decides.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", value) run_as(4) serve(_qwen3_moe(num_experts=8)) @@ -1192,37 +1196,44 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( @pytest.mark.parametrize( - ("distribution", "checkpoint", "whole_model_kwargs", "error", "match"), + ("distribution", "checkpoint", "options", "error", "match"), [ ( DistributionPlan(4, dense="fsdp"), "causal-lm", - False, + None, ValueError, "is for 4 ranks, but this run has 2", ), ( DistributionPlan(2, dense="fsdp"), "quantized", - False, + None, NotImplementedError, "cannot dequantize", ), ( DistributionPlan(2, dense="fsdp"), "vlm", - False, + None, NotImplementedError, "builds a causal LM, which this llava checkpoint is not", ), - ("fsdp2", "causal-lm", False, TypeError, "expected a DistributionPlan or None"), + ("fsdp2", "causal-lm", None, TypeError, "expected a DistributionPlan or None"), ( DistributionPlan(2), "causal-lm", - True, + "hf_config", TypeError, "returned with a plan that shards the model", ), + ( + DistributionPlan(2), + "causal-lm", + {"cpu_offload": True}, + TypeError, + r"\['cpu_offload'\] only apply to the FSDP2 load", + ), ], ids=[ "another-world", @@ -1230,6 +1241,7 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( "vlm-checkpoint", "not-a-plan", "hf-config-on-a-whole-load", + "fsdp2-option-on-a-whole-load", ], ) def test_a_load_refuses_what_it_cannot_load( @@ -1241,22 +1253,24 @@ def test_a_load_refuses_what_it_cannot_load( tmp_path, distribution, checkpoint, - whole_model_kwargs, + options, error, match, ): """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would 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.""" + language model alone; anything but a plan; and the sharded loads' config, or an option only the + FSDP2 load takes, which from_pretrained would take for a model argument and fail far from here, + or drop.""" run_as(2) 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 {} + if options == "hf_config": + options = {"auto_class": auto_class, "hf_config": config} with pytest.raises(error, match=match): - load_pretrained(str(tmp_path), distribution=distribution, **whole) + load_pretrained(str(tmp_path), distribution=distribution, **(options or {})) assert loaders == {} assert auto_class.calls == [] From de1521d9440ab63fa4af37126bf896052df87bb9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 02:37:58 +0000 Subject: [PATCH 61/99] Place a whole load on a GPU only where there is one, and read knobs alike A distributed whole-model load defaulted each rank's device_map to its GPU, so a run on a group without GPUs, such as a gloo group on CPU, which loaded before plans existed, failed placing the model on CUDA. The default now applies only where CUDA is available; elsewhere the load places the model as from_pretrained does. A model type left to read is read as the model class reads its config, as the plan decision and a config path already are, so a model class of the caller's own loads the same in one process as on several. MODELOPT_FSDP_SIZE takes surrounding spaces, as MODELOPT_EP_DEGREE now does. New rows fail without the change: two gloo ranks loading whole with no device_map, a caller's model class whose model type is read, and a padded MODELOPT_FSDP_SIZE. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 10 +++---- .../unit/torch/utils/test_model_load_utils.py | 27 +++++++++++-------- .../torch/utils/test_unplaced_source_keys.py | 11 ++++---- 3 files changed, 26 insertions(+), 22 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ad790d076f2..4041141801e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -612,7 +612,7 @@ def load_pretrained( 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. + whole model, ``device_map`` defaulting to the rank's GPU when there is one. 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, and a refusal on any rank, of its own arguments, of a @@ -818,9 +818,7 @@ def handed() -> DistributionPlan | None: ) 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 - ).model_type + _read_config(ckpt_path, auto_class, trust_remote_code, **hub_options).model_type ) load_kwargs = ( {} @@ -832,7 +830,7 @@ def handed() -> DistributionPlan | None: # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: - if plan.world_size > 1: + if plan.world_size > 1 and torch.cuda.is_available(): # 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()}"}) @@ -1341,7 +1339,7 @@ def _gpu_memory() -> int | None: def _dense_from_env(world_size: int) -> str: """``MODELOPT_FSDP_SIZE`` as a dense layout, or ``"auto"``.""" - value = env.fsdp_size().lower() + value = env.fsdp_size().strip().lower() if value == "auto": return "auto" if value == "1": diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 89e2100d1d0..0dd2c8958a0 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -609,9 +609,10 @@ def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path): ("4", "fsdp"), ("WORLD", "fsdp"), ("", "fsdp"), + (" world ", "fsdp"), ("2", None), ], - ids=["1", "4", "WORLD", "exported-empty", "part-of-the-world"], + ids=["1", "4", "WORLD", "exported-empty", "padded", "part-of-the-world"], ) def test_the_dense_knob_replicates_or_shards_over_the_whole_world( run_as, serve, gpu_memory, monkeypatch, tmp_path, value, dense @@ -1067,11 +1068,11 @@ 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): +def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed, placed): model, refused = None, None with bounded_collectives(): 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") + model = load_pretrained(ckpt, distribution=handed[rank], **placed) except Exception as error: refused = str(error) if isinstance(handed[0], DistributionPlan): @@ -1082,20 +1083,24 @@ def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed): @pytest.mark.parametrize( - "handed", + ("handed", "placed"), [ - (DistributionPlan(2), DistributionPlan(2, dense="fsdp")), - (DistributionPlan(2), None), - ("fsdp2", DistributionPlan(2)), + ((DistributionPlan(2), DistributionPlan(2, dense="fsdp")), {"device_map": "cpu"}), + ((DistributionPlan(2), None), {"device_map": "cpu"}), + (("fsdp2", DistributionPlan(2)), {"device_map": "cpu"}), + ((DistributionPlan(2), DistributionPlan(2)), {}), ], - ids=["another-plan", "none", "not-a-plan"], + ids=["another-plan", "none", "not-a-plan", "placed-by-default"], ) -def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, handed): +def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, handed, placed): """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.""" + join: both would hang. A distribution rank 0 refuses fails every rank. Placed by default, each + rank takes the model onto its GPU, and a run without GPUs, as these ranks are, loads it whole.""" LlamaForCausalLM(LlamaConfig(**_TINY)).save_pretrained(tmp_path) - cpu_dist_workers.run(_every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path), handed=handed) + cpu_dist_workers.run( + _every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path), handed=handed, placed=placed + ) def _fetched_in_step(ckpt_path, **hub_options): diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 4ddbffeedec..46a5e7ef634 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -241,14 +241,15 @@ class of the caller's own too.""" assert model._modelopt_unplaced_source_keys == [STRAY] -def test_a_model_class_of_the_callers_own_reads_its_config_itself(tmp_path): - """Its model type given, a single process reads no config for a model class of the caller's - own: the class reads the checkpoint's through its own config class, which AutoConfig need not - know, and the plan of a single process is the same whatever the model.""" +@pytest.mark.parametrize("model_type", [_OwnConfig.model_type, None], ids=["given", "read"]) +def test_a_model_class_of_the_callers_own_reads_its_config_itself(tmp_path, model_type): + """A single process plans nothing for a model class of the caller's own, and a model type left + to read is read as the class reads its config, through its own config class, which AutoConfig + need not know: the plan of a single process is the same whatever the model.""" config = _OwnConfig(**_tiny_llama_config().to_diff_dict()) _save(tmp_path, config, {**LlamaForCausalLM(config).state_dict(), STRAY: torch.ones(32)}) - model = _load(tmp_path, _OwnModel, model_type=_OwnConfig.model_type) + model = _load(tmp_path, _OwnModel, model_type=model_type) assert model._modelopt_unplaced_source_keys == [STRAY] assert distribution_plan(model) == DistributionPlan() From 7a42a639ae4173b08032d8f4038c0f77b131bcb3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 02:39:31 +0000 Subject: [PATCH 62/99] Answer the rank-GPU placement test as a host with GPUs The default device_map now applies only where CUDA is available, so the test of that default says each rank has a GPU; a run without one is covered by the two-rank gloo load. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 0dd2c8958a0..1b608c9d557 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -914,12 +914,13 @@ def process_group(monkeypatch): ids=["defaults_to_the_rank_gpu", "caller_placement_wins"], ) def test_a_replicated_plan_loads_the_whole_model_on_every_rank( - run_as, serve, loaders, auto_class, tmp_path, device_map, placed + run_as, serve, loaders, auto_class, monkeypatch, 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.""" run_as(4, rank=1) serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) # each rank has a GPU kwargs = {} if device_map is None else {"device_map": device_map} model = load_pretrained(str(tmp_path), auto_class, **kwargs) From f04db97bc2a2ecca766ed4ef94704a3f04ab8e12 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 07:01:14 +0000 Subject: [PATCH 63/99] Tighten the load routing's comments and docstrings Shorter wording for the routing's own comments and docstrings, and one exit path per mode in _refused_together, with no change in what any load does or refuses. The automatic layout's placement refusal asks MODELOPT_FSDP_SIZE's one parse, _dense_from_env, rather than reading the variable again, so a padded auto gets that refusal's message too. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/fsdp2_utils.py | 6 +-- .../torch/utils/plugins/model_load_utils.py | 50 +++++++------------ 2 files changed, 19 insertions(+), 37 deletions(-) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index 059d29abdc4..f939fbb36fd 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -100,10 +100,8 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False replicated and are moved onto the shards' device, which ``fully_shard`` does not do for the params it ignores. - 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. + It records no distribution plan: :func:`load_pretrained` does, and a caller wrapping a model + itself records one (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). """ # Importing the quantization package registers every mode and plugin; only a wrap needs it. from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 4041141801e..37d4d8c5e6a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -472,28 +472,21 @@ def _route(plan: DistributionPlan) -> str | None: @contextmanager def _refused_together(label: str, alone: bool): - """Raise on every rank if the block raises on any, through :func:`dist.all_reduce`. - - A rank refusing a load alone would leave the others waiting in the load's next collective (the - checkpoint's fetch, the sharded load). A process of a launch that set no group up (``alone``) has - no peers to tell. - """ + """Raise on every rank if the block raises on any, so no rank waits for one that refused.""" error = None try: yield - # A SystemExit too: the others would wait for this rank all the same. - except BaseException as e: + except BaseException as e: # A SystemExit too: the others would wait for this rank as well. error = e - if alone: - if error is not None: - raise error - return def refuse() -> None: if error is not None: raise error - dist.all_reduce(label, refuse, combine=lambda _: None) + if alone: # A process of a launch that set no group up has no peers to tell. + refuse() + else: + dist.all_reduce(label, refuse, combine=lambda _: None) def _read_config(path, auto_class, trust_remote_code: bool = False, **hub_options): @@ -566,8 +559,7 @@ def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_o f"or distribution=DistributionPlan({dist.size()}, dense='replicate')." ) if not dist.is_initialized() and "RANK" not in os.environ: - # FSDP2 shards through a process group even on one rank; a launch without its group was - # refused before the load. + # FSDP2 shards through a process group even on one rank. 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`." @@ -605,11 +597,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). 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. + ``broadcast_chunk_size`` ``kwargs`` may set). A single process sets a group of one up for it, + which needs ``torchrun --nproc_per_node=1``. 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 when there is one. @@ -661,8 +651,7 @@ def load_pretrained( from modelopt.torch.models.hf import prepare_model_for_loading # 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. + # --use_fsdp2, say): each loads as a single process does, the whole model with no plan. alone = dist.is_distributed_launch() and not dist.is_initialized() if alone and isinstance(distribution, DistributionPlan) and distribution.world_size > 1: raise RuntimeError( @@ -704,11 +693,8 @@ def handed() -> DistributionPlan | 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, and before - # the checkpoint is fetched. + # Every rank takes rank 0's distribution, so every rank runs the same load and collectives; + # rank 0 checks it before it sends, so one it refuses raises on every rank, before the fetch. plan = handed() if alone else dist.broadcast_from("distribution", handed) hub_options = { k: kwargs.pop(k) for k in ("token", "cache_dir", "local_files_only") if k in kwargs @@ -740,9 +726,8 @@ def handed() -> DistributionPlan | None: "plan decided from your config." ) if plan is not None and _route(plan) is not None: - # A handed plan names its load before anything is fetched, and the arguments alone - # decide what that load refuses. A config passed for it becomes its hf_config; - # subfolder is where the checkpoint is. + # A handed plan's arguments alone decide what its load refuses, before the fetch; a + # config passed for it becomes its hf_config, and subfolder is where the checkpoint is. _refuse_sharded_load_options( _route(plan), auto_class, @@ -793,7 +778,7 @@ def handed() -> DistributionPlan | None: "model; pass a plan with dense='fsdp' to shard it." ) placed = _whole_load_options(kwargs) - if decided and route == _FSDP2 and env.fsdp_size().lower() == "auto" and placed: + if decided and route == _FSDP2 and _dense_from_env(plan.world_size) == "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( @@ -804,8 +789,7 @@ def handed() -> DistributionPlan | None: f"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. + # The sharded load builds from hf_config, a copy, as building writes to the config. del kwargs["config"] if "hf_config" not in kwargs: kwargs["hf_config"] = copy.deepcopy(config) From 8d9120b0795c2ba9c3ef741bd2f4446b8988ee21 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 07:26:47 +0000 Subject: [PATCH 64/99] Say which loads read no knob and get no device_map default load_pretrained's docstring now says that a single process, or a group of one, reads neither MODELOPT_ knob and takes the FSDP2 load only from a plan passed as distribution, and that a launch without its group has each process load with no device_map default, as a single process does. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 37d4d8c5e6a..31ffc9c1745 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -597,9 +597,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 sets a group of one up for it, - which needs ``torchrun --nproc_per_node=1``. It reads the checkpoint's tensors as stored, so it - refuses a quantized checkpoint. + ``broadcast_chunk_size`` ``kwargs`` may set). A single process, or a group of one, reads no + knob and takes this load only from a plan passed as ``distribution``; it sets a group of one up + for it, which needs ``torchrun --nproc_per_node=1``. 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 when there is one. @@ -608,11 +609,11 @@ def load_pretrained( every rank loads the plan rank 0 was handed, and a refusal on any rank, of its own arguments, of a checkpoint it cannot see or of the config it reads, raises on every rank. 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. + single process does, the whole model with no plan to agree on and no ``device_map`` default; 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. 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 From 5cac859936fbd1ef1907b5fe1663ed385f0727e8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 07:33:50 +0000 Subject: [PATCH 65/99] Fail every rank when a whole-model load fails on one The whole-model load runs no collective, so a rank whose from_pretrained failed alone, out of memory on its GPU say, left the others to wait for it in the next collective, calibration's or the export's. The load now runs inside the same failure-carrying all_reduce as the refusals before it. test_a_whole_load_that_fails_on_one_rank_fails_every_rank, on two gloo ranks, fails without it. set_distribution_plan's docstring reads more plainly. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 2 +- .../torch/utils/plugins/model_load_utils.py | 14 ++++++--- .../unit/torch/utils/test_model_load_utils.py | 31 +++++++++++++++++++ 3 files changed, 41 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 7febd2fe36a..238b1ea7e69 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -356,7 +356,7 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non Code handed a submodule, such as a VLM's language model, then reads the same plan as code handed the whole model; a module added afterwards records none. ``load_pretrained`` records the plan it - loads by once the model is built; code that distributes a model itself records it here before + loaded by, once the model is built; code that distributes a model itself records it here before calibrating or exporting the model. The plan says how the model is placed over the run; the :class:`ParallelState` a quantized module holds names the process groups its collectives run on. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 31ffc9c1745..5936d3c88e7 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -820,11 +820,15 @@ def handed() -> DistributionPlan | None: # accelerate could pick the same idle GPU. kwargs.setdefault("device_map", {"": f"cuda:{dist.local_rank()}"}) 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, + # It runs no collective, so a rank failing in it alone would leave the others in the next. + with ( + _refused_together("whole-model load", alone), + 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 diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 1b608c9d557..a5cc9ad1ba7 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1160,6 +1160,37 @@ def test_a_load_refused_on_one_rank_fails_every_rank( ) +class _LoadsWholeExceptOnRank1: + """A model class whose whole-model load fails on rank 1 alone, as an out-of-memory GPU would.""" + + @classmethod + def from_pretrained(cls, ckpt_path, **kwargs): + if dist.rank() == 1: + raise OSError("rank 1 could not load the model") + return torch.nn.Linear(1, 1) + + +def _one_rank_fails_to_load_job(rank, size, ckpt): + refused = None + with ( + bounded_collectives(), + mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), + ): + try: # not pytest.raises: its failure would end the pooled worker instead of the job + load_pretrained(ckpt, _LoadsWholeExceptOnRank1, distribution=DistributionPlan(2)) + except Exception as error: + refused = str(error) + # Rank 1's failure on both ranks, rather than rank 0 going on to the next collective alone. + assert refused is not None and "could not load the model" in refused, (rank, refused) + + +def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, tmp_path): + """A whole-model load runs no collective, so a rank failing in it alone would leave the others + to wait for it in the next one, calibration's or the export's.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run(_one_rank_fails_to_load_job, ckpt=str(tmp_path)) + + @pytest.mark.parametrize( ("change", "error", "match"), [ From 25c1017d2f70f50ccb76414bab4b79de27f0da67 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 07:52:24 +0000 Subject: [PATCH 66/99] Refuse a replicating plan on a model sharded after it was recorded load_pretrained records the plan it loaded by, and a plan is recorded once, so a model it returned and the caller then FSDP2-wraps keeps a plan that replicates it, which the export would believe. distribution_plan() now refuses that model, and the docstrings say to wrap a model loaded with from_pretrained instead. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 11 +++++++++++ modelopt/torch/utils/fsdp2_utils.py | 6 ++++-- modelopt/torch/utils/plugins/model_load_utils.py | 3 ++- tests/unit/torch/utils/test_distributed.py | 9 ++++++++- 4 files changed, 25 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 238b1ea7e69..07599bcb209 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -395,8 +395,19 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: model it was added to, so read the plan off that model. A run is its live process group: before a launch sets its group up, each of its processes runs alone with a model of its own, so a model with no plan reads as a single process's there. + + A model FSDP2-sharded under a plan that replicates its dense part raises too: it was sharded + after its plan was recorded, so the plan no longer says how it is placed. """ plan = getattr(model, _PLAN_ATTR, None) + if plan is not None and not plan.is_fsdp_dense and is_fsdp2_model(model): + raise RuntimeError( + f"{type(model).__name__} is FSDP2-sharded, but its recorded plan, {plan!r}, " + "replicates it: it was sharded after its plan was recorded, and a plan is recorded " + "once. To shard a model yourself, load it with from_pretrained rather than " + "load_pretrained, which records the plan it loaded by, then record yours with " + "set_distribution_plan() once it is sharded." + ) if plan is not None: return plan sharded = is_fsdp2_model(model) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index f939fbb36fd..1c25b21d0f8 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -100,8 +100,10 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False replicated and are moved onto the shards' device, which ``fully_shard`` does not do for the params it ignores. - It records no distribution plan: :func:`load_pretrained` does, and a caller wrapping a model - itself records one (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). + It records no distribution plan: a caller wrapping a model itself records one once it is wrapped + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). Wrap a model loaded with + ``from_pretrained``, not one :func:`load_pretrained` returned: that one has the plan it was + loaded by recorded already, and a plan is recorded once. """ # Importing the quantization package registers every mode and plugin; only a wrap needs it. from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 5936d3c88e7..77a3a254d63 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -613,7 +613,8 @@ def load_pretrained( 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. + and export read it back. It is recorded once, so a model to shard yourself (:func:`fsdp2_wrap`) + is one loaded with ``from_pretrained``, not here. 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 diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py index f292843c89a..ad07e4bf5dc 100644 --- a/tests/unit/torch/utils/test_distributed.py +++ b/tests/unit/torch/utils/test_distributed.py @@ -143,13 +143,20 @@ def _no_plan_on_several_ranks_job(rank, size): assert distribution_plan(model) is plan model = _plan_model() set_distribution_plan(model, DistributionPlan(world_size=size)) + fully_shard(model, mesh=init_device_mesh("cpu", (size,))) + _, error = _outcome(lambda: distribution_plan(model)) + assert isinstance(error, RuntimeError), error + assert "sharded after its plan was recorded" in str(error), error + model = _plan_model() + set_distribution_plan(model, DistributionPlan(world_size=size)) model.add_module("added", nn.Linear(2, 2)) _, error = _outcome(lambda: distribution_plan(model.added)) assert isinstance(error, RuntimeError), error def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(cpu_dist_workers): - """On two ranks a model needs a recorded plan, whole, FSDP2-sharded or added to later.""" + """On two ranks a model needs a recorded plan, whole, FSDP2-sharded or added to later, and one + sharded after its plan was recorded is refused rather than read as replicated.""" cpu_dist_workers.run(_no_plan_on_several_ranks_job) From 6ab33c7d5d424328e2640a515c2d53a4f53d6fb1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:00:14 +0000 Subject: [PATCH 67/99] Take config, variant or hf_config passed as None as not passed None is their default, but the FSDP2 load refused it as an option it does not take, and a load without a plan refused hf_config=None as a config without its plan. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 5 +++++ tests/unit/torch/utils/test_model_load_utils.py | 12 +++++++++--- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 77a3a254d63..45bc84cedb7 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -683,6 +683,11 @@ def load_pretrained( f"{os.environ['WORLD_SIZE']}-rank launch, whose process group is not set up; load " "inside distributed_run(), which sets it up." ) + # None is these options' own default, so passed explicitly it is no option: the sharded loads + # refuse the options they do not take, and would refuse it. + for name in ("config", "variant", "hf_config"): + if name in kwargs and kwargs[name] is None: + del kwargs[name] # The config to load with, or, as from_pretrained also takes it, a path to read it from. config = kwargs.get("config") config = config if isinstance(config, (PretrainedConfig, str, os.PathLike)) else None diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index a5cc9ad1ba7..89727b9acb1 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -962,10 +962,16 @@ def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( assert loaders == {} +@pytest.mark.parametrize( + "unset", + [{}, {"config": None}, {"variant": None}, {"hf_config": None}], + ids=["none_passed", "config_none", "variant_none", "hf_config_none"], +) def test_a_load_it_decides_reads_the_config_once( - run_as, loaders, process_group, monkeypatch, tmp_path + run_as, loaders, process_group, monkeypatch, tmp_path, unset ): - """The config the plan was decided from is the one the FSDP2 load gets, not a second read.""" + """The config the plan was decided from is the one the FSDP2 load gets, not a second read. An + option passed as None, its default, is as if not passed, not one the FSDP2 load refuses.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") run_as(2, rank=1) config, reads = LlamaConfig(**_TINY), [] @@ -973,7 +979,7 @@ def test_a_load_it_decides_reads_the_config_once( transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a) or config ) - model = load_pretrained(str(tmp_path)) + model = load_pretrained(str(tmp_path), **unset) assert len(reads) == 1 assert loaders["fsdp2"]["hf_config"] is config From 2966a0ff53e7f3f4fee9adf6d42a39c6b3390fc3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:07:30 +0000 Subject: [PATCH 68/99] Refuse the FSDP2 load on a rank without a GPU with the other checks A rank without one failed inside the load, where the others waited for it in its broadcasts; in a group the caller set up, nothing had checked it before. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 8 ++++ .../unit/torch/utils/test_model_load_utils.py | 42 +++++++++++++------ 2 files changed, 37 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 45bc84cedb7..52393ceaf22 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -564,6 +564,14 @@ def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_o 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`." ) + if not torch.cuda.is_available(): + # Refused here, with the other ranks, rather than inside the load, where they would wait + # for this rank in its broadcasts. + raise RuntimeError( + f"The {route} load shards onto each rank's GPU, and this rank has none. To load it " + "whole on every rank, set MODELOPT_FSDP_SIZE=1 or pass " + f"distribution=DistributionPlan({dist.size()}, dense='replicate')." + ) return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 89727b9acb1..c0db6e3556a 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -897,7 +897,13 @@ def from_pretrained(cls, ckpt_path, **kwargs): @pytest.fixture -def process_group(monkeypatch): +def a_gpu(monkeypatch): + """This rank has a GPU, which the FSDP2 load shards onto; the stand-in loaders never use it.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + +@pytest.fixture +def process_group(monkeypatch, a_gpu): """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) @@ -1117,15 +1123,18 @@ def _fetched_in_step(ckpt_path, **hub_options): def _loaded_in_step(*args, **kwargs): - """The FSDP2 load, as the collectives it runs.""" + """The FSDP2 load, as the collectives it runs, which a rank without a GPU fails before.""" + if not torch.cuda.is_available(): + raise RuntimeError("no GPU to shard onto") torch.distributed.barrier() return _Loaded() -def _one_rank_refused_job(rank, size, ckpt, refused_on_1, match): +def _one_rank_refused_job(rank, size, ckpt, refused_on_1, gpu_on_1, match): refused = None with ( bounded_collectives(), + mock.patch.object(torch.cuda, "is_available", return_value=rank == 0 or gpu_on_1), mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", _loaded_in_step), ): @@ -1143,26 +1152,33 @@ def _one_rank_refused_job(rank, size, ckpt, refused_on_1, match): @pytest.mark.parametrize( - ("refused_on_1", "match"), + ("refused_on_1", "gpu_on_1", "match"), [ - ({"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}, "both name"), - ({"dtype": torch.float32}, "keeps their dtype"), + ({"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}, True, "both name"), + ({"dtype": torch.float32}, True, "keeps their dtype"), ( {"hf_config": LlamaConfig(**_TINY, quantization_config={"quant_method": "fp8"})}, + True, "cannot dequantize", ), + ({}, False, "this rank has none"), ], - ids=["two-configs", "dtype", "quantized-config"], + ids=["two-configs", "dtype", "quantized-config", "no-gpu"], ) def test_a_load_refused_on_one_rank_fails_every_rank( - cpu_dist_workers, tmp_path, refused_on_1, match + cpu_dist_workers, tmp_path, refused_on_1, gpu_on_1, match ): - """Rank 1 alone is refused, by its arguments before the checkpoint is fetched or by the config - it was handed before the FSDP2 load: every rank raises the refusal, rather than leave rank 0 - waiting in the fetch's or the load's collectives, which stand-ins on a bounded group run here.""" + """Rank 1 alone is refused, by its arguments before the checkpoint is fetched, or by the config + it was handed or its lack of a GPU before the FSDP2 load: every rank raises the refusal, rather + than leave rank 0 waiting in the fetch's or the load's collectives, which stand-ins on a bounded + group run here.""" LlamaConfig(**_TINY).save_pretrained(tmp_path) cpu_dist_workers.run( - _one_rank_refused_job, ckpt=str(tmp_path), refused_on_1=refused_on_1, match=match + _one_rank_refused_job, + ckpt=str(tmp_path), + refused_on_1=refused_on_1, + gpu_on_1=gpu_on_1, + match=match, ) @@ -1382,7 +1398,7 @@ def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched( @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 + monkeypatch, serve, loaders, a_gpu, 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.""" From abb748b1eeb1b80a255fdc185af7e227a9e1c41e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:19:12 +0000 Subject: [PATCH 69/99] Default a rank's device_map only to a GPU of its own A distributed whole-model load defaulted every rank to cuda: whenever CUDA was available, so on a host with fewer GPUs than ranks a rank was handed a device it does not have. The default, and the FSDP2 load's refusal, now ask whether the rank has the GPU its local rank numbers. The changelog says a whole-model load on several ranks goes onto each rank's GPU. The two-rank load test masks CUDA, so it no longer depends on the host's GPUs. A model class given by name is refused by name, an ep_size that is no whole number is refused rather than raise a TypeError or be truncated, and the loading hook's config is chosen by None rather than truthiness. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/utils/plugins/model_load_utils.py | 31 +++++--- .../unit/torch/utils/test_model_load_utils.py | 78 ++++++++++++++----- 3 files changed, 80 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 189df7305a2..b856699217e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +43,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 loads a model distributed under ``torchrun``, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument. Until the expert-parallel load lands, an MoE planned expert-parallel is refused; set ``MODELOPT_EP_DEGREE=1`` to load it. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now loads a model distributed under ``torchrun``, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument; a model loaded whole on several ranks goes onto each rank's own GPU unless ``device_map`` places it. Until the expert-parallel load lands, an MoE planned expert-parallel is refused; set ``MODELOPT_EP_DEGREE=1`` to load it. **Backward Breaking Changes** diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 52393ceaf22..d0187dec101 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -499,9 +499,9 @@ def _read_config(path, auto_class, trust_remote_code: bool = False, **hub_option def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide.""" if auto_class not in (None, AutoModelForCausalLM): + name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( - f"The {route} load builds its model class itself; auto_class={auto_class.__name__} " - "does not apply to it." + f"The {route} load builds its model class itself; auto_class={name} does not apply to it." ) # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. for name in ("dtype", "torch_dtype"): @@ -564,17 +564,23 @@ def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_o 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`." ) - if not torch.cuda.is_available(): + if not _has_own_gpu(): # Refused here, with the other ranks, rather than inside the load, where they would wait # for this rank in its broadcasts. raise RuntimeError( - f"The {route} load shards onto each rank's GPU, and this rank has none. To load it " + f"The {route} load shards onto each rank's own GPU, and this rank has none: it sees " + f"{torch.cuda.device_count()} GPU(s) as local rank {dist.local_rank()}. To load it " "whole on every rank, set MODELOPT_FSDP_SIZE=1 or pass " f"distribution=DistributionPlan({dist.size()}, dense='replicate')." ) return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} +def _has_own_gpu() -> bool: + """Whether this rank has a GPU of its own: the one its local rank numbers, as setup() binds.""" + return torch.cuda.is_available() and dist.local_rank() < torch.cuda.device_count() + + def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: """The FSDP2 load, by the arguments :func:`_sharded_load_kwargs` checked.""" if not dist.is_initialized(): @@ -611,7 +617,7 @@ def load_pretrained( 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 when there is one. + whole model, ``device_map`` defaulting to the rank's own GPU when it has one. 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, and a refusal on any rank, of its own arguments, of a @@ -816,7 +822,8 @@ def handed() -> DistributionPlan | None: 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 ( + chosen = hf_config if hf_config is not None else config + model_type = getattr(chosen, "model_type", None) or ( _read_config(ckpt_path, auto_class, trust_remote_code, **hub_options).model_type ) load_kwargs = ( @@ -829,9 +836,10 @@ def handed() -> DistributionPlan | None: # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: - if plan.world_size > 1 and torch.cuda.is_available(): + if plan.world_size > 1 and _has_own_gpu(): # Each rank takes the whole model onto its own GPU. Left to "auto", every rank's - # accelerate could pick the same idle GPU. + # accelerate could pick the same idle GPU; a rank with none loads as from_pretrained + # places it. kwargs.setdefault("device_map", {"": f"cuda:{dist.local_rank()}"}) auto_class = auto_class or AutoModelForCausalLM # It runs no collective, so a rank failing in it alone would leave the others in the next. @@ -1362,7 +1370,9 @@ def _ep_degree(ep_size) -> int | str: return "auto" try: degree = int(value) - except ValueError: + except (ValueError, TypeError): # TypeError: an argument that is no number at all, as [2] + degree = 0 + if isinstance(value, float) and value != degree: # int() truncates a fraction degree = 0 if degree < 1: name = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" @@ -1399,8 +1409,9 @@ def _decide_distribution( 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. + name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( - f"A sharded load builds its model class itself, so auto_class={auto_class.__name__} " + f"A sharded load builds its model class itself, so 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)." diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index c0db6e3556a..3e1334ab036 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -501,6 +501,16 @@ def test_an_ep_degree_that_does_not_divide_the_world_names_where_it_came_from( resolve_distribution_plan(str(tmp_path), **arguments) +@pytest.mark.parametrize("value", [[2], 2.9], ids=["a-list", "a-fraction"]) +def test_an_ep_size_that_is_no_whole_number_is_refused(run_as, serve, tmp_path, value): + """An argument int() fails on with a TypeError, or would truncate, is no degree either.""" + run_as(4) + serve(LlamaConfig(**_TINY)) + + with pytest.raises(ValueError, match=r"ep_size=.*: expected auto, 0 \(auto\) or an EP degree"): + resolve_distribution_plan(str(tmp_path), ep_size=value) + + @pytest.mark.parametrize("value", ["world", "-2"], ids=["word", "below-1"]) @pytest.mark.parametrize( ("world_size", "source"), @@ -898,8 +908,10 @@ def from_pretrained(cls, ckpt_path, **kwargs): @pytest.fixture def a_gpu(monkeypatch): - """This rank has a GPU, which the FSDP2 load shards onto; the stand-in loaders never use it.""" + """This rank has a GPU of its own, which the FSDP2 load shards onto; the stand-in loaders never + use it.""" monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 8) @pytest.fixture @@ -915,23 +927,29 @@ def process_group(monkeypatch, a_gpu): @pytest.mark.parametrize( - ("device_map", "placed"), - [(None, {"": "cuda:1"}), ("cpu", "cpu")], - ids=["defaults_to_the_rank_gpu", "caller_placement_wins"], + ("device_map", "gpus", "placed"), + [ + (None, 4, {"device_map": {"": "cuda:1"}}), + ("cpu", 4, {"device_map": "cpu"}), + (None, 1, {}), + ], + ids=["defaults_to_the_rank_gpu", "caller_placement_wins", "no_gpu_of_its_own"], ) def test_a_replicated_plan_loads_the_whole_model_on_every_rank( - run_as, serve, loaders, auto_class, monkeypatch, tmp_path, device_map, placed + run_as, serve, loaders, auto_class, monkeypatch, tmp_path, device_map, gpus, placed ): """Each rank takes the whole model onto its own GPU; "auto" would let every rank's accelerate - pick the same idle GPU.""" + pick the same idle GPU. A rank with none of its own, on a host with fewer GPUs than ranks, is + left to from_pretrained's placement rather than handed a device it does not have.""" run_as(4, rank=1) serve(LlamaConfig(**_TINY)) - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) # each rank has a GPU + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: gpus) kwargs = {} if device_map is None else {"device_map": device_map} model = load_pretrained(str(tmp_path), auto_class, **kwargs) - assert auto_class.calls == [{"device_map": placed}] + assert auto_class.calls == [placed] assert distribution_plan(model) == DistributionPlan(4) assert loaders == {} @@ -1070,6 +1088,22 @@ class may not build (an encoder-decoder loaded through AutoModel). The decision assert loaders == {} +@pytest.mark.parametrize("handed", [True, False], ids=["handed", "decided"]) +def test_a_model_class_given_by_name_is_refused_by_name( + run_as, serve, loaders, monkeypatch, tmp_path, handed +): + """A sharded load refuses a model class it does not build, given by name too, rather than fail + on the name having no __name__ of its own.""" + run_as(2) + serve(LlamaConfig(**_TINY)) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} + + with pytest.raises(NotImplementedError, match="auto_class='AutoModelForCausalLM'"): + load_pretrained(str(tmp_path), "AutoModelForCausalLM", **plan) + 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) @@ -1083,7 +1117,8 @@ def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed, placed): model, refused = None, None - with bounded_collectives(): + # Ranks without GPUs, whatever the host has: the pooled workers keep no CUDA context. + with bounded_collectives(), mock.patch.object(torch.cuda, "is_available", return_value=False): try: # not pytest.raises: its failure would end the pooled worker instead of the job model = load_pretrained(ckpt, distribution=handed[rank], **placed) except Exception as error: @@ -1124,17 +1159,19 @@ def _fetched_in_step(ckpt_path, **hub_options): def _loaded_in_step(*args, **kwargs): """The FSDP2 load, as the collectives it runs, which a rank without a GPU fails before.""" - if not torch.cuda.is_available(): + if dist.local_rank() >= torch.cuda.device_count(): raise RuntimeError("no GPU to shard onto") torch.distributed.barrier() return _Loaded() -def _one_rank_refused_job(rank, size, ckpt, refused_on_1, gpu_on_1, match): +def _one_rank_refused_job(rank, size, ckpt, refused_on_1, gpus_on_1, match): refused = None + gpus = size if rank == 0 else gpus_on_1 # what this rank's host shows it with ( bounded_collectives(), - mock.patch.object(torch.cuda, "is_available", return_value=rank == 0 or gpu_on_1), + mock.patch.object(torch.cuda, "is_available", return_value=gpus > 0), + mock.patch.object(torch.cuda, "device_count", return_value=gpus), mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", _loaded_in_step), ): @@ -1152,21 +1189,22 @@ def _one_rank_refused_job(rank, size, ckpt, refused_on_1, gpu_on_1, match): @pytest.mark.parametrize( - ("refused_on_1", "gpu_on_1", "match"), + ("refused_on_1", "gpus_on_1", "match"), [ - ({"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}, True, "both name"), - ({"dtype": torch.float32}, True, "keeps their dtype"), + ({"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}, 2, "both name"), + ({"dtype": torch.float32}, 2, "keeps their dtype"), ( {"hf_config": LlamaConfig(**_TINY, quantization_config={"quant_method": "fp8"})}, - True, + 2, "cannot dequantize", ), - ({}, False, "this rank has none"), + ({}, 0, "this rank has none"), + ({}, 1, "this rank has none"), ], - ids=["two-configs", "dtype", "quantized-config", "no-gpu"], + ids=["two-configs", "dtype", "quantized-config", "no-gpu", "fewer-gpus-than-ranks"], ) def test_a_load_refused_on_one_rank_fails_every_rank( - cpu_dist_workers, tmp_path, refused_on_1, gpu_on_1, match + cpu_dist_workers, tmp_path, refused_on_1, gpus_on_1, match ): """Rank 1 alone is refused, by its arguments before the checkpoint is fetched, or by the config it was handed or its lack of a GPU before the FSDP2 load: every rank raises the refusal, rather @@ -1177,7 +1215,7 @@ def test_a_load_refused_on_one_rank_fails_every_rank( _one_rank_refused_job, ckpt=str(tmp_path), refused_on_1=refused_on_1, - gpu_on_1=gpu_on_1, + gpus_on_1=gpus_on_1, match=match, ) From ded136278fdc6586565718c1ff6bb102db1f0a3f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:25:27 +0000 Subject: [PATCH 70/99] Refuse a model recorded as replicated in fsdp2_wrap, not in distribution_plan The check in distribution_plan() walked the whole model on every read of a replicated plan, which the export does per weight. The wrap is where a model gets sharded behind its plan, so fsdp2_wrap refuses one whose recorded plan replicates it, and distribution_plan() is an attribute read again. The single-process FSDP2 load's comment also says the group it sets up outlives the load outside distributed_run. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 11 ------ modelopt/torch/utils/fsdp2_utils.py | 16 ++++++-- .../torch/utils/plugins/model_load_utils.py | 3 +- tests/unit/torch/utils/test_distributed.py | 9 +---- tests/unit/torch/utils/test_fsdp2_utils.py | 37 ++++++++++++++++++- .../unit/torch/utils/test_model_load_utils.py | 19 ---------- 6 files changed, 52 insertions(+), 43 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 07599bcb209..238b1ea7e69 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -395,19 +395,8 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: model it was added to, so read the plan off that model. A run is its live process group: before a launch sets its group up, each of its processes runs alone with a model of its own, so a model with no plan reads as a single process's there. - - A model FSDP2-sharded under a plan that replicates its dense part raises too: it was sharded - after its plan was recorded, so the plan no longer says how it is placed. """ plan = getattr(model, _PLAN_ATTR, None) - if plan is not None and not plan.is_fsdp_dense and is_fsdp2_model(model): - raise RuntimeError( - f"{type(model).__name__} is FSDP2-sharded, but its recorded plan, {plan!r}, " - "replicates it: it was sharded after its plan was recorded, and a plan is recorded " - "once. To shard a model yourself, load it with from_pretrained rather than " - "load_pretrained, which records the plan it loaded by, then record yours with " - "set_distribution_plan() once it is sharded." - ) if plan is not None: return plan sharded = is_fsdp2_model(model) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index 1c25b21d0f8..6a71409d5cc 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -26,6 +26,7 @@ from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed.tensor import DTensor +from modelopt.torch.utils.distributed import _PLAN_ATTR from modelopt.torch.utils.logging import warn_rank_0 __all__ = ["fsdp2_wrap"] @@ -101,10 +102,19 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False params it ignores. It records no distribution plan: a caller wrapping a model itself records one once it is wrapped - (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). Wrap a model loaded with - ``from_pretrained``, not one :func:`load_pretrained` returned: that one has the plan it was - loaded by recorded already, and a plan is recorded once. + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). It refuses a model whose + recorded plan replicates it, as the plan :func:`load_pretrained` records on a model it loads + whole does: a plan is recorded once, so wrap a model loaded with ``from_pretrained`` instead. """ + recorded = getattr(model, _PLAN_ATTR, None) + if recorded is not None and not recorded.is_fsdp_dense: + # Sharded, the model would read as replicated: the export would gather it as if it were. + raise RuntimeError( + f"{type(model).__name__} is recorded as distributed by {recorded!r}, which replicates " + "it, and a plan is recorded once. To shard a model yourself, load it with " + "from_pretrained rather than load_pretrained, which records the plan it loaded by, " + "then record yours with set_distribution_plan() once it is wrapped." + ) # Importing the quantization package registers every mode and plugin; only a wrap needs it. from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d0187dec101..3b5ddf2c4e8 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -584,7 +584,8 @@ def _has_own_gpu() -> bool: def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: """The FSDP2 load, by the arguments :func:`_sharded_load_kwargs` checked.""" if not dist.is_initialized(): - # A single process's group of one; a run (distributed_run) destroys it on exit. + # A single process's group of one, which the sharded model needs for as long as it lives: + # a run (distributed_run) destroys it on exit, and outside one it stays up. dist.setup() return parallel_load_and_prepare_fsdp2( ckpt_path, diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py index ad07e4bf5dc..f292843c89a 100644 --- a/tests/unit/torch/utils/test_distributed.py +++ b/tests/unit/torch/utils/test_distributed.py @@ -143,20 +143,13 @@ def _no_plan_on_several_ranks_job(rank, size): assert distribution_plan(model) is plan model = _plan_model() set_distribution_plan(model, DistributionPlan(world_size=size)) - fully_shard(model, mesh=init_device_mesh("cpu", (size,))) - _, error = _outcome(lambda: distribution_plan(model)) - assert isinstance(error, RuntimeError), error - assert "sharded after its plan was recorded" in str(error), error - model = _plan_model() - set_distribution_plan(model, DistributionPlan(world_size=size)) model.add_module("added", nn.Linear(2, 2)) _, error = _outcome(lambda: distribution_plan(model.added)) assert isinstance(error, RuntimeError), error def test_a_model_with_no_plan_raises_in_a_run_of_several_ranks(cpu_dist_workers): - """On two ranks a model needs a recorded plan, whole, FSDP2-sharded or added to later, and one - sharded after its plan was recorded is refused rather than read as replicated.""" + """On two ranks a model needs a recorded plan, whole, FSDP2-sharded or added to later.""" cpu_dist_workers.run(_no_plan_on_several_ranks_job) diff --git a/tests/unit/torch/utils/test_fsdp2_utils.py b/tests/unit/torch/utils/test_fsdp2_utils.py index b1bcadc933a..6855cb74a31 100644 --- a/tests/unit/torch/utils/test_fsdp2_utils.py +++ b/tests/unit/torch/utils/test_fsdp2_utils.py @@ -13,12 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for ``modelopt.torch.utils.fsdp2_utils``: which parameters stay out of FSDP2's shards.""" +"""Tests for ``modelopt.torch.utils.fsdp2_utils``: which parameters stay out of FSDP2's shards, and +the distribution plan a wrap leaves to its caller.""" import pytest import torch import torch.nn as nn +import modelopt.torch.utils.distributed as dist +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils import fsdp2_utils +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan, set_distribution_plan from modelopt.torch.utils.fsdp2_utils import _off_dtype_params @@ -61,3 +66,33 @@ def test_off_dtype_params_dominant_is_by_numel(params, expected): off = _off_dtype_params(model) kept = {n for n, p in model.named_parameters() if p not in off} assert kept == {expected} + + +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.""" + 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(fsdp2_utils, "fully_shard", lambda module, **kwargs: None) + monkeypatch.setattr(dist, "size", lambda group=None: 4) + + fsdp2_utils.fsdp2_wrap(model) + + with pytest.raises(RuntimeError, match="no distribution plan recorded"): + distribution_plan(model) + + +def test_fsdp2_wrap_refuses_a_model_recorded_as_replicated(monkeypatch): + """A plan is recorded once, so a model recorded as replicated, as load_pretrained records a model + it loads whole, is refused rather than sharded behind its plan, which the export would believe.""" + model = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + set_distribution_plan(model, DistributionPlan()) + monkeypatch.setattr( + LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) + ) + monkeypatch.setattr(fsdp2_utils, "fully_shard", lambda *a, **k: pytest.fail("sharded")) + + with pytest.raises(RuntimeError, match="which replicates it, and a plan is recorded once"): + fsdp2_utils.fsdp2_wrap(model) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 3e1334ab036..8307754f653 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1553,22 +1553,3 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( with pytest.raises(error, match=match): load_pretrained(str(tmp_path), distribution=distribution) assert loaders == {} - - -# --------------------------------------------------------------------------------- fsdp2_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.""" - 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(dist, "size", lambda group=None: 4) - - model_load_utils.fsdp2_wrap(model) - - with pytest.raises(RuntimeError, match="no distribution plan recorded"): - distribution_plan(model) From cc7d2551f729ee0dafbf39b6f5f53e2f635c6ee4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:41:55 +0000 Subject: [PATCH 71/99] Match the load refusals by the option, knob or value they name The refusal tests matched sentence fragments of the messages; each now matches the option, knob or value that tells which refusal fired, so a rewording needs no test edit. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_fsdp2_utils.py | 2 +- .../unit/torch/utils/test_model_load_utils.py | 71 +++++++++++-------- 2 files changed, 41 insertions(+), 32 deletions(-) diff --git a/tests/unit/torch/utils/test_fsdp2_utils.py b/tests/unit/torch/utils/test_fsdp2_utils.py index 6855cb74a31..20b17944ed8 100644 --- a/tests/unit/torch/utils/test_fsdp2_utils.py +++ b/tests/unit/torch/utils/test_fsdp2_utils.py @@ -94,5 +94,5 @@ def test_fsdp2_wrap_refuses_a_model_recorded_as_replicated(monkeypatch): ) monkeypatch.setattr(fsdp2_utils, "fully_shard", lambda *a, **k: pytest.fail("sharded")) - with pytest.raises(RuntimeError, match="which replicates it, and a plan is recorded once"): + with pytest.raises(RuntimeError, match="which replicates it"): fsdp2_utils.fsdp2_wrap(model) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 8307754f653..e9acf36a55e 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -507,7 +507,7 @@ def test_an_ep_size_that_is_no_whole_number_is_refused(run_as, serve, tmp_path, run_as(4) serve(LlamaConfig(**_TINY)) - with pytest.raises(ValueError, match=r"ep_size=.*: expected auto, 0 \(auto\) or an EP degree"): + with pytest.raises(ValueError, match=r"ep_size=.*EP degree"): resolve_distribution_plan(str(tmp_path), ep_size=value) @@ -1043,7 +1043,7 @@ def test_hf_config_is_refused_without_the_plan_it_came_with( run_as(2, rank=1) serve(LlamaConfig(**_TINY)) - with pytest.raises(TypeError, match="returned with a plan that shards the model: pass that"): + with pytest.raises(TypeError, match=r"hf_config.*resolve_distribution_plan"): load_pretrained(str(tmp_path), hf_config=LlamaConfig(**{**_TINY, "num_hidden_layers": 2})) assert loaders == {} @@ -1083,7 +1083,7 @@ class may not build (an encoder-decoder loaded through AutoModel). The decision assert len(auto_class.calls) == 1 else: monkeypatch.setenv("MODELOPT_FSDP_SIZE", fsdp_size) - with pytest.raises(NotImplementedError, match="builds its model class itself"): + with pytest.raises(NotImplementedError, match="auto_class=Auto "): load_pretrained(str(tmp_path), auto_class) assert loaders == {} @@ -1109,7 +1109,7 @@ def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, run_as(2) serve(Qwen3MoeConfig(num_experts=8, moe_intermediate_size=8, **_TINY)) - with pytest.raises(NotImplementedError, match="shards the routed experts"): + with pytest.raises(NotImplementedError, match=r"ep_size=2.*MODELOPT_EP_DEGREE=1"): load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) assert loaders == {} assert auto_class.calls == [] @@ -1127,7 +1127,7 @@ def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed, placed): 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) + assert "expected a DistributionPlan" in str(refused), (rank, refused) @pytest.mark.parametrize( @@ -1191,15 +1191,19 @@ def _one_rank_refused_job(rank, size, ckpt, refused_on_1, gpus_on_1, match): @pytest.mark.parametrize( ("refused_on_1", "gpus_on_1", "match"), [ - ({"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}, 2, "both name"), - ({"dtype": torch.float32}, 2, "keeps their dtype"), + ( + {"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}, + 2, + "config and hf_config", + ), + ({"dtype": torch.float32}, 2, "dtype=torch.float32"), ( {"hf_config": LlamaConfig(**_TINY, quantization_config={"quant_method": "fp8"})}, 2, - "cannot dequantize", + "dequantize", ), - ({}, 0, "this rank has none"), - ({}, 1, "this rank has none"), + ({}, 0, "as local rank 1"), + ({}, 1, "as local rank 1"), ], ids=["two-configs", "dtype", "quantized-config", "no-gpu", "fewer-gpus-than-ranks"], ) @@ -1255,18 +1259,18 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, ("change", "error", "match"), [ ({"kwargs": {"device_map": "auto"}}, TypeError, r"\['device_map'\] do not apply"), - ({"kwargs": {"dtype": torch.bfloat16}}, TypeError, "keeps their dtype"), + ({"kwargs": {"dtype": torch.bfloat16}}, TypeError, r"dtype=torch\.bfloat16"), ( {"kwargs": {"torch_dtype": torch.bfloat16}}, TypeError, - r"keeps their dtype, so torch_dtype=torch\.bfloat16", + r"torch_dtype=torch\.bfloat16", ), - ({"auto_class": AutoModel}, NotImplementedError, "builds its model class itself"), - ({"mtp": True}, NotImplementedError, "cannot build the MTP head"), + ({"auto_class": AutoModel}, NotImplementedError, "auto_class=AutoModel"), + ({"mtp": True}, NotImplementedError, "MTP head"), ( {"kwargs": {"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}}, TypeError, - "config and hf_config both name", + "config and hf_config", ), ], ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs"], @@ -1300,36 +1304,36 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( "causal-lm", None, ValueError, - "is for 4 ranks, but this run has 2", + "is for 4 ranks", ), ( DistributionPlan(2, dense="fsdp"), "quantized", None, NotImplementedError, - "cannot dequantize", + "dequantize", ), ( DistributionPlan(2, dense="fsdp"), "vlm", None, NotImplementedError, - "builds a causal LM, which this llava checkpoint is not", + "llava", ), - ("fsdp2", "causal-lm", None, TypeError, "expected a DistributionPlan or None"), + ("fsdp2", "causal-lm", None, TypeError, "expected a DistributionPlan"), ( DistributionPlan(2), "causal-lm", "hf_config", TypeError, - "returned with a plan that shards the model", + r"hf_config.*resolve_distribution_plan", ), ( DistributionPlan(2), "causal-lm", {"cpu_offload": True}, TypeError, - r"\['cpu_offload'\] only apply to the FSDP2 load", + r"\['cpu_offload'\] only apply", ), ], ids=[ @@ -1382,7 +1386,7 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( 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"): + with pytest.raises(TypeError, match=r"\['device_map'\] place.*MODELOPT_FSDP_SIZE=1"): load_pretrained(str(tmp_path), device_map="auto") assert loaders == {} @@ -1391,24 +1395,29 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( ("distribution", "options", "error", "match"), [ (DistributionPlan(4), {}, ValueError, "is for 4 ranks"), - (DistributionPlan(2), {"hf_config": LlamaConfig(**_TINY)}, TypeError, "plan that shards"), + ( + DistributionPlan(2), + {"hf_config": LlamaConfig(**_TINY)}, + TypeError, + r"hf_config.*resolve_distribution_plan", + ), ( DistributionPlan(2, dense="fsdp"), {"device_map": "auto"}, TypeError, - r"\['device_map'\] do not apply to the FSDP2 load", + r"\['device_map'\] do not apply", ), ( DistributionPlan(2, dense="fsdp"), {"auto_class": AutoModel}, NotImplementedError, - "auto_class=AutoModel does not apply", + "auto_class=AutoModel", ), ( DistributionPlan(2, dense="fsdp"), {"dtype": torch.float32}, TypeError, - "keeps their dtype", + r"dtype=torch\.float32", ), ], ids=[ @@ -1454,7 +1463,7 @@ def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( plan = DistributionPlan(dense="fsdp") if quantized: - with pytest.raises(NotImplementedError, match="cannot dequantize"): + with pytest.raises(NotImplementedError, match="dequantize"): load_pretrained(str(tmp_path), distribution=plan) assert live == [] return @@ -1509,21 +1518,21 @@ def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( DistributionPlan(8), {}, RuntimeError, - "is for 8 ranks, whose process group this 8-rank launch has not", + r"is for 8 ranks.*distributed_run", ), ( DistributionPlan(dense="fsdp"), {}, RuntimeError, - r"shards through this 8-rank launch's .* DistributionPlan\(8, dense='fsdp'\)", + r"DistributionPlan\(8, dense='fsdp'\)", ), ( None, {"MODELOPT_FSDP_SIZE": "world"}, RuntimeError, - "MODELOPT_FSDP_SIZE=world ask to shard", + "MODELOPT_FSDP_SIZE=world", ), - (None, {"MODELOPT_EP_DEGREE": "4"}, RuntimeError, "MODELOPT_EP_DEGREE=4 ask to shard"), + (None, {"MODELOPT_EP_DEGREE": "4"}, RuntimeError, "MODELOPT_EP_DEGREE=4"), (None, {"MODELOPT_EP_DEGREE": "eigth"}, ValueError, "MODELOPT_EP_DEGREE='eigth'"), ], ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob", "ep-knob-not-a-degree"], From 391d5c49435ada412701783ec3b2648d95269e34 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:44:45 +0000 Subject: [PATCH 72/99] Take a load option passed as None as not passed, and file the breaking change device_map=None, as hf_ptq passes it, max_memory, offload_folder or quantization_config at from_pretrained's own default ask nothing, but the FSDP2 load refused them as options it does not take, and the automatic layout as placing a whole model. The CHANGELOG files what changes for a load in a group of several ranks under Backward Breaking Changes, with the knobs or plan that keep it whole. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 3 +- .../torch/utils/plugins/model_load_utils.py | 18 +++++++--- .../unit/torch/utils/test_model_load_utils.py | 35 ++++++++++++++----- 3 files changed, 41 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b856699217e..c92d407a33f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,7 +43,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 loads a model distributed under ``torchrun``, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument; a model loaded whole on several ranks goes onto each rank's own GPU unless ``device_map`` places it. Until the expert-parallel load lands, an MoE planned expert-parallel is refused; set ``MODELOPT_EP_DEGREE=1`` to load it. +- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now loads a model distributed under ``torchrun``, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument. **Backward Breaking Changes** @@ -52,6 +52,7 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded and refuses placement arguments, an MoE planned expert-parallel is refused until the expert-parallel load lands, a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard, and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3b5ddf2c4e8..d034c54e584 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -449,8 +449,12 @@ 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"} + """The load ``kwargs`` that place a whole model, which the FSDP2 load refuses. + + One passed as None, from_pretrained's own default, places nothing. + """ + passed = {name for name, value in kwargs.items() if value is not None} + names = passed - _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) @@ -512,7 +516,9 @@ def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: 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] - {"dtype", "torch_dtype"}) + # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. + passed = {name for name, value in kwargs.items() if value is not None} + unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) if unknown: whole = ( " To place the model yourself, have every rank load it whole: MODELOPT_FSDP_SIZE=1, or " @@ -531,7 +537,8 @@ def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_o from modelopt.torch.models.hf import checkpoint_has_mtp _refuse_sharded_load_options(route, auto_class, kwargs) - kwargs = {k: v for k, v in kwargs.items() if k not in ("dtype", "torch_dtype")} + # Past the refusal, anything else is a dtype the load keeps or an option at its default, None. + kwargs = {k: v for k, v in kwargs.items() if k in _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. @@ -633,7 +640,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``, and refuses any other keyword - rather than ignore it. With neither a plan nor ``MODELOPT_FSDP_SIZE``, rank 0's GPU memory + rather than ignore it, but for one passed as ``None``, ``from_pretrained``'s own default, which + asks nothing. 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. The model is weighed in the checkpoint's own diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index e9acf36a55e..8cba08d4e43 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -955,17 +955,28 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( @pytest.mark.parametrize( - "dtype", - [{"dtype": "auto"}, {"dtype": None}, {"torch_dtype": "auto"}], - ids=["auto", "none", "torch_dtype"], + "defaults", + [ + {"dtype": "auto"}, + {"dtype": None}, + {"torch_dtype": "auto"}, + # As hf_ptq passes them: placement options at from_pretrained's own default. + { + "device_map": None, + "max_memory": None, + "offload_folder": None, + "quantization_config": None, + }, + ], + ids=["auto", "none", "torch_dtype", "placement-none"], ) def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( - run_as, loaders, process_group, monkeypatch, tmp_path, dtype + run_as, loaders, process_group, monkeypatch, tmp_path, defaults ): """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). It keeps the checkpoint's dtype, as from_pretrained does for "auto" and for None, under either - name of the argument.""" + name of the argument, and takes a placement option passed as None as not passed.""" run_as(2, rank=1) served, reads = LlamaConfig(**_TINY), [] monkeypatch.setattr( @@ -974,7 +985,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, local_files_only=True, **dtype + str(tmp_path), distribution=plan, cpu_offload=True, local_files_only=True, **defaults ) call = loaders.pop("fsdp2") @@ -1376,16 +1387,22 @@ def test_a_load_refuses_what_it_cannot_load( assert auto_class.calls == [] +@pytest.mark.parametrize("device_map", ["auto", None], ids=["placed", "none"]) def test_an_automatic_layout_says_so_when_it_refuses_a_placement( - run_as, serve, loaders, monkeypatch, tmp_path + run_as, serve, loaders, process_group, monkeypatch, tmp_path, device_map ): """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) + layout was decided rather than asked for, and how to keep the load whole. device_map=None, as + hf_ptq passes it, places nothing, and the model is sharded.""" + run_as(2, rank=1) serve(LlamaConfig(**_TINY)) monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: 1) + if device_map is None: + load_pretrained(str(tmp_path), AutoModelForCausalLM, device_map=None, dtype="auto") + assert "fsdp2" in loaders + return with pytest.raises(TypeError, match=r"\['device_map'\] place.*MODELOPT_FSDP_SIZE=1"): load_pretrained(str(tmp_path), device_map="auto") assert loaders == {} From cc787a2f0cc449c1708e52a6eeebd1923f473215 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:56:55 +0000 Subject: [PATCH 73/99] Say device_map=None keeps a whole load on the CPU; take mp_policy=None on one The docstring read as if every rank's whole model went onto its GPU; that default holds only when no device_map is passed, and device_map=None, from_pretrained's own, loads on the CPU for the caller to place, which a test row now pins. mp_policy or cpu_offload passed as None, their default, no longer refuses a whole-model load, which made the same call load or raise by rank 0's GPU memory. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 9 +++++- .../unit/torch/utils/test_model_load_utils.py | 29 ++++++++++++------- 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d034c54e584..010d2a4ca5e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -625,7 +625,8 @@ def load_pretrained( 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 own GPU when it has one. + whole model, onto the rank's own GPU when it has one and no ``device_map`` is passed; + ``device_map=None``, ``from_pretrained``'s own, loads it on the CPU for the caller to place. 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, and a refusal on any rank, of its own arguments, of a @@ -802,6 +803,12 @@ def handed() -> DistributionPlan | None: # Options only the FSDP2 load takes would reach from_pretrained, which has no use for them. shared = {"trust_remote_code", "attn_implementation", "hf_config"} # hf_config: see above own = _SHARDED_LOAD_OPTIONS[_FSDP2] - shared + if route is None: + # mp_policy or cpu_offload passed as None, their default, asks nothing of a whole load; + # broadcast_chunk_size=None does ask, to broadcast all of a source's layers at once. + for name in own - {"broadcast_chunk_size"}: + if name in kwargs and kwargs[name] is None: + del kwargs[name] if route is None and (sharded := sorted(set(kwargs) & own)): raise TypeError( f"{sharded} only apply to the FSDP2 load, but {plan!r} gives every rank the whole " diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 8cba08d4e43..df97bb0cc78 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -927,27 +927,36 @@ def process_group(monkeypatch, a_gpu): @pytest.mark.parametrize( - ("device_map", "gpus", "placed"), + ("passed", "gpus", "placed"), [ - (None, 4, {"device_map": {"": "cuda:1"}}), - ("cpu", 4, {"device_map": "cpu"}), - (None, 1, {}), + ({}, 4, {"device_map": {"": "cuda:1"}}), + ({"device_map": "cpu"}, 4, {"device_map": "cpu"}), + ({"device_map": None}, 4, {"device_map": None}), + ({}, 1, {}), + ({"mp_policy": None, "cpu_offload": None}, 4, {"device_map": {"": "cuda:1"}}), + ], + ids=[ + "defaults_to_the_rank_gpu", + "caller_placement_wins", + "none_loads_on_the_cpu", + "no_gpu_of_its_own", + "fsdp2_options_at_none", ], - ids=["defaults_to_the_rank_gpu", "caller_placement_wins", "no_gpu_of_its_own"], ) def test_a_replicated_plan_loads_the_whole_model_on_every_rank( - run_as, serve, loaders, auto_class, monkeypatch, tmp_path, device_map, gpus, placed + run_as, serve, loaders, auto_class, monkeypatch, tmp_path, passed, gpus, placed ): """Each rank takes the whole model onto its own GPU; "auto" would let every rank's accelerate - pick the same idle GPU. A rank with none of its own, on a host with fewer GPUs than ranks, is - left to from_pretrained's placement rather than handed a device it does not have.""" + pick the same idle GPU. device_map=None is from_pretrained's own: a load on the CPU, which the + caller places, as hf_ptq does. A rank with none of its own, on a host with fewer GPUs than + ranks, is left to from_pretrained's placement rather than handed a device it does not have. The + FSDP2 load's own options at None, their default, ask nothing and do not reach from_pretrained.""" run_as(4, rank=1) serve(LlamaConfig(**_TINY)) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "device_count", lambda: gpus) - kwargs = {} if device_map is None else {"device_map": device_map} - model = load_pretrained(str(tmp_path), auto_class, **kwargs) + model = load_pretrained(str(tmp_path), auto_class, **passed) assert auto_class.calls == [placed] assert distribution_plan(model) == DistributionPlan(4) From 28c1a944ce83a714b9e5ddf08d72ab43274217f4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 17:50:38 +0000 Subject: [PATCH 74/99] Point the docstrings' references to fsdp2_wrap at fsdp2_utils, where it lives The loaders import fsdp2_wrap, so a bare reference in their docstrings named no documented function, and the wrap's own reference to load_pretrained named none in its new module. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/fsdp2_utils.py | 6 ++++-- modelopt/torch/utils/plugins/model_load_utils.py | 8 +++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index 6a71409d5cc..e5c84b0f510 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -103,8 +103,10 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False It records no distribution plan: a caller wrapping a model itself records one once it is wrapped (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). It refuses a model whose - recorded plan replicates it, as the plan :func:`load_pretrained` records on a model it loads - whole does: a plan is recorded once, so wrap a model loaded with ``from_pretrained`` instead. + recorded plan replicates it, as the plan + :func:`~modelopt.torch.utils.plugins.model_load_utils.load_pretrained` records on a model it + loads whole does: a plan is recorded once, so wrap a model loaded with ``from_pretrained`` + instead. """ recorded = getattr(model, _PLAN_ATTR, None) if recorded is not None and not recorded.is_fsdp_dense: diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 010d2a4ca5e..cd727f7df65 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -636,8 +636,9 @@ def load_pretrained( 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. It is recorded once, so a model to shard yourself (:func:`fsdp2_wrap`) - is one loaded with ``from_pretrained``, not here. + and export read it back. It is recorded once, so a model to shard yourself + (:func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap`) is one loaded with ``from_pretrained``, + not here. 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 @@ -1207,7 +1208,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). - It records no distribution plan, as :func:`fsdp2_wrap` does not; :func:`load_pretrained` does. + It records no distribution plan, as :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap` does + not; :func:`load_pretrained` does. """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) weight_map = indexed_weight_map(resolved_path) From 0cd4793b9cb6edb45ab1fa47d5e801a76c86f294 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 17:52:31 +0000 Subject: [PATCH 75/99] Sort the wrap tests' imports, and name load_pretrained plainly in the wrap's docstring The wrap's module does not import load_pretrained, so its docstring names it in plain text, as the shorter line keeps the source budget. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/fsdp2_utils.py | 6 ++---- tests/unit/torch/utils/test_fsdp2_utils.py | 6 +++++- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index e5c84b0f510..58b8d95b190 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -103,10 +103,8 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False It records no distribution plan: a caller wrapping a model itself records one once it is wrapped (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). It refuses a model whose - recorded plan replicates it, as the plan - :func:`~modelopt.torch.utils.plugins.model_load_utils.load_pretrained` records on a model it - loads whole does: a plan is recorded once, so wrap a model loaded with ``from_pretrained`` - instead. + recorded plan replicates it, as the plan ``load_pretrained`` records on a model it loads whole + does: a plan is recorded once, so wrap a model loaded with ``from_pretrained`` instead. """ recorded = getattr(model, _PLAN_ATTR, None) if recorded is not None and not recorded.is_fsdp_dense: diff --git a/tests/unit/torch/utils/test_fsdp2_utils.py b/tests/unit/torch/utils/test_fsdp2_utils.py index 20b17944ed8..2df797311ca 100644 --- a/tests/unit/torch/utils/test_fsdp2_utils.py +++ b/tests/unit/torch/utils/test_fsdp2_utils.py @@ -23,7 +23,11 @@ import modelopt.torch.utils.distributed as dist from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils import fsdp2_utils -from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan, set_distribution_plan +from modelopt.torch.utils.distributed import ( + DistributionPlan, + distribution_plan, + set_distribution_plan, +) from modelopt.torch.utils.fsdp2_utils import _off_dtype_params From 36b16aa53c3a902722099a63c36a16646404a904 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:36:15 +0000 Subject: [PATCH 76/99] Refuse a sharded load of a VLM whose language model is quantized The FSDP2 load refused a quantized checkpoint by its top-level quantization_config alone, but a VLM such as Kimi-K2.5 keeps its language model's under text_config, and its remote code builds it as a causal LM, so the load read the packed weights as stored and left the model's weights uninitialized. The check now looks under text_config and llm_config too. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 6 ++++-- .../unit/torch/utils/test_model_load_utils.py | 20 +++++++++++++++---- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index cd727f7df65..af0e461706e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -552,8 +552,10 @@ def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_o 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. + nested = (getattr(hf_config, name, None) for name in ("text_config", "llm_config")) + if any(getattr(c, "quantization_config", None) is not None for c in (hf_config, *nested)): + # It reads the tensors as stored: quantized weights, a VLM's language model's included, + # would load as if they were not. raise NotImplementedError( f"The {route} load cannot dequantize this {model_type} checkpoint; load it unsharded." ) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index df97bb0cc78..3eae8d1a8e1 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1333,6 +1333,13 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( NotImplementedError, "dequantize", ), + ( + DistributionPlan(2, dense="fsdp"), + "quantized-language-model", + {"trust_remote_code": True}, + NotImplementedError, + "dequantize", + ), ( DistributionPlan(2, dense="fsdp"), "vlm", @@ -1359,6 +1366,7 @@ def test_a_sharded_load_refuses_what_it_cannot_honor( ids=[ "another-world", "quantized-checkpoint", + "quantized-language-model", "vlm-checkpoint", "not-a-plan", "hf-config-on-a-whole-load", @@ -1379,14 +1387,18 @@ def test_a_load_refuses_what_it_cannot_load( match, ): """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would - 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, or an option only the - FSDP2 load takes, which from_pretrained would take for a model argument and fail far from here, - or drop.""" + read as stored, a VLM's whose remote code builds it as a causal LM with its language model + quantized included; a VLM, which FSDP2 would build through AutoModelForCausalLM, refused or as + its language model alone; anything but a plan; and the sharded loads' config, or an option only + the FSDP2 load takes, which from_pretrained would take for a model argument and fail far from + here, or drop.""" run_as(2) config = serve(transformers.LlavaConfig() if checkpoint == "vlm" else LlamaConfig(**_TINY)) if checkpoint == "quantized": config.quantization_config = {"quant_method": "fp8"} + if checkpoint == "quantized-language-model": # Kimi-K2.5's layout + config.auto_map = _REMOTE_VLM + config.text_config = LlamaConfig(**_TINY, quantization_config={"quant_method": "fp8"}) if options == "hf_config": options = {"auto_class": auto_class, "hf_config": config} From f42065c1dfd4335eb8b614a5facbb6dc7efb2996 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:36:40 +0000 Subject: [PATCH 77/99] Reword the load refusal test's docstring for the quantized language model Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 3eae8d1a8e1..63e5d22f256 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1387,11 +1387,11 @@ def test_a_load_refuses_what_it_cannot_load( match, ): """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would - read as stored, a VLM's whose remote code builds it as a causal LM with its language model - quantized included; a VLM, which FSDP2 would build through AutoModelForCausalLM, refused or as - its language model alone; anything but a plan; and the sharded loads' config, or an option only - the FSDP2 load takes, which from_pretrained would take for a model argument and fail far from - here, or drop.""" + read as stored, even when only a VLM's language model is quantized and its remote code builds it + as a causal LM; a VLM, which FSDP2 would build through AutoModelForCausalLM, refused or as its + language model alone; anything but a plan; and the sharded loads' config, or an option only the + FSDP2 load takes, which from_pretrained would take for a model argument and fail far from here, + or drop.""" run_as(2) config = serve(transformers.LlavaConfig() if checkpoint == "vlm" else LlamaConfig(**_TINY)) if checkpoint == "quantized": From 9cd575c0fe6ebeb021c24aac14362616c1a16771 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:44:28 +0000 Subject: [PATCH 78/99] Refuse to wrap a model that records any distribution plan fsdp2_wrap refused only a model recorded as replicated, so one load_pretrained had already sharded passed the check and reached FSDP2's own error at its first layer. It now refuses a model that, or any of whose submodules, records a plan, and names the plan and where it is recorded. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/fsdp2_utils.py | 28 ++++++++++++---------- tests/unit/torch/utils/test_fsdp2_utils.py | 20 ++++++++++++---- 2 files changed, 31 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index 58b8d95b190..ba8f5c8c5d6 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -102,19 +102,23 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False params it ignores. It records no distribution plan: a caller wrapping a model itself records one once it is wrapped - (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). It refuses a model whose - recorded plan replicates it, as the plan ``load_pretrained`` records on a model it loads whole - does: a plan is recorded once, so wrap a model loaded with ``from_pretrained`` instead. + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). It refuses a model that, or + any of whose submodules, already records a plan, as a model ``load_pretrained`` returned does: a + plan is recorded once, so wrap a model loaded with ``from_pretrained`` instead. """ - recorded = getattr(model, _PLAN_ATTR, None) - if recorded is not None and not recorded.is_fsdp_dense: - # Sharded, the model would read as replicated: the export would gather it as if it were. - raise RuntimeError( - f"{type(model).__name__} is recorded as distributed by {recorded!r}, which replicates " - "it, and a plan is recorded once. To shard a model yourself, load it with " - "from_pretrained rather than load_pretrained, which records the plan it loaded by, " - "then record yours with set_distribution_plan() once it is wrapped." - ) + for name, module in model.named_modules(): + recorded = getattr(module, _PLAN_ATTR, None) + if recorded is not None: + # The record would no longer hold: a model recorded as replicated would read as + # replicated while sharded, which the export would believe, and one recorded as sharded + # is wrapped already. + where = f"{type(module).__name__} {name!r}" if name else type(module).__name__ + raise RuntimeError( + f"{where} is already recorded as distributed by {recorded!r}, and a plan is " + "recorded once. To shard a model yourself, load it with from_pretrained rather " + "than load_pretrained, which records the plan it loaded by, then record yours with " + "set_distribution_plan() once it is wrapped." + ) # Importing the quantization package registers every mode and plugin; only a wrap needs it. from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector diff --git a/tests/unit/torch/utils/test_fsdp2_utils.py b/tests/unit/torch/utils/test_fsdp2_utils.py index 2df797311ca..c792df630fc 100644 --- a/tests/unit/torch/utils/test_fsdp2_utils.py +++ b/tests/unit/torch/utils/test_fsdp2_utils.py @@ -88,15 +88,25 @@ def test_fsdp2_wrap_leaves_the_plan_to_the_load(monkeypatch): distribution_plan(model) -def test_fsdp2_wrap_refuses_a_model_recorded_as_replicated(monkeypatch): - """A plan is recorded once, so a model recorded as replicated, as load_pretrained records a model - it loads whole, is refused rather than sharded behind its plan, which the export would believe.""" +@pytest.mark.parametrize( + ("plan", "recorded_on"), + [ + (DistributionPlan(), "model"), + (DistributionPlan(4, dense="fsdp"), "model"), + (DistributionPlan(), "submodule"), + ], + ids=["replicated", "sharded", "in-a-submodule"], +) +def test_fsdp2_wrap_refuses_a_model_with_a_plan_recorded(monkeypatch, plan, recorded_on): + """A plan is recorded once, so a model that records one, as load_pretrained records one on every + model it loads, is refused rather than sharded behind it: a replicated model would read as + replicated, which the export would believe, and a sharded one would be wrapped twice.""" model = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) - set_distribution_plan(model, DistributionPlan()) + set_distribution_plan(model if recorded_on == "model" else model[0], plan) monkeypatch.setattr( LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) ) monkeypatch.setattr(fsdp2_utils, "fully_shard", lambda *a, **k: pytest.fail("sharded")) - with pytest.raises(RuntimeError, match="which replicates it"): + with pytest.raises(RuntimeError, match="already recorded as distributed"): fsdp2_utils.fsdp2_wrap(model) From 62c855f121ae4bfd4f33128dcbcbb980b9dbad62 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:44:38 +0000 Subject: [PATCH 79/99] Warn on a rank of a replicated load that has no GPU of its own Such a rank loads the whole model where from_pretrained places it, on the CPU, while its peers load onto their GPUs. It now says so with a warning of its own, rather than one on rank 0, which may have a GPU and nothing to say. The load is otherwise unchanged. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 9 +++++++++ tests/unit/torch/utils/test_model_load_utils.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index af0e461706e..3804d457230 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -19,6 +19,7 @@ import logging import os import re +import warnings from collections import defaultdict from contextlib import contextmanager from itertools import chain @@ -860,6 +861,14 @@ def handed() -> DistributionPlan | None: # accelerate could pick the same idle GPU; a rank with none loads as from_pretrained # places it. kwargs.setdefault("device_map", {"": f"cuda:{dist.local_rank()}"}) + elif plan.world_size > 1 and "device_map" not in kwargs: + # This rank says so itself: rank 0 may have a GPU of its own, and nothing to say. + warnings.warn( + f"Rank {dist.rank()} has no GPU of its own: it sees {torch.cuda.device_count()} " + f"GPU(s) as local rank {dist.local_rank()}, so it loads the whole model where " + "from_pretrained places it, on the CPU, while ranks with one load it onto theirs.", + stacklevel=2, + ) auto_class = auto_class or AutoModelForCausalLM # It runs no collective, so a rank failing in it alone would leave the others in the next. with ( diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 63e5d22f256..81234139a5a 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -963,6 +963,22 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( assert loaders == {} +def test_a_rank_without_a_gpu_of_its_own_says_so( + run_as, serve, loaders, auto_class, monkeypatch, tmp_path +): + """On a host with fewer GPUs than ranks, a rank of a replicated load left without one loads the + model where from_pretrained places it while its peers load onto theirs: it warns itself, as rank + 0 may have a GPU of its own and nothing to say.""" + run_as(4, rank=1) + serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + + with pytest.warns(UserWarning, match="local rank 1"): + load_pretrained(str(tmp_path), auto_class) + assert auto_class.calls == [{}] + + @pytest.mark.parametrize( "defaults", [ From 125a5b99bace61f1c6ffba380f867e7459109318 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:56:03 +0000 Subject: [PATCH 80/99] Let a Qwen3.5 VLM through to the FSDP2 load; offer the knob only for a decided plan The sharded load refused any checkpoint whose whole model is not built through AutoModelForCausalLM, but the FSDP2 load builds the text-only VLMs it supports (Qwen3.5) through their image-text class, as hf_ptq's --use_fsdp2 relies on. It now refuses a checkpoint only when that is not the class the load would build. Its refusals offered MODELOPT_FSDP_SIZE=1 as a way to load whole on every rank, which a plan passed as distribution= ignores. They offer it only for a plan the load decided, and a replicated plan in either case. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 47 ++++++++++++------- .../unit/torch/utils/test_model_load_utils.py | 42 +++++++++++++++-- 2 files changed, 69 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3804d457230..9a4708e1b9c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -501,7 +501,13 @@ def _read_config(path, auto_class, trust_remote_code: bool = False, **hub_option return AutoConfig.from_pretrained(path, trust_remote_code=trust_remote_code, **hub_options) -def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: +def _whole_instead(decided: bool) -> str: + """How to have every rank load the model whole: the knob counts only for a decided plan.""" + plan = f"distribution=DistributionPlan({dist.size()}, dense='replicate')" + return f"MODELOPT_FSDP_SIZE=1 or {plan}" if decided else plan + + +def _refuse_sharded_load_options(route, auto_class, kwargs, decided: bool) -> None: """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide.""" if auto_class not in (None, AutoModelForCausalLM): name = getattr(auto_class, "__name__", repr(auto_class)) @@ -522,8 +528,8 @@ def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) 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')." + " To place the model yourself, have every rank load it whole: " + f"{_whole_instead(decided)}." if route == _FSDP2 else "" ) @@ -533,11 +539,13 @@ def _refuse_sharded_load_options(route, auto_class, kwargs) -> None: ) -def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_options) -> dict: +def _sharded_load_kwargs( + route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided: bool +) -> dict: """What the sharded load ``route`` names takes, refusing what it cannot honor.""" from modelopt.torch.models.hf import checkpoint_has_mtp - _refuse_sharded_load_options(route, auto_class, kwargs) + _refuse_sharded_load_options(route, auto_class, kwargs, decided) # Past the refusal, anything else is a dtype the load keeps or an option at its default, None. kwargs = {k: v for k, v in kwargs.items() if k in _SHARDED_LOAD_OPTIONS[route]} if checkpoint_has_mtp(model_type, ckpt_path): @@ -560,13 +568,17 @@ def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_o raise NotImplementedError( 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. + # It builds as build_meta_causal_lm does: a text-only VLM it supports through its image-text + # class, anything else through AutoModelForCausalLM, which refuses most VLMs' configs and, for + # some, builds the language model alone, whose other towers the export would carry unquantized. + builds = ( + AutoModelForImageTextToText if is_fsdp2_text_only_vlm(hf_config) else AutoModelForCausalLM + ) + if _auto_class(hf_config, trust_remote_code=trust_remote_code) is not builds: raise NotImplementedError( - 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')." + f"The {route} load builds a causal LM, or a text-only VLM it supports, which this " + f"{model_type} checkpoint is not. To load it whole on every rank, pass its model class " + f"as auto_class, with {_whole_instead(decided)}." ) if not dist.is_initialized() and "RANK" not in os.environ: # FSDP2 shards through a process group even on one rank. @@ -580,8 +592,7 @@ def _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_o raise RuntimeError( f"The {route} load shards onto each rank's own GPU, and this rank has none: it sees " f"{torch.cuda.device_count()} GPU(s) as local rank {dist.local_rank()}. To load it " - "whole on every rank, set MODELOPT_FSDP_SIZE=1 or pass " - f"distribution=DistributionPlan({dist.size()}, dense='replicate')." + f"whole on every rank: {_whole_instead(decided)}." ) return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} @@ -620,8 +631,9 @@ def load_pretrained( 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 + * ``dense="fsdp"``: a causal LM, or a text-only VLM the FSDP2 load supports (Qwen3.5), + 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, or a group of one, reads no knob and takes this load only from a plan passed as ``distribution``; it sets a group of one up for it, which needs ``torchrun --nproc_per_node=1``. It reads the checkpoint's tensors as @@ -771,6 +783,7 @@ def handed() -> DistributionPlan | None: for k, v in kwargs.items() if k != "subfolder" and not (k == "config" and config is not None) }, + decided=False, ) # Rank 0 fetches the checkpoint, and each rank then checks it can see it: one that cannot fails # every rank, rather than leave the others in the load's next collective. @@ -849,7 +862,9 @@ def handed() -> DistributionPlan | None: load_kwargs = ( {} if route is None - else _sharded_load_kwargs(route, ckpt_path, auto_class, model_type, kwargs, hub_options) + else _sharded_load_kwargs( + route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided + ) ) if route is not None: model = _load_sharded(ckpt_path, load_kwargs) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 81234139a5a..8bb75e60d33 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1022,6 +1022,40 @@ def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( assert loaders == {} +@pytest.mark.parametrize("model_type", ["qwen3_5", "qwen3_5_moe"]) +def test_a_text_only_vlm_the_fsdp2_load_builds_reaches_it( + run_as, serve, loaders, process_group, tmp_path, model_type +): + """The FSDP2 load builds Qwen3.5's VLMs through their image-text class, as hf_ptq's --use_fsdp2 + does, so a plan that shards one reaches the load rather than being refused as no causal LM.""" + try: + config = serve(AutoConfig.for_model(model_type)) + except ValueError: + pytest.skip("Qwen3.5 requires a recent Transformers") + run_as(2, rank=1) + + load_pretrained(str(tmp_path), distribution=DistributionPlan(2, dense="fsdp")) + + assert loaders.pop("fsdp2")["hf_config"] is config + + +@pytest.mark.parametrize("handed", [True, False], ids=["handed", "decided"]) +def test_a_sharded_load_offers_the_knob_only_for_a_plan_it_decided( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, handed +): + """MODELOPT_FSDP_SIZE counts only for a plan the load decides: to keep a model whole under a + plan of its own, a caller passes a replicated one, which is all the refusal offers it.""" + run_as(2, rank=1) + serve(LlamaConfig(**_TINY)) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} + + with pytest.raises(TypeError, match=r"\['device_map'\] do not apply") as refused: + load_pretrained(str(tmp_path), device_map="auto", **plan) + assert ("MODELOPT_FSDP_SIZE=1" in str(refused.value)) is not handed + assert loaders == {} + + @pytest.mark.parametrize( "unset", [{}, {"config": None}, {"variant": None}, {"hf_config": None}], @@ -1404,10 +1438,10 @@ def test_a_load_refuses_what_it_cannot_load( ): """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would read as stored, even when only a VLM's language model is quantized and its remote code builds it - as a causal LM; a VLM, which FSDP2 would build through AutoModelForCausalLM, refused or as its - language model alone; anything but a plan; and the sharded loads' config, or an option only the - FSDP2 load takes, which from_pretrained would take for a model argument and fail far from here, - or drop.""" + as a causal LM; a VLM but a text-only one it supports, which FSDP2 would build through + AutoModelForCausalLM, refused or as its language model alone; anything but a plan; and the + sharded loads' config, or an option only the FSDP2 load takes, which from_pretrained would take + for a model argument and fail far from here, or drop.""" run_as(2) config = serve(transformers.LlavaConfig() if checkpoint == "vlm" else LlamaConfig(**_TINY)) if checkpoint == "quantized": From c6f9d570616c82c6bf7f8db76c985c907c9b3581 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:58:54 +0000 Subject: [PATCH 81/99] File fsdp2_wrap's refusal of a loaded model as breaking, and tidy two texts load_pretrained now records a plan on every model it returns, a single process's too, and fsdp2_wrap refuses a model that records one, so wrapping a model load_pretrained returned no longer works; the changelog says so and what to do instead. The warning of a rank without a GPU of its own no longer claims its peers have one, and the FSDP2 loader's docstring says a direct caller records its own plan. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + modelopt/torch/utils/plugins/model_load_utils.py | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c92d407a33f..24c23d672a0 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -53,6 +53,7 @@ Changelog - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. - ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded and refuses placement arguments, an MoE planned expert-parallel is refused until the expert-parallel load lands, a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard, and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 9a4708e1b9c..19a18a4d8c1 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -881,7 +881,7 @@ def handed() -> DistributionPlan | None: warnings.warn( f"Rank {dist.rank()} has no GPU of its own: it sees {torch.cuda.device_count()} " f"GPU(s) as local rank {dist.local_rank()}, so it loads the whole model where " - "from_pretrained places it, on the CPU, while ranks with one load it onto theirs.", + "from_pretrained places it without a device_map, on the CPU.", stacklevel=2, ) auto_class = auto_class or AutoModelForCausalLM @@ -1235,7 +1235,9 @@ def parallel_load_and_prepare_fsdp2( pass ``None`` to broadcast all of a source's layers at once). It records no distribution plan, as :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap` does - not; :func:`load_pretrained` does. + not: :func:`load_pretrained` records one, and a caller of this loader records its own + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) before calibrating or + exporting the model. """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) weight_map = indexed_weight_map(resolved_path) From ac1d6df6558da427803556b51cf63ecc60b826ff Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:48:41 +0000 Subject: [PATCH 82/99] Find a recorded plan with one walk, for set_distribution_plan and fsdp2_wrap alike fsdp2_wrap repeated set_distribution_plan's walk for a module that records a plan, read its private attribute, and raised a RuntimeError where set_distribution_plan raises a ValueError for the same record. Both now use _recorded_plan, and fsdp2_wrap raises a ValueError too, so a caller catches one type. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/distributed.py | 24 +++++++++++++------- modelopt/torch/utils/fsdp2_utils.py | 26 ++++++++++------------ tests/unit/torch/utils/test_fsdp2_utils.py | 3 ++- 3 files changed, 30 insertions(+), 23 deletions(-) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 238b1ea7e69..8a5dbdce815 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -351,6 +351,16 @@ def is_fsdp_dense(self) -> bool: _PLAN_ATTR = "_modelopt_distribution_plan" +def _recorded_plan(model: torch.nn.Module) -> tuple[str, DistributionPlan] | None: + """Which of ``model`` and its submodules records a plan, named, and that plan; else ``None``.""" + for name, module in model.named_modules(): + recorded = getattr(module, _PLAN_ATTR, None) + if recorded is not None: + where = f"{type(module).__name__} {name!r}" if name else type(module).__name__ + return where, recorded + return None + + def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> None: """Record that ``model`` is distributed as ``plan``, on it and on every submodule it has now. @@ -374,14 +384,12 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non f"{plan!r} is for {plan.world_size} ranks, but the run it is recorded in has " f"{size()}; a plan's world_size is the ranks that hold the model." ) - for name, module in model.named_modules(): - recorded = getattr(module, _PLAN_ATTR, None) - if recorded is not None: - where = f"{type(module).__name__} {name!r}" if name else type(module).__name__ - raise ValueError( - f"{where} is already recorded as distributed by {recorded!r}; a model is " - f"distributed and its plan recorded once, so {plan!r} cannot be recorded over it." - ) + if (found := _recorded_plan(model)) is not None: + where, recorded = found + raise ValueError( + f"{where} is already recorded as distributed by {recorded!r}; a model is distributed " + f"and its plan recorded once, so {plan!r} cannot be recorded over it." + ) for module in model.modules(): setattr(module, _PLAN_ATTR, plan) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index ba8f5c8c5d6..1289a9e0fea 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -26,7 +26,7 @@ from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed.tensor import DTensor -from modelopt.torch.utils.distributed import _PLAN_ATTR +from modelopt.torch.utils.distributed import _recorded_plan from modelopt.torch.utils.logging import warn_rank_0 __all__ = ["fsdp2_wrap"] @@ -106,19 +106,17 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False any of whose submodules, already records a plan, as a model ``load_pretrained`` returned does: a plan is recorded once, so wrap a model loaded with ``from_pretrained`` instead. """ - for name, module in model.named_modules(): - recorded = getattr(module, _PLAN_ATTR, None) - if recorded is not None: - # The record would no longer hold: a model recorded as replicated would read as - # replicated while sharded, which the export would believe, and one recorded as sharded - # is wrapped already. - where = f"{type(module).__name__} {name!r}" if name else type(module).__name__ - raise RuntimeError( - f"{where} is already recorded as distributed by {recorded!r}, and a plan is " - "recorded once. To shard a model yourself, load it with from_pretrained rather " - "than load_pretrained, which records the plan it loaded by, then record yours with " - "set_distribution_plan() once it is wrapped." - ) + if (found := _recorded_plan(model)) is not None: + # The record would no longer hold: a model recorded as replicated would read as replicated + # while sharded, which the export would believe, and one recorded as sharded is wrapped + # already. A ValueError, as set_distribution_plan raises for the same record. + where, recorded = found + raise ValueError( + f"{where} is already recorded as distributed by {recorded!r}, and a plan is recorded " + "once. To shard a model yourself, load it with from_pretrained rather than " + "load_pretrained, which records the plan it loaded by, then record yours with " + "set_distribution_plan() once it is wrapped." + ) # Importing the quantization package registers every mode and plugin; only a wrap needs it. from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector diff --git a/tests/unit/torch/utils/test_fsdp2_utils.py b/tests/unit/torch/utils/test_fsdp2_utils.py index c792df630fc..71e6fda5f24 100644 --- a/tests/unit/torch/utils/test_fsdp2_utils.py +++ b/tests/unit/torch/utils/test_fsdp2_utils.py @@ -108,5 +108,6 @@ def test_fsdp2_wrap_refuses_a_model_with_a_plan_recorded(monkeypatch, plan, reco ) monkeypatch.setattr(fsdp2_utils, "fully_shard", lambda *a, **k: pytest.fail("sharded")) - with pytest.raises(RuntimeError, match="already recorded as distributed"): + # A ValueError, as set_distribution_plan raises for the same record: one type to catch. + with pytest.raises(ValueError, match="already recorded as distributed"): fsdp2_utils.fsdp2_wrap(model) From 1364da21227345b0df6b547b531c9d042cf69902 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:16:08 +0000 Subject: [PATCH 83/99] Take cpu_offload=False on a whole load; refuse a measured layout's placement in one place A whole-model load refused cpu_offload=False, the FSDP2 load's own default and what hf_ptq passes when --cpu_offload is not set, while it took cpu_offload=None. Each FSDP2-only option now asks nothing at its own default as well as at None, broadcast_chunk_size excepted. The refusal of placement options under a layout rank 0's GPU memory chose moved into _refuse_sharded_load_options, which already decides which options place a whole model, so the two can no longer disagree, and its way out is _whole_instead's. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 63 +++++++++---------- .../unit/torch/utils/test_model_load_utils.py | 6 +- 2 files changed, 34 insertions(+), 35 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 19a18a4d8c1..c6117a51572 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -447,17 +447,10 @@ 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. - - One passed as None, from_pretrained's own default, places nothing. - """ - passed = {name for name, value in kwargs.items() if value is not None} - names = passed - _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) +# The values an FSDP2-only option asks nothing of a whole load with: the FSDP2 load's own default, +# and None where that leaves it unset. broadcast_chunk_size has none: None asks for all of a +# source's layers at once. +_ASKS_NOTHING = {"cpu_offload": (False, None), "mp_policy": (None,)} def _route(plan: DistributionPlan) -> str | None: @@ -507,8 +500,25 @@ def _whole_instead(decided: bool) -> str: return f"MODELOPT_FSDP_SIZE=1 or {plan}" if decided else plan -def _refuse_sharded_load_options(route, auto_class, kwargs, decided: bool) -> None: - """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide.""" +def _refuse_sharded_load_options( + route, auto_class, kwargs, decided: bool, measured: bool = False +) -> None: + """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide. + + ``measured`` says rank 0's GPU memory chose the sharded layout, which nothing asked for. + """ + # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. + passed = {name for name, value in kwargs.items() if value is not None} + unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) + dtypes = [name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")] + if measured and (placed := sorted({*unknown, *dtypes})): + # Say so, rather than refuse the options below as if the caller had asked for this 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 the {route} load " + f"shards it. To load it whole on every rank and place it yourself: " + f"{_whole_instead(decided)}." + ) if auto_class not in (None, AutoModelForCausalLM): name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( @@ -523,9 +533,6 @@ def _refuse_sharded_load_options(route, auto_class, kwargs, decided: bool) -> No f"so {name}={value!r} does not apply to it; FSDP2's mp_policy sets the dtype it " "computes in." ) - # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. - passed = {name for name, value in kwargs.items() if value is not None} - unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) if unknown: whole = ( " To place the model yourself, have every rank load it whole: " @@ -540,12 +547,12 @@ def _refuse_sharded_load_options(route, auto_class, kwargs, decided: bool) -> No def _sharded_load_kwargs( - route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided: bool + route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided: bool, measured: bool ) -> dict: """What the sharded load ``route`` names takes, refusing what it cannot honor.""" from modelopt.torch.models.hf import checkpoint_has_mtp - _refuse_sharded_load_options(route, auto_class, kwargs, decided) + _refuse_sharded_load_options(route, auto_class, kwargs, decided, measured) # Past the refusal, anything else is a dtype the load keeps or an option at its default, None. kwargs = {k: v for k, v in kwargs.items() if k in _SHARDED_LOAD_OPTIONS[route]} if checkpoint_has_mtp(model_type, ckpt_path): @@ -821,27 +828,15 @@ def handed() -> DistributionPlan | None: shared = {"trust_remote_code", "attn_implementation", "hf_config"} # hf_config: see above own = _SHARDED_LOAD_OPTIONS[_FSDP2] - shared if route is None: - # mp_policy or cpu_offload passed as None, their default, asks nothing of a whole load; - # broadcast_chunk_size=None does ask, to broadcast all of a source's layers at once. - for name in own - {"broadcast_chunk_size"}: - if name in kwargs and kwargs[name] is None: + for name, nothing in _ASKS_NOTHING.items(): + if name in kwargs and kwargs[name] in nothing: del kwargs[name] if route is None and (sharded := sorted(set(kwargs) & own)): raise TypeError( f"{sharded} only apply to the FSDP2 load, but {plan!r} gives every rank the whole " "model; pass a plan with dense='fsdp' to shard it." ) - placed = _whole_load_options(kwargs) - if decided and route == _FSDP2 and _dense_from_env(plan.world_size) == "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 " - f"layout, 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 pass " - f"distribution=DistributionPlan({plan.world_size}, dense='replicate')." - ) + measured = decided and route == _FSDP2 and _dense_from_env(plan.world_size) == "auto" if route is not None and config is not None: # The sharded load builds from hf_config, a copy, as building writes to the config. del kwargs["config"] @@ -863,7 +858,7 @@ def handed() -> DistributionPlan | None: {} if route is None else _sharded_load_kwargs( - route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided + route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided, measured ) ) if route is not None: diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 8bb75e60d33..68c8969a6f5 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -934,6 +934,8 @@ def process_group(monkeypatch, a_gpu): ({"device_map": None}, 4, {"device_map": None}), ({}, 1, {}), ({"mp_policy": None, "cpu_offload": None}, 4, {"device_map": {"": "cuda:1"}}), + # As hf_ptq passes --cpu_offload when it is not set: the FSDP2 load's own default. + ({"cpu_offload": False}, 4, {"device_map": {"": "cuda:1"}}), ], ids=[ "defaults_to_the_rank_gpu", @@ -941,6 +943,7 @@ def process_group(monkeypatch, a_gpu): "none_loads_on_the_cpu", "no_gpu_of_its_own", "fsdp2_options_at_none", + "cpu_offload_at_its_default", ], ) def test_a_replicated_plan_loads_the_whole_model_on_every_rank( @@ -950,7 +953,8 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( pick the same idle GPU. device_map=None is from_pretrained's own: a load on the CPU, which the caller places, as hf_ptq does. A rank with none of its own, on a host with fewer GPUs than ranks, is left to from_pretrained's placement rather than handed a device it does not have. The - FSDP2 load's own options at None, their default, ask nothing and do not reach from_pretrained.""" + FSDP2 load's own options left unset or at its defaults ask nothing and do not reach + from_pretrained.""" run_as(4, rank=1) serve(LlamaConfig(**_TINY)) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) From 4b372e0f7e8ca59c78ec15daab86a8370c694a56 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:31:08 +0000 Subject: [PATCH 84/99] Name the way back when the FSDP2 load refuses a quantized or MTP checkpoint Under a layout rank 0's GPU memory chose, which nothing asked for, the FSDP2 load refused a quantized checkpoint, or one carrying an MTP head, with only "load it unsharded". Both refusals now name the way to load it whole, as the load's other refusals do, and the changelog lists them. The docstring says which FSDP2-only options a whole load takes as asking nothing, and the dtype check is written once. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/utils/plugins/model_load_utils.py | 28 ++++++++++--------- .../unit/torch/utils/test_model_load_utils.py | 24 ++++++++++++++++ 3 files changed, 40 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 24c23d672a0..8fa86fc2798 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,7 +52,7 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. -- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded and refuses placement arguments, an MoE planned expert-parallel is refused until the expert-parallel load lands, a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard, and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments and a quantized checkpoint or one carrying an MTP head, an MoE planned expert-parallel is refused until the expert-parallel load lands, a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard, and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. - ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c6117a51572..f7ba08e4868 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -510,6 +510,7 @@ def _refuse_sharded_load_options( # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. passed = {name for name, value in kwargs.items() if value is not None} unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) + # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. dtypes = [name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")] if measured and (placed := sorted({*unknown, *dtypes})): # Say so, rather than refuse the options below as if the caller had asked for this load. @@ -524,15 +525,13 @@ def _refuse_sharded_load_options( raise NotImplementedError( f"The {route} load builds its model class itself; auto_class={name} does not apply to it." ) - # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. - for name in ("dtype", "torch_dtype"): - value = kwargs.get(name) - 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." - ) + if dtypes: # A sharded load keeps the checkpoint's dtype. + name = dtypes[0] + raise TypeError( + f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, so " + f"{name}={kwargs[name]!r} does not apply to it; FSDP2's mp_policy sets the dtype it " + "computes in." + ) if unknown: whole = ( " To place the model yourself, have every rank load it whole: " @@ -559,8 +558,8 @@ def _sharded_load_kwargs( # 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." + f"The {route} load cannot build the MTP head this {model_type} checkpoint carries. To " + f"load it whole on every rank: {_whole_instead(decided)}." ) trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) @@ -573,7 +572,8 @@ def _sharded_load_kwargs( # It reads the tensors as stored: quantized weights, a VLM's language model's included, # would load as if they were not. raise NotImplementedError( - f"The {route} load cannot dequantize this {model_type} checkpoint; load it unsharded." + f"The {route} load cannot dequantize this {model_type} checkpoint. To load it whole on " + f"every rank: {_whole_instead(decided)}." ) # It builds as build_meta_causal_lm does: a text-only VLM it supports through its image-text # class, anything else through AutoModelForCausalLM, which refuses most VLMs' configs and, for @@ -674,7 +674,9 @@ def load_pretrained( :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, or a path to one as ``from_pretrained`` takes, is the one the plan is decided from, and a sharded - load takes it as its ``hf_config``. + load takes it as its ``hf_config``. A whole-model load refuses the options only the FSDP2 load + takes, but for ``cpu_offload`` at ``False`` or ``None`` and ``mp_policy`` at ``None``, which ask + nothing of it; ``broadcast_chunk_size=None`` asks for all of a source's layers at once. Whatever the placement: diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 68c8969a6f5..a404a701a27 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1483,6 +1483,30 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( assert loaders == {} +@pytest.mark.parametrize( + ("checkpoint", "match"), + [("quantized", "dequantize"), ("mtp", "MTP head")], + ids=["quantized", "mtp"], +) +def test_a_measured_layout_names_the_way_back_when_it_refuses_a_checkpoint( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, checkpoint, match +): + """Nothing asked for the layout rank 0's GPU memory chose, so a checkpoint the FSDP2 load cannot + take is refused with the way to load it whole, as every refusal of that load says.""" + run_as(2, rank=1) + config = serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: 1) + if checkpoint == "quantized": + config.quantization_config = {"quant_method": "fp8"} + else: + monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) + + with pytest.raises(NotImplementedError, match=match) as refused: + load_pretrained(str(tmp_path)) + assert "MODELOPT_FSDP_SIZE=1" in str(refused.value) + assert loaders == {} + + @pytest.mark.parametrize( ("distribution", "options", "error", "match"), [ From 8c489b9324d539df1bf83d0f9b650efc5a07c5c5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:48:06 +0000 Subject: [PATCH 85/99] Take a rank's GPU as the one it is bound to, as torch's DeviceMesh does _has_own_gpu() read a rank's GPU as the one its local rank numbers, so a launch that exposes one GPU to each process, each bound to cuda:0, read every rank but the first as having none: the FSDP2 load refused it, and a whole load left it on the CPU with a warning. A missing LOCAL_RANK on later nodes did the same. _own_gpu_index() takes the device the process has set up once CUDA is initialized, else the one its local rank numbers, and the device_map default and the FSDP2 load place onto it. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 30 +++++++++--- .../unit/torch/utils/test_model_load_utils.py | 48 +++++++++++++++++++ 2 files changed, 71 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f7ba08e4868..462d9de5b22 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -604,9 +604,24 @@ def _sharded_load_kwargs( return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} +def _own_gpu_index() -> int | None: + """This rank's GPU, as torch's DeviceMesh takes it, or ``None`` for a rank with none there. + + Once the process has set CUDA up (setup(), a device-bound group, the caller's own set_device), + the device it set; before, the one its local rank numbers. A launch that exposes one GPU to + each process, each bound to cuda:0, so has a GPU on every rank, whatever its local rank. + """ + if not torch.cuda.is_available(): + return None + if torch.cuda.is_initialized(): + return torch.cuda.current_device() + index = dist.local_rank() + return index if index < torch.cuda.device_count() else None + + def _has_own_gpu() -> bool: - """Whether this rank has a GPU of its own: the one its local rank numbers, as setup() binds.""" - return torch.cuda.is_available() and dist.local_rank() < torch.cuda.device_count() + """Whether this rank has a GPU of its own (:func:`_own_gpu_index`).""" + return _own_gpu_index() is not None def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: @@ -617,7 +632,7 @@ def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: dist.setup() return parallel_load_and_prepare_fsdp2( ckpt_path, - torch.device("cuda", dist.local_rank()), + torch.device("cuda", _own_gpu_index()), dist.rank(), dist.size(), **load_kwargs, @@ -647,8 +662,9 @@ def load_pretrained( 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, onto the rank's own GPU when it has one and no ``device_map`` is passed; - ``device_map=None``, ``from_pretrained``'s own, loads it on the CPU for the caller to place. + whole model, onto the rank's own GPU (the one it is bound to, else the one its local rank + numbers) when it has one and no ``device_map`` is passed; ``device_map=None``, + ``from_pretrained``'s own, loads it on the CPU for the caller to place. 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, and a refusal on any rank, of its own arguments, of a @@ -868,11 +884,11 @@ def handed() -> DistributionPlan | None: # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: - if plan.world_size > 1 and _has_own_gpu(): + if plan.world_size > 1 and (index := _own_gpu_index()) is not None: # Each rank takes the whole model onto its own GPU. Left to "auto", every rank's # accelerate could pick the same idle GPU; a rank with none loads as from_pretrained # places it. - kwargs.setdefault("device_map", {"": f"cuda:{dist.local_rank()}"}) + kwargs.setdefault("device_map", {"": f"cuda:{index}"}) elif plan.world_size > 1 and "device_map" not in kwargs: # This rank says so itself: rank 0 may have a GPU of its own, and nothing to say. warnings.warn( diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index a404a701a27..3aabe5cc8f4 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -967,6 +967,24 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( assert loaders == {} +def test_a_rank_bound_to_its_one_visible_gpu_loads_onto_it( + run_as, serve, loaders, auto_class, monkeypatch, tmp_path +): + """A launch that exposes one GPU to each process binds every rank to cuda:0, whatever its local + rank: each takes the whole model onto the GPU it is bound to, rather than read as having + none.""" + run_as(4, rank=1) + serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + + load_pretrained(str(tmp_path), auto_class) + + assert auto_class.calls == [{"device_map": {"": "cuda:0"}}] + + def test_a_rank_without_a_gpu_of_its_own_says_so( run_as, serve, loaders, auto_class, monkeypatch, tmp_path ): @@ -1298,6 +1316,36 @@ def test_a_load_refused_on_one_rank_fails_every_rank( ) +def _one_visible_gpu_per_process_job(rank, size, ckpt): + devices = [] + + def load(ckpt_path, device, *args, **kwargs): + devices.append(device) + torch.distributed.barrier() + return _Loaded() + + with ( + bounded_collectives(), + # Each rank sees only its own GPU, which its process bound as cuda:0, whatever its local + # rank. + mock.patch.object(torch.cuda, "is_available", return_value=True), + mock.patch.object(torch.cuda, "device_count", return_value=1), + mock.patch.object(torch.cuda, "is_initialized", return_value=True), + mock.patch.object(torch.cuda, "current_device", return_value=0), + mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), + mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", load), + ): + load_pretrained(ckpt, distribution=DistributionPlan(2, dense="fsdp")) + assert devices == [torch.device("cuda", 0)], (rank, devices) + + +def test_one_visible_gpu_per_process_shards_on_every_rank(cpu_dist_workers, tmp_path): + """Two ranks, each shown only its own GPU and bound to it as cuda:0, both shard onto it: rank 1, + whose local rank numbers no GPU it sees, is not refused as having none.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run(_one_visible_gpu_per_process_job, ckpt=str(tmp_path)) + + class _LoadsWholeExceptOnRank1: """A model class whose whole-model load fails on rank 1 alone, as an out-of-memory GPU would.""" From 6b6ecb261e6be4f8b46a13101e6f9e170486e5f1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:56:25 +0000 Subject: [PATCH 86/99] Decide an EP degree left to load_pretrained as 1; take broadcast_chunk_size=8 as nothing With MODELOPT_EP_DEGREE unset, the plan decision picked an expert-parallel degree for an MoE, which load_pretrained then refused, as no load here shards the routed experts yet: a multi-rank load of an MoE with no knob set failed. It now decides such a degree as 1 and loads; a degree MODELOPT_EP_DEGREE asks for, as a plan handed in, is still refused. A whole-model load also takes broadcast_chunk_size=8, the FSDP2 load's own default, as asking nothing, so a caller can forward every FSDP2 option at its default whatever the plan. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/utils/plugins/model_load_utils.py | 21 ++++++++++----- .../unit/torch/utils/test_model_load_utils.py | 27 +++++++++++++++++++ 3 files changed, 42 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8fa86fc2798..e81fe4aa66f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,7 +52,7 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. -- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments and a quantized checkpoint or one carrying an MTP head, an MoE planned expert-parallel is refused until the expert-parallel load lands, a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard, and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints and ones carrying an MTP head; an MoE loads without expert parallelism, and is refused when ``MODELOPT_EP_DEGREE`` asks for it, until the expert-parallel load lands; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. - ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 462d9de5b22..e19dd74404d 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -447,10 +447,10 @@ def record_unplaced_source_keys( } ), } -# The values an FSDP2-only option asks nothing of a whole load with: the FSDP2 load's own default, -# and None where that leaves it unset. broadcast_chunk_size has none: None asks for all of a -# source's layers at once. -_ASKS_NOTHING = {"cpu_offload": (False, None), "mp_policy": (None,)} +# The values an FSDP2-only option asks nothing of a whole load with: the FSDP2 load's own default +# (parallel_load_and_prepare_fsdp2's), and None where that leaves it unset. A None +# broadcast_chunk_size asks for something: all of a source's layers at once. +_ASKS_NOTHING = {"cpu_offload": (False, None), "mp_policy": (None,), "broadcast_chunk_size": (8,)} def _route(plan: DistributionPlan) -> str | None: @@ -666,7 +666,8 @@ def load_pretrained( numbers) when it has one and no ``device_map`` is passed; ``device_map=None``, ``from_pretrained``'s own, loads it on the CPU for the caller to place. - A plan with ``ep_size > 1`` raises: no load here shards the routed experts. On a distributed run + A plan with ``ep_size > 1`` raises: no load here shards the routed experts, so it decides an EP + degree left to it as 1, and refuses one ``MODELOPT_EP_DEGREE`` asks for. On a distributed run every rank loads the plan rank 0 was handed, and a refusal on any rank, of its own arguments, of a checkpoint it cannot see or of the config it reads, raises on every rank. A launch whose processes set no process group up (a launcher that only exports ``WORLD_SIZE``) has each process load as a @@ -691,8 +692,10 @@ def load_pretrained( builds from it as given, so without that plan it is refused. A ``config`` passed for the load, or a path to one as ``from_pretrained`` takes, is the one the plan is decided from, and a sharded load takes it as its ``hf_config``. A whole-model load refuses the options only the FSDP2 load - takes, but for ``cpu_offload`` at ``False`` or ``None`` and ``mp_policy`` at ``None``, which ask - nothing of it; ``broadcast_chunk_size=None`` asks for all of a source's layers at once. + takes, but at the FSDP2 load's own defaults (``cpu_offload=False``, ``mp_policy=None``, + ``broadcast_chunk_size=8``) or ``cpu_offload=None``, which ask nothing of it, so a caller can + forward them whatever the plan; ``broadcast_chunk_size=None`` asks for all of a source's layers + at once. Whatever the placement: @@ -831,6 +834,10 @@ def handed() -> DistributionPlan | None: config=config, auto_class=auto_class, trust_remote_code=trust_remote_code, + # No load here shards the routed experts yet: an EP degree left to the decision is 1, + # rather than a layout this function would refuse; one MODELOPT_EP_DEGREE asks for, as + # a plan handed in, is still refused. + ep_size=1 if _ep_degree(None) == "auto" else None, **hub_options, ) if _route(plan) is not None: diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 3aabe5cc8f4..47c17c9dd7b 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -936,6 +936,7 @@ def process_group(monkeypatch, a_gpu): ({"mp_policy": None, "cpu_offload": None}, 4, {"device_map": {"": "cuda:1"}}), # As hf_ptq passes --cpu_offload when it is not set: the FSDP2 load's own default. ({"cpu_offload": False}, 4, {"device_map": {"": "cuda:1"}}), + ({"broadcast_chunk_size": 8}, 4, {"device_map": {"": "cuda:1"}}), ], ids=[ "defaults_to_the_rank_gpu", @@ -944,6 +945,7 @@ def process_group(monkeypatch, a_gpu): "no_gpu_of_its_own", "fsdp2_options_at_none", "cpu_offload_at_its_default", + "broadcast_chunk_size_at_its_default", ], ) def test_a_replicated_plan_loads_the_whole_model_on_every_rank( @@ -1196,6 +1198,31 @@ def test_a_model_class_given_by_name_is_refused_by_name( assert loaders == {} +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize("knob", [None, "4"], ids=["left-to-decide", "asked-for"]) +def test_a_load_it_decides_keeps_the_experts_whole( + run_as, serve, loaders, monkeypatch, tmp_path, knob +): + """No load here shards the routed experts yet, so an MoE whose EP degree is left to the decision + is decided without expert parallelism and loads, where an automatic degree would be one the + load refuses; a degree MODELOPT_EP_DEGREE asks for is still refused, not ignored.""" + run_as(4) + serve(_qwen3_moe()) + if knob is not None: + monkeypatch.setenv("MODELOPT_EP_DEGREE", knob) + monkeypatch.setattr( + transformers.AutoModelForCausalLM, + "from_pretrained", + classmethod(lambda cls, *a, **k: torch.nn.Linear(1, 1)), + ) + + if knob is not None: + with pytest.raises(NotImplementedError, match=r"ep_size=4.*MODELOPT_EP_DEGREE=1"): + load_pretrained(str(tmp_path)) + return + assert distribution_plan(load_pretrained(str(tmp_path))) == DistributionPlan(4) + + 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 6ed852067abc614d20e9f0d37d04a9881820b3e1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:30:17 +0000 Subject: [PATCH 87/99] Read a blank distribution knob as unset, normalized where it is read MODELOPT_EP_DEGREE or MODELOPT_FSDP_SIZE set to whitespace only stripped to an empty value both parses refused. env.py now strips both knobs and reads a blank one as auto, as it reads an empty one, so the parses need no strip of their own. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/env.py | 8 ++++---- modelopt/torch/utils/plugins/model_load_utils.py | 4 ++-- tests/unit/torch/utils/test_model_load_utils.py | 13 ++++++++----- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/env.py b/modelopt/torch/utils/env.py index 7e27e7857c3..89f869c3efb 100644 --- a/modelopt/torch/utils/env.py +++ b/modelopt/torch/utils/env.py @@ -35,10 +35,10 @@ def ep_degree() -> str: - """Expert-parallel degree knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``.""" - return os.environ.get("MODELOPT_EP_DEGREE") or "auto" + """Expert-parallel degree knob for ``resolve_distribution_plan``; blank/unset -> ``"auto"``.""" + return (os.environ.get("MODELOPT_EP_DEGREE") or "").strip() or "auto" def fsdp_size() -> str: - """Dense-sharding knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``.""" - return os.environ.get("MODELOPT_FSDP_SIZE") or "auto" + """Dense-sharding knob for ``resolve_distribution_plan``; blank/unset -> ``"auto"``.""" + return (os.environ.get("MODELOPT_FSDP_SIZE") or "").strip() or "auto" diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index e19dd74404d..0f22e77d784 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1416,7 +1416,7 @@ def _gpu_memory() -> int | None: def _dense_from_env(world_size: int) -> str: """``MODELOPT_FSDP_SIZE`` as a dense layout, or ``"auto"``.""" - value = env.fsdp_size().strip().lower() + value = env.fsdp_size().lower() if value == "auto": return "auto" if value == "1": @@ -1432,7 +1432,7 @@ def _dense_from_env(world_size: int) -> str: def _ep_degree(ep_size) -> int | str: """``ep_size`` as an EP degree or ``"auto"``; ``None`` reads ``MODELOPT_EP_DEGREE``.""" value = env.ep_degree() if ep_size is None else ep_size - if str(value).strip().lower() in ("auto", "0"): + if str(value).lower() in ("auto", "0"): return "auto" try: degree = int(value) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 47c17c9dd7b..277aeb42781 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -536,10 +536,12 @@ def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from( @needs_qwen3_moe_ep_plan @pytest.mark.parametrize( - "value", ["0", "", "AUTO", " auto "], ids=["zero", "exported-empty", "capitals", "padded"] + "value", + ["0", "", " ", "AUTO", " auto "], + ids=["zero", "exported-empty", "exported-blank", "capitals", "padded"], ) def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value): - """And so is one exported empty, and auto in any case or padded: the plan decides.""" + """And so is one exported empty or blank, and auto in any case or padded: the plan decides.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", value) run_as(4) serve(_qwen3_moe(num_experts=8)) @@ -619,16 +621,17 @@ def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path): ("4", "fsdp"), ("WORLD", "fsdp"), ("", "fsdp"), + (" ", "fsdp"), (" world ", "fsdp"), ("2", None), ], - ids=["1", "4", "WORLD", "exported-empty", "padded", "part-of-the-world"], + ids=["1", "4", "WORLD", "exported-empty", "exported-blank", "padded", "part-of-the-world"], ) def test_the_dense_knob_replicates_or_shards_over_the_whole_world( run_as, serve, gpu_memory, monkeypatch, tmp_path, value, dense ): - """Exported empty, the knob is auto, which shards a model too big for one GPU. Sharding over - part of the world is refused: no layout does that.""" + """Exported empty or blank, the knob is auto, which shards a model too big for one GPU. Sharding + over part of the world is refused: no layout does that.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", value) run_as(4) serve(LlamaConfig(**_TINY)) From 10d85bda9f42cf4ae3c080700d9cecaddc899225 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:30:35 +0000 Subject: [PATCH 88/99] Say a sharded load takes no VLM class, but builds a text-only VLM's itself A sharded load supports a text-only VLM, which it builds through its image-text class, yet refuses any auto_class but AutoModelForCausalLM. The docstring now says so: leave auto_class out for such a VLM. 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 0f22e77d784..d01df2af0fc 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -679,7 +679,9 @@ def load_pretrained( (:func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap`) is one loaded with ``from_pretrained``, not here. - A sharded load builds the model class and places it itself. Besides its own options it takes + A sharded load builds the model class and places it itself: a text-only VLM through its + image-text class, anything else through ``AutoModelForCausalLM``, the one ``auto_class`` it + takes, so leave ``auto_class`` out for a VLM. Besides its own options it takes ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other keyword rather than ignore it, but for one passed as ``None``, ``from_pretrained``'s own default, which asks nothing. With neither a plan nor ``MODELOPT_FSDP_SIZE``, rank 0's GPU memory From 91f7ab44209a675dcba738cc04d6610182904d75 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:01:19 +0000 Subject: [PATCH 89/99] Say device_map=None places nothing, so it keeps no measured layout whole The docstring said a call that places the model loads only where it fits one GPU, which reads as covering device_map=None. An option at from_pretrained's own default places nothing: the whole load keeps device_map=None on the CPU, while a layout rank 0's GPU memory chose shards the model. Say so. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d01df2af0fc..1de0bba6122 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -687,17 +687,18 @@ def load_pretrained( asks nothing. 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. The model is weighed in the checkpoint's own - dtype (its config's, else bfloat16), not one the load asks for: hand a load that widens it, - ``dtype=torch.float32`` say, its plan. ``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, or - a path to one as ``from_pretrained`` takes, is the one the plan is decided from, and a sharded - load takes it as its ``hf_config``. A whole-model load refuses the options only the FSDP2 load - takes, but at the FSDP2 load's own defaults (``cpu_offload=False``, ``mp_policy=None``, - ``broadcast_chunk_size=8``) or ``cpu_offload=None``, which ask nothing of it, so a caller can - forward them whatever the plan; ``broadcast_chunk_size=None`` asks for all of a source's layers - at once. + or set ``MODELOPT_FSDP_SIZE=1`` to keep it whole. An option at ``from_pretrained``'s own + default, ``device_map=None`` included, places nothing and does not keep the load whole. The + model is weighed in the checkpoint's own dtype (its config's, else bfloat16), not one the load + asks for: hand a load that widens it, ``dtype=torch.float32`` say, its plan. ``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, or a path to one as ``from_pretrained`` takes, is the one the + plan is decided from, and a sharded load takes it as its ``hf_config``. A whole-model load + refuses the options only the FSDP2 load takes, but at the FSDP2 load's own defaults + (``cpu_offload=False``, ``mp_policy=None``, ``broadcast_chunk_size=8``) or ``cpu_offload=None``, + which ask nothing of it, so a caller can forward them whatever the plan; + ``broadcast_chunk_size=None`` asks for all of a source's layers at once. Whatever the placement: From bb07c3bc823a59823bbb83a740d47fd83a596aeb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 23:01:29 +0000 Subject: [PATCH 90/99] Name the way out of the EP-knob and auto_class refusals A launch without a group that MODELOPT_EP_DEGREE asks to shard was told to load inside distributed_run(), where the plan's EP degree is refused next, as no load here shards the routed experts: the refusal now names MODELOPT_EP_DEGREE=1 as well. A sharded load's refusal of an auto_class was the one with no way out: it now says to leave the class out, or to load the model whole with it, and both auto_class refusals say which class the load builds, a text-only VLM's image-text class included. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 23 +++++++++++++------ .../unit/torch/utils/test_model_load_utils.py | 13 +++++++++-- 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1de0bba6122..c4ff36d7dae 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -523,7 +523,9 @@ def _refuse_sharded_load_options( if auto_class not in (None, AutoModelForCausalLM): name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( - f"The {route} load builds its model class itself; auto_class={name} does not apply to it." + f"The {route} load builds its model class itself, a causal LM's or a text-only VLM's " + f"image-text class, so it takes no auto_class={name}: leave it out to have the load " + f"pick, or load the model whole with it: {_whole_instead(decided)}." ) if dtypes: # A sharded load keeps the checkpoint's dtype. name = dtypes[0] @@ -745,14 +747,21 @@ def load_pretrained( # The knobs ask to shard over the launch, rather than leave each process the whole model. A # value that is not an EP degree is refused, by the plan decision's own parse. ep, world = _ep_degree(None), int(os.environ["WORLD_SIZE"]) - asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if ep != "auto" and ep > 1 else [] + experts = ep != "auto" and ep > 1 + asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if experts else [] if _dense_from_env(world) == "fsdp": asked.append(f"MODELOPT_FSDP_SIZE={env.fsdp_size()}") if asked: + # The group alone would not do for the EP knob: no load here shards the experts. 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." + "inside distributed_run(), which sets it up" + + ( + ", with MODELOPT_EP_DEGREE=1: no load here shards the routed experts." + if experts + else "." + ) ) # None is these options' own default, so passed explicitly it is no option: the sharded loads # refuse the options they do not take, and would refuse it. @@ -1480,10 +1489,10 @@ def _decide_distribution( # than build the model of the decision's own pick of class, which may not build at all. name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( - f"A sharded load builds its model class itself, so 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)." + f"A sharded load builds its model class itself, a causal LM's or a text-only VLM's " + f"image-text class, so it takes no auto_class={name} while ep_size={ep_size!r}, " + f"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. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 277aeb42781..99eadbbd837 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1417,7 +1417,11 @@ def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, TypeError, r"torch_dtype=torch\.bfloat16", ), - ({"auto_class": AutoModel}, NotImplementedError, "auto_class=AutoModel"), + ( + {"auto_class": AutoModel}, + NotImplementedError, + r"auto_class=AutoModel: leave it out.*DistributionPlan\(2, dense='replicate'\)", + ), ({"mtp": True}, NotImplementedError, "MTP head"), ( {"kwargs": {"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}}, @@ -1726,7 +1730,12 @@ def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( RuntimeError, "MODELOPT_FSDP_SIZE=world", ), - (None, {"MODELOPT_EP_DEGREE": "4"}, RuntimeError, "MODELOPT_EP_DEGREE=4"), + ( + None, + {"MODELOPT_EP_DEGREE": "4"}, + RuntimeError, + r"MODELOPT_EP_DEGREE=4 .*distributed_run\(\).*MODELOPT_EP_DEGREE=1", + ), (None, {"MODELOPT_EP_DEGREE": "eigth"}, ValueError, "MODELOPT_EP_DEGREE='eigth'"), ], ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob", "ep-knob-not-a-degree"], From d926a4c9fa3523da99f8eecb927afbf4432911b3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:18:09 +0000 Subject: [PATCH 91/99] Pin the CUDA state the GPU check reads in the loader's unit tests _own_gpu_index reads torch.cuda.is_initialized() and current_device(), but the fixtures, the standalone placement tests and the pooled refusal workers pinned only availability and the device count, so a host process that had set CUDA up would hand them cuda:0 where they expect a local rank's GPU. One helper, _pin_cuda, pins all four: unbound for the local-rank cases, bound for the bound-device ones. Signed-off-by: Shengliang Xu --- .../unit/torch/utils/test_model_load_utils.py | 38 ++++++++++--------- 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 99eadbbd837..ab5253492ab 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -909,12 +909,21 @@ def from_pretrained(cls, ckpt_path, **kwargs): return Auto +def _pin_cuda(monkeypatch, gpus, bound=None): + """Pin the CUDA state a rank's GPU is read from: ``gpus`` visible, and the process bound to + cuda:``bound``, or to none yet, so its local rank picks one. Tests then read the case rather + than the host, whose process may have set CUDA up already.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: gpus > 0) + monkeypatch.setattr(torch.cuda, "device_count", lambda: gpus) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: bound is not None) + monkeypatch.setattr(torch.cuda, "current_device", lambda: bound or 0) + + @pytest.fixture def a_gpu(monkeypatch): """This rank has a GPU of its own, which the FSDP2 load shards onto; the stand-in loaders never use it.""" - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "device_count", lambda: 8) + _pin_cuda(monkeypatch, 8) @pytest.fixture @@ -962,8 +971,7 @@ def test_a_replicated_plan_loads_the_whole_model_on_every_rank( from_pretrained.""" run_as(4, rank=1) serve(LlamaConfig(**_TINY)) - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "device_count", lambda: gpus) + _pin_cuda(monkeypatch, gpus) model = load_pretrained(str(tmp_path), auto_class, **passed) @@ -980,10 +988,7 @@ def test_a_rank_bound_to_its_one_visible_gpu_loads_onto_it( none.""" run_as(4, rank=1) serve(LlamaConfig(**_TINY)) - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) - monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) - monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + _pin_cuda(monkeypatch, 1, bound=0) load_pretrained(str(tmp_path), auto_class) @@ -998,8 +1003,7 @@ def test_a_rank_without_a_gpu_of_its_own_says_so( 0 may have a GPU of its own and nothing to say.""" run_as(4, rank=1) serve(LlamaConfig(**_TINY)) - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + _pin_cuda(monkeypatch, 1) with pytest.warns(UserWarning, match="local rank 1"): load_pretrained(str(tmp_path), auto_class) @@ -1292,11 +1296,11 @@ def _one_rank_refused_job(rank, size, ckpt, refused_on_1, gpus_on_1, match): gpus = size if rank == 0 else gpus_on_1 # what this rank's host shows it with ( bounded_collectives(), - mock.patch.object(torch.cuda, "is_available", return_value=gpus > 0), - mock.patch.object(torch.cuda, "device_count", return_value=gpus), + pytest.MonkeyPatch.context() as cuda, mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", _loaded_in_step), ): + _pin_cuda(cuda, gpus) # the pooled worker's own CUDA state is whatever an earlier job left try: # not pytest.raises: its failure would end the pooled worker instead of the job load_pretrained( ckpt, @@ -1356,15 +1360,13 @@ def load(ckpt_path, device, *args, **kwargs): with ( bounded_collectives(), - # Each rank sees only its own GPU, which its process bound as cuda:0, whatever its local - # rank. - mock.patch.object(torch.cuda, "is_available", return_value=True), - mock.patch.object(torch.cuda, "device_count", return_value=1), - mock.patch.object(torch.cuda, "is_initialized", return_value=True), - mock.patch.object(torch.cuda, "current_device", return_value=0), + pytest.MonkeyPatch.context() as cuda, mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", load), ): + # Each rank sees only its own GPU, which its process bound as cuda:0, whatever its local + # rank. + _pin_cuda(cuda, 1, bound=0) load_pretrained(ckpt, distribution=DistributionPlan(2, dense="fsdp")) assert devices == [torch.device("cuda", 0)], (rank, devices) From 83a491666e885010691b66530ecc91ced678638d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:18:13 +0000 Subject: [PATCH 92/99] Parse the EP knob where a rank that cannot parse it fails every rank A load the plan decision decides read MODELOPT_EP_DEGREE on every rank just before resolve_distribution_plan, outside any collective refusal: a rank whose knob did not parse raised alone while rank 0 went on into the decision's collectives and waited there. The knob is now parsed in the load-arguments check, whose refusal raises on every rank. A two-rank test with the knob malformed on rank 1 alone fails on the group's bound without this. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 11 +++---- .../unit/torch/utils/test_model_load_utils.py | 29 +++++++++++++++---- 2 files changed, 29 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c4ff36d7dae..184de085d93 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -787,9 +787,14 @@ def handed() -> DistributionPlan | None: 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 {} + decided = plan is None and dist.size() > 1 # Each rank's arguments are its own, so each rank checks its own, and a refusal on one fails # every rank before the checkpoint is fetched. with _refused_together("load arguments", alone): + # No load here shards the routed experts yet: an EP degree left to the decision is 1, and + # one MODELOPT_EP_DEGREE asks for is refused. Parsed here, so a knob one rank cannot parse + # fails every rank rather than leave the others in the decision's collectives. + ep_size = 1 if decided and _ep_degree(None) == "auto" else None if kwargs.pop("output_loading_info", False): raise ValueError( "load_pretrained returns the model alone and does not take " @@ -839,17 +844,13 @@ 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 dist.size() > 1 if decided: plan, hf_config = resolve_distribution_plan( ckpt_path, config=config, auto_class=auto_class, trust_remote_code=trust_remote_code, - # No load here shards the routed experts yet: an EP degree left to the decision is 1, - # rather than a layout this function would refuse; one MODELOPT_EP_DEGREE asks for, as - # a plan handed in, is still refused. - ep_size=1 if _ep_degree(None) == "auto" else None, + ep_size=ep_size, **hub_options, ) if _route(plan) is not None: diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index ab5253492ab..724b63a79c9 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1291,22 +1291,24 @@ def _loaded_in_step(*args, **kwargs): return _Loaded() -def _one_rank_refused_job(rank, size, ckpt, refused_on_1, gpus_on_1, match): +_FSDP2_PLAN = DistributionPlan(2, dense="fsdp") + + +def _one_rank_refused_job( + rank, size, ckpt, refused_on_1, gpus_on_1, match, knobs_on_1=None, plan=_FSDP2_PLAN +): refused = None gpus = size if rank == 0 else gpus_on_1 # what this rank's host shows it with ( bounded_collectives(), pytest.MonkeyPatch.context() as cuda, + mock.patch.dict(os.environ, (knobs_on_1 or {}) if rank == 1 else {}), mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", _loaded_in_step), ): _pin_cuda(cuda, gpus) # the pooled worker's own CUDA state is whatever an earlier job left try: # not pytest.raises: its failure would end the pooled worker instead of the job - load_pretrained( - ckpt, - distribution=DistributionPlan(2, dense="fsdp"), - **(refused_on_1 if rank == 1 else {}), - ) + load_pretrained(ckpt, distribution=plan, **(refused_on_1 if rank == 1 else {})) except Exception as error: refused = str(error) # Rank 1's refusal on both ranks, rather than rank 0 failing at the group's timeout. @@ -1350,6 +1352,21 @@ def test_a_load_refused_on_one_rank_fails_every_rank( ) +def test_a_knob_one_rank_cannot_parse_fails_every_rank(cpu_dist_workers, tmp_path): + """With no plan handed, rank 1 alone reads a malformed MODELOPT_EP_DEGREE: every rank raises it, + rather than leave rank 0 waiting in the plan decision's collectives.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run( + _one_rank_refused_job, + ckpt=str(tmp_path), + refused_on_1={}, + gpus_on_1=2, + match="MODELOPT_EP_DEGREE='eigth'", + knobs_on_1={"MODELOPT_EP_DEGREE": "eigth"}, + plan=None, + ) + + def _one_visible_gpu_per_process_job(rank, size, ckpt): devices = [] From 1c0c7b3065e5fb5e5f7d19adf32c642ef2ede21c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:18:24 +0000 Subject: [PATCH 93/99] Shorten load_pretrained's docstring and the routing's comments CONTRIBUTING asks for brief docstrings. load_pretrained's keeps the routes, the recording contract and the refusal guarantees; the extended routing guidance moves to the PR description. Comments keep their reasons in fewer lines, the FSDP2-only options are named once, and the refusal of an unknown option drops a branch for a sharded load other than FSDP2, which has none. No behaviour changes. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/fsdp2_utils.py | 11 +- .../torch/utils/plugins/model_load_utils.py | 197 +++++------------- 2 files changed, 59 insertions(+), 149 deletions(-) diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index 1289a9e0fea..419abb4ae56 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/modelopt/torch/utils/fsdp2_utils.py @@ -101,15 +101,12 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False replicated and are moved onto the shards' device, which ``fully_shard`` does not do for the params it ignores. - It records no distribution plan: a caller wrapping a model itself records one once it is wrapped - (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). It refuses a model that, or - any of whose submodules, already records a plan, as a model ``load_pretrained`` returned does: a - plan is recorded once, so wrap a model loaded with ``from_pretrained`` instead. + It records no distribution plan, and refuses a model that, or any of whose submodules, records + one, as a model ``load_pretrained`` returned does: wrap one loaded with ``from_pretrained``. """ if (found := _recorded_plan(model)) is not None: - # The record would no longer hold: a model recorded as replicated would read as replicated - # while sharded, which the export would believe, and one recorded as sharded is wrapped - # already. A ValueError, as set_distribution_plan raises for the same record. + # A model recorded as replicated would read so while sharded, misleading the export, and one + # recorded as sharded is wrapped already. A ValueError, as set_distribution_plan raises. where, recorded = found raise ValueError( f"{where} is already recorded as distributed by {recorded!r}, and a plan is recorded " diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 184de085d93..c0744a57a64 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -435,21 +435,12 @@ def record_unplaced_source_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. +_FSDP2_ONLY = frozenset({"cpu_offload", "mp_policy", "broadcast_chunk_size"}) _SHARDED_LOAD_OPTIONS = { - _FSDP2: frozenset( - { - "trust_remote_code", - "attn_implementation", - "hf_config", - "cpu_offload", - "mp_policy", - "broadcast_chunk_size", - } - ), + _FSDP2: _FSDP2_ONLY | {"trust_remote_code", "attn_implementation", "hf_config"} } -# The values an FSDP2-only option asks nothing of a whole load with: the FSDP2 load's own default -# (parallel_load_and_prepare_fsdp2's), and None where that leaves it unset. A None -# broadcast_chunk_size asks for something: all of a source's layers at once. +# The FSDP2-only values that ask nothing of a whole load: the FSDP2 load's defaults, and None where +# that leaves the option unset (a None broadcast_chunk_size asks for all of a source's layers). _ASKS_NOTHING = {"cpu_offload": (False, None), "mp_policy": (None,), "broadcast_chunk_size": (8,)} @@ -503,10 +494,7 @@ def _whole_instead(decided: bool) -> str: def _refuse_sharded_load_options( route, auto_class, kwargs, decided: bool, measured: bool = False ) -> None: - """Refuse the arguments the sharded load ``route`` cannot honor: the arguments alone decide. - - ``measured`` says rank 0's GPU memory chose the sharded layout, which nothing asked for. - """ + """Refuse what the sharded load ``route`` cannot honor; ``measured``: GPU memory chose it.""" # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. passed = {name for name, value in kwargs.items() if value is not None} unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) @@ -535,15 +523,10 @@ def _refuse_sharded_load_options( "computes in." ) if unknown: - whole = ( - " To place the model yourself, have every rank load it whole: " - f"{_whole_instead(decided)}." - 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])}.{whole}" + f"takes {sorted(_SHARDED_LOAD_OPTIONS[route])}. To place the model yourself, have " + f"every rank load it whole: {_whole_instead(decided)}." ) @@ -577,9 +560,8 @@ def _sharded_load_kwargs( f"The {route} load cannot dequantize this {model_type} checkpoint. To load it whole on " f"every rank: {_whole_instead(decided)}." ) - # It builds as build_meta_causal_lm does: a text-only VLM it supports through its image-text - # class, anything else through AutoModelForCausalLM, which refuses most VLMs' configs and, for - # some, builds the language model alone, whose other towers the export would carry unquantized. + # As build_meta_causal_lm builds: AutoModelForCausalLM refuses most VLMs' configs, and builds + # some as the language model alone, whose other towers the export would carry unquantized. builds = ( AutoModelForImageTextToText if is_fsdp2_text_only_vlm(hf_config) else AutoModelForCausalLM ) @@ -609,9 +591,7 @@ def _sharded_load_kwargs( def _own_gpu_index() -> int | None: """This rank's GPU, as torch's DeviceMesh takes it, or ``None`` for a rank with none there. - Once the process has set CUDA up (setup(), a device-bound group, the caller's own set_device), - the device it set; before, the one its local rank numbers. A launch that exposes one GPU to - each process, each bound to cuda:0, so has a GPU on every rank, whatever its local rank. + The device CUDA is set to once it is set up, else the one the rank's local rank numbers. """ if not torch.cuda.is_available(): return None @@ -629,8 +609,7 @@ def _has_own_gpu() -> bool: def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: """The FSDP2 load, by the arguments :func:`_sharded_load_kwargs` checked.""" if not dist.is_initialized(): - # A single process's group of one, which the sharded model needs for as long as it lives: - # a run (distributed_run) destroys it on exit, and outside one it stays up. + # A single process's group of one, which the sharded model needs while it lives. dist.setup() return parallel_load_and_prepare_fsdp2( ckpt_path, @@ -651,72 +630,35 @@ def load_pretrained( ) -> nn.Module: """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. 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, or a text-only VLM the FSDP2 load supports (Qwen3.5), - 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, or a group of one, reads no - knob and takes this load only from a plan passed as ``distribution``; it sets a group of one up - for it, which needs ``torchrun --nproc_per_node=1``. It reads the checkpoint's tensors as - stored, so it refuses a quantized checkpoint. + ``distribution`` is the :class:`~modelopt.torch.utils.distributed.DistributionPlan` to load by: + every rank takes rank 0's or, when rank 0 passes none, :func:`resolve_distribution_plan`'s. + + * ``dense="fsdp"``: :func:`parallel_load_and_prepare_fsdp2` shards a causal LM, or a text-only + VLM it supports, over the process group (one process needs ``torchrun --nproc_per_node=1``). + It builds and places the model itself, so it takes its own options but no ``auto_class``, + placement or ``dtype``, and 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, onto the rank's own GPU (the one it is bound to, else the one its local rank - numbers) when it has one and no ``device_map`` is passed; ``device_map=None``, - ``from_pretrained``'s own, loads it on the CPU for the caller to place. - - A plan with ``ep_size > 1`` raises: no load here shards the routed experts, so it decides an EP - degree left to it as 1, and refuses one ``MODELOPT_EP_DEGREE`` asks for. On a distributed run - every rank loads the plan rank 0 was handed, and a refusal on any rank, of its own arguments, of a - checkpoint it cannot see or of the config it reads, raises on every rank. 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 and no ``device_map`` default; 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. It is recorded once, so a model to shard yourself - (:func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap`) is one loaded with ``from_pretrained``, - not here. - - A sharded load builds the model class and places it itself: a text-only VLM through its - image-text class, anything else through ``AutoModelForCausalLM``, the one ``auto_class`` it - takes, so leave ``auto_class`` out for a VLM. Besides its own options it takes - ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other keyword - rather than ignore it, but for one passed as ``None``, ``from_pretrained``'s own default, which - asks nothing. 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. An option at ``from_pretrained``'s own - default, ``device_map=None`` included, places nothing and does not keep the load whole. The - model is weighed in the checkpoint's own dtype (its config's, else bfloat16), not one the load - asks for: hand a load that widens it, ``dtype=torch.float32`` say, its plan. ``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, or a path to one as ``from_pretrained`` takes, is the one the - plan is decided from, and a sharded load takes it as its ``hf_config``. A whole-model load - refuses the options only the FSDP2 load takes, but at the FSDP2 load's own defaults - (``cpu_offload=False``, ``mp_policy=None``, ``broadcast_chunk_size=8``) or ``cpu_offload=None``, - which ask nothing of it, so a caller can forward them whatever the plan; - ``broadcast_chunk_size=None`` asks for all of a source's layers at once. - - Whatever the placement: + ``AutoModelForCausalLM``); on a distributed run every rank loads the whole model, onto its own + GPU unless ``device_map`` is passed. + + No load here shards the routed experts: a plan with ``ep_size > 1`` is refused. A refusal on any + rank raises on every rank, before the load. The plan is recorded on the model + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) for calibration and export, so + :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap` refuses it: to shard a model yourself, + load it with ``from_pretrained``. Around the load: * 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 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. + 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. * 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. A sharded load cannot run it, so - it refuses a checkpoint the hook would build modules for. + ``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 @@ -725,8 +667,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 - # A launcher can start several processes that set no process group up (hf_ptq without - # --use_fsdp2, say): each loads as a single process does, the whole model with no plan. + # A launch whose processes set no group up (hf_ptq without --use_fsdp2): each loads alone. alone = dist.is_distributed_launch() and not dist.is_initialized() if alone and isinstance(distribution, DistributionPlan) and distribution.world_size > 1: raise RuntimeError( @@ -735,8 +676,7 @@ def load_pretrained( "which sets it up." ) if alone and isinstance(distribution, DistributionPlan) and _route(distribution) is not None: - # One rank of a launch of several cannot set the group up for itself, and a one-rank plan - # for a sharded load is no plan for the launch. + # One rank cannot set the launch's group up, and a one-rank plan is no plan for the launch. launch = os.environ["WORLD_SIZE"] raise RuntimeError( f"The {_route(distribution)} load shards through this {launch}-rank launch's process " @@ -744,27 +684,20 @@ def load_pretrained( f"plan for the launch: DistributionPlan({launch}, dense='fsdp')." ) if alone and distribution is None: - # The knobs ask to shard over the launch, rather than leave each process the whole model. A - # value that is not an EP degree is refused, by the plan decision's own parse. + # Knobs asking to shard over the launch are refused, as is a value that is no EP degree. ep, world = _ep_degree(None), int(os.environ["WORLD_SIZE"]) experts = ep != "auto" and ep > 1 asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if experts else [] if _dense_from_env(world) == "fsdp": asked.append(f"MODELOPT_FSDP_SIZE={env.fsdp_size()}") if asked: - # The group alone would not do for the EP knob: no load here shards the experts. + ep_way = ", with MODELOPT_EP_DEGREE=1: no load here shards the routed experts" 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" - + ( - ", with MODELOPT_EP_DEGREE=1: no load here shards the routed experts." - if experts - else "." - ) + f"inside distributed_run(), which sets it up{ep_way if experts else ''}." ) - # None is these options' own default, so passed explicitly it is no option: the sharded loads - # refuse the options they do not take, and would refuse it. + # None is these options' own default: passed explicitly, it is no option to refuse. for name in ("config", "variant", "hf_config"): if name in kwargs and kwargs[name] is None: del kwargs[name] @@ -780,8 +713,7 @@ def handed() -> DistributionPlan | None: _route(distribution) return distribution - # Every rank takes rank 0's distribution, so every rank runs the same load and collectives; - # rank 0 checks it before it sends, so one it refuses raises on every rank, before the fetch. + # Every rank loads rank 0's distribution, which rank 0 checks before it sends it. plan = handed() if alone else dist.broadcast_from("distribution", handed) hub_options = { k: kwargs.pop(k) for k in ("token", "cache_dir", "local_files_only") if k in kwargs @@ -810,16 +742,14 @@ def handed() -> DistributionPlan | 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 "hf_config" in kwargs and (plan is None or _route(plan) is None): - # Only a sharded load builds from hf_config, as given, so the plan must be the one it - # came with, not one decided here from another config nor one that loads it whole. + # Only a sharded load builds from hf_config, so it needs the plan it came with. raise TypeError( "hf_config is the config resolve_distribution_plan returned with a plan that " "shards the model: pass that plan as distribution=, or pass config= to have the " "plan decided from your config." ) if plan is not None and _route(plan) is not None: - # A handed plan's arguments alone decide what its load refuses, before the fetch; a - # config passed for it becomes its hf_config, and subfolder is where the checkpoint is. + # A config passed becomes the load's hf_config; subfolder is where the checkpoint is. _refuse_sharded_load_options( _route(plan), auto_class, @@ -830,13 +760,11 @@ def handed() -> DistributionPlan | None: }, decided=False, ) - # Rank 0 fetches the checkpoint, and each rank then checks it can see it: one that cannot fails - # every rank, rather than leave the others in the load's next collective. + # Rank 0 fetches the checkpoint; a rank that cannot see it then fails every rank. with _refused_together("checkpoint fetch", alone): _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) if config is not None and not isinstance(config, PretrainedConfig): - # Read as the load reads it, so the plan and the loading hook follow the config the - # model is built from, which need not be in the checkpoint's directory. + # Read as the load reads it, so the plan and the loading hook follow the model's config. trust = kwargs.get("trust_remote_code", False) config = _read_config(config, auto_class, trust, **hub_options) subfolder = kwargs.pop("subfolder", None) @@ -859,17 +787,14 @@ def handed() -> DistributionPlan | None: # One rank holds the whole model, whatever the model: there is nothing to decide or read. plan = DistributionPlan() route = _route(plan) - # Each rank reads the checkpoint's config itself, so a refusal of what it reads, or of its own - # arguments, fails every rank before the load's collectives. + # Each rank reads the config itself, so a refusal of it fails every rank before the load. with _refused_together("load checks", alone): # Options only the FSDP2 load takes would reach from_pretrained, which has no use for them. - shared = {"trust_remote_code", "attn_implementation", "hf_config"} # hf_config: see above - own = _SHARDED_LOAD_OPTIONS[_FSDP2] - shared if route is None: for name, nothing in _ASKS_NOTHING.items(): if name in kwargs and kwargs[name] in nothing: del kwargs[name] - if route is None and (sharded := sorted(set(kwargs) & own)): + if route is None and (sharded := sorted(set(kwargs) & _FSDP2_ONLY)): raise TypeError( f"{sharded} only apply to the FSDP2 load, but {plan!r} gives every rank the whole " "model; pass a plan with dense='fsdp' to shard it." @@ -905,9 +830,8 @@ def handed() -> DistributionPlan | None: keys = model._modelopt_unplaced_source_keys else: if plan.world_size > 1 and (index := _own_gpu_index()) is not None: - # Each rank takes the whole model onto its own GPU. Left to "auto", every rank's - # accelerate could pick the same idle GPU; a rank with none loads as from_pretrained - # places it. + # Each rank takes the whole model onto its own GPU: under "auto", every rank's + # accelerate could pick the same idle GPU. kwargs.setdefault("device_map", {"": f"cuda:{index}"}) elif plan.world_size > 1 and "device_map" not in kwargs: # This rank says so itself: rank 0 may have a GPU of its own, and nothing to say. @@ -1267,10 +1191,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). - It records no distribution plan, as :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap` does - not: :func:`load_pretrained` records one, and a caller of this loader records its own - (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) before calibrating or - exporting the model. + It records no distribution plan: a caller records its own, as :func:`load_pretrained` does + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) weight_map = indexed_weight_map(resolved_path) @@ -1481,13 +1403,10 @@ def _decide_distribution( 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. + # The sharded loads refuse this class: with nothing asking to shard, load it whole, unbuilt. 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. + # Refused here, rather than build the decision's own pick of class, which may not build. name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( f"A sharded load builds its model class itself, a causal LM's or a text-only VLM's " @@ -1583,9 +1502,7 @@ def resolve_distribution_plan( every later step reads it back from there (:func:`~modelopt.torch.utils.distributed.distribution_plan`). Returns ``(plan, hf_config)``, ``hf_config`` being this rank's own copy of the config (rank 0 decides the plan from its copy), - 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``. + for a sharded load to build from: :func:`load_pretrained` takes it only with a plan that shards. 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 @@ -1613,13 +1530,9 @@ def resolve_distribution_plan( :func:`ensure_local_checkpoint`, and passes the local directory: the plan is then decided for the snapshot the load reads. - ``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 building writes to it. + ``config``, the caller's own config, is decided from instead of the checkpoint's. ``auto_class`` + is the class the load builds with: with one the sharded loads refuse, and nothing asking to + shard, the plan loads the model whole, decided without a build. """ # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather # than leave the others waiting in a broadcast below for a rank that already raised. From 6d628194c07b23686559ec297272aaded48c4cc0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:47:59 +0000 Subject: [PATCH 94/99] Word the automatic layout's refusal for any option a sharded load refuses The refusal under a layout rank 0's GPU memory chose said the refused options place a whole model, but they are every keyword the sharded load does not take (use_safetensors, low_cpu_mem_usage, ...). It now says they do not apply to a sharded load, as the refusal beside it does, and names the way to load the model whole. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 9 ++++----- tests/unit/torch/utils/test_model_load_utils.py | 8 ++++---- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c0744a57a64..d840ff659d6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -500,13 +500,12 @@ def _refuse_sharded_load_options( unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. dtypes = [name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")] - if measured and (placed := sorted({*unknown, *dtypes})): + if measured and (refused := sorted({*unknown, *dtypes})): # Say so, rather than refuse the options below as if the caller had asked for this 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 the {route} load " - f"shards it. To load it whole on every rank and place it yourself: " - f"{_whole_instead(decided)}." + f"{refused} do not apply to a sharded load, but no plan or MODELOPT_FSDP_SIZE asked for " + f"a layout, and by rank 0's GPU memory the model does not fit one GPU, so the {route} " + f"load shards it. To load it whole on every rank: {_whole_instead(decided)}." ) if auto_class not in (None, AutoModelForCausalLM): name = getattr(auto_class, "__name__", repr(auto_class)) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 724b63a79c9..6223185166a 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1568,9 +1568,9 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( run_as, serve, loaders, process_group, monkeypatch, tmp_path, device_map ): """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. device_map=None, as - hf_ptq passes it, places nothing, and the model is sharded.""" + big for it, and the sharded load refuses options it does not take. The refusal says the layout + was decided rather than asked for, and how to keep the load whole. device_map=None, as hf_ptq + passes it, places nothing, and the model is sharded.""" run_as(2, rank=1) serve(LlamaConfig(**_TINY)) monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: 1) @@ -1579,7 +1579,7 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( load_pretrained(str(tmp_path), AutoModelForCausalLM, device_map=None, dtype="auto") assert "fsdp2" in loaders return - with pytest.raises(TypeError, match=r"\['device_map'\] place.*MODELOPT_FSDP_SIZE=1"): + with pytest.raises(TypeError, match=r"\['device_map'\] do not apply.*MODELOPT_FSDP_SIZE=1"): load_pretrained(str(tmp_path), device_map="auto") assert loaders == {} From b0fc305f15f63126ffd7b7f7a5be763fe9f778a6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:48:10 +0000 Subject: [PATCH 95/99] Say a load placed elsewhere passes device_map=None On a distributed run the whole-model load defaults each rank's device_map to its own GPU. A loader that places the model itself, as init_quantized_weights does, refuses any device_map, so its caller passes device_map=None, which the default leaves alone. The docstring says so. 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 d840ff659d6..0d48b954def 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -638,7 +638,8 @@ def load_pretrained( placement or ``dtype``, and refuses a quantized checkpoint. * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to ``AutoModelForCausalLM``); on a distributed run every rank loads the whole model, onto its own - GPU unless ``device_map`` is passed. + GPU unless ``device_map`` is passed. Pass ``device_map=None`` where the model is placed + otherwise, as under ``init_quantized_weights``, which refuses a device map. No load here shards the routed experts: a plan with ``ep_size > 1`` is refused. A refusal on any rank raises on every rank, before the load. The plan is recorded on the model From 3d4d1b40c46196e534126d35e2c843534d4fe20b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 03:51:10 +0000 Subject: [PATCH 96/99] Refuse a sharded load inside init_quantized_weights init_quantized_weights quantizes and compresses the model in the AutoModelForCausalLM.from_pretrained it patches. Under a live multi-rank group a load the automatic layout shards, or a handed plan's, built the model through from_config instead and returned it unquantized. The context now marks the from_pretrained it installs, and a sharded load refuses while it is installed, on every rank, before anything loads, naming the way to load the model whole. A test forces the automatic layout to FSDP2 inside the real context, and hands a sharded plan, and fails without this. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- .../torch/quantization/plugins/accelerate.py | 2 ++ .../torch/utils/plugins/model_load_utils.py | 7 ++++++ .../unit/torch/utils/test_model_load_utils.py | 22 +++++++++++++++++++ 4 files changed, 32 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e81fe4aa66f..8b335cec12c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,7 +52,7 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. -- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints and ones carrying an MTP head; an MoE loads without expert parallelism, and is refused when ``MODELOPT_EP_DEGREE`` asks for it, until the expert-parallel load lands; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE loads without expert parallelism, and is refused when ``MODELOPT_EP_DEGREE`` asks for it, until the expert-parallel load lands; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. - ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm. diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py index f1601de058d..36b08318a23 100644 --- a/modelopt/torch/quantization/plugins/accelerate.py +++ b/modelopt/torch/quantization/plugins/accelerate.py @@ -358,6 +358,8 @@ def patched_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwar **dispatch_kwargs, ) + # A sharded load_pretrained builds the model without from_pretrained, so it reads this mark. + patched_from_pretrained._modelopt_init_quantized_weights = True # type: ignore[attr-defined] try: # Apply our patches transformers.AutoModelForCausalLM.from_pretrained = classmethod(patched_from_pretrained) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 0d48b954def..659a7b2257d 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -507,6 +507,13 @@ def _refuse_sharded_load_options( f"a layout, and by rank 0's GPU memory the model does not fit one GPU, so the {route} " f"load shards it. To load it whole on every rank: {_whole_instead(decided)}." ) + if getattr(AutoModelForCausalLM.from_pretrained, "_modelopt_init_quantized_weights", False): + # Its from_pretrained quantizes and compresses the model; a sharded load never calls it. + raise NotImplementedError( + f"The {route} load builds the model without from_pretrained, so it would skip " + "init_quantized_weights' quantization; load the model whole on every rank: " + f"{_whole_instead(decided)}." + ) if auto_class not in (None, AutoModelForCausalLM): name = getattr(auto_class, "__name__", repr(auto_class)) raise NotImplementedError( diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 6223185166a..08393c1494a 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -60,7 +60,9 @@ ) import modelopt.torch.models.hf as models_hf +import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist +from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights 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 @@ -1584,6 +1586,26 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( assert loaders == {} +@pytest.mark.parametrize("handed", [False, True], ids=["automatic", "handed"]) +def test_a_sharded_load_is_refused_inside_init_quantized_weights( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, handed +): + """init_quantized_weights quantizes and compresses in the from_pretrained it patches, which the + FSDP2 load never calls: a layout that shards there, rank 0's GPU memory's or a handed plan's, + is refused rather than return the model unquantized.""" + run_as(2, rank=1) + serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: 1) + plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} + + with ( + init_quantized_weights(mtq.INT8_DEFAULT_CFG), + pytest.raises(NotImplementedError, match=r"init_quantized_weights.*dense='replicate'"), + ): + load_pretrained(str(tmp_path), device_map=None, **plan) + assert loaders == {} + + @pytest.mark.parametrize( ("checkpoint", "match"), [("quantized", "dequantize"), ("mtp", "MTP head")], From 11bd1b027adb4a38c5e6f6ce66aacf56f4f258b0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 03:51:59 +0000 Subject: [PATCH 97/99] Publish load_pretrained's routing in the docs, and say the FSDP2 load skips Hub options The routing guidance that moved out of load_pretrained's docstring is now a docs page, "Loading Hugging Face models across ranks", under the quantization guide, and the docstring links it. The docstring said the config and model loads both take the Hub options; the FSDP2 load builds the model with from_config, which takes none, so the docstring and the page now say so. set_distribution_plan's docstring is back to main's wording. Signed-off-by: Shengliang Xu --- docs/source/guides/1_quantization.rst | 1 + docs/source/guides/_distributed_loading.rst | 118 ++++++++++++++++++ modelopt/torch/utils/distributed.py | 8 +- .../torch/utils/plugins/model_load_utils.py | 11 +- 4 files changed, 129 insertions(+), 9 deletions(-) create mode 100644 docs/source/guides/_distributed_loading.rst diff --git a/docs/source/guides/1_quantization.rst b/docs/source/guides/1_quantization.rst index 38ce0956b79..1db915c6975 100644 --- a/docs/source/guides/1_quantization.rst +++ b/docs/source/guides/1_quantization.rst @@ -19,6 +19,7 @@ Below, you can find the documentation for the quantization toolkit in ModelOpt: ./_basic_quantization.rst ./_choosing_quant_methods.rst ./_pytorch_quantization.rst + ./_distributed_loading.rst ./_quant_cfg.rst ./_customized_model_quantization.rst ./_compress_quantized_models.rst diff --git a/docs/source/guides/_distributed_loading.rst b/docs/source/guides/_distributed_loading.rst new file mode 100644 index 00000000000..69c30a6d671 --- /dev/null +++ b/docs/source/guides/_distributed_loading.rst @@ -0,0 +1,118 @@ +======================================== +Loading Hugging Face models across ranks +======================================== + +:func:`load_pretrained ` loads a +Hugging Face model for ModelOpt to optimize and export, whole on every rank or distributed. It loads +by a :class:`DistributionPlan `. Every rank takes +rank 0's plan: the one rank 0 passes as ``distribution``, or, when it passes none, the one +:func:`resolve_distribution_plan ` +decides from ``MODELOPT_FSDP_SIZE``, ``MODELOPT_EP_DEGREE`` and rank 0's GPU memory. The plan is +recorded on the model, where calibration and export read it. + +.. code-block:: python + + # torchrun --nproc_per_node 8 quantize.py + from modelopt.torch.utils.distributed import distributed_run + from modelopt.torch.utils.plugins.model_load_utils import load_pretrained + + with distributed_run(): + model = load_pretrained("meta-llama/Llama-3.1-8B") # a whole copy on every rank's GPU + # With MODELOPT_FSDP_SIZE=world: FSDP2-sharded over the 8 ranks instead. + +Routes +====== + +* A plan with ``dense="fsdp"`` takes the FSDP2 load (``parallel_load_and_prepare_fsdp2``): a causal + LM, or a text-only VLM it supports (Qwen3.5), is FSDP2-sharded over the process group, each rank + reading its own layers. +* Any other plan gives every rank the whole model, through ``auto_class.from_pretrained`` + (``auto_class`` defaults to ``AutoModelForCausalLM``). + +No load shards the routed experts yet. A plan the load decides keeps the expert-parallel degree at +1, so an MoE loads without expert parallelism; an EP degree ``MODELOPT_EP_DEGREE`` asks for, or a +plan with ``ep_size > 1``, is refused. + +A refusal on any rank raises on every rank, before the load, and says how to load the model whole: +``MODELOPT_FSDP_SIZE=1`` for a plan the load decided, or +``distribution=DistributionPlan(, dense="replicate")``. + +The plan is recorded once, so ``fsdp2_wrap`` refuses a model ``load_pretrained`` returned: to shard +a model yourself, load it with ``from_pretrained``. + +Whole-model loads +================= + +* On a distributed run, each rank loads the model onto its own GPU when it has one and no + ``device_map`` is passed: the GPU its process is bound to once CUDA is set up, else the one its + local rank numbers. ``device_map=None``, ``from_pretrained``'s own default, loads it on the CPU + for the caller to place. A rank without a GPU of its own warns and loads where + ``from_pretrained`` places it. +* The options only the FSDP2 load takes are refused, except at the FSDP2 load's own defaults + (``cpu_offload=False``, ``mp_policy=None``, ``broadcast_chunk_size=8``) or + ``cpu_offload=None``, which ask nothing of a whole-model load, so a caller can forward them + whatever the plan. ``broadcast_chunk_size=None`` asks for all of a source's layers at once. + +Sharded loads +============= + +* The FSDP2 load builds the model class and places the model itself: a text-only VLM through its + image-text class, anything else through ``AutoModelForCausalLM``. Leave ``auto_class`` out. +* Besides its own options (``cpu_offload``, ``mp_policy``, ``broadcast_chunk_size``) it takes + ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other keyword + rather than ignore it. An option passed as ``None``, ``from_pretrained``'s own default, asks + nothing and is not refused. +* It reads the checkpoint's tensors as stored and keeps their dtype, so it refuses a ``dtype`` and + a quantized checkpoint. It cannot run the model type's loading hook, so it refuses a checkpoint + the hook would build modules for, such as a Nemotron-H MTP head. +* A single process, or a group of one, reads no knob: it takes the FSDP2 load only from a plan + passed as ``distribution``, and sets a group of one up for it, which needs + ``torchrun --nproc_per_node=1``. + +The automatic layout +==================== + +With neither a plan nor ``MODELOPT_FSDP_SIZE``, rank 0's GPU memory decides whether the dense part +is sharded. The model is weighed in the checkpoint's own dtype (its config's, else bfloat16), not +one the load asks for: hand a load that widens it, ``dtype=torch.float32`` say, its plan. + +A call that places the model (``device_map``, ``quantization_config``, a ``dtype``) loads only +where the model fits one GPU; where it does not, the refusal says the layout was decided rather +than asked for. Pass ``distribution=`` or set ``MODELOPT_FSDP_SIZE=1`` to keep the load whole. An +option at ``from_pretrained``'s own default, ``device_map=None`` included, places nothing and does +not keep the load whole. + +Configs +======= + +``hf_config`` is the config ``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, or a path to one, is the one the plan is decided from, and a sharded +load takes it as its ``hf_config``. + +Hub options +=========== + +A Hub model ID is made local first, so the export can read the source checkpoint from disk; the +``revision``, ``token``, ``cache_dir`` and ``local_files_only`` options go to that fetch. The config +read and a whole-model load also take all but ``revision``: an ``auto_map`` entry can name custom +code in another Hub repo, which Transformers fetches with them. The FSDP2 load builds the model +from the config with ``from_config``, which takes none of them, so such code comes with the +environment's Hub settings. + +Launches without a process group +================================ + +A launch whose processes set no process group up (a launcher that only exports ``WORLD_SIZE``, as +``hf_ptq.py`` without ``--use_fsdp2``) has each process load as a single process does: the whole +model, with no plan to agree on and no ``device_map`` default. A plan for several ranks, a sharded +load, and a knob asking to shard need the group, so they are refused there: set it up first, with +``distributed_run()``. + +Inside ``init_quantized_weights`` +================================= + +``init_quantized_weights`` (see :doc:`_compress_quantized_models`) quantizes and compresses the +model in the ``from_pretrained`` it patches, which the FSDP2 load never calls, so a sharded load is +refused inside it. A whole-model load there passes ``device_map=None``, as ``init_quantized_weights`` +computes the device map itself and refuses one. diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 8a5dbdce815..3f1dd35abeb 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -365,10 +365,10 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non """Record that ``model`` is distributed as ``plan``, on it and on every submodule it has now. Code handed a submodule, such as a VLM's language model, then reads the same plan as code handed - the whole model; a module added afterwards records none. ``load_pretrained`` records the plan it - loaded by, once the model is built; code that distributes a model itself records it here before - calibrating or exporting the model. The plan says how the model is placed over the run; the - :class:`ParallelState` a quantized module holds names the process groups its collectives run on. + the whole model; a module added afterwards records none. Code that distributes a model records + its plan here before calibrating or exporting the model. The plan says how the model is placed + over the run; the :class:`ParallelState` a quantized module holds names the process groups its + collectives run on. A model is distributed once, and its plan recorded once: recording a plan on a model that, or any of whose submodules, already records one raises, the same plan included, and leaves the diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 659a7b2257d..c06e69d2e30 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -652,15 +652,16 @@ def load_pretrained( rank raises on every rank, before the load. The plan is recorded on the model (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) for calibration and export, so :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap` refuses it: to shard a model yourself, - load it with ``from_pretrained``. Around the load: + load it with ``from_pretrained``. :doc:`/guides/_distributed_loading` has the routing in full. + Around the load: * 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 a 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 the model with ``from_config``, which takes none. 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 From c78d2e50ed6d7bd2cb65655ef1209b2b17122da2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:09:42 +0000 Subject: [PATCH 98/99] Load whole inside init_quantized_weights when nothing asks to shard Inside init_quantized_weights a multi-rank load_pretrained failed whatever it decided: a sharded layout was refused, and a replicated one injected each rank's device_map, which the context's from_pretrained refuses, naming an argument the caller never passed. As for a model class the sharded loads refuse, the plan decision now gives every rank the whole model there when nothing asks to shard, without measuring it, and the whole-model load leaves the device map to the context, which computes it. A knob or a handed plan that asks to shard is still refused. A test loads inside the real context, with rank 0's GPU memory fitting the model or not, and reaches the context's from_pretrained past its device_map check; it fails without this. Signed-off-by: Shengliang Xu --- docs/source/guides/_distributed_loading.rst | 17 ++++---- .../torch/utils/plugins/model_load_utils.py | 43 ++++++++++--------- .../unit/torch/utils/test_model_load_utils.py | 34 ++++++++++++--- 3 files changed, 60 insertions(+), 34 deletions(-) diff --git a/docs/source/guides/_distributed_loading.rst b/docs/source/guides/_distributed_loading.rst index 69c30a6d671..230b0249629 100644 --- a/docs/source/guides/_distributed_loading.rst +++ b/docs/source/guides/_distributed_loading.rst @@ -43,11 +43,11 @@ a model yourself, load it with ``from_pretrained``. Whole-model loads ================= -* On a distributed run, each rank loads the model onto its own GPU when it has one and no - ``device_map`` is passed: the GPU its process is bound to once CUDA is set up, else the one its - local rank numbers. ``device_map=None``, ``from_pretrained``'s own default, loads it on the CPU - for the caller to place. A rank without a GPU of its own warns and loads where - ``from_pretrained`` places it. +* On a distributed run, each rank loads the model onto its own GPU when it has one, no + ``device_map`` is passed and ``init_quantized_weights`` does not place it: the GPU its process is + bound to once CUDA is set up, else the one its local rank numbers. ``device_map=None``, + ``from_pretrained``'s own default, loads it on the CPU for the caller to place. A rank without a + GPU of its own warns and loads where ``from_pretrained`` places it. * The options only the FSDP2 load takes are refused, except at the FSDP2 load's own defaults (``cpu_offload=False``, ``mp_policy=None``, ``broadcast_chunk_size=8``) or ``cpu_offload=None``, which ask nothing of a whole-model load, so a caller can forward them @@ -113,6 +113,7 @@ Inside ``init_quantized_weights`` ================================= ``init_quantized_weights`` (see :doc:`_compress_quantized_models`) quantizes and compresses the -model in the ``from_pretrained`` it patches, which the FSDP2 load never calls, so a sharded load is -refused inside it. A whole-model load there passes ``device_map=None``, as ``init_quantized_weights`` -computes the device map itself and refuses one. +model in the ``from_pretrained`` it patches, which the FSDP2 load never calls. Inside it, a plan +the load decides gives every rank the whole model, with nothing to measure, and the load leaves +the device map to ``init_quantized_weights``, which computes it and refuses one it is passed. A +knob or a plan that asks to shard is refused there. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c06e69d2e30..094a2d79504 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -485,6 +485,11 @@ def _read_config(path, auto_class, trust_remote_code: bool = False, **hub_option return AutoConfig.from_pretrained(path, trust_remote_code=trust_remote_code, **hub_options) +def _init_quantized_weights() -> bool: + """Whether ``init_quantized_weights`` patched ``from_pretrained`` to quantize and place.""" + return getattr(AutoModelForCausalLM.from_pretrained, "_modelopt_init_quantized_weights", False) + + def _whole_instead(decided: bool) -> str: """How to have every rank load the model whole: the knob counts only for a decided plan.""" plan = f"distribution=DistributionPlan({dist.size()}, dense='replicate')" @@ -507,8 +512,7 @@ def _refuse_sharded_load_options( f"a layout, and by rank 0's GPU memory the model does not fit one GPU, so the {route} " f"load shards it. To load it whole on every rank: {_whole_instead(decided)}." ) - if getattr(AutoModelForCausalLM.from_pretrained, "_modelopt_init_quantized_weights", False): - # Its from_pretrained quantizes and compresses the model; a sharded load never calls it. + if _init_quantized_weights(): raise NotImplementedError( f"The {route} load builds the model without from_pretrained, so it would skip " "init_quantized_weights' quantization; load the model whole on every rank: " @@ -546,8 +550,7 @@ def _sharded_load_kwargs( # Past the refusal, anything else is a dtype the load keeps or an option at its default, None. kwargs = {k: v for k, v in kwargs.items() if k in _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. + # The head is built on a whole model only; carried over, it would export unquantized. raise NotImplementedError( f"The {route} load cannot build the MTP head this {model_type} checkpoint carries. To " f"load it whole on every rank: {_whole_instead(decided)}." @@ -560,8 +563,7 @@ def _sharded_load_kwargs( ) nested = (getattr(hf_config, name, None) for name in ("text_config", "llm_config")) if any(getattr(c, "quantization_config", None) is not None for c in (hf_config, *nested)): - # It reads the tensors as stored: quantized weights, a VLM's language model's included, - # would load as if they were not. + # Read as stored, quantized weights (a VLM's language model's too) would load as if not. raise NotImplementedError( f"The {route} load cannot dequantize this {model_type} checkpoint. To load it whole on " f"every rank: {_whole_instead(decided)}." @@ -584,8 +586,7 @@ def _sharded_load_kwargs( "started without a launcher cannot set up; launch it with `torchrun --nproc_per_node=1`." ) if not _has_own_gpu(): - # Refused here, with the other ranks, rather than inside the load, where they would wait - # for this rank in its broadcasts. + # Refused with the other ranks, not inside the load, where they would wait for this one. raise RuntimeError( f"The {route} load shards onto each rank's own GPU, and this rank has none: it sees " f"{torch.cuda.device_count()} GPU(s) as local rank {dist.local_rank()}. To load it " @@ -645,8 +646,7 @@ def load_pretrained( placement or ``dtype``, and refuses a quantized checkpoint. * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to ``AutoModelForCausalLM``); on a distributed run every rank loads the whole model, onto its own - GPU unless ``device_map`` is passed. Pass ``device_map=None`` where the model is placed - otherwise, as under ``init_quantized_weights``, which refuses a device map. + GPU unless ``device_map`` is passed or ``init_quantized_weights`` places it. No load here shards the routed experts: a plan with ``ep_size > 1`` is refused. A refusal on any rank raises on every rank, before the load. The plan is recorded on the model @@ -728,12 +728,10 @@ def handed() -> DistributionPlan | None: } revision = {"revision": kwargs.pop("revision")} if "revision" in kwargs else {} decided = plan is None and dist.size() > 1 - # Each rank's arguments are its own, so each rank checks its own, and a refusal on one fails - # every rank before the checkpoint is fetched. + # Each rank checks its own arguments; a refusal on one fails every rank before the fetch. with _refused_together("load arguments", alone): - # No load here shards the routed experts yet: an EP degree left to the decision is 1, and - # one MODELOPT_EP_DEGREE asks for is refused. Parsed here, so a knob one rank cannot parse - # fails every rank rather than leave the others in the decision's collectives. + # No load here shards the routed experts yet, so an EP degree left to the decision is 1. + # Parsed in this check, a knob one rank cannot parse fails every rank. ep_size = 1 if decided and _ep_degree(None) == "auto" else None if kwargs.pop("output_loading_info", False): raise ValueError( @@ -837,11 +835,13 @@ def handed() -> DistributionPlan | None: # A sharded load finds the unplaced keys while it places the weights. keys = model._modelopt_unplaced_source_keys else: - if plan.world_size > 1 and (index := _own_gpu_index()) is not None: + # init_quantized_weights computes the device map itself, and refuses one it is passed. + spread = plan.world_size > 1 and not _init_quantized_weights() + if spread and (index := _own_gpu_index()) is not None: # Each rank takes the whole model onto its own GPU: under "auto", every rank's # accelerate could pick the same idle GPU. kwargs.setdefault("device_map", {"": f"cuda:{index}"}) - elif plan.world_size > 1 and "device_map" not in kwargs: + elif spread and "device_map" not in kwargs: # This rank says so itself: rank 0 may have a GPU of its own, and nothing to say. warnings.warn( f"Rank {dist.rank()} has no GPU of its own: it sees {torch.cuda.device_count()} " @@ -1407,11 +1407,11 @@ 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) + (auto_class not in (None, AutoModelForCausalLM) or _init_quantized_weights()) and ep_size in ("auto", 1) and dense in ("auto", "replicate") ): - # The sharded loads refuse this class: with nothing asking to shard, load it whole, unbuilt. + # The sharded loads refuse this class and init_quantized_weights: load it whole, unbuilt. return DistributionPlan(world_size) if auto_class not in (None, AutoModelForCausalLM): # Refused here, rather than build the decision's own pick of class, which may not build. @@ -1539,8 +1539,9 @@ def resolve_distribution_plan( the snapshot the load reads. ``config``, the caller's own config, is decided from instead of the checkpoint's. ``auto_class`` - is the class the load builds with: with one the sharded loads refuse, and nothing asking to - shard, the plan loads the model whole, decided without a build. + is the class the load builds with: with one the sharded loads refuse, or inside + ``init_quantized_weights``, and nothing asking to shard, the plan loads the model whole, decided + without a build. """ # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather # than leave the others waiting in a broadcast below for a rank that already raised. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 08393c1494a..9266e601180 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -62,6 +62,7 @@ import modelopt.torch.models.hf as models_hf import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist +from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan @@ -1586,23 +1587,46 @@ def test_an_automatic_layout_says_so_when_it_refuses_a_placement( assert loaders == {} -@pytest.mark.parametrize("handed", [False, True], ids=["automatic", "handed"]) +class _ContextLoadReachedError(Exception): + """Raised where init_quantized_weights' from_pretrained builds the model, past its checks.""" + + +@pytest.mark.parametrize("memory", [1 << 50, 1], ids=["fits", "too-big"]) +def test_a_load_it_decides_inside_init_quantized_weights_is_whole_and_unplaced( + run_as, loaders, process_group, monkeypatch, tmp_path, memory +): + """Inside init_quantized_weights, with nothing asking to shard, every rank gets the whole model, + whatever rank 0's GPU memory, and the load leaves the device map to the context, which computes + it and refuses one passed: the context's from_pretrained gets past that refusal.""" + run_as(2, rank=1) + LlamaConfig(**_TINY).save_pretrained(tmp_path) # read from disk, as the context reads it too + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: memory) + monkeypatch.setattr( + accelerate_plugin, "init_empty_weights", mock.Mock(side_effect=_ContextLoadReachedError) + ) + + with init_quantized_weights(mtq.INT8_DEFAULT_CFG), pytest.raises(_ContextLoadReachedError): + load_pretrained(str(tmp_path)) + assert loaders == {} + + +@pytest.mark.parametrize("handed", [False, True], ids=["knob", "handed"]) def test_a_sharded_load_is_refused_inside_init_quantized_weights( run_as, serve, loaders, process_group, monkeypatch, tmp_path, handed ): """init_quantized_weights quantizes and compresses in the from_pretrained it patches, which the - FSDP2 load never calls: a layout that shards there, rank 0's GPU memory's or a handed plan's, - is refused rather than return the model unquantized.""" + FSDP2 load never calls: a knob or a plan asking to shard there is refused rather than return the + model unquantized.""" run_as(2, rank=1) serve(LlamaConfig(**_TINY)) - monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: 1) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} with ( init_quantized_weights(mtq.INT8_DEFAULT_CFG), pytest.raises(NotImplementedError, match=r"init_quantized_weights.*dense='replicate'"), ): - load_pretrained(str(tmp_path), device_map=None, **plan) + load_pretrained(str(tmp_path), **plan) assert loaders == {} From 0d3d0937c055ef18ed5ff19d324fa7e67a84820f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:34:20 +0000 Subject: [PATCH 99/99] Make load_pretrained's changelog entry match what it refuses The entry said only a knob refuses a model class other than a causal LM, but a layout rank 0's GPU memory chose refuses one too, and it ran to five behaviour changes in one sentence. It now lists what the sharded load refuses, says what to set to keep loading whole, and points to the new guide; the per-rank GPU placement of a whole load is a bullet of its own. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8b335cec12c..92f87a9d746 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,7 +52,8 @@ Changelog - Bump minimum required PyTorch version to 2.10. - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra. - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later. -- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE loads without expert parallelism, and is refused when ``MODELOPT_EP_DEGREE`` asks for it, until the expert-parallel load lands; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to load it whole on every rank. +- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, and the sharded load refuses placement arguments, a ``dtype``, a quantized checkpoint, one carrying an MTP head, a model class it does not build, and a call inside ``init_quantized_weights``; an MoE loads without expert parallelism, which ``MODELOPT_EP_DEGREE`` cannot ask for until the expert-parallel load lands. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(, dense="replicate")``, to keep loading it whole on every rank; the guide "Loading Hugging Face models across ranks" has the details. +- ``load_pretrained`` loading a model whole on several ranks now places it on each rank's own GPU unless ``device_map`` is passed. - ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm.