Redo of #1477: Qwen2-MoE bridge support (retarget to dev)

tests/integration/model_bridge/test_qwen2_moe_bridge.py

@@ -0,0 +1,119 @@
+"""Integration tests for the Qwen2-MoE TransformerBridge."""
+
+import copy
+
+import torch
+from transformers import Qwen2MoeConfig, Qwen2MoeForCausalLM
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.generalized_components import (
+    GatedMLPBridge,
+    MoEBridge,
+    PositionEmbeddingsAttentionBridge,
+)
+from transformer_lens.model_bridge.sources._bridge_builder import (
+    build_bridge_from_module,
+)
+
+
+def tiny_qwen2_moe_config() -> Qwen2MoeConfig:
+    return Qwen2MoeConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=96,
+        moe_intermediate_size=32,
+        shared_expert_intermediate_size=96,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        num_experts=4,
+        num_experts_per_tok=2,
+        max_position_embeddings=64,
+        decoder_sparse_step=1,
+        mlp_only_layers=[],
+    )
+
+
+def tiny_qwen2_moe_bridge() -> tuple[TransformerBridge, Qwen2MoeForCausalLM]:
+    torch.manual_seed(0)
+    cfg = tiny_qwen2_moe_config()
+    cfg._attn_implementation = "eager"
+    hf_reference = Qwen2MoeForCausalLM(cfg).eval()
+    hf_reference.config._attn_implementation = "eager"
+    for layer in hf_reference.model.layers:
+        layer.self_attn.config._attn_implementation = "eager"
+
+    hf_model = Qwen2MoeForCausalLM(copy.deepcopy(cfg)).eval()
+    hf_model.load_state_dict(hf_reference.state_dict())
+    bridge = build_bridge_from_module(
+        hf_model,
+        "Qwen2MoeForCausalLM",
+        hf_config=copy.deepcopy(cfg),
+        tokenizer=None,
+        device="cpu",
+    ).eval()
+    bridge.adapter.setup_component_testing(hf_model, bridge_model=bridge)
+    return bridge, hf_reference
+
+
+def _tokens() -> torch.Tensor:
+    return torch.tensor([[1, 2, 3, 4, 5]])
+
+
+class TestQwen2MoeBridge:
+    def test_bridge_structure(self) -> None:
+        bridge, _ = tiny_qwen2_moe_bridge()
+
+        assert len(bridge.blocks) == 2
+        assert isinstance(bridge.blocks[0].attn, PositionEmbeddingsAttentionBridge)
+        assert isinstance(bridge.blocks[0].mlp, MoEBridge)
+        assert isinstance(bridge.blocks[0].mlp.shared_expert, GatedMLPBridge)
+        assert hasattr(bridge.blocks[0].mlp, "shared_expert_gate")
+
+    def test_forward_matches_hf(self) -> None:
+        bridge, hf_reference = tiny_qwen2_moe_bridge()
+        tokens = _tokens()
+
+        with torch.no_grad():
+            bridge_out = bridge(tokens)
+            hf_out = hf_reference(tokens).logits
+
+        assert bridge_out.shape == (1, 5, 128)
+        assert not torch.isnan(bridge_out).any()
+        assert not torch.isinf(bridge_out).any()
+        max_diff = (bridge_out - hf_out).abs().max().item()
+        assert max_diff < 1e-5, f"Bridge vs HF max diff = {max_diff}"
+
+    def test_run_with_cache_captures_moe_hooks(self) -> None:
+        bridge, _ = tiny_qwen2_moe_bridge()
+        tokens = _tokens()
+        batch, seq_len = tokens.shape
+        flat_tokens = batch * seq_len
+        d_model = bridge.cfg.d_model
+        num_experts = bridge.cfg.num_experts
+
+        _, cache = bridge.run_with_cache(tokens)
+
+        for layer_idx in range(len(bridge.blocks)):
+            gate_key = f"blocks.{layer_idx}.mlp.gate.hook_out"
+            assert gate_key in cache, f"Missing cache key: {gate_key}"
+            assert cache[gate_key].shape == (flat_tokens, num_experts)
+
+            experts_key = f"blocks.{layer_idx}.mlp.experts.hook_out"
+            assert experts_key in cache, f"Missing cache key: {experts_key}"
+            assert cache[experts_key].shape == (flat_tokens, d_model)
+
+            shared_expert_key = f"blocks.{layer_idx}.mlp.shared_expert.hook_out"
+            assert shared_expert_key in cache, f"Missing cache key: {shared_expert_key}"
+            assert cache[shared_expert_key].shape == (flat_tokens, d_model)
+
+            shared_expert_gate_key = f"blocks.{layer_idx}.mlp.shared_expert_gate.hook_out"
+            assert shared_expert_gate_key in cache, f"Missing cache key: {shared_expert_gate_key}"
+            assert cache[shared_expert_gate_key].shape == (flat_tokens, 1)
+
+            mlp_out_key = f"blocks.{layer_idx}.mlp.hook_out"
+            assert mlp_out_key in cache, f"Missing cache key: {mlp_out_key}"
+            assert cache[mlp_out_key].shape == (batch, seq_len, d_model)
+
+            router_scores_key = f"blocks.{layer_idx}.mlp.hook_router_scores"
+            assert router_scores_key not in cache

tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py

@@ -0,0 +1,151 @@
+"""Unit tests for the Qwen2MoeArchitectureAdapter."""
+
+import pytest
+import torch
+from transformers import Qwen2MoeConfig
+
+from transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion import (
+    RearrangeTensorConversion,
+)
+from transformer_lens.conversion_utils.param_processing_conversion import (
+    ParamProcessingConversion,
+)
+from transformer_lens.factories.architecture_adapter_factory import (
+    ArchitectureAdapterFactory,
+)
+from transformer_lens.model_bridge.generalized_components import (
+    GatedMLPBridge,
+    LinearBridge,
+    MoEBridge,
+    PositionEmbeddingsAttentionBridge,
+)
+from transformer_lens.model_bridge.sources._bridge_builder import (
+    build_bridge_config_from_hf,
+)
+from transformer_lens.model_bridge.sources.transformers import (
+    determine_architecture_from_hf_config,
+)
+from transformer_lens.model_bridge.supported_architectures.qwen2_moe import (
+    Qwen2MoeArchitectureAdapter,
+)
+
+
+def _tiny_config() -> Qwen2MoeConfig:
+    return Qwen2MoeConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=96,
+        moe_intermediate_size=32,
+        shared_expert_intermediate_size=96,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        num_experts=4,
+        num_experts_per_tok=2,
+        max_position_embeddings=64,
+        decoder_sparse_step=1,
+        mlp_only_layers=[],
+    )
+
+
+@pytest.fixture
+def bridge_cfg():
+    return build_bridge_config_from_hf(
+        _tiny_config(),
+        "Qwen2MoeForCausalLM",
+        "qwen2-moe-test",
+        torch.float32,
+    )
+
+
+@pytest.fixture
+def adapter(bridge_cfg) -> Qwen2MoeArchitectureAdapter:
+    return Qwen2MoeArchitectureAdapter(bridge_cfg)
+
+
+class TestQwen2MoeDetection:
+    def test_model_type_detects_qwen2_moe(self) -> None:
+        assert determine_architecture_from_hf_config(_tiny_config()) == "Qwen2MoeForCausalLM"
+
+    def test_factory_selects_adapter(self, bridge_cfg) -> None:
+        selected = ArchitectureAdapterFactory.select_architecture_adapter(bridge_cfg)
+        assert isinstance(selected, Qwen2MoeArchitectureAdapter)
+
+
+class TestQwen2MoeConfigMapping:
+    def test_core_fields_map_from_hf_config(self, bridge_cfg) -> None:
+        hf_cfg = _tiny_config()
+        assert bridge_cfg.d_vocab == hf_cfg.vocab_size
+        assert bridge_cfg.d_model == hf_cfg.hidden_size
+        assert bridge_cfg.n_layers == hf_cfg.num_hidden_layers
+        assert bridge_cfg.n_heads == hf_cfg.num_attention_heads
+        assert bridge_cfg.n_key_value_heads == hf_cfg.num_key_value_heads
+        assert bridge_cfg.d_mlp == hf_cfg.intermediate_size
+
+    def test_moe_fields_map_from_hf_config(self, bridge_cfg) -> None:
+        hf_cfg = _tiny_config()
+        assert bridge_cfg.num_experts == hf_cfg.num_experts
+        assert bridge_cfg.experts_per_token == hf_cfg.num_experts_per_tok
+        assert getattr(bridge_cfg, "moe_intermediate_size") == hf_cfg.moe_intermediate_size
+        assert (
+            getattr(bridge_cfg, "shared_expert_intermediate_size")
+            == hf_cfg.shared_expert_intermediate_size
+        )
+
+
+class TestQwen2MoeComponentMapping:
+    def test_reuses_qwen2_attention_mapping(self, adapter: Qwen2MoeArchitectureAdapter) -> None:
+        blocks = adapter.component_mapping["blocks"]
+        attn = blocks.submodules["attn"]
+        assert isinstance(attn, PositionEmbeddingsAttentionBridge)
+        assert set(attn.submodules.keys()) == {"q", "k", "v", "o"}
+        assert attn.submodules["q"].name == "q_proj"
+        assert attn.submodules["k"].name == "k_proj"
+        assert attn.submodules["v"].name == "v_proj"
+        assert attn.submodules["o"].name == "o_proj"
+
+    def test_mlp_is_moe_bridge(self, adapter: Qwen2MoeArchitectureAdapter) -> None:
+        mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+        assert isinstance(mlp, MoEBridge)
+        assert mlp.name == "mlp"
+
+    def test_moe_submodules(self, adapter: Qwen2MoeArchitectureAdapter) -> None:
+        mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+        assert set(mlp.submodules.keys()) == {
+            "gate",
+            "experts",
+            "shared_expert",
+            "shared_expert_gate",
+        }
+        assert isinstance(mlp.submodules["gate"], LinearBridge)
+        assert isinstance(mlp.submodules["experts"], MoEBridge)
+        assert isinstance(mlp.submodules["shared_expert"], GatedMLPBridge)
+        assert isinstance(mlp.submodules["shared_expert_gate"], LinearBridge)
+
+    def test_moe_hf_paths(self, adapter: Qwen2MoeArchitectureAdapter) -> None:
+        mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+        assert mlp.submodules["gate"].name == "gate"
+        assert mlp.submodules["experts"].name == "experts"
+        assert mlp.submodules["shared_expert"].name == "shared_expert"
+        assert mlp.submodules["shared_expert_gate"].name == "shared_expert_gate"
+
+    def test_shared_expert_submodules(self, adapter: Qwen2MoeArchitectureAdapter) -> None:
+        mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+        shared_expert = mlp.submodules["shared_expert"]
+        assert set(shared_expert.submodules.keys()) == {"gate", "in", "out"}
+        assert shared_expert.submodules["gate"].name == "gate_proj"
+        assert shared_expert.submodules["in"].name == "up_proj"
+        assert shared_expert.submodules["out"].name == "down_proj"
+
+
+class TestQwen2MoeWeightConversions:
+    def test_kv_rearrange_uses_n_key_value_heads(
+        self, adapter: Qwen2MoeArchitectureAdapter
+    ) -> None:
+        conversions = adapter.weight_processing_conversions
+        assert conversions is not None
+        for slot in ("k", "v"):
+            conv = conversions[f"blocks.{{i}}.attn.{slot}.weight"]
+            assert isinstance(conv, ParamProcessingConversion)
+            assert isinstance(conv.tensor_conversion, RearrangeTensorConversion)
+            assert conv.tensor_conversion.axes_lengths["n"] == 2

