Fix the allocation for dbias_out when main_grad is enabled and add logging in sharding comm-overlap hook. (#63433)

Xreki · web-flow · commit 774b7f2e2384 · 2024-04-15T18:54:06.000+08:00
* Fix the allocation for dbias_out when main_grad is enabled and add logging in sharding comm-overlap hook.

* Fix memory use error.
diff --git a/paddle/phi/kernels/fusion/gpu/fused_linear_param_grad_add_kernel.cu b/paddle/phi/kernels/fusion/gpu/fused_linear_param_grad_add_kernel.cu
@@ -24,6 +24,7 @@
 #include "paddle/phi/common/amp_type_traits.h"
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/kernels/elementwise_add_kernel.h"
+#include "paddle/phi/kernels/empty_kernel.h"
 #include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
@@ -68,23 +69,44 @@ void FusedLinearParamGradAddImpl(const Context &ctx,
 
   if (!has_bias) return;
 
+  // if dbias is given, dbias_out will share memory with dbias:
+  //       dbias_tmp = sum(dout), dbias_out = dbias + dbias_tmp
+  // else: dbias_out = sum(dout)
+  DenseTensor dbias_tmp_tensor;
+  if (dbias) {
+    if (kIsMultiPrecision) {
+      dbias_tmp_tensor = phi::EmptyLike<MT, Context>(ctx, dbias.get());
+    } else {
+      dbias_tmp_tensor = phi::EmptyLike<T, Context>(ctx, dbias.get());
+    }
+  }
+  DenseTensor *dbias_tmp = !dbias ? dbias_out : &dbias_tmp_tensor;
+
   if (!fuse_bias_grad) {
     auto dout_copy = dout;
     dout_copy.Resize({M, N});
     if (kIsMultiPrecision) {
-      *dbias_out = phi::Sum<T, Context>(
-          ctx, dout_copy, {0}, phi::CppTypeToDataType<MT>::Type(), false);
+      phi::SumKernel<T, Context>(ctx,
+                                 dout_copy,
+                                 {0},
+                                 phi::CppTypeToDataType<MT>::Type(),
+                                 false,
+                                 dbias_tmp);
     } else {
-      *dbias_out = phi::Sum<T, Context>(
-          ctx, dout_copy, {0}, phi::CppTypeToDataType<T>::Type(), false);
+      phi::SumKernel<T, Context>(ctx,
+                                 dout_copy,
+                                 {0},
+                                 phi::CppTypeToDataType<T>::Type(),
+                                 false,
+                                 dbias_tmp);
     }
   }
 
   if (dbias) {
     if (kIsMultiPrecision) {
-      phi::AddKernel<MT, Context>(ctx, *dbias_out, dbias.get(), dbias_out);
+      phi::AddKernel<MT, Context>(ctx, dbias.get(), *dbias_tmp, dbias_out);
     } else {
-      phi::AddKernel<T, Context>(ctx, *dbias_out, dbias.get(), dbias_out);
+      phi::AddKernel<T, Context>(ctx, dbias.get(), *dbias_tmp, dbias_out);
     }
   }
 }
@@ -132,6 +154,10 @@ void FusedLinearParamGradAdd(const Context &ctx,
                              DenseTensor *dbias_out) {
   using MT = typename phi::dtype::MPTypeTrait<T>::Type;
 
+  if (std::is_same<T, MT>::value) {
+    multi_precision = false;
+  }
+
   bool use_addto = false;
   if (dweight_out) {
     if (dweight_out->dtype() == phi::DataType::FLOAT16) {
@@ -164,12 +190,27 @@ void FusedLinearParamGradAdd(const Context &ctx,
     }
   }
 
-  if (std::is_same<T, MT>::value) {
-    multi_precision = false;
-  }
-
   if (has_bias && dbias_out) {
-    ctx.template Alloc<T>(dbias_out);
+    if (dbias) {
+      *dbias_out = dbias.get();
+      if (multi_precision) {
+        PADDLE_ENFORCE_EQ(
+            dbias_out->dtype(),
+            phi::CppTypeToDataType<MT>::Type(),
+            phi::errors::InvalidArgument("Invaid data type error."));
+      } else {
+        PADDLE_ENFORCE_EQ(
+            dbias_out->dtype(),
+            phi::CppTypeToDataType<T>::Type(),
+            phi::errors::InvalidArgument("Invaid data type error."));
+      }
+    } else {
+      if (multi_precision) {
+        ctx.template Alloc<MT>(dbias_out);
+      } else {
+        ctx.template Alloc<T>(dbias_out);
+      }
+    }
   }
 
   int64_t K = x.dims()[x.dims().size() - 1];
diff --git a/python/paddle/distributed/fleet/utils/tensor_fusion_helper.py b/python/paddle/distributed/fleet/utils/tensor_fusion_helper.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import itertools
 import weakref
 from collections import OrderedDict
@@ -24,6 +25,8 @@
     core,
 )
 
+from .log_util import logger
+
 
 class HOOK_ACTION:
     ALL_REDUCE = 0
@@ -521,13 +524,9 @@ def add_grad(self, param, use_comm=True):
         if not self._release_grads:
             current_ptr = get_grad_address(param, self.use_main_grad)
             if self._grads_to_addr[param.name] != current_ptr:
-                raise ValueError(
-                    "The address of the grad/main_grad of the param has been changed during training, "
-                    "which is not allowed for dp/sharding overlap with pp. "
-                    "This may be caused by some non-inplace operations on the grad/main_grad. Here are some examples: "
-                    "1. The grad/main_grad of the param is changed by other operations, such as: clear_grad, "
-                    "2. Using non-inplace operations on the grad/main_grad, such as: add, sub, mul, div, etc. "
-                )
+                error_message = f"The address of the grad/main_grad of param {param.name} has been changed during training, which is not allowed for dp/sharding overlap with pp. This may be caused by some non-inplace operations on the grad/main_grad. Here are some examples: 1. The grad/main_grad of the param is changed by other operations, such as: clear_grad; 2. Using non-inplace operations on the grad/main_grad, such as: add, sub, mul, div, etc."
+                logger.error(error_message)
+                raise ValueError(error_message)
         else:
             self._copy_grad_to_buffer(param)
 
diff --git a/test/legacy_test/test_fused_linear_param_grad_add.py b/test/legacy_test/test_fused_linear_param_grad_add.py
@@ -70,7 +70,12 @@ def run_ground_truth(x, dy, dweight, dbias, multi_precision, has_bias):
         dweight += dweight_tmp
 
     if has_bias:
-        dbias_tmp = dy.reshape([-1, dy.shape[-1]]).sum(axis=0)
+        if multi_precision:
+            dbias_tmp = (
+                promote_dtype(dy).reshape([-1, dy.shape[-1]]).sum(axis=0)
+            )
+        else:
+            dbias_tmp = dy.reshape([-1, dy.shape[-1]]).sum(axis=0)
         if dbias is None:
             dbias = dbias_tmp
         else:
@@ -91,6 +96,10 @@ def run_fused_linear_param_grad_add(
     )
     if dweight is not None:
         assert dweight_new.data_ptr() == dweight.data_ptr()
+    if has_bias and dbias is not None:
+        assert (
+            dbias_new.data_ptr() == dbias.data_ptr()
+        ), f"multi_precision={multi_precision}, has_bias={has_bias}, dbias.dtype={dbias.dtype}."
     if has_bias:
         return (
             promote_dtype(dweight_new).numpy(),
@@ -152,7 +161,11 @@ def check_main(self, has_dweight, has_dbias, multi_precision, has_bias):
         self.assertEqual(len(res1), len(res2))
         for r1, r2 in zip(res1, res2):
             max_diff = np.max(np.abs(r1 - r2))
-            self.assertLess(max_diff, 1e-10)
+            self.assertLess(
+                max_diff,
+                1e-10,
+                f"Check failed when: has_dweight={has_dweight}, has_dbias={has_dbias}, multi_precision={multi_precision}, has_bias={has_bias}",
+            )
 
     def test_main(self):
         if not paddle.is_compiled_with_cuda() or paddle.is_compiled_with_rocm():
