chore: import upstream snapshot with attribution

test/amp/CMakeLists.txt

@@ -0,0 +1,59 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+if(NOT WITH_GPU AND NOT WITH_XPU)
+  list(REMOVE_ITEM TEST_OPS test_grad_scaler_out)
+endif()
+
+function(py_test_modules TARGET_NAME)
+  if(WITH_TESTING)
+    set(options SERIAL)
+    set(oneValueArgs "")
+    set(multiValueArgs MODULES DEPS ENVS)
+    cmake_parse_arguments(py_test_modules "${options}" "${oneValueArgs}"
+                          "${multiValueArgs}" ${ARGN})
+
+    string(REGEX MATCH "_deprecated\.py$" DEPRECATED_MODULES
+                 "${py_test_modules_MODULES}")
+    string(REGEX MATCH "_deprecated$" DEPRECATED_TARGET_NAME "${TARGET_NAME}")
+    set(FLAGS_PIR_MODE "")
+    if((NOT "${DEPRECATED_MODULES}" STREQUAL "")
+       OR (NOT "${DEPRECATED_TARGET_NAME}" STREQUAL ""))
+      set(FLAGS_PIR_MODE FLAGS_enable_pir_api=0)
+    endif()
+    if(WITH_COVERAGE AND NOT (WITH_INCREMENTAL_COVERAGE
+                              AND "$ENV{PADDLE_GIT_DIFF_PY_FILE}" STREQUAL ""))
+      add_test(
+        NAME ${TARGET_NAME}
+        COMMAND
+          ${CMAKE_COMMAND} -E env PYTHONPATH=${PADDLE_BINARY_DIR}/python
+          ${py_test_modules_ENVS} ${FLAGS_PIR_MODE}
+          COVERAGE_FILE=${PADDLE_BINARY_DIR}/python-coverage.data
+          ${PYTHON_EXECUTABLE} -m coverage run --branch -p
+          ${PADDLE_SOURCE_DIR}/tools/test_runner.py ${py_test_modules_MODULES}
+        WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+    else()
+      add_test(
+        NAME ${TARGET_NAME}
+        COMMAND
+          ${CMAKE_COMMAND} -E env PYTHONPATH=${PADDLE_BINARY_DIR}/python
+          ${py_test_modules_ENVS} ${FLAGS_PIR_MODE} ${PYTHON_EXECUTABLE}
+          ${PADDLE_SOURCE_DIR}/tools/test_runner.py ${py_test_modules_MODULES}
+        WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+    endif()
+
+    if(py_test_modules_SERIAL)
+      set_property(TEST ${TARGET_NAME} PROPERTY RUN_SERIAL 1)
+    endif()
+    if(WIN32)
+      set_tests_properties(${TARGET_NAME} PROPERTIES TIMEOUT 150)
+    endif()
+  endif()
+endfunction()
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP})
+endforeach()

test/amp/amp_base_models.py

@@ -0,0 +1,453 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import struct
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle import nn
+from paddle.base import core
+from paddle.framework import in_dynamic_or_pir_mode
+
+
+def copy_bits_from_float_to_uint16(f):
+    return struct.unpack('<I', struct.pack('<f', f))[0] >> 16
+
+
+def convert_float_to_uint16(in_list):
+    if in_list.dtype == np.float32:
+        new_output = []
+        for x in np.nditer(in_list):
+            new_output.append(np.uint16(copy_bits_from_float_to_uint16(x)))
+        new_output = np.reshape(new_output, in_list.shape).view(np.uint16)
+        return new_output
+    else:
+        return in_list
+
+
+def convert_uint16_to_float(in_list):
+    if in_list.dtype == np.uint16:
+        in_list = np.asarray(in_list)
+        out = np.vectorize(
+            lambda x: struct.unpack('<f', struct.pack('<I', x << 16))[0],
+            otypes=[np.float32],
+        )(in_list.flat)
+        return np.reshape(out, in_list.shape)
+    else:
+        return in_list
+
+
+_fixed_add_param = np.random.random(size=[16, 16]).astype("float32")
+
+
+def _build_optimizer(
+    use_amp,
+    amp_dtype="float16",
+    amp_level="O1",
+    amp_lists=None,
+    use_grad_clip=False,
+    use_promote=False,
+    use_master_grad=False,
+    model=None,
+):
+    if use_grad_clip:
+        grad_clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=1.0)
+    else:
+        grad_clip = None
+    if in_dynamic_or_pir_mode():
+        assert model is not None
+        parameters = model.parameters()
+    else:
+        parameters = None
+    optimizer = paddle.optimizer.AdamW(
+        learning_rate=0.01,
+        parameters=parameters,
+        grad_clip=grad_clip,
+        beta1=0.78,
+        beta2=0.836,
+        epsilon=1e-4,
+        weight_decay=0.01,
+    )
+    if not in_dynamic_or_pir_mode() and use_amp:
+        optimizer = paddle.static.amp.decorate(
+            optimizer,
+            amp_lists,
+            level=amp_level,
+            dtype=amp_dtype,
+            master_grad=use_master_grad,
+            use_promote=use_promote,
+        )
+    return optimizer
+
+
+class SimpleAddNet(nn.Layer):
+    def __init__(self, dtype):
+        super().__init__()
+        global _fixed_add_param
+        self.weight = paddle.create_parameter(
+            name="add_w",
+            shape=[16, 16],
+            dtype=dtype,
+            attr=paddle.ParamAttr(
+                initializer=paddle.nn.initializer.Assign(_fixed_add_param)
+            ),
+        )
+
+    def forward(self, x):
+        return x + self.weight
+
+
+def cast_add_param(amp_dtype):
+    global _fixed_add_param
+    if amp_dtype == "bfloat16":
+        _fixed_add_param_bf16 = convert_float_to_uint16(_fixed_add_param)
+        _fixed_add_param = convert_uint16_to_float(_fixed_add_param_bf16)
+    else:
+        pass
+
+
+def build_add_model(
+    use_amp, amp_dtype="float16", amp_level="O1", use_promote=False
+):
+    main_program = paddle.static.Program()
+    startup_program = paddle.static.Program()
+    with (
+        paddle.utils.unique_name.guard(),
+        paddle.static.program_guard(main_program, startup_program),
+    ):
+        x_dtype = "float32"
+        if use_amp and amp_level == "O2":
+            if amp_dtype == "bfloat16":
+                x_dtype = "uint16"
+            elif amp_dtype == "float16":
+                x_dtype = "float16"
+        cast_add_param(amp_dtype)
+        model = SimpleAddNet(x_dtype)
+        x = paddle.static.data(name='input', shape=[16, 16], dtype=x_dtype)
+        out = model(x)
+        loss = paddle.mean(out)
+
+        if use_amp:
+            amp_lists = paddle.static.amp.AutoMixedPrecisionLists(
+                custom_white_list=["elementwise_add"],
+                custom_black_list=["reduce_mean"],
+                dtype=amp_dtype,
+            )
+        else:
+            amp_lists = None
+        optimizer = _build_optimizer(
+            use_amp,
+            amp_dtype,
+            amp_level,
+            amp_lists,
+            use_promote=use_promote,
+        )
+        optimizer.minimize(loss)
+    feed_vars = [x]
+    fetch_vars = [loss]
+    return main_program, startup_program, optimizer, feed_vars, fetch_vars
+
+
+class SimpleConvNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2D(in_channels=1, out_channels=6, kernel_size=3)
+        self.linear = nn.Linear(in_features=96, out_features=4)
+
+    def forward(self, x):
+        out = self.conv(x)
+        out = nn.functional.relu(out.cast("float32"))
+        out = out.flatten(start_axis=1, stop_axis=3)
+        out = self.linear(out)
+        out = nn.functional.softmax(out)
+        return out
+
+
+def build_conv_model(
+    use_amp, amp_dtype="float16", amp_level="O1", use_promote=False
+):
+    if in_dynamic_or_pir_mode():
+        model = SimpleConvNet()
+        optimizer = _build_optimizer(use_amp=False, model=model)
+        if use_amp and amp_dtype == "float16":
+            scaler = paddle.amp.GradScaler(init_loss_scaling=32768.0)
+        else:
+            scaler = None
+        if use_amp and amp_level == "O2":
+            model, optimizer = paddle.amp.decorate(
+                models=model,
+                optimizers=optimizer,
+                level=amp_level,
+                dtype=amp_dtype,
+            )
+        return model, optimizer, scaler
+
+    main_program = paddle.static.Program()
+    startup_program = paddle.static.Program()
+    with (
+        paddle.utils.unique_name.guard(),
+        paddle.static.program_guard(main_program, startup_program),
+    ):
+        model = SimpleConvNet()
+        x = paddle.static.data(
+            name='input', shape=[None, 1, 6, 6], dtype='float32'
+        )
+        out = model(x)
+        loss = paddle.mean(out)
+        optimizer = _build_optimizer(
+            use_amp, amp_dtype, amp_level, use_promote=use_promote
+        )
+        optimizer.minimize(loss)
+    feed_vars = [x]
+    fetch_vars = [loss]
+    return main_program, startup_program, optimizer, feed_vars, fetch_vars
+
+
+class SimpleEmbeddingNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.vocab_size = 128
+        self.hidden_size = 16
+        self.embedding = nn.Embedding(self.vocab_size, self.hidden_size)
+        self.linear = nn.Linear(in_features=16, out_features=10)
+
+    def forward(self, x):
+        out = self.embedding(x)
+        scale = paddle.full(shape=[1], fill_value=2, dtype="int64")
+        out = out * (scale.astype("float32"))
+        out = self.linear(out)
+        out = nn.functional.dropout(out, p=0.2)
+        return out
+
+
+def build_embedding_model(
+    use_amp,
+    amp_dtype="float16",
+    amp_level="O1",
+    use_promote=False,
+    use_master_grad=False,
+):
+    main_program = paddle.static.Program()
+    startup_program = paddle.static.Program()
+    with (
+        paddle.utils.unique_name.guard(),
+        paddle.static.program_guard(main_program, startup_program),
+    ):
+        model = SimpleEmbeddingNet()
+        x = paddle.static.data(name='x', shape=[None, 32], dtype='int64')
+        out = model(x)
+        loss = paddle.mean(out)
+        if use_amp:
+            amp_lists = paddle.static.amp.AutoMixedPrecisionLists(
+                custom_white_list=["elementwise_mul"],
+                custom_black_list=["reduce_mean"],
+                dtype=amp_dtype,
+            )
+        else:
+            amp_lists = None
+        optimizer = _build_optimizer(
+            use_amp,
+            amp_dtype,
+            amp_level,
+            amp_lists,
+            True,
+            use_promote=use_promote,
+            use_master_grad=use_master_grad,
+        )
+        optimizer.minimize(loss)
+
+    feed_vars = [x]
+    fetch_vars = [loss]
+    return main_program, startup_program, optimizer, feed_vars, fetch_vars
+
+
+class SimpleMLPNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.linear0 = paddle.nn.Linear(16, 10)
+        self.linear1 = paddle.nn.Linear(10, 32)
+
+    def forward(self, x):
+        out = self.linear0(x)
+        out = nn.functional.relu(out)
+        out = self.linear1(out)
+        out = nn.functional.relu(out)
+        out = nn.functional.dropout(out, p=0.2)
+        return out
+
+
+def build_MLP_model(
+    use_amp,
+    use_grad_clip=False,
+    amp_dtype="float16",
+    amp_level="O1",
+    use_promote=False,
+    use_master_grad=False,
+):
+    main_program = paddle.static.Program()
+    startup_program = paddle.static.Program()
+    with (
+        paddle.utils.unique_name.guard(),
+        paddle.static.program_guard(main_program, startup_program),
+    ):
+        model = SimpleMLPNet()
+        x_dtype = "float32"
+        if use_amp and amp_level == "O2":
+            if amp_dtype == "bfloat16":
+                x_dtype = "uint16"
+            elif amp_dtype == "float16":
+                x_dtype = "float16"
+        x = paddle.static.data(name='x', shape=[None, 16], dtype=x_dtype)
+        out = model(x)
+        loss = paddle.mean(out)
+
+        if use_amp:
+            amp_lists = paddle.static.amp.AutoMixedPrecisionLists(
+                custom_black_list=["reduce_mean"],
+                dtype=amp_dtype,
+            )
+        else:
+            amp_lists = None
+
+        optimizer = _build_optimizer(
+            use_amp,
+            amp_dtype,
+            amp_level,
+            amp_lists,
+            use_grad_clip=use_grad_clip,
+            use_promote=use_promote,
+            use_master_grad=use_master_grad,
+        )
+        optimizer.minimize(loss)
+
+    feed_vars = [x]
+    fetch_vars = [loss]
+    return main_program, startup_program, optimizer, feed_vars, fetch_vars
+
+
+class SimpleWhileNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.linear = paddle.nn.Linear(16, 10)
+
+    def forward(self, x):
+        def cond(i, loop_len, x, result):
+            return i < loop_len
+
+        def body(i, loop_len, x, result):
+            result = self.linear(x)
+            paddle.increment(i)
+            return [i, loop_len, x, result]
+
+        i = paddle.zeros(shape=[1], dtype='int64')
+        loop_len = paddle.ones(shape=[1], dtype='int64')
+        result = paddle.zeros(
+            shape=x.shape[:-1] + self.linear.weight.shape[-1:], dtype="float32"
+        )
+        result.stop_gradient = False
+        _, _, _, results = paddle.static.nn.while_loop(
+            cond, body, [i, loop_len, x, result]
+        )
+        return results
+
+
+def build_while_model():
+    main_program = paddle.static.Program()
+    startup_program = paddle.static.Program()
+    with (
+        paddle.utils.unique_name.guard(),
+        paddle.static.program_guard(main_program, startup_program),
+    ):
+        model = SimpleWhileNet()
+        x = paddle.static.data(name='x', shape=[32, 16], dtype='float32')
+        out = model(x)
+        loss = paddle.mean(out)
+    return main_program, startup_program
+
+
+@unittest.skipIf(
+    not (core.is_compiled_with_cuda() or core.is_compiled_with_xpu()),
+    "core is not compiled with CUDA or XPU and not support amp.",
+)
+class AmpTestBase(unittest.TestCase):
+    def setUp(self):
+        self.amp_dtype = None
+        self.amp_level = None
+
+    def _check_op_calls(
+        self,
+        op_stats_dict,
+        expected_bf16_calls={},
+        expected_fp16_calls={},
+        debug_info=None,
+    ):
+        def _extract_op_call(op_calls_str, pos):
+            return int(op_calls_str.split(",")[pos])
+
+        for op_type, expected_value in expected_bf16_calls.items():
+            # print(f"[BF16] op_type={op_type}, value={value}")
+            if isinstance(op_stats_dict[op_type], str):
+                actual_value = _extract_op_call(op_stats_dict[op_type], 1)
+            else:
+                actual_value = op_stats_dict[op_type].bf16_calls
+            self.assertEqual(
+                actual_value,
+                expected_value,
+                f"[{debug_info}] The number of bf16 calls of operator < {op_type} > is expected to be {expected_value}, but received {actual_value}.",
+            )
+        for op_type, expected_value in expected_fp16_calls.items():
+            # print(f"[FP16] op_type={op_type}, value={value}")
+            if isinstance(op_stats_dict[op_type], str):
+                actual_value = _extract_op_call(op_stats_dict[op_type], 0)
+            else:
+                actual_value = op_stats_dict[op_type].fp16_calls
+            self.assertEqual(
+                actual_value,
+                expected_value,
+                f"[{debug_info}] The number of fp16 calls of operator < {op_type} > is expected to be {expected_value}, but received {actual_value}.",
+            )
+
+    def run_program(
+        self,
+        main_program,
+        startup_program,
+        optimizer,
+        feed_vars,
+        fetch_vars,
+        place,
+        exe,
+        x_np,
+        max_iters,
+        dtype,
+        level,
+    ):
+        losses = []
+        scope = paddle.static.Scope()
+        with paddle.static.scope_guard(scope):
+            exe.run(startup_program)
+            if level == 'O2':
+                optimizer.amp_init(place)
+            for iter_id in range(max_iters):
+                results = exe.run(
+                    program=main_program,
+                    feed={feed_vars[0].name: x_np},
+                    fetch_list=fetch_vars,
+                )
+                print(
+                    f"-- [AMP {dtype} {level}] iter={iter_id}, loss={results[0]}"
+                )
+                losses.append(results[0])
+        return losses

