[Pir-AutoParallel] Update Demo Unitest (#62630)

JZ-LIANG · web-flow · commit 9d593d0090bf · 2024-03-14T10:25:56.000+08:00
* update test

* update test

* hack for clone
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -738,12 +738,16 @@ def _build(self, mode):
         # TODO(zhiqiu): distributed_context is no longer used in pir_program
         # so, just return here and need to reimplement the logics below
         if self._in_pir_mode:
-            if mode != "train":
-                self._fwd_main_progs[mode] = serial_main_prog.clone(
-                    for_test=True
-                )
-            else:
-                self._fwd_main_progs[mode] = serial_main_prog
+            # TODO(ljz): pir not support clone_for_test,
+            # so we need to update the method to create eval/test program in engine.
+
+            # if mode != "train":
+            #     self._fwd_main_progs[mode] = serial_main_prog.clone(
+            #         for_test=True
+            #     )
+            # else:
+
+            self._fwd_main_progs[mode] = serial_main_prog
             return
 
         default_ctx = get_default_distributed_context()
diff --git a/test/auto_parallel/pir/test_to_static_pir_program.py b/test/auto_parallel/pir/test_to_static_pir_program.py
@@ -46,9 +46,11 @@ class DemoNet(nn.Layer):
     def __init__(self, mesh):
         super().__init__()
         self._mesh = mesh
-        self.linear_0 = nn.Linear(IMAGE_SIZE, IMAGE_SIZE)
-        self.linear_1 = nn.Linear(IMAGE_SIZE, CLASS_NUM)
-        self.relu = nn.ReLU()
+        self.linear_0 = nn.Linear(IMAGE_SIZE, IMAGE_SIZE, bias_attr=False)
+        self.linear_1 = nn.Linear(IMAGE_SIZE, CLASS_NUM, bias_attr=False)
+        self.relu_0 = nn.ReLU()
+        self.relu_1 = nn.ReLU()
+        self.relu_2 = nn.ReLU()
         # shard the weights of this layer
         self.linear_0.weight = dist.shard_tensor(
             self.linear_0.weight,
@@ -64,9 +66,11 @@ def __init__(self, mesh):
         )
 
     def forward(self, x):
-        out = self.linear_0(x)
-        out = self.relu(out)
+        out = self.relu_0(x)  # triggle backward partial allreduce
+        out = self.linear_0(out)
+        out = self.relu_1(out)
         out = self.linear_1(out)
+        out = self.relu_2(out)  # triggle forward partial allreduce
         return out
 
 
@@ -78,7 +82,33 @@ def create_data_loader():
     return loader
 
 
-class TestToStaticPirProgram(unittest.TestCase):
+class TestToStaticPirProgramEval(unittest.TestCase):
+    def test_to_static_program(self):
+        paddle.base.set_flags({'FLAGS_enable_pir_api': 1})
+        mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+        layer = DemoNet(mesh)
+        opt = None  # forward only
+        loss_fn = nn.MSELoss()
+        loader = create_data_loader()
+        dist_loader = dist.shard_dataloader(loader, meshes=[mesh])
+        dist_model = dist.to_static(layer, dist_loader, loss_fn, opt)
+
+        dist_model.eval()
+        main_program = dist_model._engine._fwd_main_progs["eval"]
+
+        for op in main_program.global_block().ops:
+            tensor = op.result(0)
+            if op.name() == 'pd_op.data':
+                self.assertTrue(tensor.is_dist_dense_tensor_type())
+                self.assertEqual(tensor.process_mesh.shape, [2])
+                self.assertEqual(tensor.process_mesh.process_ids, [0, 1])
+                self.assertEqual(tensor.dims_mapping, [-1, -1])
+                self.assertEqual(tensor.partial_dims, set())
+            elif op.name() == "builtin.parameter":
+                pass  # TODO check
+
+
+class TestToStaticPirProgramTrain(unittest.TestCase):
     def test_to_static_program(self):
         paddle.base.set_flags({'FLAGS_enable_pir_api': 1})
         mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
@@ -93,7 +123,7 @@ def test_to_static_program(self):
 
         dist_model.train()
         main_program = dist_model._engine._fwd_main_progs["train"]
-        print(main_program, flush=1)
+
         for op in main_program.global_block().ops:
             tensor = op.result(0)
             if op.name() == 'pd_op.data':
@@ -119,7 +149,6 @@ def test_to_static_program(self):
                         self.assertEqual(tensor.dims_mapping, [0, -1])
                     self.assertEqual(tensor.partial_dims, set())
 
-        # training
         # dist_model.train()
         # for batch_id, (image, label) in enumerate(dist_loader()):
         #     loss = dist_model(image, label)
