From cb00b2f2dcc106c1cc48e5ba6f354a7d16ce54f4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:21:35 +0000 Subject: [PATCH 01/10] 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 76d4dfed1d21ce95b460df834921b49a1fc20817 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:21:45 +0000 Subject: [PATCH 02/10] 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 | 235 +++++++++++++++++- 2 files changed, 320 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 6a84af3be42..29622095993 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -60,6 +60,8 @@ from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key from modelopt.torch.export.model_utils import get_language_model_from_vl +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.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe @@ -2145,6 +2147,38 @@ def _ep_aware(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs) hf_modeling_utils.get_total_byte_count = orig +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. @@ -2822,6 +2856,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()) # The dense part shards over the whole world, the experts' ranks: a forward's dense all-gathers # and expert all-to-alls then run in turn on the same ranks, so they cannot deadlock. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 11018b22be2..0a3190a3c19 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -100,6 +100,7 @@ _fsdp2_shard_dense, _load_ep_dense_distributed, _make_expert_param_matcher, + _refuse_unquantizable_experts, _resolve_target, _seed_parallel_states, _transformers_mesh_builder, @@ -2556,9 +2557,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 @@ -2829,6 +2848,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( ("source", "on_the_config"), [ @@ -3218,6 +3351,80 @@ def from_pretrained(path, config, **kwargs): ) +# ------------------------------------------------------------------------ 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 @@ -3414,11 +3621,12 @@ def _one_rank_fails_job(rank, size, ckpt, failure, failing_rank, match): path, hf_config, patches = ckpt, None, [] if failure == "unreadable-plan-provenance": # A VLM's config, handed in: the plan to keep is read from where it sits in config.json. - plan = {"base_model_ep_plan": dict(_EP_PLAN), "num_experts": 4} + plan = {"base_model_ep_plan": dict(_EP_PLAN), "num_experts": 8} hf_config = SimpleNamespace(model_type=_NO_SPEC, text_config=SimpleNamespace(**plan)) elif failure != "unreadable-config": + # 8 experts over 2 ranks: the stand-in model holds this rank's 4. hf_config = SimpleNamespace( - model_type=_NO_SPEC, base_model_ep_plan=dict(_EP_PLAN), num_experts=4 + model_type=_NO_SPEC, base_model_ep_plan=dict(_EP_PLAN), num_experts=8 ) if failure != "unseen-checkpoint": local = mock.patch.object( @@ -3445,8 +3653,12 @@ def fail(*args, **kwargs): # This rank reads a config.json that is no JSON object: it declares no plan, not an error. reads = SimpleNamespace(load=lambda f: []) patches.append(mock.patch.object(model_load_utils, "json", reads)) - if failure in ("load", "record", "unreadable-plan-provenance"): + if failure in ("load", "record", "unreadable-plan-provenance", "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) if failure == "unreadable-plan-provenance": load = _load_in_step @@ -3500,17 +3712,18 @@ def fail(*args, **kwargs): ("load", 1, "the GPU ran out of memory"), ("record", 1, "a weight file this rank cannot read"), ("unreadable-plan-provenance", 1, _InTheLoadError.__name__), + ("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. A config.json one rank cannot read - as a JSON object declares no plan to keep there, so every rank goes on into the load (stood in - by a bounded barrier too).""" + 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. A config.json one rank cannot read as a JSON object declares no plan to keep there, + so every rank goes on into the load (stood in by a bounded barrier too).""" Qwen3MoeConfig(num_experts=4, num_hidden_layers=1, hidden_size=8).save_pretrained(tmp_path) cpu_dist_workers.run( _one_rank_fails_job, From 7fc9e9ff32133f32437e022f97d5dacddfb6c415 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:35:45 +0000 Subject: [PATCH 03/10] 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 29622095993..2c9922159f3 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -60,8 +60,6 @@ from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key from modelopt.torch.export.model_utils import get_language_model_from_vl -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.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import env, moe @@ -2160,6 +2158,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) @@ -2857,7 +2862,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 0a3190a3c19..2a03630b4fe 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -24,10 +24,12 @@ import copy import dataclasses +import importlib import inspect import json import math import os +import sys import warnings from contextlib import ExitStack, contextmanager, nullcontext from types import SimpleNamespace @@ -86,6 +88,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.moe import _EP_MIN_TRANSFORMERS from modelopt.torch.utils.plugins import hf_checkpoint_utils, model_load_utils @@ -3425,6 +3428,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 a38c47c76b2462d76b6fc54bf306b46d9fea427b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:41:30 +0000 Subject: [PATCH 04/10] 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 2a03630b4fe..0ab1c553500 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2869,7 +2869,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 c917396526f138eb69b3b3fb6d671dabdaa936a2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:28:12 +0000 Subject: [PATCH 05/10] 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 2c9922159f3..fea599319c9 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2165,6 +2165,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) @@ -2893,14 +2895,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 0ab1c553500..d71f2011193 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2889,7 +2889,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", ], @@ -2901,10 +2901,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 @@ -3381,9 +3381,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 16902fa58db5f481712fb081f4b9cc44c55d306f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:28:38 +0000 Subject: [PATCH 06/10] 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 d71f2011193..d5b67b92a44 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2560,7 +2560,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 f857a052cc0fd723ab4cf20a4b13d0d6fd15e0cd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:35:45 +0000 Subject: [PATCH 07/10] 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 d5b67b92a44..a46674d13e4 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2867,7 +2867,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", @@ -2948,7 +2948,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 {}) From 65353fe69906454db5f16f13203074ed894741e3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 07:12:28 +0000 Subject: [PATCH 08/10] Say that only is_transposed is read for a non-gated experts container _transposed_experts finds a transposed gated container by its shapes, but a non-gated one (up_proj, as Nemotron-H's) matches no _fused_experts_wrapper, so only its is_transposed flag is read and its own wrapper is responsible for the layout. The docstring now says so, as promised in review. Docstring only. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/plugins/huggingface.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 60669054c0f..851b4283325 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1701,7 +1701,9 @@ def _transposed_experts(modules) -> list[str]: 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. + out, which stored transposed it never does. A non-gated container (``up_proj``, as Nemotron-H's) + matches no :func:`_fused_experts_wrapper`, so only its ``is_transposed`` is read, and its own + wrapper is responsible for the layout. """ return sorted( { From e8ea09e6f9f58f07e8d634664f3a8a008dfd1556 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 08:20:29 +0000 Subject: [PATCH 09/10] Give the expert-refusal tests' stand-in configs a model type With #2780 in main, load_and_prepare_ep looks up the model spec by hf_config.model_type before it loads. The stand-in configs of test_routed_experts_without_experts_in_their_names_are_refused and of test_an_expert_count_the_config_does_not_name_is_read_off_the_loaded_experts named none, so all 7 cases raised AttributeError instead of reaching the refusals they test. They now name _NO_SPEC, as main's other load test doubles do. Test only. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 5 ++++- 1 file changed, 4 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 a46674d13e4..a220cac9e41 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2869,6 +2869,7 @@ 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: []) config = SimpleNamespace( + model_type=_NO_SPEC, base_model_ep_plan={ "layers.*.mlp.routed": "moe_tp_experts", "layers.*.mlp.routed.gate_up_proj": "grouped_gemm", @@ -2951,7 +2952,9 @@ def next_check(experts): 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 {}) + model_type=_NO_SPEC, + base_model_ep_plan=_EP_PLAN, + **({"num_experts": named} if named else {}), ) if (loaded, held) == (4, 2): outcome = pytest.raises(_PastError) From 01f034c70e8938058d27522437916e59b89dd5fc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 17:54:04 +0000 Subject: [PATCH 10/10] Let the expert-refusal tests' stand-in record take keywords With #3137 in main, load_and_prepare_ep passes tensor_files= to record_unplaced_source_keys. The stand-ins of test_routed_experts_without_experts_in_their_names_are_refused and of test_an_expert_count_the_config_does_not_name_is_read_off_the_loaded_experts took positional arguments only, so all 7 cases raised TypeError instead of reaching the refusals they test. They take keywords too now, as main's other stand-ins do. Test only. 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 a220cac9e41..9e7190d4487 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2867,7 +2867,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: []) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args, **kwargs: []) config = SimpleNamespace( model_type=_NO_SPEC, base_model_ep_plan={ @@ -2949,7 +2949,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: []) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args, **kwargs: []) monkeypatch.setattr(model_load_utils, "_refuse_unquantizable_experts", next_check) config = SimpleNamespace( model_type=_NO_SPEC,