[PIR AMP]Adapt PIR AMP dy2st and open some uts (#62606)

0x45f · web-flow · commit ef4906f0a2c3 · 2024-03-14T18:58:52.000+08:00
diff --git a/python/paddle/jit/dy2static/convert_operators.py b/python/paddle/jit/dy2static/convert_operators.py
@@ -109,7 +109,7 @@ def convert_load(x):
             if new_var is not None:
                 return new_var
 
-        if x is paddle.amp.auto_cast:
+        if x is paddle.amp.auto_cast and not use_pir_api():
             return convert_auto_cast
 
     return x
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -20,7 +20,6 @@
 import paddle
 import paddle.pir.core as ir_static
 from paddle import _legacy_C_ops
-from paddle.amp.auto_cast import _in_amp_guard, _in_pure_fp16_guard
 from paddle.autograd.backward_utils import ValueDict
 from paddle.autograd.ir_backward import grad
 from paddle.base import core, framework
@@ -613,29 +612,17 @@ def program_id(self):
         """
         Return current train or eval program hash id.
         """
-        if _in_amp_guard() or _in_pure_fp16_guard():
-            raise NotImplementedError(
-                "Currently, AMP is not supported in PIR mode"
-            )
         if self.training:
             return self._train_program_id
         else:
             return self._infer_program_id
 
     @cached_property
     def train_program(self):
-        if _in_amp_guard() or _in_pure_fp16_guard():
-            raise NotImplementedError(
-                "Currently, AMP is not supported in PIR mode"
-            )
         return self._create_program()
 
     @cached_property
     def infer_program(self):
-        if _in_amp_guard() or _in_pure_fp16_guard():
-            raise NotImplementedError(
-                "Currently, AMP is not supported in PIR mode"
-            )
         return self._create_program(is_infer_mode=True)
 
     def _verify_program(self, main_program):
diff --git a/python/paddle/jit/sot/infer_meta.py b/python/paddle/jit/sot/infer_meta.py
@@ -21,6 +21,7 @@
     UniqueNameGenerator,
     guard as UniqueNameGuard,
 )
+from paddle.framework import use_pir_api
 from paddle.utils import flatten, is_sequence
 
 from .utils import Cache, Singleton, map_if_extend, meta_str
@@ -56,7 +57,8 @@ def from_tensor(tensor):
         # We always use float32 in simulation if AMP is enabled.
         current_amp_state = amp_state()
         if (
-            dtype == paddle.float16
+            not use_pir_api()
+            and dtype == paddle.float16
             and current_amp_state is not None
             and current_amp_state["dtype"] == "float16"
         ):
diff --git a/python/paddle/jit/sot/symbolic/compile_cache.py b/python/paddle/jit/sot/symbolic/compile_cache.py
@@ -116,7 +116,8 @@ def __call__(self, *args, **kwargs):
                 2,
                 lambda: print("[FallbackWrapper] start run SIR: \n", self.SIR),
             )
