[CustomDevice] npu support bf16 (#62604)

tianhaodongbd · web-flow · commit 82740d4653f5 · 2024-03-14T15:16:33.000+08:00
diff --git a/paddle/phi/kernels/shape_kernel.cc b/paddle/phi/kernels/shape_kernel.cc
@@ -105,7 +105,8 @@ PD_REGISTER_KERNEL(shape,
                    double,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
   kernel->OutputAt(0).SetDataType(phi::DataType::INT32);
diff --git a/python/paddle/amp/auto_cast.py b/python/paddle/amp/auto_cast.py
@@ -18,6 +18,7 @@
 import paddle
 from paddle.base import core
 from paddle.base.framework import (
+    _current_expected_place,
     _dygraph_tracer,
     dygraph_only,
     in_dynamic_or_pir_mode,
@@ -143,6 +144,14 @@ def _is_gpu_bfloat16_supported():
     return prop[0] >= 8 and cuda_version_check
 
 
+def _is_custom_device_bfloat16_supported():
+    """
+    Judge whether current custom device support bfloat16 amp.
+    """
+    place = _current_expected_place()
+    return place.get_device_type() == 'npu'
+
+
 def need_keep_fp32(layer, dtype):
     need_keep_fp32 = False
     # Highest priority. Because all the layers except BN will use bfloat16 params in bfloat16 training,
@@ -471,7 +480,7 @@ def amp_guard(
                 "current_tracer is None, maybe it is not in imperative mode."
             )
         # check device_type:
-        # NOTE: Now, amp only support gpu for float16 and bfloat16, xpu for float16, npu for float16.
+        # NOTE: Now, amp only support gpu for float16 and bfloat16, xpu for float16, npu for float16 and bfloat16.
         # Maybe we will support cpu for bfloat16.
         if enable and not (
             tracer._expected_place.is_gpu_place()
@@ -489,8 +498,10 @@ def amp_guard(
                 warnings.warn('XPUPlace only support float16 amp.')
                 enable = False
             # For custom device:
-            if tracer._expected_place.is_custom_place() and (
-                dtype == 'bfloat16'
+            if (
+                tracer._expected_place.is_custom_place()
+                and not _is_custom_device_bfloat16_supported()
+                and (dtype == 'bfloat16')
             ):
                 warnings.warn('CustomPlace only support float16 amp.')
                 enable = False
@@ -753,7 +764,11 @@ def amp_decorate(
         else:
             return models, optimizers
     # For custom device:
-    if tracer._expected_place.is_custom_place() and (dtype == 'bfloat16'):
+    if (
+        tracer._expected_place.is_custom_place()
+        and not _is_custom_device_bfloat16_supported()
+        and (dtype == 'bfloat16')
+    ):
         if optimizers is None:
             return models
         else:
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
@@ -222,8 +222,12 @@ def _comm_and_clip(
         )
         clip_var_fp16 = paddle.cast(clip_var, paddle.float16)
 
-        if not isinstance(
-            paddle.framework._current_expected_place(), paddle.CustomPlace
+        if (
+            not isinstance(
+                paddle.framework._current_expected_place(), paddle.CustomPlace
+            )
+            or paddle.framework._current_expected_place().get_device_type()
+            == 'npu'
         ):
             clip_var_bf16 = paddle.cast(clip_var, paddle.bfloat16)
         for p, g in params_grads:
