diff --git a/auto_round/modeling/unfused_moe/glm_moe.py b/auto_round/modeling/unfused_moe/glm_moe.py index e68ac9974c..5ef9fcc2c4 100644 --- a/auto_round/modeling/unfused_moe/glm_moe.py +++ b/auto_round/modeling/unfused_moe/glm_moe.py @@ -98,7 +98,10 @@ def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape router_logits = self.gate(hidden_states) - topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) + if isinstance(router_logits, tuple): # transformers >= 5.13.0 + _, topk_weights, topk_indices = router_logits + else: + topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.experts_forward(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) diff --git a/auto_round/modeling/unfused_moe/glm_moe_dsa.py b/auto_round/modeling/unfused_moe/glm_moe_dsa.py index 2fb34de0f0..3a108a8cbb 100644 --- a/auto_round/modeling/unfused_moe/glm_moe_dsa.py +++ b/auto_round/modeling/unfused_moe/glm_moe_dsa.py @@ -88,7 +88,10 @@ def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape router_logits = self.gate(hidden_states) - topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) + if isinstance(router_logits, tuple): # transformers >= 5.13.0 + _, topk_weights, topk_indices = router_logits + else: + topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.experts_forward(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) diff --git a/auto_round/modeling/unfused_moe/glm_moe_light.py b/auto_round/modeling/unfused_moe/glm_moe_light.py index c9506ecb4f..03814650f6 100644 --- a/auto_round/modeling/unfused_moe/glm_moe_light.py +++ b/auto_round/modeling/unfused_moe/glm_moe_light.py @@ -100,7 +100,10 @@ def forward(self, hidden_states): residuals = hidden_states orig_shape = hidden_states.shape router_logits = self.gate(hidden_states) - topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) + if isinstance(router_logits, tuple): # transformers >= 5.13.0 + _, topk_weights, topk_indices = router_logits + else: + topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.experts_forward(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals)