[Refactor] Model-family spec registry for export: MoE expert naming

docs/source/reference/1_modelopt_api.rst

@@ -10,5 +10,6 @@ modelopt API
    :recursive:
 
    modelopt.deploy
+   modelopt.models
    modelopt.onnx
    modelopt.torch

modelopt/models/__init__.py

@@ -0,0 +1,47 @@
+# 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, one module per HF model type.
+
+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 .registry import *
+from .specs import *
+
+# Importing the model modules registers every spec as a side effect.
+from . import (  # isort: skip
+    arctic,
+    dbrx,
+    deepseek,
+    gemma,
+    gemma4,
+    gpt_oss,
+    llama,
+    mixtral,
+    nemotron,
+    nemotron_h,
+    qwen2_moe,
+    qwen3,
+    qwen3_5_moe,
+    qwen3_moe,
+    qwen3_next,
+)

modelopt/models/arctic.py

@@ -0,0 +1,35 @@
+# 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.
+
+"""Snowflake Arctic specs (trust-remote-code model type ``arctic``)."""
+
+from .registry import register
+from .specs import ModelSpec, MoEVariant
+
+register(
+    ModelSpec(
+        model_type="arctic",
+        moe_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"),
+            ),
+        ),
+    )
+)

modelopt/models/dbrx.py

@@ -0,0 +1,36 @@
+# 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 specs (HF model type ``dbrx``)."""
+
+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
+# 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(
+    ModelSpec(
+        model_type="dbrx",
+        moe_variants=(
+            MoEVariant(
+                block_names=("DbrxFFN",),
+                expert_linear_names=("w1_linear", "w2_linear", "v1_linear"),
+            ),
+        ),
+    )
+)

modelopt/models/deepseek.py

@@ -0,0 +1,40 @@
+# 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-MoE specs (trust-remote-code model type ``deepseek``).
+
+Matches the remote-code ``DeepseekMoE`` block, not the HF-native ``deepseek_v3``
+classes.
+"""
+
+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
+# resmoothing) is validated on this model — the flag currently doubles as that
+# support gate.
+register(
+    ModelSpec(
+        model_type="deepseek",
+        moe_variants=(
+            MoEVariant(
+                block_names=("DeepseekMoE",),
+                expert_linear_names=("gate_proj", "down_proj", "up_proj"),
+                gate_up_pair=("gate_proj", "up_proj"),
+            ),
+        ),
+    )
+)

modelopt/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 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(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",)))

modelopt/models/gemma4.py

@@ -0,0 +1,35 @@
+# 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.
+
+"""Gemma4 specs (HF model type ``gemma4``)."""
+
+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
+# 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))

modelopt/models/gpt_oss.py

@@ -0,0 +1,32 @@
+# 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 specs (HF model type ``gpt_oss``)."""
+
+from .registry import register
+from .specs import ModelSpec, MoEVariant
+
+register(
+    ModelSpec(
+        model_type="gpt_oss",
+        moe_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"),
+            ),
+        ),
+    )
+)

modelopt/models/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 specs (HF model type ``llama``)."""
+
+from .registry import register
+from .specs import ModelSpec
+
+register(
+    ModelSpec(
+        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"),
+            (("LlamaMLP",), "up_proj", "down_proj"),
+        ),
+    )
+)

modelopt/models/mixtral.py

@@ -0,0 +1,42 @@
+# 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 specs (HF model type ``mixtral``)."""
+
+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.
+register(
+    ModelSpec(
+        model_type="mixtral",
+        moe_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"),
+            ),
+        ),
+    )
+)

modelopt/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 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(
+    ModelSpec(
+        model_type="nemotron",
+        weight_plus_one_norm_names=("LayerNorm1P", "NemotronLayerNorm1P"),
+    )
+)