[AutoParallel]:fix bug when auto_model use local_tensor (#69462)

blacksheep-Aristotle · web-flow · commit 8f50b879de7c · 2024-11-22T14:34:00.000+08:00
* [AutoParallel]:fix bug when auto_model use local_tensor

* [AutoParallel]:fix sharding bug
diff --git a/python/paddle/distributed/auto_parallel/intermediate/parallel_base.py b/python/paddle/distributed/auto_parallel/intermediate/parallel_base.py
@@ -139,11 +139,54 @@ def parallelize_model(self):
         if self.sharding_parallelizer is not None:
             assert callable(self.sharding_parallelizer)
             self.model = self.sharding_parallelizer(self.model)
-
+        self._shard_all_param(self.model)
         self.is_parallelized = True
 
         return self.model
 
+    def _shard_all_param(self, model):
+        param_name_to_shard_param = {}
+
+        def shard_layer_param(layer):
+            if self.pp_parallelizer is not None:
+                assert hasattr(layer, "pipeline_stage_index")
+            layer_attrs = dir(layer)
+            for layer_attr in layer_attrs:
+                if getattr(layer, layer_attr) is not None and is_tensor(
+                    getattr(layer, layer_attr)
+                ):
+                    layer_tensor = getattr(layer, layer_attr)
+                    if not layer_tensor.is_dist():
+                        tensor_name = layer_tensor.name
+                        if tensor_name in param_name_to_shard_param:
+                            setattr(
+                                layer,
+                                layer_attr,
+                                param_name_to_shard_param[tensor_name],
+                            )
+                        else:
+                            ipp = (
+                                layer.pipeline_stage_index
+                                if hasattr(layer, "pipeline_stage_index")
+                                else 0
+                            )
+                            mesh = self.get_mesh(ipp)
+                            layer_tensor = dist.shard_tensor(
+                                layer_tensor,
+                                mesh,
+                                [
+                                    dist.Replicate()
+                                    for _ in range(len(mesh._shape))
+                                ],
+                            )
+                            param_name_to_shard_param[tensor_name] = (
+                                layer_tensor
+                            )
+                            setattr(layer, layer_attr, layer_tensor)
+
+        for name, layer in model.named_sublayers():
+            shard_layer_param(layer)
+
 
 def parallelize_model_and_optimizer(model, optimizer=None):
     if not isinstance(model, ParallelModel):
diff --git a/test/auto_parallel/hybrid_strategy/parallel_api.py b/test/auto_parallel/hybrid_strategy/parallel_api.py
@@ -193,7 +193,7 @@ def check_mp(self, layer):
                         dist.Replicate(),
                         dist.Shard(0),
                     ]
