[XPU] hybrid_parallel_optimizer xpu support bf16 (#60213)

dynamicheart · web-flow · commit d155697b59e0 · 2023-12-21T18:31:29.000+08:00
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
@@ -225,12 +225,8 @@ def _comm_and_clip(
         )
         clip_var_fp16 = paddle.cast(clip_var, paddle.float16)
 
-        # bf16 is not supported on XPU now
-        if not (
-            paddle.is_compiled_with_xpu()
-            or isinstance(
-                paddle.framework._current_expected_place(), paddle.CustomPlace
-            )
+        if not isinstance(
+            paddle.framework._current_expected_place(), paddle.CustomPlace
         ):
             clip_var_bf16 = paddle.cast(clip_var, paddle.bfloat16)
         for p, g in params_grads:
@@ -241,10 +237,6 @@ def _comm_and_clip(
             if g.dtype == paddle.float16:
                 g.multiply_(clip_var_fp16)
             elif g.dtype == paddle.bfloat16:
-                if paddle.is_compiled_with_xpu():
-                    raise NotImplementedError(
-                        "BF16 is not supported on XPU now"
-                    )
                 g.multiply_(clip_var_bf16)
             else:
                 g.multiply_(clip_var)
