[Fix PIR Unittest] fix pir ut in test/legacy_test/(test_nadam_op,test_asgd_op).py (#66709)

ooooo-create · web-flow · commit 1f3abbcbf1e5 · 2024-08-01T15:26:55.000+08:00
diff --git a/test/legacy_test/test_asgd_op.py b/test/legacy_test/test_asgd_op.py
@@ -252,49 +252,107 @@ def dygraph_asgd_mp(self, mp):
 
     def static_asgd_mp(self, mp):
         paddle.enable_static()
-        paddle.seed(10)
-        np.random.seed(10)
-        exe = paddle.static.Executor('gpu')
-        train_program = paddle.static.Program()
-        startup_program = paddle.static.Program()
-        optimizer = paddle.optimizer.ASGD(batch_num=2, multi_precision=mp)
+        with paddle.pir_utils.OldIrGuard():
+            paddle.seed(10)
+            np.random.seed(10)
+            exe = paddle.static.Executor('gpu')
+            train_program = paddle.static.Program()
+            startup_program = paddle.static.Program()
+            optimizer = paddle.optimizer.ASGD(batch_num=2, multi_precision=mp)
 
-        if mp:
-            optimizer = paddle.static.amp.decorate(
-                optimizer,
-                init_loss_scaling=128.0,
-                use_dynamic_loss_scaling=True,
-                use_pure_fp16=True,
-                use_fp16_guard=False,
-            )
-        with paddle.static.program_guard(train_program, startup_program):
             if mp:
-                data = paddle.static.data(
-                    shape=[2, 2], name='X', dtype='float16'
+                optimizer = paddle.static.amp.decorate(
+                    optimizer,
+                    init_loss_scaling=128.0,
+                    use_dynamic_loss_scaling=True,
+                    use_pure_fp16=True,
+                    use_fp16_guard=False,
                 )
+            with paddle.static.program_guard(train_program, startup_program):
+                if mp:
+                    data = paddle.static.data(
+                        shape=[2, 2], name='X', dtype='float16'
+                    )
+                else:
+                    data = paddle.static.data(
+                        shape=[2, 2], name='X', dtype='float32'
+                    )
+                hidden = paddle.static.nn.fc(x=data, size=10)
+                loss = paddle.mean(hidden)
+                optimizer.minimize(loss)
+            exe.run(startup_program)
+
+            if mp:
+                optimizer.amp_init(
+                    place=paddle.CUDAPlace(0),
+                    scope=paddle.static.global_scope(),
+                )
+                x = np.random.random(size=(2, 2)).astype('float16')
             else:
-                data = paddle.static.data(
-                    shape=[2, 2], name='X', dtype='float32'
+                x = np.random.random(size=(2, 2)).astype('float32')
+            out = []
+            for idx in range(5):
+                (loss_data,) = exe.run(
+                    train_program, feed={"X": x}, fetch_list=[loss.name]
                 )
-            hidden = paddle.static.nn.fc(x=data, size=10)
-            loss = paddle.mean(hidden)
-            optimizer.minimize(loss)
-        exe.run(startup_program)
+                out.append(loss_data)
+            return out
 
-        if mp:
-            optimizer.amp_init(
-                place=paddle.CUDAPlace(0), scope=paddle.static.global_scope()
-            )
-            x = np.random.random(size=(2, 2)).astype('float16')
-        else:
-            x = np.random.random(size=(2, 2)).astype('float32')
-        out = []
-        for idx in range(5):
-            (loss_data,) = exe.run(
-                train_program, feed={"X": x}, fetch_list=[loss.name]
-            )
-            out.append(loss_data)
-        return out
+    def pir_asgd_mp(self, mp):
+        paddle.enable_static()
+        with paddle.pir_utils.IrGuard():
+            paddle.seed(10)
+            np.random.seed(10)
+            exe = paddle.static.Executor('gpu')
+            train_program = paddle.static.Program()
+            startup_program = paddle.static.Program()
+
+            with paddle.static.program_guard(train_program, startup_program):
+                model = paddle.nn.Linear(2, 10)
+                optimizer = paddle.optimizer.ASGD(
+                    batch_num=2,
+                    multi_precision=mp,
+                    parameters=model.parameters(),
+                )
+                if mp:
+                    model, optimizer = paddle.amp.decorate(
+                        models=model,
+                        optimizers=optimizer,
+                        level='O2',
+                    )
+                    scaler = paddle.amp.GradScaler(
+                        init_loss_scaling=1024, use_dynamic_loss_scaling=True
+                    )
+                    data = paddle.static.data(
+                        shape=[2, 2], name='X', dtype='float16'
+                    )
+                    with paddle.amp.auto_cast(
+                        level='O2', dtype='float16', use_promote=True
+                    ):
+                        output = model(data)
+                        loss = paddle.mean(output)
+                    scaled = scaler.scale(loss)
+                    scaler.minimize(optimizer, scaled)
+                else:
+                    data = paddle.static.data(
+                        shape=[2, 2], name='X', dtype='float32'
+                    )
+                    output = model(data)
+                    loss = paddle.mean(output)
+                    optimizer.minimize(loss)
+            exe.run(startup_program)
+
+            if mp:
+                x = np.random.random(size=(2, 2)).astype('float16')
+            else:
+                x = np.random.random(size=(2, 2)).astype('float32')
+            out = []
+            for idx in range(5):
+                (loss_data,) = exe.run(
+                    train_program, feed={"X": x}, fetch_list=[loss]
+                )
+                out.append(loss_data)
+            return out
 
     def test_main(self):
         if not paddle.is_compiled_with_cuda():
@@ -325,6 +383,16 @@ def test_main(self):
                 rtol=1e-05,
                 atol=0.1,
             )
+        "Test pir graph mode"
+        output1_pir = self.pir_asgd_mp(mp=True)
+        output2_pir = self.pir_asgd_mp(mp=False)
+        for idx in range(len(output1_st)):
+            np.testing.assert_allclose(
+                output1_pir[idx].astype('float32'),
+                output2_pir[idx].astype('float32'),
+                rtol=1e-05,
+                atol=0.1,
+            )
 
 
 class TestASGDSimple(unittest.TestCase):
diff --git a/test/legacy_test/test_nadam_op.py b/test/legacy_test/test_nadam_op.py
@@ -543,6 +543,60 @@ def static_nadam_mp(self, mp, use_amp):
 
             return out
 
+    def pir_nadam_mp(self, mp, use_amp):
+        paddle.enable_static()
+        with paddle.pir_utils.IrGuard():
+            paddle.seed(2024)
+            exe = paddle.static.Executor('gpu')
+            train_program = paddle.static.Program()
+            startup_program = paddle.static.Program()
+
+            with paddle.static.program_guard(train_program, startup_program):
+                model = paddle.nn.Linear(2, 10)
+                optimizer = paddle.optimizer.NAdam(
+                    0.1, parameters=model.parameters()
+                )
+                if use_amp:
+                    data = paddle.static.data(
+                        shape=[2, 2], name='X', dtype='float16'
+                    )
+                    model, optimizer = paddle.amp.decorate(
+                        models=model,
+                        optimizers=optimizer,
+                        level='O2',
+                        master_weight=mp,
+                    )
+                    scaler = paddle.amp.GradScaler(init_loss_scaling=128.0)
+                    with paddle.amp.auto_cast(
+                        level='O2', dtype="float16", use_promote=True
+                    ):
+                        output = model(data)
+                        loss = paddle.mean(output)
+                    scaled = scaler.scale(loss)
+                    scaler.minimize(optimizer, scaled)
+                else:
+                    data = paddle.static.data(
+                        shape=[2, 2], name='X', dtype='float32'
+                    )
+                    output = model(data)
+                    loss = paddle.mean(output)
+                    optimizer.minimize(loss)
+            exe.run(startup_program)
+
+            np.random.seed(2024)
+            if use_amp:
+                x = np.random.random(size=(2, 2)).astype('float16')
+            else:
+                x = np.random.random(size=(2, 2)).astype('float32')
+            out = []
+            for idx in range(5):
+                (loss_data,) = exe.run(
+                    train_program, feed={"X": x}, fetch_list=[loss]
+                )
+                out.append(loss_data)
+
+            return out
+
     def static_nadam_amp_o2_without_scaler(self):
         paddle.enable_static()
         paddle.seed(2024)
@@ -613,6 +667,16 @@ def test_main(self):
                 rtol=1e-05,
                 atol=0.1,
             )
+        "Test pir mode"
+        output1_pir = self.pir_nadam_mp(use_amp=True, mp=True)
+        output2_pir = self.pir_nadam_mp(use_amp=False, mp=False)
+        for idx in range(len(output1_st)):
+            np.testing.assert_allclose(
+                output1_pir[idx].astype('float32'),
+                output2_pir[idx].astype('float32'),
+                rtol=1e-05,
+                atol=0.1,
+            )
 
         for idx in range(len(output1_st)):
             np.testing.assert_allclose(
