fix embedding ColWiseParallel in qwen model

PaddlePaddle · Feb 27, 2025 · 9e93201 · 9e93201
1 parent 45b5012
commit 9e93201
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/paddlenlp/transformers/qwen/modeling_auto.py b/paddlenlp/transformers/qwen/modeling_auto.py
@@ -538,7 +538,7 @@ def __init__(self, config):
         self.recompute_granularity = config.recompute_granularity
 
         self.wte = nn.Embedding(self.vocab_size, self.embed_dim)
-        self.wte.weight = dist.shard_tensor(self.wte.weight, get_mesh(), [dist.Replicate(), dist.Shard(0)])
+        self.wte.weight = dist.shard_tensor(self.wte.weight, get_mesh(), [dist.Replicate(), dist.Shard(1)])
         self.drop = nn.Dropout(config.emb_dropout_prob)
 
         self.h = nn.LayerList(

diff --git a/paddlenlp/transformers/qwen/modeling_network.py b/paddlenlp/transformers/qwen/modeling_network.py
@@ -667,7 +667,7 @@ def auto_dist_config(self, prefix=""):
             "sp_config": {
                 "parallelize_plan": {
                     f"{prefix}qwen.wte": [
-                        dist.RowWiseParallel(),
+                        dist.ColWiseParallel(),
                         dist.SequenceParallelBegin(),
                     ],
                     f"{prefix}qwen.h.*.attn.c_attn": dist.ColWiseParallel(),
@@ -684,7 +684,7 @@ def auto_dist_config(self, prefix=""):
             },
             "mp_config": {
                 "parallelize_plan": {
-                    f"{prefix}qwen.wte": dist.RowWiseParallel(),
+                    f"{prefix}qwen.wte": dist.ColWiseParallel(),
                     f"{prefix}qwen.h.*.attn.c_attn": dist.ColWiseParallel(),
                     f"{prefix}qwen.h.*.attn.c_proj": dist.RowWiseParallel(),
                     f"{prefix}qwen.h.*.mlp.gate_up_fused_proj": dist.ColWiseParallel(),