fix(compat): handle DeepseekV2MoE→DeepseekV2Moe Transformers rename

moe_infinity/models/deepseek.py

@@ -48,25 +48,32 @@ def __init__(self, config):
         self.num_expert = config.n_routed_experts
 
         if self.config.model_type == "deepseek_v2":
+            import transformers.models.deepseek_v2.modeling_deepseek_v2 as _dsv2
             from transformers.models.deepseek_v2.modeling_deepseek_v2 import (
                 DeepseekV2MLP,
-                DeepseekV2MoEGate,
             )
 
             self.mlp_cls = DeepseekV2MLP
-            self.gate_cls = DeepseekV2MoEGate
+            # DeepseekV2MoEGate was removed in newer Transformers; fall back to
+            # the local DeepseekMoEGate which has the same raw-logits interface.
+            _gate_cls = getattr(_dsv2, "DeepseekV2MoEGate", None)
+            self.gate_cls = (
+                _gate_cls if _gate_cls is not None else DeepseekMoEGate
+            )
         if self.config.model_type == "deepseek_v3":
             from transformers.models.deepseek_v3.modeling_deepseek_v3 import (
                 DeepseekV3MLP,
             )
 
             self.mlp_cls = DeepseekV3MLP
-            # V3 upstream has no standalone gate; use V2 gate (same interface)
-            from transformers.models.deepseek_v2.modeling_deepseek_v2 import (
-                DeepseekV2MoEGate,
-            )
+            # V3 upstream has no standalone gate; use V2 gate when available,
+            # otherwise fall back to the local DeepseekMoEGate.
+            import transformers.models.deepseek_v2.modeling_deepseek_v2 as _dsv2
 
-            self.gate_cls = DeepseekV2MoEGate
+            _gate_cls = getattr(_dsv2, "DeepseekV2MoEGate", None)
+            self.gate_cls = (
+                _gate_cls if _gate_cls is not None else DeepseekMoEGate
+            )
 
         self.experts = nn.ModuleList(
             [

tests/python/ops/test_deepseek_v2_gate_consistency.py

@@ -41,12 +41,23 @@ def decorator(fn):
 
 _ensure_nvtx_stub_has_annotate()
 
+import transformers.models.deepseek_v2.modeling_deepseek_v2 as _dsv2_mod
 from transformers import DeepseekV2Config
-from transformers.models.deepseek_v2.modeling_deepseek_v2 import (
-    DeepseekV2MoE,
+
+DeepseekV2MoE = getattr(_dsv2_mod, "DeepseekV2Moe", None) or getattr(
+    _dsv2_mod, "DeepseekV2MoE", None
 )
-from transformers.models.deepseek_v2.modeling_deepseek_v2 import (
-    DeepseekV2MoEGate as MoEGate,
+if DeepseekV2MoE is None:
+    raise ImportError(
+        "Neither 'DeepseekV2Moe' nor 'DeepseekV2MoE' found in "
+        "transformers.models.deepseek_v2.modeling_deepseek_v2"
+    )
+
+MoEGate = getattr(_dsv2_mod, "DeepseekV2MoEGate", None)
+
+requires_moe_gate = pytest.mark.skipif(
+    MoEGate is None,
+    reason="DeepseekV2MoEGate removed in this Transformers version",
 )
 
 from moe_infinity.models.deepseek import DeepseekMoEBlock, DeepseekMoEGate
@@ -132,6 +143,7 @@ def wait_dispatch_local(self):
 
 
 @requires_cuda
+@requires_moe_gate
 def test_v2_gate_equivalence(seed_everything):
     """V2-Lite config: native MoEGate and simplified DeepseekMoEGate must
     produce identical routing because topk_method='greedy' with
@@ -185,6 +197,7 @@ def test_v2_gate_equivalence(seed_everything):
 
 
 @requires_cuda
+@requires_moe_gate
 def test_v2_gate_group_limited_greedy(seed_everything):
     """V2-full config: simplified gate MUST differ from native gate because
     group_limited_greedy constrains which expert groups can be selected."""