From dde68d6aca4c5c89fe871c221892c445b515f954 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:21:35 +0000 Subject: [PATCH 1/7] Share the eager-only and transposed experts rules in the HF plugin Quantization decides in two places which fused experts it cannot run on the grouped forward: experts stored transposed, which it leaves unquantized, and experts whose wrapper quantizes on the eager forward only, which keep the forced eager backend. Each rule is now one helper, _transposed_experts and _eager_only_experts, which register_fused_experts_on_the_fly and normalize_experts_impl_on_the_fly use, so the expert-parallel load can refuse exactly what quantization would get wrong. The eager-only rule is the registry lookup alone: the on-the-fly registration only adds the generic wrapper, which is not eager-only. No behaviour changes. Signed-off-by: Shengliang Xu --- .../torch/quantization/plugins/huggingface.py | 63 +++++++++++++------ 1 file changed, 43 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index e7b160840f8..60669054c0f 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1674,13 +1674,7 @@ def register_fused_experts_on_the_fly(model): visited_types.add(mod_type) wrapper = _fused_experts_wrapper(module) - if wrapper is not None and ( - getattr(module, "is_transposed", False) - # Stored `(num_experts, out_dim, in_dim)`, the first projection takes the dim down_proj - # puts out. A gated container's first projection is `2 * intermediate`, so its shapes - # give a transposed layout away whether it says so or not. - or getattr(module, wrapper._first_proj_attr).shape[-1] != module.down_proj.shape[1] - ): + if wrapper is not None and _transposed_experts([module]): # `_fused_experts_wrapper` matches on parameter SHAPE only, so it also claims a # transposed `(num_experts, in_dim, out_dim)` container. The fused-experts export # (export/moe_utils.py) then splits gate/up by slicing dim 1 at `intermediate_dim`, the @@ -1702,6 +1696,46 @@ def register_fused_experts_on_the_fly(model): QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(wrapper) +def _transposed_experts(modules) -> list[str]: + """Names of the expert classes in ``modules`` stored ``(num_experts, in_dim, out_dim)``. + + A container may say so (``is_transposed``), and a gated fused-experts one shows it either way: + stored ``(num_experts, out_dim, in_dim)``, its first projection takes the dim ``down_proj`` puts + out, which stored transposed it never does. + """ + return sorted( + { + type(module).__name__ + for module in modules + if getattr(module, "is_transposed", False) + or ( + (wrapper := _fused_experts_wrapper(module)) is not None + and getattr(module, wrapper._first_proj_attr).shape[-1] != module.down_proj.shape[1] + ) + } + ) + + +def _eager_only_experts(modules) -> list[str]: + """Names of the expert classes in ``modules`` ModelOpt quantizes on the eager forward only. + + A class with its own quantized wrapper -- GPT-OSS's quantizes down_proj's input by intercepting + the eager forward's bmm/matmul -- only works on the eager forward, and so does a fused-experts + wrapper that says so (FP8's). The registry lookup caches the dynamic class it generates, as + quantizing the module would. Those wrappers register when their module is imported; + register_fused_experts_on_the_fly, which may run after this, registers only the generic + wrapper, which is not eager-only. + """ + return sorted( + { + type(module).__name__ + for module in modules + if (wrapper := QuantModuleRegistry.get(type(module))) is not None + and (not issubclass(wrapper, _QuantFusedExperts) or wrapper._eager_only) + } + ) + + def normalize_experts_impl_on_the_fly(model): """Redirect HF fused-experts modules onto a quantization-supported backend. @@ -1741,19 +1775,8 @@ def normalize_experts_impl_on_the_fly(model): ("grouped_mm",) if _HAS_GROUPED_LINEAR else () ) fallback_impl = "grouped_mm" if _HAS_GROUPED_LINEAR else "eager" - # A class with its own quantized wrapper -- GPT-OSS's quantizes down_proj's input by intercepting - # the eager forward's bmm/matmul -- only works on the eager forward, and so does a fused-experts - # wrapper that says so (FP8's), so keep main's forced eager for such a model. Those register - # when their module is imported; register_fused_experts_on_the_fly, which may run after this, - # registers only the generic wrapper, which is not eager-only. - eager_only = sorted( - { - type(module).__name__ - for module in fused_experts - if (wrapper := QuantModuleRegistry.get(type(module))) is not None - and (not issubclass(wrapper, _QuantFusedExperts) or wrapper._eager_only) - } - ) + # Experts that quantize on the eager forward only keep main's forced eager. + eager_only = _eager_only_experts(fused_experts) def _configs(cfg): if cfg is not None: From 39e769574e5d1f22a991c94436334d01156cfdfc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:21:45 +0000 Subject: [PATCH 2/7] Refuse an expert-parallel load that would calibrate the experts wrong Once the experts load, every rank checks what calibration relies on and refuses a load that would run without an error but compute or calibrate wrong: - a plan whose expert containers match no module of the loaded model; - containers calibration would not take for routed experts by name; - a count neither the config nor the experts' weights show divisible by ep_size; - a rank holding another share of the experts than ep_size gives it; - experts quantization cannot run on the transport's grouped forward: stored transposed, or under a wrapper of their own or an eager-only one. The checks run inside the load's _refused_together block, so a rank refusing alone fails every rank; the share check no longer needs its own all_reduce. The quantization rules and _is_routed_expert are imported at the top of the module. A two-rank test row has rank 1 alone hold half its share: every rank raises its refusal, which fails without the block. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 96 ++++++++ .../unit/torch/utils/test_model_load_utils.py | 232 +++++++++++++++++- 2 files changed, 319 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 74ab4a9486d..9996fff1932 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -52,6 +52,8 @@ 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.model_calib import _is_routed_expert +from modelopt.torch.quantization.plugins.huggingface import _eager_only_experts, _transposed_experts from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe from modelopt.torch.utils.distributed import ( @@ -1602,6 +1604,38 @@ def checkpoint_read() -> bool: return plan, hf_config +def _refuse_unquantizable_experts(expert_modules) -> None: + """Refuse routed experts that expert-parallel PTQ would quantize wrong. + + It quantizes the routed experts as they are, through the grouped GEMM of their ``torch_a2a`` + forward (as for DeepSeek, GLM and MiniMax). Experts stored transposed, ``(num_experts, in_dim, + out_dim)`` like GPT-OSS's, that would leave unquantized. Experts ModelOpt quantizes through a + wrapper of their own, like Llama4's, ``mtq.quantize`` would convert to it, and its forward drops + the process group the all-to-all transport passes in. The generic fused-experts wrapper is no + bar: it quantizes through the grouped GEMM, and ``mtq.quantize`` registers it, for the whole + process, for every fused-experts class it meets, so a second load of the same class sees it. A + fused-experts wrapper that quantizes on the eager forward only, like FP8's, is refused: the + transport runs the grouped GEMM. + """ + # Quantization's own rules: it leaves transposed experts unquantized, and its torch_a2a check + # refuses eager-only ones. This refuses both at load. + transposed = _transposed_experts(expert_modules) + if transposed: + raise NotImplementedError( + f"Expert-parallel PTQ cannot quantize {', '.join(transposed)}: these experts store their " + "weights transposed, which the grouped quantization leaves unquantized. Quantize the " + "model without expert parallelism." + ) + wrapped = _eager_only_experts(expert_modules) + if wrapped: + raise NotImplementedError( + f"Expert-parallel PTQ cannot quantize {', '.join(wrapped)}: ModelOpt quantizes these " + "experts through a wrapper of their own, whose forward drops the expert-parallel " + "transport, or on their eager forward only, which the transport does not run. Quantize " + "the model without expert parallelism." + ) + + def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_state) -> None: """Seed every module's ``_parallel_state`` for calibration's amax sync. @@ -1934,6 +1968,68 @@ def gpu() -> tuple[bool, bool, str | None]: f"{type(model).__name__} matches, so calibration would reduce the experts' amax over " "ranks holding different experts. Quantize the model without expert parallelism." ) + # Calibration tells routed experts by their names, and reduces any other module's weight amax + # over the expert-parallel group, which would merge the scales of different experts. + unnamed = [name for name in expert_modules if not _is_routed_expert(name)] + if unnamed: + raise NotImplementedError( + f"Expert-parallel PTQ cannot calibrate {len(unnamed)} routed-expert container(s), e.g. " + f"{unnamed[0]}: calibration tells routed experts by an 'experts' in their names, so it " + "would merge their weight scales over ranks holding different experts. Quantize the " + "model without expert parallelism." + ) + if not num_experts: + # The config names no count for the check before the load: the experts' weights carry it + # whole as DTensors (transformers 5.16 on; 5.16.0 splits an uneven count without an error). + # A count they do not show either is refused, as the distribution plan's decision refuses + # a count it cannot read. + counts = { + param.shape[0] + for module in expert_modules.values() + for param in module.parameters(recurse=False) + if isinstance(param, DTensor) + } + if not counts or any(count % ep_degree for count in counts): + raise ValueError( + f"{plan!r}: ep_size={ep_degree} must divide the routed experts, which the all-to-all " + "transport shares out equally over the expert-parallel ranks; the config names no " + f"count, and the loaded experts' weights show {sorted(counts) or 'none'}." + ) + # Each rank must hold its share of every container's experts, num_experts // ep_size: the share the + # all-to-all transport routes tokens by. A load that leaves them whole on every rank (a plan whose + # expert weights no style shards) or shards them over another mesh (transformers 5.18 builds one + # of its own for the load, which _experts_sharded_over swaps for the expert-parallel one) would + # run without an error, computing tokens against other experts' weights and merging the scales of + # different experts. A DTensor (transformers 5.16 on) shows the whole count and this rank's share; + # a plain tensor shows the share, checked against the config's count when it names one. A model + # whose containers show neither, as containers owning no weights themselves would, is refused too: + # an unchecked share passes as a correct one. The check reads the containers' weights together, + # so it relies on a model's routed experts sharing one layout, as every transformers MoE's do. + # Every rank refuses alike, whichever ranks hold a wrong share: the guard tells every rank. + held = [ + (param.shape[0], param.to_local().shape[0]) + if isinstance(param, DTensor) + else (num_experts, param.shape[0]) + for module in expert_modules.values() + for param in module.parameters(recurse=False) + if isinstance(param, DTensor) or num_experts + ] + if not held: + raise RuntimeError( + f"{plan!r}: no routed-expert weight shows how its experts are shared out over the " + "expert-parallel ranks, so the all-to-all transport's expert-to-rank mapping cannot " + "be checked. Quantize the model without expert parallelism." + ) + wrong = sorted({(whole, share) for whole, share in held if share * ep_degree != whole}) + if wrong: + whole, share = wrong[0] + raise RuntimeError( + f"{plan!r}: this rank holds {share} of a routed-expert container's {whole} experts, " + f"not {whole // ep_degree}: they are not sharded over this rank's expert-parallel " + "group, so the all-to-all transport would send tokens to ranks without their " + "experts, and calibration would merge the scales of different experts." + ) + _refuse_unquantizable_experts(expert_modules.values()) # Each module's ParallelState for calibration's amax sync: dense modules reduce over the whole world; # routed experts reduce over their expert-data-parallel replicas, and share one input amax across the diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 55c1dc717a2..2930237f784 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -28,7 +28,7 @@ import math import os import warnings -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack, contextmanager, nullcontext from types import SimpleNamespace from unittest import mock @@ -83,6 +83,7 @@ _conversion_plan, _convert_keys, _experts_sharded_over, + _refuse_unquantizable_experts, _resolve_target, _seed_parallel_states, _transformers_mesh_builder, @@ -1885,9 +1886,27 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( # ======================================================================== the expert-parallel load # What the expert-parallel load reads off the loaded model and its config: it takes an expert-parallel -# plan, ``_seed_parallel_states`` gives every module the parallel state its amax syncs over, and -# ``_ckpt_quant_method`` reads the checkpoint's native quantization. Everything runs in one CPU process -# but the bounded two-rank tests at the end. +# plan and keeps its experts out of ModelOpt's per-expert quant wrappers, ``_seed_parallel_states`` +# gives every module the parallel state its amax syncs over, and ``_ckpt_quant_method`` reads the +# checkpoint's native quantization. Everything runs in one CPU process but the bounded two-rank tests +# at the end. + +# ------------------------------------------------------------------------- expert classes + + +class _Experts(nn.Module): + """A layer's routed experts, of a class of their own as a model's are: quantization's registry + and the load's refusals go by it.""" + + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(4, 4, 2)) + self.down_proj = nn.Parameter(torch.zeros(4, 2, 2)) + + +class _SharedExperts(nn.Module): + """The shared experts' class, which the all-to-all transport does not manage.""" + # ------------------------------------------------------------------------------- the plan @@ -2136,6 +2155,120 @@ def from_pretrained(path, **kwargs): load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) +def test_routed_experts_without_experts_in_their_names_are_refused(ep_world_of_two, monkeypatch): + """Calibration tells routed experts by an ``experts`` in their names, and reduces the weight amax + of any other module over the expert-parallel group, merging the scales of different experts.""" + model = model_with_decoder_at("model.") + for layer in model.model.layers: + layer.mlp.routed = _Experts() + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + return model + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + config = SimpleNamespace( + base_model_ep_plan={"layers.*.mlp.routed": "moe_tp_experts"}, num_experts=4 + ) + + with pytest.raises( + NotImplementedError, match=r"2 routed-expert .*model\.layers\.0\.mlp\.routed:" + ): + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + + +@pytest.mark.parametrize( + ("loaded", "held", "named"), + [(4, 2, None), (3, 3, None), (None, None, None), (4, 4, None), (None, None, 4), (None, 0, 4)], + ids=[ + "dividing", + "not-dividing", + "not-shown", + "sharded-over-another-mesh", + "never-sharded", + "owning-no-weights", + ], +) +def test_an_expert_count_the_config_does_not_name_is_read_off_the_loaded_experts( + ep_world_of_two, gloo_world_of_one, monkeypatch, loaded, held, named +): + """A config that names its expert count none of the ways the loader reads is checked against the + loaded experts' weights: transformers keeps the whole count on the DTensors it shards them into, + and 5.16.0 splits a count ep_size does not divide without an error. A count neither shows is + refused, as the distribution plan's decision refuses it. So is a rank holding another share of + the experts than ``ep_size`` gives it: sharded over another mesh than its EP group, or never + sharded, which plain weights show against the count the config names. Experts whose container + owns no weights itself (each expert's in a child module) show no share at all, and are refused + rather than passed unchecked.""" + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, Shard + + class _PastError(Exception): + """Raised by the load's next check: the experts' count passed.""" + + model = model_with_decoder_at("model.") + if held == 0: + for layer in model.model.layers: + experts, weights = layer.mlp.experts, nn.Module() + for name, param in list(experts.named_parameters(recurse=False)): + delattr(experts, name) + weights.register_parameter(name, param) + experts.add_module("weights", weights) + if loaded is not None: + mesh = init_device_mesh("cpu", (1,)) + for layer in model.model.layers: + experts = layer.mlp.experts + for name, param in list(experts.named_parameters(recurse=False)): + whole = torch.zeros(loaded, *param.shape[1:]) + # This rank's share of a count sharded over two ranks, or of another split. + shard = DTensor.from_local( + whole[:held], + mesh, + [Shard(0)], + run_check=False, + shape=whole.shape, + stride=whole.stride(), + ) + setattr(experts, name, nn.Parameter(shard)) + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + return model + + def next_check(experts): + raise _PastError + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "_refuse_unquantizable_experts", next_check) + config = SimpleNamespace( + base_model_ep_plan=_EP_PLAN, **({"num_experts": named} if named else {}) + ) + if (loaded, held) == (4, 2): + outcome = pytest.raises(_PastError) + elif held == 0: + outcome = pytest.raises(RuntimeError, match="no routed-expert weight shows how") + elif loaded == 4 or named: + outcome = pytest.raises( + RuntimeError, match="holds 4 of a routed-expert container's 4 experts" + ) + else: + shown = r"\[3\]" if loaded == 3 else "none" + outcome = pytest.raises(ValueError, match=f"weights show {shown}") + + with outcome: + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) + + @pytest.mark.parametrize("declared", [True, False], ids=["on-the-config", "a-class-default"]) def test_the_model_keeps_the_plan_its_checkpoint_declares( ep_world_of_two, gloo_world_of_one, monkeypatch, declared @@ -2274,6 +2407,80 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): assert classes == [True] +# ------------------------------------------------------------------------ quant wrappers + + +@pytest.fixture +def quant_wrappers(): + """ModelOpt quant wrappers registered for the experts' and the shared experts' classes.""" + from modelopt.torch.quantization.nn.modules.quant_module import QuantModule, QuantModuleRegistry + + class _PerExpertQuantWrapper(QuantModule): + pass + + classes = (_Experts, _SharedExperts) + QuantModuleRegistry.register({cls: f"test.{cls.__name__}" for cls in classes})( + _PerExpertQuantWrapper + ) + yield QuantModuleRegistry + for cls in classes: + QuantModuleRegistry.unregister(cls) + + +def test_routed_experts_with_a_quant_wrapper_of_their_own_are_refused(quant_wrappers): + """Quantization would convert them to it, whose forward drops the process group the all-to-all + transport passes in. The registry serves every model in the process, so no registration is + added or removed; the shared experts, which the transport does not manage, may keep a wrapper.""" + with pytest.raises(NotImplementedError, match=r"cannot quantize _Experts: ModelOpt quantizes"): + _refuse_unquantizable_experts([_Experts(), _Experts()]) + + assert _Experts in quant_wrappers + assert _SharedExperts in quant_wrappers + + +@pytest.mark.parametrize( + ("wrapper", "outcome"), + [ + ("_QuantFusedExperts", nullcontext()), + ("_QuantFP8Experts", pytest.raises(NotImplementedError, match="on their eager forward")), + ], + ids=["grouped-gemm", "eager-only"], +) +def test_routed_experts_under_a_fused_experts_wrapper(wrapper, outcome): + """mtq.quantize registers the generic fused-experts wrapper, for the whole process, for every + fused-experts class it meets, and that wrapper quantizes through the grouped GEMM: a second + expert-parallel load of the same class is not refused for it. FP8's applies its block scales on + the eager forward only, which the all-to-all transport does not run: refused.""" + from modelopt.torch.quantization.nn.modules.quant_module import QuantModuleRegistry + from modelopt.torch.quantization.plugins import huggingface + + QuantModuleRegistry.register({_Experts: "test._Experts"})(getattr(huggingface, wrapper)) + try: + with outcome: + _refuse_unquantizable_experts([_Experts(), _Experts()]) + finally: + QuantModuleRegistry.unregister(_Experts) + + +@pytest.mark.parametrize("layout", ["says-so", "by-its-shapes", "standard"]) +def test_transposed_experts_are_refused(layout): + """The grouped quantization leaves experts stored transposed (GPT-OSS's) unquantized: the model + would be quantized everywhere but its experts. Quantization tells them by their flag or, for + gated experts that do not say so, by their shapes; so does the load.""" + experts = [_Experts(), _Experts()] + for module in experts: + if layout == "says-so": + module.is_transposed = True + else: # a fused-experts container: 4 experts, gate and up 2 x 2 wide, on hidden size 2 + module.num_experts = 4 + if layout == "by-its-shapes": # (num_experts, in_dim, out_dim) + module.gate_up_proj = nn.Parameter(torch.zeros(4, 2, 4)) + refused = pytest.raises(NotImplementedError, match="cannot quantize _Experts: these experts") + + with refused if layout != "standard" else nullcontext(): + _refuse_unquantizable_experts(experts) + + # ----------------------------------------------------------------------- _seed_parallel_states @@ -2459,7 +2666,8 @@ def _meshes_in_step(*args, **kwargs): def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): path, hf_config, patches = ckpt, None, [] if failure != "unreadable-config": - hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=4) + # 8 experts over 2 ranks: the stand-in model holds this rank's 4. + hf_config = SimpleNamespace(base_model_ep_plan=dict(_EP_PLAN), num_experts=8) if failure != "unseen-checkpoint": local = mock.patch.object( model_load_utils, "ensure_local_checkpoint", side_effect=lambda p, **k: (None, p) @@ -2481,8 +2689,12 @@ def fail(*args, **kwargs): elif failure == "unseen-checkpoint" and fails_here: # Rank 0 makes the checkpoint local; a rank on a node without that storage cannot see it. patches.append(mock.patch.object(hf_checkpoint_utils.os.path, "isdir", return_value=False)) - if failure in ("load", "record"): + if failure in ("load", "record", "share"): loaded = model_with_decoder_at("model.") + if failure == "share" and fails_here: + for layer in loaded.model.layers: # 2 of the 8 experts here, not 4 + layer.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(2, 4, 2)) + layer.mlp.experts.down_proj = nn.Parameter(torch.zeros(2, 2, 2)) load = fail if failure == "load" and fails_here else (lambda path, **kwargs: loaded) loader = SimpleNamespace(__name__="Loader", from_pretrained=load) record = fail if failure == "record" and fails_here else None @@ -2522,15 +2734,17 @@ def fail(*args, **kwargs): ("no-gpu", 1, "rank(s) [1] have none while the others have theirs"), ("load", 1, "the GPU ran out of memory"), ("record", 1, "a weight file this rank cannot read"), + ("share", 1, "this rank holds 2 of a routed-expert container's 8 experts"), ], ) def test_a_load_failing_on_one_rank_fails_every_rank( cpu_dist_workers, tmp_path, failure, failing_rank, match ): """One rank fails alone: reading the config, on its transformers, on its arguments, seeing the - checkpoint rank 0 made local, without a GPU of its own where the others have theirs, or in the - load. Every rank raises that failure, rather than leave the others waiting for it in the meshes - (stood in by a bounded barrier) or the warmup's collectives.""" + checkpoint rank 0 made local, without a GPU of its own where the others have theirs, in the + load, or holding another share of the experts. Every rank raises that failure, rather than leave + the others waiting for it in the meshes (stood in by a bounded barrier) or the warmup's + collectives.""" Qwen3MoeConfig(num_experts=4, num_hidden_layers=1, hidden_size=8).save_pretrained(tmp_path) cpu_dist_workers.run( _one_rank_fails_job, From 8cfcde6261c09288ebab6ff55f2ef603d3e99494 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:35:45 +0000 Subject: [PATCH 3/7] Import the quantization package only where the EP refusals use it The HF quantization plugin loads under import_plugin's guard, so a transformers it cannot load is a warning and the other loaders keep working. Importing it at the top of model_load_utils made that a hard failure for load_pretrained, fsdp2_wrap and record_unplaced_source_keys, none of which need it. The plugin's two rules are imported inside _refuse_unquantizable_experts again, and _is_routed_expert next to the check that uses it, as fsdp2_wrap imports the quantization package where it needs it. A test imports the module with the plugin unavailable; it fails with the imports at the top. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 14 +++++++++++--- tests/unit/torch/utils/test_model_load_utils.py | 14 ++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 9996fff1932..066d9528144 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -52,8 +52,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.model_calib import _is_routed_expert -from modelopt.torch.quantization.plugins.huggingface import _eager_only_experts, _transposed_experts from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe from modelopt.torch.utils.distributed import ( @@ -1617,6 +1615,13 @@ def _refuse_unquantizable_experts(expert_modules) -> None: fused-experts wrapper that quantizes on the eager forward only, like FP8's, is refused: the transport runs the grouped GEMM. """ + # Imported here: the HF quantization plugin loads under import_plugin, so a release it cannot + # load leaves this module's other loaders working; only this refusal needs it. + from modelopt.torch.quantization.plugins.huggingface import ( + _eager_only_experts, + _transposed_experts, + ) + # Quantization's own rules: it leaves transposed experts unquantized, and its torch_a2a check # refuses eager-only ones. This refuses both at load. transposed = _transposed_experts(expert_modules) @@ -1969,7 +1974,10 @@ def gpu() -> tuple[bool, bool, str | None]: "ranks holding different experts. Quantize the model without expert parallelism." ) # Calibration tells routed experts by their names, and reduces any other module's weight amax - # over the expert-parallel group, which would merge the scales of different experts. + # over the expert-parallel group, which would merge the scales of different experts. Imported + # here, as fsdp2_wrap does: the quantization package sits above this module. + from modelopt.torch.quantization.model_calib import _is_routed_expert + unnamed = [name for name in expert_modules if not _is_routed_expert(name)] if unnamed: 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 2930237f784..032387570d1 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -24,9 +24,11 @@ import copy import dataclasses +import importlib import inspect import math import os +import sys import warnings from contextlib import ExitStack, contextmanager, nullcontext from types import SimpleNamespace @@ -75,6 +77,7 @@ 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 import plugins from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import hf_checkpoint_utils, model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( @@ -2481,6 +2484,17 @@ def test_transposed_experts_are_refused(layout): _refuse_unquantizable_experts(experts) +def test_the_loaders_import_without_the_hf_quantization_plugin(monkeypatch): + """The HF quantization plugin loads under import_plugin's guard, so a release it cannot load is + a warning; only the expert-parallel refusals need it, so the loaders still import without it.""" + monkeypatch.setitem(sys.modules, "modelopt.torch.quantization.plugins.huggingface", None) + # Re-imported below: put the module and its package attribute back afterwards. + monkeypatch.delitem(sys.modules, model_load_utils.__name__) + monkeypatch.setattr(plugins, "model_load_utils", model_load_utils) + + assert importlib.import_module(model_load_utils.__name__).load_pretrained + + # ----------------------------------------------------------------------- _seed_parallel_states From ef3f98f426df1e3ceb18733fc0b4a78e2c6418de Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:41:30 +0000 Subject: [PATCH 4/7] Shard the unnamed-container test's experts as the loader now requires entry under it before anything loads, so this test's container-only plan no longer reached the refusal it checks. Its plan now shards the container's gate_up_proj through grouped_gemm, as a real plan does. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 6 +++++- 1 file changed, 5 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 032387570d1..4d8ed7425c0 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2176,7 +2176,11 @@ def from_pretrained(path, **kwargs): monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) config = SimpleNamespace( - base_model_ep_plan={"layers.*.mlp.routed": "moe_tp_experts"}, num_experts=4 + base_model_ep_plan={ + "layers.*.mlp.routed": "moe_tp_experts", + "layers.*.mlp.routed.gate_up_proj": "grouped_gemm", + }, + num_experts=4, ) with pytest.raises( From ef18547dd808e1e71db52c14a8b2e90262dea881 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:28:12 +0000 Subject: [PATCH 5/7] Read the experts once per refusal rule, and say what the share check compares The unquantizable-experts refusal reads its experts twice, once per rule, so it takes them as a tuple: a one-shot iterable would let the second rule pass on nothing. The share check's comment says it compares counts, catching a wrong share rather than a same-size mesh, and its test row is named for the mesh of another size it builds. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 18 ++++++++++-------- .../unit/torch/utils/test_model_load_utils.py | 15 ++++++++------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 066d9528144..3d84e48e2d9 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1622,6 +1622,8 @@ def _refuse_unquantizable_experts(expert_modules) -> None: _transposed_experts, ) + # Read once per rule: a one-shot iterable would let the second rule pass on nothing. + expert_modules = tuple(expert_modules) # Quantization's own rules: it leaves transposed experts unquantized, and its torch_a2a check # refuses eager-only ones. This refuses both at load. transposed = _transposed_experts(expert_modules) @@ -2005,14 +2007,14 @@ def gpu() -> tuple[bool, bool, str | None]: ) # Each rank must hold its share of every container's experts, num_experts // ep_size: the share the # all-to-all transport routes tokens by. A load that leaves them whole on every rank (a plan whose - # expert weights no style shards) or shards them over another mesh (transformers 5.18 builds one - # of its own for the load, which _experts_sharded_over swaps for the expert-parallel one) would - # run without an error, computing tokens against other experts' weights and merging the scales of - # different experts. A DTensor (transformers 5.16 on) shows the whole count and this rank's share; - # a plain tensor shows the share, checked against the config's count when it names one. A model - # whose containers show neither, as containers owning no weights themselves would, is refused too: - # an unchecked share passes as a correct one. The check reads the containers' weights together, - # so it relies on a model's routed experts sharing one layout, as every transformers MoE's do. + # expert weights no style shards) or shards them over a mesh of another size would run without an + # error, computing tokens against other experts' weights and merging the scales of different + # experts. The check compares counts: it catches a wrong share, not a same-size mesh. A DTensor + # (transformers 5.16 on) shows the whole count and this rank's share; a plain tensor shows the + # share, checked against the config's count when it names one. A model whose containers show + # neither, as containers owning no weights themselves would, is refused too: an unchecked share + # passes as a correct one. The check reads the containers' weights together, so it relies on a + # model's routed experts sharing one layout, as every transformers MoE's do. # Every rank refuses alike, whichever ranks hold a wrong share: the guard tells every rank. held = [ (param.shape[0], param.to_local().shape[0]) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 4d8ed7425c0..e3aef8ab7f5 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2196,7 +2196,7 @@ def from_pretrained(path, **kwargs): "dividing", "not-dividing", "not-shown", - "sharded-over-another-mesh", + "sharded-over-a-mesh-of-another-size", "never-sharded", "owning-no-weights", ], @@ -2208,10 +2208,10 @@ def test_an_expert_count_the_config_does_not_name_is_read_off_the_loaded_experts loaded experts' weights: transformers keeps the whole count on the DTensors it shards them into, and 5.16.0 splits a count ep_size does not divide without an error. A count neither shows is refused, as the distribution plan's decision refuses it. So is a rank holding another share of - the experts than ``ep_size`` gives it: sharded over another mesh than its EP group, or never - sharded, which plain weights show against the count the config names. Experts whose container - owns no weights itself (each expert's in a child module) show no share at all, and are refused - rather than passed unchecked.""" + the experts than ``ep_size`` gives it: sharded over a mesh of another size than its EP group, + or never sharded, which plain weights show against the count the config names. Experts whose + container owns no weights itself (each expert's in a child module) show no share at all, and are + refused rather than passed unchecked.""" from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor, Shard @@ -2437,9 +2437,10 @@ class _PerExpertQuantWrapper(QuantModule): def test_routed_experts_with_a_quant_wrapper_of_their_own_are_refused(quant_wrappers): """Quantization would convert them to it, whose forward drops the process group the all-to-all transport passes in. The registry serves every model in the process, so no registration is - added or removed; the shared experts, which the transport does not manage, may keep a wrapper.""" + added or removed; the shared experts, which the transport does not manage, may keep a wrapper. + The experts come as a one-shot iterable, which the transposed-experts rule reads first.""" with pytest.raises(NotImplementedError, match=r"cannot quantize _Experts: ModelOpt quantizes"): - _refuse_unquantizable_experts([_Experts(), _Experts()]) + _refuse_unquantizable_experts(_Experts() for _ in range(2)) assert _Experts in quant_wrappers assert _SharedExperts in quant_wrappers From d0a4ea724d0f67c2350880c53494db16917c04e6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:28:38 +0000 Subject: [PATCH 6/7] Say in the loader tests' docstring that the loader refuses unquantizable experts It refuses them, and leaves the quant-wrapper registry alone, rather than keeping them out of ModelOpt's per-expert wrappers. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 2 +- 1 file changed, 1 insertion(+), 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 e3aef8ab7f5..e1823746a49 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -1889,7 +1889,7 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( # ======================================================================== the expert-parallel load # What the expert-parallel load reads off the loaded model and its config: it takes an expert-parallel -# plan and keeps its experts out of ModelOpt's per-expert quant wrappers, ``_seed_parallel_states`` +# plan and refuses routed experts expert-parallel PTQ would quantize wrong, ``_seed_parallel_states`` # gives every module the parallel state its amax syncs over, and ``_ckpt_quant_method`` reads the # checkpoint's native quantization. Everything runs in one CPU process but the bounded two-rank tests # at the end. From 740ba93d1a24a25a2534b91df5d567ac3321439b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:35:45 +0000 Subject: [PATCH 7/7] Stand in a list for the unplaced keys in the expert-refusal tests record_unplaced_source_keys returns a sorted list, which a later check iterates; these tests' stand-ins returned None. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index e1823746a49..5c39ce5c825 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2174,7 +2174,7 @@ def from_pretrained(path, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) config = SimpleNamespace( base_model_ep_plan={ "layers.*.mlp.routed": "moe_tp_experts", @@ -2255,7 +2255,7 @@ def next_check(experts): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) monkeypatch.setattr(model_load_utils, "_refuse_unquantizable_experts", next_check) config = SimpleNamespace( base_model_ep_plan=_EP_PLAN, **({"num_experts": named} if named else {})