[Cleanup][A-10] clean some VarType for test (#61563)

co63oc · web-flow · commit 0b7b20d801c4 · 2024-02-19T12:25:11.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -926,14 +926,14 @@ def _initialize(self, mode, init_parameters=True):
                         if scope_var and buffer_tensor._is_initialized():
                             continue
                         # for amp
-                        if dest_type == core.VarDesc.VarType.BF16:
+                        if dest_type == paddle.bfloat16:
                             buffer_tensor.set(
                                 _convert_float_to_bfloat16(
                                     self._place, buffer.numpy()
                                 ),
                                 self._place,
                             )
-                        elif dest_type == core.VarDesc.VarType.FP16:
+                        elif dest_type == paddle.float16:
                             buffer_tensor.set(
                                 np.float16(buffer.numpy()), self._place
                             )
diff --git a/python/paddle/distributed/fleet/dataset/dataset.py b/python/paddle/distributed/fleet/dataset/dataset.py
@@ -15,6 +15,7 @@
 
 from google.protobuf import text_format
 
+import paddle
 from paddle.base import core
 from paddle.base.proto import data_feed_pb2
 
@@ -182,9 +183,9 @@ def _set_use_var(self, var_list):
             if var.lod_level == 0:
                 slot_var.is_dense = True
                 slot_var.shape.extend(var.shape)
-            if var.dtype == core.VarDesc.VarType.FP32:
+            if var.dtype == paddle.float32:
                 slot_var.type = "float"
-            elif var.dtype == core.VarDesc.VarType.INT64:
+            elif var.dtype == paddle.int64:
                 slot_var.type = "uint64"
             else:
                 raise ValueError(
@@ -318,9 +319,7 @@ def _check_use_var_with_data_generator(
                                 % ele[0]
                             )
 
-                        if var_list[
-                            i
-                        ].dtype == core.VarDesc.VarType.FP32 and not all(
+                        if var_list[i].dtype == paddle.float32 and not all(
                             isinstance(ele, float) for ele in ele[1]
                         ):
                             raise TypeError(
@@ -332,8 +331,8 @@ def _check_use_var_with_data_generator(
                             )
 
                         if (
-                            var_list[i].dtype == core.VarDesc.VarType.INT64
-                            or var_list[i].dtype == core.VarDesc.VarType.INT32
+                            var_list[i].dtype == paddle.int64
+                            or var_list[i].dtype == paddle.int32
                         ) and not all(isinstance(ele, int) for ele in ele[1]):
                             raise TypeError(
                                 "var dtype mismatch error: var name = {}, var type in var_list = {}, while var in data_generator contains non-int value, which is {} \n"
diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding/offload_helper.py b/python/paddle/distributed/fleet/meta_optimizers/sharding/offload_helper.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import paddle
 from paddle.framework import core
 from paddle.utils import unique_name
 
@@ -59,7 +60,7 @@ def _insert_cast_op(self, block, idx, src_name, dst_name):
                 persistable=True,
             )
         dst_var = block.var(dst_name)
-        assert dst_var.dtype == core.VarDesc.VarType.FP16
+        assert dst_var.dtype == paddle.float16
         block._insert_op_without_sync(
             idx,
             type='cast',
diff --git a/python/paddle/distributed/fleet/utils/mix_precision_utils.py b/python/paddle/distributed/fleet/utils/mix_precision_utils.py
@@ -206,12 +206,12 @@ def unscale_method(self, optimizer):
         for group in optimizer._param_groups:
             for param in group['params']:
                 if param.main_grad is not None:
-                    assert param.main_grad.dtype == core.VarDesc.VarType.FP32
+                    assert param.main_grad.dtype == paddle.float32
                     param_grads.append(param.main_grad)
     else:
         for param in optimizer._parameter_list:
             if param.main_grad is not None:
-                assert param.main_grad.dtype == core.VarDesc.VarType.FP32
+                assert param.main_grad.dtype == paddle.float32
                 param_grads.append(param.main_grad)
 
     temp_found_inf = paddle.to_tensor(np.array([0]).astype(np.bool_))
diff --git a/python/paddle/distributed/passes/auto_parallel_amp.py b/python/paddle/distributed/passes/auto_parallel_amp.py
@@ -59,9 +59,9 @@
 
 
 def _dtype_to_str(dtype):
-    if dtype == core.VarDesc.VarType.FP16:
+    if dtype == paddle.float16:
         return 'fp16'
-    elif dtype == core.VarDesc.VarType.BF16:
+    elif dtype == paddle.bfloat16:
         return 'bf16'
     else:
         return 'fp32'
@@ -282,7 +282,7 @@ def _cast_block(self, block):
                     # deal with op with attribute 'dtype', such as 'fill_constant'
                     if (
                         op.has_attr('dtype')
-                        and op.attr('dtype') == core.VarDesc.VarType.FP32
+                        and op.attr('dtype') == paddle.float32
                     ):
                         op._set_attr('dtype', _str_to_dtype(self.amp_dtype))
                     num_cast_ops = self._insert_cast_op_forward(
@@ -321,7 +321,7 @@ def _cast_block(self, block):
                         # deal with op with attribute 'dtype', such as 'fill_constant'
                         if (
                             op.has_attr('dtype')
-                            and op.attr('dtype') == core.VarDesc.VarType.FP32
+                            and op.attr('dtype') == paddle.float32
                         ):
                             op._set_attr('dtype', _str_to_dtype(self.amp_dtype))
                         num_cast_ops = self._insert_cast_op_backward(
@@ -372,7 +372,7 @@ def _insert_cast_op_forward(
 
         for in_name in op.input_names:
             if (
-                src_dtype == core.VarDesc.VarType.FP32
+                src_dtype == paddle.float32
                 and self.amp_lists._op_keep_fp32_input(op, in_name)
             ):
                 continue
@@ -457,9 +457,8 @@ def _insert_cast_op_forward(
                         op._set_attr('in_dtype', dst_dtype)
         self._var_name_dict[op.desc.original_id()] = var_name_dict
 
-        if (
-            src_dtype == core.VarDesc.VarType.FP32
-            and dst_dtype == _str_to_dtype(self.amp_dtype)
+        if src_dtype == paddle.float32 and dst_dtype == _str_to_dtype(
+            self.amp_dtype
         ):
             for out_name in op.output_names:
                 if self.amp_lists._op_keep_fp32_output(op, out_name):
@@ -468,7 +467,7 @@ def _insert_cast_op_forward(
                     out_var = block._var_recursive(out_var_name)
                     if out_var.type not in _valid_types:
                         continue
-                    if out_var.dtype == core.VarDesc.VarType.FP32:
+                    if out_var.dtype == paddle.float32:
                         out_var.desc.set_dtype(_str_to_dtype(self.amp_dtype))
                         if op.has_attr('out_dtype'):
                             op._set_attr(
@@ -522,12 +521,10 @@ def _keep_fp32_output(op, out_name):
             return num_cast_ops
 
         for in_name in op.input_names:
-            if src_dtype == core.VarDesc.VarType.FP32 and _keep_fp32_input(
-                op, in_name
-            ):
+            if src_dtype == paddle.float32 and _keep_fp32_input(op, in_name):
                 for in_var_name in op.input(in_name):
                     in_var = block._var_recursive(in_var_name)
-                    assert in_var.dtype == core.VarDesc.VarType.FP32
+                    assert in_var.dtype == paddle.float32
                 continue
 
             for in_var_name in op.input(in_name):
@@ -559,12 +556,10 @@ def _keep_fp32_output(op, out_name):
                         )
 
         for out_name in op.output_names:
-            if src_dtype == core.VarDesc.VarType.FP32 and _keep_fp32_output(
-                op, out_name
-            ):
+            if src_dtype == paddle.float32 and _keep_fp32_output(op, out_name):
                 for out_var_name in op.output(out_name):
                     out_var = block._var_recursive(out_var_name)
-                    assert out_var.dtype == core.VarDesc.VarType.FP32
+                    assert out_var.dtype == paddle.float32
                 continue
 
             for out_var_name in op.output(out_name):
@@ -652,10 +647,7 @@ def _keep_fp32_output(op, out_name):
                 else:
                     assert out_var.dtype == dst_dtype
 
-        if (
-            op.has_attr('dtype')
-            and op.attr('dtype') == core.VarDesc.VarType.FP32
-        ):
+        if op.has_attr('dtype') and op.attr('dtype') == paddle.float32:
             op._set_attr('dtype', _str_to_dtype(self.amp_dtype))
 
         return num_cast_ops
@@ -754,7 +746,7 @@ def _update_backward_cast_ops(self):
 
         for p, g in self.params_grads:
             op = g.op
-            if g.dtype == core.VarDesc.VarType.FP32 and op.type == 'cast':
+            if g.dtype == paddle.float32 and op.type == 'cast':
                 if int(op.attr('op_role')) == int(
                     OpRole.Backward
                 ) and op.has_attr('op_role_var'):
@@ -1171,9 +1163,9 @@ def _update_loss_scaling(self, grads, found_inf):
             check_variable_and_dtype(
                 e, "x", ['float16', 'float32', 'float64'], 'update_loss_scaling'
             )
-            if e.dtype == core.VarDesc.VarType.FP16:
+            if e.dtype == paddle.float16:
                 assert (
-                    self._loss_scaling.dtype == core.VarDesc.VarType.FP32
+                    self._loss_scaling.dtype == paddle.float32
                 ), "The dtype of prev_loss_scaling should be float32 when the dtype of x is float16."
             else:
                 assert (
