add ut

Author: LittleHeroZZZX

python/paddle/amp/auto_cast.py

@@ -92,12 +92,6 @@ def clear_grad(self, set_to_zero: bool) -> None: ...
 }
 
 AMP_LEVEL = core.AmpLevel
-_g_amp_state_ = None
-
-
-def amp_state():
-    global _g_amp_state_
-    return _g_amp_state_
 
 
 class AMPGlobalState:
@@ -520,11 +514,6 @@ def amp_guard(
         in_dynamic_or_pir_mode()
     ), "We only support 'amp_guard' in dynamic or pir mode."
 
-    amp_state = locals()
-    global _g_amp_state_
-    original_state = _g_amp_state_
-    _g_amp_state_ = amp_state
-
     # check amp_level: O0-O2
     level = level.upper()
     if level not in ['O0', 'OD', 'O1', 'O2']:
@@ -576,10 +565,10 @@ def amp_guard(
         try:
             yield
         finally:
-            _g_amp_state_ = original_state
             amp_attrs._amp_level = original_amp_level
             core._set_amp_op_list(original_white_list, original_black_list)
             amp_attrs._amp_dtype = original_amp_dtype
+            amp_global_state().amp_dtype = original_amp_dtype
             if amp_level == AMP_LEVEL.O2:
                 amp_attrs._use_promote = original_use_promote
 
@@ -780,13 +769,13 @@ def param_hook(tmp_grad):
             yield
         finally:
             if tracer:
-                _g_amp_state_ = original_state
                 tracer._amp_level = original_amp_level
                 tracer._set_amp_op_list(
                     original_white_list, original_black_list
                 )
                 # set_flags(original_flags)
                 tracer._amp_dtype = original_amp_dtype
+                amp_global_state().amp_dtype = original_amp_dtype
                 if amp_level == AMP_LEVEL.O2:
                     tracer._use_promote = original_use_promote
 

test/amp/test_get_autocast_dtype.py

@@ -20,7 +20,9 @@
 @unittest.skipIf(paddle.device.get_device() == "cpu", "Skip AMP test on CPU")
 class TestAutocast(unittest.TestCase):
     def setUp(self) -> None:
+        paddle.disable_static()
         self.device_list = [None, paddle.device.get_device()]
+        self.default_dtype = "float32"
 
     def do_test(self, device, expected_type):
         self.assertTrue(paddle.get_autocast_dtype(device) == expected_type)
@@ -30,25 +32,25 @@ def do_test(self, device, expected_type):
 
     def test_amp_default(self):
         for device in self.device_list:
-            self.do_test(device, "float16")
+            self.do_test(device, self.default_dtype)
 
     def test_amp_autocast_fp16(self):
         for device in self.device_list:
             with paddle.amp.auto_cast(True, dtype="float16"):
                 self.do_test(device, "float16")
-            self.do_test(device, "float16")
+            self.do_test(device, self.default_dtype)
 
     def test_amp_autocast_bf16(self):
         for device in self.device_list:
             with paddle.amp.auto_cast(True, dtype="bfloat16"):
                 self.do_test(device, "bfloat16")
-            self.do_test(device, "float16")
+            self.do_test(device, self.default_dtype)
 
     def test_amp_autocast_false_bf16(self):
         for device in self.device_list:
             with paddle.amp.auto_cast(True, dtype="bfloat16"):
                 self.do_test(device, "bfloat16")
-            self.do_test(device, "float16")
+            self.do_test(device, self.default_dtype)
 
     def test_amp_nested_context(self):
         for device in self.device_list:
@@ -57,7 +59,14 @@ def test_amp_nested_context(self):
                 with paddle.amp.auto_cast(False, dtype="float16"):
                     self.do_test(device, "float16")
                 self.do_test(device, "bfloat16")
-            self.do_test(device, "float16")
+            self.do_test(device, self.default_dtype)
+
+
+class TestAutocastStatic(TestAutocast):
+    def setUp(self) -> None:
+        paddle.enable_static()
+        self.device_list = [None, paddle.device.get_device()]
+        self.default_dtype = "float32"
 
 
 if __name__ == "__main__":

test/amp/test_is_autocast_enabled.py

@@ -20,6 +20,7 @@
 @unittest.skipIf(paddle.device.get_device() == "cpu", "Skip AMP test on CPU")
 class TestAutocast(unittest.TestCase):
     def setUp(self) -> None:
+        paddle.disable_static()
         self.device_list = [None, paddle.device.get_device()]
 
     def test_amp_default(self):
@@ -61,5 +62,11 @@ def test_amp_nested_context(self):
             self.assertFalse(paddle.amp.is_autocast_enabled(device))
 
 
+class TestAutocastStatic(TestAutocast):
+    def setUp(self) -> None:
+        paddle.enable_static()
+        self.device_list = [None, paddle.device.get_device()]
+
+
 if __name__ == "__main__":
     unittest.main()