test/amp/test_amp_api.py

@@ -0,0 +1,460 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from contextlib import contextmanager
+
+import numpy as np
+from amp_base_models import AmpTestBase
+
+import paddle
+from paddle import nn
+from paddle.base import core
+from paddle.static import amp
+
+paddle.set_flags({"FLAGS_use_legacy_linear": True})
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
+    "Bugs on XPU3, disable temporarily",
+)
+class TestAutoCast(AmpTestBase):
+    def init_net(self):
+        self._conv = paddle.nn.Conv2D(
+            in_channels=1, out_channels=6, kernel_size=3, bias_attr=False
+        )
+        self._linear = paddle.nn.Linear(in_features=4, out_features=4)
+
+    def test_amp_OD_level(self):
+        self.init_net()
+        with paddle.amp.auto_cast(level='OD'):
+            out1 = self._conv(paddle.rand(shape=[1, 1, 6, 6], dtype='float32'))
+            out2 = out1 + paddle.rand(shape=out1.shape, dtype='float16')
+            out3 = self._linear(out2)
+
+        self.assertEqual(out1.dtype, paddle.float16)
+        self.assertEqual(out2.dtype, paddle.float32)
+        self.assertEqual(out3.dtype, paddle.float32)
+
+    def test_pir_amp_OD_level(self):
+        with (
+            paddle.pir_utils.IrGuard(),
+            paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ),
+        ):
+            self.init_net()
+            with paddle.amp.auto_cast(level='OD'):
+                out1 = self._conv(
+                    paddle.rand(shape=[1, 1, 6, 6], dtype='float32')
+                )
+                out2 = out1 + paddle.rand(shape=out1.shape, dtype='float16')
+                out3 = self._linear(out2)
+
+            self.assertEqual(out1.dtype, core.DataType.FLOAT16)
+            self.assertEqual(out2.dtype, core.DataType.FLOAT32)
+            self.assertEqual(out3.dtype, core.DataType.FLOAT32)
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
+    "Bugs on XPU3, disable temporarily",
+)
+class TestCudaAutoCast(unittest.TestCase):
+    def setUp(self):
+        self._conv = paddle.nn.Conv2D(1, 1, 3, bias_attr=False)
+        self._linear = paddle.nn.Linear(4, 4)
+
+    def _run_autocast_test(self, ctx):
+        with ctx:
+            out1 = self._conv(paddle.rand(shape=[1, 1, 6, 6], dtype='float32'))
+            out2 = out1 + paddle.rand(shape=out1.shape, dtype='float16')
+            out3 = self._linear(out2)
+
+        self.assertEqual(out1.dtype, paddle.float16)
+        self.assertEqual(out2.dtype, paddle.float16)
+        self.assertEqual(out3.dtype, paddle.float32)
+
+    def test_amp_autocast(self):
+        self._run_autocast_test(paddle.amp.autocast(device_type='cuda'))
+
+    def test_amp_autocast2(self):
+        self._run_autocast_test(
+            paddle.amp.autocast(
+                device_type='cuda',
+                enabled=True,
+                dtype=paddle.float16,
+                cache_enabled=True,
+            )
+        )
+
+    def test_autocast(self):
+        self._run_autocast_test(
+            paddle.autocast(
+                device_type='cuda',
+                enabled=True,
+                dtype=paddle.float16,
+                cache_enabled=True,
+            )
+        )
+
+    def test_cuda_amp_autocast(self):
+        self._run_autocast_test(paddle.cuda.amp.autocast())
+
+    def test_device_amp_autocast(self):
+        self._run_autocast_test(paddle.device.amp.autocast())
+
+    def test_cuda_amp_autocast_mode_autocast(self):
+        self._run_autocast_test(paddle.cuda.amp.autocast_mode.autocast())
+
+
+class SimpleConvNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self._conv = paddle.nn.Conv2D(
+            in_channels=1, out_channels=6, kernel_size=3, bias_attr=False
+        )
+        self._linear = paddle.nn.Linear(in_features=4, out_features=4)
+
+    def forward(self, x):
+        out1 = self._conv(paddle.rand(shape=[1, 1, 6, 6], dtype='float32'))
+        out2 = out1 + paddle.rand(shape=out1.shape, dtype='float16')
+        out3 = self._linear(out2)
+        return out3
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
+    "Bugs on XPU3, disable temporarily",
+)
+class TestStaticDecorate(AmpTestBase):
+    def check_results(
+        self, use_amp, dtype, level, use_promote, expected_op_calls
+    ):
+        main_program = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with (
+            paddle.utils.unique_name.guard(),
+            paddle.static.program_guard(main_program, startup_program),
+        ):
+            model = SimpleConvNet()
+            x = paddle.static.data(
+                name='input', shape=[None, 1, 6, 6], dtype='float32'
+            )
+            out = model(x)
+            loss = paddle.mean(out)
+            optimizer = paddle.optimizer.Adadelta(learning_rate=0.001)
+            optimizer = paddle.static.amp.decorate(
+                optimizer,
+                init_loss_scaling=128.0,
+                use_dynamic_loss_scaling=True,
+                level=level,
+            )
+            optimizer.minimize(loss)
+
+        feed_vars = [x]
+        fetch_vars = [loss]
+        self.assertEqual(main_program.num_blocks, 1)
+
+        amp.debugging.collect_operator_stats(main_program)
+        op_stats_list = amp.debugging._get_op_stats_list(main_program)
+
+        self._check_op_calls(
+            op_stats_list[0], expected_fp16_calls=expected_op_calls
+        )
+
+        if paddle.is_compiled_with_cuda():
+            place = paddle.CUDAPlace(0)
+        elif paddle.device.is_compiled_with_xpu():
+            place = paddle.device.XPUPlace(0)
+        else:
+            raise ValueError("Only support CUDA or XPU Place.")
+        exe = paddle.static.Executor(place)
+
+        max_iters = 2
+        x_fp32 = np.random.random(size=[1, 1, 6, 6]).astype("float32")
+        losses_o1 = self.run_program(
+            main_program,
+            startup_program,
+            optimizer,
+            feed_vars,
+            fetch_vars,
+            place,
+            exe,
+            x_fp32,
+            max_iters,
+            dtype,
+            level,
+        )
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
+    "Bugs on XPU3, disable temporarily",
+)
+class TestGradScaler(AmpTestBase):
+    def test_amp_grad_scaler(self):
+        model = paddle.nn.Conv2D(3, 2, 3)
+        optimizer = paddle.optimizer.SGD(
+            learning_rate=0.01, parameters=model.parameters()
+        )
+        scaler = paddle.amp.GradScaler()
+        data = paddle.rand([1, 3, 8, 8], dtype='float32')
+        paddle.amp.debugging.enable_operator_stats_collection()
+        with paddle.amp.auto_cast(
+            custom_black_list=['conv2d'], dtype='bfloat16'
+        ):
+            out = model(data)
+            loss = out.mean()
+        scaled = scaler.scale(loss)
+        scaled.backward()
+        scaler.minimize(optimizer, scaled)
+        optimizer.clear_grad()
+        paddle.amp.debugging.disable_operator_stats_collection()
+        op_list = paddle.base.core.get_low_precision_op_list()
+
+        self.assertEqual(scaler._enable, False)
+        self.assertEqual(scaler._use_dynamic_loss_scaling, False)
+        self.assertTrue('scale' not in op_list)
+        self.assertTrue('check_finite_and_unscale' not in op_list)
+
+    def test_pir_amp_grad_scaler(self):
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                model = paddle.nn.Conv2D(3, 2, 3)
+                optimizer = paddle.optimizer.SGD(
+                    learning_rate=0.01, parameters=model.parameters()
+                )
+                model, optimizer = paddle.amp.decorate(
+                    models=model,
+                    optimizers=optimizer,
+                )
+                scaler = paddle.amp.GradScaler()
+                data = paddle.static.data('data', [1, 3, 8, 8], dtype='float32')
+
+                with paddle.amp.auto_cast(
+                    custom_black_list=['conv2d'], dtype='bfloat16'
+                ):
+                    out = model(data)
+                    loss = out.mean()
+                scaled = scaler.scale(loss)
+                scaler.minimize(optimizer, scaled)
+
+                if paddle.is_compiled_with_cuda():
+                    place = paddle.CUDAPlace(0)
+                elif paddle.device.is_compiled_with_xpu():
+                    place = paddle.device.XPUPlace(0)
+                else:
+                    raise ValueError("Only support CUDA or XPU Place.")
+                exe = paddle.static.Executor(place)
+                exe.run(startup)
+                paddle.amp.debugging.enable_operator_stats_collection()
+                exe.run(
+                    main,
+                    feed={'data': np.random.rand(1, 3, 8, 8).astype('float32')},
+                    fetch_list=[loss],
+                )
+                paddle.amp.debugging.disable_operator_stats_collection()
+                op_list = paddle.base.core.get_low_precision_op_list()
+
+                self.assertEqual(scaler._enable, False)
+                self.assertEqual(scaler._use_dynamic_loss_scaling, False)
+                self.assertTrue('pd_op.scale' not in op_list)
+                self.assertTrue(
+                    'pd_op.check_finite_and_unscale_' not in op_list
+                )
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
+    "Bugs on XPU3, disable temporarily",
+)
+class SimpleModelIncludeSetValue(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.norm = nn.LayerNorm(3)
+
+    def forward(self, x):
+        x = x + 1
+        tmp = x * 1
+        y = self.norm(tmp)
+        x[:] = y
+
+        z = x * 1
+        return z
+
+
+# Copy from ../dygraph_to_static/dygraph_to_static_utils.py
+@contextmanager
+def pir_dygraph_guard():
+    in_dygraph_mode = paddle.in_dynamic_mode()
+    with paddle.pir_utils.IrGuard():
+        if in_dygraph_mode:
+            paddle.disable_static()
+        yield
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
+    "Bugs on XPU3, disable temporarily",
+)
+class TestDy2STWithSetValue(AmpTestBase):
+    def test_op_called_as_expected(self):
+        if paddle.framework.use_pir_api():
+            return
+        expected_fp16_calls = {
+            "cast": 1,
+            "layer_norm": 1,
+            "scale": 3,
+            "set_value": 1,
+        }
+
+        func = SimpleModelIncludeSetValue()
+        func = paddle.amp.decorate(func, level='O2')
+        func = paddle.jit.to_static(func, full_graph=True, backend=None)
+        input = paddle.randn((2, 3))
+
+        with paddle.amp.auto_cast(level='O2', use_promote=False):
+            res = func(input)
+            loss = res.sum()
+            prog = func.forward.get_concrete_program(input)[1].forward_program
+            amp.debugging.collect_operator_stats(prog)
+            op_stats_list = amp.debugging._get_op_stats_list(prog)
+        loss.backward()
+        self._check_op_calls(
+            op_stats_list[0], expected_fp16_calls=expected_fp16_calls
+        )
+
+    def test_pir_op_called_as_expected(self):
+        expected_fp16_calls = {
+            "pd_op.layer_norm": 1,
+            "pd_op.scale": 1,
+            "pd_op.scale_": 2,
+            "pd_op.set_value_with_tensor_": 1,
+        }
+
+        with pir_dygraph_guard():
+            func = SimpleModelIncludeSetValue()
+            func = paddle.amp.decorate(func, level='O2')
+            func = paddle.jit.to_static(func, full_graph=True, backend=None)
+            input = paddle.randn((2, 3))
+
+            paddle.amp.debugging.enable_operator_stats_collection()
+            with paddle.amp.auto_cast(level='O2', use_promote=False):
+                res = func(input)
+                loss = res.sum()
+                paddle.amp.debugging.disable_operator_stats_collection()
+                op_stats = paddle.base.core.get_low_precision_op_list()
+
+            loss.backward()
+            self._check_op_calls(
+                op_stats, expected_fp16_calls=expected_fp16_calls
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/amp/test_amp_decorate.py

@@ -0,0 +1,248 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+import paddle.nn.functional as F
+from paddle.base import core
+
+
+class ConvBNLayer(paddle.nn.Layer):
+    def __init__(
+        self,
+        num_channels,
+        num_filters,
+        filter_size,
+        stride=1,
+        groups=1,
+        act=None,
+    ):
+        super().__init__()
+
+        self._conv = paddle.nn.Conv2D(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=(filter_size - 1) // 2,
+            groups=groups,
+            bias_attr=None,
+        )
+
+        self._batch_norm = paddle.nn.BatchNorm(num_filters, act=act)
+
+    def forward(self, inputs):
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+
+        return y
+
+
+class Model(paddle.nn.Layer):
+    def __init__(
+        self, input_channel, hidden_size, fp16_conv=True, fp16_linear=True
+    ):
+        super().__init__()
+        self.conv = ConvBNLayer(input_channel, 8, 3)
+        self.linear = paddle.nn.Linear(8, hidden_size)
+        self.layernorm = paddle.nn.Sequential(
+            paddle.nn.LayerNorm(hidden_size),
+            paddle.nn.LayerNorm(hidden_size),
+        )
+        self.fp16_conv = fp16_conv
+        self.fp16_linear = fp16_linear
+
+    def forward(self, inputs):
+        with paddle.amp.auto_cast(enable=self.fp16_conv):
+            if not self.fp16_conv:
+                inputs = inputs.astype('float32')
+            x = self.conv(inputs)
+        with paddle.amp.auto_cast(enable=self.fp16_linear):
+            if not self.fp16_linear:
+                x = x.astype('float32')
+            x = self.linear(x)
+        x = F.relu(x)
+        x = self.layernorm(x)
+        return x
+
+
+class LayerNorm2D(paddle.nn.LayerNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x):
+        x = x.transpose([0, 2, 3, 1])
+        x = super().forward(x)
+        return x.transpose([0, 3, 1, 2])
+
+
+class CustomLayer(paddle.nn.Layer):
+    def __init__(
+        self, input_channel, hidden_size, fp16_conv=True, fp16_linear=True
+    ):
+        super().__init__()
+        self.conv = ConvBNLayer(input_channel, 8, 3)
+        self.linear = paddle.nn.Linear(8, hidden_size)
+        self.layernorm = paddle.nn.Sequential(
+            LayerNorm2D(hidden_size),
+            LayerNorm2D(hidden_size),
+        )
+        self.fp16_conv = fp16_conv
+        self.fp16_linear = fp16_linear
+
+    def forward(self, inputs):
+        with paddle.amp.auto_cast(enable=self.fp16_conv):
+            if not self.fp16_conv:
+                inputs = inputs.astype('float32')
+            x = self.conv(inputs)
+        with paddle.amp.auto_cast(enable=self.fp16_linear):
+            if not self.fp16_linear:
+                x = x.astype('float32')
+            x = self.linear(x)
+        x = F.relu(x)
+        x = self.layernorm(x)
+        return x
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+class TestAMPDecorate(unittest.TestCase):
+    def check_results(self, fp32_layers=[], fp16_layers=[]):
+        for idx in range(len(fp32_layers)):
+            for layer in fp32_layers[idx].sublayers(include_self=False):
+                self.assertTrue(
+                    layer.weight.dtype
+                    in (paddle.float32, core.DataType.FLOAT32)
+                )
+                self.assertTrue(
+                    layer.bias.dtype in (paddle.float32, core.DataType.FLOAT32)
+                )
+
+        for idx in range(len(fp16_layers)):
+            for layer in fp16_layers[idx].sublayers(include_self=False):
+                self.assertTrue(
+                    layer.weight.dtype
+                    in (paddle.float16, core.DataType.FLOAT16)
+                )
+                self.assertTrue(
+                    layer.bias.dtype in (paddle.float16, core.DataType.FLOAT16)
+                )
+
+    def test_excluded_layers(self):
+        model = Model(4, 8, fp16_conv=False)
+        model = paddle.amp.decorate(
+            models=model,
+            level='O2',
+            dtype='float16',
+            excluded_layers=model.conv,
+        )
+        with paddle.amp.auto_cast(level='O2'):
+            out = model(paddle.rand(shape=[2, 4, 8, 8], dtype='float32'))
+        self.check_results(
+            fp32_layers=[model.conv, model.layernorm],
+            fp16_layers=[model.linear],
+        )
+
+    def test_excluded_layers_attr_list(self):
+        model = Model(4, 8, fp16_conv=False, fp16_linear=False)
+        model = paddle.amp.decorate(
+            models=model,
+            level='O2',
+            dtype='float16',
+            excluded_layers=[model.conv, model.linear],
+        )
+
+        with paddle.amp.auto_cast(level='O2'):
+            out = model(paddle.rand(shape=[2, 4, 8, 8], dtype='float32'))
+
+        self.check_results(
+            fp32_layers=[model.conv, model.linear, model.layernorm]
+        )
+
+    def test_excluded_layers_attr_types(self):
+        model = Model(4, 8)
+        model = paddle.amp.decorate(
+            models=model,
+            level='O2',
+            dtype='float16',
+            excluded_layers=[paddle.nn.Conv2D, model.linear],
+        )
+
+        with paddle.amp.auto_cast(level='O2'):
+            out = model(paddle.rand(shape=[2, 4, 8, 8], dtype='float16'))
+
+        self.check_results(
+            fp32_layers=[model.conv, model.linear, model.layernorm]
+        )
+
+    def test_excluded_layers_attr_none(self):
+        model = Model(4, 8)
+        model = paddle.amp.decorate(
+            models=model,
+            level='O2',
+            dtype='float16',
+            excluded_layers=None,
+        )
+
+        with paddle.amp.auto_cast(level='O2'):
+            out = model(paddle.rand(shape=[2, 4, 8, 8], dtype='float16'))
+
+        self.check_results(
+            fp32_layers=[model.layernorm, model.conv._batch_norm],
+            fp16_layers=[model.conv._conv, model.linear],
+        )
+
+    def test_excluded_layers_custom_layer(self):
+        model = CustomLayer(4, 8)
+        model = paddle.amp.decorate(
+            models=model,
+            level='O2',
+            dtype='bfloat16',
+            excluded_layers=[paddle.nn.LayerNorm, paddle.nn.BatchNorm],
+        )
+        with paddle.amp.auto_cast(level='O2'):
+            out = model(paddle.rand(shape=[2, 4, 8, 8], dtype='float32'))
+        self.check_results(
+            fp32_layers=[model.layernorm, model.conv._batch_norm],
+        )
+
+    def test_pir(self):
+        with (
+            paddle.pir_utils.IrGuard(),
+            paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ),
+        ):
+            self.test_excluded_layers()
+            self.test_excluded_layers_attr_list()
+            self.test_excluded_layers_attr_types()
+            self.test_excluded_layers_attr_none()
+            self.test_excluded_layers_custom_layer()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/amp/test_amp_list.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.base import core
+from paddle.static.amp import AutoMixedPrecisionLists, fp16_lists
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+class TestAMPList(unittest.TestCase):
+    def setUp(self):
+        self.default_black_list = [
+            'linear_interp_v2',
+            'nearest_interp_v2',
+            'bilinear_interp_v2',
+            'bicubic_interp_v2',
+            'trilinear_interp_v2',
+        ]
+        self.custom_white_list = [
+            'lookup_table',
+            'lookup_table_v2',
+        ]
+
+    def check_if_op_in_list(self, op_list, amp_list):
+        for op in op_list:
+            self.assertTrue(op in amp_list)
+
+    def check_if_op_not_in_list(self, op_list, amp_list):
+        for op in op_list:
+            self.assertTrue(op not in amp_list)
+
+    def test_static(self):
+        amp_list = AutoMixedPrecisionLists(
+            custom_white_list=self.custom_white_list
+        )
+        self.check_if_op_in_list(self.default_black_list, amp_list.black_list)
+        self.check_if_op_in_list(self.custom_white_list, amp_list.white_list)
+        self.check_if_op_not_in_list(
+            self.custom_white_list, amp_list.black_list
+        )
+        if paddle.amp.is_float16_supported():
+            self.check_if_op_not_in_list(
+                self.custom_white_list, amp_list.black_list
+            )
+
+    def test_eager(self):
+        if not paddle.amp.is_float16_supported():
+            return
+        white_list = paddle.amp.white_list()
+        black_list = paddle.amp.black_list()
+        self.check_if_op_in_list(
+            self.default_black_list, black_list["float16"]["O2"]
+        )
+        self.check_if_op_not_in_list(['log', 'elementwise_add'], white_list)
+        with paddle.amp.auto_cast(custom_white_list={'elementwise_add'}):
+            out1 = paddle.rand([2, 3]) + paddle.rand([2, 3])
+            out2 = out1.mean()
+            out3 = paddle.log(out2)
+        self.check_if_op_not_in_list(['log', 'elementwise_add'], white_list)
+        self.assertEqual(out1.dtype, paddle.float16)
+        self.assertEqual(out2.dtype, paddle.float32)
+        self.assertEqual(out3.dtype, paddle.float32)
+
+    def test_pir(self):
+        with (
+            paddle.pir_utils.IrGuard(),
+            paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ),
+        ):
+            white_list = paddle.amp.white_list()
+            black_list = paddle.amp.black_list()
+            self.check_if_op_in_list(
+                self.default_black_list, black_list["float16"]["O2"]
+            )
+            self.check_if_op_not_in_list(['log', 'elementwise_add'], white_list)
+            with paddle.amp.auto_cast(custom_white_list={'elementwise_add'}):
+                out1 = paddle.rand([2, 3]) + paddle.rand([2, 3])
+                out2 = out1.mean()
+                out3 = paddle.log(out2)
+            self.check_if_op_not_in_list(['log', 'elementwise_add'], white_list)
+            self.assertEqual(out1.dtype, core.DataType.FLOAT16)
+            self.assertEqual(out2.dtype, core.DataType.FLOAT32)
+            self.assertEqual(out3.dtype, core.DataType.FLOAT32)
+
+    def test_apis(self):
+        def _run_check_dtype():
+            fp16_lists.check_amp_dtype(dtype="int64")
+
+        self.assertRaises(ValueError, _run_check_dtype)
+
+        for vartype in [core.VarDesc.VarType.FP16, core.VarDesc.VarType.BF16]:
+            self.assertEqual(
+                fp16_lists.get_low_precision_vartype(vartype), vartype
+            )
+        self.assertEqual(
+            fp16_lists.get_low_precision_vartype("float16"),
+            core.VarDesc.VarType.FP16,
+        )
+        self.assertEqual(
+            fp16_lists.get_low_precision_vartype("bfloat16"),
+            core.VarDesc.VarType.BF16,
+        )
+
+        def _run_get_vartype():
+            fp16_lists.get_low_precision_vartype(dtype="int64")
+
+        self.assertRaises(ValueError, _run_get_vartype)
+
+        for dtype in ["float16", "bfloat16"]:
+            self.assertEqual(
+                fp16_lists.get_low_precision_dtypestr(dtype), dtype
+            )
+        if paddle.framework.use_pir_api():
+            self.assertEqual(
+                fp16_lists.get_low_precision_dtypestr(core.DataType.FLOAT16),
+                "float16",
+            )
+            self.assertEqual(
+                fp16_lists.get_low_precision_dtypestr(core.DataType.BFLOAT16),
+                "bfloat16",
+            )
+        else:
+            self.assertEqual(
+                fp16_lists.get_low_precision_dtypestr(
+                    core.VarDesc.VarType.FP16
+                ),
+                "float16",
+            )
+            self.assertEqual(
+                fp16_lists.get_low_precision_dtypestr(
+                    core.VarDesc.VarType.BF16
+                ),
+                "bfloat16",
+            )
+
+        def _run_get_dtypestr():
+            fp16_lists.get_low_precision_dtypestr(dtype="int64")
+
+        self.assertRaises(ValueError, _run_get_dtypestr)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/amp/test_amp_master_grad.py

@@ -0,0 +1,222 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.base import core
+
+paddle.set_flags({"FLAGS_use_legacy_linear": True})
+
+
+class SimpleNet(paddle.nn.Layer):
+    def __init__(self, input_size, output_size):
+        super().__init__()
+        self.linear = paddle.nn.Linear(input_size, output_size)
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and not core.is_float16_supported(core.CUDAPlace(0)),
+    "core is not compiled with CUDA and not support the float16",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
+    "Bugs on XPU3, disable temporarily",
+)
+class TestMasterGrad(unittest.TestCase):
+    def check_results(
+        self, fp32_grads, op_list, total_steps, accumulate_batches_num
+    ):
+        for grad in fp32_grads:
+            self.assertEqual(grad.dtype, paddle.float32)
+        # fp16 calls
+        self.assertEqual(int(op_list['matmul_v2'].split(',')[0]), total_steps)
+        self.assertEqual(
+            int(op_list['adam_'].split(',')[0]),
+            2 * (total_steps / accumulate_batches_num),
+        )
+        # Since two additional casts are called when constructing master grad,
+        # the number of operators of this type +2
+        self.assertEqual(
+            int(op_list['cast'].split(',')[0]),
+            total_steps * 2 + 2,
+        )
+
+    def run_dygraph(
+        self, total_steps, accumulate_batches_num, model, optimizer
+    ):
+        model, opt = paddle.amp.decorate(
+            model, optimizers=optimizer, level='O2', master_grad=True
+        )
+        scaler = paddle.amp.GradScaler()
+        paddle.amp.debugging.enable_operator_stats_collection()
+        for i in range(total_steps):
+            x = np.random.random((2, 2)).astype('float32')
+            label = np.random.random((2, 4)).astype('float32')
+
+            with paddle.amp.auto_cast(level='O2'):
+                out = model(paddle.to_tensor(x))
+                loss = paddle.nn.functional.l1_loss(
+                    out, paddle.to_tensor(label)
+                )
+            scaled = scaler.scale(loss)
+            scaled.backward()
+            fp32_grads = [model.linear.weight.grad, model.linear.bias.grad]
+            if (i + 1) % accumulate_batches_num == 0:
+                scaler.step(opt)
+                scaler.update()
+                opt.clear_grad()
+        paddle.amp.debugging.disable_operator_stats_collection()
+        op_list = paddle.base.core.get_low_precision_op_list()
+        return fp32_grads, op_list
+
+    def test_adam_master_grad(self):
+        total_steps = 4
+        accumulate_batches_num = 2
+        model = SimpleNet(2, 4)
+        opt = paddle.optimizer.Adam(parameters=model.parameters())
+        fp32_grads, op_list = self.run_dygraph(
+            total_steps, accumulate_batches_num, model, opt
+        )
+        self.check_results(
+            fp32_grads, op_list, total_steps, accumulate_batches_num
+        )
+
+    def test_momentum_master_grad(self):
+        total_steps = 4
+        accumulate_batches_num = 1
+        model = SimpleNet(2, 4)
+        L1Decay = paddle.regularizer.L1Decay(0.0001)
+        opt = paddle.optimizer.Momentum(
+            parameters=model.parameters(), weight_decay=L1Decay
+        )
+        fp32_grads, op_list = self.run_dygraph(
+            total_steps, accumulate_batches_num, model, opt
+        )
+        for grad in fp32_grads:
+            self.assertEqual(grad.dtype, paddle.float32)
+
+    def run_pir(self, total_steps, accumulate_batches_num, model, optimizer):
+        model, opt = paddle.amp.decorate(
+            model, optimizers=optimizer, level='O2', master_grad=True
+        )
+        scaler = paddle.amp.GradScaler()
+        x = paddle.static.data('x', (2, 2), 'float32')
+        label = paddle.static.data('label', (2, 4), 'float32')
+        with paddle.amp.auto_cast(level='O2'):
+            out = model(paddle.to_tensor(x))
+            loss = paddle.nn.functional.l1_loss(out, paddle.to_tensor(label))
+        scaled = scaler.scale(loss)
+        scaler.minimize(opt, scaled)
+
+        fp32_grads = list(opt._optimizer._master_grads.values())
+        if paddle.is_compiled_with_cuda():
+            place = paddle.CUDAPlace(0)
+        elif paddle.device.is_compiled_with_xpu():
+            place = paddle.device.XPUPlace(0)
+        else:
+            raise ValueError("Only support CUDA or XPU Place.")
+        exe = paddle.static.Executor(place)
+        exe.run(paddle.static.default_startup_program())
+        paddle.amp.debugging.enable_operator_stats_collection()
+        for i in range(total_steps):
+            exe.run(
+                paddle.static.default_main_program(),
+                feed={
+                    'x': np.random.random((2, 2)).astype('float32'),
+                    'label': np.random.random((2, 4)).astype('float32'),
+                },
+                fetch_list=[loss],
+            )
+        paddle.amp.debugging.disable_operator_stats_collection()
+        op_list = paddle.base.core.get_low_precision_op_list()
+        return fp32_grads, op_list
+
+    def check_pir_results(
+        self, fp32_grads, op_list, total_steps, accumulate_batches_num
+    ):
+        for grad in fp32_grads:
+            self.assertEqual(grad.dtype, core.DataType.FLOAT32)
+        # fp16 calls
+        self.assertEqual(
+            int(op_list['pd_op.matmul'].split(',')[0]), total_steps
+        )
+        self.assertEqual(
+            int(op_list['pd_op.adam_'].split(',')[0]),
+            2 * total_steps,
+        )
+        self.assertEqual(
+            int(op_list['pd_op.cast'].split(',')[0]),
+            total_steps * 3,
+        )
+
+    def test_pir_adam_master_grad(self):
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                total_steps = 4
+                accumulate_batches_num = 2
+                model = SimpleNet(2, 4)
+                opt = paddle.optimizer.Adam(parameters=model.parameters())
+                fp32_grads, op_list = self.run_pir(
+                    total_steps, accumulate_batches_num, model, opt
+                )
+                self.check_pir_results(
+                    fp32_grads, op_list, total_steps, accumulate_batches_num
+                )
+
+    def test_pir_momentum_master_grad(self):
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                total_steps = 4
+                accumulate_batches_num = 1
+                model = SimpleNet(2, 4)
+                L1Decay = paddle.regularizer.L1Decay(0.0001)
+                opt = paddle.optimizer.Momentum(
+                    parameters=model.parameters(), weight_decay=L1Decay
+                )
+                fp32_grads, op_list = self.run_pir(
+                    total_steps, accumulate_batches_num, model, opt
+                )
+                for grad in fp32_grads:
+                    self.assertEqual(grad.dtype, core.DataType.FLOAT32)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/amp/test_amp_master_weight.py

@@ -0,0 +1,213 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from amp_base_models import AmpTestBase
+
+import paddle
+from paddle.base import core
+
+
+class SimpleNet(paddle.nn.Layer):
+    def __init__(self, input_size, output_size):
+        super().__init__()
+        weight_attr = paddle.ParamAttr(
+            name="weight", initializer=paddle.nn.initializer.Constant(value=0.5)
+        )
+        bias_attr = paddle.ParamAttr(
+            name="bias", initializer=paddle.nn.initializer.Constant(value=1.0)
+        )
+        self.linear = paddle.nn.Linear(
+            input_size, output_size, weight_attr, bias_attr
+        )
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+class TestMasterWeight(AmpTestBase):
+    def run_dygraph(self, dtype, level, use_promote, max_iters, x_data):
+        losses = []
+        model = SimpleNet(100, 100)
+        optimizer = paddle.optimizer.AdamW(
+            learning_rate=0.01,
+            parameters=model.parameters(),
+        )
+        scaler = paddle.amp.GradScaler()
+        model, optimizer = paddle.amp.decorate(
+            models=model,
+            optimizers=optimizer,
+            level=level,
+            dtype=dtype,
+        )
+
+        for i in range(max_iters):
+            with paddle.amp.auto_cast(
+                enable=True,
+                dtype=dtype,
+                level=level,
+                use_promote=use_promote,
+            ):
+                x = paddle.to_tensor(x_data, dtype='float16')
+                out = model(x)
+                loss = paddle.mean(out)
+                losses.append(loss)
+            scaled = scaler.scale(loss)
+            scaled.backward()
+            scaler.minimize(optimizer, scaled)
+            optimizer.clear_grad()
+        return losses
+
+    def run_pir(self, dtype, level, use_promote, max_iters, x_data):
+        with paddle.pir_utils.IrGuard():
+            losses = []
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                model = SimpleNet(100, 100)
+                optimizer = paddle.optimizer.AdamW(
+                    learning_rate=0.01,
+                    parameters=model.parameters(),
+                )
+                scaler = paddle.amp.GradScaler(enable=True)
+                model, optimizer = paddle.amp.decorate(
+                    models=model,
+                    optimizers=optimizer,
+                    level=level,
+                    dtype=dtype,
+                )
+                with paddle.amp.auto_cast(
+                    enable=True,
+                    dtype=dtype,
+                    level=level,
+                    use_promote=use_promote,
+                ):
+                    x = paddle.static.data('x', x_data.shape, 'float16')
+                    out = model(x)
+                    loss = paddle.mean(out)
+                scaled = scaler.scale(loss)
+                scaler.minimize(optimizer, scaled)
+            if paddle.is_compiled_with_cuda():
+                place = paddle.CUDAPlace(0)
+            elif paddle.device.is_compiled_with_xpu():
+                place = paddle.device.XPUPlace(0)
+            else:
+                raise ValueError("Only support CUDA or XPU Place.")
+            exe = paddle.static.Executor(place)
+            exe.run(startup)
+            for iter_id in range(max_iters):
+                results = exe.run(
+                    main,
+                    feed={'x': x_data},
+                    fetch_list=[loss],
+                )
+
+                losses.append(results[0])
+
+            return losses
+
+    def run_static(self, dtype, level, use_promote, max_iters, x_data):
+        paddle.enable_static()
+        with paddle.pir_utils.OldIrGuard():
+            main_program = paddle.static.Program()
+            startup_program = paddle.static.Program()
+            losses = []
+            with (
+                paddle.utils.unique_name.guard(),
+                paddle.static.program_guard(main_program, startup_program),
+            ):
+                model = SimpleNet(100, 100)
+                optimizer = paddle.optimizer.AdamW(learning_rate=0.01)
+                optimizer = paddle.static.amp.decorate(
+                    optimizer,
+                    level=level,
+                    dtype=dtype,
+                    use_promote=use_promote,
+                    master_weight=True,
+                )
+                x = paddle.static.data(
+                    name='input', shape=[100, 100], dtype='float16'
+                )
+                out = model(x)
+                loss = paddle.mean(out)
+                optimizer.minimize(loss)
+
+            if paddle.is_compiled_with_cuda():
+                place = paddle.CUDAPlace(0)
+            elif paddle.device.is_compiled_with_xpu():
+                place = paddle.device.XPUPlace(0)
+            else:
+                raise ValueError("Only support CUDA or XPU Place.")
+            exe = paddle.static.Executor(place)
+            exe.run(startup_program)
+            optimizer.amp_init(
+                place,
+                scope=paddle.static.global_scope(),
+                rewrite_master_weight=True,
+            )
+            for iter_id in range(max_iters):
+                results = exe.run(
+                    program=main_program,
+                    feed={x.name: x_data},
+                    fetch_list=[loss],
+                )
+                print(
+                    f"-- [AMP {dtype} {level}] iter={iter_id}, loss={results[0]}"
+                )
+                losses.append(results[0])
+
+        paddle.disable_static()
+        return losses
+
+    def test_master_weight(self):
+        np.random.seed(1)
+        paddle.seed(1)
+        dtype = 'float16'
+        level = 'O2'
+        use_promote = True
+        total_steps = 4
+        x_data = np.random.random(size=[100, 100]).astype("float16")
+
+        loss_dygraph = self.run_dygraph(
+            dtype, level, use_promote, total_steps, x_data
+        )
+        loss_static = self.run_static(
+            dtype, level, use_promote, total_steps, x_data
+        )
+        loss_pir = self.run_pir(dtype, level, use_promote, total_steps, x_data)
+
+        for i in range(total_steps):
+            self.assertEqual(loss_dygraph[i], loss_static[i])
+            self.assertEqual(loss_dygraph[i], loss_pir[i])
+
+
+if __name__ == '__main__':
+    unittest.main()

test/amp/test_amp_o2_embedding_model.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import random
+import unittest
+
+import numpy as np
+from amp_base_models import AmpTestBase, _build_optimizer
+
+import paddle
+from paddle import nn
+from paddle.framework import in_dynamic_or_pir_mode
+
+paddle.enable_static()
+
+_fixed_param = np.random.random(size=[64, 64]).astype("float32")
+
+
+class SimpleUnittedEmbeddingNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.vocab_size = 64
+        self.hidden_size = 64
+        global _fixed_param
+
+        self.param_attr = paddle.ParamAttr(
+            initializer=paddle.nn.initializer.Assign(_fixed_param)
+        )
+        self.embedding = nn.Embedding(
+            self.vocab_size, self.hidden_size, weight_attr=self.param_attr
+        )
+        self.linear = nn.Linear(
+            in_features=self.hidden_size,
+            out_features=self.vocab_size,
+            weight_attr=self.param_attr,
+        )
+
+    def forward(self, x):
+        out = self.embedding(x)
+        scale = paddle.full(shape=[1], fill_value=2, dtype="int64")
+        out = paddle.multiply(out, scale.astype("float32"))
+        out = self.linear(out)
+        out = nn.functional.dropout(out, p=0.2)
+        return out
+
+
+def build_unitted_embedding_model(
+    use_amp,
+    amp_dtype="float16",
+    amp_level="O1",
+    use_promote=False,
+):
+    if in_dynamic_or_pir_mode():
+        model = SimpleUnittedEmbeddingNet()
+        optimizer = _build_optimizer(use_amp=False, model=model)
+        if use_amp and amp_dtype == "float16":
+            scaler = paddle.amp.GradScaler(init_loss_scaling=32768.0)
+        else:
+            scaler = None
+        if use_amp and amp_level == "O2":
+            model, optimizer = paddle.amp.decorate(
+                models=model,
+                optimizers=optimizer,
+                level=amp_level,
+                dtype=amp_dtype,
+            )
+        return model, optimizer, scaler
+    else:
+        raise ValueError("Only support pir mode")
+
+
+class TestUnittedEmbedding(AmpTestBase):
+    def _generate_feed_x(self):
+        seed = 0
+        paddle.seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+        x = np.random.randint(1, 64, size=[1, 32]).astype("int64")
+        return x
+
+    def test_pir_compare_o1_and_o2_master_grad(self):
+        def _run(data, level, use_promote=False):
+            with paddle.pir_utils.IrGuard():
+                startup = paddle.static.Program()
+                main = paddle.static.Program()
+                with paddle.static.program_guard(main, startup):
+                    model, optimizer, scaler = build_unitted_embedding_model(
+                        use_amp=True,
+                        amp_dtype="float16",
+                        amp_level=level,
+                        use_promote=use_promote,
+                    )
+                    model.train()
+                    with paddle.amp.auto_cast(
+                        enable=True,
+                        dtype='float16',
+                        level=level,
+                        use_promote=use_promote,
+                    ):
+                        x = paddle.static.data('x', [None, 32], 'int64')
+                        out = model(x)
+                        loss = paddle.mean(out)
+                    scaled = scaler.scale(loss)
+                    scaler.minimize(optimizer, scaled)
+
+                    if paddle.is_compiled_with_cuda():
+                        place = paddle.CUDAPlace(0)
+                    elif paddle.device.is_compiled_with_xpu():
+                        place = paddle.device.XPUPlace(0)
+                    else:
+                        raise ValueError("Only support CUDA or XPU Place.")
+                    exe = paddle.static.Executor(place)
+                    exe.run(startup)
+                    exe.run(
+                        main,
+                        feed={
+                            'x': data,
+                        },
+                        fetch_list=[loss],
+                    )
+
+        x = self._generate_feed_x()
+        _run(x, 'O2', use_promote=False)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/amp/test_amp_promote.py

@@ -0,0 +1,424 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from amp_base_models import AmpTestBase, build_conv_model
+
+import paddle
+from paddle.base import core
+from paddle.static import amp
+
+paddle.set_flags({"FLAGS_use_legacy_linear": True})
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+class TestStaticAmpPromoteStats(AmpTestBase):
+    def check_promote_results(
+        self, use_amp, dtype, level, use_promote, expected_op_calls, debug_info
+    ):
+        paddle.enable_static()
+        (
+            main_program,
+            startup_program,
+            optimizer,
+            feed_vars,
+            fetch_vars,
+        ) = build_conv_model(use_amp, dtype, level, use_promote)
+        self.assertEqual(main_program.num_blocks, 1)
+
+        amp.debugging.collect_operator_stats(main_program)
+        op_stats_list = amp.debugging._get_op_stats_list(main_program)
+
+        self._check_op_calls(
+            op_stats_list[0],
+            expected_fp16_calls=expected_op_calls,
+            debug_info=debug_info,
+        )
+
+        if paddle.is_compiled_with_cuda():
+            place = paddle.CUDAPlace(0)
+        elif paddle.device.is_compiled_with_xpu():
+            place = paddle.device.XPUPlace(0)
+        else:
+            raise ValueError("Only support CUDA or XPU Place.")
+        exe = paddle.static.Executor(place)
+
+        max_iters = 2
+        x_fp32 = np.random.random(size=[1, 1, 6, 6]).astype("float32")
+        losses_o1 = self.run_program(
+            main_program,
+            startup_program,
+            optimizer,
+            feed_vars,
+            fetch_vars,
+            place,
+            exe,
+            x_fp32,
+            max_iters,
+            dtype,
+            level,
+        )
+        paddle.disable_static()
+
+    def test_static_amp_o1(self):
+        expected_fp16_calls = {
+            "conv2d": 1,
+            "elementwise_add": 0,
+            "relu": 0,
+            "matmul_v2": 1,
+            "softmax": 0,
+            "reduce_mean": 0,
+            "adamw": 0,
+        }
+        with paddle.pir_utils.OldIrGuard():
+            self.check_promote_results(
+                True,
+                'float16',
+                'O1',
+                use_promote=True,
+                expected_op_calls=expected_fp16_calls,
+                debug_info="TestStaticAmpPromoteStats/test_static_amp_o1",
+            )
+
+    def test_static_amp_o2(self):
+        expected_fp16_calls = {
+            "conv2d": 1,
+            "elementwise_add": 2,
+            "relu": 0,
+            "matmul_v2": 1,
+            "softmax": 1,
+            "reduce_mean": 1,
+            "adamw": 4,
+        }
+        with paddle.pir_utils.OldIrGuard():
+            self.check_promote_results(
+                True,
+                'float16',
+                'O2',
+                use_promote=True,
+                expected_op_calls=expected_fp16_calls,
+                debug_info="TestStaticAmpPromoteStats/test_static_amp_o2",
+            )
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+class TestEagerAmpPromoteStats(AmpTestBase):
+    def check_promote_results(
+        self, dtype, level, use_promote, expected_op_calls, debug_info
+    ):
+        model, optimizer, scaler = build_conv_model(
+            use_amp=True,
+            amp_dtype=dtype,
+            amp_level=level,
+            use_promote=use_promote,
+        )
+        model.train()
+
+        paddle.amp.debugging.enable_operator_stats_collection()
+        with paddle.amp.auto_cast(
+            enable=True, dtype=dtype, level=level, use_promote=use_promote
+        ):
+            x = paddle.rand(shape=[1, 1, 6, 6], dtype='float32')
+            out = model(x)
+            loss = paddle.mean(out)
+        scaled = scaler.scale(loss)
+        scaled.backward()
+        scaler.minimize(optimizer, scaled)
+        optimizer.clear_grad()
+        paddle.amp.debugging.disable_operator_stats_collection()
+        op_stats = paddle.base.core.get_low_precision_op_list()
+
+        self._check_op_calls(
+            op_stats,
+            expected_fp16_calls=expected_op_calls,
+            debug_info=debug_info,
+        )
+
+    def test_o2_promote_on(self):
+        expected_fp16_calls = {
+            "conv2d": 1,
+            "elementwise_add": 2,
+            "relu": 0,
+            "matmul_v2": 1,
+            "softmax": 1,
+            "reduce_mean": 1,
+            "adamw_": 4,
+        }
+        self.check_promote_results(
+            'float16',
+            'O2',
+            use_promote=True,
+            expected_op_calls=expected_fp16_calls,
+            debug_info="TestEagerAmpPromoteStats/test_o2_promote_on",
+        )
+
+    def test_o2_promote_off(self):
+        expected_fp16_calls = {
+            "conv2d": 1,
+            "elementwise_add": 2,
+            "relu": 1,
+            "matmul_v2": 1,
+            "softmax": 1,
+            "reduce_mean": 1,
+            "adamw_": 4,
+        }
+        self.check_promote_results(
+            'float16',
+            'O2',
+            use_promote=False,
+            expected_op_calls=expected_fp16_calls,
+            debug_info="TestEagerAmpPromoteStats/test_o2_promote_off",
+        )
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+class TestPirAmpPromoteStats(AmpTestBase):
+    def check_promote_results(
+        self, dtype, level, use_promote, expected_op_calls, debug_info
+    ):
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                model, optimizer, scaler = build_conv_model(
+                    use_amp=True,
+                    amp_dtype=dtype,
+                    amp_level=level,
+                    use_promote=use_promote,
+                )
+                model.train()
+
+                with paddle.amp.auto_cast(
+                    enable=True,
+                    dtype=dtype,
+                    level=level,
+                    use_promote=use_promote,
+                ):
+                    x = paddle.static.data(
+                        'x', shape=[1, 1, 6, 6], dtype='float32'
+                    )
+                    out = model(x)
+                    loss = paddle.mean(out)
+                scaled = scaler.scale(loss)
+                scaler.minimize(optimizer, scaled)
+
+                if paddle.is_compiled_with_cuda():
+                    place = paddle.CUDAPlace(0)
+                elif paddle.device.is_compiled_with_xpu():
+                    place = paddle.device.XPUPlace(0)
+                else:
+                    raise ValueError("Only support CUDA or XPU Place.")
+                exe = paddle.static.Executor(place)
+                exe.run(startup)
+                paddle.amp.debugging.enable_operator_stats_collection()
+                exe.run(
+                    main,
+                    feed={
+                        'x': np.random.random([1, 1, 6, 6]).astype('float32'),
+                    },
+                    fetch_list=[loss],
+                )
+                paddle.amp.debugging.disable_operator_stats_collection()
+                op_stats = paddle.base.core.get_low_precision_op_list()
+
+                self._check_op_calls(
+                    op_stats,
+                    expected_fp16_calls=expected_op_calls,
+                    debug_info=debug_info,
+                )
+
+    def test_o2_promote_on(self):
+        paddle.set_flags({"FLAGS_pir_apply_inplace_pass": 0})
+        expected_fp16_calls = {
+            "pd_op.conv2d": 1,
+            "pd_op.add": 2,
+            "pd_op.relu": 0,
+            "pd_op.matmul": 1,
+            "pd_op.softmax": 1,
+            "pd_op.mean": 1,
+            "pd_op.adamw_": 4,
+        }
+        self.check_promote_results(
+            'float16',
+            'O2',
+            use_promote=True,
+            expected_op_calls=expected_fp16_calls,
+            debug_info="TestEagerAmpPromoteStats/test_o2_promote_on",
+        )
+
+    def test_o2_promote_off(self):
+        paddle.set_flags({"FLAGS_pir_apply_inplace_pass": 0})
+        expected_fp16_calls = {
+            "pd_op.conv2d": 1,
+            "pd_op.add": 2,
+            "pd_op.relu": 1,
+            "pd_op.matmul": 1,
+            "pd_op.softmax": 1,
+            "pd_op.mean": 1,
+            "pd_op.adamw_": 4,
+        }
+        self.check_promote_results(
+            'float16',
+            'O2',
+            use_promote=False,
+            expected_op_calls=expected_fp16_calls,
+            debug_info="TestEagerAmpPromoteStats/test_o2_promote_off",
+        )
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+class TestEagerAmpPromoteSimple(AmpTestBase):
+    def setUp(self):
+        self._conv = paddle.nn.Conv2D(
+            in_channels=1, out_channels=6, kernel_size=3, bias_attr=False
+        )
+        self._linear = paddle.nn.Linear(in_features=4, out_features=4)
+
+    def test_o2_use_promote_on(self):
+        with paddle.amp.auto_cast(level='O2'):
+            x = paddle.rand(shape=[1, 1, 6, 6], dtype='float32')
+            conv_out = self._conv(x)
+            y = paddle.rand(shape=conv_out.shape, dtype='float16')
+            add_out = conv_out + y
+            linear_out = self._linear(add_out)
+
+        self.assertEqual(conv_out.dtype, paddle.float16)
+        self.assertEqual(add_out.dtype, paddle.float16)
+        self.assertEqual(linear_out.dtype, paddle.float32)
+
+    def test_o2_use_promote_off(self):
+        with paddle.amp.auto_cast(level='O2', use_promote=False):
+            x = paddle.rand(shape=[1, 1, 6, 6], dtype='float32')
+            conv_out = self._conv(x)
+            y = paddle.rand(shape=conv_out.shape, dtype='float16')
+            add_out = conv_out + y
+            linear_out = self._linear(add_out)
+
+        self.assertEqual(conv_out.dtype, paddle.float16)
+        self.assertEqual(add_out.dtype, paddle.float16)
+        self.assertEqual(linear_out.dtype, paddle.float16)
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_cuda()
+    and paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "run test when gpu's compute capability is at least 7.0.",
+)
+@unittest.skipIf(
+    core.is_compiled_with_xpu()
+    and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+    "run test when xpu's compute capability >= xpu3.",
+)
+class TestPirAmpPromoteSimple(AmpTestBase):
+    def init_net(self):
+        self._conv = paddle.nn.Conv2D(
+            in_channels=1, out_channels=6, kernel_size=3, bias_attr=False
+        )
+        self._linear = paddle.nn.Linear(in_features=4, out_features=4)
+
+    def test_o2_use_promote_on(self):
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                self.init_net()
+                with paddle.amp.auto_cast(level='O2'):
+                    x = paddle.rand(shape=[1, 1, 6, 6], dtype='float32')
+                    conv_out = self._conv(x)
+                    y = paddle.rand(shape=conv_out.shape, dtype='float16')
+                    add_out = conv_out + y
+                    linear_out = self._linear(add_out)
+
+            self.assertEqual(conv_out.dtype, paddle.float16)
+            self.assertEqual(add_out.dtype, paddle.float16)
+            self.assertEqual(linear_out.dtype, paddle.float32)
+
+    def test_o2_use_promote_off(self):
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                self.init_net()
+                with paddle.amp.auto_cast(level='O2', use_promote=False):
+                    x = paddle.rand(shape=[1, 1, 6, 6], dtype='float32')
+                    conv_out = self._conv(x)
+                    y = paddle.rand(shape=conv_out.shape, dtype='float16')
+                    add_out = conv_out + y
+                    linear_out = self._linear(add_out)
+
+            self.assertEqual(conv_out.dtype, paddle.float16)
+            self.assertEqual(add_out.dtype, paddle.float16)
+            self.assertEqual(linear_out.dtype, paddle.float16)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/amp/test_collect_operator_stats.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import unittest
+
+import numpy as np
+
+sys.path.append("../../amp")
+from amp_base_models import build_while_model
+
+import paddle
+
+
+class TestOpStatsEager(unittest.TestCase):
+    def _check_result(self, dtype):
+        # Returned the dict.
+        op_list = paddle.base.core.get_low_precision_op_list()
+
+        self.assertTrue('elementwise_add' in op_list)
+        self.assertTrue('conv2d' in op_list)
+
+        conv2d_called = op_list['conv2d'].split(',')
+        add_called = op_list['elementwise_add'].split(',')
+        add_num = 0
+        conv_num = 0
+        for i in range(4):
+            add_num += int(add_called[i])
+            conv_num += int(conv2d_called[i])
+
+        self.assertTrue(conv_num == 1)
+        self.assertTrue(add_num == 1)
+
+        if dtype == paddle.float16:
+            self.assertTrue(int(conv2d_called[0]) == 1)
+            self.assertTrue(int(add_called[0]) == 1)
+
+    def test_enable_disable(self):
+        conv = paddle.nn.Conv2D(3, 2, 3)
+        x = paddle.rand([10, 3, 32, 32])
+
+        paddle.amp.debugging.enable_operator_stats_collection()
+        # amp list conv2d, elementwise_add, cast (transfer_dtype)
+        with paddle.amp.auto_cast(enable=True, level='O2'):
+            out = conv(x)
+        # Print to the standard output.
+        paddle.amp.debugging.disable_operator_stats_collection()
+
+        self._check_result(dtype=out.dtype)
+
+    def test_context(self):
+        conv = paddle.nn.Conv2D(3, 2, 3)
+        x = paddle.rand([10, 3, 32, 32])
+
+        with (
+            paddle.amp.debugging.collect_operator_stats(),
+            # amp list conv2d, elementwise_add, cast (transfer_dtype)
+            paddle.amp.auto_cast(enable=True, level='O2'),
+        ):
+            out = conv(x)
+
+        self._check_result(dtype=out.dtype)
+
+
+class TestOpStatsPir(unittest.TestCase):
+    def _check_result(self, dtype):
+        # Returned the dict.
+        op_list = paddle.base.core.get_low_precision_op_list()
+
+        self.assertTrue('pd_op.add' in op_list)
+        self.assertTrue('pd_op.conv2d' in op_list)
+
+        conv2d_called = op_list['pd_op.conv2d'].split(',')
+        add_called = op_list['pd_op.add'].split(',')
+        add_num = 0
+        conv_num = 0
+        for i in range(4):
+            add_num += int(add_called[i])
+            conv_num += int(conv2d_called[i])
+
+        self.assertTrue(conv_num == 1)
+        self.assertTrue(add_num == 1)
+
+        if dtype == paddle.float16:
+            self.assertTrue(int(conv2d_called[0]) == 1)
+            self.assertTrue(int(add_called[0]) == 1)
+
+    def test_enable_disable(self):
+        if not paddle.is_compiled_with_cuda():
+            return
+        paddle.set_flags({"FLAGS_pir_apply_inplace_pass": 0})
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                conv = paddle.nn.Conv2D(3, 2, 3)
+                x = paddle.static.data('x', [10, 3, 32, 32], 'float32')
+
+                with paddle.amp.auto_cast(enable=True, level='O2'):
+                    out = conv(x)
+
+                place = paddle.CUDAPlace(0)
+                exe = paddle.static.Executor(place)
+                exe.run(startup)
+                paddle.amp.debugging.enable_operator_stats_collection()
+                exe.run(
+                    main,
+                    feed={
+                        'x': np.random.random([10, 3, 32, 32]).astype(
+                            'float32'
+                        ),
+                    },
+                    fetch_list=[out],
+                )
+                paddle.amp.debugging.disable_operator_stats_collection()
+                self._check_result(dtype=out.dtype)
+
+    def test_context(self):
+        if not paddle.is_compiled_with_cuda():
+            return
+        paddle.set_flags({"FLAGS_pir_apply_inplace_pass": 0})
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                conv = paddle.nn.Conv2D(3, 2, 3)
+                x = paddle.static.data('x', [10, 3, 32, 32], 'float32')
+                with paddle.amp.auto_cast(enable=True, level='O2'):
+                    out = conv(x)
+
+            place = paddle.CUDAPlace(0)
+            exe = paddle.static.Executor(place)
+            exe.run(startup)
+            with paddle.amp.debugging.collect_operator_stats():
+                exe.run(
+                    main,
+                    feed={
+                        'x': np.random.random([10, 3, 32, 32]).astype(
+                            'float32'
+                        ),
+                    },
+                    fetch_list=[out],
+                )
+            self._check_result(dtype=out.dtype)
+
+
+class TestOpStatsStatic(unittest.TestCase):
+    def test_while_op(self):
+        paddle.enable_static()
+        main_program, startup_program = build_while_model()
+        if paddle.framework.use_pir_api():
+            self.assertEqual(main_program.num_blocks, 1)
+        else:
+            self.assertEqual(main_program.num_blocks, 2)
+
+            paddle.static.amp.debugging.collect_operator_stats(
+                program=main_program, print_subblocks=True
+            )
+        paddle.disable_static()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/amp/test_compare_accuracy_api.py

@@ -0,0 +1,147 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.base import core
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda(),
+    "not support cpu TestEagerCompareAccuracyApi",
+)
+class TestEagerCompareAccuracyApi(unittest.TestCase):
+    def calc(self, path, dtype):
+        paddle.base.core.set_nan_inf_debug_path(path)
+        x = paddle.to_tensor(
+            [2000, 3000, 4, 0], place=core.CUDAPlace(0), dtype=dtype
+        )
+        y = paddle.to_tensor(
+            [100, 500, 2, 10000], place=core.CUDAPlace(0), dtype=dtype
+        )
+        # normal
+        z1 = x + y
+        # inf
+        z2 = x * y
+
+    def test(self):
+        paddle.set_flags(
+            {"FLAGS_check_nan_inf": 1, "FLAGS_check_nan_inf_level": 3}
+        )
+        fp32_path = "workerlog_fp32_log_dir"
+        fp16_path = "workerlog_fp16_log_dir"
+        self.calc(fp32_path, "float32")
+        self.calc(fp16_path, "float16")
+
+        out_excel = "compare_accuracy_out_excel.csv"
+        paddle.amp.debugging.compare_accuracy(
+            fp32_path,
+            fp16_path,
+            out_excel,
+            loss_scale=1,
+            dump_all_tensors=False,
+        )
+
+    def test2(self):
+        fp32_path = "workerlog_fp32_log_dir"
+        fp16_path = "workerlog_fp16_null_log_dir"
+        self.calc(fp32_path, "float32")
+        out_excel = "compare_accuracy_out_excel_2.csv"
+        paddle.amp.debugging.compare_accuracy(
+            fp32_path,
+            fp16_path,
+            out_excel,
+            loss_scale=1,
+            dump_all_tensors=False,
+        )
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda(),
+    "not support cpu TestPirCompareAccuracyApi",
+)
+class TestPirCompareAccuracyApi(unittest.TestCase):
+    def calc(self, path, dtype):
+        paddle.base.core.set_nan_inf_debug_path(path)
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data(
+                    'x',
+                    [
+                        4,
+                    ],
+                    dtype,
+                )
+                y = paddle.static.data(
+                    'y',
+                    [
+                        4,
+                    ],
+                    dtype,
+                )
+                # normal
+                z1 = x + y
+                # inf
+                z2 = x * y
+            place = paddle.CUDAPlace(0)
+            exe = paddle.static.Executor(place)
+            exe.run(startup)
+            exe.run(
+                main,
+                feed={
+                    'x': np.array([2000, 3000, 4, 0]).astype(dtype),
+                    'y': np.array([100, 500, 2, 10000]).astype(dtype),
+                },
+                fetch_list=[z2],
+            )
+
+    def test(self):
+        paddle.set_flags(
+            {"FLAGS_check_nan_inf": 1, "FLAGS_check_nan_inf_level": 3}
+        )
+        fp32_path = "workerlog_fp32_log_dir"
+        fp16_path = "workerlog_fp16_log_dir"
+        self.calc(fp32_path, "float32")
+        self.calc(fp16_path, "float16")
+
+        out_excel = "compare_accuracy_out_excel.csv"
+        paddle.amp.debugging.compare_accuracy(
+            fp32_path,
+            fp16_path,
+            out_excel,
+            loss_scale=1,
+            dump_all_tensors=False,
+        )
+
+    def test2(self):
+        fp32_path = "workerlog_fp32_log_dir"
+        fp16_path = "workerlog_fp16_null_log_dir"
+        self.calc(fp32_path, "float32")
+        out_excel = "compare_accuracy_out_excel_2.csv"
+        paddle.amp.debugging.compare_accuracy(
+            fp32_path,
+            fp16_path,
+            out_excel,
+            loss_scale=1,
+            dump_all_tensors=False,
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/amp/test_get_autocast_dtype.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+
+
+@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 = "float16"
+
+    def do_test(self, device, expected_type):
+        self.assertTrue(paddle.get_autocast_dtype(device) == expected_type)
+        self.assertTrue(paddle.get_autocast_gpu_dtype() == expected_type)
+        self.assertTrue(paddle.amp.get_autocast_dtype(device) == expected_type)
+        self.assertTrue(paddle.amp.get_autocast_gpu_dtype() == expected_type)
+        self.assertTrue(paddle.amp.get_autocast_cpu_dtype() == expected_type)
+        self.assertTrue(
+            paddle.amp.get_autocast_cpu_dtype(device) == expected_type
+        )
+
+    def test_amp_default(self):
+        for device in self.device_list:
+            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, self.default_dtype)
+
+    @unittest.skipIf(
+        not paddle.amp.is_bfloat16_supported(),
+        "Skip BF16 test if BF16 is not supported",
+    )
+    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, self.default_dtype)
+
+    @unittest.skipIf(
+        not paddle.amp.is_bfloat16_supported(),
+        "Skip BF16 test if BF16 is not supported",
+    )
+    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, self.default_dtype)
+
+    @unittest.skipIf(
+        not paddle.amp.is_bfloat16_supported(),
+        "Skip BF16 test if BF16 is not supported",
+    )
+    def test_amp_nested_context(self):
+        for device in self.device_list:
+            with paddle.amp.auto_cast(True, dtype="bfloat16"):
+                self.do_test(device, "bfloat16")
+                with paddle.amp.auto_cast(True, dtype="float16"):
+                    self.do_test(device, "float16")
+                self.do_test(device, "bfloat16")
+            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 = "float16"
+
+
+if __name__ == "__main__":
+    unittest.main()

test/amp/test_grad_scaler_out.py

@@ -0,0 +1,314 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.amp import AmpScaler, GradScaler
+
+
+class TestGradScalerParamAlias(unittest.TestCase):
+    """Test ParamAliasDecorator: PyTorch aliases, Paddle names, mixed,
+    positional args, single alias, conflict detection, and non-aliased params."""
+
+    def _assert_scaler(self, scaler, **expected):
+        """Helper to check multiple scaler attributes at once."""
+        attr_map = {
+            'enable': '_enable',
+            'init_loss_scaling': '_init_loss_scaling',
+            'incr_ratio': '_incr_ratio',
+            'decr_ratio': '_decr_ratio',
+            'incr_every_n_steps': '_incr_every_n_steps',
+            'decr_every_n_nan_or_inf': '_decr_every_n_nan_or_inf',
+            'use_dynamic_loss_scaling': '_use_dynamic_loss_scaling',
+        }
+        for key, val in expected.items():
+            self.assertEqual(getattr(scaler, attr_map[key]), val, msg=key)
+
+    def test_default_values(self):
+        self._assert_scaler(
+            GradScaler(),
+            enable=True,
+            init_loss_scaling=2.0**16,
+            incr_ratio=2.0,
+            decr_ratio=0.5,
+            incr_every_n_steps=2000,
+            decr_every_n_nan_or_inf=1,
+            use_dynamic_loss_scaling=True,
+        )
+
+    def test_pytorch_style_kwargs(self):
+        self._assert_scaler(
+            GradScaler(
+                enabled=True,
+                init_scale=1024.0,
+                growth_factor=3.0,
+                backoff_factor=0.25,
+                growth_interval=500,
+            ),
+            enable=True,
+            init_loss_scaling=1024.0,
+            incr_ratio=3.0,
+            decr_ratio=0.25,
+            incr_every_n_steps=500,
+        )
+
+    def test_mixed_kwargs(self):
+        self._assert_scaler(
+            GradScaler(
+                enable=True,
+                init_loss_scaling=1024.0,
+                growth_factor=3.0,
+                decr_ratio=0.25,
+                growth_interval=500,
+            ),
+            init_loss_scaling=1024.0,
+            incr_ratio=3.0,
+            decr_ratio=0.25,
+            incr_every_n_steps=500,
+        )
+
+    def test_single_alias_each(self):
+        self.assertFalse(GradScaler(enabled=False)._enable)
+        self.assertEqual(GradScaler(init_scale=512.0)._init_loss_scaling, 512.0)
+        self.assertEqual(GradScaler(growth_factor=5.0)._incr_ratio, 5.0)
+        self.assertEqual(GradScaler(backoff_factor=0.1)._decr_ratio, 0.1)
+        self.assertEqual(
+            GradScaler(growth_interval=100)._incr_every_n_steps, 100
+        )
+
+    def test_positional_args(self):
+        self._assert_scaler(
+            GradScaler(True, 1024.0, 3.0, 0.25, 500, 2, True),
+            enable=True,
+            init_loss_scaling=1024.0,
+            incr_ratio=3.0,
+            decr_ratio=0.25,
+            incr_every_n_steps=500,
+            decr_every_n_nan_or_inf=2,
+        )
+
+    def test_positional_with_alias_kwarg(self):
+        self._assert_scaler(
+            GradScaler(True, 1024.0, growth_factor=5.0),
+            enable=True,
+            init_loss_scaling=1024.0,
+            incr_ratio=5.0,
+        )
+
+    def test_non_aliased_with_aliases(self):
+        self._assert_scaler(
+            GradScaler(
+                enabled=True,
+                init_scale=2048.0,
+                decr_every_n_nan_or_inf=3,
+                use_dynamic_loss_scaling=False,
+            ),
+            enable=True,
+            init_loss_scaling=2048.0,
+            decr_every_n_nan_or_inf=3,
+            use_dynamic_loss_scaling=False,
+        )
+
+    def test_pytorch_vs_paddle_equivalence(self):
+        pt = GradScaler(
+            enabled=True,
+            init_scale=1024.0,
+            growth_factor=3.0,
+            backoff_factor=0.25,
+            growth_interval=500,
+        )
+        pd = GradScaler(
+            enable=True,
+            init_loss_scaling=1024.0,
+            incr_ratio=3.0,
+            decr_ratio=0.25,
+            incr_every_n_steps=500,
+        )
+        for attr in (
+            '_enable',
+            '_init_loss_scaling',
+            '_incr_ratio',
+            '_decr_ratio',
+            '_incr_every_n_steps',
+        ):
+            self.assertEqual(getattr(pt, attr), getattr(pd, attr), msg=attr)
+
+    def test_conflict_raises_error(self):
+        conflicts = [
+            {"enable": True, "enabled": False},
+            {"init_loss_scaling": 1024.0, "init_scale": 2048.0},
+            {"incr_ratio": 2.0, "growth_factor": 3.0},
+            {"decr_ratio": 0.5, "backoff_factor": 0.25},
+            {"incr_every_n_steps": 1000, "growth_interval": 2000},
+        ]
+        for kwargs in conflicts:
+            with self.assertRaises(ValueError, msg=str(kwargs)):
+                GradScaler(**kwargs)
+
+    def test_torch_positional_no_device(self):
+        # PyTorch older API: init_scale first (float -> detected as torch positional)
+        self._assert_scaler(
+            GradScaler(1024.0, 3.0, 0.25, 500, True),
+            enable=True,
+            init_loss_scaling=1024.0,
+            incr_ratio=3.0,
+            decr_ratio=0.25,
+            incr_every_n_steps=500,
+        )
+
+    def test_torch_positional_no_device_disabled(self):
+        # PyTorch: GradScaler(init_scale, ..., enabled=False)
+        scaler = GradScaler(1024.0, 2.0, 0.5, 2000, False)
+        self.assertFalse(scaler._enable)
+        self.assertEqual(scaler._init_loss_scaling, 1.0)
+
+    def test_torch_positional_with_device(self):
+        # PyTorch newer API: device string first
+        self._assert_scaler(
+            GradScaler('cuda', 1024.0, 3.0, 0.25, 500, True),
+            enable=True,
+            init_loss_scaling=1024.0,
+            incr_ratio=3.0,
+            decr_ratio=0.25,
+            incr_every_n_steps=500,
+        )
+
+    def test_torch_device_kwarg_dropped(self):
+        # device kwarg is silently ignored (no Paddle equivalent)
+        self._assert_scaler(
+            GradScaler(device='cuda', init_scale=2048.0, growth_factor=3.0),
+            init_loss_scaling=2048.0,
+            incr_ratio=3.0,
+        )
+
+    def test_torch_device_string_with_kwargs(self):
+        # device as positional string combined with torch keyword aliases
+        self._assert_scaler(
+            GradScaler('cuda', init_scale=512.0, enabled=True),
+            enable=True,
+            init_loss_scaling=512.0,
+        )
+
+    def test_torch_positional_partial(self):
+        # Only init_scale positionally, rest default
+        self._assert_scaler(
+            GradScaler(4096.0),
+            init_loss_scaling=4096.0,
+            incr_ratio=2.0,
+            decr_ratio=0.5,
+            incr_every_n_steps=2000,
+        )
+
+    def test_torch_device_string_only(self):
+        # GradScaler('cuda') — device only, all remaining params default
+        self._assert_scaler(
+            GradScaler('cuda'),
+            enable=True,
+            init_loss_scaling=2.0**16,
+            incr_ratio=2.0,
+            decr_ratio=0.5,
+            incr_every_n_steps=2000,
+        )
+
+    def test_torch_device_string_partial_positional(self):
+        # GradScaler('cuda', init_scale) — device + only first positional param
+        self._assert_scaler(
+            GradScaler('cuda', 1024.0),
+            enable=True,
+            init_loss_scaling=1024.0,
+            incr_ratio=2.0,
+            decr_ratio=0.5,
+            incr_every_n_steps=2000,
+        )
+
+    def test_torch_positional_with_kwarg_disabled(self):
+        # GradScaler(init_scale, enabled=False) — positional float + torch kwarg
+        scaler = GradScaler(1024.0, enabled=False)
+        self.assertFalse(scaler._enable)
+        self.assertEqual(scaler._init_loss_scaling, 1.0)
+
+    def test_disabled_scaler(self):
+        scaler = GradScaler(enabled=False)
+        self.assertFalse(scaler._enable)
+        self.assertEqual(scaler._init_loss_scaling, 1.0)
+
+
+class TestGradScalerPytorchCompatMethods(unittest.TestCase):
+    """Test PyTorch-compatible getter/setter methods."""
+
+    def test_is_enabled(self):
+        self.assertTrue(GradScaler(enabled=True).is_enabled())
+        self.assertFalse(GradScaler(enabled=False).is_enabled())
+
+    def test_get_scale(self):
+        self.assertEqual(GradScaler(init_scale=2048.0).get_scale(), 2048.0)
+        self.assertEqual(GradScaler(enabled=False).get_scale(), 0.0)
+
+    def test_growth_factor_get_set(self):
+        s = GradScaler(growth_factor=3.0)
+        self.assertEqual(s.get_growth_factor(), 3.0)
+        s.set_growth_factor(5.0)
+        self.assertEqual(s.get_growth_factor(), 5.0)
+        self.assertEqual(s.get_incr_ratio(), 5.0)
+
+    def test_backoff_factor_get_set(self):
+        s = GradScaler(backoff_factor=0.25)
+        self.assertEqual(s.get_backoff_factor(), 0.25)
+        s.set_backoff_factor(0.1)
+        self.assertEqual(s.get_backoff_factor(), 0.1)
+        self.assertEqual(s.get_decr_ratio(), 0.1)
+
+    def test_growth_interval_get_set(self):
+        s = GradScaler(growth_interval=500)
+        self.assertEqual(s.get_growth_interval(), 500)
+        s.set_growth_interval(100)
+        self.assertEqual(s.get_growth_interval(), 100)
+        self.assertEqual(s.get_incr_every_n_steps(), 100)
+
+
+class TestGradScalerCallPathsAndInheritance(unittest.TestCase):
+    """Test all public call paths and AmpScaler vs GradScaler defaults."""
+
+    def test_all_paths_same_class(self):
+        self.assertIs(paddle.device.amp.GradScaler, GradScaler)
+        self.assertIs(paddle.cuda.amp.GradScaler, GradScaler)
+
+    def test_gradscaler_is_subclass_of_ampscaler(self):
+        self.assertTrue(issubclass(GradScaler, AmpScaler))
+
+    def test_alias_via_device_and_cuda(self):
+        s1 = paddle.device.amp.GradScaler(init_scale=512.0, growth_factor=4.0)
+        self.assertEqual(s1._init_loss_scaling, 512.0)
+        self.assertEqual(s1._incr_ratio, 4.0)
+
+        s2 = paddle.cuda.amp.GradScaler(backoff_factor=0.3, growth_interval=800)
+        self.assertEqual(s2._decr_ratio, 0.3)
+        self.assertEqual(s2._incr_every_n_steps, 800)
+
+    def test_defaults_differ_from_ampscaler(self):
+        a, g = AmpScaler(), GradScaler()
+        # Intentionally different: GradScaler aligns with PyTorch
+        self.assertEqual(a._init_loss_scaling, 2.0**15)
+        self.assertEqual(g._init_loss_scaling, 2.0**16)
+        self.assertEqual(a._incr_every_n_steps, 1000)
+        self.assertEqual(g._incr_every_n_steps, 2000)
+        # Shared defaults stay the same
+        self.assertEqual(a._incr_ratio, g._incr_ratio)
+        self.assertEqual(a._decr_ratio, g._decr_ratio)
+        self.assertEqual(a._decr_every_n_nan_or_inf, g._decr_every_n_nan_or_inf)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/amp/test_is_autocast_enabled.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+
+
+@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):
+        for device in self.device_list:
+            self.assertFalse(paddle.is_autocast_enabled(device))
+            self.assertFalse(paddle.amp.is_autocast_enabled(device))
+
+    def test_amp_autocast_true(self):
+        for device in self.device_list:
+            with paddle.amp.auto_cast(True):
+                self.assertTrue(paddle.is_autocast_enabled(device))
+                self.assertTrue(paddle.amp.is_autocast_enabled(device))
+
+            self.assertFalse(paddle.is_autocast_enabled(device))
+            self.assertFalse(paddle.amp.is_autocast_enabled(device))
+
+    def test_amp_autocast_false(self):
+        for device in self.device_list:
+            with paddle.amp.auto_cast(False):
+                self.assertFalse(paddle.is_autocast_enabled(device))
+                self.assertFalse(paddle.amp.is_autocast_enabled(device))
+
+            self.assertFalse(paddle.is_autocast_enabled(device))
+            self.assertFalse(paddle.amp.is_autocast_enabled(device))
+
+    def test_amp_nested_context(self):
+        for device in self.device_list:
+            with paddle.amp.auto_cast(True):
+                self.assertTrue(paddle.is_autocast_enabled(device))
+                self.assertTrue(paddle.amp.is_autocast_enabled(device))
+
+                with paddle.amp.auto_cast(False):
+                    self.assertFalse(paddle.is_autocast_enabled(device))
+                    self.assertFalse(paddle.amp.is_autocast_enabled(device))
+
+                self.assertTrue(paddle.is_autocast_enabled(device))
+                self.assertTrue(paddle.amp.is_autocast_enabled(device))
+            self.assertFalse(paddle.is_autocast_enabled(device))
+            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()