-                    assert sub_layer.bias.placements is None
+                    # assert sub_layer.bias.placements is None
                 if 'down_proj' in name:
                     assert sub_layer.weight.placements == [
                         dist.Replicate(),
@@ -217,6 +217,7 @@ def parallel_model(self, layer, optimizer=None):
             if not self.sequence_parallel:
                 plan = {
                     "llama.embed_tokens": ColWiseParallel(),
+                    "llama.position_embedding": ColWiseParallel(),
                     "llama.layers.*.self_attn.q_proj": ColWiseParallel(),
                     "llama.layers.*.self_attn.k_proj": ColWiseParallel(),
                     "llama.layers.*.self_attn.v_proj": ColWiseParallel(),
@@ -230,6 +231,7 @@ def parallel_model(self, layer, optimizer=None):
                 if self.prepare_input_output:
                     plan = {
                         "llama.embed_tokens": ColWiseParallel(),
+                        "llama.position_embedding": ColWiseParallel(),
                         "llama.layers.*.self_attn.q_proj": ColWiseParallel(),
                         "llama.layers.*.self_attn.k_proj": ColWiseParallel(),
                         "llama.layers.*.self_attn.v_proj": ColWiseParallel(),
@@ -248,6 +250,10 @@ def parallel_model(self, layer, optimizer=None):
                             ColWiseParallel(),
                             SequenceParallelBegin(),
                         ],
+                        "llama.position_embedding": [
+                            ColWiseParallel(),
+                            SequenceParallelBegin(),
+                        ],
                         "llama.layers.*.self_attn.q_proj": ColWiseParallel(),
                         "llama.layers.*.self_attn.k_proj": ColWiseParallel(),
                         "llama.layers.*.self_attn.v_proj": ColWiseParallel(),
@@ -281,12 +287,18 @@ def parallel_model(self, layer, optimizer=None):
             return layer
         return layer, optimizer
 
-    def run_llama(self, to_static=0):
+    def run_llama(
+        self, share_embedding=False, position_embedding=False, to_static=0
+    ):
         if self.config.use_lazy_init:
             with LazyGuard():
-                model = LlamaForCausalLM(self.config)
+                model = LlamaForCausalLM(
+                    self.config, share_embedding, position_embedding
+                )
         else:
-            model = LlamaForCausalLM(self.config)
+            model = LlamaForCausalLM(
+                self.config, share_embedding, position_embedding
+            )
 
         lr_scheduler = paddle.optimizer.lr.LinearWarmup(
             learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
@@ -422,10 +434,12 @@ def run_llama(self, to_static=0):
                 if step >= 10:
                     break
 
-    def run_test_cases(self):
-        self.run_llama(to_static=0)
-        self.run_llama(to_static=1)
+    def run_test_cases(self, share_embedding=False, position_embedding=False):
+        self.run_llama(share_embedding, position_embedding, 0)
+        self.run_llama(share_embedding, position_embedding, 1)
 
 
 if __name__ == '__main__':
-    TestParallelAPI().run_test_cases()
+    share_embedding = int(os.getenv("test_share_embedding", "0"))
+    position_embedding = int(os.getenv("test_position_embedding", "0"))
+    TestParallelAPI().run_test_cases(share_embedding, position_embedding)
diff --git a/test/auto_parallel/hybrid_strategy/single_llama_model.py b/test/auto_parallel/hybrid_strategy/single_llama_model.py
@@ -151,7 +151,7 @@ def forward(self, hidden_states):
 
 
 class LlamaModel(nn.Layer):
-    def __init__(self, config):
+    def __init__(self, config, position_embedding=False):
         super().__init__()
         self.config = config
         self.vocab_size = self.config.vocab_size
@@ -162,6 +162,15 @@ def __init__(self, config):
             self.hidden_size,
         )
 
+        self.position_embedding = (
+            nn.Embedding(
+                self.vocab_size,
+                self.hidden_size,
+            )
+            if position_embedding
+            else None
+        )
+
         decoder_layers = []
         for i in range(self.config.num_hidden_layers):
             decoder_layers.append(LlamaDecoderLayer(self.config))
@@ -171,6 +180,12 @@ def __init__(self, config):
 
     def forward(self, input_ids):
         hidden_states = self.embed_tokens(input_ids)
+        if self.position_embedding is not None:
+            ones = paddle.ones(input_ids.shape, dtype="int64")
+            seq_length = paddle.cumsum(ones, axis=-1)
+            position_ids = seq_length - ones
+            position_embeddings = self.position_embedding(position_ids)
+            hidden_states = hidden_states + position_embeddings
 
         for idx, (decoder_layer) in enumerate(self.layers):
             hidden_states = decoder_layer(hidden_states)
@@ -181,17 +196,23 @@ def forward(self, input_ids):
 
 
 class LlamaLMHead(nn.Layer):
-    def __init__(self, config):
+    def __init__(self, config, weight=None):
         super().__init__()
         self.config = config
-
-        self.weight = self.create_parameter(
-            shape=[self.config.hidden_size, self.config.vocab_size],
-            dtype=paddle.get_default_dtype(),
-        )
+        self.transpose_y = False
+        if weight is not None:
+            self.weight = weight
+            self.transpose_y = True
+        else:
+            self.weight = self.create_parameter(
+                shape=[self.config.hidden_size, self.config.vocab_size],
+                dtype=paddle.get_default_dtype(),
+            )
 
     def forward(self, hidden_states):
-        logits = paddle.matmul(hidden_states, self.weight, transpose_y=False)
+        logits = paddle.matmul(
+            hidden_states, self.weight, transpose_y=self.transpose_y
+        )
         return logits
 
 
@@ -226,12 +247,17 @@ def forward(self, prediction_scores, masked_lm_labels):
 class LlamaForCausalLM(nn.Layer):
     enable_to_static_method = True
 
-    def __init__(self, config):
+    def __init__(self, config, share_embedding=False, position_embedding=False):
         super().__init__()
         self.config = config
 
-        self.llama = LlamaModel(self.config)
-        self.lm_head = LlamaLMHead(self.config)
+        self.llama = LlamaModel(self.config, position_embedding)
+        if share_embedding:
+            self.lm_head = LlamaLMHead(
+                self.config, self.llama.embed_tokens.weight
+            )
+        else:
+            self.lm_head = LlamaLMHead(self.config)
 
     def forward(self, input_ids=None):
         input_ids.stop_gradient = True
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_1d.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_1d.py
@@ -36,6 +36,12 @@ def setUp(self):
             "amp_dtype": ["bfloat16"],
             "amp_master_grad": ["False"],
             "sharding_stage": ["0", "1"],
+            "test_share_embedding": [
+                "1",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
         }
 
     def test_simple_net_dp2(self):
@@ -70,6 +76,12 @@ def setUp(self):
             "amp_dtype": ["bfloat16"],
             "amp_master_grad": ["False"],
             "num_hidden_layers": ["2", "4"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
         }
 
     def test_simple_net_pp2(self):
@@ -106,6 +118,12 @@ def setUp(self):
             "use_lazy_init": ["true", "false"],
             "sequence_parallel": ["true", "false"],
             "prepare_input_output": ["true", "false"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
         }
 
     def test_simple_net_mp2(self):
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_2d.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_2d.py
@@ -38,6 +38,12 @@ def setUp(self):
             "use_lazy_init": ["true"],
             "sequence_parallel": ["true"],
             "prepare_input_output": ["false"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
         }
 
     def test_simple_net_mp2_pp2(self):
@@ -74,6 +80,12 @@ def setUp(self):
             "use_lazy_init": ["true"],
             "num_hidden_layers": ["2", "4"],
             "sharding_stage": ["0"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
         }
 
     def test_simple_net_dp2_pp2(self):
@@ -111,6 +123,12 @@ def setUp(self):
             "sequence_parallel": ["true"],
             "prepare_input_output": ["false"],
             "sharding_stage": ["0"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
         }
 
     def test_simple_net_mp2_pp2(self):
diff --git a/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_3d.py b/test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_3d.py
@@ -39,6 +39,12 @@ def setUp(self):
             "sequence_parallel": ["true"],
             "prepare_input_output": ["false"],
             "sharding_stage": ["0", "1"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
         }
 
     def test_simple_net_dp2_mp2_pp2(self):
