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: diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 74ab4a9486d..3d84e48e2d9 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1602,6 +1602,47 @@ 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. + """ + # 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, + ) + + # 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) + 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 +1975,71 @@ 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. 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( + 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 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]) + 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..5c39ce5c825 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -24,11 +24,13 @@ import copy import dataclasses +import importlib import inspect import math import os +import sys import warnings -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack, contextmanager, nullcontext from types import SimpleNamespace from unittest import mock @@ -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 ( @@ -83,6 +86,7 @@ _conversion_plan, _convert_keys, _experts_sharded_over, + _refuse_unquantizable_experts, _resolve_target, _seed_parallel_states, _transformers_mesh_builder, @@ -1885,9 +1889,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 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. + +# ------------------------------------------------------------------------- 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 +2158,124 @@ 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: []) + config = SimpleNamespace( + base_model_ep_plan={ + "layers.*.mlp.routed": "moe_tp_experts", + "layers.*.mlp.routed.gate_up_proj": "grouped_gemm", + }, + 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-a-mesh-of-another-size", + "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 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 + + 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: []) + 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 +2414,92 @@ 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. + 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() for _ in range(2)) + + 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) + + +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 @@ -2459,7 +2685,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 +2708,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 +2753,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,