【pir_save_load】modify Paddlemix export model dtype is float32 bug (#66891)

xiaoguoguo626807 · web-flow · commit 8d0c726f2b2e · 2024-08-02T09:25:17.000+08:00
* add all

* add notes

* add jit.layer part

* split pir.h

* recover

* add pir_utils

* modify link bug

* modify link bug

* modify build bug

* modify assign initializer dtype bug

* modify assign
diff --git a/python/paddle/nn/initializer/assign.py b/python/paddle/nn/initializer/assign.py
@@ -77,7 +77,11 @@ def forward(
         assert isinstance(block, (framework.Block, paddle.pir.Block))
 
         # to be compatible of fp16 initializers
-        if var.dtype in [core.VarDesc.VarType.FP16, core.VarDesc.VarType.BF16]:
+        origin_dtype = var.dtype
+        if origin_dtype in [
+            core.VarDesc.VarType.FP16,
+            core.VarDesc.VarType.BF16,
+        ]:
             out_dtype = core.VarDesc.VarType.FP32
             np_value = self._value.astype("float32")
             out_var = block.create_var(
@@ -89,13 +93,13 @@ def forward(
                 type=core.VarDesc.VarType.LOD_TENSOR,
                 persistable=False,
             )
-        elif var.dtype in [core.DataType.FLOAT16, core.DataType.BFLOAT16]:
+        elif origin_dtype in [core.DataType.FLOAT16, core.DataType.BFLOAT16]:
             out_var = var
             out_dtype = core.DataType.FLOAT32
             np_value = self._value.astype("float32")
         else:
             out_var = var
-            out_dtype = var.dtype
+            out_dtype = origin_dtype
             np_value = self._value
 
         if out_dtype in (core.VarDesc.VarType.FP32, core.DataType.FLOAT32):
@@ -131,11 +135,13 @@ def forward(
                 values,
                 _current_expected_place(),
             )
-            if var.dtype in [
+            if origin_dtype in [
                 core.VarDesc.VarType.FP16,
                 core.VarDesc.VarType.BF16,
+                core.DataType.FLOAT16,
+                core.DataType.BFLOAT16,
             ]:
-                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)
@@ -147,8 +153,8 @@ def forward(
                 values,
                 _current_expected_place(),
             )
-            if var.dtype in [core.DataType.FLOAT16, core.DataType.BFLOAT16]:
-                out_var = _C_ops.cast(out_var, var.dtype)
+            if origin_dtype in [core.DataType.FLOAT16, core.DataType.BFLOAT16]:
+                out_var = _C_ops.cast(out_var, origin_dtype)
             return out_var
         else:
             op = block.append_op(
@@ -162,15 +168,18 @@ def forward(
                 stop_gradient=True,
             )
 
-            if var.dtype in [
+            if origin_dtype in [
                 core.VarDesc.VarType.FP16,
                 core.VarDesc.VarType.BF16,
             ]:
                 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
