fix(glm5.2): restrict MoE weights to local PP layers (#33793)

2026-08-12 12:32:08 -07:00
parent 4f883636a2
commit 2b4381956f
1 changed files with 3 additions and 3 deletions
@@ -2983,9 +2983,9 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):

        self._routed_experts_weights_of_layer = LazyValue(
            lambda: {
-                layer_id: layer.mlp.get_moe_weights()
-                for layer_id, layer in enumerate(self.model.layers)
-                if isinstance(layer.mlp, DeepseekV2MoE)
+                layer_id: self.model.layers[layer_id].mlp.get_moe_weights()
+                for layer_id in range(self.model.start_layer, self.model.end_layer)
+                if isinstance(self.model.layers[layer_id].mlp, DeepseekV2MoE)
            }
        )
        self.capture_aux_hidden_states = False