init shared parameters

Author: waliwali777

paddlenlp/transformers/gpt/modeling_auto.py

@@ -529,7 +529,7 @@ def __init__(self, config: GPTConfig, ipp=None):
         self.linear2 = nn.Linear(config.intermediate_size, config.hidden_size, bias_attr=True)
 
         self.linear1.weight = dist.shard_tensor(self.linear1.weight, get_mesh(ipp), [dist.Replicate(), dist.Shard(1)])
-        self.linear1.bias = dist.shard_tensor(self.linear1.bias, get_mesh(ipp), [dist.Replicate(), dist.Shard(0)])
+        self.linear1.bias = dist.shard_tensor(self.linear1.bias, get_mesh(ipp), [dist.Replicate(), dist.Replicate()])
         self.linear2.weight = dist.shard_tensor(self.linear2.weight, get_mesh(ipp), [dist.Replicate(), dist.Shard(0)])
         self.linear2.bias = dist.shard_tensor(self.linear2.bias, get_mesh(ipp), [dist.Replicate(), dist.Replicate()])
         # fix : change nn.LayerNorm(config.hidden_size, epsilon=1e-5, bias_attr=True) to GPTLayerNorm()
@@ -658,7 +658,7 @@ def __init__(
             config.hidden_size,
         )
         self.word_embeddings.weight = dist.shard_tensor(
-            self.word_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Replicate()]
+            self.word_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Shard(1)]
         )
         self.position_embeddings.weight = dist.shard_tensor(
             self.position_embeddings.weight, get_mesh(), [dist.Replicate(), dist.Shard(1)]
@@ -699,6 +699,7 @@ def forward(self, input_ids, position_ids=None, inputs_embeddings=None):
         # The 'with' block ensures the correct seed context is used
         with seed_guard_context(current_seed):
             embeddings = self.dropout(embeddings)
+        embeddings = dist.reshard(embeddings, get_mesh(), [dist.Replicate(), dist.Replicate()])
         return embeddings
 
 
@@ -1176,7 +1177,7 @@ def __init__(self, config: GPTConfig, embedding_weights=None, ipp=None):
                 shape=[config.vocab_size, config.hidden_size],
                 dtype=paddle.get_default_dtype(),
             )
-            self.weight = dist.shard_tensor(self.weight, get_mesh(self.ipp), [dist.Replicate(), dist.Shard(0)])
+            self.weight = dist.shard_tensor(self.weight, get_mesh(self.ipp), [dist.Replicate(), dist.Shard(1)])
 
     def forward(self, hidden_states, tensor_parallel_output=None):
 

paddlenlp/transformers/gpt/modeling_auto_pp.py

@@ -139,6 +139,13 @@ def manual_model_split(model, stage_idx, group, mode, pp_degree):
 
     layer_lists = model.layers
 
+    
+    shared_params_names = {
+            "gpt_shared_weight": ["embedding_0.w_0.dist", "gptlm_head_auto_0.w_0.dist"]
+        }
+
+    shared_mp = build_shared_param_map(model, shared_params_names)
+
     def _build_stage(model, stage_idx, group):
         new_model = None
         if stage_idx == 0:
@@ -151,7 +158,7 @@ def _build_stage(model, stage_idx, group):
             new_model = GPTChunk(
                 layer_lists[stage_idx * chunk_size : (stage_idx + 1) * chunk_size], is_first=False, is_last=False
             )
-        stage = PipelineStage(new_model, stage_idx, chunk_num, group=group)
+        stage = PipelineStage(new_model, stage_idx, chunk_num, group=group, shared_map=shared_mp)
         return stage
 
     stages = []
@@ -160,6 +167,27 @@ def _build_stage(model, stage_idx, group):
         stages.append(stage)
     return stages
 
+def build_shared_param_map(model, shared_params_names):
+    shared_mp = []
+    for key, pair in shared_params_names.items():
+        assert len(pair) == 2, (
+            "Only exactly two parameters are supported for sharing."
+        )
+        ori_name = pair[0]
+        sync_name = pair[1]
+        ori_param = get_param_from_name(ori_name, model)
+        sync_param = get_param_from_name(sync_name, model)
+        shared_mp.append({
+            "params": [ori_param, sync_param]
+        })
+    return shared_mp
+
+def get_param_from_name(param_name, model):
+    for param in model.parameters():
+        if param.name == param_name:
+            return param
+    raise ValueError(f"{param_name} not found in model parameters")
+
 
 def get_gpt_pp_schedule(model, n_microbatches, loss_fn, mode, pp_degree, group):
     assert mode in ["VPP", "1F1B", "FThenB"]