[qwen] Add Qwen3.5-MoE builder (port of onnxruntime-genai #2146)

modelbuilder/builder.py

@@ -389,6 +389,10 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
         from .builders.qwen import Qwen35CausalLMModel
 
         onnx_model = Qwen35CausalLMModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options)
+    elif config.architectures[0] == "Qwen3_5MoeForConditionalGeneration":
+        from .builders.qwen import Qwen35MoeTextModel
+
+        onnx_model = Qwen35MoeTextModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options)
     elif config.architectures[0] == "Qwen3VLForConditionalGeneration":
         text_config = config.text_config
         for key in text_config:

modelbuilder/builders/qwen.py

@@ -1895,6 +1895,11 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 
         super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
 
+        # Set the genai-config model_type here (rather than inside
+        # ``make_genai_config``) so subclasses (e.g. ``Qwen35MoeTextModel``)
+        # can override it cleanly.
+        self.model_type = "Qwen3_5ForConditionalGeneration"
+
         # OffsetRMSNorm: Qwen3.5 uses (1 + weight) * RMSNorm(x).
         # Pre-bake the +1 into the weight initializer so the base class's
         # SkipSimplifiedLayerNormalization can be used directly.
@@ -2802,10 +2807,11 @@ def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir):
         hf_config.save_pretrained(out_dir)
 
         # Temporarily restore the KV cache template keys and adjust attributes
-        # so the base class generates the right entries.
-        saved = {"num_layers": self.num_layers, "model_type": self.model_type}
+        # so the base class generates the right entries.  ``self.model_type``
+        # is already set in ``__init__`` (and may have been overridden by a
+        # subclass such as ``Qwen35MoeTextModel``).
+        saved = {"num_layers": self.num_layers}
         self.num_layers = len(self.layer_types)
-        self.model_type = "Qwen3_5ForConditionalGeneration"
         self.input_names["past_key_values.key"] = "past_key_values.%d.key"
         self.input_names["past_key_values.value"] = "past_key_values.%d.value"
         self.output_names["present.key"] = "present.%d.key"
@@ -2815,7 +2821,6 @@ def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir):
 
         # Restore
         self.num_layers = saved["num_layers"]
-        self.model_type = saved["model_type"]
         del self.input_names["past_key_values.key"]
         del self.input_names["past_key_values.value"]
         del self.output_names["present.key"]
@@ -2849,3 +2854,203 @@ def load_weights(self, input_path):
         return Qwen3_5ForCausalLM.from_pretrained(
             self.model_name_or_path, cache_dir=self.cache_dir, token=self.hf_token, trust_remote_code=self.hf_remote
         )
+
+
+class Qwen35MoeTextModel(Qwen35TextModel):
+    """Qwen3.5 MoE hybrid model builder.
+
+    Extends :class:`Qwen35TextModel` with Mixture-of-Experts MLP layers.
+    Each decoder layer replaces the dense MLP with:
+
+    - A router that selects top-k experts from ``num_experts`` candidates.
+    - Packed routed expert weights (``gate_up_proj`` and ``down_proj``).
+    - A shared expert (always-active) with its own sigmoid gating signal.
+
+    The attention side (GatedDeltaNet linear + gated full attention) is
+    inherited unchanged from the parent class.
+
+    Imported from `microsoft/onnxruntime-genai PR #2146
+    <https://github.com/microsoft/onnxruntime-genai/pull/2146>`_.
+    """
+
+    def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
+        # Map Qwen3.5-MoE config attributes to what the base class expects.
+        if hasattr(config, "text_config"):
+            tc = config.text_config
+            # Base class reads ``num_local_experts``; MoE config uses ``num_experts``.
+            if hasattr(tc, "num_experts") and not hasattr(tc, "num_local_experts"):
+                tc.num_local_experts = tc.num_experts
+            # Base class reads ``intermediate_size``; MoE has ``moe_intermediate_size``.
+            if not hasattr(tc, "intermediate_size") and hasattr(tc, "moe_intermediate_size"):
+                tc.intermediate_size = tc.moe_intermediate_size
+
+        super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
+
+        self.model_type = "Qwen3_5_MoeForConditionalGeneration"
+
+        # MoE attributes specific to Qwen3.5-MoE.
+        self.moe_attrs["activation_type"] = "swiglu"
+        self.moe_attrs["swiglu_fusion"] = 1
+        self.moe_attrs["normalize_routing_weights"] = True
+
+        self.moe_intermediate_size = getattr(config, "moe_intermediate_size", 512)
+        self.shared_expert_intermediate_size = getattr(config, "shared_expert_intermediate_size", self.moe_intermediate_size)
+
+        # MoE layers use MoE/QMoE ops instead of individual MatMul nodes,
+        # so remove any /mlp/ MatMul overrides that don't apply.
+        algo_config = self.quant_attrs["int4"].get("algo_config")
+        if algo_config is not None and hasattr(algo_config, "customized_weight_config"):
+            keys_to_remove = [k for k in algo_config.customized_weight_config if "/mlp/" in k]
+            for k in keys_to_remove:
+                del algo_config.customized_weight_config[k]
+
+    def load_weights(self, input_path):
+        from transformers import Qwen3_5MoeForConditionalGeneration
+
+        print("Loading Qwen3_5MoeForConditionalGeneration model...")
+        return Qwen3_5MoeForConditionalGeneration.from_pretrained(
+            self.model_name_or_path, cache_dir=self.cache_dir, token=self.hf_token, trust_remote_code=self.hf_remote
+        )
+
+    def make_layer(self, layer_id, layer):
+        """Override to use MoE instead of dense MLP."""
+        attn_module = layer.linear_attn if self.layer_types[layer_id] == "linear_attention" else layer.self_attn
+        self.make_layernorm(
+            layer_id,
+            layer.input_layernorm,
+            skip=not self.layernorm_attrs["first_layernorm"],
+            simple=self.layernorm_attrs["simple"],
+            location="input",
+        )
+        self.make_attention(layer_id, attn_module, root_input=self.layernorm_attrs["output_0"])
+        self.make_layernorm(
+            layer_id, layer.post_attention_layernorm, skip=True, simple=self.layernorm_attrs["simple"], location="post_attention"
+        )
+        self.make_moe(layer_id, layer.mlp, root_input=self.layernorm_attrs["output_0"])
+
+        self.layernorm_attrs["first_layernorm"] = False
+        if layer_id == self.num_layers - 1:
+            self.layernorm_attrs["last_layernorm"] = True
+
+    def make_moe(self, layer_id, mlp, root_input):
+        """Build MoE + shared expert subgraph for one decoder layer."""
+        basename = f"/model/layers.{layer_id}/moe"
+        op_type = self.moe_attrs["op_type"]
+        moe_weight_type = f"{'q' if op_type == 'QMoE' else ''}weight"
+
+        # --- Router (bias-free gate) ---
+        router_basename = f"{basename}/router/MatMul"
+        router_matmul_name = self.make_matmul(mlp.gate, router_basename, root_input)
+        router_reshape_name = f"{basename}/router/Reshape"
+        self.make_reshape(
+            router_reshape_name,
+            [f"{router_matmul_name}/output_0", f"/model/constants/INT64/{[-1, self.moe_attrs['num_experts']]}"],
+            dtype=self.io_dtype,
+            shape=["batch_size * sequence_length", self.moe_attrs["num_experts"]],
+        )
+
+        # --- Routed expert weights ---
+        gate_up_proj_weight = f"model.layers.{layer_id}.moe.experts.gate_up_proj.{moe_weight_type}"
+        gate_up_proj_scales = f"model.layers.{layer_id}.moe.experts.gate_up_proj.scales"
+        gate_up_proj_bias = f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias"
+        down_proj_weight = f"model.layers.{layer_id}.moe.experts.down_proj.{moe_weight_type}"
+        down_proj_scales = f"model.layers.{layer_id}.moe.experts.down_proj.scales"
+        down_proj_bias = f"model.layers.{layer_id}.moe.experts.down_proj.bias"
+
+        # Repack HF concatenated [gate|up] to ORT interleaved [g0,u0,g1,u1,...] for swiglu_fusion=1.
+        raw_gate_up = mlp.experts.gate_up_proj
+        half = raw_gate_up.shape[1] // 2
+        interleaved = torch.stack([raw_gate_up[:, :half, :], raw_gate_up[:, half:, :]], dim=2).reshape_as(raw_gate_up)
+
+        if op_type == "MoE":
+            self.make_initializer(interleaved, gate_up_proj_weight, to=self.io_dtype)
+            self.make_initializer(mlp.experts.down_proj, down_proj_weight, to=self.io_dtype)
+        else:
+            gate_up_qw_list, gate_up_sc_list = [], []
+            down_qw_list, down_sc_list = [], []
+            for i in range(self.moe_attrs["num_experts"]):
+                qw1, sc1 = self.make_qmoe_weights(interleaved[i])
+                gate_up_qw_list.append(qw1)
+                gate_up_sc_list.append(sc1)
+                qw2, sc2 = self.make_qmoe_weights(mlp.experts.down_proj[i])
+                down_qw_list.append(qw2)
+                down_sc_list.append(sc2)
+            self.make_initializer(torch.stack(gate_up_qw_list, dim=0).to(torch.uint8), gate_up_proj_weight)
+            self.make_initializer(torch.stack(down_qw_list, dim=0).to(torch.uint8), down_proj_weight)
+            self.make_initializer(torch.stack(gate_up_sc_list, dim=0), gate_up_proj_scales, to=self.io_dtype)
+            self.make_initializer(torch.stack(down_sc_list, dim=0), down_proj_scales, to=self.io_dtype)
+
+        num_e = self.moe_attrs["num_experts"]
+        self.make_initializer(torch.zeros(num_e, 2 * self.moe_intermediate_size), gate_up_proj_bias, to=self.io_dtype)
+        self.make_initializer(torch.zeros(num_e, self.hidden_size), down_proj_bias, to=self.io_dtype)
+
+        # --- MoE/QMoE op ---
+        moe_name = f"{basename}/{op_type}"
+        self.make_moe_op(
+            moe_name,
+            root_input=root_input,
+            router_probs=f"{router_reshape_name}/output_0",
+            weight1=gate_up_proj_weight,
+            scales1=gate_up_proj_scales if op_type == "QMoE" else "",
+            bias1=gate_up_proj_bias,
+            weight2=down_proj_weight,
+            scales2=down_proj_scales if op_type == "QMoE" else "",
+            bias2=down_proj_bias,
+        )
+
+        # --- Shared expert ---
+        shared_output = self.make_shared_expert(layer_id, mlp.shared_expert, mlp.shared_expert_gate, root_input)
+        combine_name = f"{basename}/Add"
+        self.make_add(
+            combine_name,
+            [f"{moe_name}/output_0", shared_output],
+            dtype=self.io_dtype,
+            shape=["batch_size", "sequence_length", self.hidden_size],
+        )
+        self.layernorm_attrs["skip_input"] = f"{combine_name}/output_0"
+
+    def make_shared_expert(self, layer_id, shared_expert, shared_expert_gate, root_input):
+        """Build the shared expert SiLU-MLP with sigmoid gating."""
+        basename = f"/model/layers.{layer_id}/shared_expert"
+
+        gate_matmul = self.make_matmul(shared_expert.gate_proj, f"{basename}/gate_proj/MatMul", root_input)
+        up_matmul = self.make_matmul(shared_expert.up_proj, f"{basename}/up_proj/MatMul", root_input)
+
+        silu_sigmoid_name = f"{basename}/gate_proj/Sigmoid"
+        self.make_sigmoid(
+            silu_sigmoid_name,
+            f"{gate_matmul}/output_0",
+            self.io_dtype,
+            shape=["batch_size", "sequence_length", self.shared_expert_intermediate_size],
+        )
+
+        silu_mul_name = f"{basename}/gate_proj/Mul"
+        self.make_mul(
+            silu_mul_name,
+            [f"{gate_matmul}/output_0", f"{silu_sigmoid_name}/output_0"],
+            dtype=self.io_dtype,
+            shape=["batch_size", "sequence_length", self.shared_expert_intermediate_size],
+        )
+
+        gate_up_mul_name = f"{basename}/Mul"
+        self.make_mul(
+            gate_up_mul_name,
+            [f"{silu_mul_name}/output_0", f"{up_matmul}/output_0"],
+            dtype=self.io_dtype,
+            shape=["batch_size", "sequence_length", self.shared_expert_intermediate_size],
+        )
+
+        down_matmul = self.make_matmul(shared_expert.down_proj, f"{basename}/down_proj/MatMul", f"{gate_up_mul_name}/output_0")
+
+        gate_matmul_name = self.make_matmul(shared_expert_gate, f"{basename}_gate/MatMul", root_input)
+        gate_sigmoid_name = f"{basename}_gate/Sigmoid"
+        self.make_sigmoid(gate_sigmoid_name, f"{gate_matmul_name}/output_0", self.io_dtype, shape=["batch_size", "sequence_length", 1])
+
+        gated_mul_name = f"{basename}/GatedMul"
+        self.make_mul(
+            gated_mul_name,
+            [f"{down_matmul}/output_0", f"{gate_sigmoid_name}/output_0"],
+            dtype=self.io_dtype,
+            shape=["batch_size", "sequence_length", self.hidden_size],
+        )
+        return f"{gated_mul_name}/output_0"

tests/fast/test_random_qwen3_5_moe.py

@@ -0,0 +1,141 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.  See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+"""Fast unit tests for the Qwen3.5-MoE ONNX builder.
+
+These tests cover the changes imported from `microsoft/onnxruntime-genai
+PR #2146 <https://github.com/microsoft/onnxruntime-genai/pull/2146>`_,
+which add support for the ``Qwen3_5MoeForConditionalGeneration`` HF
+architecture (256 routed experts + 1 shared expert with SwiGLU).
+"""
+
+import os
+import unittest
+
+from modelbuilder.ext_test_case import ExtTestCase, hide_stdout, requires_transformers
+
+QWEN3_5_MOE_MODEL_NAME = "Qwen/Qwen3.5-MoE"
+
+
+def _make_qwen3_5_moe_config(layer_types, num_hidden_layers=None):
+    """Return a minimal ``Qwen3_5MoeConfig`` suitable for offline unit tests.
+
+    Parameters
+    ----------
+    layer_types:
+        List of layer type strings, e.g. ``["full_attention", "linear_attention"]``.
+        The number of layers is inferred from this list unless
+        ``num_hidden_layers`` is provided.
+    num_hidden_layers:
+        Explicit override; defaults to ``len(layer_types)``.
+    """
+    from transformers.models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig
+
+    if num_hidden_layers is None:
+        num_hidden_layers = len(layer_types)
+
+    # partial_rotary_factor=0.25, head_dim=64 -> rdim=16, rdim_half=8.
+    # mrope_section=[2, 3, 3]: height positions at stride-1 offsets within rdim_half.
+    rope_cfg = {"type": "mrope", "rope_type": "default", "mrope_section": [2, 3, 3], "rope_theta": 10000.0, "partial_rotary_factor": 0.25}
+
+    text_config = Qwen3_5MoeTextConfig(
+        hidden_size=128,
+        num_hidden_layers=num_hidden_layers,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=64,
+        max_position_embeddings=256,
+        vocab_size=32000,
+        rms_norm_eps=1e-6,
+        layer_types=layer_types,
+        linear_num_key_heads=2,
+        linear_num_value_heads=2,
+        linear_key_head_dim=16,
+        linear_value_head_dim=16,
+        linear_conv_kernel_dim=4,
+        # MoE-specific settings (kept small for CI).
+        num_experts=4,
+        num_experts_per_tok=2,
+        moe_intermediate_size=64,
+        shared_expert_intermediate_size=64,
+    )
+    text_config.rope_scaling = rope_cfg
+    text_config.rope_parameters = rope_cfg
+
+    config = Qwen3_5MoeConfig(text_config=text_config, bos_token_id=1, eos_token_id=2)
+    config.architectures = ["Qwen3_5MoeForConditionalGeneration"]
+    return config
+
+
+class TestRandomQwen3_5Moe(ExtTestCase):
+    def _build_model(self, config, precision, provider):
+        """Create a random-weight HF MoE model and build its ONNX export.
+
+        Returns the output directory containing ``model.onnx``.
+        """
+        import torch
+        from transformers import AutoModelForImageTextToText
+
+        from modelbuilder.builder import create_model
+
+        basename = f"test_qwen3_5_moe_{precision}_{provider}_" + "_".join(config.text_config.layer_types)
+        model_dir_full = self.get_model_dir(basename)
+        output_dir, cache_dir = self.get_dirs(basename)
+
+        torch.manual_seed(42)
+        model = AutoModelForImageTextToText.from_config(config)
+        model.eval()
+        model.save_pretrained(model_dir_full)
+
+        tokenizer = self.make_word_level_tokenizer()
+        tokenizer.save_pretrained(model_dir_full)
+
+        create_model(
+            model_name=QWEN3_5_MOE_MODEL_NAME,
+            input_path=model_dir_full,
+            output_dir=output_dir,
+            precision=precision,
+            execution_provider=provider,
+            cache_dir=cache_dir,
+        )
+        return output_dir
+
+    @requires_transformers("5")
+    @hide_stdout()
+    def test_qwen3_5_moe_fp32_cpu_full_attention_build(self):
+        """Build a Qwen3.5-MoE decoder with only ``full_attention`` layers.
+
+        Verifies that ``Qwen35MoeTextModel`` registers correctly, emits MoE
+        ops in the ONNX graph, and produces a model with the expected
+        ``model_type`` in ``genai_config.json``.
+        """
+        import json
+
+        import onnx
+
+        config = _make_qwen3_5_moe_config(["full_attention", "full_attention"])
+        output_dir = self._build_model(config, "fp32", "cpu")
+
+        text_onnx_path = os.path.join(output_dir, "model.onnx")
+        self.assertExists(text_onnx_path)
+
+        onnx_model = onnx.load(text_onnx_path)
+        self.assertIsNotNone(onnx_model)
+
+        # Confirm an MoE op (MoE or QMoE) appears in the graph along with the
+        # shared-expert sigmoid gate (com.microsoft Sigmoid via make_sigmoid).
+        op_types = {node.op_type for node in onnx_model.graph.node}
+        self.assertTrue(("MoE" in op_types) or ("QMoE" in op_types), f"Expected an MoE/QMoE op in the graph, found: {sorted(op_types)}")
+
+        # genai_config.json should expose the MoE-specific model_type.
+        genai_config_path = os.path.join(output_dir, "genai_config.json")
+        self.assertExists(genai_config_path)
+        with open(genai_config_path, encoding="utf-8") as f:
+            genai_config = json.load(f)
+        self.assertEqual(genai_config["model"]["type"], "qwen3_5_moe")
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)