From 2dd96f8a4d323839f91ff04b25ca468cef3f4c54 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sun, 16 Aug 2026 14:49:46 +0000 Subject: [PATCH] fix(glm): version-robust GLM-MoE-DSA routing across transformers 5.12-5.15 (#144) SyncGlmMoeDsaMoEBlock borrowed transformers' internal GlmMoeDsaMoE.route_tokens_to_experts, which 5.15 removed after moving routing into GlmMoeDsaTopkRouter (now returns (router_logits, topk_weights, topk_indices)). _route is now adaptive: unpack the router tuple on 5.15+, else apply an inlined grouped-sigmoid-top-k (mirrors 5.12 route_tokens_to_experts). Removes the fragile dep+import. Verified: 5.12 numerical parity + method-absent path; 5.15 tuple-contract unpack + forward. Test: module-skip -> per-test skipif on parity only. Closes #144. --- moe_infinity/models/glm_moe_dsa.py | 51 ++++++++++++++++++++------- tests/python/unit/test_glm_routing.py | 10 +++--- 2 files changed, 44 insertions(+), 17 deletions(-) diff --git a/moe_infinity/models/glm_moe_dsa.py b/moe_infinity/models/glm_moe_dsa.py index ef7bed3c..614d331e 100644 --- a/moe_infinity/models/glm_moe_dsa.py +++ b/moe_infinity/models/glm_moe_dsa.py @@ -8,14 +8,13 @@ try: from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import ( GlmMoeDsaMLP, - GlmMoeDsaMoE, GlmMoeDsaTopkRouter, ) _GLM_AVAILABLE = True except ImportError: _GLM_AVAILABLE = False - GlmMoeDsaMoE = GlmMoeDsaTopkRouter = GlmMoeDsaMLP = None + GlmMoeDsaTopkRouter = GlmMoeDsaMLP = None class SyncGlmMoeDsaMoEBlock(nn.Module): @@ -59,9 +58,6 @@ def __init__(self, config): intermediate_size=config.moe_intermediate_size * config.n_shared_experts, ) - self._hf_route_tokens = getattr( - GlmMoeDsaMoE, "route_tokens_to_experts", None - ) def _route(self, hidden_flat: torch.Tensor): dev = hidden_flat.device @@ -69,13 +65,44 @@ def _route(self, hidden_flat: torch.Tensor): self.gate.e_score_correction_bias = ( self.gate.e_score_correction_bias.to(dev) ) - router_logits = self.gate(hidden_flat) - if self._hf_route_tokens is None: - raise RuntimeError( - "GLM-MoE-DSA routing requires a transformers build providing " - "GlmMoeDsaMoE.route_tokens_to_experts (removed in 5.15+)" - ) - return self._hf_route_tokens(self, router_logits) + routed = self.gate(hidden_flat) + if isinstance(routed, tuple): + # transformers >= 5.15 moved routing into GlmMoeDsaTopkRouter, which + # returns (router_logits, topk_weights, topk_indices). + _, topk_weights, topk_indices = routed + return topk_indices, topk_weights + return self._route_tokens_to_experts(routed) + + def _route_tokens_to_experts(self, router_logits: torch.Tensor): + # Mirror of transformers GlmMoeDsaMoE.route_tokens_to_experts (grouped + # sigmoid top-k). Inlined so offload routing stays correct on builds + # that no longer expose that method (see #144); sigmoid gating, not + # softmax. + router_logits = router_logits.sigmoid() + scores = router_logits + self.gate.e_score_correction_bias + group_scores = ( + scores.view(-1, self.n_group, self.n_routed_experts // self.n_group) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) + group_idx = torch.topk( + group_scores, k=self.topk_group, dim=-1, sorted=False + )[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(-1, self.n_group, self.n_routed_experts // self.n_group) + .reshape(-1, self.n_routed_experts) + ) + scores = scores.masked_fill(~score_mask.bool(), float("-inf")) + topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)[1] + topk_weights = router_logits.gather(1, topk_indices) + if self.norm_topk_prob: + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights = topk_weights / denominator + topk_weights = topk_weights * self.routed_scaling_factor + return topk_indices, topk_weights def _local_experts(self, hidden_flat, router_mask, routing_weights_mask): return torch.zeros_like(hidden_flat) diff --git a/tests/python/unit/test_glm_routing.py b/tests/python/unit/test_glm_routing.py index abb1d293..187003c5 100644 --- a/tests/python/unit/test_glm_routing.py +++ b/tests/python/unit/test_glm_routing.py @@ -12,11 +12,7 @@ GlmMoeDsaMoE as _GlmMoeDsaMoE, ) -if not hasattr(_GlmMoeDsaMoE, "route_tokens_to_experts"): - pytest.skip( - "transformers dropped GlmMoeDsaMoE.route_tokens_to_experts (5.15+)", - allow_module_level=True, - ) +_HAS_HF_ROUTER = hasattr(_GlmMoeDsaMoE, "route_tokens_to_experts") def _tiny_config(): @@ -53,6 +49,10 @@ def test_exported(): assert "SyncGlmMoeDsaMoEBlock" in m.__all__ +@pytest.mark.skipif( + not _HAS_HF_ROUTER, + reason="transformers build lacks GlmMoeDsaMoE.route_tokens_to_experts to compare against", +) def test_routing_parity(): from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import ( GlmMoeDsaMoE,