fix fp16 and bf16 dtype (#68120)

changeyoung98 · web-flow · commit f02e8453dfac · 2024-09-11T15:25:20.000+08:00
diff --git a/python/paddle/nn/initializer/kaiming.py b/python/paddle/nn/initializer/kaiming.py
@@ -116,8 +116,9 @@ def forward(
             self._seed = block.program.random_seed
 
         # to be compatible of fp16 initializers
-        if var.dtype == core.VarDesc.VarType.FP16 or (
-            var.dtype == core.VarDesc.VarType.BF16 and not self._uniform
+        origin_dtype = var.dtype
+        if origin_dtype == core.VarDesc.VarType.FP16 or (
+            origin_dtype == core.VarDesc.VarType.BF16 and not self._uniform
         ):
             out_dtype = core.VarDesc.VarType.FP32
             out_var = block.create_var(
@@ -130,13 +131,13 @@ def forward(
                 persistable=False,
             )
         elif (
-            var.dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
+            origin_dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
             and not self._uniform
         ):
             out_dtype = core.DataType.FLOAT32
             out_var = var
         else:
-            out_dtype = var.dtype
+            out_dtype = origin_dtype
             out_var = var
 
         if in_dygraph_mode():
@@ -159,10 +160,16 @@ def forward(
                     out_var.shape, 0.0, std, self._seed, out_dtype, place
                 )
 
-            if var.dtype == core.VarDesc.VarType.FP16 or (
-                var.dtype == core.VarDesc.VarType.BF16 and not self._uniform
+            if origin_dtype == core.VarDesc.VarType.FP16 or (
+                origin_dtype
+                in [
+                    core.VarDesc.VarType.BF16,
+                    core.DataType.FLOAT16,
+                    core.DataType.BFLOAT16,
+                ]
+                and not self._uniform
             ):
-                var_tmp = _C_ops.cast(out_var, var.dtype)
+                var_tmp = _C_ops.cast(out_var, origin_dtype)
                 var_tmp._share_underline_tensor_to(var)
             else:
                 out_var._share_underline_tensor_to(var)
@@ -188,10 +195,10 @@ def forward(
                 )
 
             if (
-                var.dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
+                origin_dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
                 and not self._uniform
             ):
-                return _C_ops.cast(out_var, var.dtype)
+                return _C_ops.cast(out_var, origin_dtype)
 
             return out_var
         else:
@@ -228,14 +235,17 @@ def forward(
                     stop_gradient=True,
                 )
 
-            if var.dtype == core.VarDesc.VarType.FP16 or (
-                var.dtype == core.VarDesc.VarType.BF16 and not self._uniform
+            if origin_dtype == core.VarDesc.VarType.FP16 or (
+                origin_dtype == core.VarDesc.VarType.BF16 and not self._uniform
             ):
                 block.append_op(
                     type="cast",
                     inputs={"X": out_var},
                     outputs={"Out": var},
-                    attrs={"in_dtype": out_var.dtype, "out_dtype": var.dtype},
+                    attrs={
+                        "in_dtype": out_var.dtype,
+                        "out_dtype": origin_dtype,
+                    },
                 )
 
             var.op = op
diff --git a/python/paddle/nn/initializer/xavier.py b/python/paddle/nn/initializer/xavier.py
@@ -124,8 +124,9 @@ def forward(
             else var.shape
         )
         # to be compatible of fp16 initializers
-        if var.dtype == core.VarDesc.VarType.FP16 or (
-            var.dtype == core.VarDesc.VarType.BF16 and not self._uniform
+        origin_dtype = var.dtype
+        if origin_dtype == core.VarDesc.VarType.FP16 or (
+            origin_dtype == core.VarDesc.VarType.BF16 and not self._uniform
         ):
             out_dtype = core.VarDesc.VarType.FP32
             out_var = block.create_var(
@@ -138,13 +139,13 @@ def forward(
                 persistable=False,
             )
         elif (
-            var.dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
+            origin_dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
             and not self._uniform
         ):
             out_dtype = core.DataType.FLOAT32
             out_var = var
         else:
-            out_dtype = var.dtype
+            out_dtype = origin_dtype
             out_var = var
 
         if in_dygraph_mode():
@@ -171,10 +172,16 @@ def forward(
                     place,
                 )
 
-            if var.dtype == core.VarDesc.VarType.FP16 or (
-                var.dtype == core.VarDesc.VarType.BF16 and not self._uniform
+            if origin_dtype == core.VarDesc.VarType.FP16 or (
+                origin_dtype
+                in [
+                    core.VarDesc.VarType.BF16,
+                    core.DataType.FLOAT16,
+                    core.DataType.BFLOAT16,
+                ]
+                and not self._uniform
             ):
-                out_var = _C_ops.cast(out_var, var.dtype)
+                out_var = _C_ops.cast(out_var, origin_dtype)
             if isinstance(var, framework.EagerParamBase) and var.is_dist():
                 # lazy init for dist tensor
                 out_var = (
@@ -207,10 +214,10 @@ def forward(
                 )
 
             if (
-                var.dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
+                origin_dtype in (core.DataType.FLOAT16, core.DataType.BFLOAT16)
                 and not self._uniform
             ):
-                return _C_ops.cast(out_var, var.dtype)
+                return _C_ops.cast(out_var, origin_dtype)
 
             return out_var
         else:
@@ -244,14 +251,17 @@ def forward(
                     stop_gradient=True,
                 )
 
-            if var.dtype == core.VarDesc.VarType.FP16 or (
-                var.dtype == core.VarDesc.VarType.BF16 and not self._uniform
+            if origin_dtype == core.VarDesc.VarType.FP16 or (
+                origin_dtype == core.VarDesc.VarType.BF16 and not self._uniform
             ):
                 block.append_op(
                     type="cast",
                     inputs={"X": out_var},
                     outputs={"Out": var},
-                    attrs={"in_dtype": out_var.dtype, "out_dtype": var.dtype},
+                    attrs={
+                        "in_dtype": out_var.dtype,
+                        "out_dtype": origin_dtype,
+                    },
                 )
 
             var.op = op
