【PIR API adaptor No.143、144】 Migrate margin_cross_entropy、masked_multihead_attention (PaddlePaddle#58762)

Author: Fancy-hjyp

python/paddle/incubate/nn/functional/masked_multihead_attention.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from paddle import _C_ops
-from paddle.framework import LayerHelper, in_dynamic_mode
+from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
 
 
 def masked_multihead_attention(
@@ -90,7 +90,7 @@ def masked_multihead_attention(
 
     """
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.masked_multihead_attention_(
             x,
             cache_kv,

python/paddle/nn/functional/loss.py

@@ -2314,7 +2314,7 @@ def margin_cross_entropy(
     if input_dims - 1 == label_dims:
         label = paddle.unsqueeze(label, axis=-1)
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         softmax, loss = _C_ops.margin_cross_entropy(
             logits,
             label,

test/legacy_test/test_margin_cross_entropy_op.py

@@ -18,7 +18,9 @@
 from op_test import OpTest, convert_float_to_uint16, paddle_static_guard
 
 import paddle
-from paddle.base import Program, core, program_guard
+from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
+from paddle.static import Program, program_guard
 
 
 def stable_softmax_comm(x):
@@ -148,10 +150,14 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output_with_place(core.CUDAPlace(0), atol=1e-5)
+        self.check_output_with_place(
+            core.CUDAPlace(0), atol=1e-5, check_pir=True
+        )
 
     def test_check_grad(self):
-        self.check_grad_with_place(core.CUDAPlace(0), ["Logits"], "Loss")
+        self.check_grad_with_place(
+            core.CUDAPlace(0), ["Logits"], "Loss", check_pir=True
+        )
 
 
 @unittest.skipIf(
@@ -168,6 +174,7 @@ def test_check_grad(self):
             "Loss",
             numeric_grad_delta=5e-2,
             max_relative_error=5e-2,
+            check_pir=True,
         )
 
 
@@ -179,7 +186,9 @@ def init_dtype(self):
         self.dtype = np.float16
 
     def test_check_output(self):
-        self.check_output_with_place(core.CUDAPlace(0), atol=5e-2)
+        self.check_output_with_place(
+            core.CUDAPlace(0), atol=5e-2, check_pir=True
+        )
 
     def test_check_grad(self):
         self.check_grad_with_place(
@@ -188,6 +197,7 @@ def test_check_grad(self):
             "Loss",
             numeric_grad_delta=6e-1,
             max_relative_error=6e-1,
+            check_pir=True,
         )
 
 
@@ -264,7 +274,9 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output_with_place(core.CUDAPlace(0), atol=5e-2)
+        self.check_output_with_place(
+            core.CUDAPlace(0), atol=5e-2, check_pir=True
+        )
 
     def test_check_grad(self):
         self.check_grad_with_place(
@@ -273,6 +285,7 @@ def test_check_grad(self):
             "Loss",
             numeric_grad_delta=6e-1,
             max_relative_error=6e-1,
+            check_pir=True,
         )
 
 
@@ -301,13 +314,17 @@ def init_loss_params(self):
 class TestMarginCrossEntropyOpCPU(TestMarginCrossEntropyOp):
     def test_check_output(self):
         try:
-            self.check_output_with_place(core.CPUPlace(), atol=1e-5)
+            self.check_output_with_place(
+                core.CPUPlace(), atol=1e-5, check_pir=True
+            )
         except RuntimeError:
             pass
 
     def test_check_grad(self):
         try:
-            self.check_grad_with_place(core.CPUPlace(), ["Logits"], "Loss")
+            self.check_grad_with_place(
+                core.CPUPlace(), ["Logits"], "Loss", check_pir=True
+            )
         except RuntimeError:
             pass
 
@@ -347,6 +364,7 @@ def init_dtype(self):
     def init_reduction(self):
         self.reduction = None
 
+    @test_with_pir_api
     def test_static(self):
         for place in self.places:
             self.check_static_result(place=place)
@@ -404,7 +422,7 @@ def check_static_result(self, place):
 
                 exe = paddle.base.Executor(place)
                 [loss_res, softmax_res] = exe.run(
-                    paddle.base.default_main_program(),
+                    paddle.static.default_main_program(),
                     feed={'logits': logits_np, 'label': labels_np},
                     fetch_list=[loss, softmax],
                 )

test/legacy_test/test_masked_multihead_attention_op.py

@@ -19,6 +19,7 @@
 import paddle
 from paddle.framework import core
 from paddle.incubate.nn.functional import masked_multihead_attention
+from paddle.pir_utils import test_with_pir_api
 
 
 @unittest.skipIf(
@@ -213,6 +214,7 @@ def check_main(
         paddle.enable_static()
         return paddle_naive_mmha_out, paddle_mmha_out
 
+    @test_with_pir_api
     def test_mmha_fp16(self):
         if not paddle.is_compiled_with_cuda():
             return
@@ -234,6 +236,7 @@ def test_mmha_fp16(self):
             atol=1e-3,
         )
 
+    @test_with_pir_api
     def test_mmha_qkv_out_scale(self):
         if not paddle.is_compiled_with_cuda():
             return
@@ -255,6 +258,7 @@ def test_mmha_qkv_out_scale(self):
             atol=1e-3,
         )
 
+    @test_with_pir_api
     def test_mmha_outlinear_in_scale(self):
         if not paddle.is_compiled_with_cuda():
             return
@@ -463,11 +467,12 @@ def check_main(
                     "bias_static": bias.astype(dtype),
                     "src_mask_static": src_mask.astype(dtype),
                 },
-                fetch_list=[outs],
+                fetch_list=[outs[0], outs[1]],
             )
 
         return paddle_naive_mmha_out, out_s
 
+    @test_with_pir_api
     def test_mmha_fp16(self):
         if not paddle.is_compiled_with_cuda():
             return