-            args, kwargs = self.amp_cast_inputs(args, kwargs)
+            if not use_pir_api():
+                args, kwargs = self.amp_cast_inputs(args, kwargs)
             log_do(
                 4,
                 lambda: print(
diff --git a/python/paddle/nn/layer/layers.py b/python/paddle/nn/layer/layers.py
@@ -44,6 +44,7 @@
     in_dygraph_mode,
     in_pir_mode,
     name_struct,
+    paddle_type_to_proto_type,
 )
 from paddle.base.layer_helper_base import LayerHelperBase
 from paddle.base.param_attr import ParamAttr
@@ -2223,13 +2224,18 @@ def _transform(self, t, device, dtype, blocking):
         if dtype is None:
             dtype = t.dtype
 
-        if type(dtype) is not VarDesc.VarType:
+        if not isinstance(dtype, (VarDesc.VarType, core.DataType)):
             dtype = convert_np_dtype_to_dtype_(dtype)
 
         # 1. gpu place need to determine whether the memory is sufficient for allocation:
         if t.place.is_gpu_place():
             # for gpu, minimum memory allocation unit is 256 bytes.
-            size_dtype = core.size_of_dtype(dtype)
+            proto_dtype = (
+                paddle_type_to_proto_type[dtype]
+                if isinstance(dtype, core.DataType)
+                else dtype
+            )
+            size_dtype = core.size_of_dtype(proto_dtype)
             # Note(zhangbo): Paddle GPU minimum memory allocation unit is 256 bytes, waiting_alloc_memory will compute ‘t’ occupied memory space.
             # Coefficient 1.2 is used to avoid OOM that may occur in this critical state when the memory is just enough.
             waiting_alloc_memory = (
diff --git a/test/dygraph_to_static/CMakeLists.txt b/test/dygraph_to_static/CMakeLists.txt
@@ -46,6 +46,7 @@ set_tests_properties(test_basic_api_transformation PROPERTIES TIMEOUT 240)
 set_tests_properties(test_reinforcement_learning PROPERTIES TIMEOUT 120)
 set_tests_properties(test_bmn PROPERTIES TIMEOUT 300)
 set_tests_properties(test_loop PROPERTIES TIMEOUT 180)
+set_tests_properties(test_mnist_amp PROPERTIES TIMEOUT 240)
 
 if(NOT WIN32)
   set_tests_properties(test_tsm PROPERTIES TIMEOUT 900)
diff --git a/test/dygraph_to_static/test_amp_case.py b/test/dygraph_to_static/test_amp_case.py
@@ -17,7 +17,9 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
-    test_legacy_and_pt,
+    test_ast_only,
+    test_legacy_and_pt_and_pir,
+    test_pir_only,
 )
 
 import paddle
@@ -39,10 +41,40 @@ def _run_static(self):
             st_out = static_func(x)
         np.testing.assert_allclose(dy_out.numpy(), st_out.numpy())
 
-    @test_legacy_and_pt
+    @test_legacy_and_pt_and_pir
     def test_ast_to_func(self):
         self._run_static()
 
 
+class Net(paddle.nn.Layer):
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = paddle.nn.Linear(5, 5)
+
+    def forward(self, x):
+        out = self.linear(x)
+        with paddle.amp.auto_cast(level='O2'):
+            out = self.linear(out)
+        return out
+
+
+class TestPartialAutoCast(Dy2StTestBase):
+    @test_ast_only
+    @test_pir_only
+    def test_run(self):
+        if not paddle.base.core.is_compiled_with_cuda():
+            return
+        x = paddle.randn([5, 5], 'float32')
+        net = Net()
+        net = paddle.jit.to_static(net)
+        out = net(x)
+        main = net.forward.main_program
+        cast_op_count = 0
+        for op in main.global_block().ops:
+            if op.name() == 'pd_op.cast':
+                cast_op_count += 1
+        np.testing.assert_equal(cast_op_count, 3)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/dygraph_to_static/test_mnist_amp.py b/test/dygraph_to_static/test_mnist_amp.py
@@ -16,6 +16,9 @@
 from time import time
 
 import numpy as np
+from dygraph_to_static_utils import (
+    test_legacy_and_pt_and_pir,
+)
 from test_mnist import MNIST, SEED, TestMNIST
 
 import paddle
@@ -32,6 +35,7 @@ def train_static(self):
     def train_dygraph(self):
         return self.train(to_static=False)
 
+    @test_legacy_and_pt_and_pir
     def test_mnist_to_static(self):
         dygraph_loss = self.train_dygraph()
         static_loss = self.train_static()
diff --git a/test/dygraph_to_static/test_mnist_pure_fp16.py b/test/dygraph_to_static/test_mnist_pure_fp16.py
@@ -16,6 +16,7 @@
 from time import time
 
 import numpy as np
+from dygraph_to_static_utils import test_legacy_and_pt_and_pir
 from test_mnist import MNIST, SEED, TestMNIST
 
 import paddle
@@ -31,6 +32,7 @@ def train_static(self):
     def train_dygraph(self):
         return self.train(to_static=False)
 
+    @test_legacy_and_pt_and_pir
     def test_mnist_to_static(self):
         if paddle.base.is_compiled_with_cuda():
             dygraph_loss = self.train_dygraph()
diff --git a/test/dygraph_to_static/test_resnet_amp.py b/test/dygraph_to_static/test_resnet_amp.py
@@ -19,6 +19,7 @@
 from dygraph_to_static_utils import (
     Dy2StTestBase,
     enable_to_static_guard,
+    test_default_and_pir,
     test_default_mode_only,
 )
 from test_resnet import SEED, ResNet, optimizer_setting
@@ -116,7 +117,7 @@ def train(self, to_static: bool):
         with enable_to_static_guard(to_static):
             return train()
 
-    @test_default_mode_only
+    @test_default_and_pir
     def test_resnet(self):
         static_loss = self.train(to_static=True)
         dygraph_loss = self.train(to_static=False)
diff --git a/test/dygraph_to_static/test_resnet_pure_fp16.py b/test/dygraph_to_static/test_resnet_pure_fp16.py
@@ -19,6 +19,7 @@
 from dygraph_to_static_utils import (
     Dy2StTestBase,
     enable_to_static_guard,
+    test_default_and_pir,
     test_default_mode_only,
 )
 from test_resnet import SEED, ResNet, optimizer_setting
@@ -123,7 +124,7 @@ def train(self, to_static: bool):
             build_strategy.enable_inplace = False
             return train(build_strategy)
 
-    @test_default_mode_only
+    @test_default_and_pir
     def test_resnet(self):
         if base.is_compiled_with_cuda():
             static_loss = self.train(to_static=True)