transformer_lens/factories/architecture_adapter_factory.py

@@ -67,6 +67,7 @@
     PhiArchitectureAdapter,
     PhiMoEArchitectureAdapter,
     Qwen2ArchitectureAdapter,
+    Qwen2MoeArchitectureAdapter,
     Qwen3_5ArchitectureAdapter,
     Qwen3_5MultimodalArchitectureAdapter,
     Qwen3ArchitectureAdapter,
@@ -144,6 +145,7 @@
     "PhiMoEForCausalLM": PhiMoEArchitectureAdapter,
     "QwenForCausalLM": QwenArchitectureAdapter,
     "Qwen2ForCausalLM": Qwen2ArchitectureAdapter,
+    "Qwen2MoeForCausalLM": Qwen2MoeArchitectureAdapter,
     "Qwen3ForCausalLM": Qwen3ArchitectureAdapter,
     "Qwen3MoeForCausalLM": Qwen3MoeArchitectureAdapter,
     "Qwen3NextForCausalLM": Qwen3NextArchitectureAdapter,

transformer_lens/model_bridge/sources/_bridge_builder.py

@@ -52,6 +52,7 @@
     # Hybrid/MoE architectures
     "layer_types",
     "moe_intermediate_size",
+    "shared_expert_intermediate_size",
     "norm_eps",
     "attention_bias",
     "lm_head_bias",

