add_test

Author: xuxinyi389

test/auto_parallel/CMakeLists.txt

@@ -45,6 +45,7 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
     set_tests_properties(test_pipeline_scheduler
                          PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 400)
   endif()
+  py_test_modules(test_process_mesh MODULES test_process_mesh)
   py_test_modules(test_reshard_r_to_p MODULES test_reshard_r_to_p)
   set_tests_properties(test_reshard_r_to_p
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 200)

test/auto_parallel/process_mesh_demo_unittest.py

@@ -0,0 +1,128 @@
+# 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 paddle
+import paddle.distributed as dist
+
+
+class TestProcessMesh:
+    def init_dist_env(self):
+        dist.init_parallel_env()
+        paddle.seed(2025)
+
+    def test_get_submesh_with_dim(self):
+        curr_rank = dist.get_rank()
+
+        # Test 2D mesh
+        mesh_2d = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["dp", "tp"])
+
+        # Test case 1: Get submesh for dp dimension
+        dp_mesh = mesh_2d.get_submesh_with_dim("dp")
+        if curr_rank == 0:
+            assert dp_mesh.process_ids == [0, 2]
+        elif curr_rank == 1:
+            assert dp_mesh.process_ids == [1, 3]
+
+        # Test case 2: Get submesh for tp dimension
+        tp_mesh = mesh_2d.get_submesh_with_dim("tp")
+        if curr_rank == 0:
+            assert tp_mesh.process_ids == [0, 1]
+        elif curr_rank == 1:
+            assert tp_mesh.process_ids == [0, 1]
+
+        # Test case 3: 3D mesh with 8 cards (2x2x2)
+        mesh_3d = dist.ProcessMesh(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=["pp", "dp", "tp"]
+        )
+
+        # Test each dimension
+        pp_mesh = mesh_3d.get_submesh_with_dim("pp")
+        dp_mesh = mesh_3d.get_submesh_with_dim("dp")
+        tp_mesh = mesh_3d.get_submesh_with_dim("tp")
+
+        # Verify pp dimension results
+        if curr_rank == 0:
+            assert pp_mesh.process_ids == [0, 4]
+        elif curr_rank == 1:
+            assert pp_mesh.process_ids == [1, 5]
+
+        # Verify dp dimension results
+        if curr_rank == 0:
+            assert dp_mesh.process_ids == [0, 2]
+        elif curr_rank == 1:
+            assert dp_mesh.process_ids == [1, 3]
+
+        # Verify tp dimension results
+        if curr_rank == 0:
+            assert tp_mesh.process_ids == [0, 1]
+        elif curr_rank == 1:
+            assert tp_mesh.process_ids == [0, 1]
+
+        # Test case 4: When rank is not in the mesh
+        mesh_small = dist.ProcessMesh([0, 1], dim_names=["x"])
+        if curr_rank not in [0, 1]:
+            assert mesh_small.get_submesh_with_dim("x") is None
+
+    def test_get_group(self):
+        curr_rank = dist.get_rank()
+
+        # Test case 1: Single dimension mesh without dim_name
+        mesh_1d = dist.ProcessMesh([0, 1], dim_names=["x"])
+        if curr_rank in [0, 1]:
+            group_1d = mesh_1d.get_group()
+            assert isinstance(group_1d, dist.communication.group.Group)
+
+            # Test case 2: Single dimension mesh with correct dim_name
+            group_1d_with_name = mesh_1d.get_group(dim_name="x")
+            assert isinstance(
+                group_1d_with_name, dist.communication.group.Group
+            )
+
+            # Test case 3: Single dimension mesh with wrong dim_name
+            try:
+                mesh_1d.get_group(dim_name="wrong_name")
+                raise AssertionError("Should raise ValueError")
+            except ValueError:
+                pass
+
+        # Test case 4: Multi-dimension mesh
+        mesh_2d = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["dp", "tp"])
+        if curr_rank in [0, 1, 2, 3]:
+            # Test without dim_name
+            try:
+                mesh_2d.get_group()
+                raise AssertionError("Should raise ValueError")
+            except ValueError:
+                pass
+
+            # Test with correct dim_name
+            group_2d = mesh_2d.get_group(dim_name="dp")
+            assert isinstance(group_2d, dist.communication.group.Group)
+
+            # Test with wrong dim_name
+            try:
+                mesh_2d.get_group(dim_name="wrong_name")
+                raise AssertionError("Should raise ValueError")
+            except ValueError:
+                pass
+
+    def test_process_mesh(self):
+        self.init_dist_env()
+        self.test_get_submesh_with_dim()
+        self.test_get_group()
+
+
+if __name__ == '__main__':
+    TestProcessMesh().test_process_mesh()

test/auto_parallel/test_process_mesh.py

@@ -0,0 +1,46 @@
+# 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 collective.test_communication_api_base as test_base
+
+
+class TestProcessMeshPass(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=2,
+            timeout=50,
+        )
+        self._default_envs = {
+            "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+        }
+
+    def test_process_mesh(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "process_mesh_demo_unittest.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()