chore: import upstream snapshot with attribution

test/auto_parallel/CMakeLists.txt

@@ -0,0 +1,226 @@
+# file(GLOB TEST_OPS RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" "test_*.py")
+# string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+add_subdirectory(spmd_rules)
+add_subdirectory(hybrid_strategy)
+add_subdirectory(custom_op)
+add_subdirectory(pir)
+add_subdirectory(end_to_end)
+
+if(WITH_DISTRIBUTE AND WITH_GPU)
+
+  # NOTE(zyl): unittests WITH multi cards and timeout
+  py_test_modules(test_converter MODULES test_converter)
+  set_tests_properties(test_converter PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE"
+                                                 TIMEOUT 50)
+  py_test_modules(test_high_order_grad MODULES test_high_order_grad)
+  set_tests_properties(test_high_order_grad
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 50)
+  py_test_modules(test_iterable_dataset MODULES test_iterable_dataset)
+  set_tests_properties(test_iterable_dataset
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 80)
+
+  py_test_modules(test_reshard_api MODULES test_reshard_api)
+  set_tests_properties(test_reshard_api PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE"
+                                                   TIMEOUT 150)
+  py_test_modules(test_reshard_s_to_p MODULES test_reshard_s_to_p)
+  set_tests_properties(test_reshard_s_to_p
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_reshard_p_to_s MODULES test_reshard_p_to_s)
+  set_tests_properties(test_reshard_p_to_s
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_reshard_s_to_s MODULES test_reshard_s_to_s)
+  set_tests_properties(test_reshard_s_to_s
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_reshard_r_to_s MODULES test_reshard_r_to_s)
+  set_tests_properties(test_reshard_r_to_s
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 320)
+  py_test_modules(test_reshard_p_to_r MODULES test_reshard_p_to_r)
+  set_tests_properties(test_reshard_p_to_r
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 160)
+  py_test_modules(test_reshard_s_to_r MODULES test_reshard_s_to_r)
+  set_tests_properties(test_reshard_s_to_r
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 150)
+  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)
+  py_test_modules(test_reshard_x_to_r MODULES test_reshard_x_to_r)
+  set_tests_properties(test_reshard_x_to_r
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_reshard_r_to_x MODULES test_reshard_r_to_x)
+  set_tests_properties(test_reshard_r_to_x
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_reshard_nd_mesh MODULES test_reshard_nd_mesh)
+  set_tests_properties(test_reshard_nd_mesh
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+
+  py_test_modules(test_reshard_same_status MODULES test_reshard_same_status)
+  set_tests_properties(test_reshard_same_status
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+
+  py_test_modules(test_semi_auto_parallel_basic MODULES
+                  test_semi_auto_parallel_basic)
+  set_tests_properties(test_semi_auto_parallel_basic
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 800)
+
+  py_test_modules(test_semi_auto_parallel_for_llama_subnet MODULES
+                  test_semi_auto_parallel_for_llama_subnet)
+  set_tests_properties(test_semi_auto_parallel_for_llama_subnet
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 300)
+
+  py_test_modules(test_semi_auto_parallel_softmax_basic MODULES
+                  test_semi_auto_parallel_softmax_basic)
+  set_tests_properties(test_semi_auto_parallel_softmax_basic
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+
+  py_test_modules(test_semi_auto_parallel_compare_basic MODULES
+                  test_semi_auto_parallel_compare_basic)
+  set_tests_properties(test_semi_auto_parallel_compare_basic
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+  py_test_modules(test_semi_auto_parallel_subgraph_embedding_basic MODULES
+                  test_semi_auto_parallel_subgraph_embedding_basic)
+  set_tests_properties(test_semi_auto_parallel_subgraph_embedding_basic
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+  py_test_modules(test_semi_auto_parallel_pylayer MODULES
+                  test_semi_auto_parallel_pylayer)
+  set_tests_properties(test_semi_auto_parallel_pylayer
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_semi_auto_parallel_single_strategy MODULES
+                  test_semi_auto_parallel_single_strategy)
+  set_tests_properties(test_semi_auto_parallel_single_strategy
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 400)
+  py_test_modules(test_semi_auto_parallel_sharding_strategy MODULES
+                  test_semi_auto_parallel_sharding_strategy)
+  set_tests_properties(test_semi_auto_parallel_sharding_strategy
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 200)
+  py_test_modules(test_semi_auto_parallel_fsdp MODULES
+                  test_semi_auto_parallel_fsdp)
+  set_tests_properties(test_semi_auto_parallel_fsdp
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+  py_test_modules(test_semi_auto_parallel_lazy_init MODULES
+                  test_semi_auto_parallel_lazy_init)
+  set_tests_properties(test_semi_auto_parallel_lazy_init
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+  py_test_modules(test_semi_auto_parallel_in_framework MODULES
+                  test_semi_auto_parallel_in_framework)
+  set_tests_properties(test_semi_auto_parallel_in_framework
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 300)
+  py_test_modules(test_semi_auto_parallel_dygraph_inplace MODULES
+                  test_semi_auto_parallel_dygraph_inplace)
+  set_tests_properties(test_semi_auto_parallel_dygraph_inplace
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_semi_auto_parallel_for_flex_checkpoint MODULES
+                  test_semi_auto_parallel_for_flex_checkpoint)
+  set_tests_properties(test_semi_auto_parallel_for_flex_checkpoint
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_dist_tensor_api MODULES test_dist_tensor_api)
+  set_tests_properties(test_dist_tensor_api
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 200)
+  py_test_modules(test_semi_auto_parallel_saved_tensor_hook MODULES
+                  test_semi_auto_parallel_saved_tensor_hook)
+  set_tests_properties(test_semi_auto_parallel_saved_tensor_hook
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(
+    test_semi_auto_parallel_dist_to_static MODULES
+    test_semi_auto_parallel_dist_to_static ENVS FLAGS_enable_pir_api=1)
+  set_tests_properties(test_semi_auto_parallel_dist_to_static
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 300)
+  py_test_modules(test_static_reshard_api MODULES test_static_reshard_api ENVS
+                  FLAGS_enable_pir_api=1)
+  set_tests_properties(test_static_reshard_api
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 300)
+  py_test_modules(test_dist_checkpoint_utils MODULES test_dist_checkpoint_utils)
+  set_tests_properties(test_dist_checkpoint_utils
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+
+  py_test_modules(
+    test_semi_auto_parallel_unshard_dtensor MODULES
+    test_semi_auto_parallel_unshard_dtensor ENVS FLAGS_enable_pir_api=1)
+  set_tests_properties(test_semi_auto_parallel_unshard_dtensor
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_auto_parallel_backward_test MODULES
+                  test_auto_parallel_backward_test ENVS)
+  set_tests_properties(test_auto_parallel_backward_test
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
+  py_test_modules(test_moe_utils MODULES test_moe_utils)
+  set_tests_properties(test_moe_utils PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE"
+                                                 TIMEOUT 30)
+  py_test_modules(test_object_list_communication MODULES
+                  test_object_list_communication)
+  set_tests_properties(test_object_list_communication
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 50)
+  # End of unittests WITH multi cards and timeout
+
+  # NOTE(zyl): unittests WITH multi cards and WITHOUT timeout
+  py_test_modules(test_semi_auto_parallel_moe_utils MODULES
+                  test_semi_auto_parallel_moe_utils)
+  set_tests_properties(test_semi_auto_parallel_moe_utils
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE")
+  # End of unittests WITH multi cards and WITHOUT timeout
+
+  py_test_modules(test_semi_auto_parallel_functional_in_single_card MODULES
+                  test_semi_auto_parallel_functional_in_single_card)
+  # End of unittests WITH single card and timeout
+
+  # NOTE(zyl): unittests WITH single card and WITHOUT timeout
+  py_test_modules(test_align_mode MODULES test_align_mode)
+  py_test_modules(test_tunable_variable MODULES test_tunable_variable)
+  py_test_modules(test_tunable_space MODULES test_tunable_space)
+  py_test_modules(test_recorder MODULES test_recorder)
+  py_test_modules(test_trial MODULES test_trial)
+  py_test_modules(test_cluster MODULES test_cluster)
+  py_test_modules(test_comm_cost MODULES test_comm_cost)
+  py_test_modules(test_comp_cost MODULES test_comp_cost)
+  py_test_modules(test_cluster_v2 MODULES test_cluster_v2)
+  py_test_modules(test_process_mesh_v2 MODULES test_process_mesh_v2)
+  py_test_modules(test_strategy MODULES test_strategy)
+  py_test_modules(test_cluster_partition MODULES test_cluster_partition)
+  py_test_modules(test_convert_to_process_meshes MODULES
+                  test_convert_to_process_meshes)
+  py_test_modules(test_dist_tensor MODULES test_dist_tensor ENVS
+                  FLAGS_enable_pir_api=1)
+  py_test_modules(test_api_dist_branch MODULES test_api_dist_branch)
+  py_test_modules(test_shard_tensor_api MODULES test_shard_tensor_api ENVS
+                  FLAGS_enable_pir_api=1)
+  py_test_modules(test_placement_types MODULES test_placement_types)
+  py_test_modules(test_strategy_api MODULES test_strategy_api)
+  py_test_modules(test_parallel_api MODULES test_parallel_api)
+  py_test_modules(test_dtensor_to_local_api MODULES test_dtensor_to_local_api)
+  py_test_modules(test_dtensor_from_local_api MODULES
+                  test_dtensor_from_local_api)
+  py_test_modules(test_dy_local_view_compute MODULES test_dy_local_view_compute)
+  py_test_modules(test_local_view_compute MODULES test_local_view_compute)
+  py_test_modules(test_microbatch MODULES test_microbatch)
+  py_test_modules(test_PipelineStage MODULES test_PipelineStage)
+  py_test_modules(
+    test_tp_conv
+    MODULES
+    test_tp_conv
+    ENVS
+    FLAGS_max_inplace_grad_add=4
+    FLAGS_cudnn_deterministic=1
+    FLAGS_embedding_deterministic=1
+    NVIDIA_TF32_OVERRIDE=0)
+  py_test_modules(test_PP_Schedules MODULES test_PP_Schedules)
+  py_test_modules(test_pipeline_sync_shared_parameters MODULES
+                  test_pipeline_sync_shared_parameters)
+  py_test_modules(
+    test_context_parallel
+    MODULES
+    test_context_parallel
+    ENVS
+    FLAGS_max_inplace_grad_add=4
+    FLAGS_cudnn_deterministic=1
+    FLAGS_embedding_deterministic=1
+    NVIDIA_TF32_OVERRIDE=0)
+  # End of unittests WITH single card WITHOUT timeout
+
+  py_test_modules(test_clear_param_storage_api MODULES
+                  test_clear_param_storage_api)
+
+endif()
+
+py_test_modules(test_job_schedule_profiler_range MODULES
+                test_job_schedule_profiler_range)
+
+set_pir_tests_properties()

test/auto_parallel/PP_Schedules_demo.py

@@ -0,0 +1,678 @@
+# 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 random
+import types
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import fleet
+from paddle.distributed.auto_parallel._utils import (
+    _patch_grads_for_step,
+)
+from paddle.distributed.auto_parallel.pipelining.schedules import (
+    Schedule1F1B,
+    ScheduleFThenB,
+    ScheduleVPP,
+)
+from paddle.distributed.auto_parallel.pipelining.stage import (
+    PipelineStage,
+)
+from paddle.io import DataLoader, Dataset
+
+
+def fix_seeds(seed=2025):
+    """Fix random seeds to ensure reproducibility"""
+    paddle.seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+
+class PPMyModel(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.mesh = paddle.distributed.ProcessMesh(
+            [0, 1, 2, 3], dim_names=["pp"]
+        )
+        self.num_layers = 8
+        self.num_layers_per_card = self.num_layers // 4
+
+        self.linears = nn.LayerList()
+        for i in range(self.num_layers):
+            linear = nn.Linear(8, 8, bias_attr=False)
+
+            # Mark network parameters
+            linear.weight = dist.shard_tensor(
+                linear.weight,
+                self.get_pp_mesh(i),
+                [dist.Replicate()],
+            )
+
+            self.linears.append(linear)
+
+    def get_pp_mesh(self, layer_index):
+        # layer_index=0-3 corresponds to mesh_idx 0,0,1,1,2,2,3,3
+        mesh_idx = int(layer_index / (self.num_layers / 4))
+        return self.mesh[mesh_idx]
+
+    def forward(self, x):
+        x.stop_gradient = False
+        out = x
+        for i in range(self.num_layers):
+            # Mark intermediate variables, reshard when switching devices
+            if i % self.num_layers_per_card == 0 and i > 0:
+                out = dist.reshard(out, self.get_pp_mesh(i), [dist.Replicate()])
+
+            out = self.linears[i](out)
+        return paddle.cast(out, 'float32')
+
+
+class PPMyModel_SingleStage(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.mesh = paddle.distributed.ProcessMesh(
+            [0, 1, 2, 3], dim_names=["pp"]
+        )
+        self.num_layers = 8
+        self.num_layers_per_card = 2
+        self.linears = nn.LayerList()
+        for i in range(self.num_layers):
+            linear = nn.Linear(8, 8, bias_attr=False)
+
+            linear.weight = dist.shard_tensor(
+                linear.weight,
+                self.get_pp_mesh(i),
+                [dist.Replicate()],
+            )
+
+            self.linears.append(linear)
+
+    def get_pp_mesh(self, layer_index):
+        # layer_index=0-7 maps to mesh_idx as 0,0,1,1,2,2,3,3
+        mesh_idx = int(layer_index // self.num_layers_per_card)
+        return self.mesh[mesh_idx]
+
+    def forward(self, x):
+        x.stop_gradient = False
+        out = x
+        device_id = dist.get_rank()
+        for i in range(self.num_layers):
+            if int(i // self.num_layers_per_card) == device_id:
+                out = self.linears[i](out)
+
+        return paddle.cast(out, 'float32')
+
+
+class PPMyModel_MultiStage(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.mesh = paddle.distributed.ProcessMesh(
+            [0, 1, 2, 3], dim_names=["pp"]
+        )
+        self.num_layers = 8
+        self.linears = nn.LayerList()
+        for i in range(self.num_layers):
+            linear = nn.Linear(8, 8, bias_attr=False)
+
+            linear.weight = dist.shard_tensor(
+                linear.weight,
+                self.get_pp_mesh(i),
+                [dist.Replicate()],
+            )
+
+            self.linears.append(linear)
+
+    def get_pp_mesh(self, layer_index):
+        mesh_idx = int(layer_index % 4)
+        return self.mesh[mesh_idx]
+
+    def forward(self, x):
+        # For MultiStage, we shard model layers, so forward calls _Pipeline_model_chunk's forward
+        pass
+
+
+class _Pipeline_model_chunk(nn.Layer):
+    def __init__(self, layers):
+        super().__init__()
+        self.layers = layers
+
+    def forward(self, x):
+        out = x
+        for layer in self.layers:
+            out = layer(out)
+        return out
+
+
+class PP_DP_MyModel(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        pp_mesh0 = paddle.distributed.ProcessMesh([0, 2], dim_names=["dp"])
+        pp_mesh1 = paddle.distributed.ProcessMesh([1, 3], dim_names=["dp"])
+        self.num_layers = 8
+        self.linears = nn.LayerList()
+        for i in range(self.num_layers):
+            linear = nn.Linear(8, 8, bias_attr=False)
+            if i < 4:
+                linear.weight = dist.shard_tensor(
+                    linear.weight, pp_mesh0, [dist.Replicate()]
+                )
+            else:
+                linear.weight = dist.shard_tensor(
+                    linear.weight, pp_mesh1, [dist.Replicate()]
+                )
+
+            self.linears.append(linear)
+
+    def forward(self, x):
+        x.stop_gradient = False
+        out = x
+        # Get current rank's position in pp group (0 or 1)
+        pp_rank = dist.get_rank() % 2
+
+        # Only process layers belonging to current rank
+        start_layer = (
+            4 * pp_rank
+        )  # rank 0/2 processes layers 0-3, rank 1/3 processes layers 4-7
+        end_layer = start_layer + 4
+
+        for i in range(start_layer, end_layer):
+            out = self.linears[i](out)
+
+        return paddle.cast(out, 'float32')
+
+
+class RandomDataset(Dataset):
+    def __init__(self, image_size, output_size, num_samples=1):
+        super().__init__()
+        self.image_size = image_size
+        self.num_samples = num_samples
+        self.output_size = output_size
+
+    def __getitem__(self, index):
+        input = paddle.rand([self.image_size], dtype='float32')
+        label = paddle.rand([self.output_size], dtype='float32')
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+class Test_Schedules:
+    @classmethod
+    def setUpClass(cls):
+        """Initialize test class setup"""
+        paddle.distributed.init_parallel_env()
+        cls.group = paddle.distributed.new_group([0, 1, 2, 3])
+        cls.rank = dist.get_rank()
+        cls.mesh = paddle.distributed.ProcessMesh(
+            [0, 1, 2, 3], dim_names=["pp"]
+        )
+        fleet.auto.set_mesh(cls.mesh)
+
+    def test_ScheduleFThenB(self):
+        fix_seeds()
+        self.model = PPMyModel_SingleStage()
+        self.micro_batches = 8
+        self.stage = PipelineStage(self.model, self.rank, 4, group=self.group)
+        self.stage.has_backward = True
+        loss_fn_ = nn.MSELoss()
+        schedule = ScheduleFThenB(
+            self.stage, self.micro_batches, loss_fn=loss_fn_
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=self.model.parameters()
+        )
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_step = []
+        num_iterations = 20
+
+        for iter_idx in range(num_iterations):
+            losses_by_micro_batch = []
+            for i, (data, label) in enumerate(loader):
+                schedule.step(data, target=label, losses=losses_by_micro_batch)
+                if self.rank == 3:
+                    losses_by_step.append(
+                        np.array(losses_by_micro_batch, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        return losses_by_step
+
+    def test_Schedule1F1B(self):
+        fix_seeds()
+        self.model = PPMyModel_SingleStage()
+        self.micro_batches = 8
+        self.stage = PipelineStage(self.model, self.rank, 4, group=self.group)
+        self.stage.has_backward = True
+        loss_fn_ = nn.MSELoss()
+        schedule = Schedule1F1B(
+            self.stage, self.micro_batches, loss_fn=loss_fn_
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=self.model.parameters()
+        )
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_step = []
+        num_iterations = 20
+
+        for iter_idx in range(num_iterations):
+            losses_by_micro_batch = []
+            for i, (data, label) in enumerate(loader):
+                schedule.step(data, target=label, losses=losses_by_micro_batch)
+                if self.rank == 3:
+                    losses_by_step.append(
+                        np.array(losses_by_micro_batch, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        return losses_by_step
+
+    def test_ScheduleVPP(self):
+        fix_seeds()
+        self.model = PPMyModel_MultiStage()
+        self.local_stages = 2
+        self.micro_batches = 8
+        self.stage_list = []
+        for i in range(self.local_stages):
+            stage_model = _Pipeline_model_chunk(
+                self.model.linears[self.rank + i * 4 : self.rank + i * 4 + 1]
+            )
+            self.stage_list.append(
+                PipelineStage(
+                    stage_model, self.rank + i * 4, 8, group=self.group
+                )
+            )
+            self.stage_list[i].has_backward = True
+        loss_fn_ = nn.MSELoss()
+        schedule = ScheduleVPP(
+            self.stage_list, self.micro_batches, loss_fn=loss_fn_
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=self.model.parameters()
+        )
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_micro_batch = []
+        losses_by_step = []
+        num_iterations = 20
+
+        for iter_idx in range(num_iterations):
+            for i, (data, label) in enumerate(loader):
+                schedule.step(data, target=label, losses=losses_by_micro_batch)
+                if self.rank == 3:
+                    losses_by_step.append(
+                        np.array(losses_by_micro_batch, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        return losses_by_step
+
+    def test_pp_model(self):
+        """Test pipeline parallel model using PPMyModel as the baseline"""
+        fix_seeds()
+        pp_model = PPMyModel()
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=pp_model.parameters()
+        )
+        loss_fn = nn.MSELoss()
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=1)
+        pp_losses_step = []
+        num_iterations = 20
+
+        for iter_idx in range(num_iterations):
+            pp_losses_micro_batch = []
+            for i, (data, label) in enumerate(loader):
+                output = pp_model(data)
+                loss = loss_fn(output, label)
+                pp_losses_micro_batch.append(loss.item())
+                loss.backward()
+            pp_losses_step.append(
+                np.array(pp_losses_micro_batch, dtype=np.float32).mean()
+            )
+            opt.step()
+            opt.clear_grad()
+        return pp_losses_step
+
+    def test_dp_pp(self):
+        fix_seeds()
+        global_mesh = paddle.distributed.ProcessMesh(
+            [[0, 2], [1, 3]], dim_names=["pp", "dp"]
+        )
+        fleet.auto.set_mesh(global_mesh)
+        self.model = PP_DP_MyModel()
+        pp_mesh0 = paddle.distributed.ProcessMesh([0, 2], dim_names=["dp"])
+        pp_mesh1 = paddle.distributed.ProcessMesh([1, 3], dim_names=["dp"])
+        dp_pp_pleacement = [dist.Shard(0)]
+        pp_group_1 = paddle.distributed.new_group([0, 1])
+        pp_group_2 = paddle.distributed.new_group([2, 3])
+        dp_group = paddle.distributed.new_group([1, 3])
+        self.micro_batches = 4
+        if self.rank < 2:
+            self.stage = PipelineStage(
+                self.model, self.rank % 2, 2, group=pp_group_1
+            )
+        else:
+            self.stage = PipelineStage(
+                self.model, self.rank % 2, 2, group=pp_group_2
+            )
+        self.stage.has_backward = True
+        loss_fn_ = nn.MSELoss()
+        schedule = ScheduleFThenB(
+            self.stage, self.micro_batches, loss_fn=loss_fn_
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=self.model.parameters()
+        )
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_step = []
+        num_iterations = 20
+        all_losses_in_one_step_md5sum = []
+        for iter_idx in range(num_iterations):
+            losses_by_micro_batch = []
+            for i, (data, label) in enumerate(loader):
+                # reorder data and label
+                batch_size = data.shape[0]
+                even_indices = list(range(0, batch_size, 2))
+                odd_indices = list(range(1, batch_size, 2))
+                reordered_indices = even_indices + odd_indices
+
+                reordered_data = data[reordered_indices]
+                reordered_label = label[reordered_indices]
+
+                dist_data = dist.shard_tensor(
+                    reordered_data, pp_mesh0, dp_pp_pleacement
+                )
+                dist_label = dist.shard_tensor(
+                    reordered_label, pp_mesh1, dp_pp_pleacement
+                )
+                schedule.step(
+                    dist_data, target=dist_label, losses=losses_by_micro_batch
+                )
+                # Losses from two dp paths are in Partial(AVG) state, need to do all_reduce
+                if self.rank == 1 or self.rank == 3:
+                    reduced_losses = []
+                    for item in losses_by_micro_batch:
+                        local_loss = item._local_value()
+                        dist.all_reduce(
+                            local_loss, op=dist.ReduceOp.AVG, group=dp_group
+                        )
+                        reduced_losses.append(local_loss)
+                        if iter_idx == 0:
+                            all_losses_in_one_step_md5sum.append(
+                                local_loss._md5sum()
+                            )
+
+                if self.rank == 3:
+                    # Calculate mean using reduced losses
+                    losses_by_step.append(
+                        np.array(reduced_losses, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        return losses_by_step, all_losses_in_one_step_md5sum
+
+    def test_pp_model_with_ClipGradByGlobalNorm(self):
+        """Test pipeline parallel model with ClipGradByGlobalNorm using PPMyModel as the baseline"""
+        fix_seeds()
+        pp_model = PPMyModel()
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001,
+            parameters=pp_model.parameters(),
+            grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+        )
+        loss_fn = nn.MSELoss()
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=1)
+        pp_losses_step = []
+        num_iterations = 20
+
+        for iter_idx in range(num_iterations):
+            pp_losses_micro_batch = []
+            for i, (data, label) in enumerate(loader):
+                output = pp_model(data)
+                loss = loss_fn(output, label)
+                pp_losses_micro_batch.append(loss.item())
+                loss.backward()
+            pp_losses_step.append(
+                np.array(pp_losses_micro_batch, dtype=np.float32).mean()
+            )
+            opt.step()
+            opt.clear_grad()
+        return pp_losses_step
+
+    def test_ScheduleFThenB_with_ClipGradByGlobalNorm(self):
+        fix_seeds()
+        self.model = PPMyModel_SingleStage()
+        self.micro_batches = 8
+        self.stage = PipelineStage(self.model, self.rank, 4, group=self.group)
+        self.stage.has_backward = True
+        loss_fn_ = nn.MSELoss()
+        schedule = ScheduleFThenB(
+            self.stage, self.micro_batches, loss_fn=loss_fn_
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001,
+            parameters=self.model.parameters(),
+            grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+        )
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_step = []
+        num_iterations = 20
+
+        for iter_idx in range(num_iterations):
+            losses_by_micro_batch = []
+            for i, (data, label) in enumerate(loader):
+                schedule.step(data, target=label, losses=losses_by_micro_batch)
+                if self.rank == 3:
+                    losses_by_step.append(
+                        np.array(losses_by_micro_batch, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        return losses_by_step
+
+    def test_FthenB_align_mode_of_GradientClipByGlobalNorm(self):
+        fix_seeds()
+        paddle.set_flags(
+            {'FLAGS_enable_auto_parallel_align_mode': True}
+        )  # Represents logical alignment with GradientClipByGlobalNorm that is semi-automatically parallel to the original dynamic graph, because the processing logic here is not aligned with the dynamic graph manually parallel
+        self.model = PPMyModel_SingleStage()
+        self.micro_batches = 8
+        self.stage = PipelineStage(self.model, self.rank, 4, group=self.group)
+        self.stage.has_backward = True
+        loss_fn_ = nn.MSELoss()
+        schedule = ScheduleFThenB(
+            self.stage, self.micro_batches, loss_fn=loss_fn_
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001,
+            parameters=self.model.parameters(),
+            grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+        )
+        if dist.in_auto_parallel_align_mode():  # When in auto parallel align mode, patching the optimizer step function
+            orig_step = (
+                opt.step.__func__ if hasattr(opt.step, "__func__") else opt.step
+            )
+            decorator = _patch_grads_for_step(amp_master_grad=True)
+            new_step = decorator(
+                orig_step
+            )  # When the step function is wrapped by the decorator, it initializes gradients for parameters belonging to other ranks prior to step method execution, ensuring their metadata is preserved.
+            opt.step = types.MethodType(new_step, opt)
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_step = []
+        num_iterations = 20
+
+        for iter_idx in range(num_iterations):
+            losses_by_micro_batch = []
+            for i, (data, label) in enumerate(loader):
+                schedule.step(data, target=label, losses=losses_by_micro_batch)
+                if self.rank == 3:
+                    losses_by_step.append(
+                        np.array(losses_by_micro_batch, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        paddle.set_flags({'FLAGS_enable_auto_parallel_align_mode': False})
+        return losses_by_step
+
+    def test_dp_pp_align_mode(self):
+        fix_seeds()
+        paddle.set_flags(
+            {'FLAGS_enable_auto_parallel_align_mode': True}
+        )  # Represents manual parallel alignment with dynamic graphs, mainly segmenting microbatches when aligning DP and PP mixing
+        global_mesh = paddle.distributed.ProcessMesh(
+            [[0, 2], [1, 3]], dim_names=["pp", "dp"]
+        )
+        fleet.auto.set_mesh(global_mesh)
+        self.model = PP_DP_MyModel()
+        pp_mesh0 = paddle.distributed.ProcessMesh([0, 2], dim_names=["dp"])
+        pp_mesh1 = paddle.distributed.ProcessMesh([1, 3], dim_names=["dp"])
+        dp_pp_pleacement = [dist.Shard(0)]
+        pp_group_1 = paddle.distributed.new_group([0, 1])
+        pp_group_2 = paddle.distributed.new_group([2, 3])
+        dp_group = paddle.distributed.new_group([1, 3])
+        self.micro_batches = 4
+        if self.rank < 2:
+            self.stage = PipelineStage(
+                self.model, self.rank % 2, 2, group=pp_group_1
+            )
+        else:
+            self.stage = PipelineStage(
+                self.model, self.rank % 2, 2, group=pp_group_2
+            )
+        self.stage.has_backward = True
+        loss_fn_ = nn.MSELoss()
+        schedule = ScheduleFThenB(
+            self.stage, self.micro_batches, loss_fn=loss_fn_
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=self.model.parameters()
+        )
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_step = []
+        all_losses_in_one_step_md5sum = []
+        num_iterations = 20
+        for iter_idx in range(num_iterations):
+            losses_by_micro_batch = []
+            for i, (data, label) in enumerate(loader):
+                dist_data = dist.shard_tensor(data, pp_mesh0, dp_pp_pleacement)
+                dist_label = dist.shard_tensor(
+                    label, pp_mesh1, dp_pp_pleacement
+                )
+                schedule.step(
+                    dist_data, target=dist_label, losses=losses_by_micro_batch
+                )
+                # Losses from two dp paths are in Partial(AVG) state, need to do all_reduce
+                if self.rank == 1 or self.rank == 3:
+                    reduced_losses = []
+                    for item in losses_by_micro_batch:
+                        local_loss = item._local_value()
+                        dist.all_reduce(
+                            local_loss, op=dist.ReduceOp.AVG, group=dp_group
+                        )
+                        reduced_losses.append(local_loss)
+                        if iter_idx == 0:
+                            all_losses_in_one_step_md5sum.append(
+                                local_loss._md5sum()
+                            )
+
+                if self.rank == 3:
+                    # Calculate mean using reduced losses
+                    losses_by_step.append(
+                        np.array(reduced_losses, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        paddle.set_flags({'FLAGS_enable_auto_parallel_align_mode': False})
+        return losses_by_step, all_losses_in_one_step_md5sum
+
+    def run_test(self):
+        """Compare losses between three training methods"""
+        self.setUpClass()
+        pp_losses = self.test_pp_model()
+        scheduleFThenB_losses = self.test_ScheduleFThenB()
+        schedule1f1b_losses = self.test_Schedule1F1B()
+        schedulevpp_losses = self.test_ScheduleVPP()
+        pp_model_with_ClipGradByGlobalNorm_losses = (
+            self.test_pp_model_with_ClipGradByGlobalNorm()
+        )
+        scheduleFThenB_with_ClipGradByGlobalNorm_losses = (
+            self.test_ScheduleFThenB_with_ClipGradByGlobalNorm()
+        )
+        scheduleFthenB_align_mode_losses_of_GradientClipByGlobalNorm = (
+            self.test_FthenB_align_mode_of_GradientClipByGlobalNorm()
+        )
+        dp_pp_losses, dp_pp_losses_md5sum = self.test_dp_pp()
+        dp_pp_align_mode_losses, dp_pp_align_mode_losses_md5sum = (
+            self.test_dp_pp_align_mode()
+        )
+
+        if self.rank == 3:
+            np.testing.assert_allclose(
+                pp_losses,
+                scheduleFThenB_losses,
+                rtol=1e-5,
+            )
+
+            np.testing.assert_allclose(
+                schedule1f1b_losses,
+                scheduleFThenB_losses,
+                rtol=1e-5,
+            )
+
+            np.testing.assert_allclose(
+                schedulevpp_losses,
+                scheduleFThenB_losses,
+                rtol=1e-5,
+            )
+
+            np.testing.assert_allclose(
+                dp_pp_losses,
+                scheduleFThenB_losses,
+                rtol=1e-5,
+            )
+
+            np.testing.assert_allclose(
+                pp_model_with_ClipGradByGlobalNorm_losses,
+                scheduleFThenB_with_ClipGradByGlobalNorm_losses,
+                rtol=1e-5,
+            )
+
+            np.testing.assert_allclose(
+                dp_pp_align_mode_losses,
+                dp_pp_losses,
+                rtol=1e-5,
+            )
+
+            np.testing.assert_allclose(
+                scheduleFthenB_align_mode_losses_of_GradientClipByGlobalNorm,
+                pp_model_with_ClipGradByGlobalNorm_losses,
+                rtol=1e-5,
+            )
+
+            assert dp_pp_losses_md5sum == dp_pp_align_mode_losses_md5sum
+
+
+if __name__ == '__main__':
+    Test_Schedules().run_test()

test/auto_parallel/PipelineStage_demo.py

@@ -0,0 +1,532 @@
+# 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.
+
+from __future__ import annotations
+
+import logging
+import random
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import fleet
+from paddle.distributed.auto_parallel.pipelining._backward import (
+    stage_backward,
+    stage_backward_input,
+    stage_backward_weight,
+)
+from paddle.distributed.auto_parallel.pipelining.stage import (
+    PipelineStage,
+    _RecvInfo,
+)
+from paddle.distributed.auto_parallel.pipelining.utils import (
+    PipeliningShapeError,
+    TensorMeta,
+    _detach_and_keep_grad,
+    _friendly_debug_info,
+    _get_stage_mesh,
+    _validate_tensor_metadata,
+    _validate_tensors_metadata,
+    _zero_initialize_with_meta,
+)
+from paddle.io import Dataset
+
+if TYPE_CHECKING:  # 添加类型检查块
+    from paddle.distributed.communication.group import Group
+logger = logging.getLogger(__name__)
+
+
+def fix_seeds(seed=2025):
+    """Fix random seeds to ensure reproducibility"""
+    paddle.seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+
+def _batch_p2p(p2p_ops, desc=None):
+    # TODO(zhengtianyu): 等合入Scheduler后，删除该函数
+    """Execute batch point-to-point communication operations"""
+    if len(p2p_ops) == 0:
+        return None
+    desc_str = f"{desc}, " if desc else ""
+    logger.debug("batch_p2p %s%s", desc_str, p2p_ops)
+    return dist.batch_isend_irecv(p2p_ops).pop()
+
+
+def _sorted_batch_p2p(p2p_ops, desc=None):
+    # TODO(zhengtianyu): 等合入Scheduler后，删除该函数
+    """Sort and execute batch point-to-point communication by peer rank"""
+    ops_by_peer: dict[int, list[dist.P2POp]] = defaultdict(list)
+    work_by_peer: dict[int, dist.Work] = {}
+    if len(p2p_ops) == 0:
+        return work_by_peer
+
+    for op in p2p_ops:
+        ops_by_peer[op.peer].append(op)
+
+    for peer, ops in sorted(ops_by_peer.items()):
+        work_by_peer[peer] = _batch_p2p(ops, desc=desc)
+
+    return work_by_peer
+
+
+class MyModel(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = nn.Linear(8, 8, bias_attr=False)
+        self.linear2 = nn.Linear(8, 8, bias_attr=False)
+        self.linear3 = nn.Linear(8, 8)
+        self.linear4 = nn.Linear(8, 8)
+
+    def forward(self, x, debug_str=None):
+        if hasattr(self, 'linear1'):
+            if debug_str:
+                logger.debug(f"{debug_str} linear1")
+            x = self.linear1(x)
+            x = self.linear2(x)
+        if hasattr(self, 'linear3'):
+            x = self.linear3(x)
+            x = self.linear4(x)
+        return x
+
+
+class PPMyModel(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.mesh = paddle.distributed.ProcessMesh([0, 1], dim_names=["pp"])
+        self.num_layers = 4
+        self.num_layers_per_card = self.num_layers // 2
+
+        # Create layers same as MyModel
+        self.linears = nn.LayerList()
+        for i in range(self.num_layers):
+            if i // 2 == 0:
+                linear = nn.Linear(8, 8, bias_attr=False)
+            else:
+                linear = nn.Linear(8, 8)
+
+            # Mark network parameters
+            linear.weight = dist.shard_tensor(
+                linear.weight,
+                self.get_pp_mesh(i),
+                [dist.Replicate()],
+            )
+
+            self.linears.append(linear)
+
+    def get_pp_mesh(self, layer_index):
+        # layer_index=0-3 corresponds to mesh_idx 0,0,1,1 respectively
+        mesh_idx = int(layer_index / (self.num_layers / 2))
+        return self.mesh[mesh_idx]
+
+    def forward(self, x):
+        x.stop_gradient = False
+        out = x
+
+        for i in range(self.num_layers):
+            # Mark intermediate variables, reshard when device switching is needed
+            if i % self.num_layers_per_card == 0 and i > 0:
+                out = dist.reshard(out, self.get_pp_mesh(i), [dist.Replicate()])
+
+            out = self.linears[i](out)
+
+        return paddle.cast(out, 'float32')
+
+
+class RandomDataset(Dataset):
+    def __init__(self, image_size, num_samples=1):
+        super().__init__()
+        self.image_size = image_size
+        self.num_samples = num_samples
+
+    def __getitem__(self, index):
+        # Keep dimension as [8]
+        input = paddle.rand([self.image_size], dtype='float32')
+        label = paddle.rand([8], dtype='float32')
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+def manual_model_split(
+    model: MyModel, stage_idx: int, group: Group
+) -> PipelineStage:
+    """Manually split model into pipeline stages"""
+    if stage_idx == 0:
+        del model.linear3
+        del model.linear4
+    elif stage_idx == 1:
+        del model.linear1
+        del model.linear2
+    else:
+        raise ValueError("Invalid stage index.")
+
+    return PipelineStage(model, stage_idx, 2, group=group)
+
+
+class TestPipelineStage:
+    @classmethod
+    def setUpClass(cls):
+        """Initialize test class setup"""
+        paddle.distributed.init_parallel_env()
+        cls.group = paddle.distributed.new_group([0, 1])
+        cls.rank = dist.get_rank()
+        cls.mesh = paddle.distributed.ProcessMesh([0, 1], dim_names=["pp"])
+        fleet.auto.set_mesh(cls.mesh)
+
+    def test_PipelineStage(self):
+        """Test complete pipeline including forward, backward and model comparison"""
+        fix_seeds()
+        self.model = MyModel()
+        self.micro_batches = 1  # The PipelineStage component is currently tested separately, so it is set to 1, and the micro_batches > 1 scenario will be overridden when the schedule component is tested in the future
+        self.stage = manual_model_split(self.model, self.rank, self.group)
+        self.stage.has_backward = True
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=self.model.parameters()
+        )
+        loss_fn = nn.MSELoss()
+        dataset = RandomDataset(image_size=8, num_samples=100)
+        losses = []
+        num_iterations = 20
+        data0 = paddle.zeros([8], dtype='float32')
+        label0 = paddle.zeros([8], dtype='float32')
+        data0 = paddle.to_tensor(data0).unsqueeze(0)
+        label0 = paddle.to_tensor(label0).unsqueeze(0)
+        #  Prepare infrastructure
+        if self.rank == 0:
+            output = self.stage._prepare_forward_infra(
+                self.micro_batches,
+                (data0,),
+                {
+                    "debug_str": "test debug_str",
+                },
+            )
+        else:
+            output = self.stage._prepare_forward_infra(
+                self.micro_batches,
+                (),
+                {
+                    "debug_str": "test debug_str",
+                },
+            )
+        loss = None
+        if self.stage.is_last:
+            loss = loss_fn(output[0], label0)
+        self.stage._prepare_backward_infra(self.micro_batches, loss)
+        for iter_idx in range(num_iterations):
+            data, label = dataset[iter_idx]
+            data = paddle.to_tensor(data).unsqueeze(0)
+            label = paddle.to_tensor(label).unsqueeze(0)
+            # Forward pass
+            fwd_sends_to_wait = []
+
+            # Receive operations
+            ops = self.stage.get_fwd_recv_ops(0)
+            works = _sorted_batch_p2p(ops, desc="fwd_recv")
+            for work in works.values():
+                work.wait()
+
+            # Forward computation
+            output = self.stage.forward_one_chunk(
+                0,
+                (data,),
+                {
+                    "debug_str": "test debug_str",
+                },
+            )
+            # Send operations
+            ops = self.stage.get_fwd_send_ops(0)
+            works = _sorted_batch_p2p(ops, desc="fwd_send")
+            fwd_sends_to_wait.extend(works.values())
+
+            # Wait for all send operations to complete
+            for work in fwd_sends_to_wait:
+                work.wait()
+
+            # Calculate loss if last stage
+            loss = None
+            if self.stage.is_last:
+                loss = loss_fn(output, label)
+                assert loss is not None
+                losses.append(loss.item())
+            # Backward pass
+            bwd_sends_to_wait = []
+
+            # Receive gradients
+            ops = self.stage.get_bwd_recv_ops(0)
+            works = _sorted_batch_p2p(ops, desc="bwd_recv")
+            for work in works.values():
+                work.wait()
+
+            # Backward computation
+            grads = self.stage.backward_one_chunk(
+                0, loss=loss, last_backward=True
+            )
+            assert grads is not None
+
+            # Send gradients
+            ops = self.stage.get_bwd_send_ops(0)
+            works = _sorted_batch_p2p(ops, desc="bwd_send")
+            bwd_sends_to_wait.extend(works.values())
+
+            # Wait for all send operations to complete
+            for work in bwd_sends_to_wait:
+                work.wait()
+            self.stage.clear_runtime_states()
+            opt.step()
+            opt.clear_grad()
+
+        return losses
+
+    def test_pp_model(self):
+        """Test pipeline parallel model using MyModel"""
+        fix_seeds()
+
+        pp_model = PPMyModel()
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=pp_model.parameters()
+        )
+        loss_fn = nn.MSELoss()
+
+        dataset = RandomDataset(image_size=8, num_samples=100)
+
+        pp_losses = []
+        num_iterations = 20
+
+        for iter_idx in range(num_iterations):
+            data, label = dataset[iter_idx]
+            data = paddle.to_tensor(data).unsqueeze(0)
+            label = paddle.to_tensor(label).unsqueeze(0)
+
+            output = pp_model(data)
+
+            loss = loss_fn(output, label)
+            pp_losses.append(loss.item())
+
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+
+        return pp_losses
+
+    def test_single_gpu(self):
+        """Test single GPU training with the complete model"""
+        # Only run single GPU training on rank 1
+        if self.rank == 1:
+            fix_seeds()
+            single_model = MyModel()
+            opt = paddle.optimizer.AdamW(
+                learning_rate=0.001, parameters=single_model.parameters()
+            )
+            loss_fn = nn.MSELoss()
+
+            dataset = RandomDataset(image_size=8, num_samples=100)
+
+            losses = []
+            num_iterations = 20
+
+            for iter_idx in range(num_iterations):
+                data, label = dataset[iter_idx]
+                output = single_model(data)
+
+                loss = loss_fn(output, label)
+                losses.append(loss.item())
+                loss.backward()
+
+                opt.step()
+                opt.clear_grad()
+
+            return losses
+        return None
+
+    def test_simple_func_about_schedules(self):
+        """Test local data transfer functions between stages on the same rank"""
+        if self.rank == 0:
+            # 1. Test set_local_fwd_input
+            tensor = paddle.to_tensor([1.0, 2.0, 3.0])
+            stage = PipelineStage(nn.Linear(3, 3), 1, 2, group=self.group)
+            stage.args_recv_info[0] = (_RecvInfo("test", 0, paddle.empty([3])),)
+            stage.set_local_fwd_input(tensor, 0)
+            assert stage.args_recv_info[0][0].buffer is not None
+
+            # 2. Test get_local_bwd_output
+            stage.has_backward = True
+            grad_tensor = paddle.to_tensor([4.0, 5.0, 6.0])
+            stage.bwd_cache[0] = (grad_tensor,)
+            stage.chunks = 2
+            bwd_output = stage.get_local_bwd_output(0)
+            assert bwd_output[0].equal_all(grad_tensor)
+
+            # 3. Test set_local_bwd_input
+            prev_stage = PipelineStage(nn.Linear(3, 3), 0, 2, group=self.group)
+            prev_stage.has_backward = True
+            prev_stage.grad_recv_info[0] = (
+                _RecvInfo("test", 1, paddle.empty([3])),
+            )
+            grad_input = (paddle.to_tensor([7.0, 8.0, 9.0]),)
+            prev_stage.set_local_bwd_input(grad_input, 0)
+            assert prev_stage.grad_recv_info[0][0].buffer.equal_all(
+                grad_input[0]
+            )
+
+    def test_backward_some_simple_examples(self):
+        """Test simple examples in backward"""
+        if self.rank == 0:
+            # 1. Test backward propagation with dictionary and tuple outputs
+            input_tensor = paddle.to_tensor([1.0, 2.0], stop_gradient=False)
+
+            output_dict = {
+                "out": input_tensor * 2.0,
+                "out_tensor_is_dict_grad_is_None": {"out": input_tensor * 2.0},
+                "out_tensor_is_tuple_grad_is_None": (input_tensor * 2.0,),
+            }
+            grad_dict = {
+                "out": paddle.to_tensor([0.1, 0.2]),
+                "out_tensor_is_dict_grad_is_None": None,
+                "out_tensor_is_tuple_grad_is_None": None,
+            }
+
+            input_grads = stage_backward(output_dict, grad_dict, [input_tensor])
+            expected_grad = paddle.to_tensor([2 * 0.1, 2 * 0.2])
+
+            np.testing.assert_allclose(
+                input_grads[0].numpy(), expected_grad.numpy(), rtol=1e-5
+            )
+            # 2. Test not yet implemented stage_backward_input and stage_backward_weight
+            try:
+                stage_backward_input(
+                    [input_tensor * 2.0],
+                    [paddle.to_tensor([0.1, 0.2])],
+                    [input_tensor],
+                    iter([paddle.to_tensor([1.0, 1.0])]),
+                )
+                raise AssertionError("Should raise Error")
+            except NotImplementedError as e:
+                pass
+            try:
+                stage_backward_weight(
+                    iter([paddle.to_tensor([1.0, 1.0])]),
+                    [{"params": [paddle.to_tensor([1.0, 1.0])]}],
+                )
+                raise AssertionError("Should raise Error")
+            except NotImplementedError as e:
+                pass
+
+    def test_utils_some_simple_examples(self):
+        """Test simple examples in utils"""
+        if self.rank == 0:
+            # 1. Test exceptions in _get_stage_mesh
+            try:
+                _get_stage_mesh(0, 2, style="v")
+                raise AssertionError("Should raise Error")
+            except NotImplementedError as e:
+                pass
+            try:
+                _get_stage_mesh(0, 2, style="unknown")
+                raise AssertionError("Should raise Error")
+            except ValueError as e:
+                pass
+
+            # 2. Test exceptions in _validate_tensors_metadata
+            try:
+                # Length mismatch
+                expected = [paddle.to_tensor([1.0, 2.0])]
+                actual = [paddle.to_tensor([1.0]), paddle.to_tensor([2.0])]
+                _validate_tensors_metadata("test", expected, actual)
+                raise AssertionError("Should raise Error")
+            except PipeliningShapeError as e:
+                pass
+
+            # 3. Test exceptions in _validate_tensor_metadata
+            try:
+                # Shape mismatch
+                expected = paddle.to_tensor([1.0, 2.0])
+                actual = paddle.to_tensor([1.0])
+                _validate_tensor_metadata("test", expected, actual)
+                raise AssertionError("Should raise Error")
+            except PipeliningShapeError as e:
+                pass
+
+            try:
+                # Dtype mismatch
+                expected = paddle.to_tensor([1.0, 2.0], dtype='float32')
+                actual = paddle.to_tensor([1, 2], dtype='int32')
+                _validate_tensor_metadata("test", expected, actual)
+                raise AssertionError("Should raise Error")
+            except PipeliningShapeError as e:
+                pass
+
+            # 4. Test _detach_and_keep_grad
+            a = paddle.to_tensor([2.0], stop_gradient=False)
+            b = a * 2
+            x = _detach_and_keep_grad(b)
+            assert x is b
+            assert x.stop_gradient == b.stop_gradient
+            assert (x.numpy() == b.numpy()).all()
+            x.stop_gradient = False
+            z = x * 3
+            z.backward()
+
+            assert x.grad is not None
+            assert a.grad is None
+
+            # 5. Test TensorMeta and _zero_initialize_with_meta
+            tensor = paddle.ones([4, 8])
+            dist_tensor = dist.shard_tensor(tensor, self.mesh, [dist.Shard(0)])
+            tensor_meta = TensorMeta(dist_tensor)
+            assert tensor_meta.shape == [4, 8]
+            assert tensor_meta._local_shape == [2, 8]
+
+            zero_tensor = _zero_initialize_with_meta(tensor_meta, self.mesh)
+            assert zero_tensor.shape == [4, 8]
+            assert zero_tensor.is_dist()
+            assert zero_tensor.process_mesh == self.mesh
+            assert zero_tensor.placements == [dist.Shard(0)]
+
+            # 6. Test _friendly_debug_info
+            a = {"test the input is not a tensor": 1}
+            assert _friendly_debug_info(a) == str(a)
+
+    def run_test(self):
+        """Compare losses between three training methods"""
+        self.setUpClass()
+        self.test_simple_func_about_schedules()
+        self.test_backward_some_simple_examples()
+        self.test_utils_some_simple_examples()
+        # Run three training methods
+        pipeline_losses = self.test_PipelineStage()
+        pp_losses = self.test_pp_model()
+        single_losses = self.test_single_gpu()
+
+        if self.rank == 1:
+            np.testing.assert_allclose(
+                pipeline_losses,
+                pp_losses,
+                rtol=1e-5,
+            )
+
+            np.testing.assert_allclose(
+                pipeline_losses,
+                single_losses,
+                rtol=1e-5,
+            )
+
+
+if __name__ == '__main__':
+    TestPipelineStage().run_test()

test/auto_parallel/auto_parallel_backward.py

@@ -0,0 +1,51 @@
+# 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 paddle
+import paddle.distributed as dist
+
+
+class TestBackwardAutoParallel:
+    def init_data(self):
+        self.mesh = dist.ProcessMesh([0], dim_names=['d0'])
+
+        self.x = paddle.to_tensor([[1]])
+        self.y = paddle.to_tensor([[1]])
+        self.z = paddle.to_tensor([[1]])
+
+        self.x.stop_gradient = False
+        self.y.stop_gradient = False
+        self.z.stop_gradient = False
+
+        self.z = dist.shard_tensor(self.z, self.mesh, [dist.Replicate()])
+
+    def run_test_case1(self):
+        self.init_data()
+        o = self.x * self.y
+        o = o + self.z
+        o = o.sum()
+        o.backward()
+
+    def run_test_case2(self):
+        self.init_data()
+        o = self.x + self.y
+        o = o - self.z
+        o = o.sum()
+        o.backward()
+
+
+# python -m paddle.distributed.launch --device=0 auto_parallel_backward.py
+if __name__ == '__main__':
+    TestBackwardAutoParallel().run_test_case1()
+    TestBackwardAutoParallel().run_test_case2()

test/auto_parallel/clear_param_storage_api.py

@@ -0,0 +1,59 @@
+# 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.distributed import fleet
+from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+    DygraphShardingOptimizerV2,
+)
+
+
+class TestClearParamStorage(unittest.TestCase):
+    def test_clear_param_storage(self):
+        class TestLayer(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)
+                self._w.color = {"color": "_w"}
+                self._b.color = {"color": "_b"}
+
+            @paddle.amp.debugging.check_layer_numerics
+            def forward(self, x):
+                return x * self._w + self._b
+
+        strategy = fleet.DistributedStrategy()
+        strategy.hybrid_configs = {
+            "dp_degree": 1,
+            "mp_degree": 1,
+            "pp_degree": 1,
+            "sharding_degree": 2,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+        hcg = fleet.get_hybrid_communicate_group()
+        dtype = 'float32'
+        model = TestLayer(dtype)
+
+        optimizer = paddle.optimizer.AdamW(parameters=model.parameters())
+        optimizer = DygraphShardingOptimizerV2(optimizer, hcg)
+        optimizer.clear_param_storage("_w")
+        optimizer.clear_param_storage("_b")
+        optimizer.clear_param_storage(None)
+        optimizer.reset_param_storage()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/auto_parallel/context_parallel.py

@@ -0,0 +1,155 @@
+# 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 os
+import random
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.ring_attention import (
+    shard_seq_load_balance,
+    unshard_seq_load_balance,
+)
+
+dist.init_parallel_env()
+
+
+class TestContextParallel:
+    def __init__(self):
+        self.rank = dist.get_rank()
+        self.world_size = dist.get_world_size()
+        self._sep_mesh = dist.ProcessMesh(
+            list(range(self.world_size)), dim_names=["sep"]
+        )
+
+    def set_seed(self, seed):
+        paddle.seed(seed)
+        np.random.seed(seed)
+        random.seed(seed)
+
+    def _test_cp_base(
+        self,
+        is_causal=True,
+    ):
+        mesh = dist.ProcessMesh(list(range(self.world_size)), dim_names=['sep'])
+        dist.auto_parallel.set_mesh(mesh)
+        self.set_seed(1024)
+        bs = 2
+        seq_len = 256  # flash_attn seq_len/card > 128
+        dim = 16
+        nheads = 2
+        dtype = paddle.bfloat16
+        q = paddle.rand(
+            (bs, seq_len, nheads, dim),
+            dtype=dtype,
+        )
+        k = paddle.rand(
+            (bs, seq_len, nheads, dim),
+            dtype=dtype,
+        )
+        v = paddle.rand(
+            (bs, seq_len, nheads, dim),
+            dtype=dtype,
+        )
+        q.stop_gradient = False
+        k.stop_gradient = False
+        v.stop_gradient = False
+
+        with paddle.no_grad():
+            dist.broadcast(q, src=0)
+            dist.broadcast(k, src=0)
+            dist.broadcast(v, src=0)
+        # base compute
+        output_ref = paddle.nn.functional.scaled_dot_product_attention(
+            q, k, v, is_causal=is_causal
+        )
+        loss_ref = output_ref.mean()
+        loss_ref.backward()
+
+        cp_q = q.detach().clone()
+        cp_k = k.detach().clone()
+        cp_v = v.detach().clone()
+        placements = [dist.Replicate() for _ in range(len(mesh.dim_names))]
+
+        # shard compute
+        sharded_q = dist.shard_tensor(cp_q, mesh, placements)
+        sharded_k = dist.shard_tensor(cp_k, mesh, placements)
+        sharded_v = dist.shard_tensor(cp_v, mesh, placements)
+        sharded_q = shard_seq_load_balance(sharded_q, 1)
+        sharded_k = shard_seq_load_balance(sharded_k, 1)
+        sharded_v = shard_seq_load_balance(sharded_v, 1)
+        sharded_q.stop_gradient = False
+        sharded_k.stop_gradient = False
+        sharded_v.stop_gradient = False
+
+        output_sharded = paddle.nn.functional.scaled_dot_product_attention(
+            sharded_q, sharded_k, sharded_v, is_causal=is_causal, backend='p2p'
+        )
+        loss_sharded = paddle.mean(output_sharded)
+        loss_sharded.backward()
+
+        with paddle.no_grad():
+            reorder_t = unshard_seq_load_balance(output_sharded, 1)
+        np.testing.assert_allclose(
+            loss_ref.numpy(), loss_sharded.numpy(), rtol=5e-06, atol=5e-06
+        )
+        np.testing.assert_allclose(
+            output_ref.to("float32").numpy(),
+            reorder_t.to("float32").numpy(),
+            rtol=2e-01,
+            atol=6e-02,
+        )
+
+        with paddle.no_grad():
+            reorder_q_grad = unshard_seq_load_balance(sharded_q.grad, 1)
+            reorder_k_grad = unshard_seq_load_balance(sharded_k.grad, 1)
+            reorder_v_grad = unshard_seq_load_balance(sharded_v.grad, 1)
+
+        rtol = 3e-05
+        atol = 3e-05
+        np.testing.assert_allclose(
+            q.grad.to("float32").numpy(),
+            reorder_q_grad.to("float32").numpy(),
+            rtol=rtol,
+            atol=atol,
+        )
+        np.testing.assert_allclose(
+            k.grad.to("float32").numpy(),
+            reorder_k_grad.to("float32").numpy(),
+            rtol=rtol,
+            atol=atol,
+        )
+        np.testing.assert_allclose(
+            v.grad.to("float32").numpy(),
+            reorder_v_grad.to("float32").numpy(),
+            rtol=rtol,
+            atol=atol,
+        )
+
+    def run_test_cases(self):
+        # flash attention is not supported yet for cpu
+        if os.getenv("backend") == "gpu":
+            cuda_version_main = int(paddle.version.cuda().split(".")[0])
+            device_prop_main = paddle.device.cuda.get_device_capability()[0]
+            if cuda_version_main >= 11 and device_prop_main >= 8:
+                self._test_cp_base()
+                self._test_cp_base(is_causal=False)
+
+
+if __name__ == '__main__':
+    tester = TestContextParallel()
+    tester.run_test_cases()

test/auto_parallel/converter.py

@@ -0,0 +1,104 @@
+# 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 numpy as np
+
+import paddle
+from paddle.distributed.auto_parallel.static.converter import Converter
+
+
+def test_convert():
+    rank_id = paddle.distributed.get_rank()
+    complete_tensor = np.arange(64).reshape([8, 8])
+    tensor_row = np.split(complete_tensor, 2, axis=0)
+    tensor_col = np.split(complete_tensor, 2, axis=1)
+    tensor_name = "tensor_0"
+    complete_strategy = {
+        tensor_name: {
+            "process_shape": [2],
+            "process_group": [0, 1],
+            "dims_mapping": [-1, -1],
+        }
+    }
+    row_strategy = {
+        tensor_name: {
+            "process_shape": [2],
+            "process_group": [0, 1],
+            "dims_mapping": [0, -1],
+        }
+    }
+    col_strategy = {
+        tensor_name: {
+            "process_shape": [2],
+            "process_group": [0, 1],
+            "dims_mapping": [-1, 0],
+        }
+    }
+
+    # test merge
+    tensor_dict = {tensor_name: tensor_row}
+    converter = Converter(tensor_dict, row_strategy, complete_strategy)
+    convert_tensor_dict = converter.convert()
+    assert np.equal(convert_tensor_dict[tensor_name], complete_tensor).all()
+
+    # test slice
+    tensor_dict = {tensor_name: [complete_tensor]}
+    converter = Converter(tensor_dict, complete_strategy, col_strategy)
+    convert_tensor_dict = converter.convert()
+    assert np.equal(convert_tensor_dict[tensor_name], tensor_col[rank_id]).all()
+
+    # test merge and slice
+    tensor_dict = {tensor_name: tensor_col}
+    converter = Converter(tensor_dict, col_strategy, row_strategy)
+    convert_tensor_dict = converter.convert()
+    assert np.equal(convert_tensor_dict[tensor_name], tensor_row[rank_id]).all()
+
+    # test merge and slice with prefix match
+    new_name = "tensor_1"
+    row_strategy = {
+        new_name: {
+            "process_shape": [2],
+            "process_group": [0, 1],
+            "dims_mapping": [0, -1],
+        }
+    }
+    converter = Converter(tensor_dict, col_strategy, row_strategy)
+    convert_tensor_dict = converter.convert(strict=False)
+    assert np.equal(convert_tensor_dict[new_name], tensor_row[rank_id]).all()
+
+    # test sliced_shape is 1
+    complete_tensor = np.arange(4).reshape([2, 2])
+    tensor_row = np.split(complete_tensor, 2, axis=0)
+    complete_strategy = {
+        "tensor_2": {
+            "process_shape": [2],
+            "process_group": [0, 1],
+            "dims_mapping": [-1, -1],
+        }
+    }
+    row_strategy = {
+        "tensor_2": {
+            "process_shape": [2],
+            "process_group": [0, 1],
+            "dims_mapping": [0, -1],
+        }
+    }
+    tensor_dict = {"tensor_2": [complete_tensor]}
+    converter = Converter(tensor_dict, complete_strategy, row_strategy)
+    convert_tensor_dict = converter.convert()
+    assert np.equal(convert_tensor_dict["tensor_2"], tensor_row[rank_id]).all()
+
+
+if __name__ == "__main__":
+    test_convert()

test/auto_parallel/custom_op/CMakeLists.txt

@@ -0,0 +1,16 @@
+set(LOCAL_ALL_ARCH ON)
+set(LOCAL_ALL_PLAT ON)
+if(WITH_DISTRIBUTE
+   AND WITH_GPU
+   AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_custom_op
+    MODULES
+    test_semi_auto_parallel_custom_op
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python;PADDLE_SOURCE_DIR=${PROJECT_SOURCE_DIR};WITH_ONEDNN=${WITH_ONEDNN};ONEDNN_INSTALL_DIR=${ONEDNN_INSTALL_DIR};WITH_ONEDNN=${WITH_ONEDNN};WITH_GPU=${WITH_GPU};WITH_ROCM=${WITH_ROCM};externalError_INCLUDE_DIR=${externalError_INCLUDE_DIR};PYBIND_INCLUDE_DIR=${PYBIND_INCLUDE_DIR}"
+  )
+  set_tests_properties(test_semi_auto_parallel_custom_op
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+
+endif()

test/auto_parallel/custom_op/custom_relu_op.cc

@@ -0,0 +1,138 @@
+// 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.
+
+#include <iostream>
+#include <vector>
+
+#include "paddle/extension.h"
+#include "paddle/phi/api/ext/spmd_infer.h"
+#include "paddle/phi/infermeta/spmd_rules/rules.h"
+
+#define CHECK_CPU_INPUT(x) \
+  PADDLE_ENFORCE_EQ(       \
+      x.is_cpu(), true, common::errors::Fatal(#x " must be a CPU Tensor."))
+
+template <typename data_t>
+void relu_cpu_forward_kernel(const data_t* x_data,
+                             data_t* out_data,
+                             int64_t x_numel) {
+  PADDLE_ENFORCE_NE(
+      x_data, nullptr, common::errors::Fatal("x_data is nullptr."));
+  PADDLE_ENFORCE_NE(
+      out_data, nullptr, common::errors::Fatal("out_data is nullptr."));
+  for (int64_t i = 0; i < x_numel; ++i) {
+    out_data[i] = std::max(static_cast<data_t>(0.), x_data[i]);
+  }
+}
+
+template <typename data_t>
+void relu_cpu_backward_kernel(const data_t* grad_out_data,
+                              const data_t* out_data,
+                              data_t* grad_x_data,
+                              int64_t out_numel) {
+  for (int64_t i = 0; i < out_numel; ++i) {
+    grad_x_data[i] =
+        grad_out_data[i] * (out_data[i] > static_cast<data_t>(0) ? 1. : 0.);
+  }
+}
+
+std::vector<paddle::Tensor> relu_cpu_forward(const paddle::Tensor& x) {
+  auto out = paddle::empty_like(x);
+
+  PD_DISPATCH_FLOATING_TYPES(
+      x.type(), "relu_cpu_forward", ([&] {
+        relu_cpu_forward_kernel<data_t>(
+            x.data<data_t>(), out.data<data_t>(), x.numel());
+      }));
+
+  return {out};
+}
+
+std::vector<paddle::Tensor> relu_cpu_backward(const paddle::Tensor& x,
+                                              const paddle::Tensor& out,
+                                              const paddle::Tensor& grad_out) {
+  auto grad_x = paddle::empty_like(x);
+
+  PD_DISPATCH_FLOATING_TYPES(out.type(), "relu_cpu_backward", ([&] {
+                               relu_cpu_backward_kernel<data_t>(
+                                   grad_out.data<data_t>(),
+                                   out.data<data_t>(),
+                                   grad_x.data<data_t>(),
+                                   out.size());
+                             }));
+
+  return {grad_x};
+}
+
+std::vector<paddle::Tensor> relu_cuda_forward(const paddle::Tensor& x);
+std::vector<paddle::Tensor> relu_cuda_backward(const paddle::Tensor& x,
+                                               const paddle::Tensor& out,
+                                               const paddle::Tensor& grad_out);
+
+std::vector<paddle::Tensor> ReluForward(const paddle::Tensor& x) {
+  if (x.is_cpu()) {
+    return relu_cpu_forward(x);
+  } else if (x.is_gpu()) {
+    return relu_cuda_forward(x);
+  } else {
+    PD_THROW("Not implemented.");
+  }
+}
+
+std::vector<paddle::Tensor> ReluBackward(const paddle::Tensor& x,
+                                         const paddle::Tensor& out,
+                                         const paddle::Tensor& grad_out) {
+  if (x.is_cpu()) {
+    return relu_cpu_backward(x, out, grad_out);
+  } else if (x.is_gpu()) {
+    return relu_cuda_backward(x, out, grad_out);
+  } else {
+    PD_THROW("Not implemented.");
+  }
+}
+
+phi::distributed::SpmdInfo ReluGradInferSpmd(
+    const phi::distributed::DistMetaTensor& x,
+    const phi::distributed::DistMetaTensor& out,
+    const phi::distributed::DistMetaTensor& out_grad) {
+  return phi::distributed::ElementwiseUnaryGradInferSpmd(x, out, out_grad);
+}
+
+PD_BUILD_OP(custom_relu)
+    .Inputs({"X"})
+    .Outputs({"Out"})
+    .SetKernelFn(PD_KERNEL(ReluForward))
+    .SetInferSpmdFn(
+        PD_INFER_SPMD_RULE(phi::distributed::ElementwiseUnaryInferSpmd));
+
+PD_BUILD_GRAD_OP(custom_relu)
+    .Inputs({"X", "Out", paddle::Grad("Out")})
+    .Outputs({paddle::Grad("X")})
+    .SetKernelFn(PD_KERNEL(ReluBackward))
+    .SetInferSpmdFn(PD_INFER_SPMD_RULE(ReluGradInferSpmd));
+
+PD_BUILD_OP(custom_relu_no_spmd)
+    .Inputs({"X"})
+    .Outputs({"Out"})
+    .SetKernelFn(PD_KERNEL(ReluForward));
+
+PD_BUILD_GRAD_OP(custom_relu_no_spmd)
+    .Inputs({"X", "Out", paddle::Grad("Out")})
+    .Outputs({paddle::Grad("X")})
+    .SetKernelFn(PD_KERNEL(ReluBackward));
+
+PD_REGISTER_SPMD_RULE(
+    custom_relu,
+    PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmd),
+    PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmdReverse));

test/auto_parallel/custom_op/custom_relu_op.cu

@@ -0,0 +1,93 @@
+// 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.
+
+#include "paddle/extension.h"
+
+#define CHECK_GPU_INPUT(x)                                         \
+  PADDLE_ENFORCE_EQ(                                               \
+      x.is_gpu(),                                                  \
+      true,                                                        \
+      common::errors::InvalidArgument("Input tensor `x` must be a" \
+                                      "GPU Tensor."));
+
+template <typename data_t>
+__global__ void relu_cuda_forward_kernel(const data_t* x,
+                                         data_t* y,
+                                         int64_t num) {
+  int64_t gid = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int64_t i = gid; i < num; i += blockDim.x * gridDim.x) {
+    y[i] = x[i] > static_cast<data_t>(0.) ? x[i] : static_cast<data_t>(0.);
+  }
+}
+
+template <typename data_t>
+__global__ void relu_cuda_backward_kernel(const data_t* dy,
+                                          const data_t* y,
+                                          data_t* dx,
+                                          int64_t num) {
+  int64_t gid = blockIdx.x * blockDim.x + threadIdx.x;
+  for (int64_t i = gid; i < num; i += blockDim.x * gridDim.x) {
+    dx[i] = dy[i] * (y[i] > static_cast<data_t>(0.) ? static_cast<data_t>(1.)
+                                                    : static_cast<data_t>(0.));
+  }
+}
+
+std::vector<paddle::Tensor> relu_cuda_forward(const paddle::Tensor& x) {
+  CHECK_GPU_INPUT(x);
+  auto out = paddle::empty_like(x);
+
+  PADDLE_ENFORCE_EQ(x.is_gpu(),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Input tensor `x` must be a GPU Tensor."));
+
+  int64_t numel = x.numel();
+  int64_t block = 512;
+  int64_t grid = (numel + block - 1) / block;
+  PD_DISPATCH_FLOATING_AND_HALF_TYPES(
+      x.type(), "relu_cuda_forward_kernel", ([&] {
+        relu_cuda_forward_kernel<data_t><<<grid, block, 0, x.stream()>>>(
+            x.data<data_t>(), out.data<data_t>(), numel);
+      }));
+
+  return {out};
+}
+
+std::vector<paddle::Tensor> relu_cuda_backward(const paddle::Tensor& x,
+                                               const paddle::Tensor& out,
+                                               const paddle::Tensor& grad_out) {
+  CHECK_GPU_INPUT(x);
+  CHECK_GPU_INPUT(out);
+  CHECK_GPU_INPUT(grad_out);
+  auto grad_x = paddle::empty_like(x);
+
+  PADDLE_ENFORCE_EQ(x.is_gpu(),
+                    true,
+                    common::errors::InvalidArgument(
+                        "Input tensor `x` must be a GPU Tensor."));
+
+  int64_t numel = out.numel();
+  int64_t block = 512;
+  int64_t grid = (numel + block - 1) / block;
+  PD_DISPATCH_FLOATING_AND_HALF_TYPES(
+      out.type(), "relu_cuda_backward_kernel", ([&] {
+        relu_cuda_backward_kernel<data_t><<<grid, block, 0, x.stream()>>>(
+            grad_out.data<data_t>(),
+            out.data<data_t>(),
+            grad_x.mutable_data<data_t>(x.place()),
+            numel);
+      }));
+
+  return {grad_x};
+}

test/auto_parallel/custom_op/custom_relu_setup.py

@@ -0,0 +1,31 @@
+# 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.
+
+from utils import extra_compile_args, paddle_includes
+
+from paddle.utils.cpp_extension import CUDAExtension, setup
+
+# Mac-CI don't support GPU
+Extension = CUDAExtension
+sources = ['custom_relu_op.cc', 'custom_relu_op.cu']
+
+setup(
+    name='custom_relu',
+    ext_modules=Extension(
+        sources=sources,
+        include_dirs=paddle_includes,
+        extra_compile_args=extra_compile_args,
+        verbose=True,
+    ),
+)

test/auto_parallel/custom_op/semi_auto_parallel_for_custom_op.py

@@ -0,0 +1,126 @@
+# 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 sys
+import unittest
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
+
+from semi_auto_parallel_util import SemiAutoParallelTestBase
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.static.mix_to_dist_pass import (
+    apply_mix2dist_pass,
+)
+from paddle.framework import core
+
+import custom_relu  # pylint: disable=unused-import # isort:skip
+
+assert core.contains_spmd_rule("custom_relu")
+
+
+class TestCustomOpSemiAutoParallel(SemiAutoParallelTestBase):
+    def __init__(self):
+        super().__init__()
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+
+    def check_placements(self, output, expected_placements):
+        assert output.placements == expected_placements, (
+            f"{output.placements}  vs {expected_placements}"
+        )
+
+    def test_custom_relu(self):
+        shapes = [16, 4, 4]
+        specs = ['x', None, None]
+        inputs, outputs = self.runfunc_and_check(
+            inputs_shape=shapes,
+            inputs_specs=specs,
+            op_func=custom_relu.custom_relu,
+            with_backward=True,
+        )
+        self.check_placements(outputs, [dist.Shard(0)])
+
+    def test_custom_relu_no_spmd(self):
+        shapes = [16, 4, 4]
+        specs = ['x', None, None]
+        inputs, outputs = self.runfunc_and_check(
+            inputs_shape=shapes,
+            inputs_specs=specs,
+            op_func=custom_relu.custom_relu_no_spmd,
+            with_backward=True,
+        )
+        self.check_placements(outputs, [dist.Replicate()])
+
+    def test_custom_relu_no_shard(self):
+        shapes = [16, 4, 4]
+        specs = [None, None, None]
+        inputs, outputs = self.runfunc_and_check(
+            inputs_shape=shapes,
+            inputs_specs=specs,
+            op_func=custom_relu.custom_relu,
+            with_backward=True,
+        )
+        self.check_placements(outputs, [dist.Replicate()])
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+        elif self._backend == "gpu":
+            paddle.set_device("gpu:" + str(dist.get_rank()))
+        else:
+            raise ValueError("Only support cpu or gpu backend.")
+        self.test_custom_relu_no_shard()
+        self.test_custom_relu()
+        self.test_custom_relu_no_spmd()
+
+
+class TestBuildFakeProgramWithCustomOp(unittest.TestCase):
+    def test_build_with_custom_relu(self):
+        shapes = [16, 4, 4]
+        paddle.enable_static()
+        with paddle.pir_utils.IrGuard():
+            main_program = paddle.base.Program()
+            with paddle.base.program_guard(main_program):
+                mesh = dist.ProcessMesh([0, 1], dim_names=['mp'])
+                input = paddle.static.data(
+                    name='input',
+                    shape=shapes,
+                )
+                dist_input = dist.shard_tensor(input, mesh, [dist.Shard(0)])
+                dist_out = custom_relu.custom_relu(dist_input)
+        apply_mix2dist_pass(main_program)
+
+        self.assertTrue(dist_out.is_dist_dense_tensor_type())
+        self.assertEqual(dist_out._local_shape, [16 // 2, 4, 4])
+        self.assertEqual(dist_out.dist_attr().dims_mapping, [0, -1, -1])
+        self.assertEqual(dist_out.dist_attr().process_mesh, mesh)
+        op_dist_attr = dist_out.get_defining_op().dist_attr
+        self.assertEqual(op_dist_attr.process_mesh, mesh)
+        self.assertEqual(
+            op_dist_attr.result(0).as_tensor_dist_attr().dims_mapping,
+            [0, -1, -1],
+        )
+        self.assertEqual(
+            op_dist_attr.operand(0).as_tensor_dist_attr().dims_mapping,
+            [0, -1, -1],
+        )
+
+
+if __name__ == '__main__':
+    TestCustomOpSemiAutoParallel().run_test_case()
+    TestBuildFakeProgramWithCustomOp().test_build_with_custom_relu()

test/auto_parallel/custom_op/test_semi_auto_parallel_custom_op.py

@@ -0,0 +1,52 @@
+# 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 site
+import sys
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+from paddle.utils.cpp_extension.extension_utils import run_cmd
+
+
+class TestCustomOp(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=200, nnode=1)
+        self._default_envs = {"dtype": "float32", "seed": "2023"}
+        self._changeable_envs = {"backend": ["cpu", "gpu"]}
+        cur_dir = os.path.dirname(os.path.abspath(__file__))
+        # compile, install the custom op egg into site-packages under background
+        if os.name == 'nt':
+            cmd = f'cd /d {cur_dir} && python custom_relu_setup.py install'
+        else:
+            site_dir = site.getsitepackages()[0]
+            cmd = f'cd {cur_dir} && {sys.executable} custom_relu_setup.py install --install-lib={site_dir}'
+        run_cmd(cmd)
+
+    # test dynamic auto parallel run
+    def test_dynamic_auto_parallel(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_for_custom_op.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/custom_op/utils.py

@@ -0,0 +1,50 @@
+# 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
+from pathlib import Path
+from site import getsitepackages
+
+from paddle.utils.cpp_extension.extension_utils import (
+    _get_all_paddle_includes_from_include_root,
+)
+
+# Test for extra compile args
+extra_cc_args = ['-w', '-g']
+extra_nvcc_args = ['-O3']
+extra_compile_args = {'cc': extra_cc_args, 'nvcc': extra_nvcc_args}
+
+
+def get_paddle_includes():
+    env_dict = os.environ
+    paddle_includes = []
+    paddle_includes.append(f"{env_dict.get('PADDLE_SOURCE_DIR')}")
+
+    # onednn
+    if env_dict.get("WITH_ONEDNN") == 'ON':
+        paddle_includes.append(f"{env_dict.get('ONEDNN_INSTALL_DIR')}/include")
+    if env_dict.get("WITH_GPU") == 'ON' or env_dict.get("WITH_ROCM") == 'ON':
+        paddle_includes.append(f"{env_dict.get('externalError_INCLUDE_DIR')}")
+    paddle_includes.append(f"{env_dict.get('PYBIND_INCLUDE_DIR')}")
+
+    for site_packages_path in getsitepackages():
+        paddle_include_dir = Path(site_packages_path) / "paddle/include"
+        paddle_includes.extend(
+            _get_all_paddle_includes_from_include_root(str(paddle_include_dir))
+        )
+
+    return paddle_includes
+
+
+paddle_includes = get_paddle_includes()

test/auto_parallel/dtensor_from_local_api.py

@@ -0,0 +1,77 @@
+# 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 os
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.api import dtensor_from_local
+
+
+class TestDtensorFromLocalAPI:
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._seeds = eval(os.getenv("seeds"))
+        self._backend = os.getenv("backend")
+        self._shard = eval(os.getenv("shard"))
+        self._mesh = dist.ProcessMesh([[0, 1]], dim_names=["x", "y"])
+
+    def run_test_cases(self):
+        self.test_case_forward_backward()
+
+    def test_case_forward_backward(self):
+        a = paddle.ones([10, 20])
+        a.stop_gradient = False
+
+        tensor1 = a + 3
+        assert not tensor1.is_dist()
+        tensor1.register_hook(self.check_grad_mesh(None, None))
+
+        mesh = self._mesh
+        tensor2 = dtensor_from_local(
+            tensor1, mesh, [dist.Shard(0), dist.Replicate()]
+        )
+
+        assert tensor2.is_dist()
+        assert tensor2.process_mesh == mesh
+        assert tensor2.placements == [dist.Shard(0), dist.Replicate()]
+        tensor2.register_hook(
+            self.check_grad_mesh(mesh, [dist.Shard(0), dist.Replicate()])
+        )
+
+        tensor3 = tensor2 * 3
+        tensor3.register_hook(
+            self.check_grad_mesh(mesh, [dist.Shard(0), dist.Replicate()])
+        )
+        tensor4 = tensor3 + 4
+
+        tensor4.backward()
+
+    def check_grad_mesh(self, mesh, placements):
+        def _check_mesh(grad):
+            if mesh is None and placements is None:
+                assert not grad.is_dist(), "grad.is_dist() is not False"
+            else:
+                assert grad.process_mesh == mesh, (
+                    "grad.process_mesh is not equal to mesh"
+                )
+                assert grad.placements == placements, (
+                    "grad.placements is not equal to placements"
+                )
+
+        return _check_mesh
+
+
+if __name__ == '__main__':
+    TestDtensorFromLocalAPI().run_test_cases()

test/auto_parallel/dtensor_to_local_api.py

@@ -0,0 +1,68 @@
+# 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 os
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import Partial
+from paddle.distributed.auto_parallel.api import dtensor_to_local
+
+
+class TestDtensorToLocalAPI:
+    def __init__(self):
+        self._shape = eval(os.getenv("shape"))
+        self._dtype = os.getenv("dtype")
+        self._seeds = eval(os.getenv("seeds"))
+        self._backend = os.getenv("backend")
+        self._shard = eval(os.getenv("shard"))
+        self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def run_test_cases(self):
+        self.test_case_forward_backward()
+
+    def test_case_forward_backward(self):
+        a = paddle.ones(self._shape)
+        a.stop_gradient = False
+
+        input_tensor = dist.shard_tensor(a, self._mesh, [Partial()])
+        input_tensor.register_hook(
+            self.check_grad_mesh(
+                input_tensor.process_mesh, input_tensor.placements
+            )
+        )
+
+        tensor1 = dtensor_to_local(
+            input_tensor, input_tensor.process_mesh, input_tensor.placements
+        )
+        assert not tensor1.is_dist()
+
+        tensor2 = tensor1 + 2
+        tensor3 = tensor2 * 3
+        tensor3.register_hook(self.check_grad_mesh(None, None))
+        tensor3.backward()
+
+    def check_grad_mesh(self, org_mesh, org_placements):
+        def _check_mesh(grad):
+            if hasattr(grad, "process_mesh") and hasattr(grad, "placements"):
+                assert grad.process_mesh == org_mesh
+                assert grad.placements == org_placements
+            else:
+                assert org_mesh is None and org_placements is None
+
+        return _check_mesh
+
+
+if __name__ == '__main__':
+    TestDtensorToLocalAPI().run_test_cases()

test/auto_parallel/dy_local_view_compute.py

@@ -0,0 +1,97 @@
+# 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 paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.api import (
+    dtensor_from_local,
+    dtensor_to_local,
+)
+
+
+class TestLocalViewCompute:
+    def __init__(self):
+        self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def run_test_cases(self):
+        self.test_local_view_compute()
+
+    def masked_lm_loss_func(self, pred, label, ignored_idx=-100):
+        pred_sub = pred[:, 0:1]  # shape [B,1]
+        label_float = paddle.cast(label, 'float32')  # shape [B,1]
+
+        raw_loss = paddle.abs(pred_sub - label_float)
+
+        lossmask = label != ignored_idx
+        lossmask_ = lossmask.reshape([-1]).cast('float32')
+        raw_loss_flat = raw_loss.reshape([-1]).cast('float32')
+
+        masked_lm_loss_sum = paddle.sum(raw_loss_flat * lossmask_)
+        valid_count = paddle.sum(lossmask_)
+
+        loss = masked_lm_loss_sum / (valid_count + 1e-8)
+        return loss
+
+    def local_view_compute(self, local_pred, local_label):
+        # do not use dist.shard_tensor here
+        local_pred = local_pred + 1
+        local_loss = self.masked_lm_loss_func(
+            local_pred, local_label, ignored_idx=-100
+        )
+
+        return local_loss
+
+    def test_local_view_compute(self):
+        dist.init_parallel_env()
+        cur_rank = dist.get_rank()
+
+        # prepare data and label for mask_lm_loss
+        if cur_rank == 0:
+            pred = paddle.to_tensor([[1.0, 2.0], [4.0, 4.0]], dtype='float32')
+            label = paddle.to_tensor([[1], [3]], dtype='int64')
+        elif cur_rank == 1:
+            pred = paddle.to_tensor([[2.0, 2.0], [7.0, 8.0]], dtype='float32')
+            label = paddle.to_tensor([[2], [-100]], dtype='int64')
+
+        local_result = self.local_view_compute(pred.clone(), label.clone())
+
+        dist_pred = dist.shard_tensor(pred, self._mesh, [dist.Replicate()])
+        dist_label = dist.shard_tensor(label, self._mesh, [dist.Replicate()])
+
+        local_pred = dtensor_to_local(
+            dist_pred, dist_pred.process_mesh, dist_pred.placements
+        )
+        local_label = dtensor_to_local(
+            dist_label, dist_label.process_mesh, dist_label.placements
+        )
+
+        local_pred = local_pred + 1
+        local_loss = self.masked_lm_loss_func(
+            local_pred, local_label, ignored_idx=-100
+        )
+
+        assert local_result == local_loss, "local_result != local_loss"
+
+        tensor_list = []
+        dist.all_gather(tensor_list, local_loss)
+        loss_sum = paddle.sum(paddle.stack(tensor_list))
+        dist_loss = dtensor_from_local(
+            local_loss, self._mesh, [dist.Partial(dist.ReduceType.kRedSum)]
+        )
+
+        assert loss_sum == dist_loss, "loss_sum != dist_loss"
+
+
+if __name__ == '__main__':
+    TestLocalViewCompute().run_test_cases()

test/auto_parallel/end_to_end/CMakeLists.txt

@@ -0,0 +1,20 @@
+# This file is generated by ${PADDLE_ROOT}/tools/gen_ut_cmakelists.py.
+# Please don't modify this file manually.
+# If you need to change unittests in this file, please modify testslist.csv in the current directory
+# and then run the command `python3 ${PADDLE_ROOT}/tools/gen_ut_cmakelists.py -f ${CURRENT_DIRECTORY}/testslist.csv`
+set(LOCAL_ALL_ARCH ON)
+set(LOCAL_ALL_PLAT ON)
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_e2e_co_shard_8cards MODULES test_e2e_co_shard_8cards ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_e2e_co_shard_8cards PROPERTIES TIMEOUT "120" LABELS
+                                                           "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_e2e_co_shard MODULES test_e2e_co_shard ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_e2e_co_shard PROPERTIES TIMEOUT "120" LABELS
+                                                    "RUN_TYPE=HYBRID")
+endif()

test/auto_parallel/end_to_end/argsort_co_shard.py

@@ -0,0 +1,258 @@
+# 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class ArgSortTestCase:
+    def __init__(
+        self,
+        input_shape: list[int],
+        input_placements: list[dist.Placement],
+        axis: int,
+        indices_placements: list[dist.Placement],
+        slice_funtor: Callable[[int], Any] | None = None,
+    ):
+        self.input_shape = input_shape
+        self.input_placements = input_placements
+        self.axis = axis
+        self.indices_placements = indices_placements
+        self.slice_funtor = slice_funtor
+        self.descending = False
+        self.stable = False
+
+
+class ArgSortGradTestCase:
+    def __init__(
+        self,
+        input_shape: list[int],
+        x_placements: list[dist.Placement],
+        axis: int,
+        out_grad_placements: list[dist.Placement],
+        x_grad_placements: list[dist.Placement],
+    ):
+        self.input_shape = input_shape
+        self.x_placements = x_placements
+        self.out_grad_placements = out_grad_placements
+        self.axis = axis
+        self.x_grad_placements = x_grad_placements
+        self.descending = False
+        self.stable = False
+
+
+class TestArgSortCoShard:
+    def setUp(self):
+        self.mesh = dist.ProcessMesh(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=['x', 'y', 'z']
+        )
+        self.test_cases_forward = [
+            # test flatten
+            ArgSortTestCase(
+                [16, 32, 48],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                -1,
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+            ),
+            ArgSortTestCase(
+                [16, 32, 48],
+                [
+                    dist.Shard(
+                        0,
+                    ),
+                    dist.Shard(2, shard_order=0),
+                    dist.Shard(2, shard_order=1),
+                ],
+                2,
+                [
+                    dist.Shard(
+                        0,
+                    ),
+                    dist.Replicate(),
+                    dist.Replicate(),
+                ],
+            ),
+            ArgSortTestCase(
+                [10, 32, 48, 24],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                    dist.Replicate(),
+                ],
+                1,
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                    dist.Replicate(),
+                ],
+            ),
+        ]
+        self.test_cases_backward = [
+            # test flatten
+            ArgSortGradTestCase(
+                [16, 32, 48],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                -1,
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+            ),
+            ArgSortGradTestCase(
+                [16, 32, 48],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(2),
+                ],
+                2,
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(2),
+                ],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+            ArgSortGradTestCase(
+                [10, 32, 48, 24],
+                [
+                    dist.Shard(0),
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                    dist.Replicate(),
+                ],
+                1,
+                [
+                    dist.Shard(0),
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                    dist.Replicate(),
+                ],
+                [
+                    dist.Shard(0),
+                    dist.Replicate(),
+                    dist.Replicate(),
+                    dist.Replicate(),
+                ],
+            ),
+        ]
+
+    def run_test_case_forward(self, test_case: ArgSortTestCase):
+        a = paddle.rand(test_case.input_shape, "float32")
+        input_placements = test_case.input_placements
+        input = dist.shard_tensor(a, self.mesh, input_placements)
+        out = paddle.argsort(
+            input, test_case.axis, test_case.descending, test_case.stable
+        )
+        case_info = f"input_shape: {test_case.input_shape}, input_placements: {input_placements}, axis: {test_case.axis}"
+        # Verify output shape
+        np.testing.assert_equal(
+            out.shape,
+            test_case.input_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.input_shape}, Actual: {out.shape}",
+        )
+
+        # Verify placements
+        assert out.placements
+        for actual, expected in zip(
+            out.placements, test_case.indices_placements
+        ):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.indices_placements}, Actual: {out.placements}",
+            )
+        # Verify local_value if given
+        if test_case.slice_funtor:
+            idx = dist.get_rank()
+            np.testing.assert_equal(
+                out._local_value().numpy().flatten(),
+                a[test_case.slice_funtor(idx)].numpy().flatten(),
+                err_msg=f"Local values mismatch when {case_info}.",
+            )
+
+    def run_test_case_backward(self, test_case: ArgSortGradTestCase):
+        a = paddle.rand(test_case.input_shape, "float32")
+        a.stop_gradient = False
+        input = dist.shard_tensor(a, self.mesh, test_case.x_placements)
+        out = paddle.argsort(
+            input, test_case.axis, test_case.descending, test_case.stable
+        )
+
+        out_grad = paddle.ones(out.shape, "float32")
+        out_grad = dist.shard_tensor(
+            out_grad, self.mesh, test_case.out_grad_placements
+        )
+
+        (x_grad,) = paddle.grad([out], input, [out_grad])
+
+        case_info = f"input_shape: {test_case.input_shape}, axis: {test_case.axis}, x_placements: {test_case.x_placements}, out_grad_placements: {test_case.out_grad_placements}"
+        # Verify output shape
+        np.testing.assert_equal(
+            x_grad.shape,
+            test_case.input_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.input_shape}, Actual: {x_grad.shape}",
+        )
+
+        # Verify placements
+        assert x_grad.placements
+        for actual, expected in zip(
+            x_grad.placements, test_case.x_grad_placements
+        ):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.x_grad_placements}, Actual: {x_grad.placements}",
+            )
+
+    def run_all_tests(self):
+        self.setUp()
+        for test_case in self.test_cases_forward:
+            self.run_test_case_forward(test_case)
+
+
+if __name__ == '__main__':
+    TestArgSortCoShard().run_all_tests()

test/auto_parallel/end_to_end/co_shard.py

@@ -0,0 +1,207 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestCoShard:
+    def basic_interface_case(self):
+        shard = dist.Shard(0, shard_order=0)
+        np.testing.assert_equal(shard, dist.Shard(dim=0, shard_order=0))
+
+        shard = dist.Shard(0, split_factor=2)
+        np.testing.assert_equal(shard, dist.Shard(dim=0, split_factor=2))
+
+    def run_test_case_0(self):
+        a = paddle.to_tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+        mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+
+        placements = [
+            dist.Shard(0, shard_order=0),
+            dist.Shard(0, shard_order=1),
+        ]
+        input = dist.shard_tensor(a, mesh, placements)
+
+        idx = dist.get_rank()
+        np.testing.assert_equal(
+            input._local_value().numpy().flatten(), a[idx].numpy().flatten()
+        )
+
+        reshard_placements = [dist.Replicate(), dist.Replicate()]
+        out = dist.reshard(input, mesh, reshard_placements)
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(), a.numpy().flatten()
+        )
+
+        reshard_placements = [dist.Shard(0), dist.Replicate()]
+        out = dist.reshard(input, mesh, reshard_placements)
+        new_idx = idx // 2 * 2
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(),
+            a[new_idx : new_idx + 2].numpy().flatten(),
+        )
+
+        reshard_placements = [dist.Replicate(), dist.Shard(0)]
+        out = dist.reshard(input, mesh, reshard_placements)
+        new_idx = idx % 2 * 2
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(),
+            a[new_idx : new_idx + 2].numpy().flatten(),
+        )
+
+    def run_test_case_1(self):
+        a = paddle.to_tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+        mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+        placements = [
+            dist.Shard(0, shard_order=1),
+            dist.Shard(0, shard_order=0),
+        ]
+        input = dist.shard_tensor(a, mesh, placements)
+
+        idx = dist.get_rank()
+        new_idx = idx % 2 * 2 + idx // 2
+        np.testing.assert_equal(
+            input._local_value().numpy().flatten(), a[new_idx].numpy().flatten()
+        )
+
+        reshard_placements = [dist.Replicate(), dist.Replicate()]
+        out = dist.reshard(input, mesh, reshard_placements)
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(), a.numpy().flatten()
+        )
+
+        reshard_placements = [dist.Shard(0), dist.Replicate()]
+        out = dist.reshard(input, mesh, reshard_placements)
+        new_idx = idx // 2 * 2
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(),
+            a[new_idx : new_idx + 2].numpy().flatten(),
+        )
+
+        reshard_placements = [dist.Replicate(), dist.Shard(0)]
+        out = dist.reshard(input, mesh, reshard_placements)
+        new_idx = idx % 2 * 2
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(),
+            a[new_idx : new_idx + 2].numpy().flatten(),
+        )
+
+    def run_test_case_2(self):
+        mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+
+        # dense tensor
+        a = paddle.to_tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+
+        placements = [dist.Shard(0, split_factor=2), dist.Replicate()]
+        # distributed tensor
+        input = dist.shard_tensor(a, mesh, placements)
+
+        idx = dist.get_rank()
+        if idx == 0 or idx == 1:
+            golden = np.array([[1, 2], [5, 6]])
+        else:
+            golden = np.array([[3, 4], [7, 8]])
+        np.testing.assert_equal(
+            input._local_value().numpy().flatten(), golden.flatten()
+        )
+
+        reshard_placements = [dist.Replicate(), dist.Replicate()]
+        out = dist.reshard(input, mesh, reshard_placements)
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(), a.numpy().flatten()
+        )
+
+        reshard_placements = [dist.Shard(0), dist.Replicate()]
+        out = dist.reshard(input, mesh, reshard_placements)
+        new_idx = idx // 2 * 2
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(),
+            a[new_idx : new_idx + 2].numpy().flatten(),
+        )
+
+        reshard_placements = [dist.Replicate(), dist.Shard(0)]
+        out = dist.reshard(input, mesh, reshard_placements)
+        new_idx = idx % 2 * 2
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(),
+            a[new_idx : new_idx + 2].numpy().flatten(),
+        )
+
+    def run_test_case_3(self):
+        a = paddle.to_tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+        mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+        placements = [dist.Shard(0), dist.Shard(1)]
+        input = dist.shard_tensor(a, mesh, placements)
+
+        reshard_placements = [
+            dist.Shard(0, shard_order=0),
+            dist.Shard(0, shard_order=1),
+        ]
+        out = dist.reshard(input, mesh, reshard_placements)
+        np.testing.assert_equal(
+            out._local_value().numpy().flatten(),
+            a[dist.get_rank()].numpy().flatten(),
+        )
+        np.testing.assert_equal(
+            out.placements[0], dist.Shard(dim=0, shard_order=0)
+        )
+        np.testing.assert_equal(
+            out.placements[1], dist.Shard(dim=0, shard_order=1)
+        )
+
+    def run_test_case_4(self):
+        a = paddle.to_tensor([[1, 2], [3, 4], [5, 6], [7, 8]], dtype='float32')
+        mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+        placements = [dist.Shard(0), dist.Shard(1)]
+        input = dist.shard_tensor(a, mesh, placements)
+
+        out = paddle.reshape(input, [-1])
+        np.testing.assert_equal(out.shape, [8])
+        np.testing.assert_equal(
+            out.placements[0], dist.Shard(dim=0, shard_order=0)
+        )
+        np.testing.assert_equal(
+            out.placements[1], dist.Shard(dim=0, shard_order=1)
+        )
+        np.testing.assert_equal(
+            out._local_value().numpy(), a[dist.get_rank()].numpy().flatten()
+        )
+
+        relu_out = paddle.nn.ReLU()(out)
+        np.testing.assert_equal(
+            relu_out.placements[0], dist.Shard(dim=0, shard_order=0)
+        )
+        np.testing.assert_equal(
+            relu_out.placements[1], dist.Shard(dim=0, shard_order=1)
+        )
+
+        # test fallback to shard by one dim.
+        add_out = paddle.add(relu_out, relu_out)
+        np.testing.assert_equal(add_out.placements[0], dist.Shard(dim=0))
+        np.testing.assert_equal(add_out.placements[1], dist.Replicate())
+
+    def run_test_case_main(self):
+        self.basic_interface_case()
+        self.run_test_case_0()
+        self.run_test_case_1()
+        self.run_test_case_2()
+        self.run_test_case_3()
+        self.run_test_case_4()
+
+
+if __name__ == '__main__':
+    TestCoShard().run_test_case_main()

test/auto_parallel/end_to_end/elementwise_co_shard.py

@@ -0,0 +1,79 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestElementWiseCoShard:
+    def run_unary_case_0(self):
+        mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+        placements = [
+            dist.Shard(0, shard_order=0),
+            dist.Shard(0, shard_order=1),
+        ]
+
+        x = paddle.to_tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]], dtype="float32"
+        )
+        x = dist.shard_tensor(x, mesh, placements)
+        # paddle.round
+        out = paddle.round(x)
+
+        np.testing.assert_equal(out.shape, [4, 2])
+        assert out.placements, "The output should be a DistTensor"
+        np.testing.assert_equal(
+            out.placements[0], dist.Shard(dim=0, shard_order=0)
+        )
+        np.testing.assert_equal(
+            out.placements[1], dist.Shard(dim=0, shard_order=1)
+        )
+
+    def run_unary_case_with_partial(self):
+        mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+        # TODO(ooooo): Test co_shard when matmul is supported.
+        x_placements = [
+            dist.Shard(0),
+            dist.Shard(1),
+        ]
+
+        x = paddle.to_tensor(
+            [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], dtype="float32"
+        )
+        y = paddle.to_tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]], dtype="float32"
+        )
+        x = dist.shard_tensor(x, mesh, x_placements)
+        y = dist.shard_tensor(
+            y, mesh, [dist.Replicate() for _ in range(mesh.ndim)]
+        )
+        # Generate partial placement
+        matmul_out = paddle.matmul(x, y)
+        # paddle.cast
+        out = paddle.cast(matmul_out, 'float64')
+
+        np.testing.assert_equal(out.shape, [2, 2])
+        assert out.placements, "The output should be a DistTensor"
+        np.testing.assert_equal(out.placements[0], dist.Shard(0))
+        np.testing.assert_equal(out.placements[1], dist.Partial())
+
+    def run_test_case_main(self):
+        self.run_unary_case_0()
+        self.run_unary_case_with_partial()
+
+
+if __name__ == '__main__':
+    TestElementWiseCoShard().run_test_case_main()

test/auto_parallel/end_to_end/index_select_co_shard.py

@@ -0,0 +1,334 @@
+# 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.
+from __future__ import annotations
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class IndexSelectTestCase:
+    def __init__(
+        self,
+        x_shape: list[int],
+        x_placements: list[dist.Placement],
+        index_shape: list[int],
+        index_placements: list[dist.Placement],
+        axis: int,
+        out_shape: list[int],
+        out_placements: list[dist.Placement],
+    ):
+        self.x_shape = x_shape
+        self.x_placements = x_placements
+        self.index_shape = index_shape
+        self.index_placements = index_placements
+        self.axis = axis
+        self.out_shape = out_shape
+        self.out_placements = out_placements
+
+
+class IndexSelectGradTestCase:
+    def __init__(
+        self,
+        x_shape: list[int],
+        x_placements: list[dist.Placement],
+        index_shape: list[int],
+        index_placements: list[dist.Placement],
+        axis: int,
+        out_grad_shape: list[int],
+        out_grad_placements: list[dist.Placement],
+        x_grad_placements: list[dist.Placement],
+    ):
+        self.x_shape = x_shape
+        self.x_placements = x_placements
+        self.index_shape = index_shape
+        self.index_placements = index_placements
+        self.axis = axis
+        self.out_grad_shape = out_grad_shape
+        self.out_grad_placements = out_grad_placements
+        self.x_grad_placements = x_grad_placements
+
+
+class TestIndexSelectCoShard:
+    def setUp(self):
+        self.mesh = dist.ProcessMesh(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=['x', 'y', 'z']
+        )
+        self.test_cases_forward = [
+            IndexSelectTestCase(
+                [8, 16, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [8],
+                [dist.Replicate(), dist.Replicate(), dist.Replicate()],
+                1,
+                [8, 8, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+            IndexSelectTestCase(
+                [8, 16, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [8],
+                [dist.Replicate(), dist.Replicate(), dist.Shard(0)],
+                1,
+                [8, 8, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+            ),
+            IndexSelectTestCase(
+                [8, 16, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [8],
+                [dist.Shard(0), dist.Replicate(), dist.Replicate()],
+                1,
+                [8, 8, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+            IndexSelectTestCase(
+                [8, 16, 32],
+                [dist.Replicate(), dist.Replicate(), dist.Shard(0)],
+                [8],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+                1,
+                [8, 8, 32],
+                [
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                    dist.Shard(0),
+                ],
+            ),
+            IndexSelectTestCase(
+                [8, 16, 32],
+                [dist.Shard(0), dist.Replicate(), dist.Replicate()],
+                [8],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+                1,
+                [8, 8, 32],
+                [dist.Shard(0), dist.Shard(1), dist.Replicate()],
+            ),
+        ]
+        self.test_cases_backward = [
+            IndexSelectGradTestCase(
+                [8, 16, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [8],
+                [dist.Replicate(), dist.Replicate(), dist.Replicate()],
+                1,
+                [8, 8, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+            IndexSelectGradTestCase(
+                [8, 16, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [8],
+                [dist.Replicate(), dist.Replicate(), dist.Shard(0)],
+                1,
+                [8, 8, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Partial(),
+                ],
+            ),
+            IndexSelectGradTestCase(
+                [8, 16, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [8],
+                [dist.Shard(0), dist.Replicate(), dist.Replicate()],
+                1,
+                [8, 8, 32],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+            IndexSelectGradTestCase(
+                [8, 16, 32],
+                [dist.Replicate(), dist.Replicate(), dist.Shard(0)],
+                [8],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+                1,
+                [8, 8, 32],
+                [
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                    dist.Shard(0),
+                ],
+                [dist.Partial(), dist.Partial(), dist.Shard(0)],
+            ),
+            IndexSelectGradTestCase(
+                [8, 16, 32],
+                [dist.Shard(0), dist.Replicate(), dist.Replicate()],
+                [8],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+                1,
+                [8, 8, 32],
+                [dist.Shard(0), dist.Shard(1), dist.Replicate()],
+                [dist.Shard(0), dist.Partial(), dist.Replicate()],
+            ),
+        ]
+
+    def run_test_case_forward(self, test_case: IndexSelectTestCase):
+        x = paddle.rand(test_case.x_shape, "float32")
+        x_placements = test_case.x_placements
+        x = dist.shard_tensor(x, self.mesh, x_placements)
+        index = paddle.randint(
+            0,
+            test_case.x_shape[test_case.axis],
+            test_case.index_shape,
+            dtype="int32",
+        )
+        index_placements = test_case.index_placements
+        index = dist.shard_tensor(index, self.mesh, index_placements)
+
+        out = paddle.index_select(x, index, test_case.axis)
+        case_info = f"x_shape: {test_case.x_shape}, x_placements: {x_placements}, index_shape: {test_case.index_shape}, index_placements: {index_placements}, axis: {test_case.axis}"
+        # Verify output shape
+        np.testing.assert_equal(
+            out.shape,
+            test_case.out_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.out_shape}, Actual: {out.shape}",
+        )
+
+        # Verify placements
+        assert out.placements
+        for actual, expected in zip(out.placements, test_case.out_placements):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.out_placements}, Actual: {out.placements}",
+            )
+
+    def run_test_case_backward(self, test_case: IndexSelectGradTestCase):
+        x = paddle.rand(test_case.x_shape, "float32")
+        x.stop_gradient = False
+        x_placements = test_case.x_placements
+        x = dist.shard_tensor(x, self.mesh, x_placements)
+
+        index = paddle.randint(
+            0,
+            test_case.x_shape[test_case.axis],
+            test_case.index_shape,
+            dtype="int32",
+        )
+        index_placements = test_case.index_placements
+        index = dist.shard_tensor(index, self.mesh, index_placements)
+
+        out = paddle.index_select(x, index, test_case.axis)
+
+        out_grad = paddle.ones(out.shape, "float32")
+        out_grad = dist.shard_tensor(
+            out_grad, self.mesh, test_case.out_grad_placements
+        )
+
+        (x_grad,) = paddle.grad([out], x, [out_grad])
+
+        case_info = f"x_shape: {test_case.x_shape}, x_placements: {test_case.x_placements}, index_shape: {test_case.index_shape}, index_placements: {test_case.index_placements}, axis: {test_case.axis}, out_grad_shape: {test_case.out_grad_shape}, out_grad_placements: {test_case.out_grad_placements}"
+        # Verify output shape
+        np.testing.assert_equal(
+            x_grad.shape,
+            test_case.x_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.x_shape}, Actual: {x_grad.shape}",
+        )
+
+        # Verify placements
+        assert x_grad.placements
+        for actual, expected in zip(
+            x_grad.placements, test_case.x_grad_placements
+        ):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.x_grad_placements}, Actual: {x_grad.placements}",
+            )
+
+    def run_all_tests(self):
+        self.setUp()
+        for test_case in self.test_cases_forward:
+            self.run_test_case_forward(test_case)
+
+
+if __name__ == '__main__':
+    TestIndexSelectCoShard().run_all_tests()

test/auto_parallel/end_to_end/matmul_co_shard.py

@@ -0,0 +1,170 @@
+# 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.
+from __future__ import annotations
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import Partial, Replicate, Shard
+
+
+class MatmulTestCase:
+    def __init__(
+        self,
+        x_shape: list[int],
+        x_placements: list[dist.Placement],
+        y_shape: list[int],
+        y_placements: list[dist.Placement],
+        trans_x: bool,
+        trans_y: bool,
+        output_shape: list[int],
+        output_placements: list[dist.Placement],
+    ):
+        self.x_shape = x_shape
+        self.x_placements = x_placements
+        self.y_shape = y_shape
+        self.y_placements = y_placements
+        self.trans_x = trans_x
+        self.trans_y = trans_y
+        self.output_shape = output_shape
+        self.output_placements = output_placements
+
+
+class TestMatmulCoShard:
+    def setUp(self):
+        self.mesh = dist.ProcessMesh(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=['x', 'y', 'z']
+        )
+        self.test_cases_forward = [
+            # test flatten
+            MatmulTestCase(
+                [64, 32],
+                [Shard(0, shard_order=0), Shard(0, shard_order=1), Replicate()],
+                [32, 48],
+                [Replicate(), Replicate(), Shard(1)],
+                False,
+                False,
+                [64, 48],
+                [Shard(0, shard_order=0), Shard(0, shard_order=1), Shard(1)],
+            ),
+            MatmulTestCase(
+                [64, 32],
+                [Replicate(), Replicate(), Replicate()],
+                [32, 48],
+                [Shard(0, shard_order=0), Shard(0, shard_order=1), Shard(1)],
+                False,
+                False,
+                [64, 48],
+                [Partial(), Partial(), Shard(1)],
+            ),
+            MatmulTestCase(
+                [64, 32],
+                [Shard(0, shard_order=1), Shard(0, shard_order=1), Shard(1)],
+                [32, 48],
+                [Replicate(), Replicate(), Replicate()],
+                False,
+                False,
+                [64, 48],
+                [Shard(0, shard_order=0), Shard(0, shard_order=1), Partial()],
+            ),
+            MatmulTestCase(
+                [64, 32],
+                [Shard(0, shard_order=1), Shard(0, shard_order=1), Shard(1)],
+                [32, 48],
+                [Shard(0), Replicate(), Replicate()],
+                False,
+                False,
+                [64, 48],
+                [Shard(0, shard_order=0), Shard(0, shard_order=1), Partial()],
+            ),
+            MatmulTestCase(
+                [512, 48, 64, 32],
+                [Shard(0, shard_order=1), Shard(0, shard_order=1), Shard(1)],
+                [1, 32, 48],
+                [Replicate(), Replicate(), Replicate()],
+                False,
+                False,
+                [512, 48, 64, 48],
+                [Shard(0, shard_order=0), Shard(0, shard_order=1), Shard(1)],
+            ),
+            MatmulTestCase(
+                [512, 48, 32, 64],
+                [Shard(0), Shard(2, shard_order=0), Shard(2, shard_order=1)],
+                [1, 32, 48],
+                [Replicate(), Replicate(), Shard(2)],
+                True,
+                False,
+                [512, 48, 64, 48],
+                [Shard(0), Partial(), Shard(3)],
+            ),
+            MatmulTestCase(
+                [512, 48, 64, 32],
+                [Shard(0), Shard(2, shard_order=0), Shard(2, shard_order=1)],
+                [1, 48, 32],
+                [Shard(1), Replicate(), Replicate()],
+                False,
+                True,
+                [512, 48, 64, 48],
+                [Shard(0), Shard(2, shard_order=0), Shard(2, shard_order=1)],
+            ),
+            MatmulTestCase(
+                [512, 48, 32, 64],
+                [Shard(2, shard_order=0), Shard(2, shard_order=1), Shard(3)],
+                [1, 48, 32],
+                [Shard(1, shard_order=0), Shard(1, shard_order=1), Shard(2)],
+                True,
+                True,
+                [512, 48, 64, 48],
+                [Shard(3, shard_order=0), Shard(3, shard_order=1), Shard(2)],
+            ),
+        ]
+
+    def run_test_case_forward(self, test_case: MatmulTestCase):
+        x = paddle.rand(test_case.x_shape, "float32")
+        x_placements = test_case.x_placements
+        x = dist.shard_tensor(x, self.mesh, x_placements)
+
+        y = paddle.rand(test_case.y_shape, "float32")
+        y_placements = test_case.y_placements
+        y = dist.shard_tensor(y, self.mesh, y_placements)
+
+        out = paddle.matmul(x, y, test_case.trans_x, test_case.trans_y)
+        case_info = f"x_shape: {test_case.x_shape}, x_placements: {x_placements}, y_shape: {test_case.y_shape}, y_placements: {test_case.y_placements}, trans_x: {test_case.trans_x}, trans_y: {test_case.trans_y}"
+        # Verify output shape
+        np.testing.assert_equal(
+            out.shape,
+            test_case.output_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.output_shape}, Actual: {out.shape}",
+        )
+
+        # Verify placements
+        assert out.placements
+        for actual, expected in zip(
+            out.placements, test_case.output_placements
+        ):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.output_placements}, Actual: {out.placements}",
+            )
+
+    def run_all_tests(self):
+        self.setUp()
+        for test_case in self.test_cases_forward:
+            self.run_test_case_forward(test_case)
+
+
+if __name__ == '__main__':
+    TestMatmulCoShard().run_all_tests()

test/auto_parallel/end_to_end/reshape_co_shard.py

@@ -0,0 +1,218 @@
+# 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class ReshapeTestCase:
+    def __init__(
+        self,
+        input_shape: list[int],
+        input_placements: list[dist.Placement],
+        target_shape: list[int],
+        output_placements: list[dist.Placement],
+        slice_funtor: Callable[[int], Any] | None = None,
+    ):
+        self.input_shape = input_shape
+        self.input_placements = input_placements
+        self.target_shape = target_shape
+        self.output_placements = output_placements
+        self.slice_funtor = slice_funtor
+
+
+class TestReshapeCoShard:
+    def setUp(self):
+        self.mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+        mesh_coord = lambda idx: (idx // 2, idx % 2)
+        self.test_cases = [
+            # test flatten
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0), dist.Shard(1)],
+                [192],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                lambda idx: (idx, slice(None), slice(None)),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(1), dist.Shard(2)],
+                [192],
+                [dist.Replicate(), dist.Replicate()],
+                lambda idx: slice(None),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                [192],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                lambda idx: (idx, slice(None), slice(None)),
+            ),
+            ReshapeTestCase(
+                [2, 12, 8],
+                [dist.Shard(0), dist.Shard(1)],
+                [192],
+                [dist.Shard(0), dist.Replicate()],
+                lambda idx: (mesh_coord(idx)[0], slice(None), slice(None)),
+            ),
+            # test split
+            ReshapeTestCase(
+                [192],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                [4, 6, 8],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                lambda idx: slice(idx * 48, (idx + 1) * 48),
+            ),
+            ReshapeTestCase(
+                [192],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                [6, 4, 8],
+                [dist.Replicate(), dist.Replicate()],
+                lambda idx: slice(None),
+            ),
+            # test combination
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0), dist.Shard(1)],
+                [2, 12, 8],
+                [dist.Shard(0), dist.Replicate()],
+                lambda idx: (
+                    slice(mesh_coord(idx)[0] * 2, (mesh_coord(idx)[0] + 1) * 2),
+                    slice(None),
+                    slice(None),
+                ),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                [2, 12, 8],
+                [dist.Replicate(), dist.Replicate()],
+                lambda idx: slice(None),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0), dist.Shard(1)],
+                [12, 2, 8],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                lambda idx: (idx, slice(None), slice(None)),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                [12, 2, 8],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                lambda idx: (idx, slice(None), slice(None)),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0), dist.Shard(1)],
+                [8, 6, 4],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                lambda idx: (idx, slice(None), slice(None)),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(1), dist.Shard(2)],
+                [8, 6, 4],
+                [dist.Replicate(), dist.Replicate()],
+                lambda idx: slice(None),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0), dist.Shard(2)],
+                [8, 6, 4],
+                [dist.Shard(0), dist.Replicate()],
+                lambda idx: (
+                    slice(mesh_coord(idx)[0] * 2, (mesh_coord(idx)[0] + 1) * 2),
+                    slice(None),
+                    slice(None),
+                ),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                [8, 6, 4],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                lambda idx: (idx, slice(None), slice(None)),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(2, shard_order=0), dist.Shard(2, shard_order=1)],
+                [24, 2, 4],
+                [dist.Replicate(), dist.Replicate()],
+                lambda idx: slice(None),
+            ),
+            ReshapeTestCase(
+                [4, 6, 8],
+                [dist.Shard(2, shard_order=0), dist.Shard(2, shard_order=1)],
+                [24, 4, 2],
+                [dist.Shard(1, shard_order=0), dist.Shard(1, shard_order=1)],
+                lambda idx: (
+                    slice(None),
+                    slice(None),
+                    slice(idx * 2, (idx + 1) * 2),
+                ),
+            ),
+        ]
+
+    def run_test_case(self, test_case: ReshapeTestCase):
+        paddle.seed(2025)
+        a = paddle.rand(test_case.input_shape, "float32")
+        a_numpy = a.numpy()
+        input_placements = test_case.input_placements
+        input = dist.shard_tensor(a, self.mesh, input_placements)
+        out = paddle.reshape(input, test_case.target_shape)
+        case_info = f"input_shape: {test_case.input_shape}, input_placements: {input_placements}, target_shape: {test_case.target_shape}"
+        # Verify output shape
+        np.testing.assert_equal(
+            out.shape,
+            test_case.target_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.target_shape}, Actual: {out.shape}",
+        )
+
+        # Verify placements
+        assert out.placements
+        for actual, expected in zip(
+            out.placements, test_case.output_placements
+        ):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.output_placements}, Actual: {out.placements}",
+            )
+        # Verify local_value if given
+        if test_case.slice_funtor:
+            idx = dist.get_rank()
+            np.testing.assert_allclose(
+                out._local_value().numpy().flatten(),
+                a_numpy[test_case.slice_funtor(idx)].flatten(),
+                err_msg=f"Local values mismatch when {case_info}.",
+            )
+
+    def run_all_tests(self):
+        self.setUp()
+        for test_case in self.test_cases:
+            self.run_test_case(test_case)
+
+
+if __name__ == '__main__':
+    TestReshapeCoShard().run_all_tests()

test/auto_parallel/end_to_end/softmax_co_shard.py

@@ -0,0 +1,312 @@
+# 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class SoftmaxTestCase:
+    def __init__(
+        self,
+        input_shape: list[int],
+        input_placements: list[dist.Placement],
+        axis: int,
+        output_shape: list[int],
+        output_placements: list[dist.Placement],
+        slice_funtor: Callable[[int], Any] | None = None,
+    ):
+        self.input_shape = input_shape
+        self.input_placements = input_placements
+        self.axis = axis
+        self.output_shape = output_shape
+        self.output_placements = output_placements
+        self.slice_funtor = slice_funtor
+
+
+class SoftmaxGradTestCase:
+    def __init__(
+        self,
+        input_shape: list[int],
+        axis: int,
+        output_shape: list[int],
+        output_placements: list[dist.Placement],
+        out_grad_placements: list[dist.Placement],
+        x_grad_placements: list[dist.Placement],
+    ):
+        self.input_shape = input_shape
+        self.axis = axis
+        self.output_shape = output_shape
+        self.output_placements = output_placements
+        self.out_grad_placements = out_grad_placements
+        self.x_grad_placements = x_grad_placements
+
+
+class TestSoftmaxCoShard:
+    def setUp(self):
+        self.mesh = dist.ProcessMesh(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=['x', 'y', 'z']
+        )
+        self.test_cases_forward = [
+            SoftmaxTestCase(
+                [32, 48, 128],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                0,
+                [32, 48, 128],
+                [dist.Replicate(), dist.Replicate(), dist.Shard(1)],
+            ),
+            SoftmaxTestCase(
+                [32, 48, 128],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                -3,
+                [32, 48, 128],
+                [dist.Replicate(), dist.Replicate(), dist.Shard(1)],
+            ),
+            SoftmaxTestCase(
+                [32, 48, 128],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                1,
+                [32, 48, 128],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+        ]
+        self.test_cases_backward = [
+            SoftmaxGradTestCase(
+                [32, 48, 128],
+                0,
+                [32, 48, 128],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [dist.Replicate(), dist.Replicate(), dist.Shard(1)],
+            ),
+            SoftmaxGradTestCase(
+                [32, 48, 128],
+                0,
+                [32, 48, 128],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [
+                    dist.Shard(0),
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                ],
+                [
+                    dist.Replicate(),
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                ],
+            ),
+            SoftmaxGradTestCase(
+                [32, 48, 128],
+                1,
+                [32, 48, 128],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(1),
+                ],
+                [
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                    dist.Shard(0),
+                ],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(0, shard_order=2),
+                ],
+            ),
+            SoftmaxGradTestCase(
+                [32, 48, 128],
+                1,
+                [32, 48, 128],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+                [dist.Replicate(), dist.Replicate(), dist.Shard(2)],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(2),
+                ],
+            ),
+            SoftmaxGradTestCase(
+                [32, 48, 128],
+                -1,
+                [32, 48, 128],
+                [
+                    dist.Shard(0),
+                    dist.Shard(1),
+                    dist.Replicate(),
+                ],
+                [
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                    dist.Replicate(),
+                ],
+                [
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+            SoftmaxGradTestCase(
+                [32, 48, 128],
+                -1,
+                [32, 48, 128],
+                [dist.Shard(0), dist.Shard(1), dist.Replicate()],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+            SoftmaxGradTestCase(
+                [32, 48, 128],
+                -1,
+                [32, 48, 128],
+                [
+                    dist.Shard(0),
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                ],
+                [
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                    dist.Replicate(),
+                ],
+                [
+                    dist.Shard(0),
+                    dist.Shard(1, shard_order=0),
+                    dist.Shard(1, shard_order=1),
+                ],
+            ),
+        ]
+
+    def run_test_case_forward(self, test_case: SoftmaxTestCase):
+        paddle.seed(2025)
+        a = paddle.rand(test_case.input_shape, "float32")
+        input_placements = test_case.input_placements
+        input = dist.shard_tensor(a, self.mesh, input_placements)
+        out = paddle.nn.functional.softmax(input, test_case.axis)
+        case_info = f"input_shape: {test_case.input_shape}, input_placements: {input_placements}, axis: {test_case.axis}"
+        # Verify output shape
+        np.testing.assert_equal(
+            out.shape,
+            test_case.output_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.output_shape}, Actual: {out.shape}",
+        )
+
+        # Verify placements
+        assert out.placements
+        for actual, expected in zip(
+            out.placements, test_case.output_placements
+        ):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.output_placements}, Actual: {out.placements}",
+            )
+        # Verify local_value if given
+        if test_case.slice_funtor:
+            idx = dist.get_rank()
+            np.testing.assert_equal(
+                out._local_value().numpy().flatten(),
+                a[test_case.slice_funtor(idx)].numpy().flatten(),
+                err_msg=f"Local values mismatch when {case_info}.",
+            )
+
+    def run_test_case_backward(self, test_case: SoftmaxGradTestCase):
+        a = paddle.rand(test_case.input_shape, "float32")
+        a.stop_gradient = False
+        input_placements = [dist.Replicate() for _ in range(self.mesh.ndim)]
+        input = dist.shard_tensor(a, self.mesh, input_placements)
+        out = paddle.nn.functional.softmax(input, test_case.axis)
+        out = dist.reshard(out, self.mesh, test_case.output_placements)
+
+        out_grad = paddle.ones(out.shape, "float32")
+        out_grad = dist.shard_tensor(
+            out_grad, self.mesh, test_case.out_grad_placements
+        )
+
+        (x_grad,) = paddle.grad([out], input, [out_grad])
+
+        case_info = f"input_shape: {test_case.input_shape}, axis: {test_case.axis}, out_placements: {test_case.output_placements}, out_grad_placements: {test_case.out_grad_placements}"
+        # Verify output shape
+        np.testing.assert_equal(
+            x_grad.shape,
+            test_case.input_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.input_shape}, Actual: {x_grad.shape}",
+        )
+
+        # Verify placements
+        assert x_grad.placements
+        for actual, expected in zip(
+            x_grad.placements, test_case.x_grad_placements
+        ):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.x_grad_placements}, Actual: {x_grad.placements}",
+            )
+
+    def run_all_tests(self):
+        self.setUp()
+        for test_case in self.test_cases_forward:
+            self.run_test_case_forward(test_case)
+
+
+if __name__ == '__main__':
+    TestSoftmaxCoShard().run_all_tests()

test/auto_parallel/end_to_end/test_e2e_co_shard.py

@@ -0,0 +1,38 @@
+# 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 TestReshardE2E(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=120)
+
+    def test_co_shard(self):
+        self.run_test_case("co_shard.py")
+
+    def test_reshape_co_shard(self):
+        self.run_test_case("reshape_co_shard.py")
+
+    def test_transpose_co_shard(self):
+        self.run_test_case("transpose_co_shard.py")
+
+    def test_elementwise_co_shard(self):
+        self.run_test_case("elementwise_co_shard.py")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/end_to_end/test_e2e_co_shard_8cards.py

@@ -0,0 +1,41 @@
+# 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 TestReshardE2E(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=120, nnode=1)
+
+    def test_tile_shard(self):
+        self.run_test_case("tile_co_shard.py")
+
+    def test_index_select_shard(self):
+        self.run_test_case("index_select_co_shard.py")
+
+    def test_softmax_shard(self):
+        self.run_test_case("softmax_co_shard.py")
+
+    def test_matmul_shard(self):
+        self.run_test_case("matmul_co_shard.py")
+
+    def test_argsort_shard(self):
+        self.run_test_case("argsort_co_shard.py")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/end_to_end/testslist.csv

@@ -0,0 +1,3 @@
+name,os,arch,timeout,run_type,launcher,num_port,run_serial,envs,conditions
+test_e2e_co_shard_8cards,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_e2e_co_shard,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,

test/auto_parallel/end_to_end/tile_co_shard.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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle._typing import TensorOrTensors
+
+
+class TileTestCase:
+    def __init__(
+        self,
+        x_shape: list[int],
+        x_placements: list[dist.Placement],
+        repeat_times: TensorOrTensors | Sequence[int],
+        out_shape: list[int],
+        out_placements: list[dist.Placement],
+    ):
+        self.x_shape = x_shape
+        self.x_placements = x_placements
+        self.repeat_times = repeat_times
+        self.out_shape = out_shape
+        self.out_placements = out_placements
+
+
+class TestTileCoShard:
+    def setUp(self):
+        self.mesh = dist.ProcessMesh(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=['x', 'y', 'z']
+        )
+        self.test_cases_forward = [
+            TileTestCase(
+                [8, 16, 24],
+                [
+                    dist.Shard(0),
+                    dist.Shard(2, shard_order=0),
+                    dist.Shard(2, shard_order=1),
+                ],
+                [2, 2, 1, 1],
+                [2, 16, 16, 24],
+                [
+                    dist.Replicate(),
+                    dist.Shard(3, shard_order=0),
+                    dist.Shard(3, shard_order=1),
+                ],
+            ),
+            TileTestCase(
+                [8, 16, 24],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(2),
+                ],
+                [1, 2],
+                [8, 16, 48],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Replicate(),
+                ],
+            ),
+            TileTestCase(
+                [8, 16, 24],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(2),
+                ],
+                [],
+                [8, 16, 24],
+                [
+                    dist.Shard(0, shard_order=0),
+                    dist.Shard(0, shard_order=1),
+                    dist.Shard(2),
+                ],
+            ),
+        ]
+
+    def run_test_case_forward(self, test_case: TileTestCase):
+        paddle.seed(2025)
+        x = paddle.rand(test_case.x_shape, "float32")
+        x_placements = test_case.x_placements
+        input = dist.shard_tensor(x, self.mesh, x_placements)
+        out = paddle.tile(input, test_case.repeat_times)
+        case_info = f"input_shape: {test_case.x_shape}, input_placements: {x_placements}, axis: {test_case.repeat_times}"
+        # Verify output shape
+        np.testing.assert_equal(
+            out.shape,
+            test_case.out_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.out_shape}, Actual: {out.shape}",
+        )
+
+        # Verify placements
+        assert out.placements
+        for actual, expected in zip(out.placements, test_case.out_placements):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.out_placements}, Actual: {out.placements}",
+            )
+
+    def run_all_tests(self):
+        self.setUp()
+        for test_case in self.test_cases_forward:
+            self.run_test_case_forward(test_case)
+
+
+if __name__ == '__main__':
+    TestTileCoShard().run_all_tests()

test/auto_parallel/end_to_end/transpose_co_shard.py

@@ -0,0 +1,120 @@
+# 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class TransposeTestCase:
+    def __init__(
+        self,
+        input_shape: list[int],
+        input_placements: list[dist.Placement],
+        perm: list[int],
+        output_shape: list[int],
+        output_placements: list[dist.Placement],
+        slice_funtor: Callable[[int], Any] | None = None,
+    ):
+        self.input_shape = input_shape
+        self.input_placements = input_placements
+        self.perm = perm
+        self.output_shape = output_shape
+        self.output_placements = output_placements
+        self.slice_funtor = slice_funtor
+
+
+class TestTransposeCoShard:
+    def setUp(self):
+        self.mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
+        self.test_cases = [
+            # test flatten
+            TransposeTestCase(
+                [64, 48, 36, 24],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                [1, 0, 2, 3],
+                [48, 64, 36, 24],
+                [dist.Shard(1, shard_order=0), dist.Shard(1, shard_order=1)],
+            ),
+            TransposeTestCase(
+                [64, 48, 36, 24],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+                [0, 1, 2, 3],
+                [64, 48, 36, 24],
+                [dist.Shard(0, shard_order=0), dist.Shard(0, shard_order=1)],
+            ),
+            TransposeTestCase(
+                [64, 48, 36, 24],
+                [dist.Shard(2, shard_order=0), dist.Shard(2, shard_order=1)],
+                [0, 2, 3, 1],
+                [64, 36, 24, 48],
+                [dist.Shard(1, shard_order=0), dist.Shard(1, shard_order=1)],
+            ),
+            TransposeTestCase(
+                [64, 48, 36, 24],
+                [dist.Shard(2, shard_order=0), dist.Shard(2, shard_order=1)],
+                [-1, 0, -2, 1],
+                [24, 64, 36, 48],
+                [dist.Shard(2, shard_order=0), dist.Shard(2, shard_order=1)],
+            ),
+        ]
+
+    def run_test_case(self, test_case: TransposeTestCase):
+        paddle.seed(2025)
+        a = paddle.rand(test_case.input_shape, "float32")
+        input_placements = test_case.input_placements
+        input = dist.shard_tensor(a, self.mesh, input_placements)
+        out = paddle.transpose(input, test_case.perm)
+        case_info = f"input_shape: {test_case.input_shape}, input_placements: {input_placements}, perm: {test_case.perm}"
+        # Verify output shape
+        np.testing.assert_equal(
+            out.shape,
+            test_case.output_shape,
+            err_msg=f"Output shape mismatch when {case_info}. Expected: {test_case.output_shape}, Actual: {out.shape}",
+        )
+
+        # Verify placements
+        assert out.placements
+        for actual, expected in zip(
+            out.placements, test_case.output_placements
+        ):
+            np.testing.assert_equal(
+                actual,
+                expected,
+                err_msg=f"Output placements mismatch when {case_info}. Expected: {test_case.output_placements}, Actual: {out.placements}",
+            )
+        # Verify local_value if given
+        if test_case.slice_funtor:
+            idx = dist.get_rank()
+            np.testing.assert_equal(
+                out._local_value().numpy().flatten(),
+                a[test_case.slice_funtor(idx)].numpy().flatten(),
+                err_msg=f"Local values mismatch when {case_info}.",
+            )
+
+    def run_all_tests(self):
+        self.setUp()
+        for test_case in self.test_cases:
+            self.run_test_case(test_case)
+
+
+if __name__ == '__main__':
+    TestTransposeCoShard().run_all_tests()

test/auto_parallel/high_order_grad.py

@@ -0,0 +1,149 @@
+# 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 numpy as np
+
+import paddle
+from paddle.distributed.fleet import auto
+from paddle.incubate.autograd import Hessian
+
+np.random.seed(1234)
+paddle.seed(1234)
+
+
+class FCNet:
+    def __init__(self, num_ins, num_outs, num_layers, hidden_size):
+        self.num_ins = num_ins
+        self.num_outs = num_outs
+        self.num_layers = num_layers
+        self.hidden_size = hidden_size
+        self.activation = paddle.tanh
+
+        self.weights = []
+        self.biases = []
+        for i in range(self.num_layers):
+            if i == 0:
+                lsize = self.num_ins
+                rsize = self.hidden_size
+            elif i == (self.num_layers - 1):
+                lsize = self.hidden_size
+                rsize = self.num_outs
+            else:
+                lsize = self.hidden_size
+                rsize = self.hidden_size
+
+            w = paddle.static.create_parameter(
+                shape=[lsize, rsize], dtype="float32", is_bias=False
+            )
+            b = paddle.static.create_parameter(
+                shape=[rsize], dtype="float32", is_bias=True
+            )
+            self.weights.append(w)
+            self.biases.append(b)
+
+    def nn_func(self, ins):
+        u = ins
+        for i in range(self.num_layers - 1):
+            u = paddle.nn.functional.linear(u, self.weights[i], self.biases[i])
+            u = self.activation(u)
+        u = paddle.nn.functional.linear(u, self.weights[-1], self.biases[-1])
+        return u
+
+
+class LaplaceModel(paddle.nn.Layer):
+    def __init__(self, num_ins=2, num_outs=1, num_layers=5, hidden_size=20):
+        super().__init__()
+        self.net = FCNet(
+            num_ins=num_ins,
+            num_outs=num_outs,
+            num_layers=num_layers,
+            hidden_size=hidden_size,
+        )
+
+    def forward(self, inputs, bc_index):
+        inputs.stop_gradient = False
+        outputs = self.net.nn_func(inputs)
+        # eq_loss
+        hes = Hessian(self.net.nn_func, inputs, is_batched=True)
+        eq_loss = paddle.norm(hes[:, 0, 0] + hes[:, 1, 1], p=2)
+        # bc_loss
+        bc_u = paddle.index_select(outputs, bc_index)
+        return eq_loss, bc_u
+
+
+class LaplaceDataset(paddle.io.Dataset):
+    def __init__(self, num_sample):
+        self.num_sample = num_sample
+
+    def __getitem__(self, index):
+        x = np.linspace(0, 0.9, 10)
+        y = np.linspace(0, 0.9, 10)
+        np.random.seed(index)  # Optional: Ensure reproducibility
+        bc_value = np.random.rand(36).reshape(36, 1).astype('float32')
+
+        domain_space = []
+        bc_index = []
+        for j in range(len(y)):
+            for i in range(len(x)):
+                domain_space.append([x[i], y[j]])
+                if i == 0 or i == 9 or j == 0 or j == 9:
+                    bc_index.append(i + 10 * j)
+        domain_space = np.array(domain_space, dtype='float32')
+        bc_index = np.array(bc_index, dtype='int64')
+        # Return a single input point and its related information based on the index
+        idx = index % len(domain_space)
+        return domain_space[idx], bc_index, bc_value
+
+    def __len__(self):
+        return self.num_sample
+
+
+def loss_func(eq_loss, bc_u, bc_value):
+    bc_diff = bc_u - bc_value
+    bc_loss = paddle.norm(bc_diff, p=2)
+    loss = eq_loss + bc_loss
+    return loss
+
+
+def main():
+    paddle.enable_static()
+    # dataset
+    train_dataset = LaplaceDataset(10)
+    # optimizer
+    optimizer = paddle.optimizer.Adam(learning_rate=0.001)
+    # model
+    laplace = LaplaceModel()
+
+    dist_strategy = auto.Strategy()
+    dist_strategy.auto_mode = "semi"
+
+    engine = auto.Engine(
+        laplace, loss=loss_func, optimizer=optimizer, strategy=dist_strategy
+    )
+    engine.fit(train_dataset, train_sample_split=2, batch_size=None)
+
+    dist_context = engine.dist_context
+    block = engine.main_program.global_block()
+    ops = block.ops
+    for op in ops:
+        if op.type == 'p_norm':
+            op_dist_attr = dist_context.get_op_dist_attr_for_program(op)
+            assert op_dist_attr.impl_type == 'p_norm'
+        if 'x' in op.input_arg_names:
+            out_name = op.output_arg_names[0]
+            assert block.vars[out_name].shape[0] == 50
+
+
+if __name__ == "__main__":
+    main()

test/auto_parallel/hybrid_strategy/CMakeLists.txt

@@ -0,0 +1,193 @@
+# This file is generated by ${PADDLE_ROOT}/tools/gen_ut_cmakelists.py.
+# Please don't modify this file manually.
+# If you need to change unittests in this file, please modify testslist.csv in the current directory
+# and then run the command `python3 ${PADDLE_ROOT}/tools/gen_ut_cmakelists.py -f ${CURRENT_DIRECTORY}/testslist.csv`
+set(LOCAL_ALL_ARCH ON)
+set(LOCAL_ALL_PLAT ON)
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_hybrid_strategy MODULES
+    test_semi_auto_parallel_hybrid_strategy ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_semi_auto_parallel_hybrid_strategy
+                       PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_save_load_state_dict MODULES test_save_load_state_dict ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_save_load_state_dict
+                       PROPERTIES TIMEOUT "400" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_flexcheckpoint_merge MODULES test_flexcheckpoint_merge ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_flexcheckpoint_merge
+                       PROPERTIES TIMEOUT "120" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_c_cross_entropy MODULES
+    test_semi_auto_parallel_c_cross_entropy ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_semi_auto_parallel_c_cross_entropy
+                       PROPERTIES TIMEOUT "120" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_cross_mesh_reshard MODULES test_cross_mesh_reshard ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_cross_mesh_reshard PROPERTIES TIMEOUT "120" LABELS
+                                                          "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_llama_model_amp MODULES
+    test_semi_auto_parallel_llama_model_amp ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_semi_auto_parallel_llama_model_amp
+                       PROPERTIES TIMEOUT "180" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_hybrid_sharding_strategy MODULES
+    test_semi_auto_parallel_hybrid_sharding_strategy ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_semi_auto_parallel_hybrid_sharding_strategy
+                       PROPERTIES TIMEOUT "120" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_global_mesh_reshard MODULES test_global_mesh_reshard ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_global_mesh_reshard PROPERTIES TIMEOUT "120" LABELS
+                                                           "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_global_input MODULES
+    test_semi_auto_parallel_global_input ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_semi_auto_parallel_global_input
+                       PROPERTIES TIMEOUT "120" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_multi_inputs MODULES
+    test_semi_auto_parallel_multi_inputs ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_semi_auto_parallel_multi_inputs
+                       PROPERTIES TIMEOUT "120" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_llama_model_vpp MODULES
+    test_semi_auto_parallel_llama_model_vpp ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_semi_auto_parallel_llama_model_vpp
+                       PROPERTIES TIMEOUT "180" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_parallel_llama_model_pir
+    MODULES
+    test_semi_auto_parallel_llama_model_pir
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python;FLAGS_enable_pir_api=1"
+  )
+  set_tests_properties(test_semi_auto_parallel_llama_model_pir
+                       PROPERTIES TIMEOUT "180" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_pir_reshard_nd_mesh_func MODULES test_pir_reshard_nd_mesh_func ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_pir_reshard_nd_mesh_func
+                       PROPERTIES TIMEOUT "60" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_llama_acc_align
+    MODULES
+    test_semi_auto_llama_acc_align
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python;FLAGS_enable_pir_api=1"
+  )
+  set_tests_properties(test_semi_auto_llama_acc_align
+                       PROPERTIES TIMEOUT "300" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_llama_save_load
+    MODULES
+    test_semi_auto_llama_save_load
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python;FLAGS_enable_pir_api=1"
+  )
+  set_tests_properties(test_semi_auto_llama_save_load
+                       PROPERTIES TIMEOUT "180" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_parallel_api_with_llama_1d MODULES test_parallel_api_with_llama_1d
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_parallel_api_with_llama_1d
+                       PROPERTIES TIMEOUT "400" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_parallel_api_with_llama_2d MODULES test_parallel_api_with_llama_2d
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_parallel_api_with_llama_2d
+                       PROPERTIES TIMEOUT "400" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_parallel_api_with_llama_2d_sep MODULES
+    test_parallel_api_with_llama_2d_sep ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_parallel_api_with_llama_2d_sep
+                       PROPERTIES TIMEOUT "400" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_parallel_api_with_llama_3d MODULES test_parallel_api_with_llama_3d
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_parallel_api_with_llama_3d
+                       PROPERTIES TIMEOUT "800" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_to_distributed_api_for_llama MODULES test_to_distributed_api_for_llama
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_to_distributed_api_for_llama
+                       PROPERTIES TIMEOUT "180" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_parallel_api_with_llama_lora MODULES test_parallel_api_with_llama_lora
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_parallel_api_with_llama_lora
+                       PROPERTIES TIMEOUT "360" LABELS "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_process_mesh MODULES test_process_mesh ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_process_mesh PROPERTIES TIMEOUT "150" LABELS
+                                                    "RUN_TYPE=HYBRID")
+endif()
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_get_group_in_different_hybrid_configs MODULES
+    test_get_group_in_different_hybrid_configs ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python")
+  set_tests_properties(test_get_group_in_different_hybrid_configs
+                       PROPERTIES TIMEOUT "150" LABELS "RUN_TYPE=HYBRID")
+endif()

test/auto_parallel/hybrid_strategy/parallel_api.py

@@ -0,0 +1,618 @@
+# 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 logging
+import os
+import random
+from dataclasses import dataclass
+from functools import reduce
+
+import numpy as np
+from single_llama_model import LlamaForCausalLM, LlamaPretrainingCriterion
+from single_lora_model import LoRAModel
+
+import paddle
+import paddle.distributed as dist
+from paddle import LazyGuard
+from paddle.distributed.auto_parallel.intermediate.parallelize import (
+    parallelize_model,
+    parallelize_optimizer,
+)
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+
+def is_pp_enable():
+    global_mesh = dist.auto_parallel.get_mesh()
+    return "pp" in global_mesh.dim_names
+
+
+def get_mesh(pp_idx=None):
+    global_mesh = dist.auto_parallel.get_mesh()
+    assert global_mesh is not None, "global_mesh is not initialized!"
+    if pp_idx is None:
+        return global_mesh
+    if is_pp_enable():
+        mesh = global_mesh.get_mesh_with_dim("pp")[pp_idx]
+        return mesh
+    else:
+        return global_mesh
+
+
+class Config:
+    vocab_size = 8192
+    hidden_size = 512
+    intermediate_size = 2048
+    seq_length = 512
+    num_hidden_layers = 2
+    num_attention_heads = 8
+    rms_norm_eps = 1e-6
+    use_lazy_init = False
+    context_parallel = False
+    sep_parallel = False
+
+
+@dataclass
+class LoRaConfig:
+    r = 8
+    lora_alpha = 8
+    lora_dropout = 0.0
+    rslora = False
+    lora_plus_scale = 1.0
+    pissa = False
+    use_quick_lora = False
+    lora_use_mixer = False
+    use_mora = False
+    trainable_bias = False
+    trainable_modules = None
+    target_modules = [
+        ".*q_proj.*",
+        ".*v_proj.*",
+        ".*k_proj.*",
+        ".*o_proj.*",
+        ".*qkv_proj.*",
+        ".*gate_proj.*",
+        ".*down_proj.*",
+        ".*up_proj.*",
+        ".*gate_up_fused_proj.*",
+    ]
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, num_samples=100):
+        super().__init__()
+        self.seq_len = seq_len
+        self.num_samples = num_samples
+
+    def __getitem__(self, index):
+        input = np.random.uniform(size=[self.seq_len]).astype("int64")
+        label = (np.random.uniform(size=[self.seq_len]) * 10).astype("int64")
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+def create_optimizer(model, lr_scheduler):
+    decay_parameters = [
+        p.name
+        for n, p in model.named_parameters()
+        if not any(nd in n for nd in ["bias", "norm"])
+    ]
+
+    def apply_decay_param_fun(x):
+        return x in decay_parameters
+
+    # test global_clip in auto_parallel
+    if os.getenv("use_param_group") == "true":
+        param_group = {}
+        param_group["params"] = list(model.parameters())
+        param_group["weight_decay"] = 0.01
+        param_group["grad_clip"] = paddle.nn.ClipGradByGlobalNorm(1.0)
+        optimizer = paddle.optimizer.adamw.AdamW(
+            learning_rate=lr_scheduler,
+            apply_decay_param_fun=apply_decay_param_fun,
+            parameters=[param_group],
+        )
+    else:
+        optimizer = paddle.optimizer.adamw.AdamW(
+            learning_rate=lr_scheduler,
+            apply_decay_param_fun=apply_decay_param_fun,
+            parameters=model.parameters(),
+            weight_decay=0.01,
+            grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+        )
+    return optimizer
+
+
+class TestParallelAPI:
+    def __init__(self):
+        self.config = Config()
+        self.lora_config = LoRaConfig()
+        self.dp = int(os.getenv("dp"))
+        self.mp = int(os.getenv("mp"))
+        self.pp = int(os.getenv("pp"))
+        self.sep = int(os.getenv("sep", "1"))
+        if os.getenv("use_lazy_init") == "true":
+            self.config.use_lazy_init = True
+        self.gradient_accumulation_steps = int(os.getenv("acc_step"))
+
+        self.amp = False
+        self.amp_dtype = "float16"
+        self.amp_level = "O1"
+        self.amp_master_grad = False
+        if os.getenv("amp") == "true":
+            self.amp = True
+        if os.getenv("amp_dtype") in ["float16", "bfloat16"]:
+            self.amp_dtype = os.getenv("amp_dtype")
+        if os.getenv("amp_level") in ["O0", "O1", "O2"]:
+            self.amp_level = os.getenv("amp_level")
+        if os.getenv("amp_master_grad") == "true":
+            self.amp_master_grad = True
+        self.level = os.getenv("sharding_stage", "0")
+        self.sequence_parallel = False
+        if os.getenv("sequence_parallel") == "true":
+            self.sequence_parallel = True
+        self.config.context_parallel = False
+        if os.getenv("context_parallel", "false") == "true":
+            self.config.context_parallel = True
+        self.config.sep_parallel = False
+        if os.getenv("sep_parallel", "false") == "true":
+            self.config.sep_parallel = True
+        self.prepare_input_output = False
+        if os.getenv("prepare_input_output") == "true":
+            self.sequence_parallel = True
+        if self.sep > 1:
+            assert (
+                self.config.context_parallel is True
+                and self.config.sep_parallel is False
+            ) or (
+                self.config.context_parallel is False
+                and self.config.sep_parallel is True
+            ), (
+                "when sep > 1, either context_parallel or sep_parallel should be true"
+            )
+        num_hidden_layers = os.getenv("num_hidden_layers")
+        if num_hidden_layers:
+            self.config.num_hidden_layers = int(num_hidden_layers)
+
+        self.one_api = False
+        if os.getenv("one_api") == "true":
+            self.one_api = True
+
+        seed = int(os.getenv("seed", 2024))
+        self.share_embedding = int(os.getenv("test_share_embedding", "0"))
+        self.position_embedding = int(os.getenv("test_position_embedding", "0"))
+        self.test_lora = int(os.getenv("test_lora", "0"))
+        np.random.seed(seed)
+        random.seed(seed)
+        paddle.seed(seed)
+        self.init_dist_env()
+
+    def init_dist_env(self):
+        mesh_dims = [
+            ("dp", self.dp),
+            ("pp", self.pp),
+            ("mp", self.mp),
+            ("sep", self.sep),
+        ]
+        if self.pp * self.mp == 1:
+            mesh_dims = [("dp", self.dp)]
+        dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
+        mesh_shape = [mesh_dim[1] for mesh_dim in mesh_dims]
+        mesh_arr = np.arange(
+            0, reduce(lambda x, y: x * y, mesh_shape, 1)
+        ).reshape(mesh_shape)
+        global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
+        dist.auto_parallel.set_mesh(global_mesh)
+
+    def check_mp(self, layer):
+        if self.mp == 1:
+            return
+        for name, sub_layer in layer.named_sublayers():
+            if len(sub_layer.sublayers()) == 0:
+                if 'q_proj' in name or 'k_proj' in name or 'v_proj' in name:
+                    assert sub_layer.weight.placements == [
+                        dist.Replicate(),
+                        dist.Shard(1),
+                        dist.Replicate(),  # cp
+                    ]
+                    assert sub_layer.bias.placements == [
+                        dist.Replicate(),
+                        dist.Shard(0),
+                        dist.Replicate(),  # cp
+                    ]
+                    if self.test_lora:
+                        assert sub_layer.lora_B.placements == [
+                            dist.Replicate(),
+                            dist.Shard(1),
+                            dist.Replicate(),  # cp
+                        ]
+                if 'gate_proj' in name or 'up_proj' in name:
+                    assert sub_layer.weight.placements == [
+                        dist.Replicate(),
+                        dist.Shard(1),
+                        dist.Replicate(),  # cp
+                    ]
+                    if self.test_lora:
+                        assert sub_layer.lora_B.placements == [
+                            dist.Replicate(),
+                            dist.Shard(1),
+                            dist.Replicate(),  # cp
+                        ]
+                if (
+                    'embed_tokens' in name or 'lm_head' in name
+                ) and not self.share_embedding:
+                    assert sub_layer.weight.placements == [
+                        dist.Replicate(),
+                        dist.Shard(1),
+                        dist.Replicate(),  # cp
+                    ]
+                if 'o_proj' in name:
+                    assert sub_layer.weight.placements == [
+                        dist.Replicate(),
+                        dist.Shard(0),
+                        dist.Replicate(),  # cp
+                    ], f'{name} , {sub_layer.weight.name} , {sub_layer.weight}'
+                    if self.test_lora:
+                        assert sub_layer.lora_A.placements == [
+                            dist.Replicate(),
+                            dist.Shard(0),
+                            dist.Replicate(),  # cp
+                        ]
+                    # assert sub_layer.bias.placements is None
+                if 'down_proj' in name:
+                    assert sub_layer.weight.placements == [
+                        dist.Replicate(),
+                        dist.Shard(0),
+                        dist.Replicate(),  # cp
+                    ]
+                    if self.test_lora:
+                        assert sub_layer.lora_A.placements == [
+                            dist.Replicate(),
+                            dist.Shard(0),
+                            dist.Replicate(),  # cp
+                        ]
+
+    def check_lora(self, layer):
+        if not self.test_lora:
+            return
+        for name, sub_layer in layer.named_sublayers():
+            if len(sub_layer.sublayers()) == 0:
+                if 'q_proj' in name or 'k_proj' in name or 'v_proj' in name:
+                    assert sub_layer.weight.stop_gradient
+                    assert not sub_layer.lora_A.stop_gradient
+                    assert not sub_layer.lora_B.stop_gradient
+                if 'gate_proj' in name or 'up_proj' in name:
+                    assert sub_layer.weight.stop_gradient
+                    assert not sub_layer.lora_A.stop_gradient
+                    assert not sub_layer.lora_B.stop_gradient
+                if (
+                    'embed_tokens' in name or 'lm_head' in name
+                ) and not self.share_embedding:
+                    assert sub_layer.weight.stop_gradient
+                if 'o_proj' in name:
+                    assert sub_layer.weight.stop_gradient, (
+                        f'{name} , {sub_layer.weight.name} , {sub_layer.weight}'
+                    )
+                    assert not sub_layer.lora_A.stop_gradient
+                    assert not sub_layer.lora_B.stop_gradient
+                    # assert sub_layer.bias.stop_gradient is None
+                if 'down_proj' in name:
+                    assert sub_layer.weight.stop_gradient
+                    assert not sub_layer.lora_A.stop_gradient
+                    assert not sub_layer.lora_B.stop_gradient
+
+    def parallel_model(self, layer):
+        dp_config = None
+        mp_config = None
+        pp_config = None
+        cp_config = None
+        prefix = "model." if self.test_lora else ""
+        if self.pp > 1:
+            # decoders_per_rank = self.config.num_hidden_layers // self.pp
+            # split_spec = {
+            #     ff"{prefix}llama.layers.{i * decoders_per_rank - 1}": SplitPoint.END
+            #     for i in range(1, self.pp)
+            # }
+            pp_config = {
+                'split_spec': f"{prefix}llama.layers",
+                "global_spec": f"{prefix}llama.global_layer",
+            }
+        if self.dp > 1:
+            dp_config = {'sharding_level': self.level}
+        if self.mp > 1:
+            if not self.sequence_parallel:
+                plan = {
+                    f"{prefix}llama.embed_tokens": dist.ColWiseParallel(
+                        gather_output=True
+                    ),
+                    f"{prefix}llama.position_embedding": dist.ColWiseParallel(),
+                    f"{prefix}llama.layers.*.self_attn.q_proj": dist.ColWiseParallel(
+                        gather_output=True
+                    ),
+                    f"{prefix}llama.layers.*.self_attn.q_proj.lora_B": dist.ColWiseParallel(),
+                    f"{prefix}llama.layers.*.self_attn.k_proj": dist.ColWiseParallel(
+                        gather_output=True
+                    ),
+                    f"{prefix}llama.layers.*.self_attn.k_proj.lora_B": dist.ColWiseParallel(),
+                    f"{prefix}llama.layers.*.self_attn.v_proj": dist.ColWiseParallel(
+                        gather_output=True
+                    ),
+                    f"{prefix}llama.layers.*.self_attn.v_proj.lora_B": dist.ColWiseParallel(),
+                    f"{prefix}llama.layers.*.self_attn.o_proj": dist.RowWiseParallel(
+                        is_input_parallel=False
+                    ),
+                    f"{prefix}llama.layers.*.self_attn.o_proj.lora_A": dist.RowWiseParallel(),
+                    f"{prefix}llama.layers.*.mlp.gate_proj": dist.ColWiseParallel(),
+                    f"{prefix}llama.layers.*.mlp.gate_proj.lora_B": dist.ColWiseParallel(),
+                    f"{prefix}llama.layers.*.mlp.up_proj": dist.ColWiseParallel(),
+                    f"{prefix}llama.layers.*.mlp.up_proj.lora_B": dist.ColWiseParallel(),
+                    f"{prefix}llama.layers.*.mlp.down_proj": dist.RowWiseParallel(),
+                    f"{prefix}llama.layers.*.mlp.down_proj.lora_A": dist.RowWiseParallel(),
+                    f"{prefix}lm_head.weight": dist.ColWiseParallel(),
+                }
+            else:
+                if self.prepare_input_output:
+                    plan = {
+                        f"{prefix}llama.embed_tokens": dist.ColWiseParallel(),
+                        f"{prefix}llama.position_embedding": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.self_attn.q_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.self_attn.k_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.self_attn.v_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.self_attn.o_proj": dist.RowWiseParallel(),
+                        f"{prefix}llama.layers.*.mlp.gate_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.mlp.up_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.mlp.down_proj": dist.RowWiseParallel(),
+                        f"{prefix}lm_head.weight": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.input_layernorm": dist.SequenceParallelEnable(),
+                        f"{prefix}llama.layers.*.post_attention_layernorm": dist.SequenceParallelEnable(),
+                        f"{prefix}llama.norm": dist.SequenceParallelEnable(),
+                    }
+                else:
+                    plan = {
+                        f"{prefix}llama.embed_tokens": [
+                            dist.ColWiseParallel(),
+                            dist.SequenceParallelBegin(),
+                        ],
+                        f"{prefix}llama.position_embedding": [
+                            dist.ColWiseParallel(),
+                            dist.SequenceParallelBegin(),
+                        ],
+                        f"{prefix}llama.layers.*.self_attn.q_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.self_attn.k_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.self_attn.v_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.self_attn.o_proj": dist.RowWiseParallel(),
+                        f"{prefix}llama.layers.*.self_attn": dist.SequenceParallelDisable(),
+                        f"{prefix}llama.layers.*.mlp.gate_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.mlp.up_proj": dist.ColWiseParallel(),
+                        f"{prefix}llama.layers.*.mlp.down_proj": dist.RowWiseParallel(),
+                        f"{prefix}llama.layers.*.mlp": dist.SequenceParallelDisable(
+                            need_transpose=False
+                        ),
+                        f"{prefix}lm_head.weight": dist.ColWiseParallel(),
+                        f"{prefix}lm_head": dist.SequenceParallelEnd(),
+                    }
+            mp_config = {'parallelize_plan': plan}
+        if self.sep > 1:
+            if not (
+                self.config.context_parallel is True
+                and (
+                    os.getenv("backend") != "gpu"
+                    or not self.amp
+                    or int(paddle.version.cuda().split(".")[0]) < 11
+                    or paddle.device.cuda.get_device_capability()[0] < 8
+                )
+            ):
+                bck = 'p2p'
+                if self.config.context_parallel is True:
+                    bck = 'p2p'
+                elif self.config.sep_parallel is True:
+                    bck = 'all2all'
+                else:
+                    logging.error(
+                        f"when sep > 1, should set context_parallel or sep_parallel, but got sep_parallel={self.config.sep_parallel}, context_parallel={self.context_parallel}"
+                    )
+                plan = {
+                    f"{prefix}llama": dist.PrepareContextParallel(backend=bck),
+                    f"{prefix}llama.layers.*.self_attn.sdpa": dist.ContextParallel(
+                        backend=bck
+                    ),
+                }
+                cp_config = {'parallelize_plan': plan}
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+
+        config = {
+            'dp_config': dp_config,
+            'mp_config': mp_config,
+            'pp_config': pp_config,
+            'cp_config': cp_config,
+        }
+
+        if self.one_api:
+            optimizer = create_optimizer(layer, lr_scheduler)
+            model, optimizer = dist.parallelize(
+                layer,
+                optimizer,
+                config=config,
+            )
+        else:
+            layer = parallelize_model(
+                layer,
+                config=config,
+            )
+            optimizer = create_optimizer(layer, lr_scheduler)
+            optimizer = parallelize_optimizer(
+                optimizer,
+                config=config,
+            )
+        self.check_mp(layer)
+        self.check_lora(layer)
+        return layer, optimizer, lr_scheduler
+
+    def run_llama(self, to_static=0):
+        if self.config.use_lazy_init:
+            with LazyGuard():
+                model = LlamaForCausalLM(
+                    self.config, self.share_embedding, self.position_embedding
+                )
+        else:
+            model = LlamaForCausalLM(
+                self.config, self.share_embedding, self.position_embedding
+            )
+        if self.test_lora:
+            if self.config.use_lazy_init:
+                with LazyGuard():
+                    model = LoRAModel(model, self.lora_config)
+            else:
+                model = LoRAModel(model, self.lora_config)
+        model, optimizer, lr_scheduler = self.parallel_model(model)
+
+        criterion = LlamaPretrainingCriterion(self.config)
+
+        if self.config.use_lazy_init:
+            for param in model.parameters():
+                assert not param._is_initialized()
+                param.initialize()
+
+        if self.amp and not to_static:
+            model, optimizer = paddle.amp.decorate(
+                models=model,
+                optimizers=optimizer,
+                level=self.amp_level,
+                dtype=self.amp_dtype,
+                master_grad=self.amp_master_grad,
+            )
+
+        train_dataset = RandomDataset(self.config.seq_length)
+        train_sampler = BatchSampler(
+            train_dataset,
+            batch_size=2,
+            shuffle=True,
+            drop_last=True,
+        )
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_sampler=train_sampler,
+            num_workers=0,
+        )
+
+        if self.pp == 1:
+            meshes = [get_mesh(0)]
+        elif self.pp > 1:
+            meshes = [get_mesh(0), get_mesh(-1)]
+        else:
+            raise ValueError("pp should be greater or equal to 1")
+
+        dist_loader = dist.shard_dataloader(
+            dataloader=train_dataloader,
+            meshes=meshes,
+            shard_dims="dp",
+        )
+
+        global_step = 1
+        tr_loss = float(0)
+
+        if not to_static:
+            model.train()
+            scaler = None
+            if self.amp and self.amp_dtype == "float16":
+                scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
+                scaler = dist.shard_scaler(scaler)
+
+            for step, inputs in enumerate(dist_loader()):
+                input_ids, labels = inputs
+                custom_black_list = [
+                    "reduce_sum",
+                    "c_softmax_with_cross_entropy",
+                ]
+                custom_white_list = []
+                if self.amp_level == "O2":
+                    custom_white_list.extend(
+                        ["lookup_table", "lookup_table_v2"]
+                    )
+                with paddle.amp.auto_cast(
+                    self.amp,
+                    custom_black_list=set(custom_black_list),
+                    custom_white_list=set(custom_white_list),
+                    level=self.amp_level,
+                    dtype=self.amp_dtype,
+                ):
+                    logits = model(input_ids)
+                    tr_loss_step = criterion(logits, labels)
+
+                if self.gradient_accumulation_steps > 1:
+                    tr_loss_step /= self.gradient_accumulation_steps
+                if scaler is not None:
+                    scaler.scale(tr_loss_step).backward()
+                else:
+                    tr_loss_step.backward()
+                tr_loss += tr_loss_step
+
+                if global_step % self.gradient_accumulation_steps == 0:
+                    logging.info(
+                        f"step: {global_step // self.gradient_accumulation_steps}  loss: {tr_loss.numpy()}"
+                    )
+                    if scaler is not None:
+                        scaler.step(optimizer)
+                        scaler.update()
+                    else:
+                        optimizer.step()
+                    optimizer.clear_grad()
+                    lr_scheduler.step()
+                    tr_loss = 0
+
+                global_step += 1
+                if global_step // self.gradient_accumulation_steps >= 3:
+                    break
+        else:
+            strategy = dist.Strategy()
+            if self.gradient_accumulation_steps > 1:
+                strategy.pipeline.accumulate_steps = (
+                    self.gradient_accumulation_steps
+                )
+
+            if self.amp:
+                amp = strategy.amp
+                amp.enable = self.amp
+                amp.dtype = self.amp_dtype
+                amp.level = self.amp_level.lower()
+                if self.amp_master_grad:
+                    amp.use_master_grad = True
+
+            dist_model = dist.to_static(
+                model,
+                dist_loader,
+                criterion,
+                optimizer,
+                strategy=strategy,
+            )
+
+            dist_model.train()
+            for step, inputs in enumerate(dist_loader()):
+                input_ids, labels = inputs
+                loss = dist_model(input_ids, labels)
+                logging.info(f"step: {step}  loss: {loss}")
+                if step >= 3:
+                    break
+
+    def run_test_cases(self):
+        self.run_llama(0)
+        if self.sep == 1:
+            # sep now only support dynamic mode
+            self.run_llama(1)
+
+
+if __name__ == '__main__':
+    TestParallelAPI().run_test_cases()

test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh.py

@@ -0,0 +1,375 @@
+# 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 paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.static.pir_pass import (
+    ReshardPasses,
+)
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+
+class TestReshardNdMesh:
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._seed = eval(os.getenv("seed"))
+        self.BATCH_SIZE = 2
+        self.SEQ_LEN = 4
+        self.HIDDEN_SIZE = 8
+        self._backend = os.getenv("backend")
+        self._mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+
+    def validate(
+        self, op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value
+    ):
+        # tgt_* are tuples, format: (process_ids, dims_mapping, partial_status)
+        operand_dist_attr = op.dist_attr.operand(0).as_tensor_dist_attr()
+        result_dist_attr = op.dist_attr.result(0).as_tensor_dist_attr()
+
+        assert operand_dist_attr.process_mesh.process_ids == tgt_operand[0]
+        assert operand_dist_attr.dims_mapping == tgt_operand[1]
+        assert operand_dist_attr.partial_status == tgt_operand[2]
+
+        assert result_dist_attr.process_mesh.process_ids == tgt_result[0]
+        assert result_dist_attr.dims_mapping == tgt_result[1]
+        assert result_dist_attr.partial_status == tgt_result[2]
+
+        in_value = op.operand_source(0)
+        out_value = op.result(0)
+        assert in_value.is_dist_dense_tensor_type()
+        assert out_value.is_dist_dense_tensor_type()
+        assert in_value.dist_attr().process_mesh.process_ids == tgt_in_value[0]
+        assert in_value.dist_attr().dims_mapping == tgt_in_value[1]
+        assert in_value.dist_attr().partial_status == tgt_in_value[2]
+
+        assert (
+            out_value.dist_attr().process_mesh.process_ids == tgt_out_value[0]
+        )
+        assert out_value.dist_attr().dims_mapping == tgt_out_value[1]
+        assert out_value.dist_attr().partial_status == tgt_out_value[2]
+
+    def create_program(self, input_shape, input_placements, output_placements):
+        paddle.enable_static()
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+            place = paddle.CPUPlace()
+        elif self._backend == "gpu":
+            place = paddle.CUDAPlace(dist.get_rank())
+
+        with paddle.pir_utils.IrGuard():
+            main_program = paddle.base.Program()
+            with paddle.base.program_guard(main_program):
+                input = paddle.ones(name='input', shape=input_shape)
+                dist_input = dist.shard_tensor(
+                    input, self._mesh, input_placements
+                )
+                dist_out = paddle._C_ops.reshard(
+                    dist_input, self._mesh, output_placements
+                )
+            dist_program = main_program.clone()
+            set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)
+            ReshardPasses.apply_reshard_pass(dist_program)
+
+        return main_program, dist_program
+
+    def run_pp_to_rr_case(self):
+        # [Partial(), Partial()] --> [Replicate(), Replicate()]
+        # ops: all_reduce sum + all_reduce sum
+        main_program, dist_program = self.create_program(
+            [self.BATCH_SIZE, self.SEQ_LEN, self.HIDDEN_SIZE],
+            [
+                dist.Partial(dist.ReduceType.kRedSum),
+                dist.Partial(dist.ReduceType.kRedSum),
+            ],
+            [dist.Replicate(), dist.Replicate()],
+        )
+
+        new_ops = dist_program.global_block().ops
+        new_ops_name = [op.name() for op in dist_program.global_block().ops]
+
+        rank_id = dist.get_rank()
+        assert dist_program.global_block().ops[-2].name() == "pd_op.all_reduce"
+        assert (
+            dist_program.global_block().ops[-2].int_attr("reduce_type")
+            == dist.ReduceOp.SUM
+        )
+        assert dist_program.global_block().ops[-1].name() == "pd_op.all_reduce"
+        assert (
+            dist_program.global_block().ops[-1].int_attr("reduce_type")
+            == dist.ReduceOp.SUM
+        )
+
+        # check the first allreduce_sum
+        op = new_ops[-2]
+        if rank_id == 0 or rank_id == 2:
+            process_ids = [0, 2]
+        elif rank_id == 1 or rank_id == 3:
+            process_ids = [1, 3]
+        tgt_operand = (process_ids, [-1, -1, -1], {0: dist.ReduceType.kRedSum})
+        tgt_result = (process_ids, [-1, -1, -1], {})
+        tgt_in_value = (
+            self._mesh.process_ids,
+            [-1, -1, -1],
+            {0: dist.ReduceType.kRedSum, 1: dist.ReduceType.kRedSum},
+        )
+        tgt_out_value = (
+            self._mesh.process_ids,
+            [-1, -1, -1],
+            {1: dist.ReduceType.kRedSum},
+        )
+        self.validate(op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value)
+
+        # check the second allreduce_sum
+        op = new_ops[-1]
+        if rank_id == 0 or rank_id == 1:
+            process_ids = [0, 1]
+        elif rank_id == 2 or rank_id == 3:
+            process_ids = [2, 3]
+        tgt_operand = (process_ids, [-1, -1, -1], {0: dist.ReduceType.kRedSum})
+        tgt_result = (process_ids, [-1, -1, -1], {})
+        tgt_in_value = (
+            self._mesh.process_ids,
+            [-1, -1, -1],
+            {1: dist.ReduceType.kRedSum},
+        )
+        tgt_out_value = (self._mesh.process_ids, [-1, -1, -1], {})
+        self.validate(op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value)
+
+    def run_pr_to_rs_case(self):
+        # [Partial(), Replicate()] --> [Replicate(), Shard(1)]
+        # all_reduce sum + slice
+        main_program, dist_program = self.create_program(
+            [self.BATCH_SIZE, self.SEQ_LEN, self.HIDDEN_SIZE],
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Replicate()],
+            [dist.Replicate(), dist.Shard(1)],
+        )
+
+        new_ops = dist_program.global_block().ops
+        new_ops_name = [op.name() for op in dist_program.global_block().ops]
+
+        check_all_reduce_sum = any(
+            op.name() == "pd_op.all_reduce"
+            and op.int_attr("reduce_type") == dist.ReduceOp.SUM
+            for op in dist_program.global_block().ops
+        )
+        assert check_all_reduce_sum
+        rank_id = dist.get_rank()
+        assert new_ops_name[-1] == "pd_op.slice"
+
+        # check the allreduce_sum
+        op = new_ops[new_ops_name.index("pd_op.all_reduce")]
+        if rank_id == 0 or rank_id == 2:
+            process_ids = [0, 2]
+        elif rank_id == 1 or rank_id == 3:
+            process_ids = [1, 3]
+        tgt_operand = (process_ids, [-1, -1, -1], {0: dist.ReduceType.kRedSum})
+        tgt_result = (process_ids, [-1, -1, -1], {})
+        tgt_in_value = (
+            self._mesh.process_ids,
+            [-1, -1, -1],
+            {0: dist.ReduceType.kRedSum},
+        )
+        tgt_out_value = (self._mesh.process_ids, [-1, -1, -1], {})
+        self.validate(op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value)
+
+        # check the second slice
+        op = new_ops[-1]
+        if rank_id == 0 or rank_id == 1:
+            process_ids = [0, 1]
+        elif rank_id == 2 or rank_id == 3:
+            process_ids = [2, 3]
+        tgt_operand = (process_ids, [-1, -1, -1], {})
+        tgt_result = (process_ids, [-1, 0, -1], {})
+        tgt_in_value = (self._mesh.process_ids, [-1, -1, -1], {})
+        tgt_out_value = (self._mesh.process_ids, [-1, 1, -1], {})
+
+    def run_pr_to_ss_case(self):
+        self.create_program(
+            [self.BATCH_SIZE, self.SEQ_LEN, self.HIDDEN_SIZE],
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Replicate()],
+            [dist.Shard(0), dist.Shard(1)],
+        )
+
+    def run_ss_to_ss_case(self):
+        # [Shard(0), Shard(1)] --> [Shard(1), Shard(0)]
+        # all_gather+all_gather+slice+slice
+        main_program, dist_program = self.create_program(
+            [self.BATCH_SIZE, self.SEQ_LEN, self.HIDDEN_SIZE],
+            [dist.Shard(0), dist.Shard(1)],
+            [dist.Shard(1), dist.Shard(0)],
+        )
+
+        new_ops = dist_program.global_block().ops
+        new_ops_name = [op.name() for op in dist_program.global_block().ops]
+
+        all_gather_ops = []
+        slice_ops = []
+        for i, op in enumerate(new_ops):
+            if op.name() == "pd_op.all_gather":
+                all_gather_ops.append(op)
+            elif op.name() == "pd_op.slice":
+                slice_ops.append(op)
+        assert len(all_gather_ops) == 2
+        assert len(slice_ops) == 2
+        rank_id = dist.get_rank()
+
+        # check the first all_gather
+        op = all_gather_ops[0]
+        if rank_id == 0 or rank_id == 1:
+            process_ids = [0, 1]
+        elif rank_id == 2 or rank_id == 3:
+            process_ids = [2, 3]
+        tgt_operand = (
+            process_ids,
+            [-1, 0, -1],
+            {},
+        )  # process_ids, dims_mapping, partial_status
+        tgt_result = (process_ids, [-1, -1, -1], {})
+        tgt_in_value = (self._mesh.process_ids, [0, 1, -1], {})
+        tgt_out_value = (self._mesh.process_ids, [0, -1, -1], {})
+        self.validate(op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value)
+
+        # check the second all_gather
+        op = all_gather_ops[1]
+        if rank_id == 0 or rank_id == 2:
+            process_ids = [0, 2]
+        elif rank_id == 1 or rank_id == 3:
+            process_ids = [1, 3]
+        tgt_operand = (
+            process_ids,
+            [0, -1, -1],
+            {},
+        )  # process_ids, dims_mapping, partial_status
+        tgt_result = (process_ids, [-1, -1, -1], {})
+        tgt_in_value = (self._mesh.process_ids, [0, -1, -1], {})
+        tgt_out_value = (self._mesh.process_ids, [-1, -1, -1], {})
+        self.validate(op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value)
+
+        # check the first slice
+        op = slice_ops[0]
+        if rank_id == 0 or rank_id == 2:
+            process_ids = [0, 2]
+        elif rank_id == 1 or rank_id == 3:
+            process_ids = [1, 3]
+        tgt_operand = (process_ids, [-1, -1, -1], {})
+        tgt_result = (process_ids, [-1, 0, -1], {})
+        tgt_in_value = (self._mesh.process_ids, [-1, -1, -1], {})
+        tgt_out_value = (self._mesh.process_ids, [-1, 0, -1], {})
+        self.validate(op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value)
+
+        # check the second slice
+        op = slice_ops[1]
+        if rank_id == 0 or rank_id == 1:
+            process_ids = [0, 1]
+        elif rank_id == 2 or rank_id == 3:
+            process_ids = [2, 3]
+        tgt_operand = (process_ids, [-1, -1, -1], {})
+        tgt_result = (process_ids, [0, -1, -1], {})
+        tgt_in_value = (self._mesh.process_ids, [-1, 0, -1], {})
+        tgt_out_value = (self._mesh.process_ids, [1, 0, -1], {})
+        self.validate(op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value)
+
+    def run_ps_to_ps_case(self):
+        # [Partial(), Shard(0)] --> [Replicate(), Shard(1)]
+        # all_reduce sum + all_gather + slice
+        main_program, dist_program = self.create_program(
+            [self.BATCH_SIZE, self.SEQ_LEN, self.HIDDEN_SIZE],
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Shard(0)],
+            [dist.Replicate(), dist.Shard(1)],
+        )
+
+        ops = dist_program.global_block().ops
+        op_names = [op.name() for op in ops]
+        check_all_reduce_sum = any(
+            op.name() == "pd_op.all_reduce"
+            and op.int_attr("reduce_type") == dist.ReduceOp.SUM
+            for op in ops
+        )
+        assert check_all_reduce_sum
+        assert "pd_op.all_gather" in op_names
+        assert "pd_op.slice" in op_names
+
+        allgather_op = ops[op_names.index("pd_op.all_gather")]
+        allreduce_sum_op = ops[op_names.index("pd_op.all_reduce")]
+        slice_op = ops[op_names.index("pd_op.slice")]
+
+        # check the allgather
+        rank_id = dist.get_rank()
+        if rank_id in [0, 1]:
+            process_ids = [0, 1]
+        elif rank_id in [2, 3]:
+            process_ids = [2, 3]
+        tgt_operand = (process_ids, [0, -1, -1], {})
+        tgt_result = (process_ids, [-1, -1, -1], {})
+        tgt_in_value = (
+            self._mesh.process_ids,
+            [1, -1, -1],
+            {0: dist.ReduceType.kRedSum},
+        )
+        tgt_out_value = (
+            self._mesh.process_ids,
+            [-1, -1, -1],
+            {0: dist.ReduceType.kRedSum},
+        )
+        self.validate(
+            allgather_op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value
+        )
+
+        # check the allreduce_sum
+        if rank_id in [0, 2]:
+            process_ids = [0, 2]
+        elif rank_id in [1, 3]:
+            process_ids = [1, 3]
+        tgt_operand = (process_ids, [-1, -1, -1], {0: dist.ReduceType.kRedSum})
+        tgt_result = (process_ids, [-1, -1, -1], {})
+        tgt_in_value = (
+            self._mesh.process_ids,
+            [-1, -1, -1],
+            {0: dist.ReduceType.kRedSum},
+        )
+        tgt_out_value = (self._mesh.process_ids, [-1, -1, -1], {})
+        self.validate(
+            allreduce_sum_op,
+            tgt_operand,
+            tgt_result,
+            tgt_in_value,
+            tgt_out_value,
+        )
+
+        # check the slice
+        if rank_id in [0, 1]:
+            process_ids = [0, 1]
+        elif rank_id in [2, 3]:
+            process_ids = [2, 3]
+        tgt_operand = (process_ids, [-1, -1, -1], {})
+        tgt_result = (process_ids, [-1, 0, -1], {})
+        tgt_in_value = (self._mesh.process_ids, [-1, -1, -1], {})
+        tgt_out_value = (self._mesh.process_ids, [-1, 1, -1], {})
+        self.validate(
+            slice_op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value
+        )
+
+    def run_test_cases(self):
+        self.run_pp_to_rr_case()
+        self.run_pr_to_rs_case()
+        self.run_pr_to_ss_case()
+        self.run_ss_to_ss_case()
+        self.run_ps_to_ps_case()
+
+
+if __name__ == '__main__':
+    TestReshardNdMesh().run_test_cases()

test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh_cross_mesh.py

@@ -0,0 +1,179 @@
+# 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 paddle
+import paddle.distributed as dist
+from paddle.distributed.auto_parallel.static.pir_pass import (
+    ReshardPasses,
+)
+from paddle.distributed.auto_parallel.static.utils import set_all_ops_op_role
+from paddle.distributed.fleet.meta_optimizers.common import OpRole
+
+
+class TestReshardNdMeshCrossMesh:
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._seed = eval(os.getenv("seed"))
+        self.BATCH_SIZE = 2
+        self.SEQ_LEN = 4
+        self.HIDDEN_SIZE = 8
+        self._backend = os.getenv("backend")
+        self._mesh0 = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+        self._mesh1 = dist.ProcessMesh([[4, 5], [6, 7]], dim_names=["x", "y"])
+
+    def validate(
+        self, op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value
+    ):
+        # tgt_* are tuples, format: (process_ids, dims_mapping, partial_status)
+        operand_dist_attr = op.dist_attr.operand(0).as_tensor_dist_attr()
+        result_dist_attr = op.dist_attr.result(0).as_tensor_dist_attr()
+
+        assert operand_dist_attr.process_mesh.process_ids == tgt_operand[0]
+        assert operand_dist_attr.dims_mapping == tgt_operand[1]
+        assert operand_dist_attr.partial_status == tgt_operand[2]
+
+        assert result_dist_attr.process_mesh.process_ids == tgt_result[0]
+        assert result_dist_attr.dims_mapping == tgt_result[1]
+        assert result_dist_attr.partial_status == tgt_result[2]
+
+        in_value = op.operand_source(0)
+        out_value = op.result(0)
+        assert in_value.is_dist_dense_tensor_type()
+        assert out_value.is_dist_dense_tensor_type()
+        assert in_value.dist_attr().process_mesh.process_ids == tgt_in_value[0]
+        assert in_value.dist_attr().dims_mapping == tgt_in_value[1]
+        assert in_value.dist_attr().partial_status == tgt_in_value[2]
+
+        assert (
+            out_value.dist_attr().process_mesh.process_ids == tgt_out_value[0]
+        )
+        assert out_value.dist_attr().dims_mapping == tgt_out_value[1]
+        assert out_value.dist_attr().partial_status == tgt_out_value[2]
+
+    def create_program(self, input_shape, input_placements, output_placements):
+        paddle.enable_static()
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+            place = paddle.CPUPlace()
+        elif self._backend == "gpu":
+            place = paddle.CUDAPlace(dist.get_rank())
+
+        with paddle.pir_utils.IrGuard():
+            main_program = paddle.base.Program()
+            with paddle.base.program_guard(main_program):
+                input = paddle.ones(name='input', shape=input_shape)
+                dist_input = dist.shard_tensor(
+                    input, self._mesh0, input_placements
+                )
+                dist_out = paddle._C_ops.reshard(
+                    dist_input, self._mesh1, output_placements
+                )
+            dist_program = main_program.clone()
+            set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)
+            ReshardPasses.apply_reshard_pass(dist_program)
+
+        return main_program, dist_program
+
+    def run_pp_to_rr_case(self):
+        # [Partial(), Partial()] --> [Replicate(), Replicate()]
+        # ops: all_reduce sum + all_reduce sum
+        main_program, dist_program = self.create_program(
+            [self.BATCH_SIZE, self.SEQ_LEN, self.HIDDEN_SIZE],
+            [
+                dist.Partial(dist.ReduceType.kRedSum),
+                dist.Partial(dist.ReduceType.kRedSum),
+            ],
+            [dist.Replicate(), dist.Replicate()],
+        )
+
+        new_ops = dist_program.global_block().ops
+        old_ops_name = [op.name() for op in main_program.global_block().ops]
+        new_ops_name = [op.name() for op in dist_program.global_block().ops]
+
+        rank_id = dist.get_rank()
+        if rank_id in self._mesh0.process_ids:
+            assert dist_program.global_block().ops[2].name() == "pd_op.send_v2"
+        else:
+            assert dist_program.global_block().ops[2].name() == "pd_op.recv_v2"
+            assert (
+                dist_program.global_block().ops[-2].name() == "pd_op.all_reduce"
+            )
+            assert (
+                dist_program.global_block().ops[-2].int_attr("reduce_type")
+                == dist.ReduceOp.SUM
+            )
+            assert (
+                dist_program.global_block().ops[-1].name() == "pd_op.all_reduce"
+            )
+            assert (
+                dist_program.global_block().ops[-1].int_attr("reduce_type")
+                == dist.ReduceOp.SUM
+            )
+
+            # check the first allreduce_sum
+            op = new_ops[-2]
+            if rank_id == 4 or rank_id == 6:
+                process_ids = [4, 6]
+            elif rank_id == 5 or rank_id == 7:
+                process_ids = [5, 7]
+            tgt_operand = (
+                process_ids,
+                [-1, -1, -1],
+                {0: dist.ReduceType.kRedSum},
+            )
+            tgt_result = (process_ids, [-1, -1, -1], {})
+            tgt_in_value = (
+                self._mesh1.process_ids,
+                [-1, -1, -1],
+                {0: dist.ReduceType.kRedSum, 1: dist.ReduceType.kRedSum},
+            )
+            tgt_out_value = (
+                self._mesh1.process_ids,
+                [-1, -1, -1],
+                {1: dist.ReduceType.kRedSum},
+            )
+            self.validate(
+                op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value
+            )
+
+            # check the second allreduce_sum
+            op = new_ops[-1]
+            if rank_id == 4 or rank_id == 5:
+                process_ids = [4, 5]
+            elif rank_id == 6 or rank_id == 7:
+                process_ids = [6, 7]
+            tgt_operand = (
+                process_ids,
+                [-1, -1, -1],
+                {0: dist.ReduceType.kRedSum},
+            )
+            tgt_result = (process_ids, [-1, -1, -1], {})
+            tgt_in_value = (
+                self._mesh1.process_ids,
+                [-1, -1, -1],
+                {1: dist.ReduceType.kRedSum},
+            )
+            tgt_out_value = (self._mesh1.process_ids, [-1, -1, -1], {})
+            self.validate(
+                op, tgt_operand, tgt_result, tgt_in_value, tgt_out_value
+            )
+
+    def run_test_cases(self):
+        self.run_pp_to_rr_case()
+
+
+if __name__ == '__main__':
+    TestReshardNdMeshCrossMesh().run_test_cases()

test/auto_parallel/hybrid_strategy/process_mesh_demo_unittest.py

@@ -0,0 +1,138 @@
+# 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")
+        dp_mesh_ = mesh_2d["dp"]
+        assert dp_mesh == dp_mesh_
+        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")
+        tp_mesh_ = mesh_2d["tp"]
+        assert tp_mesh == tp_mesh_
+        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")
+        pp_mesh_ = mesh_3d["pp"]
+        assert pp_mesh == pp_mesh_
+        dp_mesh = mesh_3d.get_submesh_with_dim("dp")
+        dp_mesh_ = mesh_3d["dp"]
+        assert dp_mesh == dp_mesh_
+        tp_mesh = mesh_3d.get_submesh_with_dim("tp")
+        tp_mesh_ = mesh_3d["tp"]
+        assert tp_mesh == tp_mesh_
+
+        # 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
+            )
+            assert group_1d_with_name.id == group_1d.id
+            # 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/hybrid_strategy/save_load_state_dict_with_aoa_config_reverse.py

@@ -0,0 +1,100 @@
+# 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 os
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
+    load_state_dict,
+)
+
+
+class HuggingFaceModel(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.huggingface = nn.Linear(2, 2, bias_attr=False)
+
+
+class FCModel(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(1, 2, bias_attr=False)
+        self.fc2 = nn.Linear(1, 2, bias_attr=False)
+
+
+def init_hf_model_weights(model):
+    with paddle.no_grad():
+        w = paddle.to_tensor([[0, 1], [2, 3]], dtype="float16")
+        model.huggingface.weight.set_value(w)
+
+
+def save_safetensors_model(model, ckpt_path):
+    import safetensors.numpy
+
+    os.makedirs(ckpt_path, exist_ok=True)
+
+    weight_np = model.huggingface.weight.numpy()
+    file_path = os.path.join(ckpt_path, "tensor1.safetensors")
+    safetensors.numpy.save_file({"huggingface.weight": weight_np}, file_path)
+
+
+def test_save_load_with_aoa_config_reverse():
+    ckpt_path = os.getenv("ckpt_path")
+
+    dist.init_parallel_env()
+
+    hf_model = HuggingFaceModel()
+    fc_model = FCModel()
+    hf_model = paddle.amp.decorate(
+        models=hf_model, optimizers=None, level="O2", dtype="float16"
+    )
+    init_hf_model_weights(hf_model)
+
+    save_safetensors_model(hf_model, ckpt_path)
+
+    aoa_statements = [
+        "huggingface.weight -> A,B ,axis = 1 \n",
+        "A^T -> A  \n",
+        "B^T -> B \n",
+        "A -> fc1.weight ,src_dtype = 'float16', dst_dtype = 'float32' \n",
+        "B -> fc2.weight ,src_dtype = 'float16', dst_dtype = 'float32' \n",
+    ]
+    aoa_config = {"aoa_statements": aoa_statements}
+
+    load_state_dict(
+        fc_model.sharded_state_dict(),
+        ckpt_path,
+        safetensors=True,
+        aoa_config=aoa_config,
+    )
+
+    full = paddle.to_tensor([[0, 1], [2, 3]], dtype="float32")
+    A = full[:, 0].unsqueeze(0)
+    B = full[:, 1].unsqueeze(0)
+
+    assert paddle.allclose(fc_model.fc1.weight, A)
+    assert paddle.allclose(fc_model.fc2.weight, B)
+
+    aoa_config["aoa_config_reverse"] = True
+    itr = fc_model.full(aoa_config=aoa_config)
+    full_param = dict(itr)
+    (full_weight_k, full_weight_v) = next(iter(full_param.items()))
+    assert full_weight_k == "huggingface.weight"
+    assert paddle.allclose(full_weight_v, hf_model.huggingface.weight)
+
+
+if __name__ == "__main__":
+    test_save_load_with_aoa_config_reverse()

test/auto_parallel/hybrid_strategy/save_safetensors_load_fc.py

@@ -0,0 +1,223 @@
+# 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 json
+import os
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import ShardedWeight
+
+
+def get_global_tensors():
+    """Create fixed test tensors for verification."""
+    # tensor1: [[0, 1], [2, 3]]
+    tensor1 = paddle.to_tensor([[0, 1], [2, 3]], dtype='float32')
+    # tensor2: [[4, 5], [6, 7]]
+    tensor2 = paddle.to_tensor([[4, 5], [6, 7]], dtype='float32')
+    return {"tensor1": tensor1, "tensor2": tensor2}
+
+
+def save_safetensors_to_ranks(ckpt_path):
+    """Save tensors to different ranks as safetensors files."""
+    import safetensors.numpy
+
+    global_tensors = get_global_tensors()
+
+    if dist.get_rank() == 0:
+        os.makedirs(ckpt_path, exist_ok=True)
+        file_path = os.path.join(ckpt_path, "tensor1.safetensors")
+
+        tensor1_np = global_tensors["tensor1"].numpy()
+        safetensors.numpy.save_file({"tensor1": tensor1_np}, file_path)
+
+    elif dist.get_rank() == 1:
+        os.makedirs(ckpt_path, exist_ok=True)
+        file_path = os.path.join(ckpt_path, "tensor2.safetensors")
+
+        tensor2_np = global_tensors["tensor2"].numpy()
+        safetensors.numpy.save_file({"tensor2": tensor2_np}, file_path)
+
+    dist.barrier()
+
+
+def create_sharded_state_dict_for_loading():
+    """Create sharded state dict for tp loading."""
+    sharded_state_dict = {}
+
+    if dist.get_rank() == 0:
+        local_tensor1 = paddle.zeros([2, 1], dtype='float32')
+        sharded_weight1 = ShardedWeight(
+            key="tensor1",
+            local_tensor=local_tensor1,
+            local_shape=(2, 1),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+            is_flattened=False,
+        )
+        sharded_state_dict["tensor1"] = sharded_weight1
+
+        local_tensor2 = paddle.zeros([2, 1], dtype='float32')
+        sharded_weight2 = ShardedWeight(
+            key="tensor2",
+            local_tensor=local_tensor2,
+            local_shape=(2, 1),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+            is_flattened=False,
+        )
+        sharded_state_dict["tensor2"] = sharded_weight2
+
+    elif dist.get_rank() == 1:
+        local_tensor1 = paddle.zeros([2, 1], dtype='float32')
+        sharded_weight1 = ShardedWeight(
+            key="tensor1",
+            local_tensor=local_tensor1,
+            local_shape=(2, 1),
+            global_shape=(2, 2),
+            global_offset=(0, 1),
+            is_flattened=False,
+        )
+        sharded_state_dict["tensor1"] = sharded_weight1
+
+        local_tensor2 = paddle.zeros([2, 1], dtype='float32')
+        sharded_weight2 = ShardedWeight(
+            key="tensor2",
+            local_tensor=local_tensor2,
+            local_shape=(2, 1),
+            global_shape=(2, 2),
+            global_offset=(0, 1),
+            is_flattened=False,
+        )
+        sharded_state_dict["tensor2"] = sharded_weight2
+
+    return sharded_state_dict
+
+
+def test_save_safetensors_load_fc():
+    """Test saving safetensors and loading with flex checkpoint."""
+    ckpt_path = os.getenv("ckpt_path")
+    dist.init_parallel_env()
+
+    save_safetensors_to_ranks(ckpt_path)
+
+    sharded_state_dict = create_sharded_state_dict_for_loading()
+
+    from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
+        load_state_dict,
+    )
+
+    load_state_dict(sharded_state_dict, ckpt_path, safetensors=True)
+
+    loaded_tensor1 = sharded_state_dict["tensor1"].local_tensor
+    loaded_tensor2 = sharded_state_dict["tensor2"].local_tensor
+
+    if dist.get_rank() == 0:
+        # Rank 0 should have first column of both tensors
+        # tensor1: [[0], [2]] (first column)
+        # tensor2: [[4], [6]] (first column)
+        expected_tensor1 = paddle.to_tensor([[0], [2]], dtype='float32')
+        expected_tensor2 = paddle.to_tensor([[4], [6]], dtype='float32')
+
+        assert paddle.allclose(loaded_tensor1, expected_tensor1), (
+            f"Rank 0 tensor1 mismatch: got {loaded_tensor1}, expected {expected_tensor1}"
+        )
+        assert paddle.allclose(loaded_tensor2, expected_tensor2), (
+            f"Rank 0 tensor2 mismatch: got {loaded_tensor2}, expected {expected_tensor2}"
+        )
+
+    elif dist.get_rank() == 1:
+        # Rank 1 should have second column of both tensors
+        # tensor1: [[1], [3]] (second column)
+        # tensor2: [[5], [7]] (second column)
+        expected_tensor1 = paddle.to_tensor([[1], [3]], dtype='float32')
+        expected_tensor2 = paddle.to_tensor([[5], [7]], dtype='float32')
+
+        assert paddle.allclose(loaded_tensor1, expected_tensor1), (
+            f"Rank 1 tensor1 mismatch: got {loaded_tensor1}, expected {expected_tensor1}"
+        )
+        assert paddle.allclose(loaded_tensor2, expected_tensor2), (
+            f"Rank 1 tensor2 mismatch: got {loaded_tensor2}, expected {expected_tensor2}"
+        )
+
+    dist.barrier()
+
+
+def create_index_json(ckpt_path):
+    """Create model.safetensors.index.json that maps keys to their safetensors files."""
+    index_data = {
+        "weight_map": {
+            "tensor1": "tensor1.safetensors",
+            "tensor2": "tensor2.safetensors",
+        }
+    }
+    index_file_path = os.path.join(ckpt_path, "model.safetensors.index.json")
+    if dist.get_rank() == 0:
+        with open(index_file_path, "w") as f:
+            json.dump(index_data, f)
+    dist.barrier()
+
+
+def test_save_safetensors_load_fc_with_index():
+    """Test saving safetensors and loading with flex checkpoint when model.safetensors.index.json exists."""
+    ckpt_path = os.getenv("ckpt_path")
+    dist.init_parallel_env()
+
+    save_safetensors_to_ranks(ckpt_path)
+
+    # Create index json to exercise the integrity check branch
+    create_index_json(ckpt_path)
+
+    sharded_state_dict = create_sharded_state_dict_for_loading()
+
+    from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
+        load_state_dict,
+    )
+
+    load_state_dict(sharded_state_dict, ckpt_path, safetensors=True)
+
+    loaded_tensor1 = sharded_state_dict["tensor1"].local_tensor
+    loaded_tensor2 = sharded_state_dict["tensor2"].local_tensor
+
+    if dist.get_rank() == 0:
+        expected_tensor1 = paddle.to_tensor([[0], [2]], dtype='float32')
+        expected_tensor2 = paddle.to_tensor([[4], [6]], dtype='float32')
+
+        assert paddle.allclose(loaded_tensor1, expected_tensor1), (
+            f"Rank 0 tensor1 mismatch: got {loaded_tensor1}, expected {expected_tensor1}"
+        )
+        assert paddle.allclose(loaded_tensor2, expected_tensor2), (
+            f"Rank 0 tensor2 mismatch: got {loaded_tensor2}, expected {expected_tensor2}"
+        )
+
+    elif dist.get_rank() == 1:
+        expected_tensor1 = paddle.to_tensor([[1], [3]], dtype='float32')
+        expected_tensor2 = paddle.to_tensor([[5], [7]], dtype='float32')
+
+        assert paddle.allclose(loaded_tensor1, expected_tensor1), (
+            f"Rank 1 tensor1 mismatch: got {loaded_tensor1}, expected {expected_tensor1}"
+        )
+        assert paddle.allclose(loaded_tensor2, expected_tensor2), (
+            f"Rank 1 tensor2 mismatch: got {loaded_tensor2}, expected {expected_tensor2}"
+        )
+
+    dist.barrier()
+
+
+if __name__ == "__main__":
+    test_func = os.getenv("test_func", "test_save_safetensors_load_fc")
+    if test_func == "test_save_safetensors_load_fc_with_index":
+        test_save_safetensors_load_fc_with_index()
+    else:
+        test_save_safetensors_load_fc()

test/auto_parallel/hybrid_strategy/semi_auto_llama.py

@@ -0,0 +1,308 @@
+# 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
+from functools import reduce
+
+import numpy as np
+from semi_auto_parallel_llama_model import (
+    LlamaForCausalLMAuto,
+    LlamaPretrainingCriterionAuto,
+    get_mesh,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle import LazyGuard
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+
+class Config:
+    vocab_size = 32000
+    hidden_size = 4096
+    intermediate_size = 11008
+    max_position_embeddings = 2048
+    seq_length = 2048
+    num_hidden_layers = 2
+    num_attention_heads = 32
+    num_key_value_heads = 32
+    initializer_range = 0.02
+    rms_norm_eps = 1e-6
+    use_cache = True
+    use_flash_attention = False
+    sequence_parallel = False
+    rope = True
+    recompute = False
+    recompute_granularity = None
+    use_lazy_init = False
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, num_samples=100):
+        super().__init__()
+        self.seq_len = seq_len
+        self.num_samples = num_samples
+
+    def __getitem__(self, index):
+        input = np.random.uniform(size=[self.seq_len]).astype("int64")
+        label = (np.random.uniform(size=[self.seq_len]) * 10).astype("int64")
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+def create_optimizer(model, lr_scheduler):
+    decay_parameters = [
+        p.name
+        for n, p in model.named_parameters()
+        if not any(nd in n for nd in ["bias", "norm"])
+    ]
+
+    def apply_decay_param_fun(x):
+        return x in decay_parameters
+
+    # test global_clip in auto_parallel
+    if os.getenv("use_param_group") == "true":
+        param_group = {}
+        param_group["params"] = list(model.parameters())
+        param_group["weight_decay"] = 0.01
+        param_group["grad_clip"] = paddle.nn.ClipGradByGlobalNorm(1.0)
+        optimizer = paddle.optimizer.adamw.AdamW(
+            learning_rate=lr_scheduler,
+            apply_decay_param_fun=apply_decay_param_fun,
+            parameters=[param_group],
+        )
+    else:
+        optimizer = paddle.optimizer.adamw.AdamW(
+            learning_rate=lr_scheduler,
+            apply_decay_param_fun=apply_decay_param_fun,
+            parameters=model.parameters(),
+            weight_decay=0.01,
+            grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+        )
+    return optimizer
+
+
+class TestLlamaAuto:
+    def __init__(self):
+        self.config = Config()
+        self.dp = int(os.getenv("dp"))
+        self.mp = int(os.getenv("mp"))
+        self.pp = int(os.getenv("pp"))
+        self.sep = int(os.getenv("sep", "1"))
+        if os.getenv("use_sp") == "true":
+            self.config.sequence_parallel = True
+        if os.getenv("recompute") == "true":
+            self.config.recompute = True
+        self.config.recompute_granularity = os.getenv("recompute_granularity")
+        if os.getenv("use_lazy_init") == "true":
+            self.config.use_lazy_init = True
+        self.gradient_accumulation_steps = int(os.getenv("acc_step"))
+        self.amp = False
+        self.amp_dtype = "float16"
+        self.amp_level = "O1"
+        self.amp_master_grad = False
+        if os.getenv("amp") == "true":
+            self.amp = True
+        if os.getenv("amp_dtype") in ["float16", "bfloat16"]:
+            self.amp_dtype = os.getenv("amp_dtype")
+        if os.getenv("amp_level") in ["O0", "O1", "O2"]:
+            self.amp_level = os.getenv("amp_level")
+        if os.getenv("amp_master_grad") == "true":
+            self.amp_master_grad = True
+        self.config.tensor_parallel_degree = self.mp
+        self.config.pipeline_parallel_degree = self.pp
+        self.config.context_parallel_degree = 1
+        self.config.sep_parallel_degree = 1
+        if os.getenv("context_parallel", "false") == "true":
+            self.config.context_parallel_degree = self.sep
+            self.config.use_flash_attention = True
+            dist.init_parallel_env()
+        if os.getenv("sep_parallel", "false") == "true":
+            self.config.sep_parallel_degree = self.sep
+        if self.sep > 1:
+            # only one of the context_parallel and sep_parallel can be True
+            assert (
+                self.config.sep_parallel_degree
+                != self.config.context_parallel_degree
+            ), (
+                f"only one of the context_parallel and sep_parallel can be True, but get context_parallel_degree = {self.config.context_parallel_degree} and sep_parallel_degree = {self.config.sep_parallel_degree}, please check your env"
+            )
+
+        self.init_dist_env()
+
+    def init_dist_env(self):
+        order = ["dp", "pp", "mp"]
+        dp_degree = self.dp
+        mp_degree = self.mp
+        pp_degree = self.pp
+        degree = [dp_degree, pp_degree, mp_degree]
+        mesh_dims = list(filter(lambda x: x[1] > 1, list(zip(order, degree))))
+        if not mesh_dims:
+            mesh_dims = [("dp", 1)]
+        dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
+        mesh_shape = [mesh_dim[1] for mesh_dim in mesh_dims]
+        mesh_arr = np.arange(
+            0, reduce(lambda x, y: x * y, mesh_shape, 1)
+        ).reshape(mesh_shape)
+        global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
+        dist.auto_parallel.set_mesh(global_mesh)
+
+    def run_llama(self, to_static=0):
+        if self.config.use_lazy_init:
+            with LazyGuard():
+                model = LlamaForCausalLMAuto(self.config)
+            for param in model.parameters():
+                assert not param._is_initialized()
+                param.initialize()
+        else:
+            model = LlamaForCausalLMAuto(self.config)
+        criterion = LlamaPretrainingCriterionAuto(self.config)
+
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+        optimizer = create_optimizer(model, lr_scheduler)
+        if self.amp and not to_static:
+            model, optimizer = paddle.amp.decorate(
+                models=model,
+                optimizers=optimizer,
+                level=self.amp_level,
+                dtype=self.amp_dtype,
+                master_grad=self.amp_master_grad,
+            )
+        optimizer = dist.shard_optimizer(optimizer)
+
+        train_dataset = RandomDataset(self.config.seq_length)
+        train_sampler = BatchSampler(
+            train_dataset,
+            batch_size=2,
+            shuffle=True,
+            drop_last=True,
+        )
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_sampler=train_sampler,
+            num_workers=0,
+        )
+
+        if self.pp == 1:
+            meshes = [get_mesh(0)]
+        elif self.pp > 1:
+            meshes = [get_mesh(0), get_mesh(-1)]
+        else:
+            raise ValueError("pp should be greater or equal to 1")
+
+        dist_loader = dist.shard_dataloader(
+            dataloader=train_dataloader,
+            meshes=meshes,
+            shard_dims="dp",
+        )
+
+        global_step = 1
+        tr_loss = float(0)
+
+        if not to_static:
+            model.train()
+            scaler = None
+            if self.amp and self.amp_dtype == "float16":
+                scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
+                scaler = dist.shard_scaler(scaler)
+
+            for epoch_idx in range(1):
+                for step, inputs in enumerate(dist_loader()):
+                    input_ids, labels = inputs
+                    custom_black_list = [
+                        "reduce_sum",
+                        "c_softmax_with_cross_entropy",
+                    ]
+                    custom_white_list = []
+                    if self.amp_level == "O2":
+                        custom_white_list.extend(
+                            ["lookup_table", "lookup_table_v2"]
+                        )
+                    with paddle.amp.auto_cast(
+                        self.amp,
+                        custom_black_list=set(custom_black_list),
+                        custom_white_list=set(custom_white_list),
+                        level=self.amp_level,
+                        dtype=self.amp_dtype,
+                    ):
+                        logits = model(input_ids)
+                        tr_loss_step = criterion(logits, labels)
+
+                    if self.gradient_accumulation_steps > 1:
+                        tr_loss_step /= self.gradient_accumulation_steps
+                    if scaler is not None:
+                        scaler.scale(tr_loss_step).backward()
+                    else:
+                        tr_loss_step.backward()
+                    tr_loss += tr_loss_step
+
+                    if global_step % self.gradient_accumulation_steps == 0:
+                        print(
+                            f"step: {global_step // self.gradient_accumulation_steps}  loss: {tr_loss.numpy()}"
+                        )
+                        if scaler is not None:
+                            scaler.step(optimizer)
+                            scaler.update()
+                        else:
+                            optimizer.step()
+                        optimizer.clear_grad()
+                        lr_scheduler.step()
+                        tr_loss = 0
+
+                    global_step += 1
+                    if global_step // self.gradient_accumulation_steps >= 10:
+                        break
+        else:
+            strategy = dist.Strategy()
+            if self.gradient_accumulation_steps > 1:
+                strategy.pipeline.accumulate_steps = (
+                    self.gradient_accumulation_steps
+                )
+
+            if self.amp:
+                amp = strategy.amp
+                amp.enable = self.amp
+                amp.dtype = self.amp_dtype
+                amp.level = self.amp_level.lower()
+                if self.amp_master_grad:
+                    amp.use_master_grad = True
+
+            dist_model = dist.to_static(
+                model,
+                dist_loader,
+                criterion,
+                optimizer,
+                strategy=strategy,
+            )
+
+            dist_model.train()
+            for step, inputs in enumerate(dist_loader()):
+                input_ids, labels = inputs
+                loss = dist_model(input_ids, labels)
+                print(step, loss)
+
+                if step >= 10:
+                    break
+
+    def run_test_cases(self):
+        self.run_llama(to_static=0)
+        self.run_llama(to_static=1)
+
+
+if __name__ == '__main__':
+    TestLlamaAuto().run_test_cases()

test/auto_parallel/hybrid_strategy/semi_auto_llama_acc_align.py

@@ -0,0 +1,419 @@
+# 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
+
+os.environ["FLAGS_enable_pir_api"] = str(1)
+import random
+from functools import reduce
+
+import numpy as np
+from semi_auto_parallel_llama_model import (
+    LlamaForCausalLMAuto,
+    LlamaPretrainingCriterionAuto,
+    get_mesh,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+
+class Config:
+    vocab_size = 320
+    hidden_size = 8
+    intermediate_size = 64
+    max_position_embeddings = 8
+    seq_length = 8
+
+    num_hidden_layers = 2
+    num_attention_heads = 2
+    num_key_value_heads = 2
+    initializer_range = 0.02
+    rms_norm_eps = 1e-6
+    use_cache = True
+    use_flash_attention = False
+    sequence_parallel = False
+    rope = True
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, num_samples=100):
+        super().__init__()
+        self.seq_len = seq_len
+        self.num_samples = num_samples
+
+    def __getitem__(self, index):
+        input = np.full([self.seq_len], index, dtype="int64")
+        label = np.array([index] * self.seq_len)
+
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+def create_optimizer(model, lr_scheduler):
+    decay_parameters = [
+        p.name
+        for n, p in model.named_parameters()
+        if not any(nd in n for nd in ["bias", "norm"])
+    ]
+
+    def apply_decay_param_fun(x):
+        return x in decay_parameters
+
+    optimizer = paddle.optimizer.adamw.AdamW(
+        learning_rate=lr_scheduler,
+        apply_decay_param_fun=apply_decay_param_fun,
+        parameters=model.parameters(),
+        weight_decay=0.01,
+        grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+    )
+    return optimizer
+
+
+class TestLlamaAuto:
+    def __init__(self):
+        self.config = Config()
+        self.dp = int(os.getenv("dp"))
+        self.mp = int(os.getenv("mp"))
+        self.pp = int(os.getenv("pp"))
+        self.sep = int(os.getenv("sep", "1"))
+        if os.getenv("use_sp") == "true":
+            self.config.sequence_parallel = True
+
+        if os.getenv("seq_length"):
+            self.config.seq_length = int(os.getenv("seq_length"))
+        if os.getenv("hidden_size"):
+            self.config.hidden_size = int(os.getenv("hidden_size"))
+        if os.getenv("num_attention_heads"):
+            self.config.num_attention_heads = int(
+                os.getenv("num_attention_heads")
+            )
+        if os.getenv("num_key_value_heads"):
+            self.config.num_key_value_heads = int(
+                os.getenv("num_key_value_heads")
+            )
+        if os.getenv("max_position_embeddings"):
+            self.config.max_position_embeddings = int(
+                os.getenv("max_position_embeddings")
+            )
+        self.strategy = dist.Strategy()
+
+        # amp config
+        amp = self.strategy._amp
+        if os.getenv("amp"):
+            amp.enable = True if os.getenv("amp") == "true" else False
+        if os.getenv("amp_dtype"):
+            amp.dtype = os.getenv("amp_dtype")
+        if os.getenv("amp_level"):
+            amp.level = os.getenv("amp_level")
+        if os.getenv("amp_master_grad"):
+            amp.use_master_grad = (
+                True if os.getenv("amp_master_grad") == "true" else False
+            )
+        if os.getenv("scale_loss"):
+            amp.init_loss_scaling = os.getenv("scale_loss")
+        if os.getenv("amp_custom_black_list"):
+            amp.custom_black_list = os.getenv("amp_custom_black_list")
+        if os.getenv("amp_custom_white_list"):
+            amp.custom_white_list = os.getenv("amp_custom_white_list")
+        self.gradient_accumulation_steps = 1
+        if os.getenv("acc_step"):
+            self.gradient_accumulation_steps = int(os.getenv("acc_step"))
+
+        if self.gradient_accumulation_steps > 1:
+            self.strategy.gradient_merge.enable = True
+            self.strategy.gradient_merge.k_steps = (
+                self.gradient_accumulation_steps
+            )
+            self.strategy.gradient_merge.avg = False
+
+        self.config.recompute = False
+
+        self.config.tensor_parallel_degree = self.mp
+        self.config.pipeline_parallel_degree = self.pp
+        self.config.context_parallel_degree = 1
+        self.config.sep_parallel_degree = 1
+        if os.getenv("context_parallel", "false") == "true":
+            self.config.context_parallel_degree = self.sep
+            self.config.use_flash_attention = True
+            dist.init_parallel_env()
+        if os.getenv("sep_parallel", "false") == "true":
+            self.config.sep_parallel_degree = self.sep
+
+        if self.sep > 1:
+            # only one of the context_parallel and sep_parallel can be True
+            assert (
+                self.config.sep_parallel_degree
+                != self.config.context_parallel_degree
+            ), (
+                f"only one of the context_parallel and sep_parallel can be True, but get context_parallel_degree = {self.config.context_parallel_degree} and sep_parallel_degree = {self.config.sep_parallel_degree}, please check your env"
+            )
+
+        self.run_step = 10
+        self.run_step_dy2static = (
+            self.run_step // self.gradient_accumulation_steps
+        )
+
+    def run_llama(self, to_static=0):
+        # model
+        model = LlamaForCausalLMAuto(self.config)
+        criterion = LlamaPretrainingCriterionAuto(self.config)
+        if self.strategy._amp.enable and self.strategy._amp.level == "O2":
+            paddle.amp.decorate(
+                models=model,
+                level=self.strategy._amp.level,
+                dtype=self.strategy._amp.dtype,
+                master_grad=self.strategy._amp.use_master_grad,
+            )
+
+        # optimizer
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+        optimizer = create_optimizer(model, lr_scheduler)
+        optimizer = dist.shard_optimizer(optimizer)
+
+        # dataloader
+        train_dataset = RandomDataset(self.config.seq_length)
+        train_sampler = BatchSampler(
+            train_dataset,
+            batch_size=4 if self.sep > 1 else 2,
+            shuffle=True,
+            drop_last=True,
+        )
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_sampler=train_sampler,
+            num_workers=0,
+        )
+        dist_loader = dist.shard_dataloader(
+            dataloader=train_dataloader,
+            meshes=[get_mesh(0), get_mesh(1)],
+            shard_dims="dp",
+        )
+
+        if to_static:
+            model = dist.to_static(
+                model, dist_loader, criterion, optimizer, strategy=self.strategy
+            )
+        model.train()
+        losses = []
+
+        for step, inputs in enumerate(dist_loader()):
+            if step >= self.run_step:
+                break
+            input_ids, labels = inputs
+            if to_static:
+                loss = model(input_ids, labels)
+                if loss is None:
+                    numpy_array = np.array([])
+                else:
+                    numpy_array = np.array(loss)
+                losses.append(numpy_array)
+                array_bytes = numpy_array.tobytes()
+            else:
+                logits = model(input_ids)
+                loss = criterion(logits, labels)
+                losses.append(np.array(loss))
+                loss.backward()
+                optimizer.step()
+                optimizer.clear_grad()
+            lr_scheduler.step()
+        return losses
+
+    def init_dist_env(self):
+        order = ["dp", "pp", "mp", "sep"]
+        dp_degree = self.dp
+        mp_degree = self.mp
+        pp_degree = self.pp
+        sep_degree = self.sep
+        degree = [dp_degree, pp_degree, mp_degree, sep_degree]
+
+        mesh_dims = list(filter(lambda x: x[1] > 1, list(zip(order, degree))))
+        if not mesh_dims:
+            mesh_dims = [("dp", 1)]
+
+        dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
+        mesh_shape = [mesh_dim[1] for mesh_dim in mesh_dims]
+        mesh_arr = np.arange(
+            0, reduce(lambda x, y: x * y, mesh_shape, 1)
+        ).reshape(mesh_shape)
+
+        global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
+        dist.auto_parallel.set_mesh(global_mesh)
+
+        paddle.seed(1024)
+        np.random.seed(1024)
+        random.seed(1024)
+
+    def run_dynamic(self):
+        model = LlamaForCausalLMAuto(self.config)
+        criterion = LlamaPretrainingCriterionAuto(self.config)
+
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+        optimizer = create_optimizer(model, lr_scheduler)
+
+        train_dataset = RandomDataset(self.config.seq_length)
+        train_sampler = BatchSampler(
+            train_dataset,
+            batch_size=4 if self.sep > 1 else 2,
+            shuffle=False,
+            drop_last=True,
+        )
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_sampler=train_sampler,
+            num_workers=0,
+        )
+        dist_loader = dist.shard_dataloader(
+            dataloader=train_dataloader,
+            meshes=[get_mesh(0), get_mesh(1)],
+            shard_dims="dp",
+        )
+
+        tr_loss = float(0)
+        tr_loss_add = float(0)
+        model.train()
+        #####
+        for step, inputs in enumerate(dist_loader()):
+            if step >= self.run_step:
+                break
+
+            input_ids, labels = inputs
+            logits = model(input_ids)
+            tr_loss_step = criterion(logits, labels)
+
+            tr_loss_step.backward()
+            tr_loss_add += tr_loss_step
+
+            if int(dist.get_rank()) in [2, 3, 6, 7]:
+                assert tr_loss_step._is_initialized()
+            else:
+                assert not tr_loss_step._is_initialized()
+
+            if (step + 1) % self.gradient_accumulation_steps == 0:
+                tr_loss_add /= self.gradient_accumulation_steps
+                tr_loss = tr_loss_add
+
+                optimizer.step()
+                optimizer.clear_grad()
+                lr_scheduler.step()
+
+                tr_loss_add = 0
+        return np.array(tr_loss)
+
+    def run_dy2static(self):
+        model = LlamaForCausalLMAuto(self.config)
+        criterion = LlamaPretrainingCriterionAuto(self.config)
+
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+        optimizer = create_optimizer(model, lr_scheduler)
+
+        train_dataset = RandomDataset(self.config.seq_length)
+        train_sampler = BatchSampler(
+            train_dataset,
+            batch_size=(
+                4 * self.gradient_accumulation_steps
+                if self.sep > 1
+                else 2 * self.gradient_accumulation_steps
+            ),
+            shuffle=False,
+            drop_last=True,
+        )
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_sampler=train_sampler,
+            num_workers=0,
+        )
+        dist_loader = dist.shard_dataloader(
+            dataloader=train_dataloader,
+            meshes=[get_mesh(0), get_mesh(1)],
+            shard_dims="dp",
+        )
+
+        strategy = None
+        if self.gradient_accumulation_steps > 1:
+            strategy = dist.Strategy()
+
+            strategy.gradient_merge.enable = True
+            strategy.gradient_merge.k_steps = self.gradient_accumulation_steps
+            strategy.gradient_merge.avg = False
+
+        dist_model = dist.to_static(
+            model, dist_loader, criterion, optimizer, strategy=strategy
+        )
+
+        dist_model.train()
+        loss = None
+        for step, inputs in enumerate(dist_loader()):
+            if step >= self.run_step_dy2static:
+                break
+
+            input_ids, labels = inputs
+            loss = dist_model(input_ids, labels)
+
+            lr_scheduler.step()
+            if int(dist.get_rank()) in [2, 3, 6, 7]:
+                assert loss is not None
+            else:
+                assert loss is None
+
+        if loss is not None:
+            loss = np.average(loss)
+        return np.array(loss)
+
+    def run_test_cases(self):
+        self.init_dist_env()
+        # context parallel with flash_attn backend not support CPU, not support float32
+        # flash_attn only support Cuda Compute Capability >= 8 and cuda version >= 11
+        if self.config.context_parallel_degree > 1 and (
+            os.getenv("backend") != "gpu"
+            or not self.strategy._amp.enable
+            or int(paddle.version.cuda().split(".")[0]) < 11
+            or paddle.device.cuda.get_device_capability()[0] < 8
+        ):
+            return
+        if self.gradient_accumulation_steps > 1:
+            dy_losses = self.run_dynamic()
+            # context parallel not support static mode
+            if self.sep > 1:
+                return
+            self.init_dist_env()
+            st_losses = self.run_dy2static()
+            if int(dist.get_rank()) in [2, 3, 6, 7]:
+                np.testing.assert_allclose(dy_losses, st_losses, atol=1e-7)
+
+        else:
+            dy_losses = self.run_llama(to_static=0)
+            # context parallel not support static mode
+            if self.sep > 1:
+                return
+            self.init_dist_env()
+            st_losses = self.run_llama(to_static=1)
+            assert len(dy_losses) == len(st_losses)
+            for idx in range(len(dy_losses)):
+                np.testing.assert_allclose(
+                    dy_losses[idx], st_losses[idx], atol=1e-7
+                )
+
+
+if __name__ == '__main__':
+    TestLlamaAuto().run_test_cases()

test/auto_parallel/hybrid_strategy/semi_auto_llama_dataloader.py

@@ -0,0 +1,281 @@
+# 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
+from functools import reduce
+
+import numpy as np
+from semi_auto_parallel_llama_model import (
+    LlamaForCausalLMAuto,
+    LlamaPretrainingCriterionAuto,
+    get_mesh,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle import LazyGuard
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+np.random.seed(2024)
+
+
+class Config:
+    vocab_size = 32000
+    hidden_size = 4096
+    intermediate_size = 11008
+    max_position_embeddings = 2048
+    seq_length = 2048
+    num_hidden_layers = 2
+    num_attention_heads = 32
+    num_key_value_heads = 32
+    initializer_range = 0.02
+    rms_norm_eps = 1e-6
+    use_cache = True
+    use_flash_attention = False
+    sequence_parallel = False
+    rope = True
+    recompute = False
+    recompute_granularity = None
+    use_lazy_init = False
+
+
+inputs = []
+labels = []
+
+for i in range(100):
+    inputs.append(
+        np.random.uniform(low=0, high=32000, size=[Config().seq_length]).astype(
+            "int64"
+        )
+    )
+    labels.append(
+        (np.random.uniform(size=[Config().seq_length]) * 10).astype("int64")
+    )
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, num_samples=100):
+        super().__init__()
+        self.seq_len = seq_len
+        self.num_samples = num_samples
+
+    def __getitem__(self, index):
+        global inputs, labels
+        return inputs[index], labels[index]
+
+    def __len__(self):
+        return self.num_samples
+
+
+def create_optimizer(model, lr_scheduler):
+    decay_parameters = [
+        p.name
+        for n, p in model.named_parameters()
+        if not any(nd in n for nd in ["bias", "norm"])
+    ]
+
+    def apply_decay_param_fun(x):
+        return x in decay_parameters
+
+    # test global_clip in auto_parallel
+    if os.getenv("use_param_group") == "true":
+        param_group = {}
+        param_group["params"] = list(model.parameters())
+        param_group["weight_decay"] = 0.01
+        param_group["grad_clip"] = paddle.nn.ClipGradByGlobalNorm(1.0)
+        optimizer = paddle.optimizer.adamw.AdamW(
+            learning_rate=lr_scheduler,
+            apply_decay_param_fun=apply_decay_param_fun,
+            parameters=[param_group],
+        )
+    else:
+        optimizer = paddle.optimizer.adamw.AdamW(
+            learning_rate=lr_scheduler,
+            apply_decay_param_fun=apply_decay_param_fun,
+            parameters=model.parameters(),
+            weight_decay=0.01,
+            grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+        )
+    return optimizer
+
+
+class TestLlamaAuto:
+    def __init__(self):
+        self.config = Config()
+        self.dp = int(os.getenv("dp"))
+        self.mp = int(os.getenv("mp"))
+        self.pp = int(os.getenv("pp"))
+        self.sep = int(os.getenv("sep", "1"))
+        if os.getenv("use_sp") == "true":
+            self.config.sequence_parallel = True
+        if os.getenv("recompute") == "true":
+            self.config.recompute = True
+        self.config.recompute_granularity = os.getenv("recompute_granularity")
+        if os.getenv("use_lazy_init") == "true":
+            self.config.use_lazy_init = True
+        self.gradient_accumulation_steps = int(os.getenv("acc_step"))
+        self.amp = False
+        self.amp_dtype = "float16"
+        self.amp_level = "O1"
+        self.amp_master_grad = False
+        if os.getenv("amp") == "true":
+            self.amp = True
+        if os.getenv("amp_dtype") in ["float16", "bfloat16"]:
+            self.amp_dtype = os.getenv("amp_dtype")
+        if os.getenv("amp_level") in ["O0", "O1", "O2"]:
+            self.amp_level = os.getenv("amp_level")
+        if os.getenv("amp_master_grad") == "true":
+            self.amp_master_grad = True
+        self.config.tensor_parallel_degree = self.mp
+        self.config.pipeline_parallel_degree = self.pp
+        self.config.context_parallel_degree = 1
+        self.config.sep_parallel_degree = 1
+        if os.getenv("context_parallel", "false") == "true":
+            self.config.context_parallel_degree = self.sep
+            self.config.use_flash_attention = True
+            dist.init_parallel_env()
+        if os.getenv("sep_parallel", "false") == "true":
+            self.config.sep_parallel_degree = self.sep
+        if self.sep > 1:
+            # only one of the context_parallel and sep_parallel can be True
+            assert (
+                self.config.sep_parallel_degree
+                != self.config.context_parallel_degree
+            ), (
+                f"only one of the context_parallel and sep_parallel can be True, but get context_parallel_degree = {self.config.context_parallel_degree} and sep_parallel_degree = {self.config.sep_parallel_degree}, please check your env"
+            )
+
+        self.init_dist_env()
+
+    def init_dist_env(self):
+        order = ["dp", "pp", "mp"]
+        dp_degree = self.dp
+        mp_degree = self.mp
+        pp_degree = self.pp
+        degree = [dp_degree, pp_degree, mp_degree]
+        mesh_dims = list(filter(lambda x: x[1] > 1, list(zip(order, degree))))
+        if not mesh_dims:
+            mesh_dims = [("dp", 1)]
+        dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
+        mesh_shape = [mesh_dim[1] for mesh_dim in mesh_dims]
+        mesh_arr = np.arange(
+            0, reduce(lambda x, y: x * y, mesh_shape, 1)
+        ).reshape(mesh_shape)
+        global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
+        dist.auto_parallel.set_mesh(global_mesh)
+
+    def run_llama(self, to_static=0):
+        if self.config.use_lazy_init:
+            with LazyGuard():
+                model = LlamaForCausalLMAuto(self.config)
+            for param in model.parameters():
+                assert not param._is_initialized()
+                param.initialize()
+        else:
+            model = LlamaForCausalLMAuto(self.config)
+        criterion = LlamaPretrainingCriterionAuto(self.config)
+
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+        optimizer = create_optimizer(model, lr_scheduler)
+        if self.amp and not to_static:
+            model, optimizer = paddle.amp.decorate(
+                models=model,
+                optimizers=optimizer,
+                level=self.amp_level,
+                dtype=self.amp_dtype,
+                master_grad=self.amp_master_grad,
+            )
+        optimizer = dist.shard_optimizer(optimizer)
+
+        train_dataset = RandomDataset(self.config.seq_length)
+        train_sampler = BatchSampler(
+            train_dataset,
+            batch_size=2,
+            shuffle=True,
+            drop_last=False,
+        )
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_sampler=train_sampler,
+            num_workers=0,
+        )
+
+        if self.pp == 1:
+            meshes = [get_mesh(0)]
+        elif self.pp > 1:
+            meshes = [get_mesh(0), get_mesh(-1)]
+        else:
+            raise ValueError("pp should be greater or equal to 1")
+
+        dist_loader = dist.shard_dataloader(
+            dataloader=train_dataloader,
+            meshes=meshes,
+            shard_dims="dp",
+        )
+
+        global_step = 1
+        tr_loss = float(0)
+
+        if not to_static:
+            model.train()
+            scaler = None
+            if self.amp and self.amp_dtype == "float16":
+                scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
+                scaler = dist.shard_scaler(scaler)
+
+            for epoch_idx in range(1):
+                for step, inputs in enumerate(dist_loader()):
+                    input_ids, labels = inputs
+                    return input_ids._local_value()._md5sum()
+                    break
+        else:
+            strategy = dist.Strategy()
+            if self.gradient_accumulation_steps > 1:
+                strategy.pipeline.accumulate_steps = (
+                    self.gradient_accumulation_steps
+                )
+
+            if self.amp:
+                amp = strategy.amp
+                amp.enable = self.amp
+                amp.dtype = self.amp_dtype
+                amp.level = self.amp_level.lower()
+                if self.amp_master_grad:
+                    amp.use_master_grad = True
+
+            dist_model = dist.to_static(
+                model,
+                dist_loader,
+                criterion,
+                optimizer,
+                strategy=strategy,
+            )
+
+            dist_model.train()
+            for step, inputs in enumerate(dist_loader()):
+                input_ids, labels = inputs
+                return input_ids._local_value()._md5sum()
+                break
+
+    def run_test_cases(self):
+        dynamic_input_md5sum = self.run_llama(to_static=0)
+        static_input_md5sum = self.run_llama(to_static=1)
+        if dist.get_rank() == 0:
+            assert dynamic_input_md5sum == static_input_md5sum
+
+
+if __name__ == '__main__':
+    TestLlamaAuto().run_test_cases()

test/auto_parallel/hybrid_strategy/semi_auto_llama_pp_gradmerge.py

@@ -0,0 +1,315 @@
+# 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
+from functools import reduce
+
+import numpy as np
+from semi_auto_parallel_llama_model import (
+    LlamaForCausalLMAuto,
+    LlamaPretrainingCriterionAuto,
+    get_mesh,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle import LazyGuard
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+
+class Config:
+    vocab_size = 32000
+    hidden_size = 4096
+    intermediate_size = 11008
+    max_position_embeddings = 2048
+    seq_length = 2048
+    num_hidden_layers = 4
+    num_attention_heads = 32
+    num_key_value_heads = 32
+    initializer_range = 0.02
+    rms_norm_eps = 1e-6
+    use_cache = True
+    use_flash_attention = False
+    sequence_parallel = False
+    rope = True
+    recompute = False
+    recompute_granularity = None
+    use_lazy_init = False
+    virtual_pp_degree = 1
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, num_samples=100):
+        super().__init__()
+        self.seq_len = seq_len
+        self.num_samples = num_samples
+
+    def __getitem__(self, index):
+        input = np.random.uniform(size=[self.seq_len]).astype("int64")
+        label = (np.random.uniform(size=[self.seq_len]) * 10).astype("int64")
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+def create_optimizer(model, lr_scheduler):
+    decay_parameters = [
+        p.name
+        for n, p in model.named_parameters()
+        if not any(nd in n for nd in ["bias", "norm"])
+    ]
+
+    def apply_decay_param_fun(x):
+        return x in decay_parameters
+
+    # test global_clip in auto_parallel
+    if os.getenv("use_param_group") == "true":
+        param_group = {}
+        param_group["params"] = list(model.parameters())
+        param_group["weight_decay"] = 0.01
+        param_group["grad_clip"] = paddle.nn.ClipGradByGlobalNorm(1.0)
+        optimizer = paddle.optimizer.adamw.AdamW(
+            learning_rate=lr_scheduler,
+            apply_decay_param_fun=apply_decay_param_fun,
+            parameters=[param_group],
+        )
+    else:
+        optimizer = paddle.optimizer.adamw.AdamW(
+            learning_rate=lr_scheduler,
+            apply_decay_param_fun=apply_decay_param_fun,
+            parameters=model.parameters(),
+            weight_decay=0.01,
+            grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+        )
+    return optimizer
+
+
+class TestLlamaAuto:
+    def __init__(self):
+        self.config = Config()
+        self.dp = int(os.getenv("dp"))
+        self.mp = int(os.getenv("mp"))
+        self.pp = int(os.getenv("pp"))
+        self.sep = int(os.getenv("sep", "1"))
+        if os.getenv("virtual_pp_degree"):
+            self.config.virtual_pp_degree = int(os.getenv("virtual_pp_degree"))
+        if os.getenv("use_sp") == "true":
+            self.config.sequence_parallel = True
+        if os.getenv("recompute") == "true":
+            self.config.recompute = True
+        self.config.recompute_granularity = os.getenv("recompute_granularity")
+        self.gradient_accumulation_steps = int(os.getenv("acc_step"))
+        self.only_static = os.getenv("only_static")
+
+        if self.config.virtual_pp_degree == 1:
+            self.schedule_mode = "1F1B"
+        elif self.config.virtual_pp_degree > 1:
+            self.schedule_mode = "VPP"
+
+        self.config.tensor_parallel_degree = self.mp
+        self.config.pipeline_parallel_degree = self.pp
+        self.config.context_parallel_degree = 1
+        self.config.sep_parallel_degree = 1
+        if os.getenv("context_parallel", "false") == "true":
+            self.config.context_parallel_degree = self.sep
+            self.config.use_flash_attention = True
+            dist.init_parallel_env()
+        if os.getenv("sep_parallel", "false") == "true":
+            self.config.sep_parallel_degree = self.sep
+        if self.sep > 1:
+            # only one of the context_parallel and sep_parallel can be True
+            assert (
+                self.config.sep_parallel_degree
+                != self.config.context_parallel_degree
+            ), (
+                f"only one of the context_parallel and sep_parallel can be True, but get context_parallel_degree = {self.config.context_parallel_degree} and sep_parallel_degree = {self.config.sep_parallel_degree}, please check your env"
+            )
+
+        self.init_dist_env()
+
+    def init_dist_env(self):
+        mesh_dims = [("pp", self.pp), ("dp", self.dp), ("mp", self.mp)]
+        if not mesh_dims:
+            mesh_dims = [("dp", 1)]
+        dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
+        mesh_shape = [mesh_dim[1] for mesh_dim in mesh_dims]
+        mesh_arr = np.arange(
+            0, reduce(lambda x, y: x * y, mesh_shape, 1)
+        ).reshape(mesh_shape)
+        global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
+        dist.auto_parallel.set_mesh(global_mesh)
+
+    def run_llama(self, to_static=0):
+        if self.only_static and to_static == 0:
+            return
+
+        if self.config.use_lazy_init:
+            with LazyGuard():
+                model = LlamaForCausalLMAuto(self.config)
+            for param in model.parameters():
+                assert not param._is_initialized()
+                param.initialize()
+        else:
+            model = LlamaForCausalLMAuto(self.config)
+        model = LlamaForCausalLMAuto(self.config)
+        criterion = LlamaPretrainingCriterionAuto(self.config)
+
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+        optimizer = create_optimizer(model, lr_scheduler)
+        optimizer = dist.shard_optimizer(optimizer)
+
+        micro_bsz = 2
+        global_bsz = micro_bsz * self.dp * self.gradient_accumulation_steps
+        run_step = 5
+        total_sample_num = run_step * global_bsz
+        global_step = 1
+        tr_loss = float(0)
+
+        if not to_static:
+            train_dataset = RandomDataset(
+                self.config.seq_length, total_sample_num
+            )
+            train_sampler = BatchSampler(
+                train_dataset,
+                batch_size=micro_bsz,
+                shuffle=True,
+                drop_last=True,
+            )
+            train_dataloader = DataLoader(
+                train_dataset,
+                batch_sampler=train_sampler,
+                num_workers=0,
+            )
+
+            model.train()
+            for epoch_idx in range(1):
+                for step, inputs in enumerate(train_dataloader):
+                    input_ids, labels = inputs
+                    logits = model(input_ids)
+                    tr_loss_step = criterion(logits, labels)
+
+                    if self.gradient_accumulation_steps > 1:
+                        tr_loss_step /= self.gradient_accumulation_steps
+
+                    tr_loss_step.backward()
+                    tr_loss += tr_loss_step
+
+                    if global_step % self.gradient_accumulation_steps == 0:
+                        print(
+                            f"step: {global_step // self.gradient_accumulation_steps}  loss: {tr_loss.numpy()}"
+                        )
+                        optimizer.step()
+                        optimizer.clear_grad()
+                        lr_scheduler.step()
+                        tr_loss = 0
+
+                    global_step += 1
+                    if global_step // self.gradient_accumulation_steps >= 10:
+                        break
+        else:
+            strategy = dist.Strategy()
+            if self.pp > 1 and self.gradient_accumulation_steps > 1:
+                strategy.pipeline.enable = True
+                strategy.pipeline.accumulate_steps = (
+                    self.gradient_accumulation_steps
+                )
+                strategy.pipeline.pp_degree = self.pp
+                strategy.pipeline.micro_batch_size = micro_bsz
+                strategy.pipeline.schedule_mode = self.schedule_mode
+                strategy.pipeline.vpp_degree = self.config.virtual_pp_degree
+                strategy.pipeline.vpp_seg_method = "LlamaDecoderLayerAuto"
+            elif self.gradient_accumulation_steps > 1:
+                strategy.gradient_merge.enable = True
+                strategy.gradient_merge.k_steps = (
+                    self.gradient_accumulation_steps
+                )
+                strategy.gradient_merge.avg = True
+
+            train_dataset = RandomDataset(
+                self.config.seq_length, total_sample_num
+            )
+            train_sampler = BatchSampler(
+                train_dataset,
+                batch_size=global_bsz,
+                shuffle=True,
+                drop_last=True,
+            )
+            train_dataloader = DataLoader(
+                train_dataset,
+                batch_sampler=train_sampler,
+                num_workers=0,
+            )
+
+            dist_loader = dist.shard_dataloader(
+                train_dataloader,
+                meshes=[get_mesh(0), get_mesh(-1)],
+                shard_dims="dp",
+            )
+            dist_model = dist.to_static(
+                model, dist_loader, criterion, optimizer, strategy=strategy
+            )
+
+            def validate_batch(batch):
+                if self.gradient_accumulation_steps == 1 or self.pp > 1:
+                    batches = [batch]
+                else:
+                    split_batches = [
+                        np.split(
+                            np.array(b), self.gradient_accumulation_steps, 0
+                        )
+                        for b in batch
+                    ]
+                    batches = []
+                    for i in range(len(split_batches[0])):
+                        micro_batch = [
+                            split_batch[i] for split_batch in split_batches
+                        ]
+                        batches.append(micro_batch)
+                return batches
+
+            dist_model.train()
+            for epoch_idx in range(1):
+                for step, inputs in enumerate(dist_loader()):
+                    batches = validate_batch(inputs)
+                    for micro_batch in batches:
+                        input_ids, labels = micro_batch
+                        tr_loss_step = dist_model(input_ids, labels)
+
+                        if (
+                            tr_loss_step is not None
+                            and self.gradient_accumulation_steps > 1
+                        ):
+                            tr_loss_step = np.sum(tr_loss_step)
+                            tr_loss_step /= self.gradient_accumulation_steps
+
+                        if tr_loss_step:
+                            tr_loss += tr_loss_step
+
+                    print(f"step: {step} loss: {np.array(tr_loss)}")
+                    lr_scheduler.step()
+                    tr_loss = float(0)
+
+                    if step >= run_step:
+                        break
+
+    def run_test_cases(self):
+        self.run_llama(to_static=0)
+        self.run_llama(to_static=1)
+
+
+if __name__ == '__main__':
+    TestLlamaAuto().run_test_cases()

test/auto_parallel/hybrid_strategy/semi_auto_llama_save_load.py

@@ -0,0 +1,325 @@
+# 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 hashlib
+import os
+import random
+import tempfile
+import time
+from functools import reduce
+
+import numpy as np
+from semi_auto_parallel_llama_model import (
+    LlamaForCausalLMAuto,
+    LlamaPretrainingCriterionAuto,
+    get_mesh,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+
+class Config:
+    vocab_size = 320
+    hidden_size = 8
+    intermediate_size = 64
+    max_position_embeddings = 8
+    seq_length = 8
+
+    num_hidden_layers = 2
+    num_attention_heads = 2
+    num_key_value_heads = 2
+    initializer_range = 0.02
+    rms_norm_eps = 1e-6
+    use_cache = True
+    use_flash_attention = False
+    sequence_parallel = False
+    rope = True
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, num_samples=100):
+        super().__init__()
+        self.seq_len = seq_len
+        self.num_samples = num_samples
+
+    def __getitem__(self, index):
+        input = np.full([self.seq_len], index, dtype="int64")
+        label = np.array([index] * 8)
+
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+def create_optimizer(model, lr_scheduler):
+    decay_parameters = [
+        p.name
+        for n, p in model.named_parameters()
+        if not any(nd in n for nd in ["bias", "norm"])
+    ]
+
+    def apply_decay_param_fun(x):
+        return x in decay_parameters
+
+    optimizer = paddle.optimizer.adamw.AdamW(
+        learning_rate=lr_scheduler,
+        apply_decay_param_fun=apply_decay_param_fun,
+        parameters=model.parameters(),
+        weight_decay=0.01,
+        grad_clip=paddle.nn.ClipGradByGlobalNorm(1.0),
+    )
+    return optimizer
+
+
+class TestLlamaAuto:
+    def __init__(self):
+        self.config = Config()
+        self.dp = int(os.getenv("dp"))
+        self.mp = int(os.getenv("mp"))
+        self.pp = int(os.getenv("pp"))
+        self.sep = int(os.getenv("sep", "1"))
+        if os.getenv("use_sp") == "true":
+            self.config.sequence_parallel = True
+        self.gradient_accumulation_steps = int(os.getenv("acc_step"))
+        self.config.recompute = False
+        self.config.context_parallel_degree = 1
+        self.config.sep_parallel_degree = 1
+        self.config.tensor_parallel_degree = self.mp
+        self.config.pipeline_parallel_degree = self.pp
+        if os.getenv("context_parallel", "false") == "true":
+            self.config.context_parallel_degree = self.sep
+            self.config.use_flash_attention = True
+            dist.init_parallel_env()
+        if os.getenv("sep_parallel", "false") == "true":
+            self.config.sep_parallel_degree = self.sep
+        if self.sep > 1:
+            # only one of the context_parallel and sep_parallel can be True
+            assert (
+                self.config.sep_parallel_degree
+                != self.config.context_parallel_degree
+            ), (
+                f"only one of the context_parallel and sep_parallel can be True, but get context_parallel_degree = {self.config.context_parallel_degree} and sep_parallel_degree = {self.config.sep_parallel_degree}, please check your env"
+            )
+
+        self.init_dist_env()
+
+    def init_dist_env(self):
+        order = ["dp", "pp", "mp"]
+        dp_degree = self.dp
+        mp_degree = self.mp
+        pp_degree = self.pp
+        degree = [dp_degree, pp_degree, mp_degree]
+        mesh_dims = list(filter(lambda x: x[1] > 1, list(zip(order, degree))))
+        if not mesh_dims:
+            mesh_dims = [("dp", 1)]
+        dim_names = [mesh_dim[0] for mesh_dim in mesh_dims]
+        mesh_shape = [mesh_dim[1] for mesh_dim in mesh_dims]
+        mesh_arr = np.arange(
+            0, reduce(lambda x, y: x * y, mesh_shape, 1)
+        ).reshape(mesh_shape)
+        global_mesh = dist.ProcessMesh(mesh_arr, dim_names)
+        dist.auto_parallel.set_mesh(global_mesh)
+        paddle.seed(1024)
+        np.random.seed(1024)
+        random.seed(1024)
+
+    def check_program_equal(self, program_a, program_b):
+        assert program_a.num_ops() == program_b.num_ops(), (
+            f'The number of ops between two programs is different: {program_a.num_ops()} vs {program_b.num_ops()}.'
+        )
+        for i in range(program_a.num_ops()):
+            a_op = program_a.global_block().ops[i]
+            b_op = program_b.global_block().ops[i]
+            # check op name
+            assert a_op.name() == b_op.name(), (
+                f'The name of {i} op in program is different: {a_op.name()} vs {b_op.name()}.'
+            )
+            # check op inputs
+            for index in range(a_op.num_operands()):
+                assert str(a_op.operand(index).source()) == str(
+                    b_op.operand(index).source()
+                ), (
+                    f'The type of {index} operand is different: {a_op.operand(index).source()} vs {b_op.operand(index).source()}'
+                )
+            # check op outputs
+            for index in range(a_op.num_results()):
+                assert str(a_op.result(index)) == str(b_op.result(index)), (
+                    f'The type of {index} result is different: {a_op.result(index)} vs {b_op.result(index)}'
+                )
+            # check op attrs
+            for k, v in a_op.attrs().items():
+                if k in ["op_callstack"]:
+                    continue
+                assert k in b_op.attrs(), (
+                    f'Can not find key of {k} attribute in other program'
+                )
+                if k == 'place':
+                    assert type(v) == type(b_op.attrs()[k]), (
+                        f'The attribute of {k} is different: {type(v)} vs {type(b_op.attrs()[k])}'
+                    )
+                else:
+                    assert v == b_op.attrs()[k], (
+                        f'The attribute of {k} is different: {v} vs {b_op.attrs()[k]}'
+                    )
+
+    def run_dy2static(self, tmp_ckpt_path):
+        model = LlamaForCausalLMAuto(self.config)
+        criterion = LlamaPretrainingCriterionAuto(self.config)
+
+        lr_scheduler = paddle.optimizer.lr.LinearWarmup(
+            learning_rate=0.0001, warmup_steps=2, start_lr=0, end_lr=0.0001
+        )
+        optimizer = create_optimizer(model, lr_scheduler)
+
+        train_dataset = RandomDataset(self.config.seq_length)
+        train_sampler = BatchSampler(
+            train_dataset,
+            batch_size=2,
+            shuffle=False,
+            drop_last=True,
+        )
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_sampler=train_sampler,
+            num_workers=0,
+        )
+        dist_loader = dist.shard_dataloader(
+            dataloader=train_dataloader,
+            meshes=[get_mesh(0), get_mesh(1)],
+            shard_dims="dp",
+        )
+
+        strategy = None
+        if self.gradient_accumulation_steps > 1:
+            strategy = dist.Strategy()
+            strategy.pipeline.accumulate_steps = (
+                self.gradient_accumulation_steps
+            )
+
+        dist_model = dist.to_static(
+            model, dist_loader, criterion, optimizer, strategy=strategy
+        )
+
+        dist_model.train()
+
+        state_dict = dist_model.state_dict()
+
+        loss_before_save = []
+        for step, inputs in enumerate(dist_loader()):
+            input_ids, labels = inputs
+            loss = dist_model(input_ids, labels)
+            lr_scheduler.step()
+            if step == 2:
+                state_dict = dist_model.state_dict()
+                dist.save_state_dict(state_dict, tmp_ckpt_path, async_save=True)
+            if step > 2:
+                numpy_array = np.array(loss)
+                array_bytes = numpy_array.tobytes()
+                loss_md5 = hashlib.md5(array_bytes).hexdigest()
+                loss_before_save.append(loss_md5)
+
+                if int(dist.get_rank()) in [2, 3, 6, 7]:
+                    assert loss is not None
+                else:
+                    assert loss is None
+
+            if step >= 9:
+                break
+
+        # check pir dist_model save&load
+        paddle.enable_static()
+        model_file_path = os.path.join(
+            tmp_ckpt_path,
+            "rank_" + str(paddle.distributed.get_rank()) + ".pd_dist_model",
+        )
+        paddle.save(
+            dist_model._engine._pir_dist_main_progs["train"], model_file_path
+        )
+        loaded_model = paddle.load(model_file_path)
+        self.check_program_equal(
+            dist_model._engine._pir_dist_main_progs["train"], loaded_model
+        )
+        paddle.disable_static()
+        paddle.distributed.barrier()
+
+        time.sleep(10)
+
+        loss_after_load = []
+        for step, inputs in enumerate(dist_loader()):
+            if step < 2:
+                continue
+            input_ids, labels = inputs
+            loss = dist_model(input_ids, labels)
+            lr_scheduler.step()
+            if step == 2:
+                state_dict = dist_model.state_dict()
+                dist.load_state_dict(state_dict, tmp_ckpt_path)
+            if step > 2:
+                numpy_array = np.array(loss)
+                array_bytes = numpy_array.tobytes()
+                loss_md5 = hashlib.md5(array_bytes).hexdigest()
+                loss_after_load.append(loss_md5)
+
+                if int(dist.get_rank()) in [2, 3, 6, 7]:
+                    assert loss is not None
+                else:
+                    assert loss is None
+            if step >= 9:
+                break
+
+        return (loss_before_save, loss_after_load)
+
+    def broadcast_ckpt_path(self, ckpt_path):
+        dist.init_parallel_env()
+        rank = dist.get_rank()
+        if rank == 0:
+            byte_array = np.frombuffer(
+                ckpt_path.encode('utf-8'), dtype=np.uint8
+            )
+            length = np.array([len(byte_array)], dtype=np.int32)
+        else:
+            length = np.array([0], dtype=np.int32)
+
+        length_tensor = paddle.to_tensor(length)
+        dist.broadcast(length_tensor, src=0)
+        length = length_tensor.numpy()[0]
+
+        if rank != 0:
+            byte_array = np.empty(length, dtype=np.uint8)
+
+        byte_array_tensor = paddle.to_tensor(byte_array)
+
+        dist.broadcast(byte_array_tensor, src=0)
+
+        global_ckpt_path = byte_array_tensor.numpy().tobytes().decode('utf-8')
+
+        return global_ckpt_path
+
+    def run_test_cases(self):
+        self.init_dist_env()
+        ckpt_path = tempfile.TemporaryDirectory()
+        tmp_ckpt_path = self.broadcast_ckpt_path(ckpt_path.name)
+        loss = self.run_dy2static(tmp_ckpt_path)
+        if int(dist.get_rank()) in [2, 3, 6, 7]:
+            assert len(loss[0]) == len(loss[1])
+            for i in range(len(loss[0])):
+                assert loss[0][i] == loss[1][i]
+        ckpt_path.cleanup()
+
+
+if __name__ == '__main__':
+    TestLlamaAuto().run_test_cases()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_2d_global_mesh_reshard.py

@@ -0,0 +1,100 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestSemiAutoParallel2DGlobalMeshReshard:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._global_mesh = dist.ProcessMesh(
+            [[0, 1], [2, 3]], dim_names=["pp", "dp"]
+        )
+        self._mesh0 = dist.ProcessMesh([0, 1], dim_names=["dp"])
+        self._mesh1 = dist.ProcessMesh([2, 3], dim_names=["dp"])
+        paddle.set_device(self._backend)
+
+    def test_basic(self):
+        input = paddle.ones(shape=[2, 3], dtype='float32')
+        input = dist.shard_tensor(
+            input, self._global_mesh, [dist.Replicate(), dist.Shard(0)]
+        )
+        input.stop_gradient = False
+        global_input = input + 1.0  # global_input: 2.0
+
+        # forward on pp0
+        input_pp0 = dist.reshard(global_input, self._mesh0, [dist.Shard(0)])
+        output = input_pp0 + 1.0  # output_pp0: 3.0
+
+        # forward on pp1
+        output = dist.reshard(output, self._mesh1, [dist.Shard(0)])
+        input_pp1 = dist.reshard(global_input, self._mesh1, [dist.Shard(0)])
+        output = input_pp1 + output  # output_pp1: 5.0
+        loss = paddle.sum(output)  # 30.0
+        np.testing.assert_allclose(loss.numpy(), 30.0, rtol=1e-06, verbose=True)
+        loss.backward()
+        np.testing.assert_allclose(
+            input.grad.numpy(),
+            np.full(shape=(2, 3), fill_value=2.0, dtype=np.float32),
+            rtol=1e-06,
+            verbose=True,
+        )
+
+    def test_split_dim1(self):
+        global_mesh = dist.ProcessMesh([[0, 1], [2, 3]])
+        mesh0 = dist.ProcessMesh([[0], [2]])
+        mesh1 = dist.ProcessMesh([[1], [3]])
+
+        input = paddle.ones(shape=[2, 3], dtype='float32')
+        input = dist.shard_tensor(
+            input, global_mesh, [dist.Shard(0), dist.Replicate()]
+        )
+        input.stop_gradient = False
+        global_input = input + 1.0  # global_input: 2.0
+
+        # forward on pp0
+        input_pp0 = dist.reshard(
+            global_input, mesh0, [dist.Shard(0), dist.Replicate()]
+        )
+        output = input_pp0 + 1.0  # output_pp0: 3.0
+
+        # forward on pp1
+        output = dist.reshard(output, mesh1, [dist.Shard(0), dist.Replicate()])
+        input_pp1 = dist.reshard(
+            global_input, mesh1, [dist.Shard(0), dist.Replicate()]
+        )
+        output = input_pp1 + output  # output_pp1: 5.0
+        loss = paddle.sum(output)  # 30.0
+        np.testing.assert_allclose(loss.numpy(), 30.0, rtol=1e-06, verbose=True)
+        loss.backward()
+        np.testing.assert_allclose(
+            input.grad.numpy(),
+            np.full(shape=(2, 3), fill_value=2.0, dtype=np.float32),
+            rtol=1e-06,
+            verbose=True,
+        )
+
+    def run_test_case(self):
+        self.test_basic()
+        self.test_split_dim1()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallel2DGlobalMeshReshard().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_3d_global_mesh_reshard.py

@@ -0,0 +1,82 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestSemiAutoParallel3DGlobalMeshReshard:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._global_mesh = dist.ProcessMesh(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=['pp', 'dp', 'mp']
+        )
+        self._mesh0 = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['dp', 'mp'])
+        self._mesh1 = dist.ProcessMesh([[4, 5], [6, 7]], dim_names=['dp', 'mp'])
+        paddle.set_device(self._backend)
+
+    def test_basic(self):
+        global_input = dist.shard_tensor(
+            paddle.ones(shape=[6, 8], dtype='float32'),
+            self._global_mesh,
+            [dist.Replicate(), dist.Replicate(), dist.Replicate()],
+        )  # 1.0
+        global_input.stop_gradient = False
+        # forward on mesh0
+        input_mesh0 = dist.reshard(
+            global_input, self._mesh0, [dist.Replicate(), dist.Replicate()]
+        )
+        output = input_mesh0 + 1.0  # 2.0
+
+        # forward on mesh1
+        output = dist.reshard(
+            output, self._mesh1, [dist.Replicate(), dist.Replicate()]
+        )
+        input_mesh1 = dist.reshard(
+            global_input, self._mesh1, [dist.Replicate(), dist.Replicate()]
+        )
+        output = output + input_mesh1  # 3.0
+        loss = paddle.sum(output)  # 144.0
+        np.testing.assert_allclose(
+            loss.numpy(), 144.0, rtol=1e-06, verbose=True
+        )
+        loss.backward()
+        np.testing.assert_allclose(
+            global_input.grad.numpy(),
+            np.full(shape=(6, 8), fill_value=2.0, dtype=np.float32),
+            rtol=1e-06,
+            verbose=True,
+        )
+
+    def test_3d_mesh_with_any_status(self):
+        dense_tensor = paddle.ones(shape=[2, 6], dtype='float32')
+        dist_tensor = dist.shard_tensor(
+            dense_tensor,
+            self._global_mesh,
+            [dist.Replicate(), dist.Shard(0), dist.Replicate()],
+        )
+        np.testing.assert_equal(dist_tensor._local_shape, [1, 6])
+
+    def run_test_case(self):
+        self.test_basic()
+        self.test_3d_mesh_with_any_status()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallel3DGlobalMeshReshard().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_c_cross_entropy_dp.py

@@ -0,0 +1,185 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import Shard, fleet
+from paddle.distributed.fleet import auto
+
+
+def set_random_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    paddle.seed(seed)
+
+
+def dygraph_parallel_cross_entropy(data, label):
+    model = fleet.meta_parallel.ParallelCrossEntropy()
+    loss = model(data, label)
+    return paddle.mean(loss)
+
+
+def dygraph_cross_entropy(data, label):
+    model = paddle.nn.CrossEntropyLoss()
+    loss = model(data, label)
+    return loss
+
+
+class MyDataset(paddle.io.Dataset):
+    def __init__(self, data, label):
+        self._data = data
+        self._label = label
+
+    def __getitem__(self, index):
+        return self._data[index], self._label[index]
+
+    def __len__(self):
+        return self._data.shape[0]
+
+
+class MyMLP(nn.Layer):
+    def __init__(self, process_mesh, placements):
+        super().__init__()
+        self.process_mesh = process_mesh
+        self.placements = placements
+
+    def forward(self, x):
+        dist.shard_tensor(
+            x, self.process_mesh, self.placements, stop_gradient=False
+        )
+        return x
+
+
+def auto_parallel_cross_entropy(data, label, process_mesh, placements):
+    with paddle.LazyGuard():
+        model = MyMLP(process_mesh, placements)
+        loss_layer = paddle.nn.CrossEntropyLoss()
+        auto.fetch("input0@GRAD", "logits_grad", logging=False)
+        auto.fetch(
+            "softmax_with_cross_entropy_0.tmp_1",
+            "loss_before_mean",
+            logging=False,
+        )
+    optimizer = paddle.optimizer.Adam(
+        learning_rate=0.001, parameters=model.parameters()
+    )
+    engine = auto.Engine(model, loss_layer, optimizer)
+    train_dataset = MyDataset(data, label)
+    log = engine.fit(train_dataset, epochs=1, batch_size=data.shape[0])
+    logit_grad = np.array(log.history["fetches"][0]["logits_grad"])
+    loss = np.array(log.history["loss"])
+    paddle.disable_static()
+    return loss, logit_grad
+
+
+class TestDpDistTraining:
+    def __init__(self):
+        self.nsample = 40
+        self.nclass = 20
+        self.seed = 100
+
+    def run_test_case(self):
+        strategy = fleet.DistributedStrategy()
+        strategy.hybrid_configs = {
+            "dp_degree": 2,
+            "mp_degree": 1,
+            "pp_degree": 1,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+
+        nsample = self.nsample
+        nclass = self.nclass
+        seed = self.seed
+
+        set_random_seed(seed)
+        rank_id = dist.get_rank()
+
+        paddle.seed(rank_id * 10)
+        random.seed(seed)
+        np.random.seed(seed)
+        check_group = dist.new_group(list(range(2)))
+        process_mesh = dist.ProcessMesh(mesh=[0, 1], dim_names=["x"])
+
+        np_label = np.random.randint(0, nclass, (nsample // 2, 1))
+        label = paddle.to_tensor(np_label, dtype="int64")
+
+        data = paddle.randn(
+            shape=[nsample // 2, nclass],
+            dtype='float32',
+        )
+        data.stop_gradient = False
+
+        integral_data = []
+        partial_data = data.clone().detach()
+        paddle.distributed.all_gather(
+            integral_data, partial_data, group=check_group
+        )
+        integral_data = paddle.concat(integral_data, axis=0)
+        integral_data = integral_data.detach().clone()
+        integral_data.stop_gradient = False
+
+        integral_label = []
+        partial_label = label.clone().detach()
+        paddle.distributed.all_gather(
+            integral_label, partial_label, group=check_group
+        )
+        integral_label = paddle.concat(integral_label, axis=0)
+        integral_label = integral_label.detach().clone()
+        integral_label.stop_gradient = False
+
+        loss_dygraph_parallel = dygraph_cross_entropy(data, label)
+        loss_auto, auto_grad = auto_parallel_cross_entropy(
+            integral_data.numpy(),
+            integral_label.numpy(),
+            process_mesh,
+            [Shard(0)],
+        )
+
+        np.testing.assert_allclose(
+            loss_dygraph_parallel.numpy(), loss_auto, rtol=1e-6
+        )
+
+        loss_dygraph_parallel.backward()
+
+        integral_grad = []
+        partial_grad = data.grad.clone().detach()
+        paddle.distributed.all_gather(
+            integral_grad, partial_grad, group=check_group
+        )
+        integral_grad = paddle.concat(integral_grad, axis=0)
+
+        integral_auto_grad = []
+        paddle.distributed.all_gather(
+            integral_auto_grad,
+            paddle.to_tensor(auto_grad),
+            group=check_group,
+        )
+        integral_auto_grad = paddle.concat(integral_auto_grad, axis=0)
+
+        parallel_grad = integral_grad.numpy()
+        auto_grad = integral_auto_grad.numpy()
+        np.testing.assert_allclose(
+            parallel_grad,
+            auto_grad,
+            rtol=1e-6,
+        )
+
+
+if __name__ == '__main__':
+    TestDpDistTraining().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_c_cross_entropy_hybrid.py

@@ -0,0 +1,170 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import Shard, fleet
+from paddle.distributed.fleet import auto
+
+
+def set_random_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    paddle.seed(seed)
+
+
+def dygraph_parallel_cross_entropy(data, label):
+    model = fleet.meta_parallel.ParallelCrossEntropy()
+    loss = model(data, label)
+    return paddle.mean(loss)
+
+
+def dygraph_cross_entropy(data, label):
+    model = paddle.nn.CrossEntropyLoss()
+    loss = model(data, label)
+    return loss
+
+
+class MyDataset(paddle.io.Dataset):
+    def __init__(self, data, label):
+        self._data = data
+        self._label = label
+
+    def __getitem__(self, index):
+        return self._data[index], self._label[index]
+
+    def __len__(self):
+        return self._data.shape[0]
+
+
+class MyMLP(nn.Layer):
+    def __init__(self, process_mesh, placements):
+        super().__init__()
+        self.process_mesh = process_mesh
+        self.placements = placements
+
+    def forward(self, x):
+        dist.shard_tensor(
+            x, self.process_mesh, self.placements, stop_gradient=False
+        )
+        return x
+
+
+def auto_parallel_cross_entropy(data, label, process_mesh, placements):
+    with paddle.LazyGuard():
+        model = MyMLP(process_mesh, placements)
+        loss_layer = paddle.nn.CrossEntropyLoss()
+        auto.fetch("input0@GRAD", "logits_grad", logging=False)
+        auto.fetch("input0", "logits", logging=False)
+        auto.fetch(
+            "softmax_with_cross_entropy_0.tmp_1",
+            "loss_before_mean",
+            logging=False,
+        )
+    optimizer = paddle.optimizer.Adam(
+        learning_rate=0.001, parameters=model.parameters()
+    )
+    engine = auto.Engine(model, loss_layer, optimizer)
+    train_dataset = MyDataset(data, label)
+    log = engine.fit(train_dataset, epochs=1, batch_size=data.shape[0])
+    logit_grad = np.array(log.history["fetches"][0]["logits_grad"])
+    loss = np.array(log.history["loss"])
+    logits = np.array(log.history["fetches"][0]["logits"])
+    paddle.disable_static()
+    return loss, logit_grad
+
+
+class TestHybridDistTraining:
+    def __init__(self):
+        self.nsample = 2
+        self.seq_len = 2
+        self.nclass = 4
+        self.seed = 100
+
+    def run_test_case(self):
+        strategy = fleet.DistributedStrategy()
+        self.model_parallel_size = 2
+        strategy.hybrid_configs = {
+            "dp_degree": 2,
+            "mp_degree": 2,
+            "pp_degree": 1,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+
+        nsample = self.nsample
+        nclass = self.nclass
+        seq_len = self.seq_len
+        seed = self.seed
+
+        set_random_seed(seed)
+        rank_id = dist.get_rank()
+
+        paddle.seed(rank_id * 10)
+        random.seed(seed)
+        np.random.seed(seed)
+        process_mesh = dist.ProcessMesh(
+            mesh=[[0, 1], [2, 3]], dim_names=["x", "y"]
+        )
+
+        integral_np_label = np.random.randint(0, nclass, (nsample, seq_len, 1))
+        integral_label = paddle.to_tensor(integral_np_label, dtype="int64")
+
+        integral_np_data = np.random.randn(nsample, seq_len, nclass).astype(
+            "float32"
+        )
+        integral_data = paddle.to_tensor(integral_np_data)
+        integral_data.stop_gradient = False
+
+        loss_dygraph = dygraph_cross_entropy(integral_data, integral_label)
+        dp_start_idx = rank_id // 2 * (nsample // 2)
+        dp_end_idx = dp_start_idx + (nsample // 2)
+        mp_start_idx = rank_id % 2 * (nclass // 2)
+        mp_end_idx = mp_start_idx + (nclass // 2)
+        # the dataloader cannot support shard on non-batch dim,
+        # so we should slice the data and label tensor manually
+        mp_sliced_np_data = integral_np_data[:, :, mp_start_idx:mp_end_idx]
+        loss_auto, auto_grad = auto_parallel_cross_entropy(
+            mp_sliced_np_data,
+            integral_np_label,
+            process_mesh,
+            [Shard(0), Shard(2)],
+        )
+        pd_loss_auto = paddle.to_tensor(loss_auto)
+        paddle.distributed.all_reduce(pd_loss_auto)
+        pd_loss_auto = pd_loss_auto / 4
+        np.testing.assert_allclose(
+            loss_dygraph.numpy(), pd_loss_auto.numpy(), rtol=1e-6
+        )
+
+        loss_dygraph.backward()
+
+        sliced_grad = integral_data.grad[
+            dp_start_idx:dp_end_idx, :, mp_start_idx:mp_end_idx
+        ]
+        partial_grad = sliced_grad.clone().detach()
+
+        np.testing.assert_allclose(
+            partial_grad.numpy(),
+            auto_grad,
+            rtol=1e-6,
+        )
+
+
+if __name__ == '__main__':
+    TestHybridDistTraining().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_c_cross_entropy_mp.py

@@ -0,0 +1,174 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import Shard, fleet
+from paddle.distributed.fleet import auto
+
+
+def set_random_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    paddle.seed(seed)
+
+
+def dygraph_parallel_cross_entropy(data, label):
+    model = fleet.meta_parallel.ParallelCrossEntropy()
+    loss = model(data, label)
+    return paddle.mean(loss)
+
+
+def dygraph_cross_entropy(data, label):
+    model = paddle.nn.CrossEntropyLoss()
+    loss = model(data, label)
+    return loss
+
+
+class MyDataset(paddle.io.Dataset):
+    def __init__(self, data, label):
+        self._data = data
+        self._label = label
+
+    def __getitem__(self, index):
+        return self._data[index], self._label[index]
+
+    def __len__(self):
+        return self._data.shape[0]
+
+
+class MyMLP(nn.Layer):
+    def __init__(self, process_mesh, placements):
+        super().__init__()
+        self.process_mesh = process_mesh
+        self.placements = placements
+
+    def forward(self, x):
+        dist.shard_tensor(
+            x, self.process_mesh, self.placements, stop_gradient=False
+        )
+        return x
+
+
+def auto_parallel_cross_entropy(data, label, process_mesh, placements):
+    with paddle.LazyGuard():
+        model = MyMLP(process_mesh, placements)
+        loss_layer = paddle.nn.CrossEntropyLoss()
+        auto.fetch("input0@GRAD", "logits_grad", logging=False)
+        auto.fetch(
+            "softmax_with_cross_entropy_0.tmp_1",
+            "loss_before_mean",
+            logging=False,
+        )
+    optimizer = paddle.optimizer.Adam(
+        learning_rate=0.001, parameters=model.parameters()
+    )
+    engine = auto.Engine(model, loss_layer, optimizer)
+    train_dataset = MyDataset(data, label)
+    log = engine.fit(train_dataset, epochs=1, batch_size=data.shape[0])
+    logit_grad = np.array(log.history["fetches"][0]["logits_grad"])
+    loss = np.array(log.history["loss"])
+    paddle.disable_static()
+    return loss, logit_grad
+
+
+class TestMpDistTraining:
+    def __init__(self):
+        self.nsample = 40
+        self.nclass = 20
+        self.seed = 100
+
+    def run_test_case(self):
+        strategy = fleet.DistributedStrategy()
+        self.model_parallel_size = 2
+        strategy.hybrid_configs = {
+            "dp_degree": 1,
+            "mp_degree": self.model_parallel_size,
+            "pp_degree": 1,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+
+        nsample = self.nsample
+        nclass = self.nclass
+        seed = self.seed
+
+        set_random_seed(seed)
+        rank_id = dist.get_rank()
+
+        paddle.seed(rank_id * 10)
+        random.seed(seed)
+        np.random.seed(seed)
+        check_group = dist.new_group(list(range(self.model_parallel_size)))
+        process_mesh = dist.ProcessMesh(mesh=[0, 1], dim_names=["x"])
+
+        np_label = np.random.randint(0, nclass, (nsample, 1))
+        label = paddle.to_tensor(np_label, dtype="int64")
+
+        data = paddle.randn(
+            shape=[nsample, nclass // self.model_parallel_size],
+            dtype='float32',
+        )
+        data.stop_gradient = False
+
+        integral_data = []
+        partial_data = data.clone().detach()
+        paddle.distributed.all_gather(
+            integral_data, partial_data, group=check_group
+        )
+        integral_data = paddle.concat(integral_data, axis=-1)
+        integral_data = integral_data.detach().clone()
+        integral_data.stop_gradient = False
+
+        loss_dygraph_parallel = dygraph_parallel_cross_entropy(data, label)
+        loss_auto, auto_grad = auto_parallel_cross_entropy(
+            data.numpy(), np_label, process_mesh, [Shard(1)]
+        )
+
+        np.testing.assert_allclose(
+            loss_dygraph_parallel.numpy(), loss_auto, rtol=1e-6
+        )
+
+        loss_dygraph_parallel.backward()
+
+        integral_grad = []
+        partial_grad = data.grad.clone().detach()
+        paddle.distributed.all_gather(
+            integral_grad, partial_grad, group=check_group
+        )
+        integral_grad = paddle.concat(integral_grad, axis=-1)
+
+        integral_auto_grad = []
+        paddle.distributed.all_gather(
+            integral_auto_grad,
+            paddle.to_tensor(auto_grad),
+            group=check_group,
+        )
+        integral_auto_grad = paddle.concat(integral_auto_grad, axis=-1)
+
+        parallel_grad = integral_grad.numpy()
+        auto_grad = integral_auto_grad.numpy()
+        np.testing.assert_allclose(
+            parallel_grad,
+            auto_grad,
+            rtol=1e-6,
+        )
+
+
+if __name__ == '__main__':
+    TestMpDistTraining().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_cross_mesh_reshard.py

@@ -0,0 +1,183 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestSemiAutoParallelCrossMeshReshard:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh0 = dist.ProcessMesh([0, 1], dim_names=["x"])
+        self._mesh1 = dist.ProcessMesh([2, 3], dim_names=["x"])
+        self._shape = (20, 20)
+        self._shard_axis = 0
+        self._out_shard_axis = 1
+        paddle.set_device(self._backend)
+
+    def test_p_to_r(self):
+        a = paddle.ones(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Partial(dist.ReduceType.kRedSum)]
+        )
+        out = dist.reshard(input_tensor, self._mesh1, [dist.Replicate()])
+
+        if dist.get_rank() in [2, 3]:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+
+    def test_p_to_s(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(
+            a, axis=self._shard_axis, num_or_sections=self._mesh1.shape[0]
+        )
+        expect_out_shape = list(self._shape)
+        expect_out_shape[self._shard_axis] = (
+            self._shape[self._shard_axis] // self._mesh1.shape[0]
+        )
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Partial(dist.ReduceType.kRedSum)]
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Shard(self._shard_axis)]
+        )
+        if dist.get_rank() in [2, 3]:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                (
+                    expect_out[0].numpy()
+                    if dist.get_rank() == 2
+                    else expect_out[1].numpy()
+                ),
+            )
+
+    def test_r_to_p(self):
+        a = paddle.ones(self._shape)
+        b = paddle.zeros(self._shape)
+
+        input_tensor = dist.shard_tensor(a, self._mesh0, [dist.Replicate()])
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Partial(dist.ReduceType.kRedSum)]
+        )
+        if dist.get_rank() == 2:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+        if dist.get_rank() == 3:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), b.numpy())
+
+    def test_r_to_s(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(
+            a, axis=self._shard_axis, num_or_sections=self._mesh1.shape[0]
+        )
+        expect_out_shape = list(self._shape)
+        expect_out_shape[self._shard_axis] = (
+            self._shape[self._shard_axis] // self._mesh1.shape[0]
+        )
+
+        input_tensor = dist.shard_tensor(a, self._mesh0, [dist.Replicate()])
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Shard(self._shard_axis)]
+        )
+        if dist.get_rank() in [2, 3]:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                (
+                    expect_out[0].numpy()
+                    if dist.get_rank() == 2
+                    else expect_out[1].numpy()
+                ),
+            )
+
+    def test_s_to_p(self):
+        a = paddle.ones(self._shape)
+        b = paddle.zeros(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(self._shard_axis)]
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Partial(dist.ReduceType.kRedSum)]
+        )
+        if dist.get_rank() == 2:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+        if dist.get_rank() == 3:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), b.numpy())
+
+    def test_s_to_r(self):
+        a = paddle.ones(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(self._shard_axis)]
+        )
+        out = dist.reshard(input_tensor, self._mesh1, [dist.Replicate()])
+        if dist.get_rank() in [2, 3]:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+
+    def test_s_to_s(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(
+            a, axis=self._out_shard_axis, num_or_sections=self._mesh1.shape[0]
+        )
+        expect_out_shape = list(self._shape)
+        expect_out_shape[self._out_shard_axis] = (
+            self._shape[self._out_shard_axis] // self._mesh1.shape[0]
+        )
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(self._shard_axis)]
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Shard(self._out_shard_axis)]
+        )
+
+        if dist.get_rank() in [2, 3]:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                (
+                    expect_out[0].numpy()
+                    if dist.get_rank() == 2
+                    else expect_out[1].numpy()
+                ),
+            )
+
+    def run_test_case(self):
+        self.test_p_to_r()
+        self.test_p_to_s()
+        self.test_r_to_p()
+        self.test_r_to_s()
+        self.test_s_to_p()
+        self.test_s_to_r()
+        self.test_s_to_s()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelCrossMeshReshard().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_for_fused_linear_param_grad_add.py

@@ -0,0 +1,106 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import _C_ops
+
+
+class TestFusedParamGradAddForSemiAutoParallel:
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+
+    def check_tensor_eq(self, a, b):
+        np1 = a.numpy()
+        np2 = b.numpy()
+        np.testing.assert_allclose(np1, np2, rtol=1e-05, verbose=True)
+
+    def test_body(self):
+        x_shape = [4, 16, 32]
+        y_shape = [4, 16, 64]
+
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+
+        x_np = np.random.random(size=x_shape).astype(self._dtype)
+        y_np = np.random.random(size=y_shape).astype(self._dtype)
+
+        def run_acc_step(x, y):
+            weight_grad = None
+            bias_grad = None
+            for _ in range(2):
+                weight_grad, bias_grad = _C_ops.fused_linear_param_grad_add(
+                    x,
+                    y,
+                    weight_grad,
+                    bias_grad,
+                    False,
+                    True,
+                )
+            return weight_grad, bias_grad
+
+        x = paddle.to_tensor(x_np)
+        y = paddle.to_tensor(y_np)
+        x.stop_gradient = True
+        y.stop_gradient = True
+
+        weight_grad, bias_grad = run_acc_step(x, y)
+
+        # test mp col split
+        x_placements = [dist.Shard(0), dist.Replicate()]
+        y_placements = [dist.Shard(0), dist.Shard(2)]
+
+        dist_x = dist.shard_tensor(x_np, self._mesh, x_placements)
+        dist_y = dist.shard_tensor(y_np, self._mesh, y_placements)
+        dist_x.stop_gradient = True
+        dist_y.stop_gradient = True
+
+        weight_grad_dist, bias_grad_dist = run_acc_step(dist_x, dist_y)
+        self.check_tensor_eq(weight_grad, weight_grad_dist)
+        self.check_tensor_eq(bias_grad, bias_grad_dist)
+
+        # test mp row split
+        x_placements = [dist.Shard(0), dist.Shard(2)]
+        y_placements = [dist.Shard(0), dist.Replicate()]
+        dist_x = dist.shard_tensor(x_np, self._mesh, x_placements)
+        dist_y = dist.shard_tensor(y_np, self._mesh, y_placements)
+        dist_x.stop_gradient = True
+        dist_y.stop_gradient = True
+        weight_grad_dist, bias_grad_dist = run_acc_step(dist_x, dist_y)
+        self.check_tensor_eq(weight_grad, weight_grad_dist)
+        self.check_tensor_eq(bias_grad, bias_grad_dist)
+
+    def test_fused_linear_param_grad_add(self):
+        self.test_body()
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+        elif self._backend == "gpu":
+            paddle.set_device("gpu:" + str(dist.get_rank()))
+        else:
+            raise ValueError("Only support cpu or gpu backend.")
+
+        self.test_fused_linear_param_grad_add()
+
+
+if __name__ == '__main__':
+    TestFusedParamGradAddForSemiAutoParallel().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_for_llama_decoder_dp_mp.py

@@ -0,0 +1,312 @@
+# 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 random
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.distributed import Replicate, Shard
+from paddle.nn.functional.flash_attention import flash_attention
+
+BATCH_NUM = 4
+BATCH_SIZE = 16
+HIDDEN_SIZE = 1024
+INTERMEDIATE_SIZE = 1024 // 3 * 8
+SEQ_LEN = 128
+N_HEAD = 8
+
+
+def create_numpy_like_random(name):
+    return paddle.ParamAttr(
+        name=name, initializer=paddle.nn.initializer.Uniform(-0.1, 0.1)
+    )
+
+
+class LlamaAttention(nn.Layer):
+    def __init__(self, param_prefix="", hidden_size=HIDDEN_SIZE, n_head=N_HEAD):
+        super().__init__()
+        weight_attr_0 = create_numpy_like_random(param_prefix + "_0")
+        weight_attr_1 = create_numpy_like_random(param_prefix + "_1")
+        self.hidden_size = hidden_size
+        self.num_heads = n_head
+        self.head_dim = hidden_size // n_head
+        self.qkv_proj = nn.Linear(hidden_size, hidden_size * 3, weight_attr_0)
+        self.o_proj = nn.Linear(hidden_size, hidden_size, weight_attr_1)
+
+    def forward(self, x):
+        mix_layer = self.qkv_proj(x)
+        target_shape = [0, 0, self.num_heads, 3 * self.head_dim]
+        mix_layer = paddle.reshape(mix_layer, target_shape)
+        mix_layer = paddle.cast(mix_layer, paddle.bfloat16)
+        query_states, key_states, value_states = paddle.split(
+            mix_layer, num_or_sections=3, axis=-1
+        )
+        attn_output, _ = flash_attention(
+            query_states, key_states, value_states, causal=True
+        )
+        attn_output = paddle.cast(attn_output, paddle.float32)
+        attn_output = attn_output.reshape(
+            [BATCH_SIZE, SEQ_LEN, self.hidden_size]
+        )
+        attn_output = self.o_proj(attn_output)
+        return attn_output
+
+
+class LlamaMlp(nn.Layer):
+    def __init__(
+        self,
+        param_prefix="",
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        weight_attr_0 = create_numpy_like_random(param_prefix + "_0")
+        bias_attr_0 = create_numpy_like_random(param_prefix + "_bias_0")
+        weight_attr_1 = create_numpy_like_random(param_prefix + "_1")
+        bias_attr_1 = create_numpy_like_random(param_prefix + "_bias_1")
+        weight_attr_2 = create_numpy_like_random(param_prefix + "_2")
+        bias_attr_2 = create_numpy_like_random(param_prefix + "_bias_2")
+
+        self.up_proj = nn.Linear(
+            hidden_size, intermediate_size, weight_attr_0, bias_attr_0
+        )
+        self.gate_proj = nn.Linear(
+            hidden_size, intermediate_size, weight_attr_1, bias_attr_1
+        )
+        self.down_proj = nn.Linear(
+            intermediate_size, hidden_size, weight_attr_2, bias_attr_2
+        )
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LlamaRMSNorm(nn.Layer):
+    def __init__(self, hidden_size=HIDDEN_SIZE):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.weight = paddle.create_parameter(
+            shape=[self.hidden_size],
+            dtype=paddle.get_default_dtype(),
+            default_initializer=nn.initializer.Constant(1.0),
+        )
+        self.variance_epsilon = 1.0
+
+    def forward(self, hidden_states):
+        with paddle.amp.auto_cast(False):
+            variance = (
+                hidden_states.astype("float32").pow(2).mean(-1, keepdim=True)
+            )
+            hidden_states = (
+                paddle.rsqrt(variance + self.variance_epsilon) * hidden_states
+            )
+        if self.weight.dtype in [paddle.float16, paddle.bfloat16]:
+            hidden_states = paddle.cast(hidden_states, self.weight.dtype)
+        return hidden_states * self.weight
+
+
+class LlamaLayerDecoder(nn.Layer):
+    def __init__(
+        self,
+        param_prefix="",
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+
+        self.self_attn = LlamaAttention(param_prefix + "_att", hidden_size)
+        self.mlp = LlamaMlp(param_prefix + "_mlp")
+        self.input_layernorm = LlamaRMSNorm(hidden_size)
+        self.post_attn_layernorm = LlamaRMSNorm(hidden_size)
+
+    def forward(self, x):
+        residual = x
+        hidden_states = self.input_layernorm(x)
+        hidden_states = self.self_attn(hidden_states)
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attn_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class TestLlamaDecoderForSemiAutoParallel:
+    def __init__(self):
+        self._dtype = os.getenv("dtype", "float32")
+        self._backend = os.getenv("backend", "gpu")
+        self._seed = eval(os.getenv("seed", "2023"))
+        paddle.set_device(self._backend)
+        self.init_single_card_net_result()
+
+    def dp_mp_shard_fn(self, layer_name, layer, process_mesh):
+        col_linear = ["qkv_proj", "gate_proj", "up_proj"]
+        row_linear = ["o_proj", "down_proj"]
+
+        def contains(a, b):
+            return b in a
+
+        is_col_linear = any(contains(layer_name, e) for e in col_linear)
+        is_row_linear = any(contains(layer_name, e) for e in row_linear)
+
+        if is_col_linear:
+            layer.weight = dist.shard_tensor(
+                layer.weight, process_mesh, [Replicate(), Shard(1)]
+            )
+            layer.bias = dist.shard_tensor(
+                layer.bias, process_mesh, [Replicate(), Shard(0)]
+            )
+
+        if is_row_linear:
+            layer.weight = dist.shard_tensor(
+                layer.weight, process_mesh, [Replicate(), Shard(0)]
+            )
+
+    def set_random_seed(self, seed):
+        random.seed(seed)
+        np.random.seed(seed)
+        paddle.seed(seed)
+
+    def init_input_data(self):
+        input = np.random.random([BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE]).astype(
+            self._dtype
+        )
+        input = paddle.to_tensor(input)
+        return input
+
+    def init_single_card_net_result(self):
+        self.set_random_seed(self._seed)
+        self.base_out, self.base_parameters = self.train_loop(
+            LlamaLayerDecoder("demo_weight")
+        )
+
+    def train_loop(self, layer, process_mesh=None, shard_input=False):
+        # run forward and backward
+
+        opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=layer.parameters()
+        )
+        for _ in range(5):
+            input = self.init_input_data()
+            if shard_input:
+                input = dist.shard_tensor(input, process_mesh, shard_input)
+            out = layer(input)
+            loss = paddle.sum(out)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        return out, layer.parameters()
+
+    def check_tensor_eq(self, a, b, rtol=1e-05, atol=0, verbose=True):
+        if a is None:
+            assert b is None
+            return
+        np1 = a.astype("float32").numpy()
+        np2 = b.astype("float32").numpy()
+        np.testing.assert_allclose(
+            np1, np2, rtol=rtol, atol=atol, verbose=verbose
+        )
+
+    def check_placements(self, output, expected_placements):
+        assert output.placements == expected_placements, (
+            f"{output.placements}  vs {expected_placements}"
+        )
+
+    def get_shard_check_hook(self, dims_mapping, check_input=False):
+        def check_func(layer, input, output=None):
+            if check_input:
+                if isinstance(input, tuple):
+                    input = input[0]
+                self.check_placements(input, dims_mapping)
+            else:
+                if isinstance(output, tuple):
+                    output = output[0]
+                self.check_placements(output, dims_mapping)
+
+        return check_func
+
+    # python -m paddle.distributed.launch --devices=0,1,2,3 semi_auto_parallel_for_llama_decoder_dp_mp.py
+    def test_dp_mp(self):
+        self.set_random_seed(self._seed)
+        dp_mp_mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+        dp_mp_layer = dist.shard_layer(
+            LlamaLayerDecoder("mp_demo_weight"), dp_mp_mesh, self.dp_mp_shard_fn
+        )
+        input_layer_norm_post_hook = self.get_shard_check_hook(
+            [dist.Shard(0), dist.Replicate()]
+        )
+        attn_pre_hook = self.get_shard_check_hook(
+            [dist.Shard(0), dist.Replicate()], True
+        )
+        attn_post_hook = self.get_shard_check_hook(
+            [dist.Shard(0), dist.Replicate()]
+        )
+        post_attn_layer_norm_pre_hook = self.get_shard_check_hook(
+            [dist.Shard(0), dist.Replicate()], True
+        )
+        post_attn_layer_norm_post_hook = self.get_shard_check_hook(
+            [dist.Shard(0), dist.Replicate()]
+        )
+        mlp_pre_hook = self.get_shard_check_hook(
+            [dist.Shard(0), dist.Replicate()], True
+        )
+        mlp_post_hook = self.get_shard_check_hook(
+            [dist.Shard(0), dist.Replicate()]
+        )
+
+        dp_mp_layer.input_layernorm.register_forward_post_hook(
+            input_layer_norm_post_hook
+        )
+
+        dp_mp_layer.self_attn.register_forward_pre_hook(attn_pre_hook)
+        dp_mp_layer.self_attn.register_forward_post_hook(attn_post_hook)
+
+        dp_mp_layer.post_attn_layernorm.register_forward_pre_hook(
+            post_attn_layer_norm_pre_hook
+        )
+        dp_mp_layer.post_attn_layernorm.register_forward_post_hook(
+            post_attn_layer_norm_post_hook
+        )
+
+        dp_mp_layer.mlp.register_forward_pre_hook(mlp_pre_hook)
+        dp_mp_layer.mlp.register_forward_post_hook(mlp_post_hook)
+
+        dp_mp_out, dp_mp_parameters = self.train_loop(
+            dp_mp_layer, dp_mp_mesh, shard_input=[Shard(0), Replicate()]
+        )
+        self.check_tensor_eq(dp_mp_out, self.base_out)
+        for param, param_base in zip(dp_mp_parameters, self.base_parameters):
+            self.check_tensor_eq(param, param_base)
+            self.check_tensor_eq(param.grad, param_base.grad)
+
+    def run_test_case(self):
+        if self._backend == "gpu":
+            cuda_version_main = int(paddle.version.cuda().split(".")[0])
+            device_prop_main = paddle.device.cuda.get_device_capability()[0]
+            if cuda_version_main >= 11 and device_prop_main >= 8:
+                self.test_dp_mp()
+
+
+if __name__ == '__main__':
+    TestLlamaDecoderForSemiAutoParallel().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_global_input.py

@@ -0,0 +1,223 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+SEQ_LEN = 4
+HIDDEN_SIZE = 8
+global_mesh = dist.ProcessMesh(
+    [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=['pp', 'dp', 'mp']
+)
+mesh0 = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['dp', 'mp'])
+mesh1 = dist.ProcessMesh([[4, 5], [6, 7]], dim_names=['dp', 'mp'])
+
+
+class MlpModel(paddle.nn.Layer):
+    def __init__(self, variable_initial_values, run_single_process=False):
+        super().__init__()
+        self.w0 = self.create_parameter(
+            shape=[HIDDEN_SIZE, HIDDEN_SIZE],
+            default_initializer=paddle.nn.initializer.Assign(
+                variable_initial_values[0]
+            ),
+        )
+        self.w1 = self.create_parameter(
+            shape=[HIDDEN_SIZE, HIDDEN_SIZE],
+            default_initializer=paddle.nn.initializer.Assign(
+                variable_initial_values[1]
+            ),
+        )
+        self.global_input = paddle.uniform(
+            shape=[SEQ_LEN, HIDDEN_SIZE],
+            dtype=paddle.float32,
+            min=-0.0001,
+            max=0.0001,
+        )
+        if run_single_process is False:
+            self.w0 = dist.shard_tensor(
+                self.w0,
+                mesh0,
+                [dist.Replicate(), dist.Shard(1)],
+            )
+            self.w1 = dist.shard_tensor(
+                self.w1,
+                mesh1,
+                [dist.Replicate(), dist.Shard(0)],
+            )
+            self.global_input = dist.shard_tensor(
+                self.global_input,
+                global_mesh,
+                [dist.Replicate(), dist.Replicate(), dist.Replicate()],
+            )
+        self.run_single_process = run_single_process
+
+    def process_global_input(self, input):
+        return input + 0.0001
+
+    def forward(self, x):
+        # x: [bs, seq_len, hidden]
+        # forward on mesh0
+        global_input = self.process_global_input(self.global_input)
+        if self.run_single_process is False:
+            global_input1 = dist.reshard(
+                global_input, mesh0, [dist.Replicate(), dist.Replicate()]
+            )
+        else:
+            global_input1 = global_input
+        x = x + global_input1
+        y = x @ self.w0
+        # forward on mesh1
+        if self.run_single_process is False:
+            y = dist.reshard(y, mesh1, [dist.Shard(0), dist.Shard(2)])
+            global_input2 = dist.reshard(
+                global_input, mesh1, [dist.Replicate(), dist.Replicate()]
+            )
+        else:
+            global_input2 = global_input
+
+        y = y + global_input2
+        z = y @ self.w1
+        return z
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, hidden, num_samples=8):
+        super().__init__()
+        self.seq_len = seq_len
+        self.hidden = hidden
+        self.num_samples = num_samples
+        self.inputs = [
+            np.random.uniform(size=[self.seq_len, self.hidden]).astype(
+                "float32"
+            )
+            for _ in range(num_samples)
+        ]
+        self.labels = [
+            np.array(index, dtype="float32") for index in range(num_samples)
+        ]
+
+    def __getitem__(self, index):
+        return self.inputs[index], self.labels[index]
+
+    def __len__(self):
+        return self.num_samples
+
+
+def create_dataloader():
+    dataset = RandomDataset(SEQ_LEN, HIDDEN_SIZE)
+    sampler = BatchSampler(
+        dataset,
+        batch_size=2,
+    )
+    dataloader = DataLoader(
+        dataset,
+        batch_sampler=sampler,
+    )
+    return dataloader
+
+
+def get_variable_initial_value(var_num=2):
+    res = []
+    for i in range(var_num):
+        res.append(
+            paddle.uniform(
+                shape=[HIDDEN_SIZE, HIDDEN_SIZE],
+                dtype=paddle.float32,
+                min=-0.0001,
+                max=0.0001,
+            )
+        )
+    return res
+
+
+def loss_fn(logits, label):
+    # logits: [bs, seq_len, hidden], label: [bs]
+    loss = paddle.nn.MSELoss(reduction="sum")
+    logits = paddle.sum(logits, axis=[1, 2])
+    return loss(logits, label)
+
+
+class TestSemiAutoParallelGlobalInput:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._run_static = eval(os.getenv("run_static"))
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        paddle.set_device(self._backend)
+        self.dataloader = create_dataloader()
+        self.variable_initial_values = get_variable_initial_value()
+        self.single_process_loss = self.get_single_process_loss()
+
+    def get_single_process_loss(self):
+        model = MlpModel(
+            variable_initial_values=self.variable_initial_values,
+            run_single_process=True,
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+        for step, (input, label) in enumerate(self.dataloader()):
+            logits = model(input)
+            loss = loss_fn(logits, label)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        return loss.numpy()
+
+    def test_basic(self):
+        model = MlpModel(variable_initial_values=self.variable_initial_values)
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+        dist_dataloader = dist.shard_dataloader(
+            dataloader=self.dataloader, meshes=[mesh0, mesh1], shard_dims="dp"
+        )
+        cur_rank = paddle.distributed.get_rank()
+        if self._run_static:
+            dist_model = dist.to_static(model, dist_dataloader, loss_fn, opt)
+            dist_model.train()
+
+            for input, label in dist_dataloader:
+                loss = dist_model(input, label)
+
+            if cur_rank in [5, 7]:
+                loss = paddle.to_tensor(loss)
+                group = paddle.distributed.new_group([5, 7])
+                dist.all_reduce(loss, group=group)
+        else:
+            dist_opt = dist.shard_optimizer(opt)
+            for input, label in dist_dataloader:
+                logits = model(input)
+                loss = loss_fn(logits, label)
+                loss.backward()
+                dist_opt.step()
+                dist_opt.clear_grad()
+        if cur_rank in [5, 7]:
+            np.testing.assert_allclose(
+                loss.numpy(), self.single_process_loss, rtol=1e-06, verbose=True
+            )
+
+    def run_test_case(self):
+        self.test_basic()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelGlobalInput().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_multi_inputs.py

@@ -0,0 +1,255 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+SEQ_LEN = 4
+HIDDEN_SIZE = 8
+global_mesh = dist.ProcessMesh(
+    [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=['pp', 'dp', 'mp']
+)
+mesh0 = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['dp', 'mp'])
+mesh1 = dist.ProcessMesh([[4, 5], [6, 7]], dim_names=['dp', 'mp'])
+
+
+class MlpModel(paddle.nn.Layer):
+    def __init__(self, variable_initial_values, run_single_process=False):
+        super().__init__()
+        self.w0 = self.create_parameter(
+            shape=[HIDDEN_SIZE, HIDDEN_SIZE],
+            default_initializer=paddle.nn.initializer.Assign(
+                variable_initial_values[0]
+            ),
+        )
+        self.w1 = self.create_parameter(
+            shape=[HIDDEN_SIZE, HIDDEN_SIZE],
+            default_initializer=paddle.nn.initializer.Assign(
+                variable_initial_values[1]
+            ),
+        )
+        if run_single_process is False:
+            self.w0 = dist.shard_tensor(
+                self.w0,
+                mesh0,
+                [dist.Replicate(), dist.Shard(1)],
+            )
+            self.w1 = dist.shard_tensor(
+                self.w1,
+                mesh1,
+                [dist.Replicate(), dist.Shard(0)],
+            )
+        self.run_single_process = run_single_process
+
+    def forward(self, input1, input2, extra_input1=None, extra_input2=None):
+        # extra_input1 and extra_input2 only used for test non_tensor input in shard_dataloader
+        x = input1 + input2
+        # x: [bs, seq_len, hidden]
+        # forward on mesh0
+        y = x @ self.w0
+        # forward on mesh1
+        if self.run_single_process is False:
+            y = dist.reshard(y, mesh1, [dist.Shard(0), dist.Shard(2)])
+        z = y @ self.w1
+        return z
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, hidden, num_samples=8):
+        super().__init__()
+        self.seq_len = seq_len
+        self.hidden = hidden
+        self.num_samples = num_samples
+        self.inputs1 = [
+            np.random.uniform(size=[self.seq_len, self.hidden]).astype(
+                "float32"
+            )
+            for _ in range(num_samples)
+        ]
+        self.inputs2 = [
+            np.random.uniform(size=[self.seq_len, self.hidden]).astype(
+                "float32"
+            )
+            for _ in range(num_samples)
+        ]
+        self.labels = [
+            np.array(index, dtype="float32") for index in range(num_samples)
+        ]
+
+    def __getitem__(self, index):
+        return {
+            "inputs": [self.inputs1[index], self.inputs2[index]],
+            "label": self.labels[index],
+        }
+
+    def __len__(self):
+        return self.num_samples
+
+
+def create_dataloader(collate_fn=None):
+    dataset = RandomDataset(SEQ_LEN, HIDDEN_SIZE)
+    sampler = BatchSampler(
+        dataset,
+        batch_size=2,
+    )
+    dataloader = DataLoader(
+        dataset,
+        batch_sampler=sampler,
+        collate_fn=collate_fn,
+    )
+    return dataloader
+
+
+def get_variable_initial_value(var_num=2):
+    res = []
+    for i in range(var_num):
+        res.append(
+            paddle.uniform(
+                shape=[HIDDEN_SIZE, HIDDEN_SIZE],
+                dtype=paddle.float32,
+                min=-0.0001,
+                max=0.0001,
+            )
+        )
+    return res
+
+
+def loss_fn(logits, label):
+    # logits: [bs, seq_len, hidden], label: [bs]
+    loss = paddle.nn.MSELoss(reduction="sum")
+    logits = paddle.sum(logits, axis=[1, 2])
+    return loss(logits, label)
+
+
+class TestSemiAutoParallelMultiInputs:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._run_static = eval(os.getenv("run_static"))
+        paddle.seed(self._seed)
+        np.random.seed(self._seed)
+        paddle.set_device(self._backend)
+        self.dataloader = create_dataloader()
+        self.variable_initial_values = get_variable_initial_value()
+        self.single_process_loss = self.get_single_process_loss()
+
+    def get_single_process_loss(self):
+        model = MlpModel(
+            variable_initial_values=self.variable_initial_values,
+            run_single_process=True,
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+        for step, data in enumerate(self.dataloader()):
+            input1, input2 = data["inputs"]
+            logits = model(input1, input2)
+            label = data["label"]
+            loss = loss_fn(logits, label)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        return loss.numpy()
+
+    def test_basic(self):
+        model = MlpModel(variable_initial_values=self.variable_initial_values)
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+        dist_dataloader = dist.shard_dataloader(
+            dataloader=self.dataloader,
+            meshes=[mesh0, mesh1],  # or [[mesh0, mesh0], mesh1]
+            shard_dims="dp",
+            input_keys=["inputs", "label"],
+        )
+        cur_rank = paddle.distributed.get_rank()
+        if self._run_static:
+            dist_model = dist.to_static(model, dist_dataloader, loss_fn, opt)
+            dist_model.train()
+
+            for step, data in enumerate(dist_dataloader()):
+                input1, input2 = data["inputs"]
+                label = data["label"]
+                loss = dist_model(input1, input2, label)
+
+            if cur_rank in [5, 7]:
+                loss = paddle.to_tensor(loss)
+                group = paddle.distributed.new_group([5, 7])
+                dist.all_reduce(loss, group=group)
+        else:
+            dist_opt = dist.shard_optimizer(opt)
+            for step, data in enumerate(dist_dataloader()):
+                input1, input2 = data["inputs"]
+                logits = model(input1, input2)
+                label = data["label"]
+                loss = loss_fn(logits, label)
+                loss.backward()
+                dist_opt.step()
+                dist_opt.clear_grad()
+        if cur_rank in [5, 7]:
+            np.testing.assert_allclose(
+                loss.numpy(), self.single_process_loss, rtol=1e-06, verbose=True
+            )
+
+    def test_non_tensor_input(self):
+        model = MlpModel(variable_initial_values=self.variable_initial_values)
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+
+        def custom_collate_fn(batch):
+            collated_batch = {
+                "inputs": [
+                    paddle.to_tensor([item["inputs"][0] for item in batch]),
+                    paddle.to_tensor([item["inputs"][1] for item in batch]),
+                    12.0,
+                ],
+                "extra_input": 12,
+                "label": paddle.to_tensor([item["label"] for item in batch]),
+            }
+            return collated_batch
+
+        self.dataloader = create_dataloader(custom_collate_fn)
+
+        dist_dataloader = dist.shard_dataloader(
+            dataloader=self.dataloader,
+            meshes=[mesh0, mesh0, mesh1],
+            shard_dims="dp",
+            input_keys=["inputs", "extra_input", "label"],
+        )
+
+        dist_opt = dist.shard_optimizer(opt)
+        for step, data in enumerate(dist_dataloader()):
+            input1, input2, extra_input1 = data["inputs"]
+            extra_input2 = data["extra_input"]
+            logits = model(input1, input2, extra_input1, extra_input2)
+            label = data["label"]
+            loss = loss_fn(logits, label)
+            loss.backward()
+            dist_opt.step()
+            dist_opt.clear_grad()
+
+    def run_test_case(self):
+        self.test_basic()
+        if not self._run_static:
+            self.test_non_tensor_input()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelMultiInputs().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_mutual_load_between_dynamic_and_static.py

@@ -0,0 +1,228 @@
+# 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 random
+
+import numpy as np
+from auto_parallel.semi_auto_parallel_dist_to_static_mlp import RandomDataset
+from auto_parallel.semi_auto_parallel_simple_net import (
+    BATCH_SIZE,
+    CLASS_NUM,
+    IMAGE_SIZE,
+    DemoNet,
+    TestSimpleNetForSemiAutoParallel,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.io import DataLoader
+
+
+class TestSemiAutoParallelMutualLoadBetweenDynamicAndStatic(
+    TestSimpleNetForSemiAutoParallel
+):
+    def __init__(self):
+        self._ckpt_path = os.environ.get("ckpt_path")
+        self._seed = os.environ.get("seed", 123)
+        self.mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def create_data_loader(self):
+        images = np.random.rand(BATCH_SIZE, IMAGE_SIZE).astype('float32')
+        labels = np.random.rand(BATCH_SIZE, CLASS_NUM).astype('float32')
+        dataset = RandomDataset(images, labels, BATCH_SIZE)
+        loader = DataLoader(dataset, batch_size=BATCH_SIZE)
+        return loader
+
+    def run_dynamic(self, layer, opt, data_loader, is_recompute=False):
+        # MSELoss only support pir, but test_save_load_state_dict.py set FLAGS_enable_pir_api=0
+        loss_fn = nn.SmoothL1Loss()
+
+        loss_list = []
+        for _ in range(5):
+            for batch_id, (image, label) in enumerate(data_loader()):
+                if is_recompute:
+                    image.stop_gradient = False
+                out = layer(image)
+                loss = loss_fn(out, label)
+                loss_list.append(loss.numpy())
+                loss.backward()
+
+                opt.step()
+                opt.clear_grad()
+        return np.array(loss_list)
+
+    def run_dy2static(self, layer, opt, data_loader):
+        # create loss
+        # MSELoss only support pir, but test_save_load_state_dict.py set FLAGS_enable_pir_api=0
+        loss_fn = nn.SmoothL1Loss()
+        dist_loader = dist.shard_dataloader(
+            dataloader=data_loader,
+            meshes=[self.mesh],
+            shard_dims=None,
+        )
+        # static training
+        dist_model = dist.to_static(layer, dist_loader, loss_fn, opt)
+        loss_list = []
+        dist_model.train()
+        for epoch in range(5):
+            for batch_id, (image, label) in enumerate(dist_loader()):
+                loss = dist_model(image, label)
+                loss_list.append(loss)
+
+        return np.array(loss_list), dist_model
+
+    def set_random_seed(self, seed):
+        random.seed(seed)
+        np.random.seed(seed)
+        paddle.seed(seed)
+
+    def test_dygraph_save_static_load(self):
+        paddle.disable_static()
+        # set seed to promise the same input for different tp rank
+        self.set_random_seed(self._seed)
+        data_loader = self.create_data_loader()
+        dist_loader = dist.shard_dataloader(
+            dataloader=data_loader,
+            meshes=[self.mesh],
+            shard_dims=None,
+        )
+        dy_layer = dist.shard_layer(
+            DemoNet("dp_mp_hybrid_strategy"), self.mesh, self.shard_fn
+        )
+        dy_opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=dy_layer.parameters()
+        )
+        dy_losses = self.run_dynamic(dy_layer, dy_opt, dist_loader)
+        saved_dy_layer_state_dict = dy_layer.state_dict()
+        ckpt_path = os.path.join(
+            self._ckpt_path, "test_dygraph_save_static_load"
+        )
+        dist.save_state_dict(saved_dy_layer_state_dict, ckpt_path)
+        dist.barrier()
+
+        dy2static_opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=dy_layer.parameters()
+        )
+        # MSELoss only support pir, but test_save_load_state_dict.py set FLAGS_enable_pir_api=0
+        loss_fn = nn.SmoothL1Loss()
+        dist_model = dist.to_static(
+            dy_layer, dist_loader, loss_fn, dy2static_opt
+        )
+        need_load_state_dict = {}
+        expected_state_dict = {}
+        with paddle.base.dygraph.guard():
+            for k, v in saved_dy_layer_state_dict.items():
+                expected_state_dict[k] = v._local_value().clone()
+                need_load_state_dict[k] = paddle.zeros_like(v)
+        dist_model.train()
+        dist_model.set_state_dict(need_load_state_dict)
+        state_dict_to_load = dist_model.state_dict(mode="param")
+        assert len(state_dict_to_load) == len(expected_state_dict)
+        for k, v in state_dict_to_load.items():
+            assert k in expected_state_dict, (
+                f"key {k} not in expected_state_dict:{expected_state_dict}"
+            )
+            assert np.any(
+                np.not_equal(
+                    v._local_value().numpy(),
+                    expected_state_dict[k].numpy(),
+                )
+            ), (
+                f"key:{k}, v:{v}, expected_state_dict[k]:{expected_state_dict[k]}"
+            )
+
+        dist.load_state_dict(state_dict_to_load, ckpt_path)
+        dist_model.set_state_dict(state_dict_to_load)
+
+        program_state_dict = dist_model.state_dict(mode="param")
+        assert len(expected_state_dict) == len(program_state_dict)
+        for k, v in program_state_dict.items():
+            assert k in expected_state_dict, (
+                f"key {k} not in expected_state_dict:{expected_state_dict}"
+            )
+            np.testing.assert_equal(
+                v._local_value().numpy(),
+                expected_state_dict[k].numpy(),
+            )
+
+    def test_static_save_dynamic_load(self):
+        paddle.disable_static()
+        # set seed to promise the same input for different tp rank
+        self.set_random_seed(self._seed)
+        data_loader = self.create_data_loader()
+
+        dy_layer = dist.shard_layer(
+            DemoNet("dp_mp_hybrid_strategy"), self.mesh, self.shard_fn
+        )
+
+        dy2static_opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=dy_layer.parameters()
+        )
+        dy2static_losses, dist_model = self.run_dy2static(
+            dy_layer, dy2static_opt, data_loader
+        )
+
+        saved_static_layer_state_dict = dist_model.state_dict("param")
+        ckpt_path = os.path.join(
+            self._ckpt_path, "test_static_save_dynamic_load"
+        )
+        dist.save_state_dict(saved_static_layer_state_dict, ckpt_path)
+        dist.barrier()
+
+        paddle.disable_static()
+        need_load_state_dict = {}
+        expected_state_dict = {}
+        with paddle.base.dygraph.guard():
+            for k, v in saved_static_layer_state_dict.items():
+                expected_state_dict[k] = v._local_value().clone()
+                need_load_state_dict[k] = paddle.zeros_like(v)
+        dy_layer.set_state_dict(need_load_state_dict)
+        state_dict_to_load = dy_layer.state_dict()
+        assert len(state_dict_to_load) == len(expected_state_dict)
+        for k, v in state_dict_to_load.items():
+            assert k in expected_state_dict, (
+                f"key {k} not in expected_state_dict:{expected_state_dict}"
+            )
+            assert np.any(
+                np.not_equal(
+                    v._local_value().numpy(),
+                    expected_state_dict[k].numpy(),
+                )
+            ), (
+                f"key:{k}, v:{v}, expected_state_dict[k]:{expected_state_dict[k]}"
+            )
+
+        dist.load_state_dict(state_dict_to_load, ckpt_path)
+        dy_layer.set_state_dict(state_dict_to_load)
+
+        state_dict = dy_layer.state_dict()
+        assert len(expected_state_dict) == len(state_dict)
+        for k, v in state_dict.items():
+            assert k in expected_state_dict, (
+                f"key {k} not in expected_state_dict:{expected_state_dict}"
+            )
+            np.testing.assert_equal(
+                v._local_value().numpy(),
+                expected_state_dict[k].numpy(),
+            )
+
+    def run_test_case(self):
+        self.test_dygraph_save_static_load()
+        self.test_static_save_dynamic_load()
+
+
+if __name__ == "__main__":
+    TestSemiAutoParallelMutualLoadBetweenDynamicAndStatic().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_nd_cross_mesh_reshard.py

@@ -0,0 +1,419 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestSemiAutoParallelNdCrossMeshReshard:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh0 = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+        self._mesh1 = dist.ProcessMesh([[4, 5], [6, 7]], dim_names=["x", "y"])
+        self._dst_rank = [4, 5, 6, 7]
+        self._shape = (20, 20)
+        paddle.set_device(self._backend)
+
+    def test_pp_to_rr(self):
+        a = paddle.ones(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a,
+            self._mesh0,
+            [
+                dist.Partial(dist.ReduceType.kRedSum),
+                dist.Partial(dist.ReduceType.kRedSum),
+            ],
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Replicate(), dist.Replicate()]
+        )
+
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+
+    def test_pp_to_ss(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=0, num_or_sections=2)
+        if dist.get_rank() in [4, 5]:
+            expect_out = paddle.split(expect_out[0], axis=1, num_or_sections=2)
+        else:
+            expect_out = paddle.split(expect_out[1], axis=1, num_or_sections=2)
+        expect_out_shape = [10, 10]
+
+        input_tensor = dist.shard_tensor(
+            a,
+            self._mesh0,
+            [
+                dist.Partial(dist.ReduceType.kRedSum),
+                dist.Partial(dist.ReduceType.kRedSum),
+            ],
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Shard(0), dist.Shard(1)]
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                expect_out[dist.get_rank() % 2].numpy(),
+            )
+
+    def test_rr_to_pp(self):
+        a = paddle.ones(self._shape)
+        b = paddle.zeros(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Replicate(), dist.Replicate()]
+        )
+        out = dist.reshard(
+            input_tensor,
+            self._mesh1,
+            [
+                dist.Partial(dist.ReduceType.kRedSum),
+                dist.Partial(dist.ReduceType.kRedSum),
+            ],
+        )
+        if dist.get_rank() == 4:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+        if dist.get_rank() in [5, 6, 7]:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), b.numpy())
+
+    def test_rr_to_ss(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=0, num_or_sections=2)
+        if dist.get_rank() in [4, 5]:
+            expect_out = paddle.split(expect_out[0], axis=1, num_or_sections=2)
+        else:
+            expect_out = paddle.split(expect_out[1], axis=1, num_or_sections=2)
+        expect_out_shape = [10, 10]
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Replicate(), dist.Replicate()]
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Shard(0), dist.Shard(1)]
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                expect_out[dist.get_rank() % 2].numpy(),
+            )
+
+    def test_ss_to_pp(self):
+        a = paddle.ones(self._shape)
+        b = paddle.zeros(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(0), dist.Shard(1)]
+        )
+        out = dist.reshard(
+            input_tensor,
+            self._mesh1,
+            [
+                dist.Partial(dist.ReduceType.kRedSum),
+                dist.Partial(dist.ReduceType.kRedSum),
+            ],
+        )
+        if dist.get_rank() == 4:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+        if dist.get_rank() in [5, 6, 7]:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), b.numpy())
+
+    def test_ss_to_rr(self):
+        a = paddle.ones(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(0), dist.Shard(1)]
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Replicate(), dist.Replicate()]
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+
+    def test_ss_to_ss(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=0, num_or_sections=2)
+        if dist.get_rank() in [4, 5]:
+            expect_out = paddle.split(expect_out[0], axis=1, num_or_sections=2)
+        else:
+            expect_out = paddle.split(expect_out[1], axis=1, num_or_sections=2)
+        expect_out_shape = [10, 10]
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(0), dist.Shard(1)]
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Shard(1), dist.Shard(0)]
+        )
+
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                expect_out[dist.get_rank() % 2].numpy(),
+            )
+
+    def test_sp_to_ps(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=1, num_or_sections=2)
+        expect_out_shape = [20, 10]
+        b = paddle.zeros(expect_out_shape)
+
+        input_tensor = dist.shard_tensor(
+            a,
+            self._mesh0,
+            [dist.Shard(0), dist.Partial(dist.ReduceType.kRedSum)],
+        )
+        out = dist.reshard(
+            input_tensor,
+            self._mesh1,
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Shard(1)],
+        )
+
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            if dist.get_rank() in [4, 5]:
+                np.testing.assert_equal(
+                    out._local_value().numpy(),
+                    expect_out[dist.get_rank() % 2].numpy(),
+                )
+            else:
+                np.testing.assert_equal(
+                    out._local_value().numpy(),
+                    b.numpy(),
+                )
+
+    def test_sp_to_rs(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=1, num_or_sections=2)
+        expect_out_shape = [20, 10]
+
+        input_tensor = dist.shard_tensor(
+            a,
+            self._mesh0,
+            [dist.Shard(0), dist.Partial(dist.ReduceType.kRedSum)],
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Replicate(), dist.Shard(1)]
+        )
+
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                expect_out[dist.get_rank() % 2].numpy(),
+            )
+
+    def test_sp_to_rp(self):
+        a = paddle.ones(self._shape)
+        b = paddle.zeros(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a,
+            self._mesh0,
+            [dist.Shard(0), dist.Partial(dist.ReduceType.kRedSum)],
+        )
+        out = dist.reshard(
+            input_tensor,
+            self._mesh1,
+            [dist.Replicate(), dist.Partial(dist.ReduceType.kRedSum)],
+        )
+
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            if dist.get_rank() % 2 == 0:
+                np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+            else:
+                np.testing.assert_equal(out._local_value().numpy(), b.numpy())
+
+    def test_sr_to_ps(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=1, num_or_sections=2)
+        expect_out_shape = [20, 10]
+        b = paddle.zeros(expect_out_shape)
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(0), dist.Replicate()]
+        )
+        out = dist.reshard(
+            input_tensor,
+            self._mesh1,
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Shard(1)],
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            if dist.get_rank() in [4, 5]:
+                np.testing.assert_equal(
+                    out._local_value().numpy(),
+                    expect_out[dist.get_rank() % 2].numpy(),
+                )
+            else:
+                np.testing.assert_equal(
+                    out._local_value().numpy(),
+                    b.numpy(),
+                )
+
+    def test_sr_to_rs(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=1, num_or_sections=2)
+        expect_out_shape = [20, 10]
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(0), dist.Replicate()]
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Replicate(), dist.Shard(1)]
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            local_rank_in_mesh = dist.get_rank() - 4
+            shard_idx = local_rank_in_mesh % 2
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                expect_out[shard_idx].numpy(),
+            )
+
+    def test_sr_to_rp(self):
+        a = paddle.ones(self._shape)
+        b = paddle.zeros(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a, self._mesh0, [dist.Shard(0), dist.Replicate()]
+        )
+        out = dist.reshard(
+            input_tensor,
+            self._mesh1,
+            [dist.Replicate(), dist.Partial(dist.ReduceType.kRedSum)],
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            if dist.get_rank() % 2 == 0:
+                np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+            else:
+                np.testing.assert_equal(out._local_value().numpy(), b.numpy())
+
+    def test_pr_to_ps(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=1, num_or_sections=2)
+        expect_out_shape = [20, 10]
+        b = paddle.zeros(expect_out_shape)
+
+        input_tensor = dist.shard_tensor(
+            a,
+            self._mesh0,
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Replicate()],
+        )
+        out = dist.reshard(
+            input_tensor,
+            self._mesh1,
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Shard(1)],
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            if dist.get_rank() in [4, 5]:
+                np.testing.assert_equal(
+                    out._local_value().numpy(),
+                    expect_out[dist.get_rank() % 2].numpy(),
+                )
+            else:
+                np.testing.assert_equal(
+                    out._local_value().numpy(),
+                    b.numpy(),
+                )
+
+    def test_pr_to_rs(self):
+        a = paddle.ones(self._shape)
+        expect_out = paddle.split(a, axis=1, num_or_sections=2)
+        expect_out_shape = [20, 10]
+
+        input_tensor = dist.shard_tensor(
+            a,
+            self._mesh0,
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Replicate()],
+        )
+        out = dist.reshard(
+            input_tensor, self._mesh1, [dist.Replicate(), dist.Shard(1)]
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            assert np.equal(out._local_shape, expect_out_shape).all()
+            np.testing.assert_equal(
+                out._local_value().numpy(),
+                expect_out[dist.get_rank() % 2].numpy(),
+            )
+
+    def test_pr_to_rp(self):
+        a = paddle.ones(self._shape)
+        b = paddle.zeros(self._shape)
+
+        input_tensor = dist.shard_tensor(
+            a,
+            self._mesh0,
+            [dist.Partial(dist.ReduceType.kRedSum), dist.Replicate()],
+        )
+        out = dist.reshard(
+            input_tensor,
+            self._mesh1,
+            [dist.Replicate(), dist.Partial(dist.ReduceType.kRedSum)],
+        )
+        if dist.get_rank() in self._dst_rank:
+            assert np.equal(out.shape, input_tensor.shape).all()
+            if dist.get_rank() % 2 == 0:
+                np.testing.assert_equal(out._local_value().numpy(), a.numpy())
+            else:
+                np.testing.assert_equal(out._local_value().numpy(), b.numpy())
+
+    def run_test_case(self):
+        self.test_pp_to_rr()
+        self.test_pp_to_ss()
+        self.test_rr_to_pp()
+        self.test_rr_to_ss()
+        self.test_ss_to_pp()
+        self.test_ss_to_rr()
+        self.test_ss_to_ss()
+        self.test_sp_to_ps()
+        self.test_sp_to_rs()
+        self.test_sp_to_rp()
+        self.test_sr_to_ps()
+        self.test_sr_to_rs()
+        self.test_sr_to_rp()
+        self.test_pr_to_ps()
+        self.test_pr_to_rs()
+        self.test_pr_to_rp()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelNdCrossMeshReshard().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_sharding_stage_1.py

@@ -0,0 +1,114 @@
+# 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 os
+
+import numpy as np
+from auto_parallel.semi_auto_parallel_dist_to_static_api import (
+    DemoNet,
+    create_data_loader,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+
+
+class TestSemiAutoParallelShardingStage1:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+
+    def check_tensor_eq(self, a, b, rtol=1e-05, atol=0, verbose=True):
+        np.testing.assert_allclose(a, b, rtol=rtol, atol=atol, verbose=verbose)
+
+    def shard_layer_fn(self, layer_name, layer, process_mesh):
+        layer.weight = dist.shard_tensor(
+            layer.weight, process_mesh, [dist.Replicate(), dist.Shard(1)]
+        )
+        layer.bias = dist.shard_tensor(
+            layer.bias, process_mesh, [dist.Replicate(), dist.Shard(0)]
+        )
+
+    def get_single_card_rst(self):
+        paddle.seed(self._seed)
+        linear = paddle.nn.Linear(10, 10)
+        batch = paddle.rand(shape=[10, 10])
+        opt = paddle.optimizer.AdamW(parameters=linear.parameters())
+        for _ in range(5):
+            loss = linear(batch)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        self.weight = linear.weight.numpy()
+        self.bias = linear.bias.numpy()
+
+    def test_sharding_stage_1_with_mp(self):
+        paddle.seed(self._seed)
+        linear = paddle.nn.Linear(10, 10)
+        linear = dist.shard_layer(linear, self._mesh, self.shard_layer_fn)
+        batch = paddle.rand(shape=[10, 10])
+        # shard the input by sharding degree
+        batch = dist.shard_tensor(batch, self._mesh, [dist.Shard(0)])
+        # shard optimizer with stage 1 fn
+        opt = paddle.optimizer.AdamW(parameters=linear.parameters())
+        opt = dist.shard_optimizer(opt, dist.ShardingStage1("x", self._mesh))
+        for _ in range(5):
+            loss = linear(batch)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        self.check_tensor_eq(self.weight, linear.weight.numpy())
+        self.check_tensor_eq(self.bias, linear.bias.numpy())
+
+    def test_sharding_stage_1_with_mp_to_static(self):
+        data_loader = create_data_loader()
+        layer = DemoNet(
+            self._mesh, "sharding_with_mp_demonet", shard_weight=True
+        )
+        opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=layer.parameters()
+        )
+        opt = dist.shard_optimizer(opt, dist.ShardingStage1("x", self._mesh))
+        loss_fn = nn.MSELoss()
+
+        dist_loader = dist.shard_dataloader(
+            dataloader=data_loader,
+            meshes=[self._mesh],
+            shard_dims=0,
+        )
+
+        dist_model = dist.to_static(layer, dist_loader, loss_fn, opt)
+
+        dist_model.train()
+        for epoch in range(2):
+            for batch_id, (image, label) in enumerate(dist_loader()):
+                loss = dist_model(image, label)
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+        elif self._backend == "gpu":
+            paddle.set_device("gpu:" + str(dist.get_rank()))
+        else:
+            raise ValueError("Only support cpu or gpu backend.")
+
+        self.get_single_card_rst()
+        self.test_sharding_stage_1_with_mp()
+        self.test_sharding_stage_1_with_mp_to_static()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelShardingStage1().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_sharding_stage_2.py

@@ -0,0 +1,114 @@
+# 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 os
+
+import numpy as np
+from auto_parallel.semi_auto_parallel_dist_to_static_api import (
+    DemoNet,
+    create_data_loader,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+
+
+class TestSemiAutoParallelShardingStage2:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+
+    def check_tensor_eq(self, a, b, rtol=1e-05, atol=0, verbose=True):
+        np.testing.assert_allclose(a, b, rtol=rtol, atol=atol, verbose=verbose)
+
+    def shard_layer_fn(self, layer_name, layer, process_mesh):
+        layer.weight = dist.shard_tensor(
+            layer.weight, process_mesh, [dist.Shard(1)]
+        )
+        layer.bias = dist.shard_tensor(
+            layer.bias, process_mesh, [dist.Shard(0)]
+        )
+
+    def get_single_card_rst(self):
+        paddle.seed(self._seed)
+        linear = paddle.nn.Linear(10, 10)
+        batch = paddle.rand(shape=[10, 10])
+        opt = paddle.optimizer.AdamW(parameters=linear.parameters())
+        for _ in range(5):
+            loss = linear(batch)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        self.weight = linear.weight.numpy()
+        self.bias = linear.bias.numpy()
+
+    def test_sharding_stage_2_with_mp(self):
+        paddle.seed(self._seed)
+        linear = paddle.nn.Linear(10, 10)
+        linear = dist.shard_layer(linear, self._mesh, self.shard_layer_fn)
+        batch = paddle.rand(shape=[10, 10])
+        # shard the input by sharding degree
+        batch = dist.shard_tensor(batch, self._mesh, [dist.Shard(0)])
+        # shard optimizer with stage 1 fn
+        opt = paddle.optimizer.AdamW(parameters=linear.parameters())
+        opt = dist.shard_optimizer(opt, dist.ShardingStage2("x", self._mesh))
+        for _ in range(5):
+            loss = linear(batch)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        self.check_tensor_eq(self.weight, linear.weight.numpy())
+        self.check_tensor_eq(self.bias, linear.bias.numpy())
+
+    def test_sharding_stage_2_with_mp_to_static(self):
+        data_loader = create_data_loader()
+        layer = DemoNet(
+            self._mesh, "sharding_with_mp_demonet", shard_weight=True
+        )
+        opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=layer.parameters()
+        )
+        opt = dist.shard_optimizer(opt, dist.ShardingStage2("x", self._mesh))
+        loss_fn = nn.MSELoss()
+
+        dist_loader = dist.shard_dataloader(
+            dataloader=data_loader,
+            meshes=[self._mesh],
+            shard_dims=0,
+        )
+
+        dist_model = dist.to_static(layer, dist_loader, loss_fn, opt)
+
+        dist_model.train()
+        for epoch in range(2):
+            for batch_id, (image, label) in enumerate(dist_loader()):
+                loss = dist_model(image, label)
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+        elif self._backend == "gpu":
+            paddle.set_device("gpu:" + str(dist.get_rank()))
+        else:
+            raise ValueError("Only support cpu or gpu backend.")
+
+        self.get_single_card_rst()
+        self.test_sharding_stage_2_with_mp()
+        self.test_sharding_stage_2_with_mp_to_static()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelShardingStage2().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_sharding_stage_3.py

@@ -0,0 +1,114 @@
+# 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 os
+
+import numpy as np
+from auto_parallel.semi_auto_parallel_dist_to_static_api import (
+    DemoNet,
+    create_data_loader,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+
+
+class TestSemiAutoParallelShardingStage3:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+
+    def check_tensor_eq(self, a, b, rtol=1e-05, atol=0, verbose=True):
+        np.testing.assert_allclose(a, b, rtol=rtol, atol=atol, verbose=verbose)
+
+    def shard_layer_fn(self, layer_name, layer, process_mesh):
+        layer.weight = dist.shard_tensor(
+            layer.weight, process_mesh, [dist.Shard(1)]
+        )
+        layer.bias = dist.shard_tensor(
+            layer.bias, process_mesh, [dist.Shard(0)]
+        )
+
+    def get_single_card_rst(self):
+        paddle.seed(self._seed)
+        linear = paddle.nn.Linear(10, 10)
+        batch = paddle.rand(shape=[10, 10])
+        opt = paddle.optimizer.AdamW(parameters=linear.parameters())
+        for _ in range(5):
+            loss = linear(batch)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        self.weight = linear.weight.numpy()
+        self.bias = linear.bias.numpy()
+
+    def test_sharding_stage_3_with_mp(self):
+        paddle.seed(self._seed)
+        linear = paddle.nn.Linear(10, 10)
+        linear = dist.shard_layer(linear, self._mesh, self.shard_layer_fn)
+        batch = paddle.rand(shape=[10, 10])
+        # shard the input by sharding degree
+        batch = dist.shard_tensor(batch, self._mesh, [dist.Shard(0)])
+        # shard optimizer with stage 1 fn
+        opt = paddle.optimizer.AdamW(parameters=linear.parameters())
+        opt = dist.shard_optimizer(opt, dist.ShardingStage3("x", self._mesh))
+        for _ in range(5):
+            loss = linear(batch)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+        self.check_tensor_eq(self.weight, linear.weight.numpy())
+        self.check_tensor_eq(self.bias, linear.bias.numpy())
+
+    def test_sharding_stage_3_with_mp_to_static(self):
+        data_loader = create_data_loader()
+        layer = DemoNet(
+            self._mesh, "sharding_with_mp_demonet", shard_weight=True
+        )
+        opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=layer.parameters()
+        )
+        opt = dist.shard_optimizer(opt, dist.ShardingStage3("x", self._mesh))
+        loss_fn = nn.MSELoss()
+
+        dist_loader = dist.shard_dataloader(
+            dataloader=data_loader,
+            meshes=[self._mesh],
+            shard_dims=0,
+        )
+
+        dist_model = dist.to_static(layer, dist_loader, loss_fn, opt)
+
+        dist_model.train()
+        for epoch in range(2):
+            for batch_id, (image, label) in enumerate(dist_loader()):
+                loss = dist_model(image, label)
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+        elif self._backend == "gpu":
+            paddle.set_device("gpu:" + str(dist.get_rank()))
+        else:
+            raise ValueError("Only support cpu or gpu backend.")
+
+        self.get_single_card_rst()
+        self.test_sharding_stage_3_with_mp()
+        self.test_sharding_stage_3_with_mp_to_static()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelShardingStage3().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_simple_net_dp_mp.py

@@ -0,0 +1,96 @@
+# 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
+
+from auto_parallel.hybrid_strategy.semi_auto_save_state_dict import (
+    check_structure_name_mapping,
+)
+from auto_parallel.semi_auto_parallel_simple_net import (
+    DemoNet,
+    TestSimpleNetForSemiAutoParallel,
+)
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestSimpleNetHybridStrategyForSemiAutoParallel(
+    TestSimpleNetForSemiAutoParallel
+):
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._ckpt_path = os.getenv("ckpt_path")
+        self._mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+
+        paddle.set_device(self._backend)
+
+        self.set_random_seed(self._seed)
+        self.init_single_card_net_result()
+
+    def test_dp_mp_demo_net(self):
+        self.set_random_seed(self._seed)
+        model = dist.shard_layer(
+            DemoNet("dp_mp_hybrid_strategy"), self._mesh, self.shard_fn
+        )
+
+        (
+            self.dp_mp_loss,
+            self.dp_mp_parameters,
+        ) = self.run_dynamic(model, shard_input=True)
+
+        self.check_tensor_eq(self.dp_mp_loss, self.base_loss, rtol=1e-04)
+        for param, param_base in zip(
+            self.dp_mp_parameters, self.base_parameters
+        ):
+            self.check_tensor_eq(param, param_base, atol=1e-06)
+            self.check_tensor_eq(param.grad, param_base.grad)
+
+        # save load
+        state_dict = model.state_dict()
+        paddle.distributed.save_state_dict(state_dict, self._ckpt_path)
+        paddle.distributed.barrier()
+        check_structure_name_mapping(self._ckpt_path, state_dict)
+        expected_local_state_dict = {}
+        need_load_state_dict = {}
+        for k, v in state_dict.items():
+            local_value = v._local_value()
+            expected_local_state_dict[k] = local_value.clone()
+            need_load_state_dict[k] = paddle.zeros_like(v)
+        model.set_state_dict(need_load_state_dict)
+        state_dict = model.state_dict()
+        for k, v in state_dict.items():
+            assert v.numpy().sum() == 0.0, f"state_dict {k} is not zero"
+            assert k in need_load_state_dict, f"state_dict {k} is not found"
+            assert need_load_state_dict[k].numpy().sum() == 0.0, (
+                f"state_dict {k} is not zero"
+            )
+
+        paddle.distributed.load_state_dict(
+            need_load_state_dict, self._ckpt_path
+        )
+        model.set_state_dict(need_load_state_dict)
+        state_dict = model.state_dict()
+        for k, v in state_dict.items():
+            assert k in expected_local_state_dict, k
+            self.check_tensor_eq(v._local_value(), expected_local_state_dict[k])
+
+    def run_test_case(self):
+        self.test_dp_mp_demo_net()
+
+
+if __name__ == '__main__':
+    TestSimpleNetHybridStrategyForSemiAutoParallel().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_simple_net_dp_mp_pp.py

@@ -0,0 +1,152 @@
+# 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
+
+from auto_parallel.hybrid_strategy.semi_auto_save_state_dict import (
+    check_structure_name_mapping,
+)
+from auto_parallel.semi_auto_parallel_simple_net import (
+    DemoNet,
+    TestSimpleNetForSemiAutoParallel,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import Replicate, Shard
+
+
+class TestSimpleNetHybridStrategyForSemiAutoParallel(
+    TestSimpleNetForSemiAutoParallel
+):
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._ckpt_path = os.getenv("ckpt_path")
+        self._mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=["x", "y"])
+        self._pp_mesh0 = dist.ProcessMesh(
+            [[0, 1], [2, 3]], dim_names=["x", "y"]
+        )
+        self._pp_mesh1 = dist.ProcessMesh(
+            [[4, 5], [6, 7]], dim_names=["x", "y"]
+        )
+        self.pp_reshard_dist_attr = (self._pp_mesh1, [Shard(0), Shard(1)])
+
+        paddle.set_device(self._backend)
+
+        self.set_random_seed(self._seed)
+        self.init_single_card_net_result()
+
+    def dp_mp_pp_shard_fn(self, layer_name, layer, process_mesh):
+        if layer_name == 'linear_0':
+            # shard_layer doesn't support cross-mesh now.
+            # input process_mesh of pp_shard_fn is useless,
+            # it's defined just for unified format.
+            layer.weight = dist.shard_tensor(
+                layer.weight, self._pp_mesh0, [Replicate(), Shard(1)]
+            )
+            if layer.bias is not None:
+                layer.bias = dist.shard_tensor(
+                    layer.bias, self._pp_mesh0, [Replicate(), Replicate()]
+                )
+        elif layer_name == 'linear_1':
+            layer.weight = dist.shard_tensor(
+                layer.weight, self._pp_mesh1, [Replicate(), Shard(0)]
+            )
+            if layer.bias is not None:
+                layer.bias = dist.shard_tensor(
+                    layer.bias, self._pp_mesh1, [Replicate(), Replicate()]
+                )
+        elif layer_name == 'norm':
+            layer.weight = dist.shard_tensor(
+                layer.weight, self._pp_mesh1, [Replicate(), Replicate()]
+            )
+            if layer.bias is not None:
+                layer.bias = dist.shard_tensor(
+                    layer.bias, self._pp_mesh1, [Replicate(), Replicate()]
+                )
+
+    def test_dp_mp_pp_demo_net(self):
+        self.set_random_seed(self._seed)
+        model = dist.shard_layer(
+            DemoNet(
+                "dp_mp_pp_hybrid_strategy",
+                is_pp=True,
+                pp_reshard_dist_attr=self.pp_reshard_dist_attr,
+            ),
+            self._pp_mesh0,
+            self.dp_mp_pp_shard_fn,
+        )
+
+        (
+            self.dp_mp_pp_loss,
+            self.dp_mp_pp_parameters,
+        ) = self.run_dynamic(model, shard_input=True, is_pp=True)
+
+        rank = dist.get_rank()
+        # TODO(GhostScreaming): DistTensor.numpy() doesn't support
+        # cross-mesh now, ReshardXToReplicated function in eager_method
+        # needs to be fixed later.
+        if rank in [0, 1, 2, 3]:
+            # linear_0 weight and bias
+            self.check_tensor_eq(
+                self.dp_mp_pp_parameters[0], self.base_parameters[0], rtol=2e-4
+            )
+
+        else:
+            self.check_tensor_eq(self.dp_mp_pp_loss, self.base_loss, rtol=1e-4)
+            self.check_tensor_eq(
+                self.dp_mp_pp_parameters[1], self.base_parameters[1], rtol=1e-4
+            )
+            self.check_tensor_eq(
+                self.dp_mp_pp_parameters[2], self.base_parameters[2], rtol=2e-5
+            )
+            self.check_tensor_eq(
+                self.dp_mp_pp_parameters[3], self.base_parameters[3], rtol=2e-4
+            )
+
+        # save load
+        state_dict = model.state_dict()
+        paddle.distributed.save_state_dict(state_dict, self._ckpt_path)
+        paddle.distributed.barrier()
+        check_structure_name_mapping(self._ckpt_path, state_dict)
+        expected_local_state_dict = {}
+        need_load_state_dict = {}
+        for k, v in state_dict.items():
+            expected_local_state_dict[k] = (
+                v._local_value().clone() if v._is_initialized() else None
+            )
+            need_load_state_dict[k] = (
+                paddle.zeros_like(v) if v._is_initialized() else v
+            )
+        model.set_state_dict(need_load_state_dict)
+        paddle.distributed.load_state_dict(
+            need_load_state_dict, self._ckpt_path
+        )
+        model.set_state_dict(need_load_state_dict)
+        state_dict = model.state_dict()
+        for k, v in state_dict.items():
+            assert k in expected_local_state_dict, k
+            if v._is_initialized():
+                self.check_tensor_eq(
+                    v._local_value(), expected_local_state_dict[k]
+                )
+
+    def run_test_case(self):
+        self.test_dp_mp_pp_demo_net()
+
+
+if __name__ == '__main__':
+    TestSimpleNetHybridStrategyForSemiAutoParallel().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_parallel_simple_net_sp.py

@@ -0,0 +1,214 @@
+# 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 numpy as np
+from auto_parallel.semi_auto_parallel_simple_net import (
+    TestSimpleNetForSemiAutoParallel,
+    create_numpy_like_random,
+)
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import Replicate, Shard
+
+BATCH_SIZE = 8
+SEQUENCE_LEN = 512
+HIDDEN_SIZE = 1024
+NUM_HEAD = 64
+HEAD_DIM = 16
+CLASS_NUM = 10
+
+np.set_printoptions(threshold=np.inf)
+
+
+class DemoNet(nn.Layer):
+    def __init__(self, param_prefix="", is_sp=False, is_dp=False):
+        super().__init__()
+
+        if is_dp:
+            self.pp0_mesh = dist.ProcessMesh(
+                [[0, 1], [2, 3]], dim_names=["dp", "mp"]
+            )
+            self.pp1_mesh = dist.ProcessMesh(
+                [[4, 5], [6, 7]], dim_names=["dp", "mp"]
+            )
+            self.placement0 = [Replicate(), Shard(1)]
+            self.placement1 = [Replicate(), Shard(0)]
+            self.sp_reshard_placement0 = [Shard(1), Shard(0)]
+            self.sp_reshard_placement1 = [Shard(1), Replicate()]
+        else:
+            self.pp0_mesh = dist.ProcessMesh([0, 1], dim_names=["mp"])
+            self.pp1_mesh = dist.ProcessMesh([2, 3], dim_names=["mp"])
+            self.placement0 = [Shard(1)]
+            self.placement1 = [Shard(0)]
+            self.sp_reshard_placement0 = [Shard(0)]
+            self.sp_reshard_placement1 = [Replicate()]
+
+        self.is_sp = is_sp
+        self.is_dp = is_dp
+
+        self.norm = nn.LayerNorm(HIDDEN_SIZE, epsilon=1e-4)
+        self.linear_0_weight = dist.shard_tensor(
+            self.create_parameter(
+                shape=[HEAD_DIM, 4 * HIDDEN_SIZE],
+                attr=create_numpy_like_random(param_prefix + "w_0"),
+                dtype=paddle.float32,
+                is_bias=False,
+            ),
+            self.pp0_mesh,
+            self.placement0,
+        )
+
+        self.linear_1_weight = dist.shard_tensor(
+            self.create_parameter(
+                shape=[4 * HIDDEN_SIZE, HEAD_DIM],
+                attr=create_numpy_like_random(param_prefix + "w_1"),
+                dtype=paddle.float32,
+                is_bias=False,
+            ),
+            self.pp0_mesh,
+            self.placement1,
+        )
+
+        self.linear_2_weight = dist.shard_tensor(
+            self.create_parameter(
+                shape=[HIDDEN_SIZE, 4 * HIDDEN_SIZE],
+                attr=create_numpy_like_random(param_prefix + "w_2"),
+                dtype=paddle.float32,
+                is_bias=False,
+            ),
+            self.pp1_mesh,
+            self.placement0,
+        )
+
+        self.linear_3_weight = dist.shard_tensor(
+            self.create_parameter(
+                shape=[4 * HIDDEN_SIZE, CLASS_NUM],
+                attr=create_numpy_like_random(param_prefix + "w_3"),
+                dtype=paddle.float32,
+                is_bias=False,
+            ),
+            self.pp1_mesh,
+            self.placement1,
+        )
+
+    def forward(self, x):
+        # Layer 0
+        tgt = paddle.transpose(x, [1, 0, 2])
+        out = paddle.reshape(x, [BATCH_SIZE, SEQUENCE_LEN, NUM_HEAD, HEAD_DIM])
+        # [BATCH_SIZE, SEQUENCE_LEN, NUM_HEAD, HEAD_DIM] -> [BATCH_SIZE, NUM_HEAD, SEQUENCE_LEN, HEAD_DIM]
+        out = paddle.transpose(out, [0, 2, 1, 3])
+        out = paddle.matmul(out, self.linear_0_weight)
+        out = paddle.matmul(out, self.linear_1_weight)
+        out = paddle.transpose(out, [2, 0, 1, 3])
+        out = paddle.reshape(out, [SEQUENCE_LEN, BATCH_SIZE, HIDDEN_SIZE])
+
+        # SP Region, should be reduce_scatter
+        if self.is_sp:
+            out = dist.reshard(out, self.pp0_mesh, self.sp_reshard_placement0)
+
+        out = out + tgt
+        out = self.norm(out)
+
+        out = dist.reshard(out, self.pp1_mesh, self.sp_reshard_placement1)
+
+        out = paddle.matmul(out, self.linear_2_weight)
+        out = paddle.matmul(out, self.linear_3_weight)
+        out = paddle.transpose(out, [1, 0, 2])
+
+        return out
+
+
+class TestSimpleNetHybridStrategyForSemiAutoParallel(
+    TestSimpleNetForSemiAutoParallel
+):
+    def __init__(self):
+        self._dtype = os.getenv("dtype")
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._is_dp = os.getenv("is_dp") == "true"
+        if self._is_dp:
+            self.pp0_mesh = dist.ProcessMesh(
+                [[0, 1], [2, 3]], dim_names=["dp", "mp"]
+            )
+
+        paddle.set_device(self._backend)
+
+        self.set_random_seed(self._seed)
+        self.init_single_card_net_result()
+
+    def init_single_card_net_result(self):
+        self.set_random_seed(self._seed)
+        self.base_loss, self.base_parameters = self.run_dynamic(
+            DemoNet("demo_weight", is_sp=False, is_dp=self._is_dp), is_sp=False
+        )
+
+    def init_input_data(self):
+        image = np.random.randn(BATCH_SIZE, SEQUENCE_LEN, HIDDEN_SIZE).astype(
+            'float32'
+        )
+        label = np.random.randn(BATCH_SIZE, SEQUENCE_LEN, CLASS_NUM).astype(
+            'float32'
+        )
+
+        return paddle.to_tensor(image), paddle.to_tensor(label)
+
+    def check_tensor_eq(self, a, b, rtol=1e-7, atol=0, verbose=True):
+        np1 = a.astype("float32").numpy()
+        np2 = b.astype("float32").numpy()
+        np.testing.assert_allclose(
+            np1, np2, rtol=rtol, atol=atol, verbose=verbose
+        )
+
+    def run_dynamic(self, layer, is_sp=False):
+        # create loss
+        loss_fn = nn.MSELoss()
+        # run forward and backward
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=layer.parameters()
+        )
+        for _ in range(5):
+            image, label = self.init_input_data()
+            if self._is_dp:
+                image = dist.shard_tensor(
+                    image, self.pp0_mesh, [Shard(0), Replicate()]
+                )
+
+            out = layer(image)
+
+            loss = loss_fn(out, label)
+            loss.backward()
+            opt.step()
+        return loss, layer.parameters()
+
+    def test_dp_mp_sp_demo_net(self):
+        self.set_random_seed(self._seed)
+        model = DemoNet("dp_mp_hybrid_strategy", is_sp=True, is_dp=self._is_dp)
+
+        (
+            self.dp_mp_sp_loss,
+            self.dp_mp_sp_parameters,
+        ) = self.run_dynamic(model, is_sp=True)
+        if dist.get_rank() in model.pp1_mesh.process_ids:
+            self.check_tensor_eq(self.dp_mp_sp_loss, self.base_loss, rtol=1e-3)
+
+    def run_test_case(self):
+        self.test_dp_mp_sp_demo_net()
+
+
+if __name__ == '__main__':
+    TestSimpleNetHybridStrategyForSemiAutoParallel().run_test_case()

test/auto_parallel/hybrid_strategy/semi_auto_save_state_dict.py

@@ -0,0 +1,212 @@
+# 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 paddle
+import paddle.distributed as dist
+from paddle.distributed import save_state_dict
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
+    ShardedWeight,
+    make_replicated_sharded_weight,
+)
+
+
+def get_global_state_dict():
+    w1 = paddle.arange(104).reshape([13, 8])
+    w2 = paddle.arange(32, 36).reshape([2, 2])
+    return {"w1": w1, "w2": w2}
+
+
+def check_structure_name_mapping(ckpt_path, state_dict):
+    metadata_file_path = os.path.join(ckpt_path, "0.metadata")
+    data_file_path = os.path.join(
+        ckpt_path, f"{paddle.distributed.get_rank()}_0.distcp"
+    )
+    assert os.path.exists(metadata_file_path), (
+        f"metadata file {metadata_file_path} is not found"
+    )
+    assert os.path.exists(data_file_path), (
+        f"data file {data_file_path} is not found"
+    )
+    metadata = paddle.load(metadata_file_path)
+    cur_rank_state_dict = paddle.load(data_file_path, keep_name_table=True)
+    local_structure_name_mapping = cur_rank_state_dict.pop(
+        "StructuredToParameterName@@"
+    )
+    assert isinstance(local_structure_name_mapping, dict), (
+        f"local_structure_name_mapping:{local_structure_name_mapping} is not dict type"
+    )
+    for structure_name, param_name in local_structure_name_mapping.items():
+        assert structure_name in state_dict, (
+            f"tensor key:{structure_name} is not found in state dict:{state_dict}"
+        )
+        assert param_name == state_dict[structure_name].name, (
+            f"param name:{param_name} is not equal to param name in state_dict:{state_dict[structure_name].name}"
+        )
+
+
+class TestSaveStateDict:
+    def __init__(self):
+        self._ckpt_path = os.getenv("ckpt_path")
+
+    def test_save_state_dict_with_one_device(self):
+        global_state_dict = get_global_state_dict()
+        keys = list(global_state_dict.keys())
+        w1, w2 = list(global_state_dict.values())
+        state_dict = dict(zip(keys, [w1, w2]))
+        save_state_dict(state_dict, self._ckpt_path)
+        check_structure_name_mapping(self._ckpt_path, state_dict)
+
+    def test_save_state_dict_with_two_devices(self):
+        global_state_dict = get_global_state_dict()
+        keys = list(global_state_dict.keys())
+        w1, w2 = list(global_state_dict.values())
+        mesh = dist.ProcessMesh([0, 1])
+        sharded_w1 = dist.shard_tensor(w1, mesh, [dist.Shard(0)])
+        sharded_w2 = dist.shard_tensor(w2, mesh, [dist.Shard(0)])
+        state_dict = dict(zip(keys, [sharded_w1, sharded_w2]))
+        save_state_dict(state_dict, self._ckpt_path)
+        paddle.distributed.barrier()
+        check_structure_name_mapping(self._ckpt_path, state_dict)
+
+    def run_test_case(self):
+        device_num = int(os.getenv("device_num"))
+        if device_num == 1:
+            self.test_save_state_dict_with_one_device()
+        elif device_num == 2:
+            self.test_save_state_dict_with_two_devices()
+
+
+class TestSaveShardedStateDict:
+    def __init__(self):
+        self._ckpt_path = os.getenv("ckpt_path_2")
+
+    def test_save_state_dict_with_one_device(self):
+        # Construct a 4x4 integer tensor as expected result:
+        # [[ 0,  1,  2,  3],
+        #  [ 4,  5,  6,  7],
+        #  [ 8,  9, 10, 11],
+        #  [12, 13, 14, 15]]
+        local_tensor = paddle.to_tensor(
+            [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]],
+            dtype='int32',
+        )
+        sharded_state_dict = {}
+        sharded_state_dict["t"] = make_replicated_sharded_weight(
+            "t", local_tensor
+        )
+        save_state_dict(sharded_state_dict, self._ckpt_path)
+
+    def test_save_state_dict_with_two_devices(self):
+        if dist.get_rank() == 0:
+            # On rank 0:
+            # The global tensor (4x4) is distributed as:
+            # [[ 0,  1,  2,  *],
+            #  [ *,  *,  *,  *],
+            #  [ *,  *,  *,  *],
+            #  [ *,  *,  *,  *]]
+            # Numbers 0,1,2 are local, '*' means not present on this rank.
+            local_tensor = paddle.to_tensor([0, 1, 2], dtype='int32')
+            sharded_weight = ShardedWeight(
+                key="t",
+                local_tensor=local_tensor,
+                local_shape=(4, 4),
+                global_shape=(4, 4),
+                global_offset=(0, 0),
+                is_flattened=True,
+                flattened_range=slice(0, 3),
+            )
+        elif dist.get_rank() == 1:
+            # On rank 1:
+            # The global tensor (4x4) is distributed as:
+            # [[ *,  *,  *,  3],
+            #  [ 4,  5,  5,  6],
+            #  [ 8,  9,  10, 11],
+            #  [ 12, 13, 14, 15]]
+            # Numbers 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15 are local, '*' means not present on this rank.
+            local_tensor = paddle.to_tensor(
+                [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], dtype='int32'
+            )
+            sharded_weight = ShardedWeight(
+                key="t",
+                local_tensor=local_tensor,
+                local_shape=(4, 4),
+                global_shape=(4, 4),
+                global_offset=(0, 0),
+                is_flattened=True,
+                flattened_range=slice(3, 16),
+            )
+
+        sharded_state_dict = {"t": sharded_weight}
+        save_state_dict(sharded_state_dict, self._ckpt_path)
+        paddle.distributed.barrier()
+
+    def run_test_case(self):
+        device_num = int(os.getenv("device_num"))
+        if device_num == 1:
+            self.test_save_state_dict_with_one_device()
+        elif device_num == 2:
+            self.test_save_state_dict_with_two_devices()
+
+
+class TestSaveShardedStateDictWithReplica:
+    def __init__(self):
+        self._ckpt_path = os.getenv("ckpt_path_3")
+
+    def test_save_state_dict_with_one_device(self):
+        # Construct a 4x4 integer tensor as expected result:
+        # [[ 0,  1,  2,  3],
+        #  [ 4,  5,  6,  7],
+        #  [ 8,  9, 10, 11],
+        #  [12, 13, 14, 15]]
+        local_tensor = paddle.to_tensor(
+            [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]],
+            dtype='int32',
+        )
+        sharded_state_dict = {}
+        sharded_state_dict["t"] = make_replicated_sharded_weight(
+            "t", local_tensor
+        )
+        save_state_dict(sharded_state_dict, self._ckpt_path, save_replicas=True)
+
+    def test_save_state_dict_with_two_devices(self):
+        # Construct a 4x4 integer tensor as expected result:
+        # [[ 0,  1,  2,  3],
+        #  [ 4,  5,  6,  7],
+        #  [ 8,  9, 10, 11],
+        #  [12, 13, 14, 15]]
+        local_tensor = paddle.to_tensor(
+            [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]],
+            dtype='int32',
+        )
+        sharded_state_dict = {}
+        sharded_state_dict["t"] = make_replicated_sharded_weight(
+            "t", local_tensor
+        )
+        save_state_dict(sharded_state_dict, self._ckpt_path, save_replicas=True)
+        paddle.distributed.barrier()
+
+    def run_test_case(self):
+        device_num = int(os.getenv("device_num"))
+        if device_num == 1:
+            self.test_save_state_dict_with_one_device()
+        elif device_num == 2:
+            self.test_save_state_dict_with_two_devices()
+
+
+if __name__ == "__main__":
+    TestSaveStateDict().run_test_case()
+    TestSaveShardedStateDict().run_test_case()
+    TestSaveShardedStateDictWithReplica().run_test_case()

test/auto_parallel/hybrid_strategy/semi_flexcheckpoint_merge.py

@@ -0,0 +1,253 @@
+# 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 os
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.io import BatchSampler, DataLoader, Dataset
+
+
+class RandomDataset(Dataset):
+    def __init__(self, seq_len, hidden, num_samples=100):
+        super().__init__()
+        self.seq_len = seq_len
+        self.hidden = hidden
+        self.num_samples = num_samples
+
+    def __getitem__(self, index):
+        input = np.random.uniform(size=[self.seq_len, self.hidden]).astype(
+            "float32"
+        )
+        return input
+
+    def __len__(self):
+        return self.num_samples
+
+
+class DistMlpModel(paddle.nn.Layer):
+    def __init__(self, mesh):
+        super().__init__()
+        self.w0 = self.create_parameter(shape=[1024, 4096])
+        self.w1 = self.create_parameter(shape=[4096, 1024])
+        self.mesh = mesh
+        self.w0 = dist.shard_tensor(
+            self.w0, mesh, [dist.Replicate(), dist.Shard(1)]
+        )
+        self.w1 = dist.shard_tensor(
+            self.w1, mesh, [dist.Replicate(), dist.Shard(0)]
+        )
+
+    def forward(self, x):
+        x = dist.shard_tensor(x, self.mesh, [dist.Shard(0), dist.Replicate()])
+        y = paddle.matmul(x, self.w0)
+        z = paddle.matmul(y, self.w1)
+        return z
+
+
+class MultiMlpModel(paddle.nn.Layer):
+    def __init__(self, mesh):
+        super().__init__()
+        self.layer1 = DistMlpModel(mesh)
+        self.layer2 = DistMlpModel(mesh)
+
+    def forward(self, x):
+        y = self.layer1(x)
+        z = self.layer2(x)
+        return z
+
+
+class SingleMlpModel(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.w0 = self.create_parameter(shape=[1024, 4096])
+        self.w1 = self.create_parameter(shape=[4096, 1024])
+
+    def forward(self, x):
+        y = paddle.matmul(x, self.w0)
+        z = paddle.matmul(y, self.w1)
+        return z
+
+
+class TestDistCheckpoint:
+    def __init__(self):
+        np.random.seed(42)
+        self.mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['dp', 'mp'])
+        self.temp_dir = os.getenv("ckpt_path")
+
+    def _get_single_loss(self, dataloader, unsharded_state_dict):
+        with paddle.LazyGuard():
+            model = SingleMlpModel()
+        model.w0.set_value(unsharded_state_dict['w0'])
+        model.w1.set_value(unsharded_state_dict['w1'])
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+
+        losses = []
+        for step, inputs in enumerate(dataloader):
+            data = inputs
+            logits = model(data)
+            loss = paddle.mean(logits)
+            losses.append(float(loss))
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+
+        return losses[0]
+
+    def _get_dist_loss(self, dataloader, sharded_state_dict):
+        with paddle.LazyGuard():
+            model = DistMlpModel(self.mesh)
+        model.w0.set_value(sharded_state_dict['w0'])
+        model.w1.set_value(sharded_state_dict['w1'])
+
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+
+        losses = []
+        for step, inputs in enumerate(dataloader):
+            data = inputs
+            logits = model(data)
+            loss = paddle.mean(logits)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+            losses.append(float(loss))
+
+        return losses[0]
+
+    def dist_checkpoint(self, offload=False, safetensors=True, unique_id=0):
+        model_path = os.path.join(self.temp_dir, f'model.{unique_id}')
+        opt_path = os.path.join(self.temp_dir, f'opt.{unique_id}')
+
+        # Test checkpoint saving
+        with paddle.LazyGuard():
+            model = DistMlpModel(self.mesh)
+        for p in model.parameters():
+            p.initialize()
+
+        dataset = RandomDataset(128, 1024)
+        sampler = BatchSampler(
+            dataset,
+            batch_size=4,
+        )
+        dataloader = DataLoader(
+            dataset,
+            batch_sampler=sampler,
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+        opt = dist.shard_optimizer(opt)
+
+        for step, inputs in enumerate(dataloader):
+            data = inputs
+            logits = model(data)
+            loss = paddle.mean(logits)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+
+        dist.save_state_dict(
+            model.state_dict(), model_path, safetensors=safetensors
+        )
+        dist.save_state_dict(
+            opt.state_dict(), opt_path, safetensors=safetensors
+        )
+
+        unsharded_state_dict = dist.load_merged_state_dict(
+            model_path, offload=offload, safetensors=safetensors
+        )
+        # Get single loss
+        single_loss = self._get_single_loss(dataloader, unsharded_state_dict)
+
+        shard_state_dict = model.state_dict()
+        dist.load_state_dict(
+            shard_state_dict, model_path, safetensors=safetensors
+        )
+
+        # Get distributed loss
+        dist_loss = self._get_dist_loss(dataloader, shard_state_dict)
+        np.testing.assert_array_equal(
+            unsharded_state_dict['w0'].numpy(), shard_state_dict['w0'].numpy()
+        )
+        np.testing.assert_array_equal(
+            unsharded_state_dict['w1'].numpy(), shard_state_dict['w1'].numpy()
+        )
+
+    def test_dist_checkpoint(self):
+        self.dist_checkpoint(True, True, 0)
+        self.dist_checkpoint(False, True, 1)
+        self.dist_checkpoint(True, False, 2)
+        self.dist_checkpoint(False, False, 3)
+
+    def count_files_in_temp_dir(self, single_path):
+        if not os.path.exists(single_path):
+            return 0
+        files = [
+            f
+            for f in os.listdir(single_path)
+            if os.path.isfile(os.path.join(single_path, f))
+        ]
+        return len(files)
+
+    def test_checkpoint_load_merge_save(self):
+        model_path = os.path.join(self.temp_dir, 'model')
+        single_path = os.path.join(self.temp_dir, 'single_model')
+
+        # Test checkpoint saving
+        with paddle.LazyGuard():
+            model = MultiMlpModel(self.mesh)
+        for p in model.parameters():
+            p.initialize()
+
+        dataset = RandomDataset(128, 1024)
+        sampler = BatchSampler(
+            dataset,
+            batch_size=4,
+        )
+        dataloader = DataLoader(
+            dataset,
+            batch_sampler=sampler,
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=model.parameters()
+        )
+        opt = dist.shard_optimizer(opt)
+
+        for step, inputs in enumerate(dataloader):
+            data = inputs
+            logits = model(data)
+            loss = paddle.mean(logits)
+            loss.backward()
+            opt.step()
+            opt.clear_grad()
+
+        dist.save_state_dict(model.state_dict(), model_path, safetensors=False)
+
+        dist.flex_checkpoint.dcp.load_state_dict.merge_sharded_state_dict(
+            model_path, single_path, offload=True, safetensors=False
+        )
+        # assert self.count_files_in_temp_dir(single_path) == 5, (
+        #     f"Expected 5 files in temp dir, but got {self.count_files_in_temp_dir(single_path)}"
+        # )
+
+
+if __name__ == '__main__':
+    TestDistCheckpoint().test_dist_checkpoint()
+    TestDistCheckpoint().test_checkpoint_load_merge_save()

test/auto_parallel/hybrid_strategy/single_llama_model.py

@@ -0,0 +1,330 @@
+# 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 paddle
+import paddle.distributed as dist
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.nn.functional.flash_attention import _math_attention
+
+
+class SDPALayer(paddle.nn.Layer):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+    def forward(self, query, key, value, **kwargs):
+        if (
+            self.config.context_parallel is True
+            or self.config.sep_parallel is True
+        ) and (
+            int(paddle.version.cuda().split(".")[0]) >= 11
+            and paddle.device.cuda.get_device_capability()[0] >= 8
+        ):
+            out = paddle.nn.functional.scaled_dot_product_attention(
+                query, key, value, **kwargs
+            )
+        else:
+            out, _ = _math_attention(
+                query,
+                key,
+                value,
+                causal=kwargs.get("is_causal", False),
+            )
+        return out
+
+
+class LlamaAttention(nn.Layer):
+    def __init__(self, config):
+        super().__init__()
+
+        self.config = config
+        self.hidden_size = self.config.hidden_size
+        self.num_heads = self.config.num_attention_heads
+        self.head_dim = self.hidden_size // self.config.num_attention_heads
+        self.sdpa = SDPALayer(config)
+
+        self.q_proj = nn.Linear(
+            self.hidden_size,
+            self.hidden_size,
+            bias_attr=True,
+        )
+
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.hidden_size,
+            bias_attr=True,
+        )
+
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.hidden_size,
+            bias_attr=True,
+        )
+
+        self.o_proj = nn.Linear(
+            self.hidden_size,
+            self.hidden_size,
+            bias_attr=True,
+        )
+
+    def forward(self, hidden_states):
+        query_states = self.q_proj(hidden_states).reshape(
+            shape=[0, 0, self.num_heads, self.head_dim]
+        )
+        key_states = self.k_proj(hidden_states).reshape(
+            shape=[0, 0, self.num_heads, self.head_dim]
+        )
+        value_states = self.v_proj(hidden_states).reshape(
+            shape=[0, 0, self.num_heads, self.head_dim]
+        )
+
+        bsz, q_len, _, _ = query_states.shape
+        outputs = self.sdpa(
+            query_states,
+            key_states,
+            value_states,
+            is_causal=True,
+        )
+        attn_output = outputs.reshape(
+            [-1, q_len, self.head_dim * self.num_heads]
+        )
+        attn_output = self.o_proj(attn_output)
+        return attn_output
+
+
+class LlamaMLP(nn.Layer):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = self.config.hidden_size
+        self.intermediate_size = self.config.intermediate_size
+
+        self.gate_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias_attr=False
+        )
+
+        self.up_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias_attr=False
+        )
+
+        self.down_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias_attr=False
+        )
+
+    def forward(self, x, test_for_list_input_output):
+        out = self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+        return out, test_for_list_input_output
+
+
+class LlamaRMSNorm(nn.Layer):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = self.config.hidden_size
+        self.weight = paddle.create_parameter(
+            shape=[self.hidden_size],
+            dtype=paddle.get_default_dtype(),
+            default_initializer=nn.initializer.Constant(1.0),
+        )
+        self.variance_epsilon = self.config.rms_norm_eps
+
+    def forward(self, hidden_states):
+        variance = hidden_states.astype("float32").pow(2).mean(-1, keepdim=True)
+        hidden_states = (
+            paddle.rsqrt(variance + self.variance_epsilon) * hidden_states
+        )
+
+        if self.weight.dtype in [paddle.float16, paddle.bfloat16]:
+            hidden_states = paddle.cast(hidden_states, self.weight.dtype)
+
+        return hidden_states * self.weight
+
+
+class LlamaDecoderLayer(nn.Layer):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.self_attn = LlamaAttention(self.config)
+        self.mlp = LlamaMLP(self.config)
+        self.input_layernorm = LlamaRMSNorm(self.config)
+        self.post_attention_layernorm = LlamaRMSNorm(self.config)
+
+    def forward(self, hidden_states, global_tensor):
+        residual = hidden_states + global_tensor
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states, _ = self.mlp(hidden_states, "ONLY_FOR_TEST")
+        hidden_states = residual + hidden_states
+
+        return (hidden_states,)
+
+
+class GlobalOutputNet(nn.Layer):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.config = config
+
+    def forward(self, input):
+        return (
+            input
+            if input is not None
+            else paddle.rand([self.config.hidden_size], dtype="float32")
+        )
+
+
+class LlamaModel(nn.Layer):
+    def __init__(self, config, position_embedding=False):
+        super().__init__()
+        self.config = config
+        self.vocab_size = self.config.vocab_size
+        self.hidden_size = self.config.hidden_size
+
+        self.embed_tokens = nn.Embedding(
+            self.vocab_size,
+            self.hidden_size,
+        )
+
+        self.position_embedding = (
+            nn.Embedding(
+                self.vocab_size,
+                self.hidden_size,
+            )
+            if position_embedding
+            else None
+        )
+
+        self.global_layer = GlobalOutputNet(self.config)
+
+        decoder_layers = []
+        for i in range(self.config.num_hidden_layers):
+            decoder_layers.append(LlamaDecoderLayer(self.config))
+
+        self.layers = nn.LayerList(decoder_layers)
+        self.norm = LlamaRMSNorm(self.config)
+
+    def forward(self, input_ids):
+        hidden_states = self.embed_tokens(input_ids)
+        if self.position_embedding is not None:
+            ones = paddle.ones(input_ids.shape, dtype="int64")
+            seq_length = paddle.cumsum(ones, axis=-1)
+            position_ids = seq_length - ones
+            position_embeddings = self.position_embedding(position_ids)
+            hidden_states = hidden_states + position_embeddings
+
+        global_tensor = self.global_layer(None)
+
+        for idx, (decoder_layer) in enumerate(self.layers):
+            tuple_hidden_states = decoder_layer(
+                hidden_states=hidden_states, global_tensor=global_tensor
+            )
+            hidden_states = tuple_hidden_states[0]
+
+        hidden_states = self.norm(hidden_states)
+
+        return hidden_states
+
+
+class LlamaLMHead(nn.Layer):
+    def __init__(self, config, weight=None):
+        super().__init__()
+        self.config = config
+        self.transpose_y = False
+        if weight is not None:
+            self.weight = weight
+            self.transpose_y = True
+        else:
+            self.weight = self.create_parameter(
+                shape=[self.config.hidden_size, self.config.vocab_size],
+                dtype=paddle.get_default_dtype(),
+            )
+
+    def forward(self, hidden_states):
+        logits = paddle.matmul(
+            hidden_states, self.weight, transpose_y=self.transpose_y
+        )
+        return logits
+
+
+class LlamaPretrainingCriterion(paddle.nn.Layer):
+    def __init__(self, config):
+        super().__init__()
+        self.ignore_index = getattr(config, "ignore_index", -100)
+        self.config = config
+        self.loss_func = paddle.nn.CrossEntropyLoss(
+            reduction="none", ignore_index=self.ignore_index
+        )
+
+    def forward(self, prediction_scores, masked_lm_labels):
+        if isinstance(prediction_scores, paddle.Tensor):
+            masked_lm_loss = self.loss_func(
+                prediction_scores.astype("float32")._use_gpudnn(False),
+                masked_lm_labels.unsqueeze(2),
+            )
+        else:
+            masked_lm_loss = self.loss_func(
+                prediction_scores.astype("float32"),
+                masked_lm_labels.unsqueeze(2),
+            )
+        if paddle.device.is_compiled_with_xpu():
+
+            def LocalLoss(x, mask):
+                masked_lm_loss = paddle.masked_select(x, mask).astype("float32")
+                loss = paddle.mean(masked_lm_loss).unsqueeze(0)
+                return loss.unsqueeze(0)
+
+            loss_func = dist.local_map(
+                LocalLoss,
+                [[dist.Shard(0), dist.Replicate()]],
+                [[dist.Shard(0), dist.Replicate()], None],
+                masked_lm_loss.process_mesh,
+                True,
+            )
+            loss = loss_func(masked_lm_loss, masked_lm_loss > 0)
+            loss = loss.mean()
+            return loss
+
+        masked_lm_loss = paddle.masked_select(
+            masked_lm_loss, masked_lm_loss > 0
+        ).astype("float32")
+        loss = paddle.mean(masked_lm_loss)
+        return loss
+
+
+class LlamaForCausalLM(nn.Layer):
+    enable_to_static_method = True
+
+    def __init__(self, config, share_embedding=False, position_embedding=False):
+        super().__init__()
+        self.config = config
+
+        self.llama = LlamaModel(self.config, position_embedding)
+        if share_embedding:
+            self.lm_head = LlamaLMHead(
+                self.config, self.llama.embed_tokens.weight
+            )
+        else:
+            self.lm_head = LlamaLMHead(self.config)
+
+    def forward(self, input_ids=None):
+        input_ids.stop_gradient = True
+
+        hidden_states = self.llama(input_ids)
+        logits = self.lm_head(hidden_states)
+
+        return logits

test/auto_parallel/hybrid_strategy/single_lora_model.py

@@ -0,0 +1,448 @@
+# 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 logging
+import math
+import re
+
+import numpy as np
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+
+
+class LoRALinear(nn.Linear):
+    # LoRA implemented in a dense layer
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        r: int = 0,
+        lora_alpha: int = 1,
+        lora_dropout: float = 0.0,
+        use_quick_lora: bool = False,
+        rslora: bool = False,
+        lora_plus_scale: float = 1.0,
+        pissa: bool = False,
+        lora_use_mixer: bool = False,
+        use_mora: bool = False,
+        **kwargs,
+    ):
+        nn.Linear.__init__(self, in_features, out_features, **kwargs)
+        if not isinstance(r, int) or r <= 0:
+            raise ValueError("Lora rank r should be a positive integer")
+        self.use_mora = use_mora
+        self.r = r
+        self.lora_alpha = lora_alpha
+        # Optional dropout
+        if lora_dropout > 0.0:
+            self.lora_dropout = nn.Dropout(p=lora_dropout)
+        else:
+            self.lora_dropout = lambda x: x
+        # Mark the weight as unmerged
+        self.merged = False
+        self.pissa = pissa
+        self.lora_use_mixer = lora_use_mixer
+
+        # Actual trainable parameters
+        if use_mora:  # reset the rank and create high rank matrix
+            self.in_features = in_features
+            self.out_features = out_features
+            new_r = int(math.sqrt((in_features + out_features) * r) + 0.5)
+            new_r = new_r // 2 * 2
+            self.r = new_r
+            self.lora_A = self.create_parameter(
+                shape=[self.r, self.r],
+                dtype=self._dtype,
+                is_bias=False,
+                default_initializer=nn.initializer.Constant(value=0.0),
+            )
+            self.cos = None
+            self.sin = None
+            # Count the number of tiles
+            self.rb1 = (
+                self.in_features // self.r
+                if self.in_features % self.r == 0
+                else self.in_features // self.r + 1
+            )
+            self.rb2 = (
+                self.out_features // self.r
+                if self.out_features % self.r == 0
+                else self.out_features // self.r + 1
+            )
+            self.rope_init()
+        else:
+            self.lora_A = self.create_parameter(
+                shape=[in_features, r],
+                dtype=self._dtype,
+                is_bias=False,
+            )
+            if self.lora_use_mixer:
+                self.lora_AB = self.create_parameter(
+                    shape=[r, r],
+                    dtype=self._dtype,
+                    is_bias=False,
+                )
+            self.lora_B = self.create_parameter(
+                shape=[r, out_features],
+                dtype=self._dtype,
+                is_bias=False,
+                attr=paddle.ParamAttr(
+                    initializer=paddle.nn.initializer.Constant(value=0.0),
+                    learning_rate=lora_plus_scale,
+                ),
+            )
+        self.apply_pissa = False
+        if use_mora or pissa:
+            self.scaling = 1.0
+        elif not rslora:
+            self.scaling = self.lora_alpha / self.r
+        else:
+            self.scaling = self.lora_alpha / math.sqrt(self.r)
+
+        # Freezing the pre-trained weight matrix
+        self.weight.stop_gradient = True
+        self._use_quick_lora = use_quick_lora and lora_dropout == 0.0
+        self.disable_lora = False
+
+    def pissa_init(self, rank):
+        weight = self.weight
+        dtype = weight.dtype
+        if dtype != paddle.float32:
+            weight = weight.astype(paddle.float32)
+
+        U, S, Vh = paddle.linalg.svd(weight.data, full_matrices=False)
+        Ur = U[:, :rank]
+        Sr = S[:rank]
+        Vhr = Vh[:rank]
+
+        lora_A = Ur @ paddle.diag(paddle.sqrt(Sr))
+        lora_B = paddle.diag(paddle.sqrt(Sr)) @ Vhr
+        self.lora_A.set_value(lora_A.astype(dtype))
+        self.lora_B.set_value(lora_B.astype(dtype))
+        res = weight.data - lora_A @ lora_B
+        weight = res.astype(dtype)
+        self.weight.set_value(weight)
+
+    def rope_init(self):
+        if self.cos is None or self.sin is None:
+            inv_freq = 1.0 / (
+                10000
+                ** (paddle.arange(0, self.r, 2, dtype=paddle.float32) / self.r)
+            )
+            t = paddle.arange(self.rb1, dtype=paddle.float32)
+            freqs = t.unsqueeze(1) @ inv_freq.unsqueeze(0)
+            emb = paddle.concat([freqs, freqs], axis=-1)
+            self.cos = paddle.unsqueeze(paddle.cos(emb), axis=0).astype(
+                self._dtype
+            )
+            self.sin = paddle.unsqueeze(paddle.sin(emb), axis=0).astype(
+                self._dtype
+            )
+
+    @property
+    def use_quick_lora(self):
+        return self._use_quick_lora and self.training and not self.merged
+
+    def _apply_mora(self, x):
+        r = self.r
+
+        # Calculate grouping
+        sum_inter = self.in_features // r
+
+        # padding
+        if self.in_features % r != 0:
+            pad_size = r - self.in_features % r
+            x = paddle.concat([x, x[..., :pad_size]], axis=-1)
+            sum_inter += 1
+
+        # reshape the input to apply RoPE
+        in_x = x.reshape([*x.shape[:-1], sum_inter, r])
+
+        # apply RoPE rotation
+        rh_in_x = paddle.concat(
+            [-in_x[..., r // 2 :], in_x[..., : r // 2]], axis=-1
+        )
+        in_x = in_x * self.cos + rh_in_x * self.sin
+
+        # matmul with high rank matrix
+        out_x = in_x @ self.lora_A
+
+        # reshape the output
+        out_x = out_x.reshape([*x.shape[:-1], -1])[..., : self.out_features]
+        if out_x.shape[-1] < self.out_features:
+            repeat_time = self.out_features // out_x.shape[-1]
+            if self.out_features % out_x.shape[-1] != 0:
+                repeat_time += 1
+            out_x = paddle.concat([out_x] * repeat_time, axis=-1)[
+                ..., : self.out_features
+            ]
+
+        return out_x
+
+    def get_delta_weight(self, lora_A=None, lora_B=None, lora_AB=None):
+        # compute the delta weight，which is used to merge weights
+        if self.lora_use_mixer:
+            lora_A = lora_A if lora_A is not None else self.lora_A
+            lora_B = lora_B if lora_B is not None else self.lora_B
+            lora_AB = lora_AB if lora_AB is not None else self.lora_AB
+            delta_weight = lora_A @ lora_AB @ lora_B * self.scaling
+        elif self.use_mora:
+            lora_A = lora_A if lora_A is not None else self.lora_A
+            r = self.r
+            # compute padding
+            pad_size = (
+                r - self.in_features % r if self.in_features % r != 0 else 0
+            )
+            # initialize weights
+            w = paddle.zeros(
+                [self.in_features + pad_size, self.in_features],
+                dtype=lora_A.dtype,
+            )
+
+            # create the weights after rotation
+            aw2 = paddle.concat(
+                [lora_A[:, r // 2 :], -lora_A[:, : r // 2]], axis=-1
+            )
+            # apply RoPE
+            for i in range(self.rb1 - 1):
+                w[i * r : (i + 1) * r, i * r : (i + 1) * r] = (
+                    aw2 * self.sin[:, i] + lora_A * self.cos[:, i]
+                )
+            # Process the last chunk that may be incomplete
+            i = self.rb1 - 1
+            w[i * r :, i * r :] = (
+                aw2 * self.sin[:, i] + lora_A * self.cos[:, i]
+            )[:, : r - pad_size]
+            # padding
+            if pad_size > 0:
+                w[i * r :, :pad_size] = (
+                    aw2 * self.sin[:, i] + lora_A * self.cos[:, i]
+                )[:, r - pad_size :]
+            # reshape the weights
+            if self.in_features < self.out_features:
+                w = paddle.concat([w] * self.rb2, axis=0)[: self.out_features]
+            else:
+                w = w[: self.out_features]
+            final_weight = w
+            delta_weight = final_weight.T
+        else:
+            lora_A = lora_A if lora_A is not None else self.lora_A
+            lora_B = lora_B if lora_B is not None else self.lora_B
+            delta_weight = lora_A @ lora_B * self.scaling
+
+        return delta_weight
+
+    def merge(self):
+        if not self.merged:
+            delta_weight = self.get_delta_weight()
+            new_weight = self.weight + delta_weight
+            self.weight.set_value(new_weight)
+            self.merged = True
+
+    def unmerge(self):
+        if self.merged:
+            delta_weight = self.get_delta_weight()
+            new_weight = self.weight - delta_weight
+            self.weight.set_value(new_weight)
+            self.merged = False
+
+    def forward(self, input: paddle.Tensor, *args, **kwargs):
+        if not self.apply_pissa and self.pissa:
+            self.pissa_init(self.r)
+            self.apply_pissa = True
+        if self.disable_lora or self.merged:
+            result = F.linear(
+                x=input, weight=self.weight, bias=self.bias, name=self.name
+            )
+        elif self.use_mora:
+            result = F.linear(
+                x=input, weight=self.weight, bias=self.bias, name=self.name
+            )
+            input = self.lora_dropout(input)
+            mora_out = self._apply_mora(input)
+            result += mora_out
+        else:
+            result = F.linear(
+                x=input, weight=self.weight, bias=self.bias, name=self.name
+            )
+            if self.lora_use_mixer:
+                result += (
+                    self.lora_dropout(input)
+                    @ self.lora_A
+                    @ self.lora_AB
+                    @ self.lora_B
+                ) * self.scaling
+            else:
+                result += (
+                    self.lora_dropout(input) @ self.lora_A @ self.lora_B
+                ) * self.scaling
+        return result
+
+    def extra_repr(self):
+        name = f", name={self.name}" if self.name else ""
+        return f"in_features={self.weight.shape[0]}, out_features={self.weight.shape[1]}, rank={self.r}{name}"
+
+
+lora_layers = {
+    "LoRALinear": LoRALinear,
+}
+LoRALinear = lora_layers["LoRALinear"]
+AVAILABLE_LAYERS = [
+    LoRALinear,
+]
+
+
+class LoRAModel(nn.Layer):
+    def __init__(self, model, lora_config) -> None:
+        super().__init__()
+        self.model = self.get_lora_model(model, lora_config)
+
+        self.lora_config = lora_config
+        logging.info("Mark only lora and trainable_module as trainable.")
+        self.mark_only_lora_as_trainable()
+
+    def forward(self, input_ids):
+        return self.model(input_ids)
+
+    def _find_and_replace_module(self, model, module_name, lora_config):
+        parent_module = model
+        attribute_chain = module_name.split(".")
+        for name in attribute_chain[:-1]:
+            parent_module = getattr(parent_module, name)
+        module = getattr(parent_module, attribute_chain[-1])
+        lora_module = None
+        if isinstance(module, nn.Linear):
+            lora_module = LoRALinear(
+                in_features=module.weight.shape[0],
+                out_features=module.weight.shape[1],
+                r=lora_config.r,
+                lora_alpha=lora_config.lora_alpha,
+                lora_dropout=lora_config.lora_dropout,
+                rslora=lora_config.rslora,
+                lora_plus_scale=lora_config.lora_plus_scale,
+                pissa=lora_config.pissa,
+                bias_attr=False if module.bias is None else None,
+                use_quick_lora=lora_config.use_quick_lora,
+                lora_use_mixer=lora_config.lora_use_mixer,
+                use_mora=lora_config.use_mora,
+            )
+        if lora_module is None:
+            raise ValueError(
+                f"LoRA strategy only supports paddle.nn.Linear or paddle.distributed.fleet.meta_parallel.ColumnParallelLinear or paddlenlp.transformers.sequence_utils. {module}({module_name} {type(module).__name__}) is not supported。"
+            )
+        lora_module.weight = module.weight
+        if module.bias is not None:
+            lora_module.bias = module.bias
+        setattr(parent_module, attribute_chain[-1], lora_module)
+
+    def print_trainable_parameters(self) -> None:
+        freeze_numel = 0
+        trainable_numel = 0
+        for _, weight in self.model.state_dict().items():
+            if weight.stop_gradient:
+                freeze_numel += np.prod(weight.shape)
+            else:
+                trainable_numel += np.prod(weight.shape)
+        logging.debug(
+            f"Frozen parameters: {freeze_numel:.2e} || Trainable parameters:{trainable_numel:.2e} || Total parameters:{freeze_numel + trainable_numel:.2e}|| Trainable:{trainable_numel / (freeze_numel + trainable_numel):.2%}"
+        )
+
+    def mark_only_lora_as_trainable(self) -> None:
+        for _, layer in self.model.named_sublayers():
+            if isinstance(layer, LoRALinear):
+                for name, weight in layer.state_dict().items():
+                    if (
+                        self.lora_config.trainable_bias in ["lora", "all"]
+                        and "bias" in name
+                    ):
+                        weight.stop_gradient = False
+                    elif "lora" in name:
+                        weight.stop_gradient = False
+                    else:
+                        weight.stop_gradient = True
+            else:
+                for name, weight in layer.state_dict().items():
+                    if (
+                        self.lora_config.trainable_bias == "all"
+                        and "bias" in name
+                    ):
+                        weight.stop_gradient = False
+                    else:
+                        weight.stop_gradient = True
+        if self.lora_config.trainable_modules is not None:
+            for name, weight in self.model.state_dict().items():
+                if any(
+                    re.fullmatch(trainable_module, name)
+                    for trainable_module in self.lora_config.trainable_modules
+                ):
+                    weight.stop_gradient = False
+
+    def get_lora_model(self, model, lora_config):
+        if lora_config.target_modules is None:
+            return model
+        elif isinstance(lora_config.target_modules, str):
+            target_modules = [lora_config.target_modules]
+        else:
+            target_modules = lora_config.target_modules
+        for target_module in target_modules:
+            for i in model.named_sublayers():
+                module_name = i[0]
+                if re.fullmatch(target_module, module_name):
+                    self._find_and_replace_module(
+                        model, module_name, lora_config
+                    )
+        return model
+
+    def train(self):
+        self.training = True
+        self.model.training = True
+        for layer in self.model.sublayers():
+            layer.training = True
+            layer.train()
+
+    def eval(self):
+        self.training = False
+        self.model.training = False
+        for layer in self.model.sublayers():
+            layer.training = False
+            layer.eval()
+
+    def disable_lora(self):
+        for _, layer in self.model.named_sublayers():
+            if any(
+                isinstance(layer, lora_layer) for lora_layer in AVAILABLE_LAYERS
+            ):
+                layer.disable_lora = True
+
+    def enable_lora(self):
+        for _, layer in self.model.named_sublayers():
+            if any(
+                isinstance(layer, lora_layer) for lora_layer in AVAILABLE_LAYERS
+            ):
+                layer.disable_lora = False
+
+    def merge(self):
+        for _, layer in self.model.named_sublayers():
+            if any(
+                isinstance(layer, lora_layer) for lora_layer in AVAILABLE_LAYERS
+            ):
+                layer.merge()
+
+    def unmerge(self):
+        for _, layer in self.model.named_sublayers():
+            if any(
+                isinstance(layer, lora_layer) for lora_layer in AVAILABLE_LAYERS
+            ):
+                layer.unmerge()

test/auto_parallel/hybrid_strategy/test_cross_mesh_reshard.py

@@ -0,0 +1,67 @@
+# 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 collective.test_communication_api_base as test_base
+
+
+class TestSemiAutoParallelCrossMeshReshard(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=4,
+            timeout=120,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_simple_net_cross_mesh_reshard(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_cross_mesh_reshard.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestSemiAutoParallelNdCrossMeshReshard(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=200, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_simple_net_hybrid_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_nd_cross_mesh_reshard.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_flexcheckpoint_merge.py

@@ -0,0 +1,43 @@
+# 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 tempfile
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestDistCheckpointMerge(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=50, nnode=1)
+        self._default_envs = {}
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_merge_checkpoint(self):
+        ckpt_path = tempfile.TemporaryDirectory()
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "semi_flexcheckpoint_merge.py",
+                user_defined_envs=envs,
+            )
+
+        ckpt_path.cleanup()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_get_group_in_different_hybrid_configs.py

@@ -0,0 +1,162 @@
+# 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 TestProcessMeshDPGroupConsistency(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=200, nnode=1)
+
+    def test_dp_parallel(self):
+        """Test data parallel group creation and consistency"""
+        _default_envs = {
+            "dp": "2",
+            "mp": "1",
+            "pp": "1",
+            "parallel_type": "dp",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "test_process_mesh_group_consistency.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestProcessMeshMPGroupConsistency(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=200, nnode=1)
+
+    def test_mp_parallel(self):
+        """Test model parallel group creation and consistency"""
+        _default_envs = {
+            "dp": "1",
+            "mp": "2",
+            "pp": "1",
+            "parallel_type": "mp",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "test_process_mesh_group_consistency.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestProcessMeshPPGroupConsistency(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=200, nnode=1)
+
+    def test_pp_parallel(self):
+        """Test pipeline parallel group creation and consistency"""
+        _default_envs = {
+            "dp": "1",
+            "mp": "1",
+            "pp": "2",
+            "parallel_type": "pp",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "test_process_mesh_group_consistency.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestProcessMeshSEPGroupConsistency(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=200, nnode=1)
+
+    def test_sep_parallel(self):
+        """Test sequence parallel group creation and consistency"""
+        _default_envs = {
+            "dp": "1",
+            "mp": "1",
+            "pp": "1",
+            "sep": "2",
+            "sharding": "1",
+            "parallel_type": "sep",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "test_process_mesh_group_consistency.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestProcessMeshShardingGroupConsistency(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=200, nnode=1)
+
+    def test_sharding_parallel(self):
+        """Test sharding parallel group creation and consistency"""
+        _default_envs = {
+            "dp": "1",
+            "mp": "1",
+            "pp": "1",
+            "sep": "1",
+            "sharding": "2",
+            "parallel_type": "sharding",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "test_process_mesh_group_consistency.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()  # python run

test/auto_parallel/hybrid_strategy/test_global_mesh_reshard.py

@@ -0,0 +1,73 @@
+# 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 collective.test_communication_api_base as test_base
+
+
+class TestSemiAutoParallel2DGlobalMeshReshard(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=4,
+            timeout=120,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_2d_global_mesh_reshard(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_2d_global_mesh_reshard.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestSemiAutoParallel3DGlobalMeshReshard(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=8,
+            timeout=120,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_3d_global_mesh_reshard(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_3d_global_mesh_reshard.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_1d.py

@@ -0,0 +1,154 @@
+# 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 tempfile
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestShardingParallelAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "2",
+            "mp": "1",
+            "pp": "1",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["False"],
+            "sharding_stage": ["0", "1"],
+            "test_share_embedding": [
+                "1",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
+        }
+
+    def test_simple_net_dp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+class TestPipelineParallelAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "1",
+            "mp": "1",
+            "pp": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["False"],
+            "num_hidden_layers": ["3", "4"],
+            "test_share_embedding": [
+                "1",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
+        }
+
+    def test_simple_net_pp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+class TestTensorParallelAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "1",
+            "mp": "2",
+            "pp": "1",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true", "false"],
+            "sequence_parallel": ["true", "false"],
+            "prepare_input_output": ["true", "false"],
+            "test_share_embedding": [
+                "1",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
+        }
+
+    def test_simple_net_mp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        has_checked_lazy_init = False
+        has_checked_pure_mp = False
+        for envs in envs_list:
+            if envs['use_lazy_init'] == 'true':
+                if has_checked_lazy_init:
+                    continue
+                has_checked_lazy_init = True
+            if envs['sequence_parallel'] != 'true':
+                if has_checked_pure_mp:
+                    continue
+                has_checked_pure_mp = True
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_2d.py

@@ -0,0 +1,152 @@
+# 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 tempfile
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestMPPPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "1",
+            "mp": "2",
+            "pp": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true"],
+            "sequence_parallel": ["true"],
+            "prepare_input_output": ["false"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
+            "one_api": ["true", "false"],
+        }
+
+    def test_simple_net_mp2_pp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+class TestDPPPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "2",
+            "mp": "1",
+            "pp": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true"],
+            "num_hidden_layers": ["3", "4"],
+            "sharding_stage": ["0"],
+            "test_share_embedding": [
+                "1",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
+            "one_api": ["true", "false"],
+        }
+
+    def test_simple_net_dp2_pp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+class TestDPMPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "2",
+            "mp": "2",
+            "pp": "1",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true"],
+            "sequence_parallel": ["true"],
+            "prepare_input_output": ["false"],
+            "sharding_stage": ["0"],
+            "test_share_embedding": [
+                "1",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
+            "one_api": ["true", "false"],
+        }
+
+    def test_simple_net_dp2_tp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_2d_sep.py

@@ -0,0 +1,115 @@
+# 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 tempfile
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestDPMPCPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "1",
+            "mp": "2",
+            "pp": "1",
+            "sep": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true"],
+            "sequence_parallel": ["false"],
+            "context_parallel": ["true"],
+            "prepare_input_output": ["false"],
+            "sharding_stage": ["0"],
+            "test_share_embedding": [
+                "1",
+            ],
+            "test_position_embedding": [
+                "0",
+            ],
+            "one_api": ["true", "false"],
+        }
+
+    def test_simple_net_mp2_cp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+class TestDPMPSEPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "1",
+            "mp": "1",
+            "pp": "2",
+            "sep": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true"],
+            "sequence_parallel": ["false"],
+            "sep_parallel": ["true"],
+            "context_parallel": ["false"],
+            "prepare_input_output": ["false"],
+            "sharding_stage": ["0"],
+            "test_share_embedding": [
+                "1",
+            ],
+            "test_position_embedding": [
+                "0",
+            ],
+            "one_api": ["true", "false"],
+        }
+
+    def test_simple_net_mp2_sep2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_3d.py

@@ -0,0 +1,66 @@
+# 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 tempfile
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestDPMPPPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=180, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "2",
+            "mp": "2",
+            "pp": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["true"],
+            "sequence_parallel": ["true"],
+            "prepare_input_output": ["false"],
+            "sharding_stage": ["0", "1"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
+            "one_api": ["true", "false"],
+        }
+
+    def test_simple_net_dp2_mp2_pp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_parallel_api_with_llama_lora.py

@@ -0,0 +1,67 @@
+# 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 tempfile
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestDPMPPPAPI(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=360, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "dp": "2",
+            "mp": "2",
+            "pp": "2",
+            "acc_step": "2",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["true"],
+            "use_lazy_init": ["false"],
+            "sequence_parallel": ["false"],
+            "prepare_input_output": ["false"],
+            "sharding_stage": ["0", "1"],
+            "test_share_embedding": [
+                "0",
+            ],
+            "test_position_embedding": [
+                "1",
+            ],
+            "one_api": ["true", "false"],
+            "test_lora": ["1"],
+        }
+
+    def test_simple_lora_net_dp2_mp2_pp2(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "parallel_api.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_pir_reshard_nd_mesh_func.py

@@ -0,0 +1,61 @@
+# 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 collective.test_communication_api_base as test_base
+
+
+class TestPIRNdMeshReshard(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=4,
+            timeout=20,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "backend": "gpu",
+        }
+
+    def test_simple_net_reshard(self):
+        self.run_test_case(
+            "pir_reshard_nd_mesh.py",
+            user_defined_envs=self._default_envs,
+        )
+
+
+class TestPIRNdMeshReshardCrossMesh(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=8,
+            timeout=20,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "backend": "gpu",
+        }
+
+    def test_simple_net_reshard_cross_mesh(self):
+        self.run_test_case(
+            "pir_reshard_nd_mesh_cross_mesh.py",
+            user_defined_envs=self._default_envs,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/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()

test/auto_parallel/hybrid_strategy/test_process_mesh_group_consistency.py

@@ -0,0 +1,107 @@
+# 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 os
+
+import paddle.distributed as dist
+from paddle.distributed import fleet
+
+
+class TestProcessMeshGroupConsistency:
+    def __init__(self):
+        # Get configuration from environment variables
+        self.dp = int(os.getenv("dp", "1"))
+        self.mp = int(os.getenv("mp", "1"))
+        self.pp = int(os.getenv("pp", "1"))
+        self.sep = int(os.getenv("sep", "1"))
+        self.sharding = int(os.getenv("sharding", "1"))
+
+        # Determine which parallel type to test
+        self.parallel_type = os.getenv("parallel_type", "dp")
+
+    def init_dist_env(self):
+        """Initialize distributed environment"""
+        # Configure distributed strategy
+        dist_strategy = fleet.DistributedStrategy()
+        dist_strategy.hybrid_configs = {
+            "dp_degree": self.dp,
+            "mp_degree": self.mp,
+            "pp_degree": self.pp,
+            "sep_degree": self.sep,
+            "sharding_degree": self.sharding,
+        }
+
+        # Add corresponding configuration based on parallel type
+        if self.sep > 1:
+            dist_strategy.hybrid_configs["sep_degree"] = self.sep
+        if self.sharding > 1:
+            dist_strategy.hybrid_configs["sharding_degree"] = self.sharding
+
+        fleet.init(is_collective=True, strategy=dist_strategy)
+
+    def test_process_mesh_group_consistency(self):
+        """Test consistency between ProcessMesh created groups and HCG created groups"""
+
+        # Create corresponding ProcessMesh and get corresponding HCG group based on parallel type
+        if self.parallel_type == "dp":
+            mesh = dist.ProcessMesh([0, 1], dim_names=["dp"])
+            hcg = fleet.get_hybrid_communicate_group()
+            group = mesh.get_group(dim_name="dp")
+            hcg_group = hcg.get_data_parallel_group()
+
+        elif self.parallel_type == "mp":
+            mesh = dist.ProcessMesh([0, 1], dim_names=["mp"])
+            hcg = fleet.get_hybrid_communicate_group()
+            group = mesh.get_group(dim_name="mp")
+            hcg_group = hcg.get_model_parallel_group()
+
+        elif self.parallel_type == "pp":
+            mesh = dist.ProcessMesh([0, 1], dim_names=["pp"])
+            hcg = fleet.get_hybrid_communicate_group()
+            group = mesh.get_group(dim_name="pp")
+            hcg_group = hcg.get_pipe_parallel_group()
+
+        elif self.parallel_type == "sep":
+            mesh = dist.ProcessMesh([0, 1], dim_names=["sep"])
+            hcg = fleet.get_hybrid_communicate_group()
+            group = mesh.get_group(dim_name="sep")
+            hcg_group = hcg.get_sep_parallel_group()
+
+        elif self.parallel_type == "sharding":
+            mesh = dist.ProcessMesh([0, 1], dim_names=["sharding"])
+            hcg = fleet.get_hybrid_communicate_group()
+            group = mesh.get_group(dim_name="sharding")
+            hcg_group = hcg.get_sharding_parallel_group()
+
+        else:
+            raise ValueError(f"Unsupported parallel type: {self.parallel_type}")
+
+        # Verify that group ranks are consistent
+        group_ranks = group.ranks
+        hcg_group_ranks = hcg_group.ranks
+        assert set(group_ranks) == set(hcg_group_ranks)
+
+        # Verify that group IDs are consistent
+        group_id = group.id
+        hcg_group_id = hcg_group.id
+        assert group_id == hcg_group_id
+
+    def run_test_cases(self):
+        """Run test cases"""
+        self.init_dist_env()
+        self.test_process_mesh_group_consistency()
+
+
+if __name__ == "__main__":
+    TestProcessMeshGroupConsistency().run_test_cases()

test/auto_parallel/hybrid_strategy/test_save_load_state_dict.py

@@ -0,0 +1,183 @@
+# 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 tempfile
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+# should be set to FLAGS_enable_pir_api=1
+os.environ['FLAGS_enable_pir_api'] = '0'
+
+
+class TestSaveLoadStateDict(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        self._default_envs = {}
+        self._changeable_envs = {"device_num": ["1", "2", "4", "8"]}
+
+    def test_reshard(self):
+        # save with 1 device
+        ckpt_path = tempfile.TemporaryDirectory()
+        ckpt_path_2 = tempfile.TemporaryDirectory()
+        ckpt_path_3 = tempfile.TemporaryDirectory()
+        super().setUp(num_of_devices=1, timeout=120, nnode=1)
+        self.run_test_case(
+            "semi_auto_save_state_dict.py",
+            user_defined_envs={
+                "device_num": "1",
+                "ckpt_path": ckpt_path.name,
+                "ckpt_path_2": ckpt_path_2.name,
+                "ckpt_path_3": ckpt_path_3.name,
+            },
+        )
+
+        # load with 1, 2, 4, 8 devices
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            envs["ckpt_path"] = ckpt_path.name
+            envs["ckpt_path_2"] = ckpt_path_2.name
+            envs["ckpt_path_3"] = ckpt_path_3.name
+            super().setUp(
+                num_of_devices=int(envs["device_num"]),
+                timeout=180,
+                nnode=1,
+            )
+            self.run_test_case(
+                "semi_auto_load_state_dict.py",
+                user_defined_envs=envs,
+            )
+        ckpt_path.cleanup()
+        ckpt_path_2.cleanup()
+        ckpt_path_3.cleanup()
+
+        # save with 2 devices
+        ckpt_path = tempfile.TemporaryDirectory()
+        ckpt_path_2 = tempfile.TemporaryDirectory()
+        ckpt_path_3 = tempfile.TemporaryDirectory()
+        super().setUp(num_of_devices=2, timeout=120, nnode=1)
+        self.run_test_case(
+            "semi_auto_save_state_dict.py",
+            user_defined_envs={
+                "device_num": "2",
+                "ckpt_path": ckpt_path.name,
+                "ckpt_path_2": ckpt_path_2.name,
+                "ckpt_path_3": ckpt_path_3.name,
+            },
+        )
+        # load with 1, 2, 4, 8 devices
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            envs["ckpt_path"] = ckpt_path.name
+            envs["ckpt_path_2"] = ckpt_path_2.name
+            envs["ckpt_path_3"] = ckpt_path_3.name
+            super().setUp(
+                num_of_devices=int(envs["device_num"]),
+                timeout=180,
+                nnode=1,
+            )
+            self.run_test_case(
+                "semi_auto_load_state_dict.py",
+                user_defined_envs=envs,
+            )
+        ckpt_path.cleanup()
+        ckpt_path_2.cleanup()
+        ckpt_path_3.cleanup()
+
+    def test_mutual_load_between_dynamic_and_static(self):
+        changeable_envs = {"device_num": ["2"]}
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, changeable_envs
+        )
+
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            super().setUp(
+                num_of_devices=int(envs["device_num"]),
+                timeout=180,
+                nnode=1,
+            )
+            self.run_test_case(
+                "semi_auto_parallel_mutual_load_between_dynamic_and_static.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+    def test_save_safetensors_load_fc(self):
+        """Test saving safetensors files and loading with flex checkpoint."""
+        ckpt_path = tempfile.TemporaryDirectory()
+        super().setUp(num_of_devices=2, timeout=120, nnode=1)
+        self.run_test_case(
+            "save_safetensors_load_fc.py",
+            user_defined_envs={
+                "device_num": "2",
+                "ckpt_path": ckpt_path.name,
+            },
+        )
+        ckpt_path.cleanup()
+
+    def test_save_safetensors_load_fc_with_index(self):
+        """Test saving safetensors files and loading with flex checkpoint when model.safetensors.index.json exists."""
+        ckpt_path = tempfile.TemporaryDirectory()
+        super().setUp(num_of_devices=2, timeout=120, nnode=1)
+        self.run_test_case(
+            "save_safetensors_load_fc.py",
+            user_defined_envs={
+                "device_num": "2",
+                "ckpt_path": ckpt_path.name,
+                "test_func": "test_save_safetensors_load_fc_with_index",
+            },
+        )
+        ckpt_path.cleanup()
+
+    def test_save_load_state_dict_with_aoa_config_reverse(self):
+        """Test saving state dict and loading with flex checkpoint."""
+        ckpt_path = tempfile.TemporaryDirectory()
+        super().setUp(num_of_devices=1, timeout=60, nnode=1)
+        self.run_test_case(
+            "save_load_state_dict_with_aoa_config_reverse.py",
+            user_defined_envs={
+                "device_num": "1",
+                "ckpt_path": ckpt_path.name,
+            },
+        )
+        ckpt_path.cleanup()
+
+    def test_save_load_with_error_message(self):
+        """Test logger missing key and unexpected keys."""
+        ckpt_path = tempfile.TemporaryDirectory()
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, {"device_num": ["1", "2"]}
+        )
+        for envs in envs_list:
+            envs["ckpt_path"] = ckpt_path.name
+            super().setUp(
+                num_of_devices=int(envs["device_num"]),
+                timeout=60,
+                nnode=1,
+            )
+            self.run_test_case(
+                "test_save_load_with_error_message.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+if __name__ == '__main__':
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_save_load_with_error_message.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2026 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 paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
+    load_state_dict,
+)
+
+
+class HuggingFaceModel(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.huggingface = nn.Linear(2, 2, bias_attr=False)
+
+
+class FCModel(nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(1, 2, bias_attr=False)
+        self.fc2 = nn.Linear(1, 2, bias_attr=False)
+
+
+def init_hf_model_weights(model):
+    with paddle.no_grad():
+        w = paddle.to_tensor([[0, 1], [2, 3]], dtype="float16")
+        model.huggingface.weight.set_value(w)
+
+
+def save_safetensors_model(model, ckpt_path):
+    import safetensors.numpy
+
+    os.makedirs(ckpt_path, exist_ok=True)
+    weight_np = model.huggingface.weight.numpy()
+    file_path = os.path.join(ckpt_path, "tensor1.safetensors")
+    safetensors.numpy.save_file({"huggingface.weight": weight_np}, file_path)
+
+
+def test_save_load_with_missing_key_and_unexpected_keys():
+    # ckpt_path = os.getenv("ckpt_path")
+    ckpt_path = "/home/paddle/test/auto_parallel/hybrid_strategy/test_file"
+    dist.init_parallel_env()
+
+    hf_model = HuggingFaceModel()
+    fc_model = FCModel()
+    hf_model = paddle.amp.decorate(
+        models=hf_model, optimizers=None, level="O2", dtype="float16"
+    )
+    init_hf_model_weights(hf_model)
+
+    save_safetensors_model(hf_model, ckpt_path)
+
+    aoa_statements = []
+    aoa_config = {"aoa_statements": aoa_statements}
+
+    try:
+        load_state_dict(
+            fc_model.sharded_state_dict(),
+            ckpt_path,
+            safetensors=True,
+            aoa_config=aoa_config,
+        )
+        raise AssertionError
+    except Exception as e:
+        pass
+
+
+def test_save_load_with_mapping_key_to_safetensors_file():
+    ckpt_path = os.getenv("ckpt_path")
+    dist.init_parallel_env()
+
+    hf_model = HuggingFaceModel()
+    fc_model = FCModel()
+    hf_model = paddle.amp.decorate(
+        models=hf_model, optimizers=None, level="O2", dtype="float16"
+    )
+    init_hf_model_weights(hf_model)
+
+    save_safetensors_model(hf_model, ckpt_path)
+    import json
+
+    # set the key to the wrong safetensors file tensor2.safetensors, the fact file is tensor1.safetensors
+    index = {
+        "metadata": {"total_size": 8},
+        "weight_map": {"huggingface.weight": "tensor2.safetensors"},
+    }
+    with open(
+        os.path.join(ckpt_path, "model.safetensors.index.json"), "w"
+    ) as f:
+        json.dump(index, f)
+
+    aoa_statements = [
+        "huggingface.weight -> A,B ,axis = 1 \n",
+        "A^T -> A  \n",
+        "B^T -> B \n",
+        "A -> fc1.weight ,src_dtype = 'float16', dst_dtype = 'float32' \n",
+        "B -> fc2.weight ,src_dtype = 'float16', dst_dtype = 'float32' \n",
+    ]
+    aoa_config = {"aoa_statements": aoa_statements}
+
+    try:
+        load_state_dict(
+            fc_model.sharded_state_dict(),
+            ckpt_path,
+            safetensors=True,
+            aoa_config=aoa_config,
+        )
+        raise AssertionError
+    except Exception as e:
+        pass
+
+
+def test():
+    test_save_load_with_missing_key_and_unexpected_keys()
+    test_save_load_with_mapping_key_to_safetensors_file()
+
+
+if __name__ == "__main__":
+    test()

test/auto_parallel/hybrid_strategy/test_semi_auto_llama_acc_align.py

@@ -0,0 +1,140 @@
+# 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 collective.test_communication_api_base as test_base
+
+
+class TestSemiAutoParallelLlamaACCTest(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=200, nnode=1)
+
+    def test_simple_net_hybrid_strategy_acc(self):
+        _default_envs = {
+            "dp": "2",
+            "mp": "2",
+            "pp": "2",
+            "acc_step": "1",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_llama_acc_align.py",
+                user_defined_envs=envs,
+            )
+
+    def test_simple_net_hybrid_strategy_acc_grad_merge(self):
+        _default_envs = {
+            "dp": "2",
+            "mp": "2",
+            "pp": "2",
+            "acc_step": "2",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_llama_acc_align.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestSemiAutoParallelLlamaCPTest(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=200, nnode=1)
+
+    def test_dp2pp2cp2_hybrid_strategy_acc(self):
+        _default_envs = {
+            "dp": "2",
+            "mp": "2",
+            "pp": "1",
+            "sep": "2",
+            "acc_step": "1",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": ["bfloat16"],
+            "amp_master_grad": ["false"],
+            "seq_length": ["1024"],
+            "hidden_size": ["2048"],
+            "num_attention_heads": ["8"],
+            "num_key_value_heads": ["8"],
+            "context_parallel": ["true"],
+            "max_position_embeddings": ["2048"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_llama_acc_align.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestSemiAutoParallelLlamaSEPTest(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=200, nnode=1)
+
+    def test_dp2pp2sep2_hybrid_strategy_acc(self):
+        # now not support mp with sep
+        # TODO: something wrong with this case
+        _default_envs = {
+            "dp": "4",
+            "mp": "1",
+            "pp": "1",
+            "sep": "2",
+            "acc_step": "1",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
+        }
+        _changeable_envs = {
+            "backend": ["gpu"],
+            "sep_parallel": ["true"],
+            "context_parallel": ["false"],
+        }
+        envs_list = test_base.gen_product_envs_list(
+            _default_envs, _changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_llama_acc_align.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()  # python run

test/auto_parallel/hybrid_strategy/test_semi_auto_llama_save_load.py

@@ -0,0 +1,50 @@
+# 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 collective.test_communication_api_base as test_base
+
+
+class TestSemiAutoParallelLlamaSaveLoadTest(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=200, nnode=1)
+        self._default_envs = {
+            "dp": "2",
+            "mp": "2",
+            "pp": "2",
+            "acc_step": "1",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+        }
+
+    def test_simple_net_hybrid_strategy_save_load(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_llama_save_load.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()  # python run

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_c_cross_entropy.py

@@ -0,0 +1,58 @@
+# 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
+
+os.environ["PARALLEL_CROSS_ENTROPY"] = "true"
+
+import os
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+os.environ['FLAGS_enable_pir_api'] = '0'
+
+
+class TestParallelCrossEntropy(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=200, nnode=1)
+
+    def test_dp(self):
+        self.run_test_case(
+            "semi_auto_parallel_c_cross_entropy_dp.py",
+        )
+
+    def test_mp(self):
+        self.run_test_case(
+            "semi_auto_parallel_c_cross_entropy_mp.py",
+        )
+
+    def test_mp_pir(self):
+        os.environ["FLAGS_enable_pir_in_executor"] = "True"
+        self.test_mp()
+        os.environ["FLAGS_enable_pir_in_executor"] = "False"
+
+
+class TestParallelCrossEntropyHybrid(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=200, nnode=1)
+
+    def test_hybrid(self):
+        self.run_test_case(
+            "semi_auto_parallel_c_cross_entropy_hybrid.py",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_global_input.py

@@ -0,0 +1,60 @@
+# 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 os
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+os.environ['FLAGS_enable_pir_api'] = '1'
+
+
+class TestSemiAutoParallelGlobalInput(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=8,
+            timeout=120,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "1024",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_dynamic(self):
+        self._default_envs.update({"run_static": "0"})
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_global_input.py",
+                user_defined_envs=envs,
+            )
+
+    def test_static(self):
+        self._default_envs.update({"run_static": "1"})
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_global_input.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_hybrid_sharding_strategy.py

@@ -0,0 +1,69 @@
+# 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 os
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+os.environ['FLAGS_enable_pir_api'] = '0'
+
+
+class TestSemiAutoParallelInShardingStrategy(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=4,
+            timeout=120,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_sharding_stage_1_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_sharding_stage_1.py",
+                user_defined_envs=envs,
+            )
+
+    def test_sharding_stage_2_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_sharding_stage_2.py",
+                user_defined_envs=envs,
+            )
+
+    def test_sharding_stage_3_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_sharding_stage_3.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_hybrid_strategy.py

@@ -0,0 +1,133 @@
+# 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 tempfile
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+os.environ['FLAGS_enable_pir_api'] = '0'
+
+
+class TestSemiAutoParallelDPMPStrategy(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=120, nnode=1)
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_simple_net_hybrid_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "semi_auto_parallel_simple_net_dp_mp.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+    def test_fused_linear_param_grad_add(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_for_fused_linear_param_grad_add.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestSemiAutoParallelHybridStrategy(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=8,
+            timeout=120,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_simple_net_hybrid_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "semi_auto_parallel_simple_net_dp_mp_pp.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+class TestSemiAutoParallelHybridStrategyWithSP(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=4,
+            timeout=120,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["gpu"], "is_dp": ["false"]}
+
+    def test_simple_net_mp_pp_sp(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "semi_auto_parallel_simple_net_sp.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+    def test_simple_net_dp_mp_pp_sp(self):
+        super().setUp(
+            num_of_devices=8,
+            timeout=120,
+            nnode=1,
+        )
+        self._changeable_envs = {"backend": ["gpu"], "is_dp": ["true"]}
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            ckpt_path = tempfile.TemporaryDirectory()
+            envs["ckpt_path"] = ckpt_path.name
+            self.run_test_case(
+                "semi_auto_parallel_simple_net_sp.py",
+                user_defined_envs=envs,
+            )
+            ckpt_path.cleanup()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_llama_model_amp.py

@@ -0,0 +1,85 @@
+# 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 collective.test_communication_api_base as test_base
+
+
+class TestSemiAutoParallelLlama3DAMPTest(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=200, nnode=1)
+        self._default_envs = {"dp": "2", "mp": "2", "pp": "2", "acc_step": "2"}
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "use_sp": ["true"],
+            "use_param_group": ["true"],
+            "recompute": ["true"],
+            "recompute_granularity": ["full"],
+            "amp": ["true"],
+            "amp_level": ["O1"],
+            "amp_dtype": [
+                "float16",
+            ],
+            "amp_master_grad": [
+                "False",
+            ],
+        }
+
+    def test_simple_net_hybrid_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_llama.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestSemiAutoParallelLlama3DAMPMasterGradTest(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=200, nnode=1)
+        self._default_envs = {"dp": "2", "mp": "2", "pp": "2", "acc_step": "2"}
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "use_sp": ["true"],
+            "use_param_group": ["true"],
+            "recompute": ["true"],
+            "recompute_granularity": ["full"],
+            "amp": ["true"],
+            "amp_level": ["O2"],
+            "amp_dtype": [
+                "float16",
+            ],
+            "amp_master_grad": [
+                "True",
+            ],
+        }
+
+    def test_simple_net_hybrid_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_llama.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_llama_model_pir.py

@@ -0,0 +1,48 @@
+# 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 collective.test_communication_api_base as test_base
+
+
+class Test3DSemiAutoParallelStaticPirDecorate(
+    test_base.CommunicationTestDistBase
+):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=8,
+            timeout=300,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "FLAGS_enable_pir_api": "1",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_mlp(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        # self._log_dir.name = "./log"
+        for envs in envs_list:
+            self.run_test_case(
+                "../pir/mlp_demo_3d.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_llama_model_vpp.py

@@ -0,0 +1,64 @@
+# 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 sys
+import unittest
+
+sys.path.append("../../")
+import os
+
+import collective.test_communication_api_base as test_base
+
+os.environ['FLAGS_enable_pir_api'] = '1'
+
+
+class TestSemiAutoParallelLlama3DVPP(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=8, timeout=200, nnode=1)
+        self._default_envs = {
+            "seed": "2023",
+            "dp": "2",
+            "mp": "2",
+            "pp": "2",
+            "FLAGS_embedding_deterministic": "1",
+            "FLAGS_cudnn_deterministic": "1",
+            "acc_step": "4",
+            "only_static": "true",
+        }
+        self._changeable_envs = {
+            "backend": ["gpu"],
+            "use_sp": ["true"],
+            "use_param_group": ["true"],
+            "recompute": ["true"],
+            "recompute_granularity": ["full"],
+            # TODO: Temporarily turn off the vpp test in PIR mode. There will be
+            # a hang issue, which will be fixed later.
+            # "virtual_pp_degree": ["2"],
+            "virtual_pp_degree": ["1"],
+        }
+
+    def test_simple_net_hybrid_strategy(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        self._log_dir.name = "./vpp_log"
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_llama_pp_gradmerge.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_multi_inputs.py

@@ -0,0 +1,60 @@
+# 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 os
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+os.environ['FLAGS_enable_pir_api'] = '1'
+
+
+class TestSemiAutoParallelMultiInputs(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=8,
+            timeout=120,
+            nnode=1,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "1024",
+        }
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_dynamic(self):
+        self._default_envs.update({"run_static": "0"})
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_multi_inputs.py",
+                user_defined_envs=envs,
+            )
+
+    def test_static(self):
+        self._default_envs.update({"run_static": "1"})
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_multi_inputs.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/test_to_distributed_api_for_llama.py

@@ -0,0 +1,50 @@
+# 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 collective.test_communication_api_base as test_base
+
+import paddle
+
+
+class TestToDistributedApiForLlamaBasic(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        self._num_of_devices = 8
+        super().setUp(
+            num_of_devices=self._num_of_devices,
+            timeout=120,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+            "num_of_devices": self._num_of_devices,
+        }
+        self._changeable_envs = {"backend": ["cpu", "gpu"]}
+
+    def test_llama(self):
+        envs_list = test_base.gen_product_envs_list(
+            {"dtype": "float32", "seed": "2023"}, {"backend": ["gpu"]}
+        )
+        cuda_version_main = int(paddle.version.cuda().split(".")[0])
+        device_prop_main = paddle.device.cuda.get_device_capability()[0]
+        if cuda_version_main >= 11 and device_prop_main >= 8:
+            for envs in envs_list:
+                self.run_test_case(
+                    "to_distributed_api_for_llama.py",
+                    user_defined_envs=envs,
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/auto_parallel/hybrid_strategy/testslist.csv

@@ -0,0 +1,24 @@
+name,os,arch,timeout,run_type,launcher,num_port,run_serial,envs,conditions
+test_semi_auto_parallel_hybrid_strategy,LINUX,GPU,300,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_save_load_state_dict,LINUX,GPU,400,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_flexcheckpoint_merge,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_semi_auto_parallel_c_cross_entropy,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_cross_mesh_reshard,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_semi_auto_parallel_llama_model_amp,LINUX,GPU,180,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_semi_auto_parallel_hybrid_sharding_strategy,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_global_mesh_reshard,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_semi_auto_parallel_global_input,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_semi_auto_parallel_multi_inputs,LINUX,GPU,120,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_semi_auto_parallel_llama_model_vpp,LINUX,GPU,180,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_semi_auto_parallel_llama_model_pir,LINUX,GPU,180,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..;FLAGS_enable_pir_api=1,
+test_pir_reshard_nd_mesh_func,LINUX,GPU,60,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_semi_auto_llama_acc_align,LINUX,GPU,300,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..;FLAGS_enable_pir_api=1,
+test_semi_auto_llama_save_load,LINUX,GPU,180,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..;FLAGS_enable_pir_api=1,
+test_parallel_api_with_llama_1d,LINUX,GPU,400,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_parallel_api_with_llama_2d,LINUX,GPU,400,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_parallel_api_with_llama_2d_sep,LINUX,GPU,400,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_parallel_api_with_llama_3d,LINUX,GPU,800,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_to_distributed_api_for_llama,LINUX,GPU,180,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_parallel_api_with_llama_lora,LINUX,GPU,360,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_process_mesh,LINUX,GPU,150,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,
+test_get_group_in_different_hybrid_configs,LINUX,GPU,150,HYBRID,test_runner.py,,,http_proxy=;https_proxy=;PYTHONPATH=../..,

test/auto_parallel/hybrid_strategy/to_distributed_api_for_llama.py

@@ -0,0 +1,643 @@
+# 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 math
+import os
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.distributed import to_distributed
+from paddle.distributed.auto_parallel.high_level_api import ToDistributedConfig
+
+EPOCHS = 1
+VOCAB_SIZE = 8000
+BATCH_NUM = 2
+BATCH_SIZE = 4
+HIDDEN_SIZE = 2048
+INTERMEDIATE_SIZE = 4096
+SEQ_LENGTH = 1024
+N_HEAD = 32
+NUM_HIDDEN_LAYERS = 4
+
+
+def create_numpy_like_random(name):
+    return paddle.ParamAttr(
+        name=name, initializer=paddle.nn.initializer.Uniform(-0.1, 0.1)
+    )
+
+
+class LlamaRotaryEmbedding(nn.Layer):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000):
+        super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        # [dim / 2]
+        self.inv_freq = 1.0 / (
+            self.base
+            ** (
+                paddle.cast(paddle.arange(0, self.dim, 2), dtype="float32")
+                / self.dim
+            )
+        )
+        self._set_cos_sin_cache(seq_len=max_position_embeddings)
+
+    def _set_cos_sin_cache(self, seq_len):
+        self.max_seq_len_cached = seq_len
+        # [seq_len]
+        t = paddle.arange(seq_len, dtype="float32")
+        # [seq_len, dim/2]
+        freqs = paddle.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        # [seq_len, dim]
+        emb = paddle.concat([freqs, freqs], axis=-1)
+        # [1, seqlen, 1, dim]
+        self.cos_cached = emb.cos()[None, :, None, :]
+        self.sin_cached = emb.sin()[None, :, None, :]
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        cos = self.cos_cached[:, :seq_len, :, :]
+        sin = self.sin_cached[:, :seq_len, :, :]
+        return (
+            cos.cast(x.dtype) if cos.dtype != x.dtype else cos,
+            sin.cast(x.dtype) if sin.dtype != x.dtype else sin,
+        )
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return paddle.concat([-x2, x1], axis=-1)  # shape is the same as x
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
+    if position_ids is None:
+        # Note: Only for LlamaForCausalLMPipe model pretraining
+        cos = cos[:, : q.shape[1], :, :]  # [bs, seq_len, 1, dim]
+        sin = sin[:, : q.shape[1], :, :]  # [bs, seq_len, 1, dim]
+    else:
+        cos = cos.squeeze(axis=[0, 2])  # [seq_len, dim]
+        sin = sin.squeeze(axis=[0, 2])  # [seq_len, dim]
+        cos = cos[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]
+        sin = sin[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def scaled_dot_product_attention(
+    query_states,
+    key_states,
+    value_states,
+    attention_mask,
+):
+    bsz, q_len, num_heads, head_dim = query_states.shape
+    _, kv_seq_len, _, _ = value_states.shape
+
+    #  [ bz, seqlen, nhead, head_dim] -> [bs, nhead, seq_len, head_dim]
+    query_states = paddle.transpose(query_states, [0, 2, 1, 3])
+    # merge with the next transpose
+    key_states = paddle.transpose(key_states, [0, 2, 1, 3])
+    value_states = paddle.transpose(value_states, [0, 2, 1, 3])
+
+    # matmul and divide by sqrt(head_dim)
+    attn_weights = paddle.matmul(
+        query_states / math.sqrt(head_dim), key_states.transpose([0, 1, 3, 2])
+    )
+
+    attention_mask = attention_mask.reshape([bsz, 1, q_len, kv_seq_len])
+
+    attn_weights = attn_weights + attention_mask
+    if not paddle.in_dynamic_mode():
+        attn_weights = F.softmax(attn_weights, axis=-1, dtype="float32").astype(
+            query_states.dtype
+        )
+    else:
+        with paddle.amp.auto_cast(False):
+            attn_weights = F.softmax(
+                attn_weights, axis=-1, dtype="float32"
+            ).astype(query_states.dtype)
+
+    attn_output = paddle.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose([0, 2, 1, 3])
+
+    attn_output = attn_output.reshape([bsz, q_len, head_dim * num_heads])
+
+    return attn_output
+
+
+class LlamaAttention(nn.Layer):
+    def __init__(self, param_prefix="", hidden_size=HIDDEN_SIZE, n_head=N_HEAD):
+        super().__init__()
+        weight_attr_0 = create_numpy_like_random(param_prefix + "_0")
+        weight_attr_1 = create_numpy_like_random(param_prefix + "_1")
+        weight_attr_2 = create_numpy_like_random(param_prefix + "_2")
+        weight_attr_3 = create_numpy_like_random(param_prefix + "_3")
+        self.hidden_size = hidden_size
+        self.num_heads = n_head
+        self.head_dim = hidden_size // n_head
+        self.q_proj = nn.Linear(
+            hidden_size, hidden_size, weight_attr_0, bias_attr=False
+        )
+        self.k_proj = nn.Linear(
+            hidden_size, hidden_size, weight_attr_1, bias_attr=False
+        )
+        self.v_proj = nn.Linear(
+            hidden_size, hidden_size, weight_attr_2, bias_attr=False
+        )
+        self.o_proj = nn.Linear(
+            hidden_size, hidden_size, weight_attr_3, bias_attr=False
+        )
+        self.rotary_emb = LlamaRotaryEmbedding(
+            self.head_dim, max_position_embeddings=SEQ_LENGTH, base=10000
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        position_ids=None,
+        attention_mask=None,
+    ):
+        # mix_layer = self.qkv_proj(x)
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # target_shape = [0, 0, self.num_heads, 3 * self.head_dim]
+        target_query_shape = [0, 0, self.num_heads, self.head_dim]
+        target_key_value_shape = [0, 0, self.num_heads, self.head_dim]
+
+        # mix_layer = paddle.reshape(mix_layer, target_shape)
+        query_states = query_states.reshape(shape=target_query_shape)
+        key_states = key_states.reshape(shape=target_key_value_shape)
+        value_states = value_states.reshape(shape=target_key_value_shape)
+        kv_seq_len = key_states.shape[-3]
+
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin, position_ids
+        )
+
+        output = scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+        )
+
+        attn_output = output
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output
+
+
+class LlamaMlp(nn.Layer):
+    def __init__(
+        self,
+        param_prefix="",
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        weight_attr_0 = create_numpy_like_random(param_prefix + "_0")
+        weight_attr_1 = create_numpy_like_random(param_prefix + "_1")
+        weight_attr_2 = create_numpy_like_random(param_prefix + "_2")
+
+        self.gate_proj = nn.Linear(
+            hidden_size, intermediate_size, weight_attr_1, bias_attr=False
+        )
+        self.up_proj = nn.Linear(
+            hidden_size, intermediate_size, weight_attr_0, bias_attr=False
+        )
+        self.down_proj = nn.Linear(
+            intermediate_size, hidden_size, weight_attr_2, bias_attr=False
+        )
+
+    def forward(self, x):
+        x = paddle.nn.functional.swiglu(self.gate_proj(x), self.up_proj(x))
+        out = self.down_proj(x)
+        return out
+
+
+class LlamaRMSNorm(nn.Layer):
+    def __init__(self, hidden_size=HIDDEN_SIZE):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.weight = paddle.create_parameter(
+            shape=[self.hidden_size],
+            dtype=paddle.get_default_dtype(),
+            default_initializer=nn.initializer.Constant(1.0),
+        )
+        self.variance_epsilon = 1.0
+
+    def forward(self, hidden_states):
+        with paddle.amp.auto_cast(False):
+            hidden_states = hidden_states.astype("float32")
+            variance = hidden_states.pow(2).mean(-1, keepdim=True)
+            hidden_states = (
+                paddle.rsqrt(variance + self.variance_epsilon) * hidden_states
+            )
+        if self.weight.dtype in [paddle.float16, paddle.bfloat16]:
+            hidden_states = paddle.cast(hidden_states, self.weight.dtype)
+        return hidden_states * self.weight
+
+
+class LlamaDecoderLayer(nn.Layer):
+    def __init__(
+        self,
+        param_prefix="",
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+
+        self.self_attn = LlamaAttention(param_prefix + "_att", hidden_size)
+        self.mlp = LlamaMlp(param_prefix + "_mlp")
+        self.input_layernorm = LlamaRMSNorm(hidden_size)
+        self.post_attn_layernorm = LlamaRMSNorm(hidden_size)
+
+    def forward(
+        self,
+        hidden_states,
+        position_ids=None,
+        attention_mask=None,
+    ):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states, position_ids, attention_mask
+        )
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attn_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+def _expand_2d_mask(mask, dtype, tgt_length):
+    """
+    Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`.
+    """
+    batch_size, src_length = mask.shape[0], mask.shape[-1]
+    tgt_length = tgt_length if tgt_length is not None else src_length
+
+    mask = mask[:, None, None, :].astype("bool")
+    mask.stop_gradient = True
+    expanded_mask = mask.expand([batch_size, 1, tgt_length, src_length])
+
+    return expanded_mask
+
+
+def _make_causal_mask(input_ids_shape, past_key_values_length):
+    """
+    Make casual mask used for self-attention
+    """
+    batch_size, target_length = input_ids_shape  # target_length: seq_len
+
+    mask = paddle.tril(
+        paddle.ones((target_length, target_length), dtype="bool")
+    )
+
+    if past_key_values_length > 0:
+        # [tgt_len, tgt_len + past_len]
+        mask = paddle.concat(
+            [
+                paddle.ones(
+                    [target_length, past_key_values_length], dtype="bool"
+                ),
+                mask,
+            ],
+            axis=-1,
+        )
+
+    # [bs, 1, tgt_len, tgt_len + past_len]
+    return mask[None, None, :, :].expand(
+        [batch_size, 1, target_length, target_length + past_key_values_length]
+    )
+
+
+def _prepare_decoder_attention_mask(
+    attention_mask, input_shape, past_key_values_length, dtype
+):
+    if attention_mask is not None:
+        if len(attention_mask.shape) == 2:
+            expanded_attn_mask = _expand_2d_mask(
+                attention_mask, dtype, tgt_length=input_shape[-1]
+            )
+            if input_shape[-1] > 1:
+                combined_attention_mask = _make_causal_mask(
+                    input_shape,
+                    past_key_values_length=past_key_values_length,
+                )
+                expanded_attn_mask = (
+                    expanded_attn_mask & combined_attention_mask
+                )
+        # [bsz, seq_len, seq_len] -> [bsz, 1, seq_len, seq_len]
+        elif len(attention_mask.shape) == 3:
+            expanded_attn_mask = attention_mask.unsqueeze(1).astype("bool")
+        else:
+            expanded_attn_mask = attention_mask
+    else:
+        expanded_attn_mask = _make_causal_mask(
+            input_shape, past_key_values_length=past_key_values_length
+        )
+    expanded_attn_mask = paddle.where(
+        expanded_attn_mask, 0.0, paddle.finfo(dtype).min
+    ).astype(dtype)
+    return expanded_attn_mask
+
+
+class LlamaModel(nn.Layer):
+    def __init__(
+        self,
+        param_prefix="",
+        vocab_size=VOCAB_SIZE,
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+    ):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+
+        self.embed_tokens = nn.Embedding(
+            vocab_size,
+            hidden_size,
+        )
+
+        self.layers = nn.LayerList(
+            [
+                LlamaDecoderLayer(param_prefix + "_decoder_" + str(i))
+                for i in range(NUM_HIDDEN_LAYERS)
+            ]
+        )
+        self.norm = LlamaRMSNorm(hidden_size)
+
+    def forward(
+        self,
+        input_ids=None,
+        position_ids=None,
+        attention_mask=None,
+    ):
+        batch_size, seq_length = input_ids.shape
+
+        inputs_embeds = self.embed_tokens(input_ids)
+
+        # embed positions
+        attention_mask = paddle.ones(
+            (batch_size, seq_length), dtype=paddle.bool
+        )
+        if position_ids is None:
+            position_ids = paddle.arange(seq_length, dtype="int64").expand(
+                (batch_size, seq_length)
+            )
+
+        attention_mask = _prepare_decoder_attention_mask(
+            attention_mask,
+            (batch_size, seq_length),
+            0,
+            inputs_embeds.dtype,
+        )  # [bs, 1, seq_len, seq_len]
+
+        hidden_states = inputs_embeds
+
+        for idx, (decoder_layer) in enumerate(self.layers):
+            layer_outputs = decoder_layer(
+                hidden_states,
+                position_ids,
+                attention_mask,
+            )
+
+            hidden_states = layer_outputs
+
+        hidden_states = self.norm(hidden_states)
+
+        return hidden_states
+
+
+class LlamaPretrainingCriterion(paddle.nn.Layer):
+    """
+    Criterion for Llama.
+    It calculates the final loss.
+    """
+
+    def __init__(
+        self,
+        param_prefix="",
+        vocab_size=VOCAB_SIZE,
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+    ):
+        super().__init__()
+        self.ignore_index = -100
+        self.loss_func = paddle.nn.CrossEntropyLoss(
+            reduction="none", ignore_index=self.ignore_index
+        )
+
+    def forward(self, prediction_scores, masked_lm_labels):
+        with paddle.amp.auto_cast(False):
+            masked_lm_loss = self.loss_func(
+                prediction_scores.astype("float32"),
+                masked_lm_labels.unsqueeze(2),
+            )
+
+            binary_sequence = paddle.where(
+                masked_lm_loss > 0,
+                paddle.ones_like(masked_lm_loss),
+                paddle.zeros_like(masked_lm_loss),
+            )
+            count = paddle.sum(binary_sequence)
+            if count == 0:
+                loss = paddle.sum(masked_lm_loss * binary_sequence)
+            else:
+                loss = paddle.sum(masked_lm_loss * binary_sequence) / count
+
+        return loss
+
+
+class LlamaLMHead(paddle.nn.Layer):
+    def __init__(
+        self,
+        param_prefix="",
+        vocab_size=VOCAB_SIZE,
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+
+        self.weight = self.create_parameter(
+            shape=[hidden_size, vocab_size],
+            dtype=paddle.get_default_dtype(),
+        )
+
+    def forward(self, hidden_states, tensor_parallel_output=None):
+        logits = paddle.matmul(hidden_states, self.weight, transpose_y=False)
+        return logits
+
+
+class LlamaForCausalLM(paddle.nn.Layer):
+    def __init__(
+        self,
+        param_prefix="",
+        vocab_size=VOCAB_SIZE,
+        hidden_size=HIDDEN_SIZE,
+        intermediate_size=INTERMEDIATE_SIZE,
+    ):
+        super().__init__()
+        self.llama = LlamaModel(
+            param_prefix + "_llama", vocab_size, hidden_size, intermediate_size
+        )
+        self.lm_head = LlamaLMHead(
+            param_prefix + "_lm_head",
+            vocab_size,
+            hidden_size,
+            intermediate_size,
+        )
+        self.criterion = LlamaPretrainingCriterion(
+            param_prefix + "_criterion",
+            vocab_size,
+            hidden_size,
+            intermediate_size,
+        )
+
+    def forward(
+        self,
+        input_ids=None,
+        position_ids=None,
+        attention_mask=None,
+        labels=None,
+    ):
+        outputs = self.llama(
+            input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+        )
+
+        logits = self.lm_head(outputs)
+
+        loss = None
+        if labels is not None:
+            loss = self.criterion(logits, labels)
+
+        return (loss, logits)
+
+
+class RandomDataset(paddle.io.Dataset):
+    def __init__(self, inputs, labels, num_samples):
+        self.inputs = inputs
+        self.labels = labels
+        self.num_samples = num_samples
+
+    def __getitem__(self, idx):
+        return self.inputs[idx], self.labels[idx]
+
+    def __len__(self):
+        return self.num_samples
+
+
+class TestLlamaDecoderForSemiAutoParallel:
+    def __init__(self):
+        self._dtype = os.getenv("dtype", "float32")
+        self._backend = os.getenv("backend", "gpu")
+        self._seed = eval(os.getenv("seed", "2023"))
+
+        self._device_num = os.getenv("num_of_devices", 8)
+        self._node_num = 1
+
+        np.random.seed(self._seed)
+        paddle.seed(self._seed)
+        self._model = LlamaForCausalLM("demo_llama")
+
+        # ensure that input data between dp is different and data within dp is the same
+        self._mesh = dist.ProcessMesh(
+            [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dim_names=["pp", "dp", "mp"]
+        )
+        if "dp" in self._mesh.dim_names:
+            dp_seed = self._mesh.get_rank_by_dim_and_process_id(
+                "dp", dist.get_rank()
+            )
+        else:
+            dp_seed = 0
+        np.random.seed(self._seed + dp_seed)
+        paddle.seed(self._seed + dp_seed)
+        self._input_seqs = np.random.randint(
+            low=0, high=1024, size=(BATCH_SIZE * BATCH_NUM, SEQ_LENGTH)
+        ).astype("int64")
+        self._labels = np.random.randint(
+            low=0, high=1024, size=(BATCH_SIZE * BATCH_NUM, SEQ_LENGTH)
+        ).astype("int64")
+        self._dataset = RandomDataset(
+            self._input_seqs, self._labels, BATCH_SIZE * BATCH_NUM
+        )
+        self._sampler = paddle.io.BatchSampler(
+            self._dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True
+        )
+        self._loader = paddle.io.DataLoader(
+            self._dataset, batch_sampler=self._sampler
+        )
+        self._opt = paddle.optimizer.SGD(
+            learning_rate=0.1, parameters=self._model.parameters()
+        )
+
+        paddle.set_device(self._backend)
+
+    def test_to_distributed_api(self):
+        # # config: sequence_parallel
+        dist_config = ToDistributedConfig()
+        dist_config.sequence_parallel = True
+
+        # # wrap model by using **to_distributed**
+        dist_model, dist_opt, dist_loader = to_distributed(
+            self._model,
+            self._opt,
+            self._loader,
+            self._device_num,
+            self._node_num,
+            dist_config,
+        )
+
+        for epoch in range(EPOCHS):
+            dist_model.train()
+            for i, data in enumerate(dist_loader()):
+                inputs, labels = data
+                loss, _ = dist_model(inputs, labels=labels)
+                loss.backward()
+                dist_opt.step()
+                dist_opt.clear_grad()
+
+    def run_test_case(self):
+        if self._backend == "gpu":
+            cuda_version_main = int(paddle.version.cuda().split(".")[0])
+            device_prop_main = paddle.device.cuda.get_device_capability()[0]
+            if cuda_version_main >= 11 and device_prop_main >= 8:
+                self.test_to_distributed_api()
+
+
+if __name__ == '__main__':
+    TestLlamaDecoderForSemiAutoParallel().run_test_case()

test/auto_parallel/iterable_dataset.py

@@ -0,0 +1,159 @@
+# 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 os
+import tempfile
+
+import numpy as np
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.distributed.fleet import auto
+
+paddle.enable_static()
+global_process_mesh = auto.ProcessMesh(mesh=[0, 1])
+PP_MESH_0 = auto.ProcessMesh([0])
+PP_MESH_1 = auto.ProcessMesh([1])
+batch_size = 2
+batch_num = 10
+hidden_size = 1024
+sequence_len = 512
+image_size = hidden_size
+class_num = 10
+
+paddle.seed(44)
+
+
+class MyDataset(paddle.io.IterableDataset):
+    def __init__(self, num_samples):
+        self.num_samples = num_samples
+
+    def __iter__(self):
+        for i in range(self.num_samples):
+            input = np.random.uniform(size=image_size).astype("float32")
+            label = np.random.randint(0, class_num - 1, dtype="int64")
+            yield input, label
+
+
+class MyDataset1(paddle.io.Dataset):
+    def __init__(self, num_samples):
+        self.num_samples = num_samples
+        self.data = []
+        for i in range(self.num_samples):
+            input1 = np.random.uniform(size=image_size).astype("float32")
+            label1 = np.array(
+                np.random.randint(0, class_num - 1, dtype="int64")
+            )
+            input2 = np.random.uniform(size=image_size).astype("float32")
+            label2 = np.array(
+                np.random.randint(0, class_num - 1, dtype="int64")
+            )
+            input = np.stack((input1, input2))
+            label = np.stack((label1, label2))
+            self.data.append((input, label))
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+    def __len__(self):
+        return len(self.data)
+
+
+class MLPLayer(nn.Layer):
+    def __init__(
+        self,
+        hidden_size=1024,
+        intermediate_size=4 * 1024,
+        dropout_ratio=0.1,
+        initializer_range=0.02,
+    ):
+        super().__init__()
+        d_model = hidden_size
+        dim_feedforward = intermediate_size
+        weight_attr = paddle.ParamAttr(
+            initializer=nn.initializer.Normal(mean=0.0, std=initializer_range)
+        )
+        bias_attr = None
+
+        self.linear0 = nn.Linear(
+            d_model, dim_feedforward, weight_attr, bias_attr=bias_attr
+        )
+        self.linear1 = nn.Linear(
+            dim_feedforward, d_model, weight_attr, bias_attr=bias_attr
+        )
+        self.linear2 = nn.Linear(d_model, 1, weight_attr, bias_attr=bias_attr)
+        self.norm = nn.LayerNorm(d_model, epsilon=1e-5)
+        self.dropout = nn.Dropout(dropout_ratio, mode="upscale_in_train")
+
+    def forward(self, input):
+        out = auto.shard_op(self.norm, PP_MESH_0)(input)
+        out = self.linear0(out)
+        out = F.gelu(out, approximate=True)
+        out = auto.shard_op(self.linear1, PP_MESH_1)(out)
+        out = self.dropout(out)
+        out = self.linear2(out)
+        self.out = out
+        return out
+
+
+def train(fetch):
+    mlp = MLPLayer(
+        hidden_size=hidden_size,
+        intermediate_size=4 * hidden_size,
+        dropout_ratio=0.1,
+        initializer_range=0.02,
+    )
+    loss = paddle.nn.CrossEntropyLoss()
+    optimizer = paddle.optimizer.Adam(
+        learning_rate=0.00001,
+        beta1=0.9,
+        beta2=0.999,
+        epsilon=1e-08,
+        grad_clip=None,
+    )
+
+    dist_strategy = auto.Strategy()
+    dist_strategy.auto_mode = "semi"
+    dist_strategy.split_data = True
+
+    # init engine
+    engine = auto.Engine(
+        mlp, loss, optimizer, paddle.metric.Accuracy(), strategy=dist_strategy
+    )
+
+    # train
+    train_dataset = MyDataset(batch_num * batch_size)
+    engine.fit(train_dataset, epochs=2, batch_size=batch_size)
+
+    train_dataset1 = MyDataset1(batch_size * batch_num)
+    engine.fit(train_dataset1, epochs=2, batch_size=None)
+
+    # eval
+    eval_dataset = MyDataset(batch_size)
+    engine.evaluate(eval_dataset, batch_size=batch_size)
+
+    # predict
+    test_dataset = MyDataset(batch_size)
+    engine.predict(test_dataset, batch_size=batch_size)
+
+    # save
+    temp_dir = tempfile.TemporaryDirectory()
+    model_filename = os.path.join(temp_dir.name, 'mlp_inf')
+    engine.save(model_filename, training=False)
+    temp_dir.cleanup()
+
+
+if __name__ == "__main__":
+    train(fetch=True)

test/auto_parallel/launch.py

@@ -0,0 +1,23 @@
+# 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 os
+
+from paddle.distributed.fleet import launch
+from paddle.distributed.fleet.launch_utils import run_with_coverage
+
+if __name__ == "__main__":
+    if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
+        run_with_coverage(True)
+    launch.launch()

test/auto_parallel/local_view_compute.py

@@ -0,0 +1,220 @@
+# 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 random
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import ProcessMesh, fleet, get_rank, shard_dataloader
+from paddle.io import BatchSampler, DataLoader, DistributedBatchSampler
+
+base_lr = 0.001  # Learning rate
+l2_decay = 1e-5  # Weight decay
+
+epoch = 5  # Number of training epochs
+batch_num = 100  # Number of batches per epoch
+batch_size = 32  # Batch size for training
+class_dim = 10
+global_local_loss_list = []
+
+
+class RandomDataset(paddle.io.Dataset):
+    def __init__(self, images, labels):
+        self.num_samples = len(images)
+        self.images = images
+        self.labels = labels
+
+    def __getitem__(self, idx):
+        # image = np.random.random([256]).astype('float32')
+        # label = np.random.randint(0, class_dim - 1, (1, )).astype('int64')
+        image = self.images[idx]
+        label = self.labels[idx]
+        return image, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+class SimpleNet(paddle.nn.Layer):
+    def __init__(self, input_size, inner_size, output_size):
+        super().__init__()
+        self.linear1 = paddle.nn.Linear(input_size, inner_size)
+        self.linear2 = paddle.nn.Linear(inner_size, input_size)
+        self.linear3 = paddle.nn.Linear(input_size, output_size)
+        self.relu = paddle.nn.ReLU()
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.linear2(x)
+        x = self.linear3(x)
+        x = self.relu(x)
+        return x
+
+
+def masked_lm_loss_func(pred, label, global_local_loss_list_item=None):
+    """自定义损失函数，基于rank进行掩码"""
+    lossmask = paddle.zeros_like(label).astype('float32')
+    if dist.get_rank() == 0:
+        lossmask[:8] = 1
+    else:
+        lossmask[8:16] = 1
+
+    pred_sub = pred[:, 0:1]  # shape [B,1]
+    # NOTE(Pan Zhaowu): Using float64 as golden to provide more
+    # persuasive result.
+    label_float = paddle.cast(label, 'float64')  # shape [B,1]
+    raw_loss = paddle.abs(pred_sub - label_float)
+    lossmask_ = lossmask.reshape([-1]).cast('float64')
+    raw_loss_flat = raw_loss.reshape([-1]).cast('float64')
+
+    masked_lm_loss_sum = paddle.sum(raw_loss_flat * lossmask_)
+    valid_count = paddle.sum(lossmask_)
+
+    loss = masked_lm_loss_sum / (valid_count + 1e-8)
+    if global_local_loss_list_item is not None:
+        np.testing.assert_allclose(
+            global_local_loss_list_item,
+            loss.numpy(),
+            rtol=1e-8,
+        )
+    return loss
+
+
+class TestLocalViewCompute:
+    def __init__(self):
+        self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def set_random_seed(self):
+        np.random.seed(2025)
+        paddle.seed(2025)
+        random.seed(2025)
+
+    def create_dataset(self):
+        images = np.random.rand(batch_num * batch_size * 2, 256).astype(
+            'float32'
+        )
+        labels = np.random.randint(
+            0, class_dim - 1, (batch_num * batch_size * 2, 1)
+        ).astype('int64')
+        datasets = RandomDataset(images, labels)
+        return datasets
+
+    def run_test_cases(self):
+        # run_dy_hand_get_local_loss
+        self.set_random_seed()
+        dataset = self.create_dataset()
+        dist_strategy = fleet.DistributedStrategy()
+        dist_strategy.hybrid_configs = {
+            "dp_degree": 2,
+            "mp_degree": 1,
+            "pp_degree": 1,
+        }
+
+        fleet.init(is_collective=True, strategy=dist_strategy)
+        model = SimpleNet(
+            input_size=256, inner_size=102400, output_size=class_dim
+        )
+        clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=1.0)
+        optimizer = paddle.optimizer.AdamW(
+            learning_rate=base_lr,
+            weight_decay=l2_decay,
+            parameters=model.parameters(),
+            grad_clip=clip,
+        )
+
+        model = fleet.distributed_model(model)
+        optimizer = fleet.distributed_optimizer(optimizer)
+
+        sampler = DistributedBatchSampler(
+            dataset,
+            rank=get_rank(),
+            batch_size=batch_size // 2,
+            shuffle=False,
+            drop_last=True,
+        )
+        train_loader = DataLoader(
+            dataset, batch_sampler=sampler, num_workers=1, shuffle=False
+        )
+
+        model.train()
+        for batch_id, data in enumerate(train_loader()):
+            if batch_id > 10:
+                break
+
+            img, label = data
+
+            out = model(img)
+
+            avg_loss = masked_lm_loss_func(out, label)
+            avg_loss.backward()
+            optimizer.step()
+            model.clear_gradients()
+            global_local_loss_list.append(avg_loss.numpy())
+
+        # run_dy_semi_auto
+        self.set_random_seed()
+        dataset = self.create_dataset()
+        world_process_mesh = ProcessMesh([0, 1], dim_names=["dp"])
+        model = SimpleNet(
+            input_size=256, inner_size=102400, output_size=class_dim
+        )
+        optimizer = paddle.optimizer.AdamW(
+            learning_rate=base_lr,
+            weight_decay=l2_decay,
+            parameters=model.parameters(),
+            grad_clip=clip,
+        )
+
+        sampler = BatchSampler(
+            dataset, batch_size=batch_size, shuffle=False, drop_last=True
+        )
+        train_loader = DataLoader(
+            dataset, batch_sampler=sampler, num_workers=1, shuffle=False
+        )
+
+        dist_dataloader = shard_dataloader(
+            dataloader=train_loader, meshes=world_process_mesh, shard_dims="dp"
+        )
+
+        model.train()
+        process_mesh = ProcessMesh([0, 1], dim_names=["dp"])
+        out_placements = [dist.Partial(dist.ReduceType.kRedAvg)]
+
+        for batch_id, data in enumerate(dist_dataloader()):
+            if batch_id > 10:
+                break
+
+            img, label = data
+
+            out = model(img)
+            loss_func = dist.local_map(
+                masked_lm_loss_func,
+                out_placements=out_placements,
+                in_placements=[None, None],
+                process_mesh=process_mesh,
+            )
+            avg_loss = loss_func(
+                out,
+                label,
+                global_local_loss_list_item=global_local_loss_list[batch_id],
+            )
+            avg_loss.backward()
+            optimizer.step()
+            model.clear_gradients()
+
+
+if __name__ == '__main__':
+    TestLocalViewCompute().run_test_cases()

test/auto_parallel/microbatch_demo.py

@@ -0,0 +1,239 @@
+# 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
+from paddle.distributed.auto_parallel.pipelining.microbatch import (
+    TensorChunkSpec,
+    merge_chunks,
+    split_args_kwargs_into_chunks,
+)
+
+
+class TestMicrobatch:
+    def __init__(self):
+        paddle.seed(2024)
+        paddle.distributed.init_parallel_env()
+        self.batch_size = 8
+        self.feature_size = 4
+        self.tensor = paddle.randn([self.batch_size, self.feature_size])
+        self.rank = paddle.distributed.get_rank()
+
+    def test_tensor_chunk_spec(self):
+        # Test creation and string representation of TensorChunkSpec
+        spec = TensorChunkSpec(0)
+        assert spec.split_axis == 0
+        assert str(spec) == "TensorChunkSpec(0)"
+        assert "TensorChunkSpec(0)" in repr(spec)
+
+    def test_split_args_kwargs(self):
+        # Test basic parameter splitting
+        args = (self.tensor,)
+        kwargs = {"input": self.tensor}
+        num_chunks = 2
+
+        args_split, kwargs_split = split_args_kwargs_into_chunks(
+            args, kwargs, num_chunks
+        )
+
+        assert len(args_split) == num_chunks
+        assert len(kwargs_split) == num_chunks
+        assert args_split[0][0].shape[0] == self.batch_size // num_chunks
+
+        # Test splitting with non-tensor parameters
+        args = (self.tensor, 42, "string")
+        kwargs = {"tensor": self.tensor, "number": 42}
+        num_chunks = 2
+
+        args_split, kwargs_split = split_args_kwargs_into_chunks(
+            args, kwargs, num_chunks
+        )
+
+        # Verify non-tensor parameters remain unchanged in each chunk
+        assert args_split[0][1] == 42
+        assert args_split[0][2] == "string"
+        assert kwargs_split[0]["number"] == 42
+
+        # Test splitting with custom specification
+        tensor_2d = paddle.randn([4, 6])
+        args = (tensor_2d,)
+        args_chunk_spec = (TensorChunkSpec(1),)  # Split on second dimension
+
+        args_split, _ = split_args_kwargs_into_chunks(
+            args, None, 2, args_chunk_spec
+        )
+
+        assert args_split[0][0].shape[1] == 3
+
+    def test_merge_chunks(self):
+        # Test merging chunks
+        chunk1 = paddle.randn([4, 4])
+        chunk2 = paddle.randn([4, 4])
+        chunks = [chunk1, chunk2]
+        chunk_spec = [TensorChunkSpec(0)]
+
+        merged = merge_chunks(chunks, chunk_spec)
+        assert merged.shape[0] == 8
+
+        # Test merging chunks containing non-tensor values
+        chunks = [(paddle.randn([4, 4]), 42)] * 2
+        chunk_spec = [TensorChunkSpec(0), None]
+
+        merged = merge_chunks(chunks, chunk_spec)
+        assert merged[1] == 42
+
+        # Test error cases
+        try:
+            # Test error when tensor size is smaller than number of chunks
+            small_tensor = paddle.randn([1, 4])
+            split_args_kwargs_into_chunks((small_tensor,), None, 2)
+            raise AssertionError("Expected ValueError")
+        except ValueError:
+            pass
+
+        try:
+            # Test error when parameter count doesn't match chunk_spec length
+            split_args_kwargs_into_chunks(
+                (self.tensor,),
+                None,
+                2,
+                (TensorChunkSpec(0), TensorChunkSpec(1)),
+            )
+            raise AssertionError("Expected ValueError")
+        except AssertionError:
+            pass
+
+        # test merge empty chunks
+        empty_chunks = []
+        result = merge_chunks(empty_chunks, None)
+        assert result == []
+
+        # test tensor size smaller than chunks number
+        small_tensor = paddle.randn([1, 4])
+        try:
+            split_args_kwargs_into_chunks((small_tensor,), None, 2)
+            raise AssertionError("Expected ValueError")
+        except ValueError:
+            pass
+
+        # test merge non-tensor with tensor spec
+        chunks = [(42,), (42,)]
+        chunk_spec = (TensorChunkSpec(0),)
+        result = merge_chunks(chunks, chunk_spec)
+        assert result[0] == 42
+
+    def test_nested_structure(self):
+        # test nested tensor
+        nested_tensor = [
+            [paddle.randn([4, 2]), paddle.randn([4, 2])],
+            [paddle.randn([4, 2]), paddle.randn([4, 2])],
+        ]
+
+        args = (nested_tensor,)
+        kwargs = {"nested": nested_tensor}
+
+        args_split, kwargs_split = split_args_kwargs_into_chunks(
+            args, kwargs, 2
+        )
+
+        assert len(args_split) == 2
+        assert len(args_split[0][0]) == 2
+        assert len(args_split[0][0][0]) == 2
+        assert args_split[0][0][0][0].shape == [2, 2]
+
+        assert len(kwargs_split) == 2
+        assert len(kwargs_split[0]["nested"]) == 2
+        assert len(kwargs_split[0]["nested"][0]) == 2
+        assert kwargs_split[0]["nested"][0][0].shape == [2, 2]
+
+        merged_args = merge_chunks(
+            args_split,
+            [
+                [TensorChunkSpec(0), TensorChunkSpec(0)],
+                [TensorChunkSpec(0), TensorChunkSpec(0)],
+            ],
+        )
+
+        assert merged_args[0][0][0].shape == [4, 2]
+        assert merged_args[0][1][1].shape == [4, 2]
+
+        assert len(merged_args[0]) == 2
+        assert len(merged_args[0][0]) == 2
+
+    def test_dist_tensor_split_and_merge(self):
+        # test dist tensor split and merge
+        base_tensor = self.tensor
+        dense_tensor, _ = split_args_kwargs_into_chunks(
+            (base_tensor,),
+            None,
+            2,
+        )
+        mesh = paddle.distributed.ProcessMesh([0, 1], dim_names=["dp"])
+        dist_tensor = paddle.distributed.shard_tensor(
+            self.tensor,
+            mesh,
+            [paddle.distributed.Shard(0)],
+        )
+        dist_tensor_split, _ = split_args_kwargs_into_chunks(
+            (dist_tensor,),
+            None,
+            2,
+        )
+        if self.rank == 0:
+            is_equal = (
+                dist_tensor_split[0][0]
+                ._local_value()
+                .equal_all(dense_tensor[0][0][:2])
+            )
+            assert is_equal.item()
+            is_equal = (
+                dist_tensor_split[1][0]
+                ._local_value()
+                .equal_all(dense_tensor[0][0][2:])
+            )
+            assert is_equal.item()
+        else:
+            is_equal = (
+                dist_tensor_split[0][0]
+                ._local_value()
+                .equal_all(dense_tensor[1][0][:2])
+            )
+            assert is_equal.item()
+            is_equal = (
+                dist_tensor_split[1][0]
+                ._local_value()
+                .equal_all(dense_tensor[1][0][2:])
+            )
+            assert is_equal.item()
+        chunk1 = dist_tensor_split[0][0]
+        chunk2 = dist_tensor_split[1][0]
+        chunk_spec = [TensorChunkSpec(0)]
+        merged_chunk = merge_chunks([chunk1, chunk2], chunk_spec)
+        if self.rank == 0:
+            is_equal = merged_chunk._local_value().equal_all(base_tensor[:4])
+            assert is_equal.item()
+        else:
+            is_equal = merged_chunk._local_value().equal_all(base_tensor[4:])
+            assert is_equal.item()
+
+    def run_all_tests(self):
+        """Run all test methods"""
+        self.test_tensor_chunk_spec()
+        self.test_split_args_kwargs()
+        self.test_merge_chunks()
+        self.test_nested_structure()
+        self.test_dist_tensor_split_and_merge()
+
+
+if __name__ == "__main__":
+    TestMicrobatch().run_all_tests()

test/auto_parallel/object_list_communication_demo.py

@@ -0,0 +1,146 @@
+# 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 _recv_info:
+    def __init__(self, tensor, placements):
+        self.obj_size = 10
+        self.obj_type1 = paddle.distributed.Shard(0)
+        self.obj_type2 = paddle.distributed.Replicate()
+        self.obj_type3 = paddle.distributed.Partial()
+        self.obj_list = [
+            paddle.distributed.Shard(0),
+            paddle.distributed.Replicate(),
+            paddle.distributed.Partial(),
+        ]
+        self.dtype = paddle.int64
+
+
+class TestObjectListCommunication:
+    def init_dist_env(self):
+        dist.init_parallel_env()
+        paddle.seed(2025)
+
+    def test_object_list_communication(self):
+        """Test object list communication functionalities including parameter validation,
+        group operations and normal communication process"""
+        self.init_dist_env()
+        curr_rank = dist.get_rank()
+
+        # Test case 1: Parameter validation - empty list
+        if curr_rank == 0:
+            try:
+                dist.send_object_list([], dst=1)
+                raise AssertionError("Should raise ValueError")
+            except ValueError:
+                pass
+        else:
+            try:
+                dist.recv_object_list([], src=0)
+                raise AssertionError("Should raise ValueError")
+            except ValueError:
+                pass
+
+        # Test case 2: Group operations - rank not in group
+        excluded_group = dist.new_group([2, 3])
+        send_list = ["test"]
+        recv_list = [None]
+        if curr_rank == 0:
+            # test the dst is not in the group
+            dist.send_object_list(["test"], dst=1, group=excluded_group)
+        elif curr_rank == 1:
+            # test the src is not in the group
+            dist.recv_object_list([None], src=0, group=excluded_group)
+            assert recv_list[0] is None
+
+        excluded_group_1 = dist.new_group([0, 1])
+        if curr_rank == 0:
+            dist.send_object_list(send_list, dst=1, group=excluded_group_1)
+        elif curr_rank == 1:
+            dist.recv_object_list(recv_list, src=0, group=excluded_group_1)
+            assert recv_list[0] == "test"
+
+        # Test case 3: Group operations - parameter conflicts
+        if curr_rank == 0:
+            try:
+                dist.send_object_list(["test"], dst=1, dst_in_group=1)
+                raise AssertionError("Should raise ValueError")
+            except ValueError:
+                pass
+        elif curr_rank == 1:
+            try:
+                dist.recv_object_list([None], src=0, src_in_group=0)
+                raise AssertionError("Should raise ValueError")
+            except ValueError:
+                pass
+
+        # Test case 4: Group operations - using src_in_group/dst_in_group
+        test_group = dist.new_group([0, 1])
+        if curr_rank == 0:
+            data = ["test_group_dst"]
+            dist.send_object_list(data, group=test_group, dst_in_group=1)
+        elif curr_rank == 1:
+            data = [None]
+            dist.recv_object_list(data, group=test_group, src_in_group=0)
+            assert data[0] == "test_group_dst"
+
+        # Test case 5: Normal communication process
+        if curr_rank == 0:
+            data = [
+                42,  # integer
+                "hello",  # string
+                {"key": "value"},  # dictionary
+            ]
+            dist.send_object_list(data, dst=1)
+        elif curr_rank == 1:
+            data = [None] * 3
+            dist.recv_object_list(data, src=0)
+
+            assert data[0] == 42
+            assert data[1] == "hello"
+            assert data[2] == {"key": "value"}
+
+        # Test case 6: Test objects with distributed attributes
+        curr_rank = dist.get_rank()
+
+        if curr_rank == 0:
+            data1 = _recv_info(None, None)
+            data = [data1]
+            dist.send_object_list(data, dst=1)
+        elif curr_rank == 1:
+            data = [None]
+            dist.recv_object_list(data, src=0)
+
+            assert isinstance(data[0], _recv_info)
+            assert type(data[0].obj_size) == int
+            assert data[0].obj_size == 10
+
+            assert isinstance(data[0].obj_type1, paddle.distributed.Shard)
+            assert isinstance(data[0].obj_type2, paddle.distributed.Replicate)
+            assert isinstance(data[0].obj_type3, paddle.distributed.Partial)
+
+            assert data[0].dtype == paddle.int64
+
+            assert len(data[0].obj_list) == 3
+            assert isinstance(data[0].obj_list[0], paddle.distributed.Shard)
+            assert isinstance(data[0].obj_list[1], paddle.distributed.Replicate)
+            assert isinstance(data[0].obj_list[2], paddle.distributed.Partial)
+
+
+if __name__ == '__main__':
+    TestObjectListCommunication().test_object_list_communication()

test/auto_parallel/pipeline_sync_shared_parameters_unittest.py

@@ -0,0 +1,374 @@
+# 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 random
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import fleet
+from paddle.distributed.auto_parallel.pipelining.schedules import (
+    Schedule1F1B,
+    ScheduleFThenB,
+    ScheduleVPP,
+)
+from paddle.distributed.auto_parallel.pipelining.stage import PipelineStage
+from paddle.io import DataLoader, Dataset
+
+
+def fix_seeds(seed=2025):
+    """Fix random seeds to ensure reproducibility"""
+    paddle.seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+
+class PPModel(nn.Layer):
+    def __init__(self, name_prefix="", schedule="FThenB", shared_parameters={}):
+        super().__init__(name_scope=name_prefix)
+        self.name_prefix = name_prefix
+        self.mesh = paddle.distributed.ProcessMesh(
+            [0, 1, 2, 3], dim_names=["pp"]
+        )
+        self.num_layers = 8
+        self.num_layers_per_card = self.num_layers // 4
+        # Store the names of each pair of shared parameters.
+        self.shared_parameters = shared_parameters
+
+        self.linears = nn.LayerList()
+        for i in range(self.num_layers):
+            linear = nn.Linear(8, 8, bias_attr=False)
+
+            # Different models have distinct parameter name spaces to avoid naming conflicts.
+            linear.weight.name = f"{self.name_prefix}_linear_{i}_weight"
+
+            # Mark network parameters
+            linear.weight = dist.shard_tensor(
+                linear.weight,
+                (
+                    self.get_pp_mesh(i)
+                    if schedule != "VPP"
+                    else self.get_vpp_mesh(i)
+                ),
+                [dist.Replicate()],
+            )
+
+            self.linears.append(linear)
+
+        # Store the parameters to be shared under different model names.
+        self.model_shared_param_mp = {}
+
+        # Build `model_shared_param_mp`.
+        self.set_shared_param()
+
+    def set_shared_param(self):
+        for pair in self.shared_parameters:
+            assert len(pair) == 2
+            ori_name = pair[0]
+            sync_name = pair[1]
+            ori_param = None
+            for _, linear in enumerate(self.linears):
+                if ori_name == linear.weight.name:
+                    ori_param = linear.weight
+            assert ori_param is not None
+            self.model_shared_param_mp[sync_name] = ori_param
+
+    def get_pp_mesh(self, layer_index):
+        mesh_idx = int(layer_index / (self.num_layers / 4))
+        return self.mesh[mesh_idx]
+
+    def get_vpp_mesh(self, layer_index):
+        mesh_idx = int(layer_index % 4)
+        return self.mesh[mesh_idx]
+
+    def forward(self, x):
+        x.stop_gradient = False
+        out = x
+        for i in range(self.num_layers):
+            # Mark intermediate variables, reshard when switching devices
+            cur_mesh = self.get_pp_mesh(i)
+            if i % self.num_layers_per_card == 0 and i > 0:
+                out = dist.reshard(out, cur_mesh, [dist.Replicate()])
+            weight = self.linears[i].weight
+            if weight.name in self.model_shared_param_mp:
+                weight = dist.reshard(
+                    self.model_shared_param_mp[weight.name],
+                    cur_mesh,
+                    [dist.Replicate()],
+                )
+                out = paddle.matmul(out, weight)
+            else:
+                out = self.linears[i](out)
+        return paddle.cast(out, 'float32')
+
+
+class SingleStage(nn.Layer):
+    def __init__(self, layers):
+        super().__init__()
+        self.layers = layers
+
+    def forward(self, x):
+        x.stop_gradient = False
+        out = x
+        for i in range(len(self.layers)):
+            out = self.layers[i](out)
+        return paddle.cast(out, 'float32')
+
+
+class RandomDataset(Dataset):
+    def __init__(self, image_size, output_size, num_samples=1):
+        super().__init__()
+        self.image_size = image_size
+        self.num_samples = num_samples
+        self.output_size = output_size
+
+    def __getitem__(self, index):
+        input = paddle.rand([self.image_size], dtype='float32')
+        label = paddle.rand([self.output_size], dtype='float32')
+        return input, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+def _get_param_from_name(param_name, model):
+    for param in model.parameters():
+        if param.name == param_name:
+            return param
+    return None
+
+
+def build_shared_parameters(shared_params_names, model):
+    # Find the two shared parameters and build shared parameter information.
+    shared_mp = []
+    for pair in shared_params_names:
+        assert len(pair) == 2
+        ori_name = pair[0]
+        sync_name = pair[1]
+        ori_param = _get_param_from_name(ori_name, model)
+        sync_param = _get_param_from_name(sync_name, model)
+        # Note: Users must strictly maintain the format of the data structure here.
+        shared_mp.append({"params": [ori_param, sync_param]})
+    return shared_mp
+
+
+rtol = 1e-5
+
+
+class TestSharedParameters:
+    @classmethod
+    def setUpClass(cls):
+        """Initialize test class setup"""
+        paddle.distributed.init_parallel_env()
+        cls.group = paddle.distributed.new_group([0, 1, 2, 3])
+        cls.rank = dist.get_rank()
+        cls.mesh = paddle.distributed.ProcessMesh(
+            [0, 1, 2, 3], dim_names=["pp"]
+        )
+        fleet.auto.set_mesh(cls.mesh)
+
+    def test_single_schedule(self, sing_schedule="FThenB"):
+        """Test pipeline parallel model with shared parameters using FThenB/1F1B strategy"""
+        fix_seeds()
+        name_prefix = "pp_" + sing_schedule
+        self.model = PPModel(name_prefix=name_prefix)
+
+        self.micro_batches = 8
+        shared_params_names = [
+            [
+                f"{name_prefix}_linear_0_weight.dist",
+                f"{name_prefix}_linear_7_weight.dist",
+            ]
+        ]
+        # Pre-build shared parameter information.
+        shared_mp = build_shared_parameters(shared_params_names, self.model)
+
+        num_layers_per_card = 2
+        cur_rank = dist.get_rank()
+        stage_layers = SingleStage(
+            self.model.linears[
+                cur_rank * num_layers_per_card : (cur_rank + 1)
+                * num_layers_per_card
+            ]
+        )
+
+        self.stage = PipelineStage(
+            stage_layers,
+            self.rank,
+            4,
+            group=self.group,
+            shared_parameters=shared_mp,
+        )
+
+        self.stage.has_backward = True
+        loss_fn_ = nn.MSELoss()
+        if sing_schedule == "FThenB":
+            schedule = ScheduleFThenB(
+                self.stage, self.micro_batches, loss_fn=loss_fn_
+            )
+        elif sing_schedule == "1F1B":
+            schedule = Schedule1F1B(
+                self.stage, self.micro_batches, loss_fn=loss_fn_
+            )
+        else:
+            raise ValueError(
+                f"Unknown schedule type: {sing_schedule}. "
+                f"Currently `test_single_schedule` supported types are 'FThenB' and '1F1B'."
+            )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=self.model.parameters()
+        )
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_step = []
+        num_iterations = 20
+
+        for _ in range(num_iterations):
+            losses_by_micro_batch = []
+            for _, (data, label) in enumerate(loader):
+                schedule.step(data, target=label, losses=losses_by_micro_batch)
+                if self.rank == 3:
+                    losses_by_step.append(
+                        np.array(losses_by_micro_batch, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        return losses_by_step
+
+    def test_multi_schedule(self, multi_schedule="VPP"):
+        """Test pipeline parallel with shared parameters model using VPP strategy"""
+        fix_seeds()
+        name_prefix = "pp_" + multi_schedule
+        self.model = PPModel(name_prefix=name_prefix, schedule="VPP")
+        self.local_stages = 2
+        self.micro_batches = 8
+        self.stage_list = []
+
+        shared_params_names = [
+            [
+                f"{name_prefix}_linear_0_weight.dist",
+                f"{name_prefix}_linear_7_weight.dist",
+            ]
+        ]
+        # Pre-build shared parameter information.
+        shared_mp = build_shared_parameters(shared_params_names, self.model)
+
+        cur_rank = dist.get_rank()
+        for i in range(self.local_stages):
+            stage_layers = SingleStage(
+                self.model.linears[cur_rank + i * 4 : cur_rank + i * 4 + 1]
+            )
+            # Note: In VPP mode, the same `shared_mp` is used for building multiple
+            # stages to avoid redundant group creation.
+            self.stage_list.append(
+                PipelineStage(
+                    stage_layers,
+                    cur_rank + i * 4,
+                    8,
+                    group=self.group,
+                    shared_parameters=shared_mp,
+                )
+            )
+            self.stage_list[i].has_backward = True
+
+        loss_fn_ = nn.MSELoss()
+        schedule = ScheduleVPP(
+            self.stage_list, self.micro_batches, loss_fn=loss_fn_
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=self.model.parameters()
+        )
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=8)
+        losses_by_micro_batch = []
+        losses_by_step = []
+        num_iterations = 20
+
+        for _ in range(num_iterations):
+            for _, (data, label) in enumerate(loader):
+                schedule.step(data, target=label, losses=losses_by_micro_batch)
+                if self.rank == 3:
+                    losses_by_step.append(
+                        np.array(losses_by_micro_batch, dtype=np.float32).mean()
+                    )
+            opt.step()
+            opt.clear_grad()
+        return losses_by_step
+
+    def test_pp_model(self):
+        """Test pipeline parallel model using PPModel as the baseline"""
+        fix_seeds()
+        name_prefix = "pp_model"
+        shared_params_names = [
+            [
+                f"{name_prefix}_linear_0_weight.dist",
+                f"{name_prefix}_linear_7_weight.dist",
+            ]
+        ]
+        pp_model = PPModel(
+            name_prefix=name_prefix, shared_parameters=shared_params_names
+        )
+        opt = paddle.optimizer.AdamW(
+            learning_rate=0.001, parameters=pp_model.parameters()
+        )
+        loss_fn = nn.MSELoss()
+        dataset = RandomDataset(image_size=8, output_size=8, num_samples=8)
+        loader = DataLoader(dataset, batch_size=1)
+        pp_losses_step = []
+        num_iterations = 20
+
+        for _ in range(num_iterations):
+            pp_losses_micro_batch = []
+            for _, (data, label) in enumerate(loader):
+                output = pp_model(data)
+                loss = loss_fn(output, label)
+                pp_losses_micro_batch.append(loss.item())
+                loss.backward()
+            pp_losses_step.append(
+                np.array(pp_losses_micro_batch, dtype=np.float32).mean()
+            )
+            opt.step()
+            opt.clear_grad()
+        return pp_losses_step
+
+    def run_test(self):
+        """Compare shared params losses between three training methods"""
+        self.setUpClass()
+        pp_losses = self.test_pp_model()
+        pp_FThenB_losses = self.test_single_schedule(sing_schedule="FThenB")
+        pp_1F1B_losses = self.test_single_schedule(sing_schedule="1F1B")
+        pp_vpp_losses = self.test_multi_schedule(multi_schedule="VPP")
+
+        if self.rank == 3:
+            np.testing.assert_allclose(
+                pp_losses,
+                pp_FThenB_losses,
+                rtol=rtol,
+            )
+            np.testing.assert_allclose(
+                pp_losses,
+                pp_1F1B_losses,
+                rtol=rtol,
+            )
+            np.testing.assert_allclose(
+                pp_losses,
+                pp_vpp_losses,
+                rtol=rtol,
+            )
+
+
+if __name__ == '__main__':
+    TestSharedParameters().run_test()