[AutoParallel] Handle GradClipByGlobalNorm under AutoParallel Semantic (#65259)

JZ-LIANG · winter-wang · web-flow · commit 2551ac514c84 · 2024-06-27T10:27:30.000+08:00
* special deal with the inplace case in auto parallel pass.

* reshard bug

* bugfix

* unitest

* refine

* trigger ci

---------

Co-authored-by: winter-wang &lt;1030748926@qq.com&gt;
diff --git a/python/paddle/distributed/auto_parallel/placement_type.py b/python/paddle/distributed/auto_parallel/placement_type.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 from typing import cast
 
+import paddle
 from paddle.base.core import Partial, Replicate, Shard
 
 
@@ -28,7 +29,11 @@ def to_placements(dim_map, mesh, partial_idx=[]):
     Returns:
         List[Placement]: a list contains some `paddle.distributed.Placement`.
     """
-    placements = [Replicate() for _ in range(len(mesh.mesh.shape))]
+    if isinstance(mesh, paddle.base.libpaddle.ProcessMesh):
+        shape = mesh.shape
+    else:
+        shape = mesh.mesh.shape
+    placements = [Replicate() for _ in range(len(shape))]
 
     for s in partial_idx:
         placements[s] = Partial()
diff --git a/python/paddle/nn/clip.py b/python/paddle/nn/clip.py
@@ -816,6 +816,27 @@ def _pir_clip(self, params_grads):
         sum_square_list = []
         sum_square_list_fp16 = []
         sum_square_list_fp32 = []
+
+        auto_parallel_pp = False
+        pp_meshes = set()
+        pp_stage0_mesh = None
+        for p, g in params_grads:
+            if p.is_dist_dense_tensor_type():
+                pp_meshes.add(p.dist_attr().process_mesh)
+                if 0 in p.dist_attr().process_mesh.process_ids:
+                    if pp_stage0_mesh is None:
+                        pp_stage0_mesh = p.dist_attr().process_mesh
+                    else:
+                        assert p.dist_attr().process_mesh == pp_stage0_mesh
+
+        if len(pp_meshes) > 1:
+            from paddle.distributed.auto_parallel.placement_type import (
+                to_placements,
+            )
+
+            auto_parallel_pp = True
+            assert pp_stage0_mesh is not None
+
         for p, g in params_grads:
             if g is None:
                 continue
@@ -828,6 +849,19 @@ def _pir_clip(self, params_grads):
                 merge_grad = get_tensor_from_selected_rows(merge_grad)
 
             sum_square = _squared_l2_norm(merge_grad)
+            if (
+                auto_parallel_pp
+                and sum_square.dist_attr().process_mesh != pp_stage0_mesh
+            ):
+                sum_square = paddle.distributed.reshard(
+                    sum_square,
+                    pp_stage0_mesh,
+                    to_placements(
+                        sum_square.dist_attr().dims_mapping,
+                        sum_square.dist_attr().process_mesh,
+                        sum_square.dist_attr().partial_dims,
+                    ),
+                )
             if (
                 sum_square.dtype == DataType.FLOAT16
                 or sum_square.dtype == DataType.BFLOAT16
@@ -867,7 +901,7 @@ def async_add_n(var_list):
         global_norm_var = async_add_n(global_norm_var)
         global_norm_var = paddle.sqrt(global_norm_var)
         max_global_norm = paddle.full(
-            shape=[], dtype=global_norm_var.dtype, fill_value=self.clip_norm
+            shape=[1], dtype=global_norm_var.dtype, fill_value=self.clip_norm
         )
 
         need_clip = False
@@ -895,6 +929,21 @@ def async_add_n(var_list):
                     if clip_var.dtype != g.dtype
                     else clip_var
                 )
+                if (
+                    auto_parallel_pp
+                    and clip_input.dist_attr().process_mesh
+                    != g.dist_attr().process_mesh
+                ):
+                    clip_input = paddle.distributed.reshard(
+                        clip_input,
+                        g.dist_attr().process_mesh,
+                        to_placements(
+                            clip_input.dist_attr().dims_mapping,
+                            clip_input.dist_attr().process_mesh,
+                            clip_input.dist_attr().partial_dims,
+                        ),
+                    )
+
                 new_grad = paddle.multiply(g, clip_input)
                 params_and_grads.append((p, new_grad))
             else:
diff --git a/test/auto_parallel/hybrid_strategy/semi_auto_llama_acc_align.py b/test/auto_parallel/hybrid_strategy/semi_auto_llama_acc_align.py
@@ -78,7 +78,7 @@ def apply_decay_param_fun(x):
         apply_decay_param_fun=apply_decay_param_fun,
         parameters=model.parameters(),
         weight_decay=0.01,
-        grad_clip=None,
+        grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
     )
     return optimizer
 
@@ -151,7 +151,6 @@ def run_dynamic(self):
             input_ids, labels = inputs
             logits = model(input_ids)
             tr_loss_step = criterion(logits, labels)
-
             if self.gradient_accumulation_steps > 1:
                 tr_loss_step /= self.gradient_accumulation_steps
 
@@ -223,6 +222,7 @@ def run_dy2static(self):
                 assert loss is not None
             else:
                 assert loss is None
+
         numpy_array = np.array(loss)
         array_bytes = numpy_array.tobytes()
         loss_md5 = hashlib.md5(array_bytes).hexdigest()
diff --git a/test/auto_parallel/hybrid_strategy/test_semi_auto_llama_acc_align.py b/test/auto_parallel/hybrid_strategy/test_semi_auto_llama_acc_align.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import unittest
 
 import collective.test_communication_api_base as test_base
 
@@ -40,3 +41,7 @@ def test_simple_net_hybrid_strategy_acc(self):
                 "semi_auto_llama_acc_align.py",
                 user_defined_envs=envs,
             )
+
+
+if __name__ == "__main__":
+    unittest.main()  # python run
