Fix PP for Qwen3 MoE (#6709)

This commit is contained in:
iLeGend
2025-05-29 14:06:18 +08:00
committed by GitHub
parent 844a8f42c7
commit e06b076105
2 changed files with 60 additions and 4 deletions

View File

@@ -812,9 +812,9 @@ class Qwen3MoeForCausalLM(nn.Module):
logger.warning(f"Parameter {name} not found in params_dict")
self.routed_experts_weights_of_layer = {
layer_id: layer.mlp.get_moe_weights()
for layer_id, layer in enumerate(self.model.layers)
if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock)
layer_id: self.model.layers[layer_id].mlp.get_moe_weights()
for layer_id in range(self.start_layer, self.end_layer)
if isinstance(self.model.layers[layer_id].mlp, Qwen3MoeSparseMoeBlock)
}
@classmethod