From e2a32bb42a7781f8a8b59a033b087ad9bbd92f5b Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 14 Jul 2026 22:49:10 +0000 Subject: [PATCH 01/16] squash: pr1828 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 86 +++++++++++++++++ modelopt/torch/export/layer_utils.py | 32 +++---- modelopt/torch/export/modeling/__init__.py | 26 +++++ modelopt/torch/export/modeling/base.py | 50 ++++++++++ .../export/modeling/families/__init__.py | 18 ++++ .../torch/export/modeling/families/dbrx.py | 27 ++++++ .../export/modeling/families/deepseek.py | 27 ++++++ .../torch/export/modeling/families/gemma.py | 29 ++++++ .../torch/export/modeling/families/gptoss.py | 28 ++++++ .../torch/export/modeling/families/mixtral.py | 39 ++++++++ .../export/modeling/families/nemotron.py | 29 ++++++ .../torch/export/modeling/families/qwen.py | 33 +++++++ modelopt/torch/export/modeling/registry.py | 49 ++++++++++ .../unit/torch/export/test_modeling_specs.py | 94 +++++++++++++++++++ 14 files changed, 549 insertions(+), 18 deletions(-) create mode 100644 modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md create mode 100644 modelopt/torch/export/modeling/__init__.py create mode 100644 modelopt/torch/export/modeling/base.py create mode 100644 modelopt/torch/export/modeling/families/__init__.py create mode 100644 modelopt/torch/export/modeling/families/dbrx.py create mode 100644 modelopt/torch/export/modeling/families/deepseek.py create mode 100644 modelopt/torch/export/modeling/families/gemma.py create mode 100644 modelopt/torch/export/modeling/families/gptoss.py create mode 100644 modelopt/torch/export/modeling/families/mixtral.py create mode 100644 modelopt/torch/export/modeling/families/nemotron.py create mode 100644 modelopt/torch/export/modeling/families/qwen.py create mode 100644 modelopt/torch/export/modeling/registry.py create mode 100644 tests/unit/torch/export/test_modeling_specs.py diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md new file mode 100644 index 00000000000..800f0da743a --- /dev/null +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -0,0 +1,86 @@ +# Export: Model-Specific Logic Refactor — Action Plan + +**Goal:** Give the unified HF export path a per-model-family data registry +(`modelopt/torch/export/modeling/`), so that supporting a new model family means +adding one declarative spec file instead of editing if/elif chains across the +export engine. + +**Scope:** The unified HF export path (`unified_export_hf.py` + its helpers) only. + +- The **Megatron** path (`plugins/mcore_*`) already follows a per-family registry + pattern and is untouched. +- The **TRT-LLM** checkpoint path (`model_config_export.py`, the `build_*` + functions in `layer_utils.py`, `tensorrt_llm_utils.py`) is legacy: the plan is + to move it out as-is (separate track), not to refactor it. Its per-model + branches stay where they are. + +## 1. Where the HF path stands after PR #1939 + +PR #1939 gave the HF path registry-based **module dispatch**: `ExportModuleRegistry` +and `PrepareMoEInputsRegistry` (`registry.py`, `hf_export_handlers.py`) select +*which handler processes a module* by class/predicate match, replacing the if/elif +chains that used to live in `unified_export_hf.py`. + +What is still hardcoded is the **per-family data** those handlers (and other HF-path +helpers) consume. Inventory of model-specific logic reachable from the HF path: + +| Item | Location | Kind | +|---|---|---| +| MoE expert linear names (Qwen/DeepSeek/Mixtral/DBRX/GptOss/NemotronH/Gemma4) | `layer_utils.get_expert_linear_names` | data (if/elif chain) | +| Duplicate expert-naming table + iterable-experts support gate | `layer_utils.get_experts_list` | data (if/elif chain) | +| MoE block class-name list | `layer_utils.is_moe` | data (shared with TRT-LLM path) | +| AWQ `pre_quant_scale` fusion rules (Llama/Qwen3) | `quant_utils.PQS_FUSE_MODULE_MAPPING` | data (table) | +| weight+1 layernorm class names (Gemma RMSNorm, LayerNorm1P) | `quant_utils._layernorm_uses_weight_plus_one` | data | +| Handler match keys (Llama4TextExperts, GptOssExperts, DbrxExperts, QuantMoELinear) | `hf_export_handlers.py` | dispatch (stays: structural, per-module) | +| Fused-expert gated/non-gated split (`gate_up_proj` vs `up_proj`) | `moe_utils.py` | data + structure | +| dummy-forward special cases (Whisper input, Nemotron-VL tower) | `unified_export_hf.requantize_resmooth_fused_llm_layers` | behavior | +| VLM language-tower extraction, DiffusionGemma tied-key reorder | `model_utils.py` | behavior | + +## 2. Design + +Two layers: + +- **Engine** (existing files, organized by operation): owns all algorithms and the + module walk; consults the registry for per-family values. +- **Modeling library** (`modeling/`, organized by model family): declarative + per-family **data only**. No export logic, no imports from the parent package + (only `torch.nn` + stdlib), so it sits at the bottom of the dependency graph. + +```text +modeling/ + base.py # ModelSpec dataclass — the per-family contract + registry.py # register() + lookups; returns None when unmatched + __init__.py # re-exports; importing it registers all families + families/ # one small file per family; import == registration +``` + +Call sites follow a fallback-first shape, which keeps migration incremental and +behavior-preserving — a family not in the registry behaves exactly as before: + +```python +spec = match_moe_block(module) +if spec is not None and spec.expert_linear_names is not None: + return list(spec.expert_linear_names) +# ... legacy branch preserved as fallback ... +``` + +Note on naming: `modeling/registry.py` (per-family **data**, "what are this +family's values") is distinct from the top-level `registry.py` from PR #1939 +(per-module **dispatch**, "which handler processes this module"). The two layers +compose: handlers look up family data through `modeling/`. + +## 3. Migration plan + +Each step is one PR with a fallback to the legacy path and an equivalence check +against existing export tests. + +| Step | What | Status | +|---|---|---| +| **P1** | Registry skeleton + MoE expert naming: `get_expert_linear_names` and `get_experts_list` read `spec.expert_linear_names` / `spec.has_iterable_experts`. The #1 "add a MoE model" shotgun-surgery driver. | this PR | +| **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated across specs. | planned | +| **P3** | `is_moe` class-name list → `spec.moe_block_names` (shared consumer: keep fallback until the TRT-LLM path moves out). | planned | +| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs. | planned | +| **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Dead code found during inventory (`MODEL_NAME_TO_TYPE`, `get_model_type`, `adjust_attn_amax_values`, `update_experts_avg_prequant_scale`) is deleted on that track too. | separate track | + +**Guardrails:** fallback-first; one data category per PR; the engine keeps the +algorithms — families supply values only, never fork functions. diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index d5f1fb2330d..29a6b589104 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -60,6 +60,7 @@ RgLruConfig, ) from .model_config_utils import pad_weights +from .modeling import match_moe_block from .postprocess import view_as_float8_e4m3fn_if_needed, view_as_uint8_if_needed from .quant_utils import ( get_activation_scaling_factor, @@ -93,25 +94,14 @@ def get_experts_list( """ experts_list = [] - # Define linear layer names for different model types - if "mixtralforcausallm" in model_type: - linear_names = ["w1", "w2", "w3"] - elif any( - qwen_variant in model_type - for qwen_variant in [ - "qwenmoeforcausallm", - "qwen2moeforcausallm", - "qwen3moeforcausallm", - "qwen3nextforcausallm", - ] - ): - linear_names = ["gate_proj", "down_proj", "up_proj"] - elif "nemotronhforcausallm" in model_type: - linear_names = ["up_proj", "down_proj"] - elif "gemma4" in model_type: - linear_names = ["gate_proj", "down_proj", "up_proj"] - else: + # Expert linear names are per-family data (modeling/families/*). The grouped export + # path only supports families whose experts are iterable per-expert sub-modules + # (Mixtral, Qwen MoE, NemotronH, Gemma4); stacked/fused layouts (DBRX, GptOss, ...) + # raise NotImplementedError here and are handled by other paths. + spec = match_moe_block(module) + if spec is None or not spec.has_iterable_experts or spec.expert_linear_names is None: raise NotImplementedError(f" {model_type} not supported") + linear_names = list(spec.expert_linear_names) # Common logic for all supported model types experts_list.extend( @@ -992,6 +982,12 @@ def module_match_name_list(module, name_list): if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"): return [first_proj_attr, "down_proj"] + # Resolve expert names from the model-family registry; fall back to the mapping + # below when no family spec matches the MoE block. + spec = match_moe_block(module) + if spec is not None and spec.expert_linear_names is not None: + return list(spec.expert_linear_names) + if module_match_name_list( module, [ diff --git a/modelopt/torch/export/modeling/__init__.py b/modelopt/torch/export/modeling/__init__.py new file mode 100644 index 00000000000..5718219e50e --- /dev/null +++ b/modelopt/torch/export/modeling/__init__.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Model-family export descriptors. + +Holds per-model export data organized by model family. The export code resolves a +``ModelSpec`` via the registry lookups and reads its fields; an unmatched lookup +returns ``None`` so callers fall back to their default behavior. +""" + +# Importing families registers every family spec as a side effect. +from . import families +from .base import * +from .registry import * diff --git a/modelopt/torch/export/modeling/base.py b/modelopt/torch/export/modeling/base.py new file mode 100644 index 00000000000..72c64c0ebcd --- /dev/null +++ b/modelopt/torch/export/modeling/base.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model-family export descriptor. + +A ``ModelSpec`` declares how one model family differs from the generic export path, +so the export code can read these values instead of branching on model names. Each +spec holds per-model data only, not export logic. +""" + +from dataclasses import dataclass + +__all__ = ["ModelSpec"] + + +@dataclass +class ModelSpec: + """Per-model-family export data. + + A spec is resolved from a model sub-module via its matching keys and read for its + per-model data fields. ``moe_block_names`` is the matching key: MoE block + class-name substrings compared case-insensitively against + ``type(module).__name__`` (e.g. ``"Qwen3MoeSparseMoeBlock"``). + """ + + name: str + """Unique name of the model-family spec.""" + + moe_block_names: tuple[str, ...] = () + """Matching key: MoE block class-name substrings (case-insensitive).""" + + expert_linear_names: tuple[str, ...] | None = None + """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``.""" + + has_iterable_experts: bool = False + """True when experts are per-expert iterable sub-modules (Mixtral, Qwen MoE, + NemotronH, Gemma4) and can be grouped by ``get_experts_list``; False for stacked + or fused layouts (DBRX, GptOss).""" diff --git a/modelopt/torch/export/modeling/families/__init__.py b/modelopt/torch/export/modeling/families/__init__.py new file mode 100644 index 00000000000..21fc28cab87 --- /dev/null +++ b/modelopt/torch/export/modeling/families/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model-family export specs. Importing each module registers its specs.""" + +from . import dbrx, deepseek, gemma, gptoss, mixtral, nemotron, qwen diff --git a/modelopt/torch/export/modeling/families/dbrx.py b/modelopt/torch/export/modeling/families/dbrx.py new file mode 100644 index 00000000000..5d53ddb753a --- /dev/null +++ b/modelopt/torch/export/modeling/families/dbrx.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DBRX family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="dbrx", + moe_block_names=("DBRXMoeSparseMoeBlock",), + expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), + ) +) diff --git a/modelopt/torch/export/modeling/families/deepseek.py b/modelopt/torch/export/modeling/families/deepseek.py new file mode 100644 index 00000000000..fa0a0589f60 --- /dev/null +++ b/modelopt/torch/export/modeling/families/deepseek.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DeepSeek family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="deepseek", + moe_block_names=("DeepseekMoE",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + ) +) diff --git a/modelopt/torch/export/modeling/families/gemma.py b/modelopt/torch/export/modeling/families/gemma.py new file mode 100644 index 00000000000..9ffe7754eec --- /dev/null +++ b/modelopt/torch/export/modeling/families/gemma.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gemma family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="gemma4_moe", + # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. + moe_block_names=("Gemma4TextDecoderLayer",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + has_iterable_experts=True, + ) +) diff --git a/modelopt/torch/export/modeling/families/gptoss.py b/modelopt/torch/export/modeling/families/gptoss.py new file mode 100644 index 00000000000..1c1cc342e1f --- /dev/null +++ b/modelopt/torch/export/modeling/families/gptoss.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPT-OSS family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="gptoss", + # GPT-OSS fuses gate and up into a single gate_up_proj. + moe_block_names=("GptOssMoE",), + expert_linear_names=("gate_up_proj", "down_proj"), + ) +) diff --git a/modelopt/torch/export/modeling/families/mixtral.py b/modelopt/torch/export/modeling/families/mixtral.py new file mode 100644 index 00000000000..ade34693a2b --- /dev/null +++ b/modelopt/torch/export/modeling/families/mixtral.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mixtral family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +# Mixtral with iterable experts uses w1/w2/w3. Fused experts (transformers 5.0+) are +# detected from their per-expert quantizer attributes and need no naming override here. +register( + ModelSpec( + name="mixtral", + moe_block_names=("MixtralSparseMoeBlock",), + expert_linear_names=("w1", "w2", "w3"), + has_iterable_experts=True, + ) +) + +# Older transformers naming for Mixtral. +register( + ModelSpec( + name="mixtral_mcore", + moe_block_names=("MixtralMoeSparseMoeBlock",), + expert_linear_names=("linear_fc1", "linear_fc2"), + ) +) diff --git a/modelopt/torch/export/modeling/families/nemotron.py b/modelopt/torch/export/modeling/families/nemotron.py new file mode 100644 index 00000000000..8d3cb6edd5d --- /dev/null +++ b/modelopt/torch/export/modeling/families/nemotron.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Nemotron family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="nemotron_h", + # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). + moe_block_names=("NemotronHMOE",), + expert_linear_names=("up_proj", "down_proj"), + has_iterable_experts=True, + ) +) diff --git a/modelopt/torch/export/modeling/families/qwen.py b/modelopt/torch/export/modeling/families/qwen.py new file mode 100644 index 00000000000..a4e6da84a9a --- /dev/null +++ b/modelopt/torch/export/modeling/families/qwen.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="qwen_moe", + moe_block_names=( + "Qwen2MoeSparseMoeBlock", + "Qwen3MoeSparseMoeBlock", + "Qwen3NextSparseMoeBlock", + "Qwen3_5MoeSparseMoeBlock", + ), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + has_iterable_experts=True, + ) +) diff --git a/modelopt/torch/export/modeling/registry.py b/modelopt/torch/export/modeling/registry.py new file mode 100644 index 00000000000..1a7700fe6c0 --- /dev/null +++ b/modelopt/torch/export/modeling/registry.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Registry that resolves a model sub-module to its ``ModelSpec``. + +Families register their specs at import time (see ``families/``). Lookups return +``None`` when nothing matches, so callers can fall back to their default behavior. +""" + +import torch.nn as nn + +from .base import ModelSpec + +__all__ = [ + "match_moe_block", + "register", +] + +_SPECS: list[ModelSpec] = [] + + +def register(spec: ModelSpec) -> ModelSpec: + """Register a model-family spec and return it.""" + _SPECS.append(spec) + return spec + + +def match_moe_block(module: nn.Module) -> ModelSpec | None: + """Return the spec matching ``module``'s class name against ``moe_block_names``. + + Case-insensitive substring match against ``type(module).__name__``. + """ + cls_name = type(module).__name__.lower() + for spec in _SPECS: + if any(name.lower() in cls_name for name in spec.moe_block_names): + return spec + return None diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py new file mode 100644 index 00000000000..6f985d2853d --- /dev/null +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the model-family export spec registry (modelopt.torch.export.modeling).""" + +import pytest +import torch.nn as nn + +from modelopt.torch.export.layer_utils import get_expert_linear_names, get_experts_list +from modelopt.torch.export.modeling import match_moe_block + + +class _FakeQwen3MoeSparseMoeBlock(nn.Module): + pass + + +class _QuantMixtralSparseMoeBlock(nn.Module): + """Dynamically generated quantized classes keep the family name as a substring.""" + + +class _UnknownMoeBlock(nn.Module): + pass + + +def test_match_moe_block_by_substring(): + spec = match_moe_block(_FakeQwen3MoeSparseMoeBlock()) + assert spec is not None + assert spec.name == "qwen_moe" + assert spec.expert_linear_names == ("gate_proj", "down_proj", "up_proj") + assert spec.has_iterable_experts + + +def test_match_moe_block_matches_quantized_class_name(): + spec = match_moe_block(_QuantMixtralSparseMoeBlock()) + assert spec is not None + assert spec.name == "mixtral" + + +def test_match_moe_block_unmatched_returns_none(): + assert match_moe_block(_UnknownMoeBlock()) is None + + +def test_get_expert_linear_names_falls_back_when_unmatched(): + # No family spec matches — the legacy default (w1/w2/w3) must be preserved. + assert get_expert_linear_names(_UnknownMoeBlock()) == ["w1", "w2", "w3"] + + +class _FakeNemotronHMOE(nn.Module): + def __init__(self): + super().__init__() + + class _Expert(nn.Module): + def __init__(self): + super().__init__() + self.up_proj = nn.Linear(8, 16) + self.down_proj = nn.Linear(16, 8) + + self.experts = nn.ModuleList([_Expert() for _ in range(3)]) + + +def test_get_experts_list_groups_by_spec_linear_names(): + module = _FakeNemotronHMOE() + groups = get_experts_list(module, "nemotronhforcausallm") + assert len(groups) == 2 # up_proj group + down_proj group + assert all(len(group) == 3 for group in groups) + assert groups[0][0] is module.experts[0].up_proj + assert groups[1][2] is module.experts[2].down_proj + + +def test_get_experts_list_rejects_non_iterable_families(): + # DBRX matches a spec but is not an iterable-experts layout; grouped export + # must keep rejecting it (legacy behavior). + class _FakeDBRXMoeSparseMoeBlock(nn.Module): + def __init__(self): + super().__init__() + self.experts = nn.ModuleList() + + with pytest.raises(NotImplementedError): + get_experts_list(_FakeDBRXMoeSparseMoeBlock(), "dbrxforcausallm") + + with pytest.raises(NotImplementedError): + get_experts_list(_UnknownMoeBlock(), "unknownforcausallm") From c7f474ba1b96d65804c95bfc5cd73c7c54e61ac0 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 14 Jul 2026 23:45:53 +0000 Subject: [PATCH 02/16] move more model logic Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 14 ++++---- modelopt/torch/export/layer_utils.py | 4 +-- modelopt/torch/export/modeling/base.py | 6 ++++ .../export/modeling/families/__init__.py | 2 +- .../torch/export/modeling/families/arctic.py | 28 +++++++++++++++ .../torch/export/modeling/families/dbrx.py | 9 +++++ .../torch/export/modeling/families/llama.py | 30 ++++++++++++++++ .../torch/export/modeling/families/qwen.py | 12 +++++++ modelopt/torch/export/modeling/registry.py | 11 ++++++ modelopt/torch/export/quant_utils.py | 30 +++++++--------- .../unit/torch/export/test_modeling_specs.py | 35 +++++++++++++++++-- 11 files changed, 151 insertions(+), 30 deletions(-) create mode 100644 modelopt/torch/export/modeling/families/arctic.py create mode 100644 modelopt/torch/export/modeling/families/llama.py diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index 800f0da743a..4bcb3065c4e 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -26,10 +26,10 @@ helpers) consume. Inventory of model-specific logic reachable from the HF path: | Item | Location | Kind | |---|---|---| -| MoE expert linear names (Qwen/DeepSeek/Mixtral/DBRX/GptOss/NemotronH/Gemma4) | `layer_utils.get_expert_linear_names` | data (if/elif chain) | -| Duplicate expert-naming table + iterable-experts support gate | `layer_utils.get_experts_list` | data (if/elif chain) | -| MoE block class-name list | `layer_utils.is_moe` | data (shared with TRT-LLM path) | -| AWQ `pre_quant_scale` fusion rules (Llama/Qwen3) | `quant_utils.PQS_FUSE_MODULE_MAPPING` | data (table) | +| MoE expert linear names (Qwen/DeepSeek/Mixtral/DBRX/GptOss/NemotronH/Gemma4) | `layer_utils.get_expert_linear_names` | data — **migrated (P1)** | +| Duplicate expert-naming table + iterable-experts support gate | `layer_utils.get_experts_list` | data — **migrated (P1)** | +| MoE block class-name list | `layer_utils.is_moe` | data — **migrated (P3)** | +| AWQ `pre_quant_scale` fusion rules (Llama/Qwen3) | `quant_utils.PQS_FUSE_MODULE_MAPPING` | data — **migrated (P2)** | | weight+1 layernorm class names (Gemma RMSNorm, LayerNorm1P) | `quant_utils._layernorm_uses_weight_plus_one` | data | | Handler match keys (Llama4TextExperts, GptOssExperts, DbrxExperts, QuantMoELinear) | `hf_export_handlers.py` | dispatch (stays: structural, per-module) | | Fused-expert gated/non-gated split (`gate_up_proj` vs `up_proj`) | `moe_utils.py` | data + structure | @@ -77,9 +77,9 @@ against existing export tests. | Step | What | Status | |---|---|---| | **P1** | Registry skeleton + MoE expert naming: `get_expert_linear_names` and `get_experts_list` read `spec.expert_linear_names` / `spec.has_iterable_experts`. The #1 "add a MoE model" shotgun-surgery driver. | this PR | -| **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated across specs. | planned | -| **P3** | `is_moe` class-name list → `spec.moe_block_names` (shared consumer: keep fallback until the TRT-LLM path moves out). | planned | -| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs. | planned | +| **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen families). | this PR | +| **P3** | `is_moe` explicit class-name list → `spec.moe_block_names` (arctic/dbrx_ffn are identification-only specs: no expert naming, so expert-name lookups keep the engine default). The generic `*SparseMoeBlock`/`*MoeLayer` conventions and the structural router+experts check stay in the engine. | this PR | +| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the top-level dispatch registry (#1939). | planned | | **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Dead code found during inventory (`MODEL_NAME_TO_TYPE`, `get_model_type`, `adjust_attn_amax_values`, `update_experts_avg_prequant_scale`) is deleted on that track too. | separate track | **Guardrails:** fallback-first; one data category per PR; the engine keeps the diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index 29a6b589104..f84d797e0e1 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -306,8 +306,8 @@ def is_moe(module: nn.Module) -> bool: # Auto-detect common MoE patterns if name.endswith("sparsemoeblock") or "moelayer" in name: return True - # Explicit matches for non-standard naming - if any(key in name for key in ["arcticmoe", "deepseekmoe", "dbrxffn", "nemotronhmoe"]): + # Non-standard MoE block names are per-family data (modeling/families/*). + if match_moe_block(module) is not None: return True # Structural detection: modules with router + experts (e.g. Gemma4TextDecoderLayer) return ( diff --git a/modelopt/torch/export/modeling/base.py b/modelopt/torch/export/modeling/base.py index 72c64c0ebcd..5ed96584ee2 100644 --- a/modelopt/torch/export/modeling/base.py +++ b/modelopt/torch/export/modeling/base.py @@ -48,3 +48,9 @@ class ModelSpec: """True when experts are per-expert iterable sub-modules (Mixtral, Qwen MoE, NemotronH, Gemma4) and can be grouped by ``get_experts_list``; False for stacked or fused layouts (DBRX, GptOss).""" + + pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = () + """AWQ ``pre_quant_scale`` fusion rules, each a ``(module_class_substrings, + fuse_into, fuse_from)`` triple: for a module whose class name contains one of the + substrings, the pre_quant_scale on ``fuse_from`` is folded into ``fuse_into`` + (e.g. attention ``o_proj`` -> ``v_proj``, MLP ``down_proj`` -> ``up_proj``).""" diff --git a/modelopt/torch/export/modeling/families/__init__.py b/modelopt/torch/export/modeling/families/__init__.py index 21fc28cab87..38804baaac9 100644 --- a/modelopt/torch/export/modeling/families/__init__.py +++ b/modelopt/torch/export/modeling/families/__init__.py @@ -15,4 +15,4 @@ """Per-model-family export specs. Importing each module registers its specs.""" -from . import dbrx, deepseek, gemma, gptoss, mixtral, nemotron, qwen +from . import arctic, dbrx, deepseek, gemma, gptoss, llama, mixtral, nemotron, qwen diff --git a/modelopt/torch/export/modeling/families/arctic.py b/modelopt/torch/export/modeling/families/arctic.py new file mode 100644 index 00000000000..c53ca9b35bd --- /dev/null +++ b/modelopt/torch/export/modeling/families/arctic.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Arctic family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +# MoE-block identification only (non-standard block name for is_moe); expert naming +# intentionally unset so expert-name lookups keep the engine default. +register( + ModelSpec( + name="arctic", + moe_block_names=("ArcticMoE",), + ) +) diff --git a/modelopt/torch/export/modeling/families/dbrx.py b/modelopt/torch/export/modeling/families/dbrx.py index 5d53ddb753a..4f8a043beb2 100644 --- a/modelopt/torch/export/modeling/families/dbrx.py +++ b/modelopt/torch/export/modeling/families/dbrx.py @@ -25,3 +25,12 @@ expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), ) ) + +# HF DbrxFFN: MoE-block identification only (non-standard block name for is_moe); +# expert naming intentionally unset so expert-name lookups keep the engine default. +register( + ModelSpec( + name="dbrx_ffn", + moe_block_names=("DbrxFFN",), + ) +) diff --git a/modelopt/torch/export/modeling/families/llama.py b/modelopt/torch/export/modeling/families/llama.py new file mode 100644 index 00000000000..3f83e7559f9 --- /dev/null +++ b/modelopt/torch/export/modeling/families/llama.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Llama family export specs.""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="llama", + # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. + pqs_fuse_rules=( + (("LlamaAttention",), "v_proj", "o_proj"), + (("LlamaMLP",), "up_proj", "down_proj"), + ), + ) +) diff --git a/modelopt/torch/export/modeling/families/qwen.py b/modelopt/torch/export/modeling/families/qwen.py index a4e6da84a9a..9f866e7063f 100644 --- a/modelopt/torch/export/modeling/families/qwen.py +++ b/modelopt/torch/export/modeling/families/qwen.py @@ -31,3 +31,15 @@ has_iterable_experts=True, ) ) + +# Qwen3 (dense + MoE) AWQ pre_quant_scale fusion: fold o_proj into v_proj, +# down_proj into up_proj. +register( + ModelSpec( + name="qwen3", + pqs_fuse_rules=( + (("Qwen3Attention", "Qwen3MoeAttention"), "v_proj", "o_proj"), + (("Qwen3MLP", "Qwen3MoeMLP"), "up_proj", "down_proj"), + ), + ) +) diff --git a/modelopt/torch/export/modeling/registry.py b/modelopt/torch/export/modeling/registry.py index 1a7700fe6c0..6a98323d710 100644 --- a/modelopt/torch/export/modeling/registry.py +++ b/modelopt/torch/export/modeling/registry.py @@ -24,6 +24,7 @@ from .base import ModelSpec __all__ = [ + "iter_pqs_fuse_rules", "match_moe_block", "register", ] @@ -37,6 +38,16 @@ def register(spec: ModelSpec) -> ModelSpec: return spec +def iter_pqs_fuse_rules(): + """Yield every ``(module_class_substrings, fuse_into, fuse_from)`` AWQ fusion rule. + + Aggregated across all registered specs (the consumer matches each model module + against the substrings, so the order across families does not matter). + """ + for spec in _SPECS: + yield from spec.pqs_fuse_rules + + def match_moe_block(module: nn.Module) -> ModelSpec | None: """Return the spec matching ``module``'s class name against ``moe_block_names``. diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..04ddaa2c6c9 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -70,6 +70,7 @@ QUANTIZATION_W4A8_NVFP4_FP8, QUANTIZATION_W4A16_NVFP4, ) +from .modeling import iter_pqs_fuse_rules logger = logging.getLogger(__name__) @@ -1143,19 +1144,14 @@ def _update_svdquant(modules, new_pre_quant_scale): finish_stats_collection(module.weight_quantizer) -# Format: (list of target modules, tuple of (linear_to_fuse_into, linear_from_with_scale)) -PQS_FUSE_MODULE_MAPPING = [ - # Attention: Fuse o_proj's pre_quant_scale into v_proj's output dimension - # Mathematical equivalence: - # Before: o_proj_out = [attn @ (v_proj_in @ v_proj.W^T)^T * scale] @ o_proj.W^T - # After: o_proj_out = [attn @ (v_proj_in @ (v_proj.W * scale)^T)^T] @ o_proj.W^T - (["LlamaAttention", "Qwen3Attention", "Qwen3MoeAttention"], ("v_proj", "o_proj")), - # MLP: Fuse down_proj's pre_quant_scale into up_proj's output dimension - # Mathematical equivalence: - # Before: down_proj_out = {[act_fn(self.gate_proj(x)) * up_proj(x)] * scale} @ down_proj.W^T - # After: down_proj_out = {[act_fn(self.gate_proj(x)) * (up_proj(x) * scale)]} @ down_proj.W^T - (["LlamaMLP", "Qwen3MLP", "Qwen3MoeMLP"], ("up_proj", "down_proj")), -] +# AWQ pre_quant_scale fusion rules are per-model data and live in modeling/families/*: +# - Attention: fold o_proj's pre_quant_scale into v_proj's output dimension. +# Before: o_proj_out = [attn @ (v_proj_in @ v_proj.W^T)^T * scale] @ o_proj.W^T +# After: o_proj_out = [attn @ (v_proj_in @ (v_proj.W * scale)^T)^T] @ o_proj.W^T +# - MLP: fold down_proj's pre_quant_scale into up_proj's output dimension. +# Before: down_proj_out = {[act_fn(gate_proj(x)) * up_proj(x)] * scale} @ down_proj.W^T +# After: down_proj_out = {[act_fn(gate_proj(x)) * (up_proj(x) * scale)]} @ down_proj.W^T +# Each rule is (module_class_substrings, fuse_into, fuse_from); see ``iter_pqs_fuse_rules``. def fuse_prequant_to_linear(model: torch.nn.Module, fuse_grouped_heads=False): @@ -1171,12 +1167,10 @@ def fuse_prequant_to_linear(model: torch.nn.Module, fuse_grouped_heads=False): """ # Fuse pre_quant_scale to the linear weights for _, module in model.named_modules(): - for module_map in PQS_FUSE_MODULE_MAPPING: - target_module_list = module_map[0] - linear_pair = module_map[1] + for target_module_list, fuse_into, fuse_from in iter_pqs_fuse_rules(): if any(module_name in type(module).__name__ for module_name in target_module_list): - linear_fuse_into = module.get_submodule(linear_pair[0]) - linear_pqs_from = module.get_submodule(linear_pair[1]) + linear_fuse_into = module.get_submodule(fuse_into) + linear_pqs_from = module.get_submodule(fuse_from) if hasattr(linear_pqs_from, "input_quantizer") and hasattr( linear_pqs_from.input_quantizer, "_pre_quant_scale" ): diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index 6f985d2853d..a5f55ceb6b7 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -18,8 +18,8 @@ import pytest import torch.nn as nn -from modelopt.torch.export.layer_utils import get_expert_linear_names, get_experts_list -from modelopt.torch.export.modeling import match_moe_block +from modelopt.torch.export.layer_utils import get_expert_linear_names, get_experts_list, is_moe +from modelopt.torch.export.modeling import iter_pqs_fuse_rules, match_moe_block class _FakeQwen3MoeSparseMoeBlock(nn.Module): @@ -92,3 +92,34 @@ def __init__(self): with pytest.raises(NotImplementedError): get_experts_list(_UnknownMoeBlock(), "unknownforcausallm") + + +class _FakeArcticMoE(nn.Module): + pass + + +class _FakeDbrxFFN(nn.Module): + pass + + +def test_is_moe_matches_registered_non_standard_names(): + # Non-standard MoE block names (no *SparseMoeBlock suffix, no router/experts + # attributes) resolve through the family registry. + assert is_moe(_FakeArcticMoE()) + assert is_moe(_FakeDbrxFFN()) + assert not is_moe(_UnknownMoeBlock()) + + +def test_pqs_fuse_rules_match_legacy_mapping(): + # Aggregated per-family rules must reproduce the legacy PQS_FUSE_MODULE_MAPPING. + rules = { + (frozenset(substrings), fuse_into, fuse_from) + for substrings, fuse_into, fuse_from in iter_pqs_fuse_rules() + } + legacy = { + (frozenset({"LlamaAttention"}), "v_proj", "o_proj"), + (frozenset({"LlamaMLP"}), "up_proj", "down_proj"), + (frozenset({"Qwen3Attention", "Qwen3MoeAttention"}), "v_proj", "o_proj"), + (frozenset({"Qwen3MLP", "Qwen3MoeMLP"}), "up_proj", "down_proj"), + } + assert rules == legacy From 8d81535a09b2232be2406b9f8474858d85cccf51 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 15 Jul 2026 23:35:07 +0000 Subject: [PATCH 03/16] top level dir Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- docs/source/reference/1_modelopt_api.rst | 1 + .../{torch/export => }/modeling/__init__.py | 13 ++--- modelopt/{torch/export => }/modeling/base.py | 12 ++--- modelopt/modeling/models/__init__.py | 37 ++++++++++++++ .../families => modeling/models}/arctic.py | 2 +- .../families => modeling/models}/dbrx.py | 2 +- .../families => modeling/models}/deepseek.py | 6 ++- .../gemma.py => modeling/models/gemma4.py} | 4 +- .../gptoss.py => modeling/models/gpt_oss.py} | 4 +- .../families => modeling/models}/llama.py | 2 +- .../families => modeling/models}/mixtral.py | 2 +- .../models/nemotron_h.py} | 2 +- .../models/qwen2_moe.py} | 14 +++++- modelopt/modeling/models/qwen3.py | 30 ++++++++++++ modelopt/modeling/models/qwen3_5_moe.py | 28 +++++++++++ .../qwen.py => modeling/models/qwen3_moe.py} | 24 +++------ modelopt/modeling/models/qwen3_next.py | 28 +++++++++++ .../{torch/export => }/modeling/registry.py | 17 +++++-- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 49 +++++++++++-------- modelopt/torch/export/layer_utils.py | 10 ++-- modelopt/torch/export/quant_utils.py | 4 +- .../unit/torch/export/test_modeling_specs.py | 28 ++++++----- 22 files changed, 232 insertions(+), 87 deletions(-) rename modelopt/{torch/export => }/modeling/__init__.py (61%) rename modelopt/{torch/export => }/modeling/base.py (85%) create mode 100644 modelopt/modeling/models/__init__.py rename modelopt/{torch/export/modeling/families => modeling/models}/arctic.py (93%) rename modelopt/{torch/export/modeling/families => modeling/models}/dbrx.py (96%) rename modelopt/{torch/export/modeling/families => modeling/models}/deepseek.py (84%) rename modelopt/{torch/export/modeling/families/gemma.py => modeling/models/gemma4.py} (93%) rename modelopt/{torch/export/modeling/families/gptoss.py => modeling/models/gpt_oss.py} (92%) rename modelopt/{torch/export/modeling/families => modeling/models}/llama.py (95%) rename modelopt/{torch/export/modeling/families => modeling/models}/mixtral.py (96%) rename modelopt/{torch/export/modeling/families/nemotron.py => modeling/models/nemotron_h.py} (94%) rename modelopt/{torch/export/modeling/families/__init__.py => modeling/models/qwen2_moe.py} (67%) create mode 100644 modelopt/modeling/models/qwen3.py create mode 100644 modelopt/modeling/models/qwen3_5_moe.py rename modelopt/{torch/export/modeling/families/qwen.py => modeling/models/qwen3_moe.py} (62%) create mode 100644 modelopt/modeling/models/qwen3_next.py rename modelopt/{torch/export => }/modeling/registry.py (76%) diff --git a/docs/source/reference/1_modelopt_api.rst b/docs/source/reference/1_modelopt_api.rst index caf5697737b..8c1f5ab21c1 100644 --- a/docs/source/reference/1_modelopt_api.rst +++ b/docs/source/reference/1_modelopt_api.rst @@ -10,5 +10,6 @@ modelopt API :recursive: modelopt.deploy + modelopt.modeling modelopt.onnx modelopt.torch diff --git a/modelopt/torch/export/modeling/__init__.py b/modelopt/modeling/__init__.py similarity index 61% rename from modelopt/torch/export/modeling/__init__.py rename to modelopt/modeling/__init__.py index 5718219e50e..6ce8d7ec3b7 100644 --- a/modelopt/torch/export/modeling/__init__.py +++ b/modelopt/modeling/__init__.py @@ -13,14 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Model-family export descriptors. +"""Per-model descriptors, organized by HF model type. -Holds per-model export data organized by model family. The export code resolves a -``ModelSpec`` via the registry lookups and reads its fields; an unmatched lookup -returns ``None`` so callers fall back to their default behavior. +Holds declarative per-model data (no algorithms), one module per HF model type under +``models/``, mirroring ``transformers.models``. Consumers (e.g. the unified HF export +path) resolve a ``ModelSpec`` via the registry lookups and read its fields; an +unmatched lookup returns ``None`` so callers fall back to their default behavior. """ -# Importing families registers every family spec as a side effect. -from . import families +# Importing models registers every spec as a side effect. +from . import models from .base import * from .registry import * diff --git a/modelopt/torch/export/modeling/base.py b/modelopt/modeling/base.py similarity index 85% rename from modelopt/torch/export/modeling/base.py rename to modelopt/modeling/base.py index 5ed96584ee2..fe8618418ed 100644 --- a/modelopt/torch/export/modeling/base.py +++ b/modelopt/modeling/base.py @@ -13,11 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-model-family export descriptor. +"""Per-model descriptor. -A ``ModelSpec`` declares how one model family differs from the generic export path, -so the export code can read these values instead of branching on model names. Each -spec holds per-model data only, not export logic. +A ``ModelSpec`` declares how one model type differs from the generic code path, +so consumers (e.g. export) can read these values instead of branching on model +names. Each spec holds per-model data only, no logic. """ from dataclasses import dataclass @@ -27,7 +27,7 @@ @dataclass class ModelSpec: - """Per-model-family export data. + """Per-model data. A spec is resolved from a model sub-module via its matching keys and read for its per-model data fields. ``moe_block_names`` is the matching key: MoE block @@ -36,7 +36,7 @@ class ModelSpec: """ name: str - """Unique name of the model-family spec.""" + """Unique spec name; the HF model type where one exists (e.g. ``"qwen3_moe"``).""" moe_block_names: tuple[str, ...] = () """Matching key: MoE block class-name substrings (case-insensitive).""" diff --git a/modelopt/modeling/models/__init__.py b/modelopt/modeling/models/__init__.py new file mode 100644 index 00000000000..cc711b81ae9 --- /dev/null +++ b/modelopt/modeling/models/__init__.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model specs, one module per HF model type. Importing each module registers its specs. + +Module names follow the HF ``transformers.models`` layout +(https://github.com/huggingface/transformers/tree/main/src/transformers/models); +trust-remote-code models (``arctic``, ``deepseek``) use their config ``model_type``. +""" + +from . import ( + arctic, + dbrx, + deepseek, + gemma4, + gpt_oss, + llama, + mixtral, + nemotron_h, + qwen2_moe, + qwen3, + qwen3_5_moe, + qwen3_moe, + qwen3_next, +) diff --git a/modelopt/torch/export/modeling/families/arctic.py b/modelopt/modeling/models/arctic.py similarity index 93% rename from modelopt/torch/export/modeling/families/arctic.py rename to modelopt/modeling/models/arctic.py index c53ca9b35bd..1cd769ea1fe 100644 --- a/modelopt/torch/export/modeling/families/arctic.py +++ b/modelopt/modeling/models/arctic.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Arctic family export specs.""" +"""Snowflake Arctic specs (trust-remote-code model type ``arctic``).""" from ..base import ModelSpec from ..registry import register diff --git a/modelopt/torch/export/modeling/families/dbrx.py b/modelopt/modeling/models/dbrx.py similarity index 96% rename from modelopt/torch/export/modeling/families/dbrx.py rename to modelopt/modeling/models/dbrx.py index 4f8a043beb2..2f006561b54 100644 --- a/modelopt/torch/export/modeling/families/dbrx.py +++ b/modelopt/modeling/models/dbrx.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""DBRX family export specs.""" +"""DBRX specs (HF model type ``dbrx``).""" from ..base import ModelSpec from ..registry import register diff --git a/modelopt/torch/export/modeling/families/deepseek.py b/modelopt/modeling/models/deepseek.py similarity index 84% rename from modelopt/torch/export/modeling/families/deepseek.py rename to modelopt/modeling/models/deepseek.py index fa0a0589f60..ba12f778754 100644 --- a/modelopt/torch/export/modeling/families/deepseek.py +++ b/modelopt/modeling/models/deepseek.py @@ -13,7 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""DeepSeek family export specs.""" +"""DeepSeek-MoE specs (trust-remote-code model type ``deepseek``). + +Matches the remote-code ``DeepseekMoE`` block, not the HF-native ``deepseek_v3`` +classes. +""" from ..base import ModelSpec from ..registry import register diff --git a/modelopt/torch/export/modeling/families/gemma.py b/modelopt/modeling/models/gemma4.py similarity index 93% rename from modelopt/torch/export/modeling/families/gemma.py rename to modelopt/modeling/models/gemma4.py index 9ffe7754eec..9b6284145ea 100644 --- a/modelopt/torch/export/modeling/families/gemma.py +++ b/modelopt/modeling/models/gemma4.py @@ -13,14 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Gemma family export specs.""" +"""Gemma4 specs (HF model type ``gemma4``).""" from ..base import ModelSpec from ..registry import register register( ModelSpec( - name="gemma4_moe", + name="gemma4", # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. moe_block_names=("Gemma4TextDecoderLayer",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/torch/export/modeling/families/gptoss.py b/modelopt/modeling/models/gpt_oss.py similarity index 92% rename from modelopt/torch/export/modeling/families/gptoss.py rename to modelopt/modeling/models/gpt_oss.py index 1c1cc342e1f..8ded20dd3ae 100644 --- a/modelopt/torch/export/modeling/families/gptoss.py +++ b/modelopt/modeling/models/gpt_oss.py @@ -13,14 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GPT-OSS family export specs.""" +"""GPT-OSS specs (HF model type ``gpt_oss``).""" from ..base import ModelSpec from ..registry import register register( ModelSpec( - name="gptoss", + name="gpt_oss", # GPT-OSS fuses gate and up into a single gate_up_proj. moe_block_names=("GptOssMoE",), expert_linear_names=("gate_up_proj", "down_proj"), diff --git a/modelopt/torch/export/modeling/families/llama.py b/modelopt/modeling/models/llama.py similarity index 95% rename from modelopt/torch/export/modeling/families/llama.py rename to modelopt/modeling/models/llama.py index 3f83e7559f9..fe557e8becf 100644 --- a/modelopt/torch/export/modeling/families/llama.py +++ b/modelopt/modeling/models/llama.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Llama family export specs.""" +"""Llama specs (HF model type ``llama``).""" from ..base import ModelSpec from ..registry import register diff --git a/modelopt/torch/export/modeling/families/mixtral.py b/modelopt/modeling/models/mixtral.py similarity index 96% rename from modelopt/torch/export/modeling/families/mixtral.py rename to modelopt/modeling/models/mixtral.py index ade34693a2b..b3b8c76911b 100644 --- a/modelopt/torch/export/modeling/families/mixtral.py +++ b/modelopt/modeling/models/mixtral.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Mixtral family export specs.""" +"""Mixtral specs (HF model type ``mixtral``).""" from ..base import ModelSpec from ..registry import register diff --git a/modelopt/torch/export/modeling/families/nemotron.py b/modelopt/modeling/models/nemotron_h.py similarity index 94% rename from modelopt/torch/export/modeling/families/nemotron.py rename to modelopt/modeling/models/nemotron_h.py index 8d3cb6edd5d..638bf248097 100644 --- a/modelopt/torch/export/modeling/families/nemotron.py +++ b/modelopt/modeling/models/nemotron_h.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Nemotron family export specs.""" +"""Nemotron-H specs (HF model type ``nemotron_h``).""" from ..base import ModelSpec from ..registry import register diff --git a/modelopt/torch/export/modeling/families/__init__.py b/modelopt/modeling/models/qwen2_moe.py similarity index 67% rename from modelopt/torch/export/modeling/families/__init__.py rename to modelopt/modeling/models/qwen2_moe.py index 38804baaac9..65f5fe755d9 100644 --- a/modelopt/torch/export/modeling/families/__init__.py +++ b/modelopt/modeling/models/qwen2_moe.py @@ -13,6 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-model-family export specs. Importing each module registers its specs.""" +"""Qwen2-MoE specs (HF model type ``qwen2_moe``).""" -from . import arctic, dbrx, deepseek, gemma, gptoss, llama, mixtral, nemotron, qwen +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="qwen2_moe", + moe_block_names=("Qwen2MoeSparseMoeBlock",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + has_iterable_experts=True, + ) +) diff --git a/modelopt/modeling/models/qwen3.py b/modelopt/modeling/models/qwen3.py new file mode 100644 index 00000000000..78f2388c40a --- /dev/null +++ b/modelopt/modeling/models/qwen3.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3 (dense) specs (HF model type ``qwen3``).""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="qwen3", + # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. + pqs_fuse_rules=( + (("Qwen3Attention",), "v_proj", "o_proj"), + (("Qwen3MLP",), "up_proj", "down_proj"), + ), + ) +) diff --git a/modelopt/modeling/models/qwen3_5_moe.py b/modelopt/modeling/models/qwen3_5_moe.py new file mode 100644 index 00000000000..178a3677518 --- /dev/null +++ b/modelopt/modeling/models/qwen3_5_moe.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3.5-MoE specs (HF model type ``qwen3_5_moe``).""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="qwen3_5_moe", + moe_block_names=("Qwen3_5MoeSparseMoeBlock",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + has_iterable_experts=True, + ) +) diff --git a/modelopt/torch/export/modeling/families/qwen.py b/modelopt/modeling/models/qwen3_moe.py similarity index 62% rename from modelopt/torch/export/modeling/families/qwen.py rename to modelopt/modeling/models/qwen3_moe.py index 9f866e7063f..bbcf73d60e0 100644 --- a/modelopt/torch/export/modeling/families/qwen.py +++ b/modelopt/modeling/models/qwen3_moe.py @@ -13,33 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Qwen family export specs.""" +"""Qwen3-MoE specs (HF model type ``qwen3_moe``).""" from ..base import ModelSpec from ..registry import register register( ModelSpec( - name="qwen_moe", - moe_block_names=( - "Qwen2MoeSparseMoeBlock", - "Qwen3MoeSparseMoeBlock", - "Qwen3NextSparseMoeBlock", - "Qwen3_5MoeSparseMoeBlock", - ), + name="qwen3_moe", + moe_block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, - ) -) - -# Qwen3 (dense + MoE) AWQ pre_quant_scale fusion: fold o_proj into v_proj, -# down_proj into up_proj. -register( - ModelSpec( - name="qwen3", + # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( - (("Qwen3Attention", "Qwen3MoeAttention"), "v_proj", "o_proj"), - (("Qwen3MLP", "Qwen3MoeMLP"), "up_proj", "down_proj"), + (("Qwen3MoeAttention",), "v_proj", "o_proj"), + (("Qwen3MoeMLP",), "up_proj", "down_proj"), ), ) ) diff --git a/modelopt/modeling/models/qwen3_next.py b/modelopt/modeling/models/qwen3_next.py new file mode 100644 index 00000000000..f0aa8705a1c --- /dev/null +++ b/modelopt/modeling/models/qwen3_next.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3-Next specs (HF model type ``qwen3_next``).""" + +from ..base import ModelSpec +from ..registry import register + +register( + ModelSpec( + name="qwen3_next", + moe_block_names=("Qwen3NextSparseMoeBlock",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + has_iterable_experts=True, + ) +) diff --git a/modelopt/torch/export/modeling/registry.py b/modelopt/modeling/registry.py similarity index 76% rename from modelopt/torch/export/modeling/registry.py rename to modelopt/modeling/registry.py index 6a98323d710..77c32a5cadc 100644 --- a/modelopt/torch/export/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -15,14 +15,21 @@ """Registry that resolves a model sub-module to its ``ModelSpec``. -Families register their specs at import time (see ``families/``). Lookups return +Model modules register their specs at import time (see ``models/``). Lookups return ``None`` when nothing matches, so callers can fall back to their default behavior. + +Matching is by class-name string only, so this package stays dependency-free (any +``nn.Module`` — or any object — can be passed to the lookups without importing torch +here). """ -import torch.nn as nn +from typing import TYPE_CHECKING from .base import ModelSpec +if TYPE_CHECKING: + import torch.nn as nn + __all__ = [ "iter_pqs_fuse_rules", "match_moe_block", @@ -33,7 +40,7 @@ def register(spec: ModelSpec) -> ModelSpec: - """Register a model-family spec and return it.""" + """Register a model spec and return it.""" _SPECS.append(spec) return spec @@ -42,13 +49,13 @@ def iter_pqs_fuse_rules(): """Yield every ``(module_class_substrings, fuse_into, fuse_from)`` AWQ fusion rule. Aggregated across all registered specs (the consumer matches each model module - against the substrings, so the order across families does not matter). + against the substrings, so the order across specs does not matter). """ for spec in _SPECS: yield from spec.pqs_fuse_rules -def match_moe_block(module: nn.Module) -> ModelSpec | None: +def match_moe_block(module: "nn.Module") -> ModelSpec | None: """Return the spec matching ``module``'s class name against ``moe_block_names``. Case-insensitive substring match against ``type(module).__name__``. diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index 4bcb3065c4e..c4d0c9d780b 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -1,13 +1,13 @@ # Export: Model-Specific Logic Refactor — Action Plan -**Goal:** Give the unified HF export path a per-model-family data registry -(`modelopt/torch/export/modeling/`), so that supporting a new model family means -adding one declarative spec file instead of editing if/elif chains across the -export engine. +**Goal:** Give the unified HF export path a per-model data registry +(top-level `modelopt/modeling/`, organized by HF model type), so that supporting a +new model means adding one declarative spec file instead of editing if/elif chains +across the export engine. **Scope:** The unified HF export path (`unified_export_hf.py` + its helpers) only. -- The **Megatron** path (`plugins/mcore_*`) already follows a per-family registry +- The **Megatron** path (`plugins/mcore_*`) already follows a per-model registry pattern and is untouched. - The **TRT-LLM** checkpoint path (`model_config_export.py`, the `build_*` functions in `layer_utils.py`, `tensorrt_llm_utils.py`) is legacy: the plan is @@ -21,7 +21,7 @@ and `PrepareMoEInputsRegistry` (`registry.py`, `hf_export_handlers.py`) select *which handler processes a module* by class/predicate match, replacing the if/elif chains that used to live in `unified_export_hf.py`. -What is still hardcoded is the **per-family data** those handlers (and other HF-path +What is still hardcoded is the **per-model data** those handlers (and other HF-path helpers) consume. Inventory of model-specific logic reachable from the HF path: | Item | Location | Kind | @@ -41,21 +41,28 @@ helpers) consume. Inventory of model-specific logic reachable from the HF path: Two layers: - **Engine** (existing files, organized by operation): owns all algorithms and the - module walk; consults the registry for per-family values. -- **Modeling library** (`modeling/`, organized by model family): declarative - per-family **data only**. No export logic, no imports from the parent package - (only `torch.nn` + stdlib), so it sits at the bottom of the dependency graph. + module walk; consults the registry for per-model values. +- **Modeling library** (top-level `modelopt/modeling/`, organized by HF model type): + declarative per-model **data only**. No export logic, stdlib-only imports (not even + torch), so it sits at the bottom of the dependency graph and any modelopt subsystem + can depend on it. ```text -modeling/ - base.py # ModelSpec dataclass — the per-family contract +modelopt/modeling/ + base.py # ModelSpec dataclass — the per-model contract registry.py # register() + lookups; returns None when unmatched - __init__.py # re-exports; importing it registers all families - families/ # one small file per family; import == registration + __init__.py # re-exports; importing it registers all specs + models/ # one small file per HF model type (mirrors transformers.models); + # import == registration ``` +Model type names mirror +[`transformers.models`](https://github.com/huggingface/transformers/tree/main/src/transformers/models) +(e.g. `qwen3_moe.py`, `gpt_oss.py`, `nemotron_h.py`); trust-remote-code models +(`arctic`, `deepseek`) use their config `model_type`. + Call sites follow a fallback-first shape, which keeps migration incremental and -behavior-preserving — a family not in the registry behaves exactly as before: +behavior-preserving — a model not in the registry behaves exactly as before: ```python spec = match_moe_block(module) @@ -64,10 +71,10 @@ if spec is not None and spec.expert_linear_names is not None: # ... legacy branch preserved as fallback ... ``` -Note on naming: `modeling/registry.py` (per-family **data**, "what are this -family's values") is distinct from the top-level `registry.py` from PR #1939 +Note on naming: `modelopt/modeling/registry.py` (per-model **data**, "what are +this model's values") is distinct from the export-path `registry.py` from PR #1939 (per-module **dispatch**, "which handler processes this module"). The two layers -compose: handlers look up family data through `modeling/`. +compose: handlers look up model data through `modelopt.modeling`. ## 3. Migration plan @@ -77,10 +84,10 @@ against existing export tests. | Step | What | Status | |---|---|---| | **P1** | Registry skeleton + MoE expert naming: `get_expert_linear_names` and `get_experts_list` read `spec.expert_linear_names` / `spec.has_iterable_experts`. The #1 "add a MoE model" shotgun-surgery driver. | this PR | -| **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen families). | this PR | +| **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen3 specs). | this PR | | **P3** | `is_moe` explicit class-name list → `spec.moe_block_names` (arctic/dbrx_ffn are identification-only specs: no expert naming, so expert-name lookups keep the engine default). The generic `*SparseMoeBlock`/`*MoeLayer` conventions and the structural router+experts check stay in the engine. | this PR | -| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the top-level dispatch registry (#1939). | planned | +| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). | planned | | **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Dead code found during inventory (`MODEL_NAME_TO_TYPE`, `get_model_type`, `adjust_attn_amax_values`, `update_experts_avg_prequant_scale`) is deleted on that track too. | separate track | **Guardrails:** fallback-first; one data category per PR; the engine keeps the -algorithms — families supply values only, never fork functions. +algorithms — model specs supply values only, never fork functions. diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index f84d797e0e1..262235a5d54 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -28,6 +28,7 @@ except Exception: warn("Cannot find transformers package. Hugginface modules cannot be exported.") +from modelopt.modeling import match_moe_block from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import import_plugin @@ -60,7 +61,6 @@ RgLruConfig, ) from .model_config_utils import pad_weights -from .modeling import match_moe_block from .postprocess import view_as_float8_e4m3fn_if_needed, view_as_uint8_if_needed from .quant_utils import ( get_activation_scaling_factor, @@ -94,7 +94,7 @@ def get_experts_list( """ experts_list = [] - # Expert linear names are per-family data (modeling/families/*). The grouped export + # Expert linear names are per-model data (modelopt/modeling/models/*). The grouped export # path only supports families whose experts are iterable per-expert sub-modules # (Mixtral, Qwen MoE, NemotronH, Gemma4); stacked/fused layouts (DBRX, GptOss, ...) # raise NotImplementedError here and are handled by other paths. @@ -306,7 +306,7 @@ def is_moe(module: nn.Module) -> bool: # Auto-detect common MoE patterns if name.endswith("sparsemoeblock") or "moelayer" in name: return True - # Non-standard MoE block names are per-family data (modeling/families/*). + # Non-standard MoE block names are per-model data (modelopt/modeling/models/*). if match_moe_block(module) is not None: return True # Structural detection: modules with router + experts (e.g. Gemma4TextDecoderLayer) @@ -982,8 +982,8 @@ def module_match_name_list(module, name_list): if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"): return [first_proj_attr, "down_proj"] - # Resolve expert names from the model-family registry; fall back to the mapping - # below when no family spec matches the MoE block. + # Resolve expert names from the model spec registry (modelopt/modeling); fall back + # to the mapping below when no spec matches the MoE block. spec = match_moe_block(module) if spec is not None and spec.expert_linear_names is not None: return list(spec.expert_linear_names) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 04ddaa2c6c9..7a7e2dc3dbc 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -25,6 +25,7 @@ import torch.nn as nn from modelopt import __version__ +from modelopt.modeling import iter_pqs_fuse_rules from modelopt.torch.quantization.model_calib import ( enable_stats_collection, finish_stats_collection, @@ -70,7 +71,6 @@ QUANTIZATION_W4A8_NVFP4_FP8, QUANTIZATION_W4A16_NVFP4, ) -from .modeling import iter_pqs_fuse_rules logger = logging.getLogger(__name__) @@ -1144,7 +1144,7 @@ def _update_svdquant(modules, new_pre_quant_scale): finish_stats_collection(module.weight_quantizer) -# AWQ pre_quant_scale fusion rules are per-model data and live in modeling/families/*: +# AWQ pre_quant_scale fusion rules are per-model data and live in modelopt/modeling/models/*: # - Attention: fold o_proj's pre_quant_scale into v_proj's output dimension. # Before: o_proj_out = [attn @ (v_proj_in @ v_proj.W^T)^T * scale] @ o_proj.W^T # After: o_proj_out = [attn @ (v_proj_in @ (v_proj.W * scale)^T)^T] @ o_proj.W^T diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index a5f55ceb6b7..c2d61337b49 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -13,13 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the model-family export spec registry (modelopt.torch.export.modeling).""" +"""Unit tests for the per-model spec registry (modelopt.modeling).""" import pytest import torch.nn as nn +from modelopt.modeling import iter_pqs_fuse_rules, match_moe_block from modelopt.torch.export.layer_utils import get_expert_linear_names, get_experts_list, is_moe -from modelopt.torch.export.modeling import iter_pqs_fuse_rules, match_moe_block class _FakeQwen3MoeSparseMoeBlock(nn.Module): @@ -37,7 +37,7 @@ class _UnknownMoeBlock(nn.Module): def test_match_moe_block_by_substring(): spec = match_moe_block(_FakeQwen3MoeSparseMoeBlock()) assert spec is not None - assert spec.name == "qwen_moe" + assert spec.name == "qwen3_moe" assert spec.expert_linear_names == ("gate_proj", "down_proj", "up_proj") assert spec.has_iterable_experts @@ -53,7 +53,7 @@ def test_match_moe_block_unmatched_returns_none(): def test_get_expert_linear_names_falls_back_when_unmatched(): - # No family spec matches — the legacy default (w1/w2/w3) must be preserved. + # No spec matches — the legacy default (w1/w2/w3) must be preserved. assert get_expert_linear_names(_UnknownMoeBlock()) == ["w1", "w2", "w3"] @@ -79,7 +79,7 @@ def test_get_experts_list_groups_by_spec_linear_names(): assert groups[1][2] is module.experts[2].down_proj -def test_get_experts_list_rejects_non_iterable_families(): +def test_get_experts_list_rejects_non_iterable_layouts(): # DBRX matches a spec but is not an iterable-experts layout; grouped export # must keep rejecting it (legacy behavior). class _FakeDBRXMoeSparseMoeBlock(nn.Module): @@ -104,22 +104,26 @@ class _FakeDbrxFFN(nn.Module): def test_is_moe_matches_registered_non_standard_names(): # Non-standard MoE block names (no *SparseMoeBlock suffix, no router/experts - # attributes) resolve through the family registry. + # attributes) resolve through the model spec registry. assert is_moe(_FakeArcticMoE()) assert is_moe(_FakeDbrxFFN()) assert not is_moe(_UnknownMoeBlock()) def test_pqs_fuse_rules_match_legacy_mapping(): - # Aggregated per-family rules must reproduce the legacy PQS_FUSE_MODULE_MAPPING. + # Aggregated per-model rules, flattened to (class_substring, fuse_into, fuse_from) + # triples, must reproduce the legacy PQS_FUSE_MODULE_MAPPING. rules = { - (frozenset(substrings), fuse_into, fuse_from) + (substring, fuse_into, fuse_from) for substrings, fuse_into, fuse_from in iter_pqs_fuse_rules() + for substring in substrings } legacy = { - (frozenset({"LlamaAttention"}), "v_proj", "o_proj"), - (frozenset({"LlamaMLP"}), "up_proj", "down_proj"), - (frozenset({"Qwen3Attention", "Qwen3MoeAttention"}), "v_proj", "o_proj"), - (frozenset({"Qwen3MLP", "Qwen3MoeMLP"}), "up_proj", "down_proj"), + ("LlamaAttention", "v_proj", "o_proj"), + ("LlamaMLP", "up_proj", "down_proj"), + ("Qwen3Attention", "v_proj", "o_proj"), + ("Qwen3MoeAttention", "v_proj", "o_proj"), + ("Qwen3MLP", "up_proj", "down_proj"), + ("Qwen3MoeMLP", "up_proj", "down_proj"), } assert rules == legacy From 2c60c7ebfbb95b395c0a096df733a629225790f1 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 15 Jul 2026 23:40:51 +0000 Subject: [PATCH 04/16] rename Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/base.py | 6 ++++-- modelopt/modeling/models/arctic.py | 2 +- modelopt/modeling/models/dbrx.py | 4 ++-- modelopt/modeling/models/deepseek.py | 2 +- modelopt/modeling/models/gemma4.py | 2 +- modelopt/modeling/models/gpt_oss.py | 2 +- modelopt/modeling/models/llama.py | 2 +- modelopt/modeling/models/mixtral.py | 4 ++-- modelopt/modeling/models/nemotron_h.py | 2 +- modelopt/modeling/models/qwen2_moe.py | 2 +- modelopt/modeling/models/qwen3.py | 2 +- modelopt/modeling/models/qwen3_5_moe.py | 2 +- modelopt/modeling/models/qwen3_moe.py | 2 +- modelopt/modeling/models/qwen3_next.py | 2 +- tests/unit/torch/export/test_modeling_specs.py | 4 ++-- 15 files changed, 21 insertions(+), 19 deletions(-) diff --git a/modelopt/modeling/base.py b/modelopt/modeling/base.py index fe8618418ed..95c52c3d59d 100644 --- a/modelopt/modeling/base.py +++ b/modelopt/modeling/base.py @@ -35,8 +35,10 @@ class ModelSpec: ``type(module).__name__`` (e.g. ``"Qwen3MoeSparseMoeBlock"``). """ - name: str - """Unique spec name; the HF model type where one exists (e.g. ``"qwen3_moe"``).""" + model_type: str + """The HF model type this spec belongs to (``config.model_type``, e.g. + ``"qwen3_moe"``). Not necessarily unique: a model type may register several specs + for different module layouts (e.g. two ``mixtral`` MoE-block variants).""" moe_block_names: tuple[str, ...] = () """Matching key: MoE block class-name substrings (case-insensitive).""" diff --git a/modelopt/modeling/models/arctic.py b/modelopt/modeling/models/arctic.py index 1cd769ea1fe..5ea6a3fa5e4 100644 --- a/modelopt/modeling/models/arctic.py +++ b/modelopt/modeling/models/arctic.py @@ -22,7 +22,7 @@ # intentionally unset so expert-name lookups keep the engine default. register( ModelSpec( - name="arctic", + model_type="arctic", moe_block_names=("ArcticMoE",), ) ) diff --git a/modelopt/modeling/models/dbrx.py b/modelopt/modeling/models/dbrx.py index 2f006561b54..f4e0594508d 100644 --- a/modelopt/modeling/models/dbrx.py +++ b/modelopt/modeling/models/dbrx.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="dbrx", + model_type="dbrx", moe_block_names=("DBRXMoeSparseMoeBlock",), expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), ) @@ -30,7 +30,7 @@ # expert naming intentionally unset so expert-name lookups keep the engine default. register( ModelSpec( - name="dbrx_ffn", + model_type="dbrx", moe_block_names=("DbrxFFN",), ) ) diff --git a/modelopt/modeling/models/deepseek.py b/modelopt/modeling/models/deepseek.py index ba12f778754..8fb9026bcba 100644 --- a/modelopt/modeling/models/deepseek.py +++ b/modelopt/modeling/models/deepseek.py @@ -24,7 +24,7 @@ register( ModelSpec( - name="deepseek", + model_type="deepseek", moe_block_names=("DeepseekMoE",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), ) diff --git a/modelopt/modeling/models/gemma4.py b/modelopt/modeling/models/gemma4.py index 9b6284145ea..061e48bb191 100644 --- a/modelopt/modeling/models/gemma4.py +++ b/modelopt/modeling/models/gemma4.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="gemma4", + model_type="gemma4", # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. moe_block_names=("Gemma4TextDecoderLayer",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/models/gpt_oss.py b/modelopt/modeling/models/gpt_oss.py index 8ded20dd3ae..576aaae68c8 100644 --- a/modelopt/modeling/models/gpt_oss.py +++ b/modelopt/modeling/models/gpt_oss.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="gpt_oss", + model_type="gpt_oss", # GPT-OSS fuses gate and up into a single gate_up_proj. moe_block_names=("GptOssMoE",), expert_linear_names=("gate_up_proj", "down_proj"), diff --git a/modelopt/modeling/models/llama.py b/modelopt/modeling/models/llama.py index fe557e8becf..f1800b3c2cb 100644 --- a/modelopt/modeling/models/llama.py +++ b/modelopt/modeling/models/llama.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="llama", + model_type="llama", # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( (("LlamaAttention",), "v_proj", "o_proj"), diff --git a/modelopt/modeling/models/mixtral.py b/modelopt/modeling/models/mixtral.py index b3b8c76911b..ae2947d3b5c 100644 --- a/modelopt/modeling/models/mixtral.py +++ b/modelopt/modeling/models/mixtral.py @@ -22,7 +22,7 @@ # detected from their per-expert quantizer attributes and need no naming override here. register( ModelSpec( - name="mixtral", + model_type="mixtral", moe_block_names=("MixtralSparseMoeBlock",), expert_linear_names=("w1", "w2", "w3"), has_iterable_experts=True, @@ -32,7 +32,7 @@ # Older transformers naming for Mixtral. register( ModelSpec( - name="mixtral_mcore", + model_type="mixtral", moe_block_names=("MixtralMoeSparseMoeBlock",), expert_linear_names=("linear_fc1", "linear_fc2"), ) diff --git a/modelopt/modeling/models/nemotron_h.py b/modelopt/modeling/models/nemotron_h.py index 638bf248097..651f330fbe3 100644 --- a/modelopt/modeling/models/nemotron_h.py +++ b/modelopt/modeling/models/nemotron_h.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="nemotron_h", + model_type="nemotron_h", # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). moe_block_names=("NemotronHMOE",), expert_linear_names=("up_proj", "down_proj"), diff --git a/modelopt/modeling/models/qwen2_moe.py b/modelopt/modeling/models/qwen2_moe.py index 65f5fe755d9..0af58722d4c 100644 --- a/modelopt/modeling/models/qwen2_moe.py +++ b/modelopt/modeling/models/qwen2_moe.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="qwen2_moe", + model_type="qwen2_moe", moe_block_names=("Qwen2MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, diff --git a/modelopt/modeling/models/qwen3.py b/modelopt/modeling/models/qwen3.py index 78f2388c40a..23dd54800fc 100644 --- a/modelopt/modeling/models/qwen3.py +++ b/modelopt/modeling/models/qwen3.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="qwen3", + model_type="qwen3", # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( (("Qwen3Attention",), "v_proj", "o_proj"), diff --git a/modelopt/modeling/models/qwen3_5_moe.py b/modelopt/modeling/models/qwen3_5_moe.py index 178a3677518..607ef704365 100644 --- a/modelopt/modeling/models/qwen3_5_moe.py +++ b/modelopt/modeling/models/qwen3_5_moe.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="qwen3_5_moe", + model_type="qwen3_5_moe", moe_block_names=("Qwen3_5MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, diff --git a/modelopt/modeling/models/qwen3_moe.py b/modelopt/modeling/models/qwen3_moe.py index bbcf73d60e0..6be4e94bc10 100644 --- a/modelopt/modeling/models/qwen3_moe.py +++ b/modelopt/modeling/models/qwen3_moe.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="qwen3_moe", + model_type="qwen3_moe", moe_block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, diff --git a/modelopt/modeling/models/qwen3_next.py b/modelopt/modeling/models/qwen3_next.py index f0aa8705a1c..93c04f68463 100644 --- a/modelopt/modeling/models/qwen3_next.py +++ b/modelopt/modeling/models/qwen3_next.py @@ -20,7 +20,7 @@ register( ModelSpec( - name="qwen3_next", + model_type="qwen3_next", moe_block_names=("Qwen3NextSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index c2d61337b49..00fd307f829 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -37,7 +37,7 @@ class _UnknownMoeBlock(nn.Module): def test_match_moe_block_by_substring(): spec = match_moe_block(_FakeQwen3MoeSparseMoeBlock()) assert spec is not None - assert spec.name == "qwen3_moe" + assert spec.model_type == "qwen3_moe" assert spec.expert_linear_names == ("gate_proj", "down_proj", "up_proj") assert spec.has_iterable_experts @@ -45,7 +45,7 @@ def test_match_moe_block_by_substring(): def test_match_moe_block_matches_quantized_class_name(): spec = match_moe_block(_QuantMixtralSparseMoeBlock()) assert spec is not None - assert spec.name == "mixtral" + assert spec.model_type == "mixtral" def test_match_moe_block_unmatched_returns_none(): From ddbf88f15475e435900d98fc15e2951654cba907 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 15 Jul 2026 23:44:06 +0000 Subject: [PATCH 05/16] categorized spec Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/__init__.py | 8 +-- modelopt/modeling/base.py | 34 +++---------- modelopt/modeling/export.py | 49 +++++++++++++++++++ modelopt/modeling/models/arctic.py | 4 +- modelopt/modeling/models/dbrx.py | 6 +-- modelopt/modeling/models/deepseek.py | 4 +- modelopt/modeling/models/gemma4.py | 4 +- modelopt/modeling/models/gpt_oss.py | 4 +- modelopt/modeling/models/llama.py | 4 +- modelopt/modeling/models/mixtral.py | 6 +-- modelopt/modeling/models/nemotron_h.py | 4 +- modelopt/modeling/models/qwen2_moe.py | 4 +- modelopt/modeling/models/qwen3.py | 4 +- modelopt/modeling/models/qwen3_5_moe.py | 4 +- modelopt/modeling/models/qwen3_moe.py | 4 +- modelopt/modeling/models/qwen3_next.py | 4 +- modelopt/modeling/registry.py | 30 ++++++++---- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 9 ++-- 18 files changed, 117 insertions(+), 69 deletions(-) create mode 100644 modelopt/modeling/export.py diff --git a/modelopt/modeling/__init__.py b/modelopt/modeling/__init__.py index 6ce8d7ec3b7..503a01955f2 100644 --- a/modelopt/modeling/__init__.py +++ b/modelopt/modeling/__init__.py @@ -16,12 +16,14 @@ """Per-model descriptors, organized by HF model type. Holds declarative per-model data (no algorithms), one module per HF model type under -``models/``, mirroring ``transformers.models``. Consumers (e.g. the unified HF export -path) resolve a ``ModelSpec`` via the registry lookups and read its fields; an -unmatched lookup returns ``None`` so callers fall back to their default behavior. +``models/``, mirroring ``transformers.models``. Each modelopt subsystem has its own +``ModelSpec`` subclass (e.g. ``ExportSpec``); consumers resolve a spec via the +registry lookups and read its fields; an unmatched lookup returns ``None`` so callers +fall back to their default behavior. """ # Importing models registers every spec as a side effect. from . import models from .base import * +from .export import * from .registry import * diff --git a/modelopt/modeling/base.py b/modelopt/modeling/base.py index 95c52c3d59d..e49b1875a6c 100644 --- a/modelopt/modeling/base.py +++ b/modelopt/modeling/base.py @@ -13,11 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-model descriptor. +"""Common base for per-model descriptors. -A ``ModelSpec`` declares how one model type differs from the generic code path, -so consumers (e.g. export) can read these values instead of branching on model -names. Each spec holds per-model data only, no logic. +Each modelopt subsystem declares its per-model data in its own ``ModelSpec`` +subclass (e.g. ``ExportSpec``; quantization / speculative-decoding specs to follow), +so consumers can read these values instead of branching on model names. Specs hold +per-model data only, no logic. """ from dataclasses import dataclass @@ -27,32 +28,13 @@ @dataclass class ModelSpec: - """Per-model data. + """Base class for per-model data specs. - A spec is resolved from a model sub-module via its matching keys and read for its - per-model data fields. ``moe_block_names`` is the matching key: MoE block - class-name substrings compared case-insensitively against - ``type(module).__name__`` (e.g. ``"Qwen3MoeSparseMoeBlock"``). + Subclasses add the data fields of one modelopt subsystem; a model registers one + spec instance per subsystem it customizes (see ``models/``). """ model_type: str """The HF model type this spec belongs to (``config.model_type``, e.g. ``"qwen3_moe"``). Not necessarily unique: a model type may register several specs for different module layouts (e.g. two ``mixtral`` MoE-block variants).""" - - moe_block_names: tuple[str, ...] = () - """Matching key: MoE block class-name substrings (case-insensitive).""" - - expert_linear_names: tuple[str, ...] | None = None - """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``.""" - - has_iterable_experts: bool = False - """True when experts are per-expert iterable sub-modules (Mixtral, Qwen MoE, - NemotronH, Gemma4) and can be grouped by ``get_experts_list``; False for stacked - or fused layouts (DBRX, GptOss).""" - - pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = () - """AWQ ``pre_quant_scale`` fusion rules, each a ``(module_class_substrings, - fuse_into, fuse_from)`` triple: for a module whose class name contains one of the - substrings, the pre_quant_scale on ``fuse_from`` is folded into ``fuse_into`` - (e.g. attention ``o_proj`` -> ``v_proj``, MLP ``down_proj`` -> ``up_proj``).""" diff --git a/modelopt/modeling/export.py b/modelopt/modeling/export.py new file mode 100644 index 00000000000..b9ae25616f5 --- /dev/null +++ b/modelopt/modeling/export.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model data consumed by the unified HF export path.""" + +from dataclasses import dataclass + +from .base import ModelSpec + +__all__ = ["ExportSpec"] + + +@dataclass +class ExportSpec(ModelSpec): + """Per-model data for the unified HF export path. + + Resolved from a model sub-module via ``moe_block_names``, the matching key: MoE + block class-name substrings compared case-insensitively against + ``type(module).__name__`` (e.g. ``"Qwen3MoeSparseMoeBlock"``). + """ + + moe_block_names: tuple[str, ...] = () + """Matching key: MoE block class-name substrings (case-insensitive).""" + + expert_linear_names: tuple[str, ...] | None = None + """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``.""" + + has_iterable_experts: bool = False + """True when experts are per-expert iterable sub-modules (Mixtral, Qwen MoE, + NemotronH, Gemma4) and can be grouped by ``get_experts_list``; False for stacked + or fused layouts (DBRX, GptOss).""" + + pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = () + """AWQ ``pre_quant_scale`` fusion rules, each a ``(module_class_substrings, + fuse_into, fuse_from)`` triple: for a module whose class name contains one of the + substrings, the pre_quant_scale on ``fuse_from`` is folded into ``fuse_into`` + (e.g. attention ``o_proj`` -> ``v_proj``, MLP ``down_proj`` -> ``up_proj``).""" diff --git a/modelopt/modeling/models/arctic.py b/modelopt/modeling/models/arctic.py index 5ea6a3fa5e4..45785fa13cd 100644 --- a/modelopt/modeling/models/arctic.py +++ b/modelopt/modeling/models/arctic.py @@ -15,13 +15,13 @@ """Snowflake Arctic specs (trust-remote-code model type ``arctic``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register # MoE-block identification only (non-standard block name for is_moe); expert naming # intentionally unset so expert-name lookups keep the engine default. register( - ModelSpec( + ExportSpec( model_type="arctic", moe_block_names=("ArcticMoE",), ) diff --git a/modelopt/modeling/models/dbrx.py b/modelopt/modeling/models/dbrx.py index f4e0594508d..eb6e32e6f20 100644 --- a/modelopt/modeling/models/dbrx.py +++ b/modelopt/modeling/models/dbrx.py @@ -15,11 +15,11 @@ """DBRX specs (HF model type ``dbrx``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="dbrx", moe_block_names=("DBRXMoeSparseMoeBlock",), expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), @@ -29,7 +29,7 @@ # HF DbrxFFN: MoE-block identification only (non-standard block name for is_moe); # expert naming intentionally unset so expert-name lookups keep the engine default. register( - ModelSpec( + ExportSpec( model_type="dbrx", moe_block_names=("DbrxFFN",), ) diff --git a/modelopt/modeling/models/deepseek.py b/modelopt/modeling/models/deepseek.py index 8fb9026bcba..eb6bb3dc748 100644 --- a/modelopt/modeling/models/deepseek.py +++ b/modelopt/modeling/models/deepseek.py @@ -19,11 +19,11 @@ classes. """ -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="deepseek", moe_block_names=("DeepseekMoE",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/models/gemma4.py b/modelopt/modeling/models/gemma4.py index 061e48bb191..55323c7287e 100644 --- a/modelopt/modeling/models/gemma4.py +++ b/modelopt/modeling/models/gemma4.py @@ -15,11 +15,11 @@ """Gemma4 specs (HF model type ``gemma4``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="gemma4", # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. moe_block_names=("Gemma4TextDecoderLayer",), diff --git a/modelopt/modeling/models/gpt_oss.py b/modelopt/modeling/models/gpt_oss.py index 576aaae68c8..1c1bd06d504 100644 --- a/modelopt/modeling/models/gpt_oss.py +++ b/modelopt/modeling/models/gpt_oss.py @@ -15,11 +15,11 @@ """GPT-OSS specs (HF model type ``gpt_oss``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="gpt_oss", # GPT-OSS fuses gate and up into a single gate_up_proj. moe_block_names=("GptOssMoE",), diff --git a/modelopt/modeling/models/llama.py b/modelopt/modeling/models/llama.py index f1800b3c2cb..4e8da2e6ad4 100644 --- a/modelopt/modeling/models/llama.py +++ b/modelopt/modeling/models/llama.py @@ -15,11 +15,11 @@ """Llama specs (HF model type ``llama``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="llama", # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( diff --git a/modelopt/modeling/models/mixtral.py b/modelopt/modeling/models/mixtral.py index ae2947d3b5c..4c069cde53b 100644 --- a/modelopt/modeling/models/mixtral.py +++ b/modelopt/modeling/models/mixtral.py @@ -15,13 +15,13 @@ """Mixtral specs (HF model type ``mixtral``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register # Mixtral with iterable experts uses w1/w2/w3. Fused experts (transformers 5.0+) are # detected from their per-expert quantizer attributes and need no naming override here. register( - ModelSpec( + ExportSpec( model_type="mixtral", moe_block_names=("MixtralSparseMoeBlock",), expert_linear_names=("w1", "w2", "w3"), @@ -31,7 +31,7 @@ # Older transformers naming for Mixtral. register( - ModelSpec( + ExportSpec( model_type="mixtral", moe_block_names=("MixtralMoeSparseMoeBlock",), expert_linear_names=("linear_fc1", "linear_fc2"), diff --git a/modelopt/modeling/models/nemotron_h.py b/modelopt/modeling/models/nemotron_h.py index 651f330fbe3..96429eedc39 100644 --- a/modelopt/modeling/models/nemotron_h.py +++ b/modelopt/modeling/models/nemotron_h.py @@ -15,11 +15,11 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="nemotron_h", # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). moe_block_names=("NemotronHMOE",), diff --git a/modelopt/modeling/models/qwen2_moe.py b/modelopt/modeling/models/qwen2_moe.py index 0af58722d4c..c6b2e391666 100644 --- a/modelopt/modeling/models/qwen2_moe.py +++ b/modelopt/modeling/models/qwen2_moe.py @@ -15,11 +15,11 @@ """Qwen2-MoE specs (HF model type ``qwen2_moe``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="qwen2_moe", moe_block_names=("Qwen2MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/models/qwen3.py b/modelopt/modeling/models/qwen3.py index 23dd54800fc..95e12af549c 100644 --- a/modelopt/modeling/models/qwen3.py +++ b/modelopt/modeling/models/qwen3.py @@ -15,11 +15,11 @@ """Qwen3 (dense) specs (HF model type ``qwen3``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="qwen3", # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( diff --git a/modelopt/modeling/models/qwen3_5_moe.py b/modelopt/modeling/models/qwen3_5_moe.py index 607ef704365..f265cc1e32a 100644 --- a/modelopt/modeling/models/qwen3_5_moe.py +++ b/modelopt/modeling/models/qwen3_5_moe.py @@ -15,11 +15,11 @@ """Qwen3.5-MoE specs (HF model type ``qwen3_5_moe``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="qwen3_5_moe", moe_block_names=("Qwen3_5MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/models/qwen3_moe.py b/modelopt/modeling/models/qwen3_moe.py index 6be4e94bc10..831e3b54b0c 100644 --- a/modelopt/modeling/models/qwen3_moe.py +++ b/modelopt/modeling/models/qwen3_moe.py @@ -15,11 +15,11 @@ """Qwen3-MoE specs (HF model type ``qwen3_moe``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="qwen3_moe", moe_block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/models/qwen3_next.py b/modelopt/modeling/models/qwen3_next.py index 93c04f68463..bc220954851 100644 --- a/modelopt/modeling/models/qwen3_next.py +++ b/modelopt/modeling/models/qwen3_next.py @@ -15,11 +15,11 @@ """Qwen3-Next specs (HF model type ``qwen3_next``).""" -from ..base import ModelSpec +from ..export import ExportSpec from ..registry import register register( - ModelSpec( + ExportSpec( model_type="qwen3_next", moe_block_names=("Qwen3NextSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index 77c32a5cadc..62239160974 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Registry that resolves a model sub-module to its ``ModelSpec``. +"""Registry of per-model specs, queried by spec type. Model modules register their specs at import time (see ``models/``). Lookups return ``None`` when nothing matches, so callers can fall back to their default behavior. @@ -23,45 +23,57 @@ here). """ -from typing import TYPE_CHECKING +from collections.abc import Iterator +from typing import TYPE_CHECKING, TypeVar from .base import ModelSpec +from .export import ExportSpec if TYPE_CHECKING: import torch.nn as nn __all__ = [ "iter_pqs_fuse_rules", + "iter_specs", "match_moe_block", "register", ] +SpecT = TypeVar("SpecT", bound=ModelSpec) + _SPECS: list[ModelSpec] = [] -def register(spec: ModelSpec) -> ModelSpec: +def register(spec: SpecT) -> SpecT: """Register a model spec and return it.""" _SPECS.append(spec) return spec +def iter_specs(spec_cls: type[SpecT]) -> Iterator[SpecT]: + """Yield every registered spec of type ``spec_cls`` (in registration order).""" + for spec in _SPECS: + if isinstance(spec, spec_cls): + yield spec + + def iter_pqs_fuse_rules(): """Yield every ``(module_class_substrings, fuse_into, fuse_from)`` AWQ fusion rule. - Aggregated across all registered specs (the consumer matches each model module - against the substrings, so the order across specs does not matter). + Aggregated across all registered export specs (the consumer matches each model + module against the substrings, so the order across specs does not matter). """ - for spec in _SPECS: + for spec in iter_specs(ExportSpec): yield from spec.pqs_fuse_rules -def match_moe_block(module: "nn.Module") -> ModelSpec | None: - """Return the spec matching ``module``'s class name against ``moe_block_names``. +def match_moe_block(module: "nn.Module") -> ExportSpec | None: + """Return the export spec matching ``module``'s class name against ``moe_block_names``. Case-insensitive substring match against ``type(module).__name__``. """ cls_name = type(module).__name__.lower() - for spec in _SPECS: + for spec in iter_specs(ExportSpec): if any(name.lower() in cls_name for name in spec.moe_block_names): return spec return None diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index c4d0c9d780b..dab92031df5 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -49,11 +49,14 @@ Two layers: ```text modelopt/modeling/ - base.py # ModelSpec dataclass — the per-model contract - registry.py # register() + lookups; returns None when unmatched + base.py # ModelSpec — common base class (model_type) + export.py # ExportSpec(ModelSpec) — data for the unified HF export path + # (future: quantization / speculative-decoding spec classes) + registry.py # register() + lookups, queried by spec type; None when unmatched __init__.py # re-exports; importing it registers all specs models/ # one small file per HF model type (mirrors transformers.models); - # import == registration + # import == registration; a model registers one spec instance per + # subsystem it customizes ``` Model type names mirror From c5c68b2db6dcdb16c73aa5ccb60362fb590e148e Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 16 Jul 2026 00:14:01 +0000 Subject: [PATCH 06/16] better matching Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/export.py | 7 +-- modelopt/modeling/matching.py | 39 +++++++++++++++++ modelopt/modeling/registry.py | 16 ++++--- tests/unit/torch/export/test_layer_utils.py | 7 +-- .../unit/torch/export/test_modeling_specs.py | 43 ++++++++++++------- 5 files changed, 84 insertions(+), 28 deletions(-) create mode 100644 modelopt/modeling/matching.py diff --git a/modelopt/modeling/export.py b/modelopt/modeling/export.py index b9ae25616f5..2d34be80f4b 100644 --- a/modelopt/modeling/export.py +++ b/modelopt/modeling/export.py @@ -27,12 +27,13 @@ class ExportSpec(ModelSpec): """Per-model data for the unified HF export path. Resolved from a model sub-module via ``moe_block_names``, the matching key: MoE - block class-name substrings compared case-insensitively against - ``type(module).__name__`` (e.g. ``"Qwen3MoeSparseMoeBlock"``). + block class names (e.g. ``"Qwen3MoeSparseMoeBlock"``) compared case-insensitively + against the class names in the module's MRO (see ``matching.match_class_names``). """ moe_block_names: tuple[str, ...] = () - """Matching key: MoE block class-name substrings (case-insensitive).""" + """Matching key: MoE block class names, matched against the module's MRO + (case-insensitive exact names, not substrings).""" expert_linear_names: tuple[str, ...] | None = None """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``.""" diff --git a/modelopt/modeling/matching.py b/modelopt/modeling/matching.py new file mode 100644 index 00000000000..7b47cab03ce --- /dev/null +++ b/modelopt/modeling/matching.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Module-to-name matching shared by the spec registry (and, later, dispatch registries). + +Same semantics as the export dispatch registry's string keys +(``modelopt.torch.export.registry``): a name matches when it equals the ``__name__`` +of a class in ``type(module).__mro__``. Dynamically generated quantized classes are +subclasses of the original module class, so they match through their base; exact-name +comparison avoids substring false positives. Comparison is case-insensitive because +some registered names predate this registry and their casing was never exercised by +the legacy substring matching. + +Only inspects ``type(module).__mro__``, so this module stays stdlib-only. +""" + +__all__ = ["match_class_names"] + + +def match_class_names(module, names: tuple[str, ...]) -> bool: + """Return True if any of ``names`` equals a class name in ``module``'s MRO. + + Case-insensitive exact-name comparison against ``cls.__name__`` for every class + in ``type(module).__mro__``. + """ + mro_names = {cls.__name__.lower() for cls in type(module).__mro__} + return any(name.lower() in mro_names for name in names) diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index 62239160974..47eb01ed8cf 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -18,9 +18,9 @@ Model modules register their specs at import time (see ``models/``). Lookups return ``None`` when nothing matches, so callers can fall back to their default behavior. -Matching is by class-name string only, so this package stays dependency-free (any -``nn.Module`` — or any object — can be passed to the lookups without importing torch -here). +Matching is by class-name string only (see ``matching``), so this package stays +dependency-free (any ``nn.Module`` — or any object — can be passed to the lookups +without importing torch here). """ from collections.abc import Iterator @@ -28,6 +28,7 @@ from .base import ModelSpec from .export import ExportSpec +from .matching import match_class_names if TYPE_CHECKING: import torch.nn as nn @@ -68,12 +69,13 @@ def iter_pqs_fuse_rules(): def match_moe_block(module: "nn.Module") -> ExportSpec | None: - """Return the export spec matching ``module``'s class name against ``moe_block_names``. + """Return the export spec whose ``moe_block_names`` matches ``module``. - Case-insensitive substring match against ``type(module).__name__``. + Case-insensitive exact-name match against the class names in ``module``'s MRO + (see ``matching.match_class_names``); quantized wrapper classes match through + their original base class. """ - cls_name = type(module).__name__.lower() for spec in iter_specs(ExportSpec): - if any(name.lower() in cls_name for name in spec.moe_block_names): + if match_class_names(module, spec.moe_block_names): return spec return None diff --git a/tests/unit/torch/export/test_layer_utils.py b/tests/unit/torch/export/test_layer_utils.py index 93742c8fe7e..0cd5b25d24b 100644 --- a/tests/unit/torch/export/test_layer_utils.py +++ b/tests/unit/torch/export/test_layer_utils.py @@ -33,8 +33,9 @@ class _FakeMoeLayer(nn.Module): """Name contains 'moelayer' — detected by naming convention.""" -class _FakeArcticMoe(nn.Module): - """Name contains 'arcticmoe' — detected by explicit match.""" +class ArcticMoE(nn.Module): + """Non-standard MoE block name — detected via the model spec registry (exact + MRO class name, so the fake must carry the real name).""" class _StructuralMoeModule(nn.Module): @@ -64,7 +65,7 @@ def __init__(self): @pytest.mark.parametrize( "module_cls", - [_FakeSparseMoeBlock, _FakeMoeLayer, _FakeArcticMoe], + [_FakeSparseMoeBlock, _FakeMoeLayer, ArcticMoE], ) def test_is_moe_name_based(module_cls): assert is_moe(module_cls()) diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index 00fd307f829..1bd284bd695 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -22,28 +22,32 @@ from modelopt.torch.export.layer_utils import get_expert_linear_names, get_experts_list, is_moe -class _FakeQwen3MoeSparseMoeBlock(nn.Module): +class Qwen3MoeSparseMoeBlock(nn.Module): pass -class _QuantMixtralSparseMoeBlock(nn.Module): - """Dynamically generated quantized classes keep the family name as a substring.""" +class MixtralSparseMoeBlock(nn.Module): + pass + + +class QuantMixtralSparseMoeBlock(MixtralSparseMoeBlock): + """Quantized classes subclass the original module class; matching goes through the MRO.""" class _UnknownMoeBlock(nn.Module): pass -def test_match_moe_block_by_substring(): - spec = match_moe_block(_FakeQwen3MoeSparseMoeBlock()) +def test_match_moe_block_by_class_name(): + spec = match_moe_block(Qwen3MoeSparseMoeBlock()) assert spec is not None assert spec.model_type == "qwen3_moe" assert spec.expert_linear_names == ("gate_proj", "down_proj", "up_proj") assert spec.has_iterable_experts -def test_match_moe_block_matches_quantized_class_name(): - spec = match_moe_block(_QuantMixtralSparseMoeBlock()) +def test_match_moe_block_matches_quantized_class_via_mro(): + spec = match_moe_block(QuantMixtralSparseMoeBlock()) assert spec is not None assert spec.model_type == "mixtral" @@ -57,7 +61,7 @@ def test_get_expert_linear_names_falls_back_when_unmatched(): assert get_expert_linear_names(_UnknownMoeBlock()) == ["w1", "w2", "w3"] -class _FakeNemotronHMOE(nn.Module): +class NemotronHMOE(nn.Module): def __init__(self): super().__init__() @@ -71,7 +75,7 @@ def __init__(self): def test_get_experts_list_groups_by_spec_linear_names(): - module = _FakeNemotronHMOE() + module = NemotronHMOE() groups = get_experts_list(module, "nemotronhforcausallm") assert len(groups) == 2 # up_proj group + down_proj group assert all(len(group) == 3 for group in groups) @@ -82,34 +86,43 @@ def test_get_experts_list_groups_by_spec_linear_names(): def test_get_experts_list_rejects_non_iterable_layouts(): # DBRX matches a spec but is not an iterable-experts layout; grouped export # must keep rejecting it (legacy behavior). - class _FakeDBRXMoeSparseMoeBlock(nn.Module): + class DBRXMoeSparseMoeBlock(nn.Module): def __init__(self): super().__init__() self.experts = nn.ModuleList() with pytest.raises(NotImplementedError): - get_experts_list(_FakeDBRXMoeSparseMoeBlock(), "dbrxforcausallm") + get_experts_list(DBRXMoeSparseMoeBlock(), "dbrxforcausallm") with pytest.raises(NotImplementedError): get_experts_list(_UnknownMoeBlock(), "unknownforcausallm") -class _FakeArcticMoE(nn.Module): +class ArcticMoE(nn.Module): pass -class _FakeDbrxFFN(nn.Module): +class DbrxFFN(nn.Module): pass def test_is_moe_matches_registered_non_standard_names(): # Non-standard MoE block names (no *SparseMoeBlock suffix, no router/experts # attributes) resolve through the model spec registry. - assert is_moe(_FakeArcticMoE()) - assert is_moe(_FakeDbrxFFN()) + assert is_moe(ArcticMoE()) + assert is_moe(DbrxFFN()) assert not is_moe(_UnknownMoeBlock()) +def test_match_is_exact_name_not_substring(): + # A class whose name merely CONTAINS a registered name must not match; only + # exact MRO class names do (quantized wrappers match via their base class). + class MyArcticMoEWrapper(nn.Module): + pass + + assert match_moe_block(MyArcticMoEWrapper()) is None + + def test_pqs_fuse_rules_match_legacy_mapping(): # Aggregated per-model rules, flattened to (class_substring, fuse_into, fuse_from) # triples, must reproduce the legacy PQS_FUSE_MODULE_MAPPING. From 0fea97ce587158bbe4ac542117edf9cd692d6841 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 16 Jul 2026 00:59:10 +0000 Subject: [PATCH 07/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/models/arctic.py | 6 +- modelopt/modeling/models/dbrx.py | 8 ++- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 21 +++--- modelopt/torch/export/layer_utils.py | 66 +++++-------------- tests/unit/torch/export/test_layer_utils.py | 12 ++-- .../unit/torch/export/test_modeling_specs.py | 22 ++++++- 6 files changed, 63 insertions(+), 72 deletions(-) diff --git a/modelopt/modeling/models/arctic.py b/modelopt/modeling/models/arctic.py index 45785fa13cd..73bede55bce 100644 --- a/modelopt/modeling/models/arctic.py +++ b/modelopt/modeling/models/arctic.py @@ -18,11 +18,13 @@ from ..export import ExportSpec from ..registry import register -# MoE-block identification only (non-standard block name for is_moe); expert naming -# intentionally unset so expert-name lookups keep the engine default. register( ExportSpec( model_type="arctic", + # Non-standard block name (for is_moe identification). moe_block_names=("ArcticMoE",), + # ArcticMLP experts use Mixtral-style w1/w2/w3 naming (previously served by + # the engine's implicit w1/w2/w3 default, now declared explicitly). + expert_linear_names=("w1", "w2", "w3"), ) ) diff --git a/modelopt/modeling/models/dbrx.py b/modelopt/modeling/models/dbrx.py index eb6e32e6f20..d8ffc131bb3 100644 --- a/modelopt/modeling/models/dbrx.py +++ b/modelopt/modeling/models/dbrx.py @@ -26,11 +26,15 @@ ) ) -# HF DbrxFFN: MoE-block identification only (non-standard block name for is_moe); -# expert naming intentionally unset so expert-name lookups keep the engine default. +# HF DbrxFFN (non-standard block name, identified for is_moe). Expert names refer to +# the quantized layout: _QuantDbrxExpertGLU converts the fused w1/v1/w2 parameters +# into per-expert w1_linear/v1_linear/w2_linear ModuleLists on experts.mlp (see +# modelopt/torch/quantization/plugins/huggingface.py), which is what the DBRX +# prepare handler fills amax values on. register( ExportSpec( model_type="dbrx", moe_block_names=("DbrxFFN",), + expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), ) ) diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index dab92031df5..adc50db351d 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -64,15 +64,15 @@ Model type names mirror (e.g. `qwen3_moe.py`, `gpt_oss.py`, `nemotron_h.py`); trust-remote-code models (`arctic`, `deepseek`) use their config `model_type`. -Call sites follow a fallback-first shape, which keeps migration incremental and -behavior-preserving — a model not in the registry behaves exactly as before: - -```python -spec = match_moe_block(module) -if spec is not None and spec.expert_linear_names is not None: - return list(spec.expert_linear_names) -# ... legacy branch preserved as fallback ... -``` +During migration, call sites kept the legacy branches as a fallback behind the +spec lookup. Once the specs covered every family the legacy chains served, the +chains — and the silent ``w1/w2/w3`` guess for unknown models — were deleted: +expert-name resolution is now *structural detection -> spec -> raise*, so a new +MoE model fails loudly, asking for a spec, instead of inheriting another model's +naming. Generic detection that is not per-model data (the ``*SparseMoeBlock`` / +``*MoeLayer`` conventions and the router+experts structural check in ``is_moe``, +the fused-experts quantizer probe in ``get_expert_linear_names``) stays in the +engine, ahead of or beside the spec lookup. Note on naming: `modelopt/modeling/registry.py` (per-model **data**, "what are this model's values") is distinct from the export-path `registry.py` from PR #1939 @@ -92,5 +92,6 @@ against existing export tests. | **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). | planned | | **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Dead code found during inventory (`MODEL_NAME_TO_TYPE`, `get_model_type`, `adjust_attn_amax_values`, `update_experts_avg_prequant_scale`) is deleted on that track too. | separate track | -**Guardrails:** fallback-first; one data category per PR; the engine keeps the +**Guardrails:** one data category per PR (fallback-first while a category is +partially migrated, explicit-error once specs cover it); the engine keeps the algorithms — model specs supply values only, never fork functions. diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index 262235a5d54..a2152f4443e 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -94,14 +94,14 @@ def get_experts_list( """ experts_list = [] - # Expert linear names are per-model data (modelopt/modeling/models/*). The grouped export - # path only supports families whose experts are iterable per-expert sub-modules - # (Mixtral, Qwen MoE, NemotronH, Gemma4); stacked/fused layouts (DBRX, GptOss, ...) - # raise NotImplementedError here and are handled by other paths. + # The grouped export path only supports layouts whose experts are iterable + # per-expert sub-modules (spec.has_iterable_experts); stacked/fused layouts + # (DBRX, GptOss, ...) raise NotImplementedError here and are handled by other + # paths. Name resolution itself is shared with get_expert_linear_names. spec = match_moe_block(module) - if spec is None or not spec.has_iterable_experts or spec.expert_linear_names is None: + if spec is None or not spec.has_iterable_experts: raise NotImplementedError(f" {model_type} not supported") - linear_names = list(spec.expert_linear_names) + linear_names = get_expert_linear_names(module) # Common logic for all supported model types experts_list.extend( @@ -964,16 +964,13 @@ def get_stacked_scaling_factors(experts, get_function, module_name): def get_expert_linear_names(module: nn.Module) -> list[str]: - """Get the list of linear names for the experts.""" - - def module_match_name_list(module, name_list): - """Check if the module name matches any of the names in the list. - - e.g. module_match_name_list(QuantQwen3MoeSparseMoeBlock, ['Qwen3MoeSparseMoeBlock']) -> True - - """ - return any(name.lower() in type(module).__name__.lower() for name in name_list) + """Get the list of linear names for the experts. + Resolution order: structural detection of fused-expert layouts first, then the + model spec registry (modelopt/modeling). Raises NotImplementedError when neither + resolves, so a new MoE model fails loudly instead of silently inheriting another + model's naming. + """ # Structural detection: after _export_fused_experts, fused expert modules # have per-expert submodules with gate_proj/up_proj/down_proj. # Also handles models that originally used this naming (Qwen, DeepSeek, etc.). @@ -982,44 +979,15 @@ def module_match_name_list(module, name_list): if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"): return [first_proj_attr, "down_proj"] - # Resolve expert names from the model spec registry (modelopt/modeling); fall back - # to the mapping below when no spec matches the MoE block. spec = match_moe_block(module) if spec is not None and spec.expert_linear_names is not None: return list(spec.expert_linear_names) - if module_match_name_list( - module, - [ - "Qwen2MoeSparseMoeBlock", - "Qwen3MoeSparseMoeBlock", - "Qwen3NextSparseMoeBlock", - "Qwen3_5MoeSparseMoeBlock", - "DeepseekMoE", - ], - ): - return ["gate_proj", "down_proj", "up_proj"] - elif module_match_name_list(module, ["MixtralSparseMoeBlock"]): - # Old-style Mixtral (iterable experts) uses w1/w2/w3. - # Fused Mixtral (transformers 5.0+) is already handled by the - # structural first-projection quantizer check above. - return ["w1", "w2", "w3"] - elif module_match_name_list(module, ["MixtralMoeSparseMoeBlock"]): - # Older transformers naming for Mixtral - return ["linear_fc1", "linear_fc2"] - elif module_match_name_list(module, ["DBRXMoeSparseMoeBlock"]): - return ["w1_linear", "w2_linear", "v1_linear"] - elif module_match_name_list(module, ["GptOssMoE"]): - return ["gate_up_proj", "down_proj"] - elif module_match_name_list(module, ["NemotronHMOE"]): - # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). - return ["up_proj", "down_proj"] - elif module_match_name_list(module, ["Gemma4TextDecoderLayer"]): - # Gemma4 MoE experts are unfused into per-expert nn.Linear layers - return ["gate_proj", "down_proj", "up_proj"] - else: - # assuming w1, w2, w3 by default - return ["w1", "w2", "w3"] + raise NotImplementedError( + f"Cannot resolve expert linear names for MoE block {type(module).__name__!r}. " + "Register an ExportSpec with expert_linear_names for this model under " + "modelopt/modeling/models/." + ) def set_expert_quantizer_amax( diff --git a/tests/unit/torch/export/test_layer_utils.py b/tests/unit/torch/export/test_layer_utils.py index 0cd5b25d24b..538de5d605e 100644 --- a/tests/unit/torch/export/test_layer_utils.py +++ b/tests/unit/torch/export/test_layer_utils.py @@ -88,20 +88,20 @@ def test_is_moe_partial_structural(): # --------------------------------------------------------------------------- -class _FakeGemma4TextDecoderLayer(nn.Module): +class Gemma4TextDecoderLayer(nn.Module): pass -class _FakeMixtralSparseMoeBlock(nn.Module): +class MixtralSparseMoeBlock(nn.Module): pass -class _FakeNemotronHMOE(nn.Module): +class NemotronHMOE(nn.Module): pass def test_get_expert_linear_names_gemma4(): - assert get_expert_linear_names(_FakeGemma4TextDecoderLayer()) == [ + assert get_expert_linear_names(Gemma4TextDecoderLayer()) == [ "gate_proj", "down_proj", "up_proj", @@ -109,8 +109,8 @@ def test_get_expert_linear_names_gemma4(): def test_get_expert_linear_names_mixtral(): - assert get_expert_linear_names(_FakeMixtralSparseMoeBlock()) == ["w1", "w2", "w3"] + assert get_expert_linear_names(MixtralSparseMoeBlock()) == ["w1", "w2", "w3"] def test_get_expert_linear_names_nemotron(): - assert get_expert_linear_names(_FakeNemotronHMOE()) == ["up_proj", "down_proj"] + assert get_expert_linear_names(NemotronHMOE()) == ["up_proj", "down_proj"] diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index 1bd284bd695..5c0f8f9cf67 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -56,9 +56,25 @@ def test_match_moe_block_unmatched_returns_none(): assert match_moe_block(_UnknownMoeBlock()) is None -def test_get_expert_linear_names_falls_back_when_unmatched(): - # No spec matches — the legacy default (w1/w2/w3) must be preserved. - assert get_expert_linear_names(_UnknownMoeBlock()) == ["w1", "w2", "w3"] +def test_get_expert_linear_names_raises_when_unmatched(): + # No spec matches and no fused-expert structure — must fail loudly instead of + # guessing another model's naming (the legacy w1/w2/w3 default was removed). + with pytest.raises(NotImplementedError, match="expert linear names"): + get_expert_linear_names(_UnknownMoeBlock()) + + +def test_get_expert_linear_names_from_specs(): + class ArcticMoE(nn.Module): + pass + + class DbrxFFN(nn.Module): + pass + + # Arctic keeps the w1/w2/w3 naming it previously got from the engine default. + assert get_expert_linear_names(ArcticMoE()) == ["w1", "w2", "w3"] + # DbrxFFN resolves the quantized per-expert ModuleList names (previously it fell + # through to the w1/w2/w3 default, which never existed on the quantized module). + assert get_expert_linear_names(DbrxFFN()) == ["w1_linear", "w2_linear", "v1_linear"] class NemotronHMOE(nn.Module): From ed2eb7a5c8c0f1c4d40d6e2d794997d5a27066e5 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 16 Jul 2026 01:44:02 +0000 Subject: [PATCH 08/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/__init__.py | 10 ++-- modelopt/modeling/base.py | 13 +++-- modelopt/modeling/export.py | 29 ++++------ modelopt/modeling/models/arctic.py | 6 +-- modelopt/modeling/models/dbrx.py | 10 ++-- modelopt/modeling/models/deepseek.py | 10 ++-- modelopt/modeling/models/gemma4.py | 6 +-- modelopt/modeling/models/gpt_oss.py | 6 +-- modelopt/modeling/models/mixtral.py | 10 ++-- modelopt/modeling/models/nemotron_h.py | 6 +-- modelopt/modeling/models/qwen2_moe.py | 6 +-- modelopt/modeling/models/qwen3_5_moe.py | 6 +-- modelopt/modeling/models/qwen3_moe.py | 11 +++- modelopt/modeling/models/qwen3_next.py | 6 +-- modelopt/modeling/moe.py | 54 +++++++++++++++++++ modelopt/modeling/registry.py | 9 ++-- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 29 ++++++++-- modelopt/torch/export/layer_utils.py | 2 +- 18 files changed, 156 insertions(+), 73 deletions(-) create mode 100644 modelopt/modeling/moe.py diff --git a/modelopt/modeling/__init__.py b/modelopt/modeling/__init__.py index 503a01955f2..5bec9d0c81f 100644 --- a/modelopt/modeling/__init__.py +++ b/modelopt/modeling/__init__.py @@ -16,14 +16,16 @@ """Per-model descriptors, organized by HF model type. Holds declarative per-model data (no algorithms), one module per HF model type under -``models/``, mirroring ``transformers.models``. Each modelopt subsystem has its own -``ModelSpec`` subclass (e.g. ``ExportSpec``); consumers resolve a spec via the -registry lookups and read its fields; an unmatched lookup returns ``None`` so callers -fall back to their default behavior. +``models/``, mirroring ``transformers.models``. Architecture facts live in topic +specs shared across subsystems (``MoESpec``); per-subsystem policy lives in +subsystem specs (``ExportSpec``). Consumers resolve a spec via the registry lookups +and read its fields; an unmatched lookup returns ``None`` so callers fall back to +their default behavior. """ # Importing models registers every spec as a side effect. from . import models from .base import * from .export import * +from .moe import * from .registry import * diff --git a/modelopt/modeling/base.py b/modelopt/modeling/base.py index e49b1875a6c..5177892fe20 100644 --- a/modelopt/modeling/base.py +++ b/modelopt/modeling/base.py @@ -15,10 +15,15 @@ """Common base for per-model descriptors. -Each modelopt subsystem declares its per-model data in its own ``ModelSpec`` -subclass (e.g. ``ExportSpec``; quantization / speculative-decoding specs to follow), -so consumers can read these values instead of branching on model names. Specs hold -per-model data only, no logic. +``ModelSpec`` subclasses come in two kinds, both registered per model so consumers +read values instead of branching on model names: + +- **topic specs** hold architecture facts shared across subsystems (e.g. ``MoESpec``: + what a model's MoE blocks are); +- **subsystem specs** hold one subsystem's per-model policy (e.g. ``ExportSpec``; + quantization / speculative-decoding specs to follow). + +Specs hold per-model data only, no logic. """ from dataclasses import dataclass diff --git a/modelopt/modeling/export.py b/modelopt/modeling/export.py index 2d34be80f4b..362aeaf22bc 100644 --- a/modelopt/modeling/export.py +++ b/modelopt/modeling/export.py @@ -13,7 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-model data consumed by the unified HF export path.""" +"""Per-model policy of the unified HF export path. + +Architecture facts (MoE block classes, expert naming) live in ``MoESpec``; this spec +holds decisions that belong to the export/quantization algorithms only. +""" from dataclasses import dataclass @@ -24,27 +28,12 @@ @dataclass class ExportSpec(ModelSpec): - """Per-model data for the unified HF export path. - - Resolved from a model sub-module via ``moe_block_names``, the matching key: MoE - block class names (e.g. ``"Qwen3MoeSparseMoeBlock"``) compared case-insensitively - against the class names in the module's MRO (see ``matching.match_class_names``). - """ - - moe_block_names: tuple[str, ...] = () - """Matching key: MoE block class names, matched against the module's MRO - (case-insensitive exact names, not substrings).""" - - expert_linear_names: tuple[str, ...] | None = None - """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``.""" - - has_iterable_experts: bool = False - """True when experts are per-expert iterable sub-modules (Mixtral, Qwen MoE, - NemotronH, Gemma4) and can be grouped by ``get_experts_list``; False for stacked - or fused layouts (DBRX, GptOss).""" + """Per-model policy for the unified HF export path.""" pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = () """AWQ ``pre_quant_scale`` fusion rules, each a ``(module_class_substrings, fuse_into, fuse_from)`` triple: for a module whose class name contains one of the substrings, the pre_quant_scale on ``fuse_from`` is folded into ``fuse_into`` - (e.g. attention ``o_proj`` -> ``v_proj``, MLP ``down_proj`` -> ``up_proj``).""" + (e.g. attention ``o_proj`` -> ``v_proj``, MLP ``down_proj`` -> ``up_proj``). + A rule is a validated mathematical-equivalence claim for that model's modules, + which is why it is declared per model rather than applied generically.""" diff --git a/modelopt/modeling/models/arctic.py b/modelopt/modeling/models/arctic.py index 73bede55bce..32c9cb6e306 100644 --- a/modelopt/modeling/models/arctic.py +++ b/modelopt/modeling/models/arctic.py @@ -15,14 +15,14 @@ """Snowflake Arctic specs (trust-remote-code model type ``arctic``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="arctic", # Non-standard block name (for is_moe identification). - moe_block_names=("ArcticMoE",), + block_names=("ArcticMoE",), # ArcticMLP experts use Mixtral-style w1/w2/w3 naming (previously served by # the engine's implicit w1/w2/w3 default, now declared explicitly). expert_linear_names=("w1", "w2", "w3"), diff --git a/modelopt/modeling/models/dbrx.py b/modelopt/modeling/models/dbrx.py index d8ffc131bb3..f0b6fe07d6f 100644 --- a/modelopt/modeling/models/dbrx.py +++ b/modelopt/modeling/models/dbrx.py @@ -15,13 +15,13 @@ """DBRX specs (HF model type ``dbrx``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="dbrx", - moe_block_names=("DBRXMoeSparseMoeBlock",), + block_names=("DBRXMoeSparseMoeBlock",), expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), ) ) @@ -32,9 +32,9 @@ # modelopt/torch/quantization/plugins/huggingface.py), which is what the DBRX # prepare handler fills amax values on. register( - ExportSpec( + MoESpec( model_type="dbrx", - moe_block_names=("DbrxFFN",), + block_names=("DbrxFFN",), expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), ) ) diff --git a/modelopt/modeling/models/deepseek.py b/modelopt/modeling/models/deepseek.py index eb6bb3dc748..c54507d9e57 100644 --- a/modelopt/modeling/models/deepseek.py +++ b/modelopt/modeling/models/deepseek.py @@ -19,13 +19,17 @@ classes. """ -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register +# DeepseekMoE experts ARE structurally iterable (ModuleList of DeepseekMLP), but +# has_iterable_experts stays False until the grouped export path (get_experts_list +# resmoothing) is validated on this model — the flag currently doubles as that +# support gate. register( - ExportSpec( + MoESpec( model_type="deepseek", - moe_block_names=("DeepseekMoE",), + block_names=("DeepseekMoE",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), ) ) diff --git a/modelopt/modeling/models/gemma4.py b/modelopt/modeling/models/gemma4.py index 55323c7287e..d26cfa9f2a5 100644 --- a/modelopt/modeling/models/gemma4.py +++ b/modelopt/modeling/models/gemma4.py @@ -15,14 +15,14 @@ """Gemma4 specs (HF model type ``gemma4``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="gemma4", # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. - moe_block_names=("Gemma4TextDecoderLayer",), + block_names=("Gemma4TextDecoderLayer",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, ) diff --git a/modelopt/modeling/models/gpt_oss.py b/modelopt/modeling/models/gpt_oss.py index 1c1bd06d504..828d1a0dcec 100644 --- a/modelopt/modeling/models/gpt_oss.py +++ b/modelopt/modeling/models/gpt_oss.py @@ -15,14 +15,14 @@ """GPT-OSS specs (HF model type ``gpt_oss``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="gpt_oss", # GPT-OSS fuses gate and up into a single gate_up_proj. - moe_block_names=("GptOssMoE",), + block_names=("GptOssMoE",), expert_linear_names=("gate_up_proj", "down_proj"), ) ) diff --git a/modelopt/modeling/models/mixtral.py b/modelopt/modeling/models/mixtral.py index 4c069cde53b..457fb4b7c2e 100644 --- a/modelopt/modeling/models/mixtral.py +++ b/modelopt/modeling/models/mixtral.py @@ -15,15 +15,15 @@ """Mixtral specs (HF model type ``mixtral``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register # Mixtral with iterable experts uses w1/w2/w3. Fused experts (transformers 5.0+) are # detected from their per-expert quantizer attributes and need no naming override here. register( - ExportSpec( + MoESpec( model_type="mixtral", - moe_block_names=("MixtralSparseMoeBlock",), + block_names=("MixtralSparseMoeBlock",), expert_linear_names=("w1", "w2", "w3"), has_iterable_experts=True, ) @@ -31,9 +31,9 @@ # Older transformers naming for Mixtral. register( - ExportSpec( + MoESpec( model_type="mixtral", - moe_block_names=("MixtralMoeSparseMoeBlock",), + block_names=("MixtralMoeSparseMoeBlock",), expert_linear_names=("linear_fc1", "linear_fc2"), ) ) diff --git a/modelopt/modeling/models/nemotron_h.py b/modelopt/modeling/models/nemotron_h.py index 96429eedc39..3052e11fd88 100644 --- a/modelopt/modeling/models/nemotron_h.py +++ b/modelopt/modeling/models/nemotron_h.py @@ -15,14 +15,14 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="nemotron_h", # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). - moe_block_names=("NemotronHMOE",), + block_names=("NemotronHMOE",), expert_linear_names=("up_proj", "down_proj"), has_iterable_experts=True, ) diff --git a/modelopt/modeling/models/qwen2_moe.py b/modelopt/modeling/models/qwen2_moe.py index c6b2e391666..442308ba2b4 100644 --- a/modelopt/modeling/models/qwen2_moe.py +++ b/modelopt/modeling/models/qwen2_moe.py @@ -15,13 +15,13 @@ """Qwen2-MoE specs (HF model type ``qwen2_moe``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="qwen2_moe", - moe_block_names=("Qwen2MoeSparseMoeBlock",), + block_names=("Qwen2MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, ) diff --git a/modelopt/modeling/models/qwen3_5_moe.py b/modelopt/modeling/models/qwen3_5_moe.py index f265cc1e32a..90493006a3e 100644 --- a/modelopt/modeling/models/qwen3_5_moe.py +++ b/modelopt/modeling/models/qwen3_5_moe.py @@ -15,13 +15,13 @@ """Qwen3.5-MoE specs (HF model type ``qwen3_5_moe``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="qwen3_5_moe", - moe_block_names=("Qwen3_5MoeSparseMoeBlock",), + block_names=("Qwen3_5MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, ) diff --git a/modelopt/modeling/models/qwen3_moe.py b/modelopt/modeling/models/qwen3_moe.py index 831e3b54b0c..06dff3d8692 100644 --- a/modelopt/modeling/models/qwen3_moe.py +++ b/modelopt/modeling/models/qwen3_moe.py @@ -16,14 +16,21 @@ """Qwen3-MoE specs (HF model type ``qwen3_moe``).""" from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="qwen3_moe", - moe_block_names=("Qwen3MoeSparseMoeBlock",), + block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, + ) +) + +register( + ExportSpec( + model_type="qwen3_moe", # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( (("Qwen3MoeAttention",), "v_proj", "o_proj"), diff --git a/modelopt/modeling/models/qwen3_next.py b/modelopt/modeling/models/qwen3_next.py index bc220954851..9ec5819b149 100644 --- a/modelopt/modeling/models/qwen3_next.py +++ b/modelopt/modeling/models/qwen3_next.py @@ -15,13 +15,13 @@ """Qwen3-Next specs (HF model type ``qwen3_next``).""" -from ..export import ExportSpec +from ..moe import MoESpec from ..registry import register register( - ExportSpec( + MoESpec( model_type="qwen3_next", - moe_block_names=("Qwen3NextSparseMoeBlock",), + block_names=("Qwen3NextSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), has_iterable_experts=True, ) diff --git a/modelopt/modeling/moe.py b/modelopt/modeling/moe.py new file mode 100644 index 00000000000..1aa339deeba --- /dev/null +++ b/modelopt/modeling/moe.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model MoE architecture facts. + +Unlike the subsystem specs (e.g. ``ExportSpec``), a ``MoESpec`` describes what a +model's MoE blocks *are* — which class, what the expert projections are called — +so any modelopt subsystem (export, quantization, speculative decoding, ...) can +read it instead of keeping its own per-model MoE table. +""" + +from dataclasses import dataclass + +from .base import ModelSpec + +__all__ = ["MoESpec"] + + +@dataclass +class MoESpec(ModelSpec): + """MoE architecture facts for one model (or one of its MoE-block variants). + + Resolved from a model sub-module via ``block_names``, the matching key: MoE block + class names (e.g. ``"Qwen3MoeSparseMoeBlock"``) compared case-insensitively + against the class names in the module's MRO (see ``matching.match_class_names``). + """ + + block_names: tuple[str, ...] = () + """Matching key: MoE block class names, matched against the module's MRO + (case-insensitive exact names, not substrings).""" + + expert_linear_names: tuple[str, ...] | None = None + """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``. + For layouts modelopt rewrites (e.g. quantized DBRX), these are the names on the + rewritten module.""" + + has_iterable_experts: bool = False + """True when experts are per-expert iterable sub-modules (Mixtral, Qwen MoE, + NemotronH, Gemma4) and can be grouped by ``get_experts_list``; False for stacked + or fused layouts (DBRX, GptOss). NOTE: currently also doubles as the grouped-export + support gate, so it is conservatively False for structurally iterable but + unvalidated models (see ``deepseek``).""" diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index 47eb01ed8cf..b63ecdf8b48 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -29,6 +29,7 @@ from .base import ModelSpec from .export import ExportSpec from .matching import match_class_names +from .moe import MoESpec if TYPE_CHECKING: import torch.nn as nn @@ -68,14 +69,14 @@ def iter_pqs_fuse_rules(): yield from spec.pqs_fuse_rules -def match_moe_block(module: "nn.Module") -> ExportSpec | None: - """Return the export spec whose ``moe_block_names`` matches ``module``. +def match_moe_block(module: "nn.Module") -> MoESpec | None: + """Return the MoE spec whose ``block_names`` matches ``module``. Case-insensitive exact-name match against the class names in ``module``'s MRO (see ``matching.match_class_names``); quantized wrapper classes match through their original base class. """ - for spec in iter_specs(ExportSpec): - if match_class_names(module, spec.moe_block_names): + for spec in iter_specs(MoESpec): + if match_class_names(module, spec.block_names): return spec return None diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index adc50db351d..0bd2f03f628 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -50,13 +50,16 @@ Two layers: ```text modelopt/modeling/ base.py # ModelSpec — common base class (model_type) - export.py # ExportSpec(ModelSpec) — data for the unified HF export path - # (future: quantization / speculative-decoding spec classes) + moe.py # MoESpec(ModelSpec) — topic spec: MoE architecture facts shared + # across subsystems (future: attention/norm topic specs) + export.py # ExportSpec(ModelSpec) — subsystem spec: HF-export-path policy + # (future: quantization / speculative-decoding subsystem specs) + matching.py # module <-> class-name matching core (MRO exact-name) registry.py # register() + lookups, queried by spec type; None when unmatched __init__.py # re-exports; importing it registers all specs models/ # one small file per HF model type (mirrors transformers.models); - # import == registration; a model registers one spec instance per - # subsystem it customizes + # import == registration; a model registers one instance per spec + # kind it customizes ``` Model type names mirror @@ -90,8 +93,26 @@ against existing export tests. | **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen3 specs). | this PR | | **P3** | `is_moe` explicit class-name list → `spec.moe_block_names` (arctic/dbrx_ffn are identification-only specs: no expert naming, so expert-name lookups keep the engine default). The generic `*SparseMoeBlock`/`*MoeLayer` conventions and the structural router+experts check stay in the engine. | this PR | | **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). | planned | +| **P5** | Cross-subsystem pilot: unify the 3 copies of linear-fusion-group knowledge (see §4) into a spec field; delete `model_calib.py`'s private import of `export._GATE_UP_PAIRS`. `modelopt/modeling` is already top-level and stdlib-only, so quantization can consume it directly — no promotion or re-export shim needed. | planned | +| **P6** | Migrate remaining quantization-side data: default disabled-quantizer patterns, on-the-fly conversion gates, AutoQuantize grouping rules (see §4). | planned | | **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Dead code found during inventory (`MODEL_NAME_TO_TYPE`, `get_model_type`, `adjust_attn_amax_values`, `update_experts_avg_prequant_scale`) is deleted on that track too. | separate track | **Guardrails:** one data category per PR (fallback-first while a category is partially migrated, explicit-error once specs cover it); the engine keeps the algorithms — model specs supply values only, never fork functions. + +## 4. Beyond export: per-model data in quantization (P5/P6 inventory) + +The same three kinds of model-specific logic exist on the quantization side. Only +kind (a) migrates into `modelopt/modeling`; (b) stays in each subsystem's module +registry (a spec may hold pointer data, never the surgery code); (c) stays in the +engine behind structural checks. + +| Item | Location | Kind | +|---|---|---| +| Linear fusion groups (q/k/v, gate/up, `w1/w3`) — **duplicated 3x** | `export/layer_utils._GATE_UP_PAIRS`, `quantization/utils/shared_input.SHARED_PATTERNS`, `quantization/algorithms.quant_grouping_rules` | data — P5 | +| Model-class gates for on-the-fly conversion (`"DbrxForCausalLM"`, `("Step3p5ForCausalLM", ...)`) | `quantization/plugins/huggingface.py` | data — P6 | +| Default disabled-quantizer patterns (`*router*`, `*vision_tower*`; per-model, NVBug-gated) | `modelopt_recipes/.../default_disabled_quantizers.yaml` | data — P6 | +| AutoQuantize grouping regexes (llama q/k/v, Mixtral `w1/w2/w3`, NemotronH mixer) | `quantization/algorithms.py` | data — P6 | +| Quant wrapper classes (`_QuantDbrxExperts` splits `w1/v1/w2` into per-expert linears) | `quantization/plugins/huggingface.py` via `QuantModuleRegistry` | dispatch (stays) | +| Structural MoE detection (`gate`+`experts`+`top_k` attrs; 3-D `gate_up_proj` -> gated) | `quantization/plugins/huggingface.py` | behavior (stays) | diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index a2152f4443e..0bcb53ee14e 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -985,7 +985,7 @@ def get_expert_linear_names(module: nn.Module) -> list[str]: raise NotImplementedError( f"Cannot resolve expert linear names for MoE block {type(module).__name__!r}. " - "Register an ExportSpec with expert_linear_names for this model under " + "Register a MoESpec with expert_linear_names for this model under " "modelopt/modeling/models/." ) From 10d4682f25200d32b674f0138f9a620cb78b51f5 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 16 Jul 2026 02:06:10 +0000 Subject: [PATCH 09/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/__init__.py | 4 +- modelopt/modeling/base.py | 45 ------- modelopt/modeling/export.py | 39 ------ modelopt/modeling/matching.py | 39 ------ modelopt/modeling/models/__init__.py | 2 + modelopt/modeling/models/arctic.py | 3 +- modelopt/modeling/models/dbrx.py | 2 +- modelopt/modeling/models/deepseek.py | 3 +- modelopt/modeling/models/gemma.py | 31 +++++ modelopt/modeling/models/gemma4.py | 3 +- modelopt/modeling/models/gpt_oss.py | 2 +- modelopt/modeling/models/llama.py | 2 +- modelopt/modeling/models/mixtral.py | 4 +- modelopt/modeling/models/nemotron.py | 30 +++++ modelopt/modeling/models/nemotron_h.py | 2 +- modelopt/modeling/models/qwen2_moe.py | 3 +- modelopt/modeling/models/qwen3.py | 2 +- modelopt/modeling/models/qwen3_5_moe.py | 3 +- modelopt/modeling/models/qwen3_moe.py | 4 +- modelopt/modeling/models/qwen3_next.py | 3 +- modelopt/modeling/moe.py | 54 -------- modelopt/modeling/registry.py | 56 +++++++-- modelopt/modeling/specs.py | 116 ++++++++++++++++++ .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 24 ++-- modelopt/torch/export/layer_utils.py | 12 +- modelopt/torch/export/quant_utils.py | 9 +- modelopt/torch/quantization/model_calib.py | 5 +- .../unit/torch/export/test_modeling_specs.py | 38 +++++- 28 files changed, 309 insertions(+), 231 deletions(-) delete mode 100644 modelopt/modeling/base.py delete mode 100644 modelopt/modeling/export.py delete mode 100644 modelopt/modeling/matching.py create mode 100644 modelopt/modeling/models/gemma.py create mode 100644 modelopt/modeling/models/nemotron.py delete mode 100644 modelopt/modeling/moe.py create mode 100644 modelopt/modeling/specs.py diff --git a/modelopt/modeling/__init__.py b/modelopt/modeling/__init__.py index 5bec9d0c81f..e106af55666 100644 --- a/modelopt/modeling/__init__.py +++ b/modelopt/modeling/__init__.py @@ -25,7 +25,5 @@ # Importing models registers every spec as a side effect. from . import models -from .base import * -from .export import * -from .moe import * from .registry import * +from .specs import * diff --git a/modelopt/modeling/base.py b/modelopt/modeling/base.py deleted file mode 100644 index 5177892fe20..00000000000 --- a/modelopt/modeling/base.py +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Common base for per-model descriptors. - -``ModelSpec`` subclasses come in two kinds, both registered per model so consumers -read values instead of branching on model names: - -- **topic specs** hold architecture facts shared across subsystems (e.g. ``MoESpec``: - what a model's MoE blocks are); -- **subsystem specs** hold one subsystem's per-model policy (e.g. ``ExportSpec``; - quantization / speculative-decoding specs to follow). - -Specs hold per-model data only, no logic. -""" - -from dataclasses import dataclass - -__all__ = ["ModelSpec"] - - -@dataclass -class ModelSpec: - """Base class for per-model data specs. - - Subclasses add the data fields of one modelopt subsystem; a model registers one - spec instance per subsystem it customizes (see ``models/``). - """ - - model_type: str - """The HF model type this spec belongs to (``config.model_type``, e.g. - ``"qwen3_moe"``). Not necessarily unique: a model type may register several specs - for different module layouts (e.g. two ``mixtral`` MoE-block variants).""" diff --git a/modelopt/modeling/export.py b/modelopt/modeling/export.py deleted file mode 100644 index 362aeaf22bc..00000000000 --- a/modelopt/modeling/export.py +++ /dev/null @@ -1,39 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Per-model policy of the unified HF export path. - -Architecture facts (MoE block classes, expert naming) live in ``MoESpec``; this spec -holds decisions that belong to the export/quantization algorithms only. -""" - -from dataclasses import dataclass - -from .base import ModelSpec - -__all__ = ["ExportSpec"] - - -@dataclass -class ExportSpec(ModelSpec): - """Per-model policy for the unified HF export path.""" - - pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = () - """AWQ ``pre_quant_scale`` fusion rules, each a ``(module_class_substrings, - fuse_into, fuse_from)`` triple: for a module whose class name contains one of the - substrings, the pre_quant_scale on ``fuse_from`` is folded into ``fuse_into`` - (e.g. attention ``o_proj`` -> ``v_proj``, MLP ``down_proj`` -> ``up_proj``). - A rule is a validated mathematical-equivalence claim for that model's modules, - which is why it is declared per model rather than applied generically.""" diff --git a/modelopt/modeling/matching.py b/modelopt/modeling/matching.py deleted file mode 100644 index 7b47cab03ce..00000000000 --- a/modelopt/modeling/matching.py +++ /dev/null @@ -1,39 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Module-to-name matching shared by the spec registry (and, later, dispatch registries). - -Same semantics as the export dispatch registry's string keys -(``modelopt.torch.export.registry``): a name matches when it equals the ``__name__`` -of a class in ``type(module).__mro__``. Dynamically generated quantized classes are -subclasses of the original module class, so they match through their base; exact-name -comparison avoids substring false positives. Comparison is case-insensitive because -some registered names predate this registry and their casing was never exercised by -the legacy substring matching. - -Only inspects ``type(module).__mro__``, so this module stays stdlib-only. -""" - -__all__ = ["match_class_names"] - - -def match_class_names(module, names: tuple[str, ...]) -> bool: - """Return True if any of ``names`` equals a class name in ``module``'s MRO. - - Case-insensitive exact-name comparison against ``cls.__name__`` for every class - in ``type(module).__mro__``. - """ - mro_names = {cls.__name__.lower() for cls in type(module).__mro__} - return any(name.lower() in mro_names for name in names) diff --git a/modelopt/modeling/models/__init__.py b/modelopt/modeling/models/__init__.py index cc711b81ae9..6b0cb2f9315 100644 --- a/modelopt/modeling/models/__init__.py +++ b/modelopt/modeling/models/__init__.py @@ -24,10 +24,12 @@ arctic, dbrx, deepseek, + gemma, gemma4, gpt_oss, llama, mixtral, + nemotron, nemotron_h, qwen2_moe, qwen3, diff --git a/modelopt/modeling/models/arctic.py b/modelopt/modeling/models/arctic.py index 32c9cb6e306..347b6d2bb42 100644 --- a/modelopt/modeling/models/arctic.py +++ b/modelopt/modeling/models/arctic.py @@ -15,8 +15,8 @@ """Snowflake Arctic specs (trust-remote-code model type ``arctic``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec register( MoESpec( @@ -26,5 +26,6 @@ # ArcticMLP experts use Mixtral-style w1/w2/w3 naming (previously served by # the engine's implicit w1/w2/w3 default, now declared explicitly). expert_linear_names=("w1", "w2", "w3"), + gate_up_pair=("w1", "w3"), ) ) diff --git a/modelopt/modeling/models/dbrx.py b/modelopt/modeling/models/dbrx.py index f0b6fe07d6f..82f84e7c3cd 100644 --- a/modelopt/modeling/models/dbrx.py +++ b/modelopt/modeling/models/dbrx.py @@ -15,8 +15,8 @@ """DBRX specs (HF model type ``dbrx``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec register( MoESpec( diff --git a/modelopt/modeling/models/deepseek.py b/modelopt/modeling/models/deepseek.py index c54507d9e57..ff0624b13a9 100644 --- a/modelopt/modeling/models/deepseek.py +++ b/modelopt/modeling/models/deepseek.py @@ -19,8 +19,8 @@ classes. """ -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec # DeepseekMoE experts ARE structurally iterable (ModuleList of DeepseekMLP), but # has_iterable_experts stays False until the grouped export path (get_experts_list @@ -31,5 +31,6 @@ model_type="deepseek", block_names=("DeepseekMoE",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), ) ) diff --git a/modelopt/modeling/models/gemma.py b/modelopt/modeling/models/gemma.py new file mode 100644 index 00000000000..2abc89932bc --- /dev/null +++ b/modelopt/modeling/models/gemma.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gemma 1/2/3 specs (HF model types ``gemma``/``gemma2``/``gemma3``). + +Grouped in one module because each generation contributes the same single fact; +Gemma4 has its own module (``gemma4.py``) for its MoE spec. +""" + +from ..registry import register +from ..specs import NormSpec + +# Gemma RMSNorms store weight - 1 (the effective scale is weight + 1); scale-folding +# engines (AWQ pre_quant_scale fusion into the norm) must account for the +1. +# NOTE: Gemma4RMSNorm is intentionally absent to preserve legacy behavior; add it +# once the +1 handling is validated on Gemma4. +register(NormSpec(model_type="gemma", weight_plus_one_norm_names=("GemmaRMSNorm",))) +register(NormSpec(model_type="gemma2", weight_plus_one_norm_names=("Gemma2RMSNorm",))) +register(NormSpec(model_type="gemma3", weight_plus_one_norm_names=("Gemma3RMSNorm",))) diff --git a/modelopt/modeling/models/gemma4.py b/modelopt/modeling/models/gemma4.py index d26cfa9f2a5..bab1570437f 100644 --- a/modelopt/modeling/models/gemma4.py +++ b/modelopt/modeling/models/gemma4.py @@ -15,8 +15,8 @@ """Gemma4 specs (HF model type ``gemma4``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec register( MoESpec( @@ -24,6 +24,7 @@ # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. block_names=("Gemma4TextDecoderLayer",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), has_iterable_experts=True, ) ) diff --git a/modelopt/modeling/models/gpt_oss.py b/modelopt/modeling/models/gpt_oss.py index 828d1a0dcec..20f33f476a8 100644 --- a/modelopt/modeling/models/gpt_oss.py +++ b/modelopt/modeling/models/gpt_oss.py @@ -15,8 +15,8 @@ """GPT-OSS specs (HF model type ``gpt_oss``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec register( MoESpec( diff --git a/modelopt/modeling/models/llama.py b/modelopt/modeling/models/llama.py index 4e8da2e6ad4..4048db65737 100644 --- a/modelopt/modeling/models/llama.py +++ b/modelopt/modeling/models/llama.py @@ -15,8 +15,8 @@ """Llama specs (HF model type ``llama``).""" -from ..export import ExportSpec from ..registry import register +from ..specs import ExportSpec register( ExportSpec( diff --git a/modelopt/modeling/models/mixtral.py b/modelopt/modeling/models/mixtral.py index 457fb4b7c2e..bc7aa19592f 100644 --- a/modelopt/modeling/models/mixtral.py +++ b/modelopt/modeling/models/mixtral.py @@ -15,8 +15,8 @@ """Mixtral specs (HF model type ``mixtral``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec # Mixtral with iterable experts uses w1/w2/w3. Fused experts (transformers 5.0+) are # detected from their per-expert quantizer attributes and need no naming override here. @@ -26,6 +26,8 @@ block_names=("MixtralSparseMoeBlock",), expert_linear_names=("w1", "w2", "w3"), has_iterable_experts=True, + # w1 = gate, w3 = up, w2 = down (Mixtral convention). + gate_up_pair=("w1", "w3"), ) ) diff --git a/modelopt/modeling/models/nemotron.py b/modelopt/modeling/models/nemotron.py new file mode 100644 index 00000000000..aba932fa4e9 --- /dev/null +++ b/modelopt/modeling/models/nemotron.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Nemotron specs (HF model type ``nemotron``); Nemotron-H lives in ``nemotron_h.py``.""" + +from ..registry import register +from ..specs import NormSpec + +# LayerNorm1P stores weight - 1 (zero-centered gamma). Both the plain Megatron-style +# class name and the HF Nemotron port are listed; modules exposing a +# ``zero_centered_gamma`` attribute are additionally caught by the engine's +# structural fallback. +register( + NormSpec( + model_type="nemotron", + weight_plus_one_norm_names=("LayerNorm1P", "NemotronLayerNorm1P"), + ) +) diff --git a/modelopt/modeling/models/nemotron_h.py b/modelopt/modeling/models/nemotron_h.py index 3052e11fd88..95697eda89f 100644 --- a/modelopt/modeling/models/nemotron_h.py +++ b/modelopt/modeling/models/nemotron_h.py @@ -15,8 +15,8 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec register( MoESpec( diff --git a/modelopt/modeling/models/qwen2_moe.py b/modelopt/modeling/models/qwen2_moe.py index 442308ba2b4..78b8c9e3abc 100644 --- a/modelopt/modeling/models/qwen2_moe.py +++ b/modelopt/modeling/models/qwen2_moe.py @@ -15,14 +15,15 @@ """Qwen2-MoE specs (HF model type ``qwen2_moe``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec register( MoESpec( model_type="qwen2_moe", block_names=("Qwen2MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), has_iterable_experts=True, ) ) diff --git a/modelopt/modeling/models/qwen3.py b/modelopt/modeling/models/qwen3.py index 95e12af549c..85c2dd6342b 100644 --- a/modelopt/modeling/models/qwen3.py +++ b/modelopt/modeling/models/qwen3.py @@ -15,8 +15,8 @@ """Qwen3 (dense) specs (HF model type ``qwen3``).""" -from ..export import ExportSpec from ..registry import register +from ..specs import ExportSpec register( ExportSpec( diff --git a/modelopt/modeling/models/qwen3_5_moe.py b/modelopt/modeling/models/qwen3_5_moe.py index 90493006a3e..7eaadfa580d 100644 --- a/modelopt/modeling/models/qwen3_5_moe.py +++ b/modelopt/modeling/models/qwen3_5_moe.py @@ -15,14 +15,15 @@ """Qwen3.5-MoE specs (HF model type ``qwen3_5_moe``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec register( MoESpec( model_type="qwen3_5_moe", block_names=("Qwen3_5MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), has_iterable_experts=True, ) ) diff --git a/modelopt/modeling/models/qwen3_moe.py b/modelopt/modeling/models/qwen3_moe.py index 06dff3d8692..d5ac4dd4019 100644 --- a/modelopt/modeling/models/qwen3_moe.py +++ b/modelopt/modeling/models/qwen3_moe.py @@ -15,15 +15,15 @@ """Qwen3-MoE specs (HF model type ``qwen3_moe``).""" -from ..export import ExportSpec -from ..moe import MoESpec from ..registry import register +from ..specs import ExportSpec, MoESpec register( MoESpec( model_type="qwen3_moe", block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), has_iterable_experts=True, ) ) diff --git a/modelopt/modeling/models/qwen3_next.py b/modelopt/modeling/models/qwen3_next.py index 9ec5819b149..f137ae64a6c 100644 --- a/modelopt/modeling/models/qwen3_next.py +++ b/modelopt/modeling/models/qwen3_next.py @@ -15,14 +15,15 @@ """Qwen3-Next specs (HF model type ``qwen3_next``).""" -from ..moe import MoESpec from ..registry import register +from ..specs import MoESpec register( MoESpec( model_type="qwen3_next", block_names=("Qwen3NextSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), has_iterable_experts=True, ) ) diff --git a/modelopt/modeling/moe.py b/modelopt/modeling/moe.py deleted file mode 100644 index 1aa339deeba..00000000000 --- a/modelopt/modeling/moe.py +++ /dev/null @@ -1,54 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Per-model MoE architecture facts. - -Unlike the subsystem specs (e.g. ``ExportSpec``), a ``MoESpec`` describes what a -model's MoE blocks *are* — which class, what the expert projections are called — -so any modelopt subsystem (export, quantization, speculative decoding, ...) can -read it instead of keeping its own per-model MoE table. -""" - -from dataclasses import dataclass - -from .base import ModelSpec - -__all__ = ["MoESpec"] - - -@dataclass -class MoESpec(ModelSpec): - """MoE architecture facts for one model (or one of its MoE-block variants). - - Resolved from a model sub-module via ``block_names``, the matching key: MoE block - class names (e.g. ``"Qwen3MoeSparseMoeBlock"``) compared case-insensitively - against the class names in the module's MRO (see ``matching.match_class_names``). - """ - - block_names: tuple[str, ...] = () - """Matching key: MoE block class names, matched against the module's MRO - (case-insensitive exact names, not substrings).""" - - expert_linear_names: tuple[str, ...] | None = None - """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``. - For layouts modelopt rewrites (e.g. quantized DBRX), these are the names on the - rewritten module.""" - - has_iterable_experts: bool = False - """True when experts are per-expert iterable sub-modules (Mixtral, Qwen MoE, - NemotronH, Gemma4) and can be grouped by ``get_experts_list``; False for stacked - or fused layouts (DBRX, GptOss). NOTE: currently also doubles as the grouped-export - support gate, so it is conservatively False for structurally iterable but - unvalidated models (see ``deepseek``).""" diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index b63ecdf8b48..ea686d2181b 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -13,32 +13,32 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Registry of per-model specs, queried by spec type. +"""Registry that resolves a model sub-module to its spec. Model modules register their specs at import time (see ``models/``). Lookups return ``None`` when nothing matches, so callers can fall back to their default behavior. -Matching is by class-name string only (see ``matching``), so this package stays -dependency-free (any ``nn.Module`` — or any object — can be passed to the lookups -without importing torch here). +Matching is by class-name string only, so this package stays dependency-free (any +``nn.Module`` — or any object — can be passed to the lookups without importing torch +here). """ from collections.abc import Iterator from typing import TYPE_CHECKING, TypeVar -from .base import ModelSpec -from .export import ExportSpec -from .matching import match_class_names -from .moe import MoESpec +from .specs import ExportSpec, ModelSpec, MoESpec, NormSpec if TYPE_CHECKING: import torch.nn as nn __all__ = [ + "iter_gate_up_pairs", "iter_pqs_fuse_rules", "iter_specs", + "match_class_names", "match_moe_block", "register", + "weight_plus_one_norm_names", ] SpecT = TypeVar("SpecT", bound=ModelSpec) @@ -59,6 +59,22 @@ def iter_specs(spec_cls: type[SpecT]) -> Iterator[SpecT]: yield spec +def match_class_names(module, names: tuple[str, ...]) -> bool: + """Return True if any of ``names`` equals a class name in ``module``'s MRO. + + Case-insensitive exact-name comparison against ``cls.__name__`` for every class + in ``type(module).__mro__`` — the same semantics as the export dispatch + registry's string keys (``modelopt.torch.export.registry``). Dynamically + generated quantized classes are subclasses of the original module class, so they + match through their base; exact-name comparison avoids substring false + positives. Comparison is case-insensitive because some registered names predate + this registry and their casing was never exercised by the legacy substring + matching. + """ + mro_names = {cls.__name__.lower() for cls in type(module).__mro__} + return any(name.lower() in mro_names for name in names) + + def iter_pqs_fuse_rules(): """Yield every ``(module_class_substrings, fuse_into, fuse_from)`` AWQ fusion rule. @@ -69,12 +85,32 @@ def iter_pqs_fuse_rules(): yield from spec.pqs_fuse_rules +def iter_gate_up_pairs() -> Iterator[tuple[str, str]]: + """Yield the distinct (gate, up) expert-projection pairs across all MoE specs. + + Deduplicated because consumers apply every pair opportunistically to every + iterable-experts module (getattr-guarded), matching the legacy engine behavior — + unknown models still benefit if their naming matches any registered pair. + """ + seen = set() + for spec in iter_specs(MoESpec): + pair = spec.gate_up_pair + if pair is not None and pair not in seen: + seen.add(pair) + yield pair + + +def weight_plus_one_norm_names() -> tuple[str, ...]: + """All norm class names whose stored weight is ``w - 1``, across all norm specs.""" + return tuple(name for spec in iter_specs(NormSpec) for name in spec.weight_plus_one_norm_names) + + def match_moe_block(module: "nn.Module") -> MoESpec | None: """Return the MoE spec whose ``block_names`` matches ``module``. Case-insensitive exact-name match against the class names in ``module``'s MRO - (see ``matching.match_class_names``); quantized wrapper classes match through - their original base class. + (see ``match_class_names``); quantized wrapper classes match through their + original base class. """ for spec in iter_specs(MoESpec): if match_class_names(module, spec.block_names): diff --git a/modelopt/modeling/specs.py b/modelopt/modeling/specs.py new file mode 100644 index 00000000000..e0d2027d8c9 --- /dev/null +++ b/modelopt/modeling/specs.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-model descriptor classes. + +``ModelSpec`` subclasses come in two kinds, both registered per model so consumers +read values instead of branching on model names: + +- **topic specs** hold architecture facts shared across subsystems (``MoESpec``: + what a model's MoE blocks are; ``NormSpec``: norm-layer conventions); +- **subsystem specs** hold one subsystem's per-model policy (``ExportSpec``; + quantization / speculative-decoding specs to follow). + +Specs hold per-model data only, no logic. +""" + +from dataclasses import dataclass + +__all__ = ["ExportSpec", "MoESpec", "ModelSpec", "NormSpec"] + + +@dataclass +class ModelSpec: + """Base class for per-model data specs. + + Subclasses add the data fields of one topic or subsystem; a model registers one + spec instance per kind it customizes (see ``models/``). + """ + + model_type: str + """The HF model type this spec belongs to (``config.model_type``, e.g. + ``"qwen3_moe"``). Not necessarily unique: a model type may register several specs + for different module layouts (e.g. two ``mixtral`` MoE-block variants).""" + + +@dataclass +class MoESpec(ModelSpec): + """MoE architecture facts for one model (or one of its MoE-block variants). + + Unlike the subsystem specs, this describes what a model's MoE blocks *are* — + which class, what the expert projections are called — so any modelopt subsystem + (export, quantization, speculative decoding, ...) can read it instead of keeping + its own per-model MoE table. + + Resolved from a model sub-module via ``block_names``, the matching key: MoE block + class names (e.g. ``"Qwen3MoeSparseMoeBlock"``) compared case-insensitively + against the class names in the module's MRO (see ``registry.match_class_names``). + """ + + block_names: tuple[str, ...] = () + """Matching key: MoE block class names, matched against the module's MRO + (case-insensitive exact names, not substrings).""" + + expert_linear_names: tuple[str, ...] | None = None + """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``. + For layouts modelopt rewrites (e.g. quantized DBRX), these are the names on the + rewritten module.""" + + has_iterable_experts: bool = False + """True when experts are per-expert iterable sub-modules (Mixtral, Qwen MoE, + NemotronH, Gemma4) and can be grouped by ``get_experts_list``; False for stacked + or fused layouts (DBRX, GptOss). NOTE: currently also doubles as the grouped-export + support gate, so it is conservatively False for structurally iterable but + unvalidated models (see ``deepseek``).""" + + gate_up_pair: tuple[str, str] | None = None + """The (gate, up) pair among ``expert_linear_names`` that serving engines fuse + into a single ``gate_up_proj``, e.g. ``("gate_proj", "up_proj")`` or + ``("w1", "w3")``. ``None`` for non-gated experts (NemotronH) and already-fused + layouts (GptOss, DBRX). Consumed by amax syncing before quantized export (see + ``sync_moe_gate_up_amax``) and by calibration grouping.""" + + +@dataclass +class NormSpec(ModelSpec): + """Normalization-layer architecture facts for one model. + + Topic spec (like ``MoESpec``): shared facts any subsystem can read. + """ + + weight_plus_one_norm_names: tuple[str, ...] = () + """Class names of norm layers whose stored weight is ``w - 1`` (the effective + scale is ``weight + 1``), e.g. Gemma's RMSNorm variants and LayerNorm1P. + Matched against a norm module's MRO (case-insensitive exact names). Engines + must account for the +1 when folding scales into the norm weight (AWQ + pre_quant_scale fusion). A structural fallback (``zero_centered_gamma``) stays + in the engine.""" + + +@dataclass +class ExportSpec(ModelSpec): + """Per-model policy for the unified HF export path. + + Architecture facts (MoE block classes, expert naming) live in ``MoESpec``; this + spec holds decisions that belong to the export/quantization algorithms only. + """ + + pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = () + """AWQ ``pre_quant_scale`` fusion rules, each a ``(module_class_substrings, + fuse_into, fuse_from)`` triple: for a module whose class name contains one of the + substrings, the pre_quant_scale on ``fuse_from`` is folded into ``fuse_into`` + (e.g. attention ``o_proj`` -> ``v_proj``, MLP ``down_proj`` -> ``up_proj``). + A rule is a validated mathematical-equivalence claim for that model's modules, + which is why it is declared per model rather than applied generically.""" diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index 0bd2f03f628..7e6a7791df9 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -30,7 +30,10 @@ helpers) consume. Inventory of model-specific logic reachable from the HF path: | Duplicate expert-naming table + iterable-experts support gate | `layer_utils.get_experts_list` | data — **migrated (P1)** | | MoE block class-name list | `layer_utils.is_moe` | data — **migrated (P3)** | | AWQ `pre_quant_scale` fusion rules (Llama/Qwen3) | `quant_utils.PQS_FUSE_MODULE_MAPPING` | data — **migrated (P2)** | -| weight+1 layernorm class names (Gemma RMSNorm, LayerNorm1P) | `quant_utils._layernorm_uses_weight_plus_one` | data | +| weight+1 layernorm class names (Gemma RMSNorm, LayerNorm1P) | `quant_utils._layernorm_uses_weight_plus_one` | data — **migrated (NormSpec)** | +| MoE gate/up fusion pairs (`_GATE_UP_PAIRS`, also privately imported by `quantization/model_calib.py`) | `layer_utils.sync_moe_gate_up_amax` | data — **migrated (MoESpec.gate_up_pair)** | +| BMM-style expert class list (`Llama4TextExperts`/`GptOssExperts`) inline copy for weight transpose | `quant_utils` (dispatch copies in `hf_export_handlers.py` stay) | data — P4 | +| VLM detection gates (`phi4mm` model_type, `nemotronparse` architecture, `"nemotron" in model_type` tower special case) | `model_utils.py`, `unified_export_hf.py` | data (gates) + behavior (extraction) — collect until worth a spec flag | | Handler match keys (Llama4TextExperts, GptOssExperts, DbrxExperts, QuantMoELinear) | `hf_export_handlers.py` | dispatch (stays: structural, per-module) | | Fused-expert gated/non-gated split (`gate_up_proj` vs `up_proj`) | `moe_utils.py` | data + structure | | dummy-forward special cases (Whisper input, Nemotron-VL tower) | `unified_export_hf.requantize_resmooth_fused_llm_layers` | behavior | @@ -49,13 +52,12 @@ Two layers: ```text modelopt/modeling/ - base.py # ModelSpec — common base class (model_type) - moe.py # MoESpec(ModelSpec) — topic spec: MoE architecture facts shared - # across subsystems (future: attention/norm topic specs) - export.py # ExportSpec(ModelSpec) — subsystem spec: HF-export-path policy - # (future: quantization / speculative-decoding subsystem specs) - matching.py # module <-> class-name matching core (MRO exact-name) - registry.py # register() + lookups, queried by spec type; None when unmatched + specs.py # ModelSpec (base) + topic specs (MoESpec: shared architecture + # facts) + subsystem specs (ExportSpec: HF-export-path policy); + # future attention/norm topic specs and quantization/spec-dec + # subsystem specs go here too + registry.py # register() + lookups queried by spec type (None when unmatched) + # + the MRO exact-name matching core (match_class_names) __init__.py # re-exports; importing it registers all specs models/ # one small file per HF model type (mirrors transformers.models); # import == registration; a model registers one instance per spec @@ -93,9 +95,9 @@ against existing export tests. | **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen3 specs). | this PR | | **P3** | `is_moe` explicit class-name list → `spec.moe_block_names` (arctic/dbrx_ffn are identification-only specs: no expert naming, so expert-name lookups keep the engine default). The generic `*SparseMoeBlock`/`*MoeLayer` conventions and the structural router+experts check stay in the engine. | this PR | | **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). | planned | -| **P5** | Cross-subsystem pilot: unify the 3 copies of linear-fusion-group knowledge (see §4) into a spec field; delete `model_calib.py`'s private import of `export._GATE_UP_PAIRS`. `modelopt/modeling` is already top-level and stdlib-only, so quantization can consume it directly — no promotion or re-export shim needed. | planned | +| **P5** | Cross-subsystem pilot: unify the remaining copies of linear-fusion-group knowledge (see §4) into spec fields. Partially done: `_GATE_UP_PAIRS` became `MoESpec.gate_up_pair` and `model_calib.py`'s private import of it is gone — the first quantization consumer of `modelopt.modeling`. Remaining: `shared_input.SHARED_PATTERNS`, `algorithms.quant_grouping_rules`. | in progress | | **P6** | Migrate remaining quantization-side data: default disabled-quantizer patterns, on-the-fly conversion gates, AutoQuantize grouping rules (see §4). | planned | -| **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Dead code found during inventory (`MODEL_NAME_TO_TYPE`, `get_model_type`, `adjust_attn_amax_values`, `update_experts_avg_prequant_scale`) is deleted on that track too. | separate track | +| **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Candidates for deletion on that track: `adjust_attn_amax_values`, `update_experts_avg_prequant_scale` (unused). NOTE: `MODEL_NAME_TO_TYPE` / `get_model_type` are NOT dead — `examples/hf_ptq/hf_ptq.py` and `multinode_ptq.py` still call them; migrate the examples before removing. | separate track | **Guardrails:** one data category per PR (fallback-first while a category is partially migrated, explicit-error once specs cover it); the engine keeps the @@ -110,7 +112,7 @@ engine behind structural checks. | Item | Location | Kind | |---|---|---| -| Linear fusion groups (q/k/v, gate/up, `w1/w3`) — **duplicated 3x** | `export/layer_utils._GATE_UP_PAIRS`, `quantization/utils/shared_input.SHARED_PATTERNS`, `quantization/algorithms.quant_grouping_rules` | data — P5 | +| Linear fusion groups (q/k/v, gate/up, `w1/w3`) — was duplicated 3x; the `_GATE_UP_PAIRS` copy is now `MoESpec.gate_up_pair` | `quantization/utils/shared_input.SHARED_PATTERNS`, `quantization/algorithms.quant_grouping_rules` (remaining) | data — P5 | | Model-class gates for on-the-fly conversion (`"DbrxForCausalLM"`, `("Step3p5ForCausalLM", ...)`) | `quantization/plugins/huggingface.py` | data — P6 | | Default disabled-quantizer patterns (`*router*`, `*vision_tower*`; per-model, NVBug-gated) | `modelopt_recipes/.../default_disabled_quantizers.yaml` | data — P6 | | AutoQuantize grouping regexes (llama q/k/v, Mixtral `w1/w2/w3`, NemotronH mixer) | `quantization/algorithms.py` | data — P6 | diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index 0bcb53ee14e..6ea87624c11 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -28,7 +28,7 @@ except Exception: warn("Cannot find transformers package. Hugginface modules cannot be exported.") -from modelopt.modeling import match_moe_block +from modelopt.modeling import iter_gate_up_pairs, match_moe_block from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import import_plugin @@ -1155,11 +1155,6 @@ def set_expert_quantizer_amax( return uncalibrated_modules -# Gate/up naming pairs for standard (unfused) MoE architectures. -# Fused variants (gate_up_proj, linear_fc1) already share a single quantizer and need no sync. -_GATE_UP_PAIRS = [("gate_proj", "up_proj"), ("w1", "w3")] - - def sync_moe_gate_up_amax(model: nn.Module) -> int: """Take element-wise max of gate and up weight quantizer amaxes per expert. @@ -1176,6 +1171,9 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: Returns: Number of expert gate/up pairs whose amaxes were synced. """ + # Gate/up naming pairs are per-model data (MoESpec.gate_up_pair); fused variants + # (gate_up_proj, linear_fc1) already share a single quantizer and need no sync. + gate_up_pairs = list(iter_gate_up_pairs()) synced = 0 for _, sub_module in model.named_modules(): if not (is_moe(sub_module) and hasattr(sub_module, "experts")): @@ -1183,7 +1181,7 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: if not hasattr(sub_module.experts, "__iter__"): continue for expert in sub_module.experts: - for gate_name, up_name in _GATE_UP_PAIRS: + for gate_name, up_name in gate_up_pairs: gate_linear = getattr(expert, gate_name, None) up_linear = getattr(expert, up_name, None) if gate_linear is None or up_linear is None: diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 7a7e2dc3dbc..4b04cd2ad92 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -25,7 +25,7 @@ import torch.nn as nn from modelopt import __version__ -from modelopt.modeling import iter_pqs_fuse_rules +from modelopt.modeling import iter_pqs_fuse_rules, match_class_names, weight_plus_one_norm_names from modelopt.torch.quantization.model_calib import ( enable_stats_collection, finish_stats_collection, @@ -1231,10 +1231,9 @@ def fuse_prequant_to_linear(model: torch.nn.Module, fuse_grouped_heads=False): def _layernorm_uses_weight_plus_one(module: torch.nn.Module) -> bool: - if any( - name in type(module).__name__ - for name in ["LayerNorm1P", "GemmaRMSNorm", "Gemma2RMSNorm", "Gemma3RMSNorm"] - ): + # Weight-plus-one norm class names are per-model data (NormSpec); the + # zero_centered_gamma attribute check is the structural fallback. + if match_class_names(module, weight_plus_one_norm_names()): return True return bool(hasattr(module, "zero_centered_gamma") and module.zero_centered_gamma) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index e03dff2e98b..38a99c3a4a2 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -87,10 +87,9 @@ def _is_calibrated_nvfp4_static(q) -> bool: def _collect_grouped_linears(model: nn.Module) -> list[list[nn.Module]]: """Collect name-based sibling groups (Q/K/V, gate/up, w1/w3) of calibrated NVFP4-static linears.""" - # Inline import: layer_utils -> quant_utils -> model_calib cycle. - from modelopt.torch.export.layer_utils import _GATE_UP_PAIRS + from modelopt.modeling import iter_gate_up_pairs - patterns: tuple[tuple[str, ...], ...] = (("q_proj", "k_proj", "v_proj"), *_GATE_UP_PAIRS) + patterns: tuple[tuple[str, ...], ...] = (("q_proj", "k_proj", "v_proj"), *iter_gate_up_pairs()) groups: list[list[nn.Module]] = [] for parent in model.modules(): for sibling_names in patterns: diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index 5c0f8f9cf67..43a07a19714 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -18,8 +18,14 @@ import pytest import torch.nn as nn -from modelopt.modeling import iter_pqs_fuse_rules, match_moe_block +from modelopt.modeling import ( + iter_gate_up_pairs, + iter_pqs_fuse_rules, + match_moe_block, + weight_plus_one_norm_names, +) from modelopt.torch.export.layer_utils import get_expert_linear_names, get_experts_list, is_moe +from modelopt.torch.export.quant_utils import _layernorm_uses_weight_plus_one class Qwen3MoeSparseMoeBlock(nn.Module): @@ -156,3 +162,33 @@ def test_pqs_fuse_rules_match_legacy_mapping(): ("Qwen3MoeMLP", "up_proj", "down_proj"), } assert rules == legacy + + +def test_gate_up_pairs_match_legacy(): + # Aggregated per-model pairs must reproduce the legacy _GATE_UP_PAIRS set. + assert set(iter_gate_up_pairs()) == {("gate_proj", "up_proj"), ("w1", "w3")} + + +def test_weight_plus_one_norm_names_cover_legacy(): + names = set(weight_plus_one_norm_names()) + assert {"GemmaRMSNorm", "Gemma2RMSNorm", "Gemma3RMSNorm", "LayerNorm1P"} <= names + + +def test_layernorm_weight_plus_one_via_specs(): + class GemmaRMSNorm(nn.Module): + pass + + class NemotronLayerNorm1P(nn.Module): + pass + + class PlainRMSNorm(nn.Module): + pass + + class ZeroCentered(nn.Module): + zero_centered_gamma = True + + assert _layernorm_uses_weight_plus_one(GemmaRMSNorm()) + assert _layernorm_uses_weight_plus_one(NemotronLayerNorm1P()) + assert not _layernorm_uses_weight_plus_one(PlainRMSNorm()) + # Structural fallback stays in the engine. + assert _layernorm_uses_weight_plus_one(ZeroCentered()) From 1bb5139af08238cc742e8b28759c07ee561182b0 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 16 Jul 2026 02:25:21 +0000 Subject: [PATCH 10/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/registry.py | 15 ++-- modelopt/torch/export/layer_utils.py | 85 +++++++++++-------- modelopt/torch/quantization/model_calib.py | 3 +- .../unit/torch/export/test_modeling_specs.py | 55 +++++++++++- 4 files changed, 115 insertions(+), 43 deletions(-) diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index ea686d2181b..75d8dfacbab 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -86,11 +86,16 @@ def iter_pqs_fuse_rules(): def iter_gate_up_pairs() -> Iterator[tuple[str, str]]: - """Yield the distinct (gate, up) expert-projection pairs across all MoE specs. - - Deduplicated because consumers apply every pair opportunistically to every - iterable-experts module (getattr-guarded), matching the legacy engine behavior — - unknown models still benefit if their naming matches any registered pair. + """Yield the distinct (gate, up) projection-name pairs across all MoE specs. + + GLOBAL-VOCABULARY semantics: consumers (currently only calibration sibling + grouping in ``quantization/model_calib.py``, which also walks dense MLPs that + no MoE spec can match) try every pair opportunistically on every module, + getattr-guarded. Adding a pair to any spec therefore changes behavior for ALL + models whose modules happen to carry those attribute names — prefer per-module + resolution (``match_moe_block(module).gate_up_pair``) wherever the module is an + identifiable MoE block. The dense-MLP case moves to a fusion-group topic spec + in a follow-up (see MODEL_SPECIFIC_REFACTOR.md P5). """ seen = set() for spec in iter_specs(MoESpec): diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index 6ea87624c11..d70348235f2 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -28,7 +28,7 @@ except Exception: warn("Cannot find transformers package. Hugginface modules cannot be exported.") -from modelopt.modeling import iter_gate_up_pairs, match_moe_block +from modelopt.modeling import match_moe_block from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import import_plugin @@ -1171,47 +1171,62 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: Returns: Number of expert gate/up pairs whose amaxes were synced. """ - # Gate/up naming pairs are per-model data (MoESpec.gate_up_pair); fused variants - # (gate_up_proj, linear_fc1) already share a single quantizer and need no sync. - gate_up_pairs = list(iter_gate_up_pairs()) synced = 0 + unmatched_block_names: set[str] = set() for _, sub_module in model.named_modules(): if not (is_moe(sub_module) and hasattr(sub_module, "experts")): continue if not hasattr(sub_module.experts, "__iter__"): continue + # The gate/up pair is this model's own data (MoESpec.gate_up_pair) — no + # cross-model guessing. A spec declaring no pair (non-gated NemotronH, + # fused GptOss/DBRX) needs no sync; an unmatched block is warned about + # once instead of silently skipped. + spec = match_moe_block(sub_module) + if spec is None: + unmatched_block_names.add(type(sub_module).__name__) + continue + if spec.gate_up_pair is None: + continue + gate_name, up_name = spec.gate_up_pair for expert in sub_module.experts: - for gate_name, up_name in gate_up_pairs: - gate_linear = getattr(expert, gate_name, None) - up_linear = getattr(expert, up_name, None) - if gate_linear is None or up_linear is None: - continue - gate_wq = getattr(gate_linear, "weight_quantizer", None) - up_wq = getattr(up_linear, "weight_quantizer", None) - if gate_wq is None or up_wq is None: - break - gate_amax = getattr(gate_wq, "amax", None) - up_amax = getattr(up_wq, "amax", None) - if gate_amax is None or up_amax is None: - break - # Meta tensors have no storage (e.g. CPU-offloaded experts that - # were never activated during calibration). Skip — there is no - # real amax data to sync. - if gate_amax.is_meta or up_amax.is_meta: - warn( - f"Skipping gate/up amax sync for expert with meta tensors " - f"(gate_amax.is_meta={gate_amax.is_meta}, " - f"up_amax.is_meta={up_amax.is_meta}). " - f"This typically means the expert was CPU-offloaded and " - f"not activated during calibration." - ) - break - if not torch.equal(gate_amax, up_amax): - shared_amax = torch.max(gate_amax, up_amax) - gate_wq.amax = shared_amax - up_wq.amax = shared_amax.clone() - synced += 1 - break + gate_linear = getattr(expert, gate_name, None) + up_linear = getattr(expert, up_name, None) + if gate_linear is None or up_linear is None: + continue + gate_wq = getattr(gate_linear, "weight_quantizer", None) + up_wq = getattr(up_linear, "weight_quantizer", None) + if gate_wq is None or up_wq is None: + continue + gate_amax = getattr(gate_wq, "amax", None) + up_amax = getattr(up_wq, "amax", None) + if gate_amax is None or up_amax is None: + continue + # Meta tensors have no storage (e.g. CPU-offloaded experts that + # were never activated during calibration). Skip — there is no + # real amax data to sync. + if gate_amax.is_meta or up_amax.is_meta: + warn( + f"Skipping gate/up amax sync for expert with meta tensors " + f"(gate_amax.is_meta={gate_amax.is_meta}, " + f"up_amax.is_meta={up_amax.is_meta}). " + f"This typically means the expert was CPU-offloaded and " + f"not activated during calibration." + ) + continue + if not torch.equal(gate_amax, up_amax): + shared_amax = torch.max(gate_amax, up_amax) + gate_wq.amax = shared_amax + up_wq.amax = shared_amax.clone() + synced += 1 + if unmatched_block_names: + warn( + f"MoE blocks {sorted(unmatched_block_names)} have no registered MoESpec; " + "gate/up weight amax sync was skipped for them. If these models have " + "gated experts with separate gate/up projections, register a MoESpec " + "with gate_up_pair under modelopt/modeling/models/ so the fused " + "gate_up_proj weight scales stay consistent when serving." + ) return synced diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 38a99c3a4a2..ac721ea8674 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -29,6 +29,7 @@ import torch.nn.functional as F from tqdm import tqdm +from modelopt.modeling import iter_gate_up_pairs from modelopt.torch.opt.config import ModeloptBaseConfig from modelopt.torch.opt.searcher import ForwardLoop from modelopt.torch.quantization.utils.layerwise_calib import ( @@ -87,8 +88,6 @@ def _is_calibrated_nvfp4_static(q) -> bool: def _collect_grouped_linears(model: nn.Module) -> list[list[nn.Module]]: """Collect name-based sibling groups (Q/K/V, gate/up, w1/w3) of calibrated NVFP4-static linears.""" - from modelopt.modeling import iter_gate_up_pairs - patterns: tuple[tuple[str, ...], ...] = (("q_proj", "k_proj", "v_proj"), *iter_gate_up_pairs()) groups: list[list[nn.Module]] = [] for parent in model.modules(): diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index 43a07a19714..204866b9372 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -24,7 +24,12 @@ match_moe_block, weight_plus_one_norm_names, ) -from modelopt.torch.export.layer_utils import get_expert_linear_names, get_experts_list, is_moe +from modelopt.torch.export.layer_utils import ( + get_expert_linear_names, + get_experts_list, + is_moe, + sync_moe_gate_up_amax, +) from modelopt.torch.export.quant_utils import _layernorm_uses_weight_plus_one @@ -192,3 +197,51 @@ class ZeroCentered(nn.Module): assert not _layernorm_uses_weight_plus_one(PlainRMSNorm()) # Structural fallback stays in the engine. assert _layernorm_uses_weight_plus_one(ZeroCentered()) + + +class _FakeQuantizer: + def __init__(self, amax): + self.amax = amax + + +def _make_gated_block(block_cls, gate_name, up_name, gate_amax, up_amax): + import torch + + class _Expert(nn.Module): + def __init__(self): + super().__init__() + setattr(self, gate_name, nn.Linear(4, 8)) + setattr(self, up_name, nn.Linear(4, 8)) + getattr(self, gate_name).weight_quantizer = _FakeQuantizer(torch.tensor(gate_amax)) + getattr(self, up_name).weight_quantizer = _FakeQuantizer(torch.tensor(up_amax)) + + block = block_cls() + block.experts = nn.ModuleList([_Expert()]) + return block + + +def test_sync_moe_gate_up_amax_uses_own_spec(): + import torch + + class Qwen3MoeSparseMoeBlock(nn.Module): + pass + + model = nn.Module() + model.moe = _make_gated_block( + Qwen3MoeSparseMoeBlock, "gate_proj", "up_proj", [1.0, 3.0], [2.0, 2.0] + ) + assert sync_moe_gate_up_amax(model) == 1 + expert = model.moe.experts[0] + assert torch.equal(expert.gate_proj.weight_quantizer.amax, torch.tensor([2.0, 3.0])) + assert torch.equal(expert.up_proj.weight_quantizer.amax, torch.tensor([2.0, 3.0])) + + +def test_sync_moe_gate_up_amax_warns_on_unmatched_block(): + class UnknownSparseMoeBlock(nn.Module): + """Passes is_moe by naming convention but has no MoESpec.""" + + model = nn.Module() + model.moe = _make_gated_block(UnknownSparseMoeBlock, "gate_proj", "up_proj", [1.0], [2.0]) + # No spec -> no cross-model guessing: nothing synced, one warning. + with pytest.warns(UserWarning, match="no registered MoESpec"): + assert sync_moe_gate_up_amax(model) == 0 From fae15c9a21e536f20df84dad7197436c4583e24f Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:41:14 +0000 Subject: [PATCH 11/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/registry.py | 52 ++++++++++++++++--- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 11 +++- modelopt/torch/export/hf_export_handlers.py | 4 +- modelopt/torch/export/layer_utils.py | 28 ++++++---- .../torch/export/plugins/vllm_fakequant_hf.py | 6 ++- modelopt/torch/export/registry.py | 3 ++ modelopt/torch/export/unified_export_hf.py | 23 ++++++-- .../unit/torch/export/test_modeling_specs.py | 44 ++++++++++++++++ 8 files changed, 145 insertions(+), 26 deletions(-) diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index 75d8dfacbab..72114c85ee2 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -32,6 +32,7 @@ import torch.nn as nn __all__ = [ + "collect_model_types", "iter_gate_up_pairs", "iter_pqs_fuse_rules", "iter_specs", @@ -110,14 +111,53 @@ def weight_plus_one_norm_names() -> tuple[str, ...]: return tuple(name for spec in iter_specs(NormSpec) for name in spec.weight_plus_one_norm_names) -def match_moe_block(module: "nn.Module") -> MoESpec | None: +def collect_model_types(config) -> set[str]: + """Collect every HF model type in a config tree (root plus nested sub-configs). + + Walks attribute values duck-typed as configs (anything exposing a string + ``model_type``), so composite models contribute their tower types too — e.g. a + VLM yields ``{"kimi_vl", "kimi_k2"}`` via ``config.text_config`` — without this + package importing transformers. Pass ``model.config``; ``None`` yields an empty + set. + """ + found: set[str] = set() + seen: set[int] = set() + + def _walk(cfg) -> None: + if id(cfg) in seen: + return + seen.add(id(cfg)) + model_type = getattr(cfg, "model_type", None) + if isinstance(model_type, str) and model_type: + found.add(model_type) + for value in vars(cfg).values(): + if isinstance(getattr(value, "model_type", None), str): + _walk(value) + + if config is not None: + _walk(config) + return found + + +def match_moe_block(module: "nn.Module", model_types: set[str] | None = None) -> MoESpec | None: """Return the MoE spec whose ``block_names`` matches ``module``. - Case-insensitive exact-name match against the class names in ``module``'s MRO - (see ``match_class_names``); quantized wrapper classes match through their - original base class. + Identification is by case-insensitive exact-name match against the class names + in ``module``'s MRO (see ``match_class_names``); quantized wrapper classes match + through their original base class. + + ``model_types`` (e.g. from ``collect_model_types(model.config)``) scopes the + result: when several specs' ``block_names`` match, one belonging to the model's + own model types wins. Scope prefers, never excludes — a class-name match outside + the scope still resolves, because remote-code towers reuse other models' module + classes under their own model_type (e.g. ``DeepseekMoE`` blocks inside a Kimi + model). """ + fallback = None for spec in iter_specs(MoESpec): if match_class_names(module, spec.block_names): - return spec - return None + if model_types and spec.model_type in model_types: + return spec + if fallback is None: + fallback = spec + return fallback diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index 7e6a7791df9..d1f83104da7 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -69,6 +69,15 @@ Model type names mirror (e.g. `qwen3_moe.py`, `gpt_oss.py`, `nemotron_h.py`); trust-remote-code models (`arctic`, `deepseek`) use their config `model_type`. +Resolution is two-level, mirroring HF's own indexing: the engine collects the +model's HF model types once per export (`collect_model_types(model.config)` — root +plus sub-configs, so VLM towers contribute e.g. `kimi_k2` explicitly) and passes +them down; `match_moe_block(module, model_types)` identifies the block by class +name in the module's MRO and uses the scope to pick the model's own spec when +class names collide across models. Scope prefers, never excludes: a tower whose +model_type has no spec still resolves by class name (remote-code forks reuse other +models' module classes). + During migration, call sites kept the legacy branches as a fallback behind the spec lookup. Once the specs covered every family the legacy chains served, the chains — and the silent ``w1/w2/w3`` guess for unknown models — were deleted: @@ -94,7 +103,7 @@ against existing export tests. | **P1** | Registry skeleton + MoE expert naming: `get_expert_linear_names` and `get_experts_list` read `spec.expert_linear_names` / `spec.has_iterable_experts`. The #1 "add a MoE model" shotgun-surgery driver. | this PR | | **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen3 specs). | this PR | | **P3** | `is_moe` explicit class-name list → `spec.moe_block_names` (arctic/dbrx_ffn are identification-only specs: no expert naming, so expert-name lookups keep the engine default). The generic `*SparseMoeBlock`/`*MoeLayer` conventions and the structural router+experts check stay in the engine. | this PR | -| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). | planned | +| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). Model_type-scoped resolution (`collect_model_types` + scoped `match_moe_block`, threaded via `ExportContext.model_types`) is already in place from this PR. | planned | | **P5** | Cross-subsystem pilot: unify the remaining copies of linear-fusion-group knowledge (see §4) into spec fields. Partially done: `_GATE_UP_PAIRS` became `MoESpec.gate_up_pair` and `model_calib.py`'s private import of it is gone — the first quantization consumer of `modelopt.modeling`. Remaining: `shared_input.SHARED_PATTERNS`, `algorithms.quant_grouping_rules`. | in progress | | **P6** | Migrate remaining quantization-side data: default disabled-quantizer patterns, on-the-fly conversion gates, AutoQuantize grouping rules (see §4). | planned | | **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Candidates for deletion on that track: `adjust_attn_amax_values`, `update_experts_avg_prequant_scale` (unused). NOTE: `MODEL_NAME_TO_TYPE` / `get_model_type` are NOT dead — `examples/hf_ptq/hf_ptq.py` and `multinode_ptq.py` still call them; migrate the examples before removing. | separate track | diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 800b51daca9..6e7d3389c3d 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -58,7 +58,7 @@ def _export_weight( def _prepare_dbrx_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None: """Fill missing input amax values for DBRX per-expert ModuleLists.""" experts_mlp = moe_module.experts.mlp - for linear_name in get_expert_linear_names(moe_module): + for linear_name in get_expert_linear_names(moe_module, ctx.model_types): if hasattr(experts_mlp, linear_name): linear_modulelist = getattr(experts_mlp, linear_name) if hasattr(linear_modulelist, "__iter__"): @@ -94,7 +94,7 @@ def _prepare_bmm_experts(name: str, moe_module: nn.Module, ctx: ExportContext) - ) def _prepare_iterable_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None: """Fill missing input amax values for iterable per-expert submodules.""" - expert_linear_names = get_expert_linear_names(moe_module) + expert_linear_names = get_expert_linear_names(moe_module, ctx.model_types) linear_name = None try: for linear_name in expert_linear_names: diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index d70348235f2..639ba7a4b26 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -28,7 +28,7 @@ except Exception: warn("Cannot find transformers package. Hugginface modules cannot be exported.") -from modelopt.modeling import match_moe_block +from modelopt.modeling import collect_model_types, match_moe_block from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import import_plugin @@ -85,12 +85,15 @@ def get_experts_list( module: torch.nn.Module, model_type: str, + model_types: set[str] | None = None, ): """Returns list of grouped experts by linear name for given module. Args: module: MoE block (e.g. MixtralSparseMoeBlock, NemotronHMOE). - model_type: `type(root_model).__name__.lower()` (may change after ModelOpt quantize). + model_type: `type(root_model).__name__.lower()`, used in error messages only. + model_types: the model's HF model types (``collect_model_types(model.config)``) + used to scope spec resolution; ``None`` matches across all specs. """ experts_list = [] @@ -98,10 +101,10 @@ def get_experts_list( # per-expert sub-modules (spec.has_iterable_experts); stacked/fused layouts # (DBRX, GptOss, ...) raise NotImplementedError here and are handled by other # paths. Name resolution itself is shared with get_expert_linear_names. - spec = match_moe_block(module) + spec = match_moe_block(module, model_types) if spec is None or not spec.has_iterable_experts: raise NotImplementedError(f" {model_type} not supported") - linear_names = get_expert_linear_names(module) + linear_names = get_expert_linear_names(module, model_types) # Common logic for all supported model types experts_list.extend( @@ -300,14 +303,18 @@ def is_mlp(module: nn.Module) -> bool: return any(key in type(module).__name__.upper() for key in ("MLP", "T5DENSE")) -def is_moe(module: nn.Module) -> bool: - """Returns whether the module is an MOE layer.""" +def is_moe(module: nn.Module, model_types: set[str] | None = None) -> bool: + """Returns whether the module is an MOE layer. + + ``model_types`` (``collect_model_types(model.config)``) scopes the spec lookup; + identification itself is unaffected by it (scope prefers, never excludes). + """ name = type(module).__name__.lower() # Auto-detect common MoE patterns if name.endswith("sparsemoeblock") or "moelayer" in name: return True # Non-standard MoE block names are per-model data (modelopt/modeling/models/*). - if match_moe_block(module) is not None: + if match_moe_block(module, model_types) is not None: return True # Structural detection: modules with router + experts (e.g. Gemma4TextDecoderLayer) return ( @@ -963,7 +970,7 @@ def get_stacked_scaling_factors(experts, get_function, module_name): return config -def get_expert_linear_names(module: nn.Module) -> list[str]: +def get_expert_linear_names(module: nn.Module, model_types: set[str] | None = None) -> list[str]: """Get the list of linear names for the experts. Resolution order: structural detection of fused-expert layouts first, then the @@ -979,7 +986,7 @@ def get_expert_linear_names(module: nn.Module) -> list[str]: if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"): return [first_proj_attr, "down_proj"] - spec = match_moe_block(module) + spec = match_moe_block(module, model_types) if spec is not None and spec.expert_linear_names is not None: return list(spec.expert_linear_names) @@ -1171,6 +1178,7 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: Returns: Number of expert gate/up pairs whose amaxes were synced. """ + model_types = collect_model_types(getattr(model, "config", None)) synced = 0 unmatched_block_names: set[str] = set() for _, sub_module in model.named_modules(): @@ -1182,7 +1190,7 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: # cross-model guessing. A spec declaring no pair (non-gated NemotronH, # fused GptOss/DBRX) needs no sync; an unmatched block is warned about # once instead of silently skipped. - spec = match_moe_block(sub_module) + spec = match_moe_block(sub_module, model_types) if spec is None: unmatched_block_names.add(type(sub_module).__name__) continue diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index acb1968e070..f3d224fe0dd 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -27,6 +27,7 @@ import torch.nn as nn import modelopt.torch.opt as mto +from modelopt.modeling import collect_model_types from modelopt.torch.quantization.conversion import quantizer_state from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer @@ -409,6 +410,7 @@ def _resmooth_experts_for_export( name_to_module = dict(model.named_modules()) if inplace else None model_type = type(model).__name__.lower() + model_types = collect_model_types(getattr(model, "config", None)) id_to_name: dict[int, str] = {id(m): n for n, m in model.named_modules()} out: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {} requant_weights: set[str] = set() @@ -466,10 +468,10 @@ def _process_group(modules: list[nn.Module]) -> None: # different tokens to each expert, so forward hooks cannot detect them as # sharing the same input tensor. for _, module in model.named_modules(): - if not is_moe(module): + if not is_moe(module, model_types): continue try: - expert_groups = get_experts_list(module, model_type) + expert_groups = get_experts_list(module, model_type, model_types) except NotImplementedError: continue for experts in expert_groups: diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 8e2cda63df9..351f5224064 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -54,6 +54,9 @@ class ExportContext: model: nn.Module dtype: torch.dtype is_modelopt_qlora: bool = False + model_types: set[str] | None = None + """The model's HF model types (root + sub-configs, via ``collect_model_types``), + used to scope per-model spec resolution in ``modelopt.modeling``.""" tied_cache: dict[int, nn.Module] = field(default_factory=dict) moe_tied_cache: dict[tuple[int, int], nn.Module] = field(default_factory=dict) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index cee64c22c05..f250afa5eb1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -30,6 +30,8 @@ from safetensors import safe_open from safetensors.torch import save_file +from modelopt.modeling import collect_model_types + from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales try: @@ -430,6 +432,7 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): # TODO: Handle DBRX MoE quantization_format = get_quantization_format(model) model_type = type(model).__name__.lower() + model_types = collect_model_types(getattr(model, "config", None)) module_names = set() # NVFP4 SVDQuant does not need pre-quant scale fusion (either into previous linear or layernorm) because @@ -445,12 +448,12 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): module_names.add(name) # For MoE models update pre_quant_scale to average pre_quant_scale amongst experts - if is_moe(module) and ( + if is_moe(module, model_types) and ( quantization_format is not QUANTIZATION_NONE and ("awq" in quantization_format or quantization_format == QUANTIZATION_NVFP4_SVDQUANT) ): # update_experts_avg_prequant_scale(module) - grouped_experts = get_experts_list(module, model_type) + grouped_experts = get_experts_list(module, model_type, model_types) for modules in grouped_experts: with fsdp2_aware_weight_update(model, modules): preprocess_linear_fusion(modules, resmooth_only=True) @@ -797,7 +800,12 @@ def _process_quantized_modules( # Per-call tied-weight dedup caches inside the context. Created fresh on # every invocation so cache state is scoped to one export and cannot leak # into a later call (see ExportContext). - ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + ctx = ExportContext( + model=model, + dtype=dtype, + is_modelopt_qlora=is_modelopt_qlora, + model_types=collect_model_types(getattr(model, "config", None)), + ) fsdp_module_to_reshard = None for name, sub_module in model.named_modules(): @@ -858,9 +866,14 @@ def _export_transformers_checkpoint( # Handle input quantizers of experts that are not calibrated. Each MoE block is # dispatched by its experts container to the matching preparation handler. - prepare_ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + prepare_ctx = ExportContext( + model=model, + dtype=dtype, + is_modelopt_qlora=is_modelopt_qlora, + model_types=collect_model_types(getattr(model, "config", None)), + ) for name, sub_module in model.named_modules(): - if is_moe(sub_module) and hasattr(sub_module, "experts"): + if is_moe(sub_module, prepare_ctx.model_types) and hasattr(sub_module, "experts"): handler = PrepareMoEInputsRegistry.match(sub_module.experts) if handler is None: # Unsupported MoE model structure diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index 204866b9372..06ab2b01eed 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -19,11 +19,14 @@ import torch.nn as nn from modelopt.modeling import ( + MoESpec, + collect_model_types, iter_gate_up_pairs, iter_pqs_fuse_rules, match_moe_block, weight_plus_one_norm_names, ) +from modelopt.modeling.registry import _SPECS from modelopt.torch.export.layer_utils import ( get_expert_linear_names, get_experts_list, @@ -245,3 +248,44 @@ class UnknownSparseMoeBlock(nn.Module): # No spec -> no cross-model guessing: nothing synced, one warning. with pytest.warns(UserWarning, match="no registered MoESpec"): assert sync_moe_gate_up_amax(model) == 0 + + +def test_collect_model_types_walks_sub_configs(): + from types import SimpleNamespace + + text_config = SimpleNamespace(model_type="kimi_k2") + config = SimpleNamespace(model_type="kimi_vl", text_config=text_config, hidden_size=4096) + assert collect_model_types(config) == {"kimi_vl", "kimi_k2"} + assert collect_model_types(None) == set() + + +def test_match_moe_block_scope_prefers_own_model_type(): + class Qwen3MoeSparseMoeBlock(nn.Module): + pass + + # A hypothetical remote-code fork registering the same block class name under + # its own model type: scope must pick the model's own spec among candidates. + fork_spec = MoESpec( + model_type="zz_fork", + block_names=("Qwen3MoeSparseMoeBlock",), + expert_linear_names=("a_proj", "b_proj"), + ) + _SPECS.append(fork_spec) + try: + assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"zz_fork"}) is fork_spec + assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"qwen3_moe"}).model_type == "qwen3_moe" + # No scope -> first registered class-name match wins (legacy order). + assert match_moe_block(Qwen3MoeSparseMoeBlock()).model_type == "qwen3_moe" + finally: + _SPECS.remove(fork_spec) + + +def test_match_moe_block_scope_never_excludes(): + class Qwen3MoeSparseMoeBlock(nn.Module): + pass + + # A tower whose model_type has no spec of its own must still resolve by class + # name (e.g. DeepseekMoE blocks inside a Kimi model). + spec = match_moe_block(Qwen3MoeSparseMoeBlock(), {"some_unknown_vlm"}) + assert spec is not None + assert spec.model_type == "qwen3_moe" From 23b7ca8dc99524ad232cab856f530b361f9df91c Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 16 Jul 2026 22:30:35 +0000 Subject: [PATCH 12/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/models/arctic.py | 18 +-- modelopt/modeling/models/dbrx.py | 18 ++- modelopt/modeling/models/deepseek.py | 12 +- modelopt/modeling/models/gemma4.py | 16 ++- modelopt/modeling/models/gpt_oss.py | 12 +- modelopt/modeling/models/mixtral.py | 31 +++--- modelopt/modeling/models/nemotron_h.py | 14 ++- modelopt/modeling/models/qwen2_moe.py | 14 ++- modelopt/modeling/models/qwen3_5_moe.py | 14 ++- modelopt/modeling/models/qwen3_moe.py | 14 ++- modelopt/modeling/models/qwen3_next.py | 14 ++- modelopt/modeling/registry.py | 87 +++++++-------- modelopt/modeling/specs.py | 81 +++++++++++--- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 30 +++-- modelopt/torch/export/layer_utils.py | 73 ++++++++----- .../torch/export/plugins/vllm_fakequant_hf.py | 3 +- modelopt/torch/export/registry.py | 5 +- modelopt/torch/export/unified_export_hf.py | 2 +- tests/unit/torch/export/test_layer_utils.py | 6 +- .../unit/torch/export/test_modeling_specs.py | 103 ++++++++++++------ 20 files changed, 351 insertions(+), 216 deletions(-) diff --git a/modelopt/modeling/models/arctic.py b/modelopt/modeling/models/arctic.py index 347b6d2bb42..f29979e1108 100644 --- a/modelopt/modeling/models/arctic.py +++ b/modelopt/modeling/models/arctic.py @@ -16,16 +16,20 @@ """Snowflake Arctic specs (trust-remote-code model type ``arctic``).""" from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant register( MoESpec( model_type="arctic", - # Non-standard block name (for is_moe identification). - block_names=("ArcticMoE",), - # ArcticMLP experts use Mixtral-style w1/w2/w3 naming (previously served by - # the engine's implicit w1/w2/w3 default, now declared explicitly). - expert_linear_names=("w1", "w2", "w3"), - gate_up_pair=("w1", "w3"), + variants=( + MoEVariant( + # Non-standard block name (for is_moe identification). + block_names=("ArcticMoE",), + # ArcticMLP experts use Mixtral-style w1/w2/w3 naming (previously served by + # the engine's implicit w1/w2/w3 default, now declared explicitly). + expert_linear_names=("w1", "w2", "w3"), + gate_up_pair=("w1", "w3"), + ), + ), ) ) diff --git a/modelopt/modeling/models/dbrx.py b/modelopt/modeling/models/dbrx.py index 82f84e7c3cd..406d632215c 100644 --- a/modelopt/modeling/models/dbrx.py +++ b/modelopt/modeling/models/dbrx.py @@ -16,15 +16,7 @@ """DBRX specs (HF model type ``dbrx``).""" from ..registry import register -from ..specs import MoESpec - -register( - MoESpec( - model_type="dbrx", - block_names=("DBRXMoeSparseMoeBlock",), - expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), - ) -) +from ..specs import MoESpec, MoEVariant # HF DbrxFFN (non-standard block name, identified for is_moe). Expert names refer to # the quantized layout: _QuantDbrxExpertGLU converts the fused w1/v1/w2 parameters @@ -34,7 +26,11 @@ register( MoESpec( model_type="dbrx", - block_names=("DbrxFFN",), - expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), + variants=( + MoEVariant( + block_names=("DbrxFFN",), + expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), + ), + ), ) ) diff --git a/modelopt/modeling/models/deepseek.py b/modelopt/modeling/models/deepseek.py index ff0624b13a9..597f7cf4381 100644 --- a/modelopt/modeling/models/deepseek.py +++ b/modelopt/modeling/models/deepseek.py @@ -20,7 +20,7 @@ """ from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant # DeepseekMoE experts ARE structurally iterable (ModuleList of DeepseekMLP), but # has_iterable_experts stays False until the grouped export path (get_experts_list @@ -29,8 +29,12 @@ register( MoESpec( model_type="deepseek", - block_names=("DeepseekMoE",), - expert_linear_names=("gate_proj", "down_proj", "up_proj"), - gate_up_pair=("gate_proj", "up_proj"), + variants=( + MoEVariant( + block_names=("DeepseekMoE",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), + ), + ), ) ) diff --git a/modelopt/modeling/models/gemma4.py b/modelopt/modeling/models/gemma4.py index bab1570437f..d3cf44e995b 100644 --- a/modelopt/modeling/models/gemma4.py +++ b/modelopt/modeling/models/gemma4.py @@ -16,15 +16,19 @@ """Gemma4 specs (HF model type ``gemma4``).""" from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant register( MoESpec( model_type="gemma4", - # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. - block_names=("Gemma4TextDecoderLayer",), - expert_linear_names=("gate_proj", "down_proj", "up_proj"), - gate_up_pair=("gate_proj", "up_proj"), - has_iterable_experts=True, + variants=( + MoEVariant( + # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. + block_names=("Gemma4TextDecoderLayer",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), + has_iterable_experts=True, + ), + ), ) ) diff --git a/modelopt/modeling/models/gpt_oss.py b/modelopt/modeling/models/gpt_oss.py index 20f33f476a8..91910969534 100644 --- a/modelopt/modeling/models/gpt_oss.py +++ b/modelopt/modeling/models/gpt_oss.py @@ -16,13 +16,17 @@ """GPT-OSS specs (HF model type ``gpt_oss``).""" from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant register( MoESpec( model_type="gpt_oss", - # GPT-OSS fuses gate and up into a single gate_up_proj. - block_names=("GptOssMoE",), - expert_linear_names=("gate_up_proj", "down_proj"), + variants=( + MoEVariant( + # GPT-OSS fuses gate and up into a single gate_up_proj. + block_names=("GptOssMoE",), + expert_linear_names=("gate_up_proj", "down_proj"), + ), + ), ) ) diff --git a/modelopt/modeling/models/mixtral.py b/modelopt/modeling/models/mixtral.py index bc7aa19592f..ed4531224ae 100644 --- a/modelopt/modeling/models/mixtral.py +++ b/modelopt/modeling/models/mixtral.py @@ -16,26 +16,27 @@ """Mixtral specs (HF model type ``mixtral``).""" from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant # Mixtral with iterable experts uses w1/w2/w3. Fused experts (transformers 5.0+) are # detected from their per-expert quantizer attributes and need no naming override here. register( MoESpec( model_type="mixtral", - block_names=("MixtralSparseMoeBlock",), - expert_linear_names=("w1", "w2", "w3"), - has_iterable_experts=True, - # w1 = gate, w3 = up, w2 = down (Mixtral convention). - gate_up_pair=("w1", "w3"), - ) -) - -# Older transformers naming for Mixtral. -register( - MoESpec( - model_type="mixtral", - block_names=("MixtralMoeSparseMoeBlock",), - expert_linear_names=("linear_fc1", "linear_fc2"), + variants=( + MoEVariant( + block_names=("MixtralSparseMoeBlock",), + expert_linear_names=("w1", "w2", "w3"), + has_iterable_experts=True, + # w1 = gate, w3 = up, w2 = down (Mixtral convention). + gate_up_pair=("w1", "w3"), + ), + # Legacy-naming layout kept from the legacy engine chain: same model + # type, different block class and projection names. + MoEVariant( + block_names=("MixtralMoeSparseMoeBlock",), + expert_linear_names=("linear_fc1", "linear_fc2"), + ), + ), ) ) diff --git a/modelopt/modeling/models/nemotron_h.py b/modelopt/modeling/models/nemotron_h.py index 95697eda89f..6bf8622c342 100644 --- a/modelopt/modeling/models/nemotron_h.py +++ b/modelopt/modeling/models/nemotron_h.py @@ -16,14 +16,18 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant register( MoESpec( model_type="nemotron_h", - # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). - block_names=("NemotronHMOE",), - expert_linear_names=("up_proj", "down_proj"), - has_iterable_experts=True, + variants=( + MoEVariant( + # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). + block_names=("NemotronHMOE",), + expert_linear_names=("up_proj", "down_proj"), + has_iterable_experts=True, + ), + ), ) ) diff --git a/modelopt/modeling/models/qwen2_moe.py b/modelopt/modeling/models/qwen2_moe.py index 78b8c9e3abc..89dd0437d7b 100644 --- a/modelopt/modeling/models/qwen2_moe.py +++ b/modelopt/modeling/models/qwen2_moe.py @@ -16,14 +16,18 @@ """Qwen2-MoE specs (HF model type ``qwen2_moe``).""" from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant register( MoESpec( model_type="qwen2_moe", - block_names=("Qwen2MoeSparseMoeBlock",), - expert_linear_names=("gate_proj", "down_proj", "up_proj"), - gate_up_pair=("gate_proj", "up_proj"), - has_iterable_experts=True, + variants=( + MoEVariant( + block_names=("Qwen2MoeSparseMoeBlock",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), + has_iterable_experts=True, + ), + ), ) ) diff --git a/modelopt/modeling/models/qwen3_5_moe.py b/modelopt/modeling/models/qwen3_5_moe.py index 7eaadfa580d..e1c92ddf1e8 100644 --- a/modelopt/modeling/models/qwen3_5_moe.py +++ b/modelopt/modeling/models/qwen3_5_moe.py @@ -16,14 +16,18 @@ """Qwen3.5-MoE specs (HF model type ``qwen3_5_moe``).""" from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant register( MoESpec( model_type="qwen3_5_moe", - block_names=("Qwen3_5MoeSparseMoeBlock",), - expert_linear_names=("gate_proj", "down_proj", "up_proj"), - gate_up_pair=("gate_proj", "up_proj"), - has_iterable_experts=True, + variants=( + MoEVariant( + block_names=("Qwen3_5MoeSparseMoeBlock",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), + has_iterable_experts=True, + ), + ), ) ) diff --git a/modelopt/modeling/models/qwen3_moe.py b/modelopt/modeling/models/qwen3_moe.py index d5ac4dd4019..aee315d7168 100644 --- a/modelopt/modeling/models/qwen3_moe.py +++ b/modelopt/modeling/models/qwen3_moe.py @@ -16,15 +16,19 @@ """Qwen3-MoE specs (HF model type ``qwen3_moe``).""" from ..registry import register -from ..specs import ExportSpec, MoESpec +from ..specs import ExportSpec, MoESpec, MoEVariant register( MoESpec( model_type="qwen3_moe", - block_names=("Qwen3MoeSparseMoeBlock",), - expert_linear_names=("gate_proj", "down_proj", "up_proj"), - gate_up_pair=("gate_proj", "up_proj"), - has_iterable_experts=True, + variants=( + MoEVariant( + block_names=("Qwen3MoeSparseMoeBlock",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), + has_iterable_experts=True, + ), + ), ) ) diff --git a/modelopt/modeling/models/qwen3_next.py b/modelopt/modeling/models/qwen3_next.py index f137ae64a6c..002bece6093 100644 --- a/modelopt/modeling/models/qwen3_next.py +++ b/modelopt/modeling/models/qwen3_next.py @@ -16,14 +16,18 @@ """Qwen3-Next specs (HF model type ``qwen3_next``).""" from ..registry import register -from ..specs import MoESpec +from ..specs import MoESpec, MoEVariant register( MoESpec( model_type="qwen3_next", - block_names=("Qwen3NextSparseMoeBlock",), - expert_linear_names=("gate_proj", "down_proj", "up_proj"), - gate_up_pair=("gate_proj", "up_proj"), - has_iterable_experts=True, + variants=( + MoEVariant( + block_names=("Qwen3NextSparseMoeBlock",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), + has_iterable_experts=True, + ), + ), ) ) diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index 72114c85ee2..44e235dd35a 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -26,17 +26,16 @@ from collections.abc import Iterator from typing import TYPE_CHECKING, TypeVar -from .specs import ExportSpec, ModelSpec, MoESpec, NormSpec +from .specs import ExportSpec, ModelSpec, MoESpec, MoEVariant, NormSpec if TYPE_CHECKING: import torch.nn as nn __all__ = [ "collect_model_types", + "get_specs", "iter_gate_up_pairs", "iter_pqs_fuse_rules", - "iter_specs", - "match_class_names", "match_moe_block", "register", "weight_plus_one_norm_names", @@ -53,27 +52,20 @@ def register(spec: SpecT) -> SpecT: return spec -def iter_specs(spec_cls: type[SpecT]) -> Iterator[SpecT]: - """Yield every registered spec of type ``spec_cls`` (in registration order).""" - for spec in _SPECS: - if isinstance(spec, spec_cls): - yield spec +def get_specs(spec_cls: type[SpecT], model_types: set[str] | None = None) -> list[SpecT]: + """Return the registered specs of type ``spec_cls``, in registration order. - -def match_class_names(module, names: tuple[str, ...]) -> bool: - """Return True if any of ``names`` equals a class name in ``module``'s MRO. - - Case-insensitive exact-name comparison against ``cls.__name__`` for every class - in ``type(module).__mro__`` — the same semantics as the export dispatch - registry's string keys (``modelopt.torch.export.registry``). Dynamically - generated quantized classes are subclasses of the original module class, so they - match through their base; exact-name comparison avoids substring false - positives. Comparison is case-insensitive because some registered names predate - this registry and their casing was never exercised by the legacy substring - matching. + This is the model_type index: with ``model_types`` (from + ``collect_model_types(model.config)``), only the specs belonging to the model's + own model types are returned, so consumers resolve per-model data without + scanning the registry. Without it, all specs of the type are returned + (aggregators, no-config compatibility). """ - mro_names = {cls.__name__.lower() for cls in type(module).__mro__} - return any(name.lower() in mro_names for name in names) + return [ + spec + for spec in _SPECS + if isinstance(spec, spec_cls) and (not model_types or spec.model_type in model_types) + ] def iter_pqs_fuse_rules(): @@ -82,7 +74,7 @@ def iter_pqs_fuse_rules(): Aggregated across all registered export specs (the consumer matches each model module against the substrings, so the order across specs does not matter). """ - for spec in iter_specs(ExportSpec): + for spec in get_specs(ExportSpec): yield from spec.pqs_fuse_rules @@ -99,16 +91,17 @@ def iter_gate_up_pairs() -> Iterator[tuple[str, str]]: in a follow-up (see MODEL_SPECIFIC_REFACTOR.md P5). """ seen = set() - for spec in iter_specs(MoESpec): - pair = spec.gate_up_pair - if pair is not None and pair not in seen: - seen.add(pair) - yield pair + for spec in get_specs(MoESpec): + for variant in spec.variants: + pair = variant.gate_up_pair + if pair is not None and pair not in seen: + seen.add(pair) + yield pair def weight_plus_one_norm_names() -> tuple[str, ...]: """All norm class names whose stored weight is ``w - 1``, across all norm specs.""" - return tuple(name for spec in iter_specs(NormSpec) for name in spec.weight_plus_one_norm_names) + return tuple(name for spec in get_specs(NormSpec) for name in spec.weight_plus_one_norm_names) def collect_model_types(config) -> set[str]: @@ -139,25 +132,23 @@ def _walk(cfg) -> None: return found -def match_moe_block(module: "nn.Module", model_types: set[str] | None = None) -> MoESpec | None: - """Return the MoE spec whose ``block_names`` matches ``module``. +def match_moe_block(module: "nn.Module", model_types: set[str] | None = None) -> MoEVariant | None: + """Return the MoE layout variant for ``module``, resolved by model type. - Identification is by case-insensitive exact-name match against the class names - in ``module``'s MRO (see ``match_class_names``); quantized wrapper classes match - through their original base class. + ``model_types`` (from ``collect_model_types(model.config)``: root plus + sub-config types, so VLM towers are covered) is a strict filter: only specs + registered under the model's own model types are considered. A model whose + model_type has no spec resolves to ``None`` even if its module class names + coincide with another model's — register a spec instead of inheriting a + neighbor's data. ``None`` or an empty set (no config available: unit tests, + the TRT-LLM path) searches all specs. - ``model_types`` (e.g. from ``collect_model_types(model.config)``) scopes the - result: when several specs' ``block_names`` match, one belonging to the model's - own model types wins. Scope prefers, never excludes — a class-name match outside - the scope still resolves, because remote-code towers reuse other models' module - classes under their own model_type (e.g. ``DeepseekMoE`` blocks inside a Kimi - model). + Within the scope, each spec's variant ``block_names`` identifies the block and + disambiguates same-model layout variants (``MoESpec.match_variant``); quantized + wrapper classes match through their original base class in the MRO. """ - fallback = None - for spec in iter_specs(MoESpec): - if match_class_names(module, spec.block_names): - if model_types and spec.model_type in model_types: - return spec - if fallback is None: - fallback = spec - return fallback + for spec in get_specs(MoESpec, model_types): + variant = spec.match_variant(module) + if variant is not None: + return variant + return None diff --git a/modelopt/modeling/specs.py b/modelopt/modeling/specs.py index e0d2027d8c9..814cf2a9dcf 100644 --- a/modelopt/modeling/specs.py +++ b/modelopt/modeling/specs.py @@ -23,12 +23,13 @@ - **subsystem specs** hold one subsystem's per-model policy (``ExportSpec``; quantization / speculative-decoding specs to follow). -Specs hold per-model data only, no logic. +Specs hold per-model data plus trivial accessors over that data; subsystem logic +never lives here. """ from dataclasses import dataclass -__all__ = ["ExportSpec", "MoESpec", "ModelSpec", "NormSpec"] +__all__ = ["ExportSpec", "MoESpec", "MoEVariant", "ModelSpec", "NormSpec", "match_class_names"] @dataclass @@ -41,27 +42,38 @@ class ModelSpec: model_type: str """The HF model type this spec belongs to (``config.model_type``, e.g. - ``"qwen3_moe"``). Not necessarily unique: a model type may register several specs - for different module layouts (e.g. two ``mixtral`` MoE-block variants).""" + ``"qwen3_moe"``). A model registers one spec instance per spec kind; same-model + layout variants nest inside the spec (see ``MoESpec.variants``).""" -@dataclass -class MoESpec(ModelSpec): - """MoE architecture facts for one model (or one of its MoE-block variants). +def match_class_names(module, names: tuple[str, ...]) -> bool: + """Return True if any of ``names`` equals a class name in ``module``'s MRO. + + Case-insensitive exact-name comparison against ``cls.__name__`` for every class + in ``type(module).__mro__`` — the same semantics as the export dispatch + registry's string keys (``modelopt.torch.export.registry``). Dynamically + generated quantized classes are subclasses of the original module class, so they + match through their base; exact-name comparison avoids substring false + positives. Comparison is case-insensitive because some registered names predate + this registry and their casing was never exercised by the legacy substring + matching. + """ + mro_names = {cls.__name__.lower() for cls in type(module).__mro__} + return any(name.lower() in mro_names for name in names) - Unlike the subsystem specs, this describes what a model's MoE blocks *are* — - which class, what the expert projections are called — so any modelopt subsystem - (export, quantization, speculative decoding, ...) can read it instead of keeping - its own per-model MoE table. - Resolved from a model sub-module via ``block_names``, the matching key: MoE block - class names (e.g. ``"Qwen3MoeSparseMoeBlock"``) compared case-insensitively - against the class names in the module's MRO (see ``registry.match_class_names``). +@dataclass +class MoEVariant: + """One concrete MoE-block layout of a model. + + A model type usually has exactly one; it has several when the same checkpoint + materializes with different module classes and projection names (e.g. Mixtral + across transformers generations). ``block_names`` tells the variants apart. """ block_names: tuple[str, ...] = () """Matching key: MoE block class names, matched against the module's MRO - (case-insensitive exact names, not substrings).""" + (case-insensitive exact names, not substrings; see ``match_class_names``).""" expert_linear_names: tuple[str, ...] | None = None """Expert linear projection names, e.g. ``("gate_proj", "down_proj", "up_proj")``. @@ -83,6 +95,45 @@ class names (e.g. ``"Qwen3MoeSparseMoeBlock"``) compared case-insensitively ``sync_moe_gate_up_amax``) and by calibration grouping.""" +@dataclass +class MoESpec(ModelSpec): + """MoE architecture facts for one model: its MoE-block layout variant(s). + + Unlike the subsystem specs, this describes what a model's MoE blocks *are* — + which class, what the expert projections are called — so any modelopt subsystem + (export, quantization, speculative decoding, ...) can read it instead of keeping + its own per-model MoE table. + """ + + variants: tuple[MoEVariant, ...] = () + """The model's MoE-block layouts; more than one when the same checkpoint + materializes differently (see ``MoEVariant``).""" + + def match_variant(self, module) -> MoEVariant | None: + """Return the variant whose ``block_names`` matches ``module``, else None.""" + for variant in self.variants: + if match_class_names(module, variant.block_names): + return variant + return None + + def expert_linear_names_for(self, module) -> tuple[str, ...] | None: + """Resolve ``module``'s expert linear names within this model. + + When every variant agrees on one naming, the module's class is irrelevant + (a spec can provide naming without the block class being known); with + several namings, the module's class picks the variant. + """ + namings = { + variant.expert_linear_names + for variant in self.variants + if variant.expert_linear_names is not None + } + if len(namings) == 1: + return next(iter(namings)) + variant = self.match_variant(module) + return variant.expert_linear_names if variant is not None else None + + @dataclass class NormSpec(ModelSpec): """Normalization-layer architecture facts for one model. diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index d1f83104da7..710d3dd8d28 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -52,10 +52,10 @@ Two layers: ```text modelopt/modeling/ - specs.py # ModelSpec (base) + topic specs (MoESpec: shared architecture - # facts) + subsystem specs (ExportSpec: HF-export-path policy); - # future attention/norm topic specs and quantization/spec-dec - # subsystem specs go here too + specs.py # ModelSpec (base) + topic specs (MoESpec/NormSpec: shared + # architecture facts; MoESpec nests one MoEVariant per concrete + # block layout) + subsystem specs (ExportSpec: HF-export-path + # policy); future topic/subsystem specs go here too registry.py # register() + lookups queried by spec type (None when unmatched) # + the MRO exact-name matching core (match_class_names) __init__.py # re-exports; importing it registers all specs @@ -69,14 +69,20 @@ Model type names mirror (e.g. `qwen3_moe.py`, `gpt_oss.py`, `nemotron_h.py`); trust-remote-code models (`arctic`, `deepseek`) use their config `model_type`. -Resolution is two-level, mirroring HF's own indexing: the engine collects the -model's HF model types once per export (`collect_model_types(model.config)` — root -plus sub-configs, so VLM towers contribute e.g. `kimi_k2` explicitly) and passes -them down; `match_moe_block(module, model_types)` identifies the block by class -name in the module's MRO and uses the scope to pick the model's own spec when -class names collide across models. Scope prefers, never excludes: a tower whose -model_type has no spec still resolves by class name (remote-code forks reuse other -models' module classes). +Resolution is by model type first, mirroring HF's own indexing: the engine +collects the model's HF model types once per export +(`collect_model_types(model.config)` — root plus sub-configs, so VLM towers +contribute e.g. `kimi_k2` explicitly) and passes them down. The scope is strict: +only the model's own specs are consulted, so a model whose model_type has no spec +fails loudly (register a spec) instead of inheriting a neighbor's data through a +coincidental class-name match. Within the scope, each `MoESpec` nests one +`MoEVariant` per concrete block layout — several when the same checkpoint +materializes with different classes and projection names (Mixtral across +transformers generations); variant `block_names` (matched against the module +MRO) identify MoE blocks (`is_moe`) and pick the variant. +`get_expert_linear_names` doesn't need the block class at all when a model's +variants agree on one naming. Without a scope (no config available: unit tests, +the TRT-LLM path), lookups search all specs by class name. During migration, call sites kept the legacy branches as a fallback behind the spec lookup. Once the specs covered every family the legacy chains served, the diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index 639ba7a4b26..ff6d447cc4b 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -28,7 +28,7 @@ except Exception: warn("Cannot find transformers package. Hugginface modules cannot be exported.") -from modelopt.modeling import collect_model_types, match_moe_block +from modelopt.modeling import MoESpec, collect_model_types, get_specs, match_moe_block from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import import_plugin @@ -84,16 +84,14 @@ def get_experts_list( module: torch.nn.Module, - model_type: str, - model_types: set[str] | None = None, + model_types: set[str], ): """Returns list of grouped experts by linear name for given module. Args: module: MoE block (e.g. MixtralSparseMoeBlock, NemotronHMOE). - model_type: `type(root_model).__name__.lower()`, used in error messages only. - model_types: the model's HF model types (``collect_model_types(model.config)``) - used to scope spec resolution; ``None`` matches across all specs. + model_types: the model's HF model types (``collect_model_types(model.config)``), + used to resolve the model's own specs. """ experts_list = [] @@ -101,9 +99,12 @@ def get_experts_list( # per-expert sub-modules (spec.has_iterable_experts); stacked/fused layouts # (DBRX, GptOss, ...) raise NotImplementedError here and are handled by other # paths. Name resolution itself is shared with get_expert_linear_names. - spec = match_moe_block(module, model_types) - if spec is None or not spec.has_iterable_experts: - raise NotImplementedError(f" {model_type} not supported") + variant = match_moe_block(module, model_types) + if variant is None or not variant.has_iterable_experts: + raise NotImplementedError( + f"MoE block {type(module).__name__!r} " + f"(model types: {sorted(model_types or [])}) not supported" + ) linear_names = get_expert_linear_names(module, model_types) # Common logic for all supported model types @@ -306,8 +307,9 @@ def is_mlp(module: nn.Module) -> bool: def is_moe(module: nn.Module, model_types: set[str] | None = None) -> bool: """Returns whether the module is an MOE layer. - ``model_types`` (``collect_model_types(model.config)``) scopes the spec lookup; - identification itself is unaffected by it (scope prefers, never excludes). + ``model_types`` (``collect_model_types(model.config)``) strictly scopes the + registry lookup to the model's own specs; the generic naming conventions and the + structural check below are unaffected by it. """ name = type(module).__name__.lower() # Auto-detect common MoE patterns @@ -970,13 +972,20 @@ def get_stacked_scaling_factors(experts, get_function, module_name): return config -def get_expert_linear_names(module: nn.Module, model_types: set[str] | None = None) -> list[str]: +def get_expert_linear_names(module: nn.Module, model_types: set[str]) -> list[str]: """Get the list of linear names for the experts. - Resolution order: structural detection of fused-expert layouts first, then the - model spec registry (modelopt/modeling). Raises NotImplementedError when neither - resolves, so a new MoE model fails loudly instead of silently inheriting another - model's naming. + Args: + module: the MoE block. + model_types: the model's HF model types (``collect_model_types(model.config)``). + + Resolution order: structural detection of fused-expert layouts first (runtime + state: transformers>=5 fused experts, and modules rewritten during export), then + the model's own specs by model type — the block class name is not needed, so a + spec can provide expert naming without declaring ``block_names``. Only when a + composite model carries several MoE towers with different namings does the block + class disambiguate. Raises NotImplementedError when nothing resolves, so a new + MoE model fails loudly instead of silently inheriting another model's naming. """ # Structural detection: after _export_fused_experts, fused expert modules # have per-expert submodules with gate_proj/up_proj/down_proj. @@ -986,12 +995,24 @@ def get_expert_linear_names(module: nn.Module, model_types: set[str] | None = No if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"): return [first_proj_attr, "down_proj"] - spec = match_moe_block(module, model_types) - if spec is not None and spec.expert_linear_names is not None: - return list(spec.expert_linear_names) + resolved = { + names + for spec in get_specs(MoESpec, model_types) + if (names := spec.expert_linear_names_for(module)) is not None + } + if len(resolved) == 1: + return list(next(iter(resolved))) + if len(resolved) > 1: + # Several MoE model types in scope resolving differently — no known real + # case; fail loudly rather than pick one. + raise NotImplementedError( + f"Ambiguous expert linear names for MoE block {type(module).__name__!r}: " + f"model types {sorted(model_types or [])} resolve to {sorted(resolved)}." + ) raise NotImplementedError( - f"Cannot resolve expert linear names for MoE block {type(module).__name__!r}. " + f"Cannot resolve expert linear names for MoE block {type(module).__name__!r} " + f"(model types: {sorted(model_types) if model_types else 'unknown'}). " "Register a MoESpec with expert_linear_names for this model under " "modelopt/modeling/models/." ) @@ -1186,17 +1207,17 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: continue if not hasattr(sub_module.experts, "__iter__"): continue - # The gate/up pair is this model's own data (MoESpec.gate_up_pair) — no - # cross-model guessing. A spec declaring no pair (non-gated NemotronH, + # The gate/up pair is this model's own data (MoEVariant.gate_up_pair) — no + # cross-model guessing. A variant declaring no pair (non-gated NemotronH, # fused GptOss/DBRX) needs no sync; an unmatched block is warned about # once instead of silently skipped. - spec = match_moe_block(sub_module, model_types) - if spec is None: + variant = match_moe_block(sub_module, model_types) + if variant is None: unmatched_block_names.add(type(sub_module).__name__) continue - if spec.gate_up_pair is None: + if variant.gate_up_pair is None: continue - gate_name, up_name = spec.gate_up_pair + gate_name, up_name = variant.gate_up_pair for expert in sub_module.experts: gate_linear = getattr(expert, gate_name, None) up_linear = getattr(expert, up_name, None) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index f3d224fe0dd..927e7bee670 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -409,7 +409,6 @@ def _resmooth_experts_for_export( name_to_module = dict(model.named_modules()) if inplace else None - model_type = type(model).__name__.lower() model_types = collect_model_types(getattr(model, "config", None)) id_to_name: dict[int, str] = {id(m): n for n, m in model.named_modules()} out: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {} @@ -471,7 +470,7 @@ def _process_group(modules: list[nn.Module]) -> None: if not is_moe(module, model_types): continue try: - expert_groups = get_experts_list(module, model_type, model_types) + expert_groups = get_experts_list(module, model_types) except NotImplementedError: continue for experts in expert_groups: diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 351f5224064..20f751fab3b 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -54,9 +54,10 @@ class ExportContext: model: nn.Module dtype: torch.dtype is_modelopt_qlora: bool = False - model_types: set[str] | None = None + model_types: set[str] = field(default_factory=set) """The model's HF model types (root + sub-configs, via ``collect_model_types``), - used to scope per-model spec resolution in ``modelopt.modeling``.""" + used to resolve per-model specs in ``modelopt.modeling``. Empty means unknown: + spec lookups then fail loudly instead of guessing.""" tied_cache: dict[int, nn.Module] = field(default_factory=dict) moe_tied_cache: dict[tuple[int, int], nn.Module] = field(default_factory=dict) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f250afa5eb1..a9294b72c35 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -453,7 +453,7 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): and ("awq" in quantization_format or quantization_format == QUANTIZATION_NVFP4_SVDQUANT) ): # update_experts_avg_prequant_scale(module) - grouped_experts = get_experts_list(module, model_type, model_types) + grouped_experts = get_experts_list(module, model_types) for modules in grouped_experts: with fsdp2_aware_weight_update(model, modules): preprocess_linear_fusion(modules, resmooth_only=True) diff --git a/tests/unit/torch/export/test_layer_utils.py b/tests/unit/torch/export/test_layer_utils.py index 538de5d605e..8573bb99bed 100644 --- a/tests/unit/torch/export/test_layer_utils.py +++ b/tests/unit/torch/export/test_layer_utils.py @@ -101,7 +101,7 @@ class NemotronHMOE(nn.Module): def test_get_expert_linear_names_gemma4(): - assert get_expert_linear_names(Gemma4TextDecoderLayer()) == [ + assert get_expert_linear_names(Gemma4TextDecoderLayer(), {"gemma4"}) == [ "gate_proj", "down_proj", "up_proj", @@ -109,8 +109,8 @@ def test_get_expert_linear_names_gemma4(): def test_get_expert_linear_names_mixtral(): - assert get_expert_linear_names(MixtralSparseMoeBlock()) == ["w1", "w2", "w3"] + assert get_expert_linear_names(MixtralSparseMoeBlock(), {"mixtral"}) == ["w1", "w2", "w3"] def test_get_expert_linear_names_nemotron(): - assert get_expert_linear_names(NemotronHMOE()) == ["up_proj", "down_proj"] + assert get_expert_linear_names(NemotronHMOE(), {"nemotron_h"}) == ["up_proj", "down_proj"] diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index 06ab2b01eed..0ba24f149b1 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -20,6 +20,7 @@ from modelopt.modeling import ( MoESpec, + MoEVariant, collect_model_types, iter_gate_up_pairs, iter_pqs_fuse_rules, @@ -53,17 +54,16 @@ class _UnknownMoeBlock(nn.Module): def test_match_moe_block_by_class_name(): - spec = match_moe_block(Qwen3MoeSparseMoeBlock()) - assert spec is not None - assert spec.model_type == "qwen3_moe" - assert spec.expert_linear_names == ("gate_proj", "down_proj", "up_proj") - assert spec.has_iterable_experts + variant = match_moe_block(Qwen3MoeSparseMoeBlock()) + assert variant is not None + assert variant.expert_linear_names == ("gate_proj", "down_proj", "up_proj") + assert variant.has_iterable_experts def test_match_moe_block_matches_quantized_class_via_mro(): - spec = match_moe_block(QuantMixtralSparseMoeBlock()) - assert spec is not None - assert spec.model_type == "mixtral" + variant = match_moe_block(QuantMixtralSparseMoeBlock()) + assert variant is not None + assert variant.expert_linear_names == ("w1", "w2", "w3") def test_match_moe_block_unmatched_returns_none(): @@ -71,24 +71,23 @@ def test_match_moe_block_unmatched_returns_none(): def test_get_expert_linear_names_raises_when_unmatched(): - # No spec matches and no fused-expert structure — must fail loudly instead of - # guessing another model's naming (the legacy w1/w2/w3 default was removed). + # No spec for these model types and no fused-expert structure — must fail loudly + # instead of guessing another model's naming (the legacy w1/w2/w3 default was + # removed). with pytest.raises(NotImplementedError, match="expert linear names"): - get_expert_linear_names(_UnknownMoeBlock()) + get_expert_linear_names(_UnknownMoeBlock(), {"some_unknown_model"}) def test_get_expert_linear_names_from_specs(): - class ArcticMoE(nn.Module): - pass - - class DbrxFFN(nn.Module): - pass - # Arctic keeps the w1/w2/w3 naming it previously got from the engine default. - assert get_expert_linear_names(ArcticMoE()) == ["w1", "w2", "w3"] - # DbrxFFN resolves the quantized per-expert ModuleList names (previously it fell + assert get_expert_linear_names(_UnknownMoeBlock(), {"arctic"}) == ["w1", "w2", "w3"] + # DBRX resolves the quantized per-expert ModuleList names (previously it fell # through to the w1/w2/w3 default, which never existed on the quantized module). - assert get_expert_linear_names(DbrxFFN()) == ["w1_linear", "w2_linear", "v1_linear"] + assert get_expert_linear_names(_UnknownMoeBlock(), {"dbrx"}) == [ + "w1_linear", + "w2_linear", + "v1_linear", + ] class NemotronHMOE(nn.Module): @@ -106,7 +105,7 @@ def __init__(self): def test_get_experts_list_groups_by_spec_linear_names(): module = NemotronHMOE() - groups = get_experts_list(module, "nemotronhforcausallm") + groups = get_experts_list(module, {"nemotron_h"}) assert len(groups) == 2 # up_proj group + down_proj group assert all(len(group) == 3 for group in groups) assert groups[0][0] is module.experts[0].up_proj @@ -116,16 +115,16 @@ def test_get_experts_list_groups_by_spec_linear_names(): def test_get_experts_list_rejects_non_iterable_layouts(): # DBRX matches a spec but is not an iterable-experts layout; grouped export # must keep rejecting it (legacy behavior). - class DBRXMoeSparseMoeBlock(nn.Module): + class DbrxFFN(nn.Module): def __init__(self): super().__init__() self.experts = nn.ModuleList() with pytest.raises(NotImplementedError): - get_experts_list(DBRXMoeSparseMoeBlock(), "dbrxforcausallm") + get_experts_list(DbrxFFN(), {"dbrx"}) with pytest.raises(NotImplementedError): - get_experts_list(_UnknownMoeBlock(), "unknownforcausallm") + get_experts_list(_UnknownMoeBlock(), {"some_unknown_model"}) class ArcticMoE(nn.Module): @@ -265,27 +264,61 @@ class Qwen3MoeSparseMoeBlock(nn.Module): # A hypothetical remote-code fork registering the same block class name under # its own model type: scope must pick the model's own spec among candidates. - fork_spec = MoESpec( - model_type="zz_fork", + fork_variant = MoEVariant( block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("a_proj", "b_proj"), ) + fork_spec = MoESpec(model_type="zz_fork", variants=(fork_variant,)) _SPECS.append(fork_spec) try: - assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"zz_fork"}) is fork_spec - assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"qwen3_moe"}).model_type == "qwen3_moe" + assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"zz_fork"}) is fork_variant + assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"qwen3_moe"}).expert_linear_names == ( + "gate_proj", + "down_proj", + "up_proj", + ) # No scope -> first registered class-name match wins (legacy order). - assert match_moe_block(Qwen3MoeSparseMoeBlock()).model_type == "qwen3_moe" + assert match_moe_block(Qwen3MoeSparseMoeBlock()).expert_linear_names == ( + "gate_proj", + "down_proj", + "up_proj", + ) finally: _SPECS.remove(fork_spec) -def test_match_moe_block_scope_never_excludes(): +def test_match_moe_block_scope_is_strict(): class Qwen3MoeSparseMoeBlock(nn.Module): pass - # A tower whose model_type has no spec of its own must still resolve by class - # name (e.g. DeepseekMoE blocks inside a Kimi model). - spec = match_moe_block(Qwen3MoeSparseMoeBlock(), {"some_unknown_vlm"}) - assert spec is not None - assert spec.model_type == "qwen3_moe" + # A model whose model_type has no spec resolves to None even when its module + # class name coincides with another model's — register a spec instead of + # inheriting a neighbor's data. + assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"some_unknown_vlm"}) is None + # No scope (no config available) searches all specs. + assert match_moe_block(Qwen3MoeSparseMoeBlock()) is not None + + +def test_get_expert_linear_names_by_model_type_only(): + # With a scope, naming resolves from the model's own spec — the block class + # name is irrelevant (a spec need not declare block_names to provide naming). + assert get_expert_linear_names(_UnknownMoeBlock(), {"qwen3_moe"}) == [ + "gate_proj", + "down_proj", + "up_proj", + ] + with pytest.raises(NotImplementedError, match="model types"): + get_expert_linear_names(_UnknownMoeBlock(), {"some_unknown_vlm"}) + + +def test_mixtral_variants_disambiguated_by_block_class(): + class MixtralMoeSparseMoeBlock(nn.Module): + """Legacy-naming Mixtral layout — same model type, different projections.""" + + assert get_expert_linear_names(MixtralMoeSparseMoeBlock(), {"mixtral"}) == [ + "linear_fc1", + "linear_fc2", + ] + # An unrecognized block class under a multi-naming model type cannot resolve. + with pytest.raises(NotImplementedError): + get_expert_linear_names(_UnknownMoeBlock(), {"mixtral"}) From 58629a6fb58c4f7bf0d8dfe0ae3a2f5f69bc8d2c Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 16 Jul 2026 22:44:36 +0000 Subject: [PATCH 13/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/models/arctic.py | 6 +- modelopt/modeling/models/dbrx.py | 6 +- modelopt/modeling/models/deepseek.py | 6 +- modelopt/modeling/models/gemma.py | 8 +- modelopt/modeling/models/gemma4.py | 6 +- modelopt/modeling/models/gpt_oss.py | 6 +- modelopt/modeling/models/llama.py | 4 +- modelopt/modeling/models/mixtral.py | 6 +- modelopt/modeling/models/nemotron.py | 4 +- modelopt/modeling/models/nemotron_h.py | 6 +- modelopt/modeling/models/qwen2_moe.py | 6 +- modelopt/modeling/models/qwen3.py | 4 +- modelopt/modeling/models/qwen3_5_moe.py | 6 +- modelopt/modeling/models/qwen3_moe.py | 12 +-- modelopt/modeling/models/qwen3_next.py | 6 +- modelopt/modeling/registry.py | 79 +++++++-------- modelopt/modeling/specs.py | 95 ++++++++++--------- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 11 ++- modelopt/torch/export/layer_utils.py | 4 +- .../unit/torch/export/test_modeling_specs.py | 8 +- 20 files changed, 147 insertions(+), 142 deletions(-) diff --git a/modelopt/modeling/models/arctic.py b/modelopt/modeling/models/arctic.py index f29979e1108..11344b730c6 100644 --- a/modelopt/modeling/models/arctic.py +++ b/modelopt/modeling/models/arctic.py @@ -16,12 +16,12 @@ """Snowflake Arctic specs (trust-remote-code model type ``arctic``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant register( - MoESpec( + ModelSpec( model_type="arctic", - variants=( + moe_variants=( MoEVariant( # Non-standard block name (for is_moe identification). block_names=("ArcticMoE",), diff --git a/modelopt/modeling/models/dbrx.py b/modelopt/modeling/models/dbrx.py index 406d632215c..e2b74f3c36e 100644 --- a/modelopt/modeling/models/dbrx.py +++ b/modelopt/modeling/models/dbrx.py @@ -16,7 +16,7 @@ """DBRX specs (HF model type ``dbrx``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant # HF DbrxFFN (non-standard block name, identified for is_moe). Expert names refer to # the quantized layout: _QuantDbrxExpertGLU converts the fused w1/v1/w2 parameters @@ -24,9 +24,9 @@ # modelopt/torch/quantization/plugins/huggingface.py), which is what the DBRX # prepare handler fills amax values on. register( - MoESpec( + ModelSpec( model_type="dbrx", - variants=( + moe_variants=( MoEVariant( block_names=("DbrxFFN",), expert_linear_names=("w1_linear", "w2_linear", "v1_linear"), diff --git a/modelopt/modeling/models/deepseek.py b/modelopt/modeling/models/deepseek.py index 597f7cf4381..4b158bc5243 100644 --- a/modelopt/modeling/models/deepseek.py +++ b/modelopt/modeling/models/deepseek.py @@ -20,16 +20,16 @@ """ from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant # DeepseekMoE experts ARE structurally iterable (ModuleList of DeepseekMLP), but # has_iterable_experts stays False until the grouped export path (get_experts_list # resmoothing) is validated on this model — the flag currently doubles as that # support gate. register( - MoESpec( + ModelSpec( model_type="deepseek", - variants=( + moe_variants=( MoEVariant( block_names=("DeepseekMoE",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/models/gemma.py b/modelopt/modeling/models/gemma.py index 2abc89932bc..2380feb3d52 100644 --- a/modelopt/modeling/models/gemma.py +++ b/modelopt/modeling/models/gemma.py @@ -20,12 +20,12 @@ """ from ..registry import register -from ..specs import NormSpec +from ..specs import ModelSpec # Gemma RMSNorms store weight - 1 (the effective scale is weight + 1); scale-folding # engines (AWQ pre_quant_scale fusion into the norm) must account for the +1. # NOTE: Gemma4RMSNorm is intentionally absent to preserve legacy behavior; add it # once the +1 handling is validated on Gemma4. -register(NormSpec(model_type="gemma", weight_plus_one_norm_names=("GemmaRMSNorm",))) -register(NormSpec(model_type="gemma2", weight_plus_one_norm_names=("Gemma2RMSNorm",))) -register(NormSpec(model_type="gemma3", weight_plus_one_norm_names=("Gemma3RMSNorm",))) +register(ModelSpec(model_type="gemma", weight_plus_one_norm_names=("GemmaRMSNorm",))) +register(ModelSpec(model_type="gemma2", weight_plus_one_norm_names=("Gemma2RMSNorm",))) +register(ModelSpec(model_type="gemma3", weight_plus_one_norm_names=("Gemma3RMSNorm",))) diff --git a/modelopt/modeling/models/gemma4.py b/modelopt/modeling/models/gemma4.py index d3cf44e995b..e66d95843c5 100644 --- a/modelopt/modeling/models/gemma4.py +++ b/modelopt/modeling/models/gemma4.py @@ -16,12 +16,12 @@ """Gemma4 specs (HF model type ``gemma4``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant register( - MoESpec( + ModelSpec( model_type="gemma4", - variants=( + moe_variants=( MoEVariant( # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. block_names=("Gemma4TextDecoderLayer",), diff --git a/modelopt/modeling/models/gpt_oss.py b/modelopt/modeling/models/gpt_oss.py index 91910969534..8a316066469 100644 --- a/modelopt/modeling/models/gpt_oss.py +++ b/modelopt/modeling/models/gpt_oss.py @@ -16,12 +16,12 @@ """GPT-OSS specs (HF model type ``gpt_oss``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant register( - MoESpec( + ModelSpec( model_type="gpt_oss", - variants=( + moe_variants=( MoEVariant( # GPT-OSS fuses gate and up into a single gate_up_proj. block_names=("GptOssMoE",), diff --git a/modelopt/modeling/models/llama.py b/modelopt/modeling/models/llama.py index 4048db65737..13c02010a78 100644 --- a/modelopt/modeling/models/llama.py +++ b/modelopt/modeling/models/llama.py @@ -16,10 +16,10 @@ """Llama specs (HF model type ``llama``).""" from ..registry import register -from ..specs import ExportSpec +from ..specs import ModelSpec register( - ExportSpec( + ModelSpec( model_type="llama", # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( diff --git a/modelopt/modeling/models/mixtral.py b/modelopt/modeling/models/mixtral.py index ed4531224ae..f314b67758a 100644 --- a/modelopt/modeling/models/mixtral.py +++ b/modelopt/modeling/models/mixtral.py @@ -16,14 +16,14 @@ """Mixtral specs (HF model type ``mixtral``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant # Mixtral with iterable experts uses w1/w2/w3. Fused experts (transformers 5.0+) are # detected from their per-expert quantizer attributes and need no naming override here. register( - MoESpec( + ModelSpec( model_type="mixtral", - variants=( + moe_variants=( MoEVariant( block_names=("MixtralSparseMoeBlock",), expert_linear_names=("w1", "w2", "w3"), diff --git a/modelopt/modeling/models/nemotron.py b/modelopt/modeling/models/nemotron.py index aba932fa4e9..d92a1bf1d89 100644 --- a/modelopt/modeling/models/nemotron.py +++ b/modelopt/modeling/models/nemotron.py @@ -16,14 +16,14 @@ """Nemotron specs (HF model type ``nemotron``); Nemotron-H lives in ``nemotron_h.py``.""" from ..registry import register -from ..specs import NormSpec +from ..specs import ModelSpec # LayerNorm1P stores weight - 1 (zero-centered gamma). Both the plain Megatron-style # class name and the HF Nemotron port are listed; modules exposing a # ``zero_centered_gamma`` attribute are additionally caught by the engine's # structural fallback. register( - NormSpec( + ModelSpec( model_type="nemotron", weight_plus_one_norm_names=("LayerNorm1P", "NemotronLayerNorm1P"), ) diff --git a/modelopt/modeling/models/nemotron_h.py b/modelopt/modeling/models/nemotron_h.py index 6bf8622c342..2a8b1f45b1d 100644 --- a/modelopt/modeling/models/nemotron_h.py +++ b/modelopt/modeling/models/nemotron_h.py @@ -16,12 +16,12 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant register( - MoESpec( + ModelSpec( model_type="nemotron_h", - variants=( + moe_variants=( MoEVariant( # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). block_names=("NemotronHMOE",), diff --git a/modelopt/modeling/models/qwen2_moe.py b/modelopt/modeling/models/qwen2_moe.py index 89dd0437d7b..2c607a51695 100644 --- a/modelopt/modeling/models/qwen2_moe.py +++ b/modelopt/modeling/models/qwen2_moe.py @@ -16,12 +16,12 @@ """Qwen2-MoE specs (HF model type ``qwen2_moe``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant register( - MoESpec( + ModelSpec( model_type="qwen2_moe", - variants=( + moe_variants=( MoEVariant( block_names=("Qwen2MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/models/qwen3.py b/modelopt/modeling/models/qwen3.py index 85c2dd6342b..37f772121f2 100644 --- a/modelopt/modeling/models/qwen3.py +++ b/modelopt/modeling/models/qwen3.py @@ -16,10 +16,10 @@ """Qwen3 (dense) specs (HF model type ``qwen3``).""" from ..registry import register -from ..specs import ExportSpec +from ..specs import ModelSpec register( - ExportSpec( + ModelSpec( model_type="qwen3", # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( diff --git a/modelopt/modeling/models/qwen3_5_moe.py b/modelopt/modeling/models/qwen3_5_moe.py index e1c92ddf1e8..a6614baae83 100644 --- a/modelopt/modeling/models/qwen3_5_moe.py +++ b/modelopt/modeling/models/qwen3_5_moe.py @@ -16,12 +16,12 @@ """Qwen3.5-MoE specs (HF model type ``qwen3_5_moe``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant register( - MoESpec( + ModelSpec( model_type="qwen3_5_moe", - variants=( + moe_variants=( MoEVariant( block_names=("Qwen3_5MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/models/qwen3_moe.py b/modelopt/modeling/models/qwen3_moe.py index aee315d7168..bf4ac4fd6b2 100644 --- a/modelopt/modeling/models/qwen3_moe.py +++ b/modelopt/modeling/models/qwen3_moe.py @@ -16,12 +16,12 @@ """Qwen3-MoE specs (HF model type ``qwen3_moe``).""" from ..registry import register -from ..specs import ExportSpec, MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant register( - MoESpec( + ModelSpec( model_type="qwen3_moe", - variants=( + moe_variants=( MoEVariant( block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), @@ -29,12 +29,6 @@ has_iterable_experts=True, ), ), - ) -) - -register( - ExportSpec( - model_type="qwen3_moe", # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. pqs_fuse_rules=( (("Qwen3MoeAttention",), "v_proj", "o_proj"), diff --git a/modelopt/modeling/models/qwen3_next.py b/modelopt/modeling/models/qwen3_next.py index 002bece6093..7ff1ada8f6c 100644 --- a/modelopt/modeling/models/qwen3_next.py +++ b/modelopt/modeling/models/qwen3_next.py @@ -16,12 +16,12 @@ """Qwen3-Next specs (HF model type ``qwen3_next``).""" from ..registry import register -from ..specs import MoESpec, MoEVariant +from ..specs import ModelSpec, MoEVariant register( - MoESpec( + ModelSpec( model_type="qwen3_next", - variants=( + moe_variants=( MoEVariant( block_names=("Qwen3NextSparseMoeBlock",), expert_linear_names=("gate_proj", "down_proj", "up_proj"), diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index 44e235dd35a..3e1853a6919 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -13,26 +13,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Registry that resolves a model sub-module to its spec. +"""Registry indexing the per-model ``ModelSpec`` by HF model type. -Model modules register their specs at import time (see ``models/``). Lookups return -``None`` when nothing matches, so callers can fall back to their default behavior. +Model modules register their spec at import time (see ``models/``); one spec per +model type. Lookups return ``None`` (or an empty list) when nothing matches, so +callers can fail loudly or fall back per their own policy. -Matching is by class-name string only, so this package stays dependency-free (any -``nn.Module`` — or any object — can be passed to the lookups without importing torch -here). +Matching is by model-type and class-name strings only, so this package stays +dependency-free (any ``nn.Module`` — or any object — can be passed to the lookups +without importing torch here). """ from collections.abc import Iterator -from typing import TYPE_CHECKING, TypeVar +from typing import TYPE_CHECKING -from .specs import ExportSpec, ModelSpec, MoESpec, MoEVariant, NormSpec +from .specs import ModelSpec, MoEVariant if TYPE_CHECKING: import torch.nn as nn __all__ = [ "collect_model_types", + "get_spec", "get_specs", "iter_gate_up_pairs", "iter_pqs_fuse_rules", @@ -41,58 +43,59 @@ "weight_plus_one_norm_names", ] -SpecT = TypeVar("SpecT", bound=ModelSpec) +_SPECS: dict[str, ModelSpec] = {} -_SPECS: list[ModelSpec] = [] - -def register(spec: SpecT) -> SpecT: - """Register a model spec and return it.""" - _SPECS.append(spec) +def register(spec: ModelSpec) -> ModelSpec: + """Register a model spec and return it. One spec per model type.""" + if spec.model_type in _SPECS: + raise ValueError(f"ModelSpec for model type {spec.model_type!r} already registered") + _SPECS[spec.model_type] = spec return spec -def get_specs(spec_cls: type[SpecT], model_types: set[str] | None = None) -> list[SpecT]: - """Return the registered specs of type ``spec_cls``, in registration order. +def get_spec(model_type: str) -> ModelSpec | None: + """Return the spec registered for ``model_type``, or ``None``.""" + return _SPECS.get(model_type) + + +def get_specs(model_types: set[str] | None = None) -> list[ModelSpec]: + """Return registered specs, in registration order. This is the model_type index: with ``model_types`` (from ``collect_model_types(model.config)``), only the specs belonging to the model's own model types are returned, so consumers resolve per-model data without - scanning the registry. Without it, all specs of the type are returned - (aggregators, no-config compatibility). + scanning the registry. Without it, all specs are returned (aggregators, + no-config compatibility). """ - return [ - spec - for spec in _SPECS - if isinstance(spec, spec_cls) and (not model_types or spec.model_type in model_types) - ] + return [spec for spec in _SPECS.values() if not model_types or spec.model_type in model_types] def iter_pqs_fuse_rules(): """Yield every ``(module_class_substrings, fuse_into, fuse_from)`` AWQ fusion rule. - Aggregated across all registered export specs (the consumer matches each model - module against the substrings, so the order across specs does not matter). + Aggregated across all registered specs (the consumer matches each model module + against the substrings, so the order across specs does not matter). """ - for spec in get_specs(ExportSpec): + for spec in get_specs(): yield from spec.pqs_fuse_rules def iter_gate_up_pairs() -> Iterator[tuple[str, str]]: - """Yield the distinct (gate, up) projection-name pairs across all MoE specs. + """Yield the distinct (gate, up) projection-name pairs across all MoE variants. GLOBAL-VOCABULARY semantics: consumers (currently only calibration sibling grouping in ``quantization/model_calib.py``, which also walks dense MLPs that - no MoE spec can match) try every pair opportunistically on every module, + no MoE variant can match) try every pair opportunistically on every module, getattr-guarded. Adding a pair to any spec therefore changes behavior for ALL models whose modules happen to carry those attribute names — prefer per-module resolution (``match_moe_block(module).gate_up_pair``) wherever the module is an - identifiable MoE block. The dense-MLP case moves to a fusion-group topic spec - in a follow-up (see MODEL_SPECIFIC_REFACTOR.md P5). + identifiable MoE block. The dense-MLP case moves to a fusion-group topic + section in a follow-up (see MODEL_SPECIFIC_REFACTOR.md P5). """ seen = set() - for spec in get_specs(MoESpec): - for variant in spec.variants: + for spec in get_specs(): + for variant in spec.moe_variants: pair = variant.gate_up_pair if pair is not None and pair not in seen: seen.add(pair) @@ -100,8 +103,8 @@ def iter_gate_up_pairs() -> Iterator[tuple[str, str]]: def weight_plus_one_norm_names() -> tuple[str, ...]: - """All norm class names whose stored weight is ``w - 1``, across all norm specs.""" - return tuple(name for spec in get_specs(NormSpec) for name in spec.weight_plus_one_norm_names) + """All norm class names whose stored weight is ``w - 1``, across all specs.""" + return tuple(name for spec in get_specs() for name in spec.weight_plus_one_norm_names) def collect_model_types(config) -> set[str]: @@ -144,11 +147,11 @@ def match_moe_block(module: "nn.Module", model_types: set[str] | None = None) -> the TRT-LLM path) searches all specs. Within the scope, each spec's variant ``block_names`` identifies the block and - disambiguates same-model layout variants (``MoESpec.match_variant``); quantized - wrapper classes match through their original base class in the MRO. + disambiguates same-model layout variants (``MoESpec.match_moe_variant``); + quantized wrapper classes match through their original base class in the MRO. """ - for spec in get_specs(MoESpec, model_types): - variant = spec.match_variant(module) + for spec in get_specs(model_types): + variant = spec.match_moe_variant(module) if variant is not None: return variant return None diff --git a/modelopt/modeling/specs.py b/modelopt/modeling/specs.py index 814cf2a9dcf..db99a2168e9 100644 --- a/modelopt/modeling/specs.py +++ b/modelopt/modeling/specs.py @@ -15,35 +15,30 @@ """Per-model descriptor classes. -``ModelSpec`` subclasses come in two kinds, both registered per model so consumers -read values instead of branching on model names: +``ModelSpec`` is the one global descriptor of a model: everything modelopt knows +about a model type lives on a single instance, resolved by ``config.model_type``. +It is composed from section mixins so each concern stays a small, separate class: -- **topic specs** hold architecture facts shared across subsystems (``MoESpec``: +- **topic sections** hold architecture facts shared across subsystems (``MoESpec``: what a model's MoE blocks are; ``NormSpec``: norm-layer conventions); -- **subsystem specs** hold one subsystem's per-model policy (``ExportSpec``; - quantization / speculative-decoding specs to follow). +- **subsystem sections** hold one subsystem's per-model policy (``ExportSpec``; + quantization / speculative-decoding sections to follow). -Specs hold per-model data plus trivial accessors over that data; subsystem logic -never lives here. +Sections hold per-model data plus trivial accessors over that data; subsystem logic +never lives here. A model file registers exactly one ``ModelSpec``, filling only the +sections it customizes (see ``models/``). """ from dataclasses import dataclass -__all__ = ["ExportSpec", "MoESpec", "MoEVariant", "ModelSpec", "NormSpec", "match_class_names"] - - -@dataclass -class ModelSpec: - """Base class for per-model data specs. - - Subclasses add the data fields of one topic or subsystem; a model registers one - spec instance per kind it customizes (see ``models/``). - """ - - model_type: str - """The HF model type this spec belongs to (``config.model_type``, e.g. - ``"qwen3_moe"``). A model registers one spec instance per spec kind; same-model - layout variants nest inside the spec (see ``MoESpec.variants``).""" +__all__ = [ + "ExportSpec", + "MoESpec", + "MoEVariant", + "ModelSpec", + "NormSpec", + "match_class_names", +] def match_class_names(module, names: tuple[str, ...]) -> bool: @@ -62,7 +57,7 @@ def match_class_names(module, names: tuple[str, ...]) -> bool: return any(name.lower() in mro_names for name in names) -@dataclass +@dataclass(kw_only=True) class MoEVariant: """One concrete MoE-block layout of a model. @@ -95,23 +90,23 @@ class MoEVariant: ``sync_moe_gate_up_amax``) and by calibration grouping.""" -@dataclass -class MoESpec(ModelSpec): - """MoE architecture facts for one model: its MoE-block layout variant(s). +@dataclass(kw_only=True) +class MoESpec: + """Topic section: MoE architecture facts — the model's MoE-block layout(s). - Unlike the subsystem specs, this describes what a model's MoE blocks *are* — - which class, what the expert projections are called — so any modelopt subsystem - (export, quantization, speculative decoding, ...) can read it instead of keeping - its own per-model MoE table. + This describes what a model's MoE blocks *are* — which class, what the expert + projections are called — so any modelopt subsystem (export, quantization, + speculative decoding, ...) can read it instead of keeping its own per-model MoE + table. """ - variants: tuple[MoEVariant, ...] = () + moe_variants: tuple[MoEVariant, ...] = () """The model's MoE-block layouts; more than one when the same checkpoint materializes differently (see ``MoEVariant``).""" - def match_variant(self, module) -> MoEVariant | None: + def match_moe_variant(self, module) -> MoEVariant | None: """Return the variant whose ``block_names`` matches ``module``, else None.""" - for variant in self.variants: + for variant in self.moe_variants: if match_class_names(module, variant.block_names): return variant return None @@ -125,21 +120,18 @@ def expert_linear_names_for(self, module) -> tuple[str, ...] | None: """ namings = { variant.expert_linear_names - for variant in self.variants + for variant in self.moe_variants if variant.expert_linear_names is not None } if len(namings) == 1: return next(iter(namings)) - variant = self.match_variant(module) + variant = self.match_moe_variant(module) return variant.expert_linear_names if variant is not None else None -@dataclass -class NormSpec(ModelSpec): - """Normalization-layer architecture facts for one model. - - Topic spec (like ``MoESpec``): shared facts any subsystem can read. - """ +@dataclass(kw_only=True) +class NormSpec: + """Topic section: normalization-layer architecture facts.""" weight_plus_one_norm_names: tuple[str, ...] = () """Class names of norm layers whose stored weight is ``w - 1`` (the effective @@ -150,12 +142,12 @@ class NormSpec(ModelSpec): in the engine.""" -@dataclass -class ExportSpec(ModelSpec): - """Per-model policy for the unified HF export path. +@dataclass(kw_only=True) +class ExportSpec: + """Subsystem section: per-model policy of the unified HF export path. Architecture facts (MoE block classes, expert naming) live in ``MoESpec``; this - spec holds decisions that belong to the export/quantization algorithms only. + section holds decisions that belong to the export/quantization algorithms only. """ pqs_fuse_rules: tuple[tuple[tuple[str, ...], str, str], ...] = () @@ -165,3 +157,16 @@ class ExportSpec(ModelSpec): (e.g. attention ``o_proj`` -> ``v_proj``, MLP ``down_proj`` -> ``up_proj``). A rule is a validated mathematical-equivalence claim for that model's modules, which is why it is declared per model rather than applied generically.""" + + +@dataclass(kw_only=True) +class ModelSpec(MoESpec, NormSpec, ExportSpec): + """The one global per-model descriptor, composed from the section mixins. + + Resolved by HF model type (see ``registry.get_spec``); a model registers exactly + one instance, filling only the sections it customizes. + """ + + model_type: str + """The HF model type this spec describes (``config.model_type``, e.g. + ``"qwen3_moe"``). Unique across the registry.""" diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index 710d3dd8d28..01c35733e09 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -52,10 +52,10 @@ Two layers: ```text modelopt/modeling/ - specs.py # ModelSpec (base) + topic specs (MoESpec/NormSpec: shared - # architecture facts; MoESpec nests one MoEVariant per concrete - # block layout) + subsystem specs (ExportSpec: HF-export-path - # policy); future topic/subsystem specs go here too + specs.py # ModelSpec — the ONE global per-model descriptor, composed + # from section mixins: topic sections (MoESpec: MoE layouts as + # MoEVariant tuples; NormSpec) + subsystem sections (ExportSpec); + # future sections mix in the same way registry.py # register() + lookups queried by spec type (None when unmatched) # + the MRO exact-name matching core (match_class_names) __init__.py # re-exports; importing it registers all specs @@ -69,6 +69,9 @@ Model type names mirror (e.g. `qwen3_moe.py`, `gpt_oss.py`, `nemotron_h.py`); trust-remote-code models (`arctic`, `deepseek`) use their config `model_type`. +Each model registers exactly ONE `ModelSpec` (registry enforces uniqueness), so +`get_spec(model_type)` is a dict lookup and consumers call methods directly on the +global spec (e.g. `spec.expert_linear_names_for(module)`). Resolution is by model type first, mirroring HF's own indexing: the engine collects the model's HF model types once per export (`collect_model_types(model.config)` — root plus sub-configs, so VLM towers diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index ff6d447cc4b..ad425b24383 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -28,7 +28,7 @@ except Exception: warn("Cannot find transformers package. Hugginface modules cannot be exported.") -from modelopt.modeling import MoESpec, collect_model_types, get_specs, match_moe_block +from modelopt.modeling import collect_model_types, get_specs, match_moe_block from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import import_plugin @@ -997,7 +997,7 @@ class disambiguate. Raises NotImplementedError when nothing resolves, so a new resolved = { names - for spec in get_specs(MoESpec, model_types) + for spec in get_specs(model_types) if (names := spec.expert_linear_names_for(module)) is not None } if len(resolved) == 1: diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index 0ba24f149b1..cf3c0d994b6 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -19,7 +19,7 @@ import torch.nn as nn from modelopt.modeling import ( - MoESpec, + ModelSpec, MoEVariant, collect_model_types, iter_gate_up_pairs, @@ -268,8 +268,8 @@ class Qwen3MoeSparseMoeBlock(nn.Module): block_names=("Qwen3MoeSparseMoeBlock",), expert_linear_names=("a_proj", "b_proj"), ) - fork_spec = MoESpec(model_type="zz_fork", variants=(fork_variant,)) - _SPECS.append(fork_spec) + fork_spec = ModelSpec(model_type="zz_fork", moe_variants=(fork_variant,)) + _SPECS[fork_spec.model_type] = fork_spec try: assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"zz_fork"}) is fork_variant assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"qwen3_moe"}).expert_linear_names == ( @@ -284,7 +284,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module): "up_proj", ) finally: - _SPECS.remove(fork_spec) + del _SPECS[fork_spec.model_type] def test_match_moe_block_scope_is_strict(): From b8dc2cf2aa10e1754a65173d6436c99d77c82b82 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:01:33 +0000 Subject: [PATCH 14/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- modelopt/modeling/models/gemma4.py | 27 ++++---- modelopt/modeling/registry.py | 68 +++++-------------- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 20 +++--- modelopt/torch/export/hf_export_handlers.py | 4 +- modelopt/torch/export/layer_utils.py | 68 ++++++++----------- .../torch/export/plugins/vllm_fakequant_hf.py | 7 +- modelopt/torch/export/registry.py | 8 +-- modelopt/torch/export/unified_export_hf.py | 14 ++-- tests/unit/torch/export/test_layer_utils.py | 6 +- .../unit/torch/export/test_modeling_specs.py | 53 ++++++++------- 10 files changed, 119 insertions(+), 156 deletions(-) diff --git a/modelopt/modeling/models/gemma4.py b/modelopt/modeling/models/gemma4.py index e66d95843c5..be473888cd7 100644 --- a/modelopt/modeling/models/gemma4.py +++ b/modelopt/modeling/models/gemma4.py @@ -18,17 +18,18 @@ from ..registry import register from ..specs import ModelSpec, MoEVariant -register( - ModelSpec( - model_type="gemma4", - moe_variants=( - MoEVariant( - # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. - block_names=("Gemma4TextDecoderLayer",), - expert_linear_names=("gate_proj", "down_proj", "up_proj"), - gate_up_pair=("gate_proj", "up_proj"), - has_iterable_experts=True, - ), - ), - ) +# Gemma4 MoE experts are unfused into per-expert nn.Linear layers. The MoE block +# lives in the text model, so the same layout is registered for both the VLM root +# type and the text type (a text-only checkpoint's root model_type is +# ``gemma4_text``, following the gemma3 precedent). +_GEMMA4_MOE_VARIANTS = ( + MoEVariant( + block_names=("Gemma4TextDecoderLayer",), + expert_linear_names=("gate_proj", "down_proj", "up_proj"), + gate_up_pair=("gate_proj", "up_proj"), + has_iterable_experts=True, + ), ) + +register(ModelSpec(model_type="gemma4", moe_variants=_GEMMA4_MOE_VARIANTS)) +register(ModelSpec(model_type="gemma4_text", moe_variants=_GEMMA4_MOE_VARIANTS)) diff --git a/modelopt/modeling/registry.py b/modelopt/modeling/registry.py index 3e1853a6919..f6ba4621bc5 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/modeling/registry.py @@ -33,7 +33,6 @@ import torch.nn as nn __all__ = [ - "collect_model_types", "get_spec", "get_specs", "iter_gate_up_pairs", @@ -59,16 +58,9 @@ def get_spec(model_type: str) -> ModelSpec | None: return _SPECS.get(model_type) -def get_specs(model_types: set[str] | None = None) -> list[ModelSpec]: - """Return registered specs, in registration order. - - This is the model_type index: with ``model_types`` (from - ``collect_model_types(model.config)``), only the specs belonging to the model's - own model types are returned, so consumers resolve per-model data without - scanning the registry. Without it, all specs are returned (aggregators, - no-config compatibility). - """ - return [spec for spec in _SPECS.values() if not model_types or spec.model_type in model_types] +def get_specs() -> list[ModelSpec]: + """Return all registered specs, in registration order (aggregators, tests).""" + return list(_SPECS.values()) def iter_pqs_fuse_rules(): @@ -107,50 +99,26 @@ def weight_plus_one_norm_names() -> tuple[str, ...]: return tuple(name for spec in get_specs() for name in spec.weight_plus_one_norm_names) -def collect_model_types(config) -> set[str]: - """Collect every HF model type in a config tree (root plus nested sub-configs). - - Walks attribute values duck-typed as configs (anything exposing a string - ``model_type``), so composite models contribute their tower types too — e.g. a - VLM yields ``{"kimi_vl", "kimi_k2"}`` via ``config.text_config`` — without this - package importing transformers. Pass ``model.config``; ``None`` yields an empty - set. - """ - found: set[str] = set() - seen: set[int] = set() - - def _walk(cfg) -> None: - if id(cfg) in seen: - return - seen.add(id(cfg)) - model_type = getattr(cfg, "model_type", None) - if isinstance(model_type, str) and model_type: - found.add(model_type) - for value in vars(cfg).values(): - if isinstance(getattr(value, "model_type", None), str): - _walk(value) - - if config is not None: - _walk(config) - return found - - -def match_moe_block(module: "nn.Module", model_types: set[str] | None = None) -> MoEVariant | None: +def match_moe_block(module: "nn.Module", model_type: str | None = None) -> MoEVariant | None: """Return the MoE layout variant for ``module``, resolved by model type. - ``model_types`` (from ``collect_model_types(model.config)``: root plus - sub-config types, so VLM towers are covered) is a strict filter: only specs - registered under the model's own model types are considered. A model whose + ``model_type`` is the model's root HF type (``model.config.model_type``) and is + a strict filter: only that model's own spec is consulted. A model whose model_type has no spec resolves to ``None`` even if its module class names coincide with another model's — register a spec instead of inheriting a - neighbor's data. ``None`` or an empty set (no config available: unit tests, - the TRT-LLM path) searches all specs. - - Within the scope, each spec's variant ``block_names`` identifies the block and - disambiguates same-model layout variants (``MoESpec.match_moe_variant``); - quantized wrapper classes match through their original base class in the MRO. + neighbor's data. ``None`` (no config available: unit tests, the TRT-LLM path) + searches all specs. Sub-model types of composite models are not walked today; + a composite whose MoE lives under a tower type registers the root type too + (see ``gemma4.py``). + + Within the spec, variant ``block_names`` identifies the block and disambiguates + same-model layout variants (``MoESpec.match_moe_variant``); quantized wrapper + classes match through their original base class in the MRO. """ - for spec in get_specs(model_types): + if model_type: + spec = get_spec(model_type) + return spec.match_moe_variant(module) if spec is not None else None + for spec in get_specs(): variant = spec.match_moe_variant(module) if variant is not None: return variant diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md index 01c35733e09..b5dabb4f232 100644 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md @@ -72,20 +72,22 @@ Model type names mirror Each model registers exactly ONE `ModelSpec` (registry enforces uniqueness), so `get_spec(model_type)` is a dict lookup and consumers call methods directly on the global spec (e.g. `spec.expert_linear_names_for(module)`). -Resolution is by model type first, mirroring HF's own indexing: the engine -collects the model's HF model types once per export -(`collect_model_types(model.config)` — root plus sub-configs, so VLM towers -contribute e.g. `kimi_k2` explicitly) and passes them down. The scope is strict: -only the model's own specs are consulted, so a model whose model_type has no spec -fails loudly (register a spec) instead of inheriting a neighbor's data through a -coincidental class-name match. Within the scope, each `MoESpec` nests one +Resolution is by model type, mirroring HF's own indexing: the engine reads the +model's root HF type (`model.config.model_type`) once per export and passes it +down. The lookup is strict: only the model's own spec is consulted, so a model +whose model_type has no spec fails loudly (register a spec) instead of inheriting +a neighbor's data through a coincidental class-name match. Sub-config model types +of composite models are not walked today — a composite whose MoE lives under a +tower type registers the root type too (`gemma4` + `gemma4_text`, following the +`gemma3`/`gemma3_text` precedent); a recursive config walk is a future extension +if a real VLM tower needs it. Within the spec, each `MoESpec` nests one `MoEVariant` per concrete block layout — several when the same checkpoint materializes with different classes and projection names (Mixtral across transformers generations); variant `block_names` (matched against the module MRO) identify MoE blocks (`is_moe`) and pick the variant. `get_expert_linear_names` doesn't need the block class at all when a model's -variants agree on one naming. Without a scope (no config available: unit tests, -the TRT-LLM path), lookups search all specs by class name. +variants agree on one naming. Without a model type (no config available: unit +tests, the TRT-LLM path), lookups search all specs by class name. During migration, call sites kept the legacy branches as a fallback behind the spec lookup. Once the specs covered every family the legacy chains served, the diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 6e7d3389c3d..6d6b571285c 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -58,7 +58,7 @@ def _export_weight( def _prepare_dbrx_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None: """Fill missing input amax values for DBRX per-expert ModuleLists.""" experts_mlp = moe_module.experts.mlp - for linear_name in get_expert_linear_names(moe_module, ctx.model_types): + for linear_name in get_expert_linear_names(moe_module, ctx.model_type): if hasattr(experts_mlp, linear_name): linear_modulelist = getattr(experts_mlp, linear_name) if hasattr(linear_modulelist, "__iter__"): @@ -94,7 +94,7 @@ def _prepare_bmm_experts(name: str, moe_module: nn.Module, ctx: ExportContext) - ) def _prepare_iterable_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None: """Fill missing input amax values for iterable per-expert submodules.""" - expert_linear_names = get_expert_linear_names(moe_module, ctx.model_types) + expert_linear_names = get_expert_linear_names(moe_module, ctx.model_type) linear_name = None try: for linear_name in expert_linear_names: diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index ad425b24383..70a7a2449ad 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -28,7 +28,7 @@ except Exception: warn("Cannot find transformers package. Hugginface modules cannot be exported.") -from modelopt.modeling import collect_model_types, get_specs, match_moe_block +from modelopt.modeling import get_spec, match_moe_block from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import import_plugin @@ -84,14 +84,14 @@ def get_experts_list( module: torch.nn.Module, - model_types: set[str], + model_type: str | None, ): """Returns list of grouped experts by linear name for given module. Args: module: MoE block (e.g. MixtralSparseMoeBlock, NemotronHMOE). - model_types: the model's HF model types (``collect_model_types(model.config)``), - used to resolve the model's own specs. + model_type: the model's HF model type (``model.config.model_type``), used to + resolve the model's own spec. """ experts_list = [] @@ -99,13 +99,12 @@ def get_experts_list( # per-expert sub-modules (spec.has_iterable_experts); stacked/fused layouts # (DBRX, GptOss, ...) raise NotImplementedError here and are handled by other # paths. Name resolution itself is shared with get_expert_linear_names. - variant = match_moe_block(module, model_types) + variant = match_moe_block(module, model_type) if variant is None or not variant.has_iterable_experts: raise NotImplementedError( - f"MoE block {type(module).__name__!r} " - f"(model types: {sorted(model_types or [])}) not supported" + f"MoE block {type(module).__name__!r} (model type: {model_type!r}) not supported" ) - linear_names = get_expert_linear_names(module, model_types) + linear_names = get_expert_linear_names(module, model_type) # Common logic for all supported model types experts_list.extend( @@ -304,19 +303,19 @@ def is_mlp(module: nn.Module) -> bool: return any(key in type(module).__name__.upper() for key in ("MLP", "T5DENSE")) -def is_moe(module: nn.Module, model_types: set[str] | None = None) -> bool: +def is_moe(module: nn.Module, model_type: str | None = None) -> bool: """Returns whether the module is an MOE layer. - ``model_types`` (``collect_model_types(model.config)``) strictly scopes the - registry lookup to the model's own specs; the generic naming conventions and the - structural check below are unaffected by it. + ``model_type`` (``model.config.model_type``) strictly scopes the registry lookup + to the model's own spec; the generic naming conventions and the structural check + below are unaffected by it. """ name = type(module).__name__.lower() # Auto-detect common MoE patterns if name.endswith("sparsemoeblock") or "moelayer" in name: return True # Non-standard MoE block names are per-model data (modelopt/modeling/models/*). - if match_moe_block(module, model_types) is not None: + if match_moe_block(module, model_type) is not None: return True # Structural detection: modules with router + experts (e.g. Gemma4TextDecoderLayer) return ( @@ -972,20 +971,21 @@ def get_stacked_scaling_factors(experts, get_function, module_name): return config -def get_expert_linear_names(module: nn.Module, model_types: set[str]) -> list[str]: +def get_expert_linear_names(module: nn.Module, model_type: str | None) -> list[str]: """Get the list of linear names for the experts. Args: module: the MoE block. - model_types: the model's HF model types (``collect_model_types(model.config)``). + model_type: the model's HF model type (``model.config.model_type``). Resolution order: structural detection of fused-expert layouts first (runtime state: transformers>=5 fused experts, and modules rewritten during export), then - the model's own specs by model type — the block class name is not needed, so a - spec can provide expert naming without declaring ``block_names``. Only when a - composite model carries several MoE towers with different namings does the block - class disambiguate. Raises NotImplementedError when nothing resolves, so a new - MoE model fails loudly instead of silently inheriting another model's naming. + the model's own spec by model type — the block class name is not needed, so a + spec can provide expert naming without declaring ``block_names``; the block + class only disambiguates same-model layout variants (see + ``MoESpec.expert_linear_names_for``). Raises NotImplementedError when nothing + resolves, so a new MoE model fails loudly instead of silently inheriting another + model's naming. """ # Structural detection: after _export_fused_experts, fused expert modules # have per-expert submodules with gate_proj/up_proj/down_proj. @@ -995,26 +995,16 @@ class disambiguate. Raises NotImplementedError when nothing resolves, so a new if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"): return [first_proj_attr, "down_proj"] - resolved = { - names - for spec in get_specs(model_types) - if (names := spec.expert_linear_names_for(module)) is not None - } - if len(resolved) == 1: - return list(next(iter(resolved))) - if len(resolved) > 1: - # Several MoE model types in scope resolving differently — no known real - # case; fail loudly rather than pick one. - raise NotImplementedError( - f"Ambiguous expert linear names for MoE block {type(module).__name__!r}: " - f"model types {sorted(model_types or [])} resolve to {sorted(resolved)}." - ) + spec = get_spec(model_type) if model_type else None + if spec is not None: + names = spec.expert_linear_names_for(module) + if names is not None: + return list(names) raise NotImplementedError( f"Cannot resolve expert linear names for MoE block {type(module).__name__!r} " - f"(model types: {sorted(model_types) if model_types else 'unknown'}). " - "Register a MoESpec with expert_linear_names for this model under " - "modelopt/modeling/models/." + f"(model type: {model_type!r}). Register a ModelSpec with moe_variants for " + "this model under modelopt/modeling/models/." ) @@ -1199,7 +1189,7 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: Returns: Number of expert gate/up pairs whose amaxes were synced. """ - model_types = collect_model_types(getattr(model, "config", None)) + model_type = getattr(getattr(model, "config", None), "model_type", None) synced = 0 unmatched_block_names: set[str] = set() for _, sub_module in model.named_modules(): @@ -1211,7 +1201,7 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: # cross-model guessing. A variant declaring no pair (non-gated NemotronH, # fused GptOss/DBRX) needs no sync; an unmatched block is warned about # once instead of silently skipped. - variant = match_moe_block(sub_module, model_types) + variant = match_moe_block(sub_module, model_type) if variant is None: unmatched_block_names.add(type(sub_module).__name__) continue diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index 927e7bee670..02688171441 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -27,7 +27,6 @@ import torch.nn as nn import modelopt.torch.opt as mto -from modelopt.modeling import collect_model_types from modelopt.torch.quantization.conversion import quantizer_state from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer @@ -409,7 +408,7 @@ def _resmooth_experts_for_export( name_to_module = dict(model.named_modules()) if inplace else None - model_types = collect_model_types(getattr(model, "config", None)) + model_type = getattr(getattr(model, "config", None), "model_type", None) id_to_name: dict[int, str] = {id(m): n for n, m in model.named_modules()} out: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {} requant_weights: set[str] = set() @@ -467,10 +466,10 @@ def _process_group(modules: list[nn.Module]) -> None: # different tokens to each expert, so forward hooks cannot detect them as # sharing the same input tensor. for _, module in model.named_modules(): - if not is_moe(module, model_types): + if not is_moe(module, model_type): continue try: - expert_groups = get_experts_list(module, model_types) + expert_groups = get_experts_list(module, model_type) except NotImplementedError: continue for experts in expert_groups: diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 20f751fab3b..e5e2131c0e4 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -54,10 +54,10 @@ class ExportContext: model: nn.Module dtype: torch.dtype is_modelopt_qlora: bool = False - model_types: set[str] = field(default_factory=set) - """The model's HF model types (root + sub-configs, via ``collect_model_types``), - used to resolve per-model specs in ``modelopt.modeling``. Empty means unknown: - spec lookups then fail loudly instead of guessing.""" + model_type: str | None = None + """The model's HF model type (``model.config.model_type``), used to resolve the + model's spec in ``modelopt.modeling``. ``None`` means unknown: spec lookups then + fail loudly instead of guessing.""" tied_cache: dict[int, nn.Module] = field(default_factory=dict) moe_tied_cache: dict[tuple[int, int], nn.Module] = field(default_factory=dict) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index a9294b72c35..232293997e2 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -30,8 +30,6 @@ from safetensors import safe_open from safetensors.torch import save_file -from modelopt.modeling import collect_model_types - from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales try: @@ -432,7 +430,7 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): # TODO: Handle DBRX MoE quantization_format = get_quantization_format(model) model_type = type(model).__name__.lower() - model_types = collect_model_types(getattr(model, "config", None)) + hf_model_type = getattr(getattr(model, "config", None), "model_type", None) module_names = set() # NVFP4 SVDQuant does not need pre-quant scale fusion (either into previous linear or layernorm) because @@ -448,12 +446,12 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): module_names.add(name) # For MoE models update pre_quant_scale to average pre_quant_scale amongst experts - if is_moe(module, model_types) and ( + if is_moe(module, hf_model_type) and ( quantization_format is not QUANTIZATION_NONE and ("awq" in quantization_format or quantization_format == QUANTIZATION_NVFP4_SVDQUANT) ): # update_experts_avg_prequant_scale(module) - grouped_experts = get_experts_list(module, model_types) + grouped_experts = get_experts_list(module, hf_model_type) for modules in grouped_experts: with fsdp2_aware_weight_update(model, modules): preprocess_linear_fusion(modules, resmooth_only=True) @@ -804,7 +802,7 @@ def _process_quantized_modules( model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora, - model_types=collect_model_types(getattr(model, "config", None)), + model_type=getattr(getattr(model, "config", None), "model_type", None), ) fsdp_module_to_reshard = None @@ -870,10 +868,10 @@ def _export_transformers_checkpoint( model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora, - model_types=collect_model_types(getattr(model, "config", None)), + model_type=getattr(getattr(model, "config", None), "model_type", None), ) for name, sub_module in model.named_modules(): - if is_moe(sub_module, prepare_ctx.model_types) and hasattr(sub_module, "experts"): + if is_moe(sub_module, prepare_ctx.model_type) and hasattr(sub_module, "experts"): handler = PrepareMoEInputsRegistry.match(sub_module.experts) if handler is None: # Unsupported MoE model structure diff --git a/tests/unit/torch/export/test_layer_utils.py b/tests/unit/torch/export/test_layer_utils.py index 8573bb99bed..044f210bff5 100644 --- a/tests/unit/torch/export/test_layer_utils.py +++ b/tests/unit/torch/export/test_layer_utils.py @@ -101,7 +101,7 @@ class NemotronHMOE(nn.Module): def test_get_expert_linear_names_gemma4(): - assert get_expert_linear_names(Gemma4TextDecoderLayer(), {"gemma4"}) == [ + assert get_expert_linear_names(Gemma4TextDecoderLayer(), "gemma4") == [ "gate_proj", "down_proj", "up_proj", @@ -109,8 +109,8 @@ def test_get_expert_linear_names_gemma4(): def test_get_expert_linear_names_mixtral(): - assert get_expert_linear_names(MixtralSparseMoeBlock(), {"mixtral"}) == ["w1", "w2", "w3"] + assert get_expert_linear_names(MixtralSparseMoeBlock(), "mixtral") == ["w1", "w2", "w3"] def test_get_expert_linear_names_nemotron(): - assert get_expert_linear_names(NemotronHMOE(), {"nemotron_h"}) == ["up_proj", "down_proj"] + assert get_expert_linear_names(NemotronHMOE(), "nemotron_h") == ["up_proj", "down_proj"] diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_modeling_specs.py index cf3c0d994b6..9264ae7ac1f 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_modeling_specs.py @@ -21,7 +21,6 @@ from modelopt.modeling import ( ModelSpec, MoEVariant, - collect_model_types, iter_gate_up_pairs, iter_pqs_fuse_rules, match_moe_block, @@ -75,15 +74,15 @@ def test_get_expert_linear_names_raises_when_unmatched(): # instead of guessing another model's naming (the legacy w1/w2/w3 default was # removed). with pytest.raises(NotImplementedError, match="expert linear names"): - get_expert_linear_names(_UnknownMoeBlock(), {"some_unknown_model"}) + get_expert_linear_names(_UnknownMoeBlock(), "some_unknown_model") def test_get_expert_linear_names_from_specs(): # Arctic keeps the w1/w2/w3 naming it previously got from the engine default. - assert get_expert_linear_names(_UnknownMoeBlock(), {"arctic"}) == ["w1", "w2", "w3"] + assert get_expert_linear_names(_UnknownMoeBlock(), "arctic") == ["w1", "w2", "w3"] # DBRX resolves the quantized per-expert ModuleList names (previously it fell # through to the w1/w2/w3 default, which never existed on the quantized module). - assert get_expert_linear_names(_UnknownMoeBlock(), {"dbrx"}) == [ + assert get_expert_linear_names(_UnknownMoeBlock(), "dbrx") == [ "w1_linear", "w2_linear", "v1_linear", @@ -105,7 +104,7 @@ def __init__(self): def test_get_experts_list_groups_by_spec_linear_names(): module = NemotronHMOE() - groups = get_experts_list(module, {"nemotron_h"}) + groups = get_experts_list(module, "nemotron_h") assert len(groups) == 2 # up_proj group + down_proj group assert all(len(group) == 3 for group in groups) assert groups[0][0] is module.experts[0].up_proj @@ -121,10 +120,10 @@ def __init__(self): self.experts = nn.ModuleList() with pytest.raises(NotImplementedError): - get_experts_list(DbrxFFN(), {"dbrx"}) + get_experts_list(DbrxFFN(), "dbrx") with pytest.raises(NotImplementedError): - get_experts_list(_UnknownMoeBlock(), {"some_unknown_model"}) + get_experts_list(_UnknownMoeBlock(), "some_unknown_model") class ArcticMoE(nn.Module): @@ -249,15 +248,6 @@ class UnknownSparseMoeBlock(nn.Module): assert sync_moe_gate_up_amax(model) == 0 -def test_collect_model_types_walks_sub_configs(): - from types import SimpleNamespace - - text_config = SimpleNamespace(model_type="kimi_k2") - config = SimpleNamespace(model_type="kimi_vl", text_config=text_config, hidden_size=4096) - assert collect_model_types(config) == {"kimi_vl", "kimi_k2"} - assert collect_model_types(None) == set() - - def test_match_moe_block_scope_prefers_own_model_type(): class Qwen3MoeSparseMoeBlock(nn.Module): pass @@ -271,8 +261,8 @@ class Qwen3MoeSparseMoeBlock(nn.Module): fork_spec = ModelSpec(model_type="zz_fork", moe_variants=(fork_variant,)) _SPECS[fork_spec.model_type] = fork_spec try: - assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"zz_fork"}) is fork_variant - assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"qwen3_moe"}).expert_linear_names == ( + assert match_moe_block(Qwen3MoeSparseMoeBlock(), "zz_fork") is fork_variant + assert match_moe_block(Qwen3MoeSparseMoeBlock(), "qwen3_moe").expert_linear_names == ( "gate_proj", "down_proj", "up_proj", @@ -294,7 +284,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module): # A model whose model_type has no spec resolves to None even when its module # class name coincides with another model's — register a spec instead of # inheriting a neighbor's data. - assert match_moe_block(Qwen3MoeSparseMoeBlock(), {"some_unknown_vlm"}) is None + assert match_moe_block(Qwen3MoeSparseMoeBlock(), "some_unknown_vlm") is None # No scope (no config available) searches all specs. assert match_moe_block(Qwen3MoeSparseMoeBlock()) is not None @@ -302,23 +292,38 @@ class Qwen3MoeSparseMoeBlock(nn.Module): def test_get_expert_linear_names_by_model_type_only(): # With a scope, naming resolves from the model's own spec — the block class # name is irrelevant (a spec need not declare block_names to provide naming). - assert get_expert_linear_names(_UnknownMoeBlock(), {"qwen3_moe"}) == [ + assert get_expert_linear_names(_UnknownMoeBlock(), "qwen3_moe") == [ "gate_proj", "down_proj", "up_proj", ] - with pytest.raises(NotImplementedError, match="model types"): - get_expert_linear_names(_UnknownMoeBlock(), {"some_unknown_vlm"}) + with pytest.raises(NotImplementedError, match="model type"): + get_expert_linear_names(_UnknownMoeBlock(), "some_unknown_vlm") def test_mixtral_variants_disambiguated_by_block_class(): class MixtralMoeSparseMoeBlock(nn.Module): """Legacy-naming Mixtral layout — same model type, different projections.""" - assert get_expert_linear_names(MixtralMoeSparseMoeBlock(), {"mixtral"}) == [ + assert get_expert_linear_names(MixtralMoeSparseMoeBlock(), "mixtral") == [ "linear_fc1", "linear_fc2", ] # An unrecognized block class under a multi-naming model type cannot resolve. with pytest.raises(NotImplementedError): - get_expert_linear_names(_UnknownMoeBlock(), {"mixtral"}) + get_expert_linear_names(_UnknownMoeBlock(), "mixtral") + + +def test_gemma4_both_root_types_resolve(): + # A gemma4 VLM's root model_type is gemma4; a text-only checkpoint's is + # gemma4_text (gemma3 precedent). Both register the same layout. + assert get_expert_linear_names(_UnknownMoeBlock(), "gemma4") == [ + "gate_proj", + "down_proj", + "up_proj", + ] + assert get_expert_linear_names(_UnknownMoeBlock(), "gemma4_text") == [ + "gate_proj", + "down_proj", + "up_proj", + ] From bdff2eaf44590762610962fa47a334e6392a0894 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:16:02 +0000 Subject: [PATCH 15/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- docs/source/reference/1_modelopt_api.rst | 2 +- modelopt/modeling/__init__.py | 29 ------------------- modelopt/{modeling => }/models/__init__.py | 14 +++++++-- modelopt/{modeling => }/models/arctic.py | 4 +-- modelopt/{modeling => }/models/dbrx.py | 4 +-- modelopt/{modeling => }/models/deepseek.py | 4 +-- modelopt/{modeling => }/models/gemma.py | 4 +-- modelopt/{modeling => }/models/gemma4.py | 4 +-- modelopt/{modeling => }/models/gpt_oss.py | 4 +-- modelopt/{modeling => }/models/llama.py | 4 +-- modelopt/{modeling => }/models/mixtral.py | 4 +-- modelopt/{modeling => }/models/nemotron.py | 4 +-- modelopt/{modeling => }/models/nemotron_h.py | 4 +-- modelopt/{modeling => }/models/qwen2_moe.py | 4 +-- modelopt/{modeling => }/models/qwen3.py | 4 +-- modelopt/{modeling => }/models/qwen3_5_moe.py | 4 +-- modelopt/{modeling => }/models/qwen3_moe.py | 4 +-- modelopt/{modeling => }/models/qwen3_next.py | 4 +-- modelopt/{modeling => models}/registry.py | 13 +++++++++ modelopt/{modeling => models}/specs.py | 0 modelopt/torch/export/layer_utils.py | 10 +++---- .../torch/export/plugins/vllm_fakequant_hf.py | 3 +- modelopt/torch/export/quant_utils.py | 4 +-- modelopt/torch/export/registry.py | 2 +- modelopt/torch/export/unified_export_hf.py | 12 ++++---- modelopt/torch/quantization/model_calib.py | 2 +- ..._modeling_specs.py => test_model_specs.py} | 17 +++++++++-- 27 files changed, 87 insertions(+), 81 deletions(-) delete mode 100644 modelopt/modeling/__init__.py rename modelopt/{modeling => }/models/__init__.py (64%) rename modelopt/{modeling => }/models/arctic.py (94%) rename modelopt/{modeling => }/models/dbrx.py (94%) rename modelopt/{modeling => }/models/deepseek.py (95%) rename modelopt/{modeling => }/models/gemma.py (95%) rename modelopt/{modeling => }/models/gemma4.py (95%) rename modelopt/{modeling => }/models/gpt_oss.py (93%) rename modelopt/{modeling => }/models/llama.py (94%) rename modelopt/{modeling => }/models/mixtral.py (95%) rename modelopt/{modeling => }/models/nemotron.py (94%) rename modelopt/{modeling => }/models/nemotron_h.py (93%) rename modelopt/{modeling => }/models/qwen2_moe.py (93%) rename modelopt/{modeling => }/models/qwen3.py (94%) rename modelopt/{modeling => }/models/qwen3_5_moe.py (93%) rename modelopt/{modeling => }/models/qwen3_moe.py (94%) rename modelopt/{modeling => }/models/qwen3_next.py (93%) rename modelopt/{modeling => models}/registry.py (90%) rename modelopt/{modeling => models}/specs.py (100%) rename tests/unit/torch/export/{test_modeling_specs.py => test_model_specs.py} (95%) diff --git a/docs/source/reference/1_modelopt_api.rst b/docs/source/reference/1_modelopt_api.rst index 8c1f5ab21c1..259594bf584 100644 --- a/docs/source/reference/1_modelopt_api.rst +++ b/docs/source/reference/1_modelopt_api.rst @@ -10,6 +10,6 @@ modelopt API :recursive: modelopt.deploy - modelopt.modeling + modelopt.models modelopt.onnx modelopt.torch diff --git a/modelopt/modeling/__init__.py b/modelopt/modeling/__init__.py deleted file mode 100644 index e106af55666..00000000000 --- a/modelopt/modeling/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Per-model descriptors, organized by HF model type. - -Holds declarative per-model data (no algorithms), one module per HF model type under -``models/``, mirroring ``transformers.models``. Architecture facts live in topic -specs shared across subsystems (``MoESpec``); per-subsystem policy lives in -subsystem specs (``ExportSpec``). Consumers resolve a spec via the registry lookups -and read its fields; an unmatched lookup returns ``None`` so callers fall back to -their default behavior. -""" - -# Importing models registers every spec as a side effect. -from . import models -from .registry import * -from .specs import * diff --git a/modelopt/modeling/models/__init__.py b/modelopt/models/__init__.py similarity index 64% rename from modelopt/modeling/models/__init__.py rename to modelopt/models/__init__.py index 6b0cb2f9315..98895211451 100644 --- a/modelopt/modeling/models/__init__.py +++ b/modelopt/models/__init__.py @@ -13,14 +13,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-model specs, one module per HF model type. Importing each module registers its specs. +"""Per-model descriptors, one module per HF model type. -Module names follow the HF ``transformers.models`` layout +Module names mirror the ``transformers.models`` layout (https://github.com/huggingface/transformers/tree/main/src/transformers/models); trust-remote-code models (``arctic``, ``deepseek``) use their config ``model_type``. +Each model module registers one global ``ModelSpec`` (see ``specs.py``); importing +this package registers them all. Consumers resolve a spec via the registry lookups +(``get_spec`` / ``match_moe_block``) and read its fields; an unmatched lookup +returns ``None`` so callers fail loudly or fall back per their own policy. """ -from . import ( +from .registry import * +from .specs import * + +# Importing the model modules registers every spec as a side effect. +from . import ( # isort: skip arctic, dbrx, deepseek, diff --git a/modelopt/modeling/models/arctic.py b/modelopt/models/arctic.py similarity index 94% rename from modelopt/modeling/models/arctic.py rename to modelopt/models/arctic.py index 11344b730c6..e914af3302f 100644 --- a/modelopt/modeling/models/arctic.py +++ b/modelopt/models/arctic.py @@ -15,8 +15,8 @@ """Snowflake Arctic specs (trust-remote-code model type ``arctic``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant register( ModelSpec( diff --git a/modelopt/modeling/models/dbrx.py b/modelopt/models/dbrx.py similarity index 94% rename from modelopt/modeling/models/dbrx.py rename to modelopt/models/dbrx.py index e2b74f3c36e..e02f0c9c01d 100644 --- a/modelopt/modeling/models/dbrx.py +++ b/modelopt/models/dbrx.py @@ -15,8 +15,8 @@ """DBRX specs (HF model type ``dbrx``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant # HF DbrxFFN (non-standard block name, identified for is_moe). Expert names refer to # the quantized layout: _QuantDbrxExpertGLU converts the fused w1/v1/w2 parameters diff --git a/modelopt/modeling/models/deepseek.py b/modelopt/models/deepseek.py similarity index 95% rename from modelopt/modeling/models/deepseek.py rename to modelopt/models/deepseek.py index 4b158bc5243..2ed7e4d84da 100644 --- a/modelopt/modeling/models/deepseek.py +++ b/modelopt/models/deepseek.py @@ -19,8 +19,8 @@ classes. """ -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant # DeepseekMoE experts ARE structurally iterable (ModuleList of DeepseekMLP), but # has_iterable_experts stays False until the grouped export path (get_experts_list diff --git a/modelopt/modeling/models/gemma.py b/modelopt/models/gemma.py similarity index 95% rename from modelopt/modeling/models/gemma.py rename to modelopt/models/gemma.py index 2380feb3d52..b045abe6fee 100644 --- a/modelopt/modeling/models/gemma.py +++ b/modelopt/models/gemma.py @@ -19,8 +19,8 @@ Gemma4 has its own module (``gemma4.py``) for its MoE spec. """ -from ..registry import register -from ..specs import ModelSpec +from .registry import register +from .specs import ModelSpec # Gemma RMSNorms store weight - 1 (the effective scale is weight + 1); scale-folding # engines (AWQ pre_quant_scale fusion into the norm) must account for the +1. diff --git a/modelopt/modeling/models/gemma4.py b/modelopt/models/gemma4.py similarity index 95% rename from modelopt/modeling/models/gemma4.py rename to modelopt/models/gemma4.py index be473888cd7..4164c452ea6 100644 --- a/modelopt/modeling/models/gemma4.py +++ b/modelopt/models/gemma4.py @@ -15,8 +15,8 @@ """Gemma4 specs (HF model type ``gemma4``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant # Gemma4 MoE experts are unfused into per-expert nn.Linear layers. The MoE block # lives in the text model, so the same layout is registered for both the VLM root diff --git a/modelopt/modeling/models/gpt_oss.py b/modelopt/models/gpt_oss.py similarity index 93% rename from modelopt/modeling/models/gpt_oss.py rename to modelopt/models/gpt_oss.py index 8a316066469..899118e0db5 100644 --- a/modelopt/modeling/models/gpt_oss.py +++ b/modelopt/models/gpt_oss.py @@ -15,8 +15,8 @@ """GPT-OSS specs (HF model type ``gpt_oss``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant register( ModelSpec( diff --git a/modelopt/modeling/models/llama.py b/modelopt/models/llama.py similarity index 94% rename from modelopt/modeling/models/llama.py rename to modelopt/models/llama.py index 13c02010a78..1c13855b126 100644 --- a/modelopt/modeling/models/llama.py +++ b/modelopt/models/llama.py @@ -15,8 +15,8 @@ """Llama specs (HF model type ``llama``).""" -from ..registry import register -from ..specs import ModelSpec +from .registry import register +from .specs import ModelSpec register( ModelSpec( diff --git a/modelopt/modeling/models/mixtral.py b/modelopt/models/mixtral.py similarity index 95% rename from modelopt/modeling/models/mixtral.py rename to modelopt/models/mixtral.py index f314b67758a..124ad67d277 100644 --- a/modelopt/modeling/models/mixtral.py +++ b/modelopt/models/mixtral.py @@ -15,8 +15,8 @@ """Mixtral specs (HF model type ``mixtral``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant # Mixtral with iterable experts uses w1/w2/w3. Fused experts (transformers 5.0+) are # detected from their per-expert quantizer attributes and need no naming override here. diff --git a/modelopt/modeling/models/nemotron.py b/modelopt/models/nemotron.py similarity index 94% rename from modelopt/modeling/models/nemotron.py rename to modelopt/models/nemotron.py index d92a1bf1d89..5f762fa20ab 100644 --- a/modelopt/modeling/models/nemotron.py +++ b/modelopt/models/nemotron.py @@ -15,8 +15,8 @@ """Nemotron specs (HF model type ``nemotron``); Nemotron-H lives in ``nemotron_h.py``.""" -from ..registry import register -from ..specs import ModelSpec +from .registry import register +from .specs import ModelSpec # LayerNorm1P stores weight - 1 (zero-centered gamma). Both the plain Megatron-style # class name and the HF Nemotron port are listed; modules exposing a diff --git a/modelopt/modeling/models/nemotron_h.py b/modelopt/models/nemotron_h.py similarity index 93% rename from modelopt/modeling/models/nemotron_h.py rename to modelopt/models/nemotron_h.py index 2a8b1f45b1d..4bd21f65eff 100644 --- a/modelopt/modeling/models/nemotron_h.py +++ b/modelopt/models/nemotron_h.py @@ -15,8 +15,8 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant register( ModelSpec( diff --git a/modelopt/modeling/models/qwen2_moe.py b/modelopt/models/qwen2_moe.py similarity index 93% rename from modelopt/modeling/models/qwen2_moe.py rename to modelopt/models/qwen2_moe.py index 2c607a51695..32c22128a06 100644 --- a/modelopt/modeling/models/qwen2_moe.py +++ b/modelopt/models/qwen2_moe.py @@ -15,8 +15,8 @@ """Qwen2-MoE specs (HF model type ``qwen2_moe``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant register( ModelSpec( diff --git a/modelopt/modeling/models/qwen3.py b/modelopt/models/qwen3.py similarity index 94% rename from modelopt/modeling/models/qwen3.py rename to modelopt/models/qwen3.py index 37f772121f2..35bd8c44d87 100644 --- a/modelopt/modeling/models/qwen3.py +++ b/modelopt/models/qwen3.py @@ -15,8 +15,8 @@ """Qwen3 (dense) specs (HF model type ``qwen3``).""" -from ..registry import register -from ..specs import ModelSpec +from .registry import register +from .specs import ModelSpec register( ModelSpec( diff --git a/modelopt/modeling/models/qwen3_5_moe.py b/modelopt/models/qwen3_5_moe.py similarity index 93% rename from modelopt/modeling/models/qwen3_5_moe.py rename to modelopt/models/qwen3_5_moe.py index a6614baae83..a795089acb8 100644 --- a/modelopt/modeling/models/qwen3_5_moe.py +++ b/modelopt/models/qwen3_5_moe.py @@ -15,8 +15,8 @@ """Qwen3.5-MoE specs (HF model type ``qwen3_5_moe``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant register( ModelSpec( diff --git a/modelopt/modeling/models/qwen3_moe.py b/modelopt/models/qwen3_moe.py similarity index 94% rename from modelopt/modeling/models/qwen3_moe.py rename to modelopt/models/qwen3_moe.py index bf4ac4fd6b2..fa626970fcb 100644 --- a/modelopt/modeling/models/qwen3_moe.py +++ b/modelopt/models/qwen3_moe.py @@ -15,8 +15,8 @@ """Qwen3-MoE specs (HF model type ``qwen3_moe``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant register( ModelSpec( diff --git a/modelopt/modeling/models/qwen3_next.py b/modelopt/models/qwen3_next.py similarity index 93% rename from modelopt/modeling/models/qwen3_next.py rename to modelopt/models/qwen3_next.py index 7ff1ada8f6c..216054eb877 100644 --- a/modelopt/modeling/models/qwen3_next.py +++ b/modelopt/models/qwen3_next.py @@ -15,8 +15,8 @@ """Qwen3-Next specs (HF model type ``qwen3_next``).""" -from ..registry import register -from ..specs import ModelSpec, MoEVariant +from .registry import register +from .specs import ModelSpec, MoEVariant register( ModelSpec( diff --git a/modelopt/modeling/registry.py b/modelopt/models/registry.py similarity index 90% rename from modelopt/modeling/registry.py rename to modelopt/models/registry.py index f6ba4621bc5..8ffbb31afe6 100644 --- a/modelopt/modeling/registry.py +++ b/modelopt/models/registry.py @@ -35,6 +35,7 @@ __all__ = [ "get_spec", "get_specs", + "hf_model_type", "iter_gate_up_pairs", "iter_pqs_fuse_rules", "match_moe_block", @@ -99,6 +100,18 @@ def weight_plus_one_norm_names() -> tuple[str, ...]: return tuple(name for spec in get_specs() for name in spec.weight_plus_one_norm_names) +def hf_model_type(model) -> str | None: + """Return the root HF model type (``model.config.model_type``), or ``None``. + + Accepts a model or a config object (duck-typed, no transformers import): a + model contributes via its ``config`` attribute, a config via its own + ``model_type``. This is the key for ``get_spec`` / ``match_moe_block``. + """ + config = getattr(model, "config", model) + model_type = getattr(config, "model_type", None) + return model_type if isinstance(model_type, str) else None + + def match_moe_block(module: "nn.Module", model_type: str | None = None) -> MoEVariant | None: """Return the MoE layout variant for ``module``, resolved by model type. diff --git a/modelopt/modeling/specs.py b/modelopt/models/specs.py similarity index 100% rename from modelopt/modeling/specs.py rename to modelopt/models/specs.py diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index 70a7a2449ad..98ec9cc6463 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -28,7 +28,7 @@ except Exception: warn("Cannot find transformers package. Hugginface modules cannot be exported.") -from modelopt.modeling import get_spec, match_moe_block +from modelopt.models import get_spec, hf_model_type, match_moe_block from modelopt.torch.utils import distributed as dist from modelopt.torch.utils import import_plugin @@ -314,7 +314,7 @@ def is_moe(module: nn.Module, model_type: str | None = None) -> bool: # Auto-detect common MoE patterns if name.endswith("sparsemoeblock") or "moelayer" in name: return True - # Non-standard MoE block names are per-model data (modelopt/modeling/models/*). + # Non-standard MoE block names are per-model data (modelopt/models/*). if match_moe_block(module, model_type) is not None: return True # Structural detection: modules with router + experts (e.g. Gemma4TextDecoderLayer) @@ -1004,7 +1004,7 @@ class only disambiguates same-model layout variants (see raise NotImplementedError( f"Cannot resolve expert linear names for MoE block {type(module).__name__!r} " f"(model type: {model_type!r}). Register a ModelSpec with moe_variants for " - "this model under modelopt/modeling/models/." + "this model under modelopt/models/." ) @@ -1189,7 +1189,7 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: Returns: Number of expert gate/up pairs whose amaxes were synced. """ - model_type = getattr(getattr(model, "config", None), "model_type", None) + model_type = hf_model_type(model) synced = 0 unmatched_block_names: set[str] = set() for _, sub_module in model.named_modules(): @@ -1243,7 +1243,7 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: f"MoE blocks {sorted(unmatched_block_names)} have no registered MoESpec; " "gate/up weight amax sync was skipped for them. If these models have " "gated experts with separate gate/up projections, register a MoESpec " - "with gate_up_pair under modelopt/modeling/models/ so the fused " + "with gate_up_pair under modelopt/models/ so the fused " "gate_up_proj weight scales stay consistent when serving." ) return synced diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index 02688171441..2e291a319f5 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -27,6 +27,7 @@ import torch.nn as nn import modelopt.torch.opt as mto +from modelopt.models import hf_model_type from modelopt.torch.quantization.conversion import quantizer_state from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer @@ -408,7 +409,7 @@ def _resmooth_experts_for_export( name_to_module = dict(model.named_modules()) if inplace else None - model_type = getattr(getattr(model, "config", None), "model_type", None) + model_type = hf_model_type(model) id_to_name: dict[int, str] = {id(m): n for n, m in model.named_modules()} out: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {} requant_weights: set[str] = set() diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 4b04cd2ad92..cf177a07026 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -25,7 +25,7 @@ import torch.nn as nn from modelopt import __version__ -from modelopt.modeling import iter_pqs_fuse_rules, match_class_names, weight_plus_one_norm_names +from modelopt.models import iter_pqs_fuse_rules, match_class_names, weight_plus_one_norm_names from modelopt.torch.quantization.model_calib import ( enable_stats_collection, finish_stats_collection, @@ -1144,7 +1144,7 @@ def _update_svdquant(modules, new_pre_quant_scale): finish_stats_collection(module.weight_quantizer) -# AWQ pre_quant_scale fusion rules are per-model data and live in modelopt/modeling/models/*: +# AWQ pre_quant_scale fusion rules are per-model data and live in modelopt/models/*: # - Attention: fold o_proj's pre_quant_scale into v_proj's output dimension. # Before: o_proj_out = [attn @ (v_proj_in @ v_proj.W^T)^T * scale] @ o_proj.W^T # After: o_proj_out = [attn @ (v_proj_in @ (v_proj.W * scale)^T)^T] @ o_proj.W^T diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index e5e2131c0e4..3137493426f 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -56,7 +56,7 @@ class ExportContext: is_modelopt_qlora: bool = False model_type: str | None = None """The model's HF model type (``model.config.model_type``), used to resolve the - model's spec in ``modelopt.modeling``. ``None`` means unknown: spec lookups then + model's spec in ``modelopt.models``. ``None`` means unknown: spec lookups then fail loudly instead of guessing.""" tied_cache: dict[int, nn.Module] = field(default_factory=dict) moe_tied_cache: dict[tuple[int, int], nn.Module] = field(default_factory=dict) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 232293997e2..d7fa6f6d395 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -30,6 +30,8 @@ from safetensors import safe_open from safetensors.torch import save_file +from modelopt.models import hf_model_type + from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales try: @@ -430,7 +432,7 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): # TODO: Handle DBRX MoE quantization_format = get_quantization_format(model) model_type = type(model).__name__.lower() - hf_model_type = getattr(getattr(model, "config", None), "model_type", None) + model_hf_type = hf_model_type(model) module_names = set() # NVFP4 SVDQuant does not need pre-quant scale fusion (either into previous linear or layernorm) because @@ -446,12 +448,12 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): module_names.add(name) # For MoE models update pre_quant_scale to average pre_quant_scale amongst experts - if is_moe(module, hf_model_type) and ( + if is_moe(module, model_hf_type) and ( quantization_format is not QUANTIZATION_NONE and ("awq" in quantization_format or quantization_format == QUANTIZATION_NVFP4_SVDQUANT) ): # update_experts_avg_prequant_scale(module) - grouped_experts = get_experts_list(module, hf_model_type) + grouped_experts = get_experts_list(module, model_hf_type) for modules in grouped_experts: with fsdp2_aware_weight_update(model, modules): preprocess_linear_fusion(modules, resmooth_only=True) @@ -802,7 +804,7 @@ def _process_quantized_modules( model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora, - model_type=getattr(getattr(model, "config", None), "model_type", None), + model_type=hf_model_type(model), ) fsdp_module_to_reshard = None @@ -868,7 +870,7 @@ def _export_transformers_checkpoint( model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora, - model_type=getattr(getattr(model, "config", None), "model_type", None), + model_type=hf_model_type(model), ) for name, sub_module in model.named_modules(): if is_moe(sub_module, prepare_ctx.model_type) and hasattr(sub_module, "experts"): diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index ac721ea8674..ec7b34c7e47 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -29,7 +29,7 @@ import torch.nn.functional as F from tqdm import tqdm -from modelopt.modeling import iter_gate_up_pairs +from modelopt.models import iter_gate_up_pairs from modelopt.torch.opt.config import ModeloptBaseConfig from modelopt.torch.opt.searcher import ForwardLoop from modelopt.torch.quantization.utils.layerwise_calib import ( diff --git a/tests/unit/torch/export/test_modeling_specs.py b/tests/unit/torch/export/test_model_specs.py similarity index 95% rename from tests/unit/torch/export/test_modeling_specs.py rename to tests/unit/torch/export/test_model_specs.py index 9264ae7ac1f..e92c7362935 100644 --- a/tests/unit/torch/export/test_modeling_specs.py +++ b/tests/unit/torch/export/test_model_specs.py @@ -13,20 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the per-model spec registry (modelopt.modeling).""" +"""Unit tests for the per-model spec registry (modelopt.models).""" import pytest import torch.nn as nn -from modelopt.modeling import ( +from modelopt.models import ( ModelSpec, MoEVariant, + hf_model_type, iter_gate_up_pairs, iter_pqs_fuse_rules, match_moe_block, weight_plus_one_norm_names, ) -from modelopt.modeling.registry import _SPECS +from modelopt.models.registry import _SPECS from modelopt.torch.export.layer_utils import ( get_expert_linear_names, get_experts_list, @@ -327,3 +328,13 @@ def test_gemma4_both_root_types_resolve(): "down_proj", "up_proj", ] + + +def test_hf_model_type_accepts_model_or_config(): + from types import SimpleNamespace + + config = SimpleNamespace(model_type="qwen3_moe") + model = SimpleNamespace(config=config) + assert hf_model_type(model) == "qwen3_moe" + assert hf_model_type(config) == "qwen3_moe" + assert hf_model_type(SimpleNamespace()) is None From 42f6c5d88c93d36a3269dc2737ce3f708e9851d3 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:17:54 +0000 Subject: [PATCH 16/16] dev Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../torch/export/MODEL_SPECIFIC_REFACTOR.md | 140 ------------------ 1 file changed, 140 deletions(-) delete mode 100644 modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md diff --git a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md b/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md deleted file mode 100644 index b5dabb4f232..00000000000 --- a/modelopt/torch/export/MODEL_SPECIFIC_REFACTOR.md +++ /dev/null @@ -1,140 +0,0 @@ -# Export: Model-Specific Logic Refactor — Action Plan - -**Goal:** Give the unified HF export path a per-model data registry -(top-level `modelopt/modeling/`, organized by HF model type), so that supporting a -new model means adding one declarative spec file instead of editing if/elif chains -across the export engine. - -**Scope:** The unified HF export path (`unified_export_hf.py` + its helpers) only. - -- The **Megatron** path (`plugins/mcore_*`) already follows a per-model registry - pattern and is untouched. -- The **TRT-LLM** checkpoint path (`model_config_export.py`, the `build_*` - functions in `layer_utils.py`, `tensorrt_llm_utils.py`) is legacy: the plan is - to move it out as-is (separate track), not to refactor it. Its per-model - branches stay where they are. - -## 1. Where the HF path stands after PR #1939 - -PR #1939 gave the HF path registry-based **module dispatch**: `ExportModuleRegistry` -and `PrepareMoEInputsRegistry` (`registry.py`, `hf_export_handlers.py`) select -*which handler processes a module* by class/predicate match, replacing the if/elif -chains that used to live in `unified_export_hf.py`. - -What is still hardcoded is the **per-model data** those handlers (and other HF-path -helpers) consume. Inventory of model-specific logic reachable from the HF path: - -| Item | Location | Kind | -|---|---|---| -| MoE expert linear names (Qwen/DeepSeek/Mixtral/DBRX/GptOss/NemotronH/Gemma4) | `layer_utils.get_expert_linear_names` | data — **migrated (P1)** | -| Duplicate expert-naming table + iterable-experts support gate | `layer_utils.get_experts_list` | data — **migrated (P1)** | -| MoE block class-name list | `layer_utils.is_moe` | data — **migrated (P3)** | -| AWQ `pre_quant_scale` fusion rules (Llama/Qwen3) | `quant_utils.PQS_FUSE_MODULE_MAPPING` | data — **migrated (P2)** | -| weight+1 layernorm class names (Gemma RMSNorm, LayerNorm1P) | `quant_utils._layernorm_uses_weight_plus_one` | data — **migrated (NormSpec)** | -| MoE gate/up fusion pairs (`_GATE_UP_PAIRS`, also privately imported by `quantization/model_calib.py`) | `layer_utils.sync_moe_gate_up_amax` | data — **migrated (MoESpec.gate_up_pair)** | -| BMM-style expert class list (`Llama4TextExperts`/`GptOssExperts`) inline copy for weight transpose | `quant_utils` (dispatch copies in `hf_export_handlers.py` stay) | data — P4 | -| VLM detection gates (`phi4mm` model_type, `nemotronparse` architecture, `"nemotron" in model_type` tower special case) | `model_utils.py`, `unified_export_hf.py` | data (gates) + behavior (extraction) — collect until worth a spec flag | -| Handler match keys (Llama4TextExperts, GptOssExperts, DbrxExperts, QuantMoELinear) | `hf_export_handlers.py` | dispatch (stays: structural, per-module) | -| Fused-expert gated/non-gated split (`gate_up_proj` vs `up_proj`) | `moe_utils.py` | data + structure | -| dummy-forward special cases (Whisper input, Nemotron-VL tower) | `unified_export_hf.requantize_resmooth_fused_llm_layers` | behavior | -| VLM language-tower extraction, DiffusionGemma tied-key reorder | `model_utils.py` | behavior | - -## 2. Design - -Two layers: - -- **Engine** (existing files, organized by operation): owns all algorithms and the - module walk; consults the registry for per-model values. -- **Modeling library** (top-level `modelopt/modeling/`, organized by HF model type): - declarative per-model **data only**. No export logic, stdlib-only imports (not even - torch), so it sits at the bottom of the dependency graph and any modelopt subsystem - can depend on it. - -```text -modelopt/modeling/ - specs.py # ModelSpec — the ONE global per-model descriptor, composed - # from section mixins: topic sections (MoESpec: MoE layouts as - # MoEVariant tuples; NormSpec) + subsystem sections (ExportSpec); - # future sections mix in the same way - registry.py # register() + lookups queried by spec type (None when unmatched) - # + the MRO exact-name matching core (match_class_names) - __init__.py # re-exports; importing it registers all specs - models/ # one small file per HF model type (mirrors transformers.models); - # import == registration; a model registers one instance per spec - # kind it customizes -``` - -Model type names mirror -[`transformers.models`](https://github.com/huggingface/transformers/tree/main/src/transformers/models) -(e.g. `qwen3_moe.py`, `gpt_oss.py`, `nemotron_h.py`); trust-remote-code models -(`arctic`, `deepseek`) use their config `model_type`. - -Each model registers exactly ONE `ModelSpec` (registry enforces uniqueness), so -`get_spec(model_type)` is a dict lookup and consumers call methods directly on the -global spec (e.g. `spec.expert_linear_names_for(module)`). -Resolution is by model type, mirroring HF's own indexing: the engine reads the -model's root HF type (`model.config.model_type`) once per export and passes it -down. The lookup is strict: only the model's own spec is consulted, so a model -whose model_type has no spec fails loudly (register a spec) instead of inheriting -a neighbor's data through a coincidental class-name match. Sub-config model types -of composite models are not walked today — a composite whose MoE lives under a -tower type registers the root type too (`gemma4` + `gemma4_text`, following the -`gemma3`/`gemma3_text` precedent); a recursive config walk is a future extension -if a real VLM tower needs it. Within the spec, each `MoESpec` nests one -`MoEVariant` per concrete block layout — several when the same checkpoint -materializes with different classes and projection names (Mixtral across -transformers generations); variant `block_names` (matched against the module -MRO) identify MoE blocks (`is_moe`) and pick the variant. -`get_expert_linear_names` doesn't need the block class at all when a model's -variants agree on one naming. Without a model type (no config available: unit -tests, the TRT-LLM path), lookups search all specs by class name. - -During migration, call sites kept the legacy branches as a fallback behind the -spec lookup. Once the specs covered every family the legacy chains served, the -chains — and the silent ``w1/w2/w3`` guess for unknown models — were deleted: -expert-name resolution is now *structural detection -> spec -> raise*, so a new -MoE model fails loudly, asking for a spec, instead of inheriting another model's -naming. Generic detection that is not per-model data (the ``*SparseMoeBlock`` / -``*MoeLayer`` conventions and the router+experts structural check in ``is_moe``, -the fused-experts quantizer probe in ``get_expert_linear_names``) stays in the -engine, ahead of or beside the spec lookup. - -Note on naming: `modelopt/modeling/registry.py` (per-model **data**, "what are -this model's values") is distinct from the export-path `registry.py` from PR #1939 -(per-module **dispatch**, "which handler processes this module"). The two layers -compose: handlers look up model data through `modelopt.modeling`. - -## 3. Migration plan - -Each step is one PR with a fallback to the legacy path and an equivalence check -against existing export tests. - -| Step | What | Status | -|---|---|---| -| **P1** | Registry skeleton + MoE expert naming: `get_expert_linear_names` and `get_experts_list` read `spec.expert_linear_names` / `spec.has_iterable_experts`. The #1 "add a MoE model" shotgun-surgery driver. | this PR | -| **P2** | `PQS_FUSE_MODULE_MAPPING` → `spec.pqs_fuse_rules`, aggregated via `iter_pqs_fuse_rules` (llama/qwen3 specs). | this PR | -| **P3** | `is_moe` explicit class-name list → `spec.moe_block_names` (arctic/dbrx_ffn are identification-only specs: no expert naming, so expert-name lookups keep the engine default). The generic `*SparseMoeBlock`/`*MoeLayer` conventions and the structural router+experts check stay in the engine. | this PR | -| **P4** | HF handlers consume specs directly; fold remaining `moe_utils` naming data into specs; share the matcher machinery with the export dispatch registry (#1939). Model_type-scoped resolution (`collect_model_types` + scoped `match_moe_block`, threaded via `ExportContext.model_types`) is already in place from this PR. | planned | -| **P5** | Cross-subsystem pilot: unify the remaining copies of linear-fusion-group knowledge (see §4) into spec fields. Partially done: `_GATE_UP_PAIRS` became `MoESpec.gate_up_pair` and `model_calib.py`'s private import of it is gone — the first quantization consumer of `modelopt.modeling`. Remaining: `shared_input.SHARED_PATTERNS`, `algorithms.quant_grouping_rules`. | in progress | -| **P6** | Migrate remaining quantization-side data: default disabled-quantizer patterns, on-the-fly conversion gates, AutoQuantize grouping rules (see §4). | planned | -| **OUT** | TRT-LLM path branches (`decoder_type` chains in `build_*`, `model_config_export.py`, `tensorrt_llm_utils.py`): frozen, moved out unchanged on a separate track. Candidates for deletion on that track: `adjust_attn_amax_values`, `update_experts_avg_prequant_scale` (unused). NOTE: `MODEL_NAME_TO_TYPE` / `get_model_type` are NOT dead — `examples/hf_ptq/hf_ptq.py` and `multinode_ptq.py` still call them; migrate the examples before removing. | separate track | - -**Guardrails:** one data category per PR (fallback-first while a category is -partially migrated, explicit-error once specs cover it); the engine keeps the -algorithms — model specs supply values only, never fork functions. - -## 4. Beyond export: per-model data in quantization (P5/P6 inventory) - -The same three kinds of model-specific logic exist on the quantization side. Only -kind (a) migrates into `modelopt/modeling`; (b) stays in each subsystem's module -registry (a spec may hold pointer data, never the surgery code); (c) stays in the -engine behind structural checks. - -| Item | Location | Kind | -|---|---|---| -| Linear fusion groups (q/k/v, gate/up, `w1/w3`) — was duplicated 3x; the `_GATE_UP_PAIRS` copy is now `MoESpec.gate_up_pair` | `quantization/utils/shared_input.SHARED_PATTERNS`, `quantization/algorithms.quant_grouping_rules` (remaining) | data — P5 | -| Model-class gates for on-the-fly conversion (`"DbrxForCausalLM"`, `("Step3p5ForCausalLM", ...)`) | `quantization/plugins/huggingface.py` | data — P6 | -| Default disabled-quantizer patterns (`*router*`, `*vision_tower*`; per-model, NVBug-gated) | `modelopt_recipes/.../default_disabled_quantizers.yaml` | data — P6 | -| AutoQuantize grouping regexes (llama q/k/v, Mixtral `w1/w2/w3`, NemotronH mixer) | `quantization/algorithms.py` | data — P6 | -| Quant wrapper classes (`_QuantDbrxExperts` splits `w1/v1/w2` into per-expert linears) | `quantization/plugins/huggingface.py` via `QuantModuleRegistry` | dispatch (stays) | -| Structural MoE detection (`gate`+`experts`+`top_k` attrs; 3-D `gate_up_proj` -> gated) | `quantization/plugins/huggingface.py` | behavior (stays) |