diff --git a/docs/source/reference/1_modelopt_api.rst b/docs/source/reference/1_modelopt_api.rst index caf5697737b..259594bf584 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.models modelopt.onnx modelopt.torch diff --git a/modelopt/models/__init__.py b/modelopt/models/__init__.py new file mode 100644 index 00000000000..98895211451 --- /dev/null +++ b/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, +) diff --git a/modelopt/models/arctic.py b/modelopt/models/arctic.py new file mode 100644 index 00000000000..e914af3302f --- /dev/null +++ b/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"), + ), + ), + ) +) diff --git a/modelopt/models/dbrx.py b/modelopt/models/dbrx.py new file mode 100644 index 00000000000..e02f0c9c01d --- /dev/null +++ b/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"), + ), + ), + ) +) diff --git a/modelopt/models/deepseek.py b/modelopt/models/deepseek.py new file mode 100644 index 00000000000..2ed7e4d84da --- /dev/null +++ b/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"), + ), + ), + ) +) diff --git a/modelopt/models/gemma.py b/modelopt/models/gemma.py new file mode 100644 index 00000000000..b045abe6fee --- /dev/null +++ b/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",))) diff --git a/modelopt/models/gemma4.py b/modelopt/models/gemma4.py new file mode 100644 index 00000000000..4164c452ea6 --- /dev/null +++ b/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)) diff --git a/modelopt/models/gpt_oss.py b/modelopt/models/gpt_oss.py new file mode 100644 index 00000000000..899118e0db5 --- /dev/null +++ b/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"), + ), + ), + ) +) diff --git a/modelopt/models/llama.py b/modelopt/models/llama.py new file mode 100644 index 00000000000..1c13855b126 --- /dev/null +++ b/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"), + ), + ) +) diff --git a/modelopt/models/mixtral.py b/modelopt/models/mixtral.py new file mode 100644 index 00000000000..124ad67d277 --- /dev/null +++ b/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"), + ), + ), + ) +) diff --git a/modelopt/models/nemotron.py b/modelopt/models/nemotron.py new file mode 100644 index 00000000000..5f762fa20ab --- /dev/null +++ b/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"), + ) +) diff --git a/modelopt/models/nemotron_h.py b/modelopt/models/nemotron_h.py new file mode 100644 index 00000000000..4bd21f65eff --- /dev/null +++ b/modelopt/models/nemotron_h.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. + +"""Nemotron-H specs (HF model type ``nemotron_h``).""" + +from .registry import register +from .specs import ModelSpec, MoEVariant + +register( + ModelSpec( + model_type="nemotron_h", + moe_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/models/qwen2_moe.py b/modelopt/models/qwen2_moe.py new file mode 100644 index 00000000000..32c22128a06 --- /dev/null +++ b/modelopt/models/qwen2_moe.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. + +"""Qwen2-MoE specs (HF model type ``qwen2_moe``).""" + +from .registry import register +from .specs import ModelSpec, MoEVariant + +register( + ModelSpec( + model_type="qwen2_moe", + moe_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/models/qwen3.py b/modelopt/models/qwen3.py new file mode 100644 index 00000000000..35bd8c44d87 --- /dev/null +++ b/modelopt/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 .registry import register +from .specs import ModelSpec + +register( + ModelSpec( + 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"), + (("Qwen3MLP",), "up_proj", "down_proj"), + ), + ) +) diff --git a/modelopt/models/qwen3_5_moe.py b/modelopt/models/qwen3_5_moe.py new file mode 100644 index 00000000000..a795089acb8 --- /dev/null +++ b/modelopt/models/qwen3_5_moe.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. + +"""Qwen3.5-MoE specs (HF model type ``qwen3_5_moe``).""" + +from .registry import register +from .specs import ModelSpec, MoEVariant + +register( + ModelSpec( + model_type="qwen3_5_moe", + moe_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/models/qwen3_moe.py b/modelopt/models/qwen3_moe.py new file mode 100644 index 00000000000..fa626970fcb --- /dev/null +++ b/modelopt/models/qwen3_moe.py @@ -0,0 +1,38 @@ +# 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-MoE specs (HF model type ``qwen3_moe``).""" + +from .registry import register +from .specs import ModelSpec, MoEVariant + +register( + ModelSpec( + model_type="qwen3_moe", + moe_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, + ), + ), + # AWQ pre_quant_scale fusion: fold o_proj into v_proj, down_proj into up_proj. + pqs_fuse_rules=( + (("Qwen3MoeAttention",), "v_proj", "o_proj"), + (("Qwen3MoeMLP",), "up_proj", "down_proj"), + ), + ) +) diff --git a/modelopt/models/qwen3_next.py b/modelopt/models/qwen3_next.py new file mode 100644 index 00000000000..216054eb877 --- /dev/null +++ b/modelopt/models/qwen3_next.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. + +"""Qwen3-Next specs (HF model type ``qwen3_next``).""" + +from .registry import register +from .specs import ModelSpec, MoEVariant + +register( + ModelSpec( + model_type="qwen3_next", + moe_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/models/registry.py b/modelopt/models/registry.py new file mode 100644 index 00000000000..8ffbb31afe6 --- /dev/null +++ b/modelopt/models/registry.py @@ -0,0 +1,138 @@ +# 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 indexing the per-model ``ModelSpec`` by HF model type. + +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 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 + +from .specs import ModelSpec, MoEVariant + +if TYPE_CHECKING: + import torch.nn as nn + +__all__ = [ + "get_spec", + "get_specs", + "hf_model_type", + "iter_gate_up_pairs", + "iter_pqs_fuse_rules", + "match_moe_block", + "register", + "weight_plus_one_norm_names", +] + +_SPECS: dict[str, ModelSpec] = {} + + +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_spec(model_type: str) -> ModelSpec | None: + """Return the spec registered for ``model_type``, or ``None``.""" + return _SPECS.get(model_type) + + +def get_specs() -> list[ModelSpec]: + """Return all registered specs, in registration order (aggregators, tests).""" + return list(_SPECS.values()) + + +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). + """ + 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 variants. + + GLOBAL-VOCABULARY semantics: consumers (currently only calibration sibling + grouping in ``quantization/model_calib.py``, which also walks dense MLPs that + 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 + section in a follow-up (see MODEL_SPECIFIC_REFACTOR.md P5). + """ + seen = set() + 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) + yield pair + + +def weight_plus_one_norm_names() -> tuple[str, ...]: + """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 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. + + ``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`` (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. + """ + 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 + return None diff --git a/modelopt/models/specs.py b/modelopt/models/specs.py new file mode 100644 index 00000000000..db99a2168e9 --- /dev/null +++ b/modelopt/models/specs.py @@ -0,0 +1,172 @@ +# 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`` 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 sections** hold architecture facts shared across subsystems (``MoESpec``: + what a model's MoE blocks are; ``NormSpec``: norm-layer conventions); +- **subsystem sections** hold one subsystem's per-model policy (``ExportSpec``; + quantization / speculative-decoding sections to follow). + +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", +] + + +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) + + +@dataclass(kw_only=True) +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; see ``match_class_names``).""" + + 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(kw_only=True) +class MoESpec: + """Topic section: MoE architecture facts — the model's MoE-block layout(s). + + 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. + """ + + moe_variants: tuple[MoEVariant, ...] = () + """The model's MoE-block layouts; more than one when the same checkpoint + materializes differently (see ``MoEVariant``).""" + + def match_moe_variant(self, module) -> MoEVariant | None: + """Return the variant whose ``block_names`` matches ``module``, else None.""" + for variant in self.moe_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.moe_variants + if variant.expert_linear_names is not None + } + if len(namings) == 1: + return next(iter(namings)) + variant = self.match_moe_variant(module) + return variant.expert_linear_names if variant is not None else None + + +@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 + 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(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 + section 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.""" + + +@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/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 800b51daca9..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): + 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) + 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 d5f1fb2330d..98ec9cc6463 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.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 @@ -83,35 +84,27 @@ def get_experts_list( module: torch.nn.Module, - model_type: 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_type: `type(root_model).__name__.lower()` (may change after ModelOpt quantize). + model_type: the model's HF model type (``model.config.model_type``), used to + resolve the model's own spec. """ 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: - raise NotImplementedError(f" {model_type} not supported") + # 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. + 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} (model type: {model_type!r}) not supported" + ) + linear_names = get_expert_linear_names(module, model_type) # Common logic for all supported model types experts_list.extend( @@ -310,14 +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) -> bool: - """Returns whether the module is an MOE layer.""" +def is_moe(module: nn.Module, model_type: str | None = None) -> bool: + """Returns whether the module is an MOE layer. + + ``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 - # 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-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) return ( @@ -973,17 +971,22 @@ def get_stacked_scaling_factors(experts, get_function, module_name): return config -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) +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_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 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. # Also handles models that originally used this naming (Qwen, DeepSeek, etc.). @@ -992,38 +995,17 @@ def module_match_name_list(module, name_list): if hasattr(module.experts, f"{first_proj_attr}_weight_quantizers"): return [first_proj_attr, "down_proj"] - 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"] + 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 type: {model_type!r}). Register a ModelSpec with moe_variants for " + "this model under modelopt/models/." + ) def set_expert_quantizer_amax( @@ -1191,11 +1173,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. @@ -1212,44 +1189,63 @@ def sync_moe_gate_up_amax(model: nn.Module) -> int: Returns: Number of expert gate/up pairs whose amaxes were synced. """ + model_type = hf_model_type(model) 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 (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. + variant = match_moe_block(sub_module, model_type) + if variant is None: + unmatched_block_names.add(type(sub_module).__name__) + continue + if variant.gate_up_pair is None: + continue + gate_name, up_name = variant.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/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 acb1968e070..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 = type(model).__name__.lower() + 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() @@ -466,7 +467,7 @@ 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_type): continue try: expert_groups = get_experts_list(module, model_type) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..cf177a07026 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.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, @@ -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 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 +# - 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" ): @@ -1237,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/export/registry.py b/modelopt/torch/export/registry.py index 8e2cda63df9..3137493426f 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -54,6 +54,10 @@ class ExportContext: model: nn.Module dtype: torch.dtype 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.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 cee64c22c05..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,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_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 @@ -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_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, 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) @@ -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_type=hf_model_type(model), + ) 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_type=hf_model_type(model), + ) 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_type) and hasattr(sub_module, "experts"): handler = PrepareMoEInputsRegistry.match(sub_module.experts) if handler is None: # Unsupported MoE model structure diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index e03dff2e98b..ec7b34c7e47 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.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 ( @@ -87,10 +88,7 @@ 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 - - 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_layer_utils.py b/tests/unit/torch/export/test_layer_utils.py index 93742c8fe7e..044f210bff5 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()) @@ -87,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(), "gemma4") == [ "gate_proj", "down_proj", "up_proj", @@ -108,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(), "mixtral") == ["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(), "nemotron_h") == ["up_proj", "down_proj"] diff --git a/tests/unit/torch/export/test_model_specs.py b/tests/unit/torch/export/test_model_specs.py new file mode 100644 index 00000000000..e92c7362935 --- /dev/null +++ b/tests/unit/torch/export/test_model_specs.py @@ -0,0 +1,340 @@ +# 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 per-model spec registry (modelopt.models).""" + +import pytest +import torch.nn as nn + +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.models.registry import _SPECS +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 + + +class Qwen3MoeSparseMoeBlock(nn.Module): + pass + + +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_class_name(): + 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(): + 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(): + assert match_moe_block(_UnknownMoeBlock()) is None + + +def test_get_expert_linear_names_raises_when_unmatched(): + # 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(), "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"] + # 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") == [ + "w1_linear", + "w2_linear", + "v1_linear", + ] + + +class NemotronHMOE(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 = NemotronHMOE() + 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 + assert groups[1][2] is module.experts[2].down_proj + + +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 DbrxFFN(nn.Module): + def __init__(self): + super().__init__() + self.experts = nn.ModuleList() + + with pytest.raises(NotImplementedError): + get_experts_list(DbrxFFN(), "dbrx") + + with pytest.raises(NotImplementedError): + get_experts_list(_UnknownMoeBlock(), "some_unknown_model") + + +class ArcticMoE(nn.Module): + pass + + +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(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. + rules = { + (substring, fuse_into, fuse_from) + for substrings, fuse_into, fuse_from in iter_pqs_fuse_rules() + for substring in substrings + } + legacy = { + ("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 + + +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()) + + +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 + + +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_variant = MoEVariant( + block_names=("Qwen3MoeSparseMoeBlock",), + expert_linear_names=("a_proj", "b_proj"), + ) + 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 == ( + "gate_proj", + "down_proj", + "up_proj", + ) + # No scope -> first registered class-name match wins (legacy order). + assert match_moe_block(Qwen3MoeSparseMoeBlock()).expert_linear_names == ( + "gate_proj", + "down_proj", + "up_proj", + ) + finally: + del _SPECS[fork_spec.model_type] + + +def test_match_moe_block_scope_is_strict(): + class Qwen3MoeSparseMoeBlock(nn.Module): + pass + + # 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 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") == [ + "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") + + +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", + ] + + +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