test/amp/test_layer_convert_dtype.py

@@ -0,0 +1,226 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.base import core
+
+
+class MyModel(paddle.nn.Layer):
+    def __init__(self, input_size, hidden_size):
+        super().__init__()
+        self.linear1 = paddle.nn.Linear(input_size, hidden_size)
+        self.linear2 = paddle.nn.Linear(hidden_size, hidden_size)
+        self.linear3 = paddle.nn.Linear(hidden_size, 1)
+        self.batchnorm = paddle.nn.Sequential(paddle.nn.BatchNorm(hidden_size))
+        register_buffer_in_temp = paddle.ones([4, 6])
+        self.register_buffer('register_buffer_in', register_buffer_in_temp)
+
+    def forward(self, inputs):
+        x = self.linear1(inputs)
+        x = F.relu(x)
+        x = self.batchnorm(x)
+        x = self.linear3(x)
+        return x
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+    "Require compiled with CUDA or XPU.",
+)
+class TestDtypeConvert(unittest.TestCase):
+    def setUp(self):
+        self.batch_size, self.input_size, self.hidden_size = 128, 128, 256
+
+    def verify_trans_dtype(
+        self, test_type=None, excluded_layers=None, corrected_dtype=None
+    ):
+        model = MyModel(self.input_size, self.hidden_size)
+        if test_type == 'float16':
+            model.float16(excluded_layers=excluded_layers)
+        elif test_type == 'bfloat16':
+            model.bfloat16(excluded_layers=excluded_layers)
+        else:
+            model.float(excluded_layers=excluded_layers)
+
+        for name, para in model.named_parameters():
+            if 'linear' in name:
+                self.assertEqual(para.dtype, corrected_dtype)
+            elif 'batchnorm' in name:
+                if excluded_layers is None:
+                    self.assertEqual(para.dtype, paddle.float32)
+                else:
+                    self.assertEqual(para.dtype, paddle.float16)
+
+    def test_excluded_layers(self):
+        self.verify_trans_dtype(
+            test_type='float16',
+            excluded_layers=[nn.Linear],
+            corrected_dtype=paddle.float32,
+        )
+        self.verify_trans_dtype(
+            test_type='float16',
+            excluded_layers=nn.Linear,
+            corrected_dtype=paddle.float32,
+        )
+
+    def test_float16(self):
+        self.verify_trans_dtype(
+            test_type='float16',
+            corrected_dtype=paddle.float16,
+        )
+
+    @unittest.skipIf(
+        not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+        "Require compiled with CUDA or XPU.",
+    )
+    @unittest.skipIf(
+        core.is_compiled_with_cuda()
+        and paddle.device.cuda.get_device_capability()[0] >= 8.0,
+        "run test when maximum gpu's compute capability is 8.0.",
+    )
+    @unittest.skipIf(
+        core.is_compiled_with_xpu()
+        and core.get_xpu_device_version(0) >= core.XPUVersion.XPU3,
+        "run test when xpu's compute capability < xpu3.",
+    )
+    def test_unsupported_bfloat16(self):
+        self.verify_trans_dtype(
+            test_type='bfloat16',
+            corrected_dtype=paddle.float32,
+        )
+
+    @unittest.skipIf(
+        not core.is_compiled_with_cuda() and not core.is_compiled_with_xpu(),
+        "Require compiled with CUDA or XPU.",
+    )
+    @unittest.skipIf(
+        core.is_compiled_with_cuda()
+        and paddle.device.cuda.get_device_capability()[0] < 8.0,
+        "run test when gpu's compute capability is at least 8.0.",
+    )
+    @unittest.skipIf(
+        core.is_compiled_with_xpu()
+        and core.get_xpu_device_version(0) < core.XPUVersion.XPU3,
+        "run test when xpu's compute capability >= xpu3.",
+    )
+    def test_supported_bfloat16(self):
+        self.verify_trans_dtype(
+            test_type='bfloat16',
+            corrected_dtype=paddle.bfloat16,
+        )
+
+    def test_float32(self):
+        paddle.set_default_dtype('float16')
+        self.verify_trans_dtype(
+            test_type='float32',
+            corrected_dtype=paddle.float32,
+        )
+        paddle.set_default_dtype('float32')
+
+    def test_excluded_layers_type_error(self):
+        self.assertRaises(
+            TypeError, self.verify_trans_dtype, excluded_layers=111
+        )
+
+
+class TestSupportedTypeInfo(unittest.TestCase):
+    def test_cpu(self):
+        res = paddle.amp.is_float16_supported('cpu')
+        self.assertEqual(res, False)
+        res = paddle.amp.is_bfloat16_supported('cpu')
+        self.assertEqual(res, core.supports_bfloat16())
+
+    @unittest.skipIf(
+        not core.is_compiled_with_cuda(), "Require compiled with CUDA."
+    )
+    def test_gpu_fp16_supported(self):
+        res = paddle.amp.is_float16_supported()
+        self.assertEqual(res, True)
+        res = paddle.amp.is_float16_supported('gpu')
+        self.assertEqual(res, True)
+        res = paddle.amp.is_float16_supported('gpu:0')
+        self.assertEqual(res, True)
+
+    @unittest.skipIf(
+        not core.is_compiled_with_cuda()
+        or paddle.device.cuda.get_device_capability()[0] >= 8.0,
+        "run test when maximum gpu's compute capability is 8.0.",
+    )
+    def test_gpu_bf16_unsupported(self):
+        res = paddle.amp.is_bfloat16_supported()
+        self.assertEqual(res, False)
+        res = paddle.amp.is_bfloat16_supported('gpu')
+        self.assertEqual(res, False)
+
+    @unittest.skipIf(
+        not core.is_compiled_with_cuda()
+        or paddle.device.cuda.get_device_capability()[0] < 8.0,
+        "run test when gpu's compute capability is at least 8.0.",
+    )
+    def test_gpu_bf16_supported(self):
+        res = paddle.amp.is_bfloat16_supported()
+        self.assertEqual(res, True)
+        res = paddle.amp.is_bfloat16_supported('gpu')
+        self.assertEqual(res, True)
+
+    def test_device_value_error(self):
+        self.assertRaises(
+            ValueError, paddle.amp.is_float16_supported, device='xxx'
+        )
+        self.assertRaises(
+            ValueError, paddle.amp.is_float16_supported, device=111
+        )
+
+    @unittest.skipIf(
+        not core.is_compiled_with_xpu()
+        or not core.get_xpu_device_version(0) >= core.XPUVersion.XPU2,
+        "run test when xpu's compute capability >= xpu2.",
+    )
+    def test_xpu_fp16_supported(self):
+        res = paddle.amp.is_float16_supported()
+        self.assertEqual(res, True)
+        res = paddle.amp.is_float16_supported('xpu')
+        self.assertEqual(res, True)
+        res = paddle.amp.is_float16_supported('xpu:0')
+        self.assertEqual(res, True)
+
+    @unittest.skipIf(
+        not core.is_compiled_with_xpu()
+        or core.get_xpu_device_version(0) >= core.XPUVersion.XPU3,
+        "run test when xpu's compute capability < xpu3.",
+    )
+    def test_xpu_bf16_unsupported(self):
+        res = paddle.amp.is_bfloat16_supported()
+        self.assertEqual(res, False)
+        res = paddle.amp.is_bfloat16_supported('xpu')
+        self.assertEqual(res, False)
+
+    @unittest.skipIf(
+        not core.is_compiled_with_xpu()
+        or not core.get_xpu_device_version(0) >= core.XPUVersion.XPU3,
+        "run test when xpu's compute capability >= xpu3.",
+    )
+    def test_xpu_bf16_supported(self):
+        res = paddle.amp.is_bfloat16_supported()
+        self.assertEqual(res, True)
+        res = paddle.amp.is_bfloat16_supported('xpu')
+        self.assertEqual(res, True)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/amp/test_pir_amp.py

@@ -0,0 +1,292 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.base import core
+
+
+class TestAmpAttrs(unittest.TestCase):
+    def test_pir_amp_attrs(self):
+        with paddle.pir_utils.IrGuard():
+            amp_attrs = core._get_amp_attrs()
+            amp_attrs._use_promote = True
+            amp_attrs._amp_level = core.AmpLevel.O2
+            amp_attrs._amp_dtype = 'float16'
+            np.testing.assert_equal(core._get_amp_attrs()._use_promote, True)
+            np.testing.assert_equal(
+                core._get_amp_attrs()._amp_level, core.AmpLevel.O2
+            )
+            np.testing.assert_equal(core._get_amp_attrs()._amp_dtype, 'float16')
+            amp_attrs._use_promote = False
+            amp_attrs._amp_level = core.AmpLevel.O0
+            amp_attrs._amp_dtype = 'float32'
+
+
+@unittest.skipIf(
+    not paddle.is_compiled_with_cuda()
+    or paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "only support device's compute capability is at least 7.0",
+)
+class TestPirAMPProgram(unittest.TestCase):
+    def test_linear_amp_o1(self):
+        if not core.is_compiled_with_cuda():
+            return
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data('x', [3, 4], 'float32')
+                linear = paddle.nn.Linear(4, 5)
+                with paddle.amp.auto_cast(
+                    level='O1', dtype='float16', use_promote=True
+                ):
+                    out1 = linear(x)
+                    out2 = paddle.mean(out1)
+
+            cast_op_count = 0
+            for op in main.global_block().ops:
+                if op.name() == 'pd_op.cast':
+                    cast_op_count += 1
+            # NOTE(Pan Zhaowu): After implementation of linear_v2, there's no
+            # need for mix-precision add op applies to intermediate result.
+            np.testing.assert_equal(out1.dtype, core.DataType.FLOAT16)
+            np.testing.assert_equal(out2.dtype, core.DataType.FLOAT32)
+            np.testing.assert_equal(cast_op_count, 4)
+            _white_list, _black_list = core._get_amp_op_list()
+            np.testing.assert_equal(len(_white_list), 0)
+            np.testing.assert_equal(len(_black_list), 0)
+
+    def test_linear_amp_bf16_o1(self):
+        if not core.is_compiled_with_cuda():
+            return
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data('x', [3, 4], 'float32')
+                linear = paddle.nn.Linear(4, 5)
+                with paddle.amp.auto_cast(
+                    level='O1', dtype='bfloat16', use_promote=True
+                ):
+                    out1 = linear(x)
+                    out2 = paddle.mean(out1)
+
+            cast_op_count = 0
+            for op in main.global_block().ops:
+                if op.name() == 'pd_op.cast':
+                    cast_op_count += 1
+            # NOTE(Pan Zhaowu): After implementation of linear_v2, there's no
+            # need for mix-precision add op applies to intermediate result.
+            np.testing.assert_equal(out1.dtype, core.DataType.BFLOAT16)
+            np.testing.assert_equal(out2.dtype, core.DataType.FLOAT32)
+            np.testing.assert_equal(cast_op_count, 4)
+            _white_list, _black_list = core._get_amp_op_list()
+            np.testing.assert_equal(len(_white_list), 0)
+            np.testing.assert_equal(len(_black_list), 0)
+
+    def test_linear_amp_o2_without_scaler(self):
+        if not core.is_compiled_with_cuda():
+            return
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data('x', [3, 4], 'float32')
+                linear = paddle.nn.Linear(4, 5)
+                optimizer = paddle.optimizer.Adam(
+                    learning_rate=0.001, parameters=linear.parameters()
+                )
+                linear, optimizer = paddle.amp.decorate(
+                    models=linear,
+                    optimizers=optimizer,
+                    level='O2',
+                    master_weight=True,
+                    master_grad=True,
+                )
+
+                with paddle.amp.auto_cast(
+                    level='O2', dtype='float16', use_promote=True
+                ):
+                    out = linear(x)
+                    loss = paddle.mean(out)
+                    optimizer.minimize(loss)
+            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)
+            place = paddle.CUDAPlace(0)
+            exe = paddle.static.Executor(place)
+            exe.run(startup)
+            result = exe.run(
+                main,
+                feed={'x': np.random.rand(3, 4).astype('float32')},
+                fetch_list=[loss],
+            )
+
+    def test_linear_amp_o2(self):
+        if not core.is_compiled_with_cuda():
+            return
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data('x', [3, 4], 'float32')
+                linear = paddle.nn.Linear(4, 5)
+                optimizer = paddle.optimizer.Adam(
+                    learning_rate=0.001, parameters=linear.parameters()
+                )
+                linear, optimizer = paddle.amp.decorate(
+                    models=linear,
+                    optimizers=optimizer,
+                    level='O2',
+                    master_weight=True,
+                    master_grad=True,
+                )
+                scaler = paddle.amp.GradScaler(
+                    init_loss_scaling=2.0**16, use_dynamic_loss_scaling=True
+                )
+
+                with paddle.amp.auto_cast(
+                    level='O2', dtype='float16', use_promote=True
+                ):
+                    out = linear(x)
+                    loss = paddle.mean(out)
+                scaled = scaler.scale(loss)
+                opt_ops, _ = scaler.minimize(
+                    optimizer, scaled, startup_program=startup
+                )
+            np.testing.assert_equal(len(opt_ops), 8)
+            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, 5)
+            place = paddle.CUDAPlace(0)
+            exe = paddle.static.Executor(place)
+            exe.run(startup)
+            result = exe.run(
+                main,
+                feed={'x': np.random.rand(3, 4).astype('float32')},
+                fetch_list=[loss],
+            )
+
+    def test_linear_amp_bf16_o2_without_scaler(self):
+        if not core.is_compiled_with_cuda():
+            return
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data('x', [3, 4], 'float32')
+                linear = paddle.nn.Linear(4, 5)
+                optimizer = paddle.optimizer.Adam(
+                    learning_rate=0.001, parameters=linear.parameters()
+                )
+                linear, optimizer = paddle.amp.decorate(
+                    models=linear,
+                    optimizers=optimizer,
+                    level='O2',
+                    dtype='bfloat16',
+                    master_weight=True,
+                    master_grad=True,
+                )
+
+                with paddle.amp.auto_cast(
+                    level='O2', dtype='bfloat16', use_promote=True
+                ):
+                    out = linear(x)
+                    loss = paddle.mean(out)
+                    optimizer.minimize(loss)
+            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)
+            place = paddle.CUDAPlace(0)
+            exe = paddle.static.Executor(place)
+            exe.run(startup)
+            result = exe.run(
+                main,
+                feed={'x': np.random.rand(3, 4).astype('float32')},
+                fetch_list=[loss],
+            )
+
+
+class Net(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.linear = paddle.nn.Linear(2, 2)
+
+    def forward(self, x):
+        out1 = self.linear(x)
+        out2 = self.linear.weight + 1
+        return out1 + out2
+
+
+@unittest.skipIf(
+    not paddle.is_compiled_with_cuda()
+    or paddle.device.cuda.get_device_capability()[0] < 7.0,
+    "only support device's compute capability is at least 7.0",
+)
+class TestPirAMPMasterGrad(unittest.TestCase):
+    def test_multi_param_grad(self):
+        with paddle.pir_utils.IrGuard():
+            startup = paddle.static.Program()
+            main = paddle.static.Program()
+            with paddle.static.program_guard(main, startup):
+                x = paddle.static.data('x', [2, 2])
+                net = Net()
+                opt = paddle.optimizer.Adam(
+                    learning_rate=0.0001, parameters=net.parameters()
+                )
+                linear, opt = paddle.amp.decorate(
+                    models=net,
+                    optimizers=opt,
+                    level='O2',
+                    dtype='float16',
+                    master_weight=False,
+                    master_grad=True,
+                )
+                with paddle.amp.auto_cast(level='O2', dtype='float16'):
+                    out = net(x)
+                    loss = paddle.mean(out)
+                opt.minimize(loss)
+
+                place = paddle.CUDAPlace(0)
+                exe = paddle.static.Executor(place)
+                exe.run(startup)
+                result = exe.run(
+                    main,
+                    feed={'x': np.random.rand(2, 2).astype('float32')},
+                    fetch_list=[loss],
+                )
+                for op in main.global_block().ops:
+                    if op.name() == 'builtin.combine':
+                        for input in [
+                            op.operand_source(0),
+                            op.operand_source(1),
+                        ]:
+                            np.testing.assert_equal(input.dtype, paddle.float32)
+                            np.testing.assert_equal(
+                                input.get_defining_op().name(), 'pd_op.cast'
+                            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/amp/test_tensor_checker.py

@@ -0,0 +1,181 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import unittest
+
+import paddle
+
+
+class TestTensorChecker(unittest.TestCase):
+    def _parse_num_nan_inf(self, e):
+        num_nan = 0
+        num_inf = 0
+        # Cannot catch the log in CUDA kernel.
+        err_str_list = (
+            str(e)
+            .replace("(", " ")
+            .replace(")", " ")
+            .replace(",", " ")
+            .split(" ")
+        )
+        for err_str in err_str_list:
+            if "num_nan" in err_str:
+                num_nan = int(err_str.split("=")[1])
+            elif "num_inf" in err_str:
+                num_inf = int(err_str.split("=")[1])
+        return num_nan, num_inf
+
+    def _generate_num_inf(self, place):
+        num_inf = 0
+        num_nan = 0
+        paddle.set_device(place)
+        # check op list
+        x = paddle.to_tensor(
+            [1, 0, 0],
+            dtype='float32',
+            stop_gradient=False,
+        )
+        y = paddle.to_tensor([0, 0, 1], dtype='float32')
+        try:
+            res = paddle.pow(x, y)
+            # test backward
+            paddle.autograd.backward([res])
+            res = paddle.divide(y, x)
+        except Exception as e:
+            num_nan, num_inf = self._parse_num_nan_inf(e)
+        return num_nan, num_inf
+
+    def test_tensor_checker(self):
+        def _assert_flag(value):
+            flags = ['FLAGS_check_nan_inf', 'FLAGS_check_nan_inf_level']
+            res = paddle.get_flags(flags)
+            assert res["FLAGS_check_nan_inf"] == value
+
+        paddle.set_flags({"FLAGS_check_nan_inf": 0})
+        paddle.seed(102)
+        checker_config = paddle.amp.debugging.TensorCheckerConfig(
+            enable=True,
+            debug_mode=paddle.amp.debugging.DebugMode.CHECK_NAN_INF_AND_ABORT,
+            checked_op_list=["elementwise_pow_grad"],
+            skipped_op_list=["elementwise_div"],
+            debug_step=[0, 3],
+        )
+        places = []
+        if (
+            os.environ.get('FLAGS_CI_both_cpu_and_gpu', 'False').lower()
+            in ['1', 'true', 'on']
+            or not paddle.is_compiled_with_cuda()
+        ):
+            places.append('cpu')
+        if paddle.is_compiled_with_cuda():
+            places.append('gpu')
+        # check seed
+        self.assertEqual(checker_config.initial_seed, 102)
+        self.assertEqual(checker_config.seed, 102)
+        _assert_flag(False)
+
+        for place in places:
+            paddle.amp.debugging.TensorCheckerConfig.current_step_id = 0
+            for iter_id in range(5):
+                paddle.amp.debugging.enable_tensor_checker(checker_config)
+                if iter_id <= 2:
+                    _assert_flag(True)
+                    self.assertEqual(
+                        iter_id + 1,
+                        paddle.amp.debugging.TensorCheckerConfig.current_step_id,
+                    )
+                    num_nan, num_inf = self._generate_num_inf(place)
+                    print(
+                        f"-- [iter_id={iter_id}, place={place}] num_nan={num_nan}, num_inf={num_inf}"
+                    )
+                    self.assertEqual(
+                        0,
+                        num_nan,
+                        f"Expected num_nan to be 0, but received {num_nan}, place={place}.",
+                    )
+                else:
+                    self.assertEqual(
+                        3,
+                        paddle.amp.debugging.TensorCheckerConfig.current_step_id,
+                    )
+                    _assert_flag(False)
+                    num_nan, num_inf = self._generate_num_inf(place)
+                    print(
+                        f"-- [iter_id={iter_id}, place={place}] num_nan={num_nan}, num_inf={num_inf}"
+                    )
+                    self.assertEqual(
+                        0,
+                        num_nan,
+                        f"Expected num_nan to be 1, but received {num_nan}, place={place}.",
+                    )
+
+                paddle.amp.debugging.disable_tensor_checker()
+                _assert_flag(False)
+
+
+class TestCheckLayerNumerics(unittest.TestCase):
+    def test_layer_checker(self):
+        class MyLayer(paddle.nn.Layer):
+            def __init__(self, dtype):
+                super().__init__()
+                self._w = self.create_parameter([2, 3], dtype=dtype)
+                self._b = self.create_parameter([2, 3], dtype=dtype)
+
+            @paddle.amp.debugging.check_layer_numerics
+            def forward(self, x):
+                return x * self._w + self._b
+
+        dtype = 'float32'
+        x = paddle.rand([10, 2, 3], dtype=dtype)
+        model = MyLayer(dtype)
+        loss = model(x)
+        adam = paddle.optimizer.Adam(parameters=model.parameters())
+        loss.backward()
+        adam.step()
+
+    def test_error_no_element(self):
+        class MyLayer(paddle.nn.Layer):
+            def __init__(self, dtype):
+                super().__init__()
+                self._w = self.create_parameter([2, 3], dtype=dtype)
+
+            @paddle.amp.debugging.check_layer_numerics
+            def forward(self):
+                return self._w
+
+        with self.assertRaises(RuntimeError):
+            dtype = 'float32'
+            model = MyLayer(dtype)
+            data = model()
+
+    def test_error_type_error(self):
+        class MyLayer(paddle.nn.Layer):
+            def __init__(self, dtype):
+                super().__init__()
+                self._w = self.create_parameter([2, 3], dtype=dtype)
+
+            @paddle.amp.debugging.check_layer_numerics
+            def forward(self, x):
+                return self._w * x
+
+        x = 1
+        with self.assertRaises(RuntimeError):
+            dtype = 'float32'
+            model = MyLayer(dtype)
+            data = model(x)
+
+
+if __name__ == '__main__':
+    unittest.main()