fix_qwen3_moe (#10801)

gongel · web-flow · commit baaa08047764 · 2025-07-04T14:45:25.000+08:00
diff --git a/paddlenlp/transformers/moe_gate.py b/paddlenlp/transformers/moe_gate.py
@@ -496,7 +496,7 @@ def topkgating(
         top_gate = top_gate * self.routed_scaling_factor
 
         # get topk mask
-        mask = paddle.zeros_like(gates).put_along_axis(top_idx, paddle.to_tensor(1.0, dtype="float32"), axis=1)
+        mask = paddle.zeros_like(gates).put_along_axis(top_idx, paddle.to_tensor(1.0, dtype=gates.dtype), axis=1)
         if hasattr(self.config, "seq_aux") and self.config.seq_aux:
             l_aux = self._cal_seq_aux_loss(gates_ori, self.top_k, top_idx)
         else:
diff --git a/paddlenlp/transformers/qwen2/modeling.py b/paddlenlp/transformers/qwen2/modeling.py
@@ -590,7 +590,9 @@ def __init__(self, config: Qwen2Config, layerwise_recompute: bool = True, skip_r
                 )
                 self.k_proj = ColumnParallelLinear(self.hidden_size, self.config.num_key_value_heads * self.head_dim, has_bias=self.has_bias, gather_output=False)  # fmt:skip
                 self.v_proj = ColumnParallelLinear(self.hidden_size, self.config.num_key_value_heads * self.head_dim, has_bias=self.has_bias, gather_output=False)  # fmt:skip
-            self.o_proj = RowParallelLinear(self.hidden_size, self.hidden_size, has_bias=False, input_is_parallel=True)
+            self.o_proj = RowParallelLinear(
+                self.num_attention_heads * self.head_dim, self.hidden_size, has_bias=False, input_is_parallel=True
+            )
         else:
             if self.fuse_attention_qkv:
                 self.qkv_proj = Linear(