transformer_lens/model_bridge/sources/transformers.py

@@ -180,7 +180,9 @@ def map_default_transformer_lens_config(hf_config):
         tl_config.eps = source_config.layer_norm_epsilon
     elif hasattr(source_config, "norm_eps"):
         tl_config.eps = source_config.norm_eps
-    if hasattr(source_config, "num_local_experts"):
+    if hasattr(source_config, "num_experts"):
+        tl_config.num_experts = source_config.num_experts
+    elif hasattr(source_config, "num_local_experts"):
         tl_config.num_experts = source_config.num_local_experts
     if hasattr(source_config, "num_experts_per_tok"):
         tl_config.experts_per_token = source_config.num_experts_per_tok
@@ -248,6 +250,7 @@ def determine_architecture_from_hf_config(hf_config):
             "phi3": "Phi3ForCausalLM",
             "qwen": "QwenForCausalLM",
             "qwen2": "Qwen2ForCausalLM",
+            "qwen2_moe": "Qwen2MoeForCausalLM",
             "qwen3": "Qwen3ForCausalLM",
             # qwen3_5 is the top-level multimodal config type; qwen3_5_text is
             # the text-only sub-config. Both map to the text-only adapter so
@@ -544,6 +547,7 @@ def boot(
         # Hybrid/MoE architectures
         "layer_types",
         "moe_intermediate_size",
+        "shared_expert_intermediate_size",
         "norm_eps",
         "attention_bias",
         "lm_head_bias",

transformer_lens/model_bridge/supported_architectures/__init__.py

@@ -177,6 +177,9 @@
 from transformer_lens.model_bridge.supported_architectures.qwen2 import (
     Qwen2ArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.qwen2_moe import (
+    Qwen2MoeArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.qwen3 import (
     Qwen3ArchitectureAdapter,
 )
@@ -268,6 +271,7 @@
     "PhiMoEArchitectureAdapter",
     "QwenArchitectureAdapter",
     "Qwen2ArchitectureAdapter",
+    "Qwen2MoeArchitectureAdapter",
     "Qwen3ArchitectureAdapter",
     "Qwen3MoeArchitectureAdapter",
     "Qwen3NextArchitectureAdapter",