FollowUp: Fix unnecessary synchronized memcpy in hybrid parallel training (#57215)

mingxu1067 · Tom-Zheng · web-flow · commit 7d1608b5f00f · 2023-09-15T13:12:44.000+08:00
* Remove synchronization in hybrid parallel optimizer

* Fix sync memcpy in TP

* Fix wrong type of inputs to tensor.multiply_.

---------

Co-authored-by: Tian Zheng (Engrg-Hardware 1) &lt;tizheng@nvidia.com&gt;
diff --git a/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py
@@ -14,8 +14,6 @@
 
 import os
 
-import numpy as np
-
 import paddle
 from paddle import framework
 from paddle.autograd import no_grad
@@ -139,65 +137,56 @@ def _dygraph_clip(self, params_grads):
                     elif g.dtype == paddle.float32:
                         sum_square_not_dist_fp32.append(sum_square)
 
+        def async_add_n(var_list):
+            return paddle.stack(var_list).sum()
+
         # global norm of distributed FP16 params_and_grads
         if len(sum_square_dist_fp16) == 0:
-            global_norm_dist_fp16 = paddle.to_tensor(
-                np.array(0.0), dtype=paddle.float32
-            )
+            global_norm_dist_fp16 = paddle.zeros((1,), dtype=paddle.float32)
         else:
-            global_norm_dist_fp16 = paddle.add_n(sum_square_dist_fp16)
+            global_norm_dist_fp16 = async_add_n(sum_square_dist_fp16)
             global_norm_dist_fp16 = paddle.cast(
                 global_norm_dist_fp16, dtype=paddle.float32
             )
 
         # global norm of non-distributed FP16 params_and_grads
         if len(sum_square_not_dist_fp16) == 0:
-            global_norm_not_dist_fp16 = paddle.to_tensor(
-                np.array(0.0), dtype=paddle.float32
-            )
+            global_norm_not_dist_fp16 = paddle.zeros((1,), dtype=paddle.float32)
         else:
-            global_norm_not_dist_fp16 = paddle.add_n(sum_square_not_dist_fp16)
+            global_norm_not_dist_fp16 = async_add_n(sum_square_not_dist_fp16)
             global_norm_not_dist_fp16 = paddle.cast(
                 global_norm_not_dist_fp16, dtype=paddle.float32
             )
 
         # global norm of distributed BF16 params_and_grads
         if len(sum_square_dist_bf16) == 0:
-            global_norm_dist_bf16 = paddle.to_tensor(
-                np.array(0.0), dtype=paddle.float32
-            )
+            global_norm_dist_bf16 = paddle.zeros((1,), dtype=paddle.float32)
         else:
-            global_norm_dist_bf16 = paddle.add_n(sum_square_dist_bf16)
+            global_norm_dist_bf16 = async_add_n(sum_square_dist_bf16)
             global_norm_dist_bf16 = paddle.cast(
                 global_norm_dist_bf16, dtype=paddle.float32
             )
 
         # global norm of non-distributed FP16 params_and_grads
         if len(sum_square_not_dist_bf16) == 0:
-            global_norm_not_dist_bf16 = paddle.to_tensor(
-                np.array(0.0), dtype=paddle.float32
-            )
+            global_norm_not_dist_bf16 = paddle.zeros((1,), dtype=paddle.float32)
         else:
-            global_norm_not_dist_bf16 = paddle.add_n(sum_square_not_dist_bf16)
+            global_norm_not_dist_bf16 = async_add_n(sum_square_not_dist_bf16)
             global_norm_not_dist_bf16 = paddle.cast(
                 global_norm_not_dist_bf16, dtype=paddle.float32
             )
 
         # global norm of distributed FP32 params_and_grads
         if len(sum_square_dist_fp32) == 0:
-            global_norm_dist_fp32 = paddle.to_tensor(
-                np.array(0.0), dtype=paddle.float32
-            )
+            global_norm_dist_fp32 = paddle.zeros((1,), dtype=paddle.float32)
         else:
-            global_norm_dist_fp32 = paddle.add_n(sum_square_dist_fp32)
+            global_norm_dist_fp32 = async_add_n(sum_square_dist_fp32)
 
         # global norm of non-distributed FP32 params_and_grads
         if len(sum_square_not_dist_fp32) == 0:
-            global_norm_not_dist_fp32 = paddle.to_tensor(
-                np.array(0.0), dtype=paddle.float32
-            )
+            global_norm_not_dist_fp32 = paddle.zeros((1,), dtype=paddle.float32)
         else:
-            global_norm_not_dist_fp32 = paddle.add_n(sum_square_not_dist_fp32)
+            global_norm_not_dist_fp32 = async_add_n(sum_square_not_dist_fp32)
 
         global_norm_var_dist = (
             global_norm_dist_fp16
@@ -224,7 +213,7 @@ def _dygraph_clip(self, params_grads):
         clip_var = paddle.divide(
             x=max_global_norm,
             y=paddle.maximum(x=global_norm_var_fp32, y=max_global_norm)
-            + paddle.to_tensor(np.array(1.0e-6), dtype=paddle.float32),
+            + paddle.full(shape=[], dtype=paddle.float32, fill_value=1.0e-6),
         )
         clip_var_fp16 = paddle.cast(clip_var, paddle.float16)
 
@@ -338,7 +327,13 @@ def _insert_sync(self, sync_var, src, mp_group, sync_mode):
             paddle.distributed.all_reduce(
                 sync_var, group=mp_group, sync_op=True
             )
-            sync_var.scale_(1.0 / mp_group.nranks)
+            sync_var.multiply_(
+                paddle.full(
+                    shape=[],
+                    dtype=sync_var.dtype,
+                    fill_value=(1.0 / mp_group.nranks),
+                )
+            )
 
     def _filter_fn(self, param, strategy):
         p_name = param.name
diff --git a/python/paddle/distributed/fleet/utils/hybrid_parallel_util.py b/python/paddle/distributed/fleet/utils/hybrid_parallel_util.py
@@ -130,7 +130,7 @@ def _broadcast_data_help(data, shape, dtype, hcg):
     src_rank = hcg.get_model_parallel_group_src_rank()
     mp_rank = hcg.get_model_parallel_rank()
 
-    shape_gpu = paddle.to_tensor(shape, dtype="int32")
+    shape_gpu = shape._copy_to(data.place, False)
     paddle.distributed.broadcast(
         shape_gpu, src=src_rank, group=model_parallel_group, sync_op=True
     )
@@ -192,7 +192,7 @@ def broadcast_input_data(hcg, *inputs, **kwargs):
                     v_gpu = v._copy_to(place, True)
                     v._clear_data()
                     v_gpu._share_buffer_to(v)
-                _broadcast_data_help(v, v.shape, v.dtype, hcg)
+                _broadcast_data_help(v, paddle.shape(v), v.dtype, hcg)
         else:
             _broadcast_object_list_help(v, hcg)
 
@@ -203,7 +203,7 @@ def broadcast_input_data(hcg, *inputs, **kwargs):
                     v_gpu = v._copy_to(place, True)
                     v._clear_data()
                     v_gpu._share_buffer_to(v)
-                _broadcast_data_help(v, v.shape, v.dtype, hcg)
+                _broadcast_data_help(v, paddle.shape(v), v.dtype, hcg)
             kwargs[k] = v
         else:
             kwargs[k] = _broadcast_object_list_help(v, hcg)
diff --git a/python/paddle/nn/clip.py b/python/paddle/nn/clip.py
@@ -678,21 +678,24 @@ def _dygraph_clip(self, params_grads):
         ):
             return params_grads
 
+        def async_add_n(var_list):
+            return paddle.stack(var_list).sum()
+
         sum_dtype = 'float64' if len(sum_square_list) > 0 else "float32"
         global_norm_var = []
         if len(sum_square_list_fp16) > 0:
-            global_norm_var_fp16 = paddle.add_n(sum_square_list_fp16)
+            global_norm_var_fp16 = async_add_n(sum_square_list_fp16)
             global_norm_var.append(global_norm_var_fp16.astype(sum_dtype))
         if len(sum_square_list_fp32) > 0:
-            global_norm_var_fp32 = paddle.add_n(sum_square_list_fp32)
+            global_norm_var_fp32 = async_add_n(sum_square_list_fp32)
             if sum_dtype == 'float32':
                 global_norm_var.append(global_norm_var_fp32)
             else:
                 global_norm_var.append(global_norm_var_fp32.astype(sum_dtype))
         if len(sum_square_list) > 0:
-            global_norm_var_fp64 = paddle.add_n(sum_square_list)
+            global_norm_var_fp64 = async_add_n(sum_square_list)
             global_norm_var.append(global_norm_var_fp64)
-        global_norm_var = paddle.add_n(global_norm_var)
+        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
