PaddlePaddle
diff --git a/‎python/paddle/distributed/__init__.py
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎python/paddle/distributed/__init__.py
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/api.py
Copy file name to clipboard
+164Lines changed: 164 additions & 0 deletions b/‎python/paddle/distributed/auto_parallel/api.py
Copy file name to clipboard
+164Lines changed: 164 additions & 0 deletions
diff --git a/‎test/auto_parallel/CMakeLists.txt
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎test/auto_parallel/CMakeLists.txt
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎test/auto_parallel/hybrid_strategy/CMakeLists.txt
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎test/auto_parallel/hybrid_strategy/CMakeLists.txt
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎test/auto_parallel/hybrid_strategy/semi_auto_parallel_sharding_stage_1.py
Copy file name to clipboard
+84Lines changed: 84 additions & 0 deletions b/‎test/auto_parallel/hybrid_strategy/semi_auto_parallel_sharding_stage_1.py
Copy file name to clipboard
+84Lines changed: 84 additions & 0 deletions
diff --git a/‎test/auto_parallel/hybrid_strategy/semi_auto_parallel_sharding_stage_3.py
Copy file name to clipboard
+84Lines changed: 84 additions & 0 deletions b/‎test/auto_parallel/hybrid_strategy/semi_auto_parallel_sharding_stage_3.py
Copy file name to clipboard
+84Lines changed: 84 additions & 0 deletions
@@ -86,6 +86,8 @@
     shard_layer,
     shard_optimizer,
     shard_scaler,
+    ShardingStage1,
+    ShardingStage3,
     to_static,
     Strategy,
     DistModel,
@@ -171,6 +173,8 @@
     "load_state_dict",
     "shard_optimizer",
     "shard_scaler",
+    "ShardingStage1",
+    "ShardingStage3",
     "to_static",
     "Strategy",
     "DistModel",
 
@@ -551,6 +551,35 @@ def replicate_layer_params_and_buffers(
         )
 
 
+def get_placement_with_sharding(param):
+    shard_axis = -1
+    for placement in param.placements:
+        if isinstance(placement, dist.Shard):
+            # the parameter can't be shard twice on different mesh now
+            # assert here in case
+            assert (
+                shard_axis == -1
+            ), "The parameter can't be shard twich even in different mesh now."
+            shard_axis = placement.get_dim()
+
+    placement_with_sharding = None
+    for dim in range(param.ndim):
+        if dim != shard_axis:
+            placement_with_sharding = dist.Shard(dim)
+
+    new_placements = param.placements
+    for mesh_axis, placement in enumerate(param.placements):
+        # we need to keep the placement replicate if the it is out of tensor's dim
+        if (
+            isinstance(placement, dist.Replicate)
+            and placement_with_sharding is not None
+        ):
+            new_placements[mesh_axis] = placement_with_sharding
+            break
+
+    return new_placements
+
+
 class _ShardOptimizer:
     def __init__(self, optimizer, shard_fn=None):
         assert (
@@ -576,6 +605,13 @@ def __init__(self, optimizer, shard_fn=None):
         self._inner_opt = optimizer
         self._shard_fn = shard_fn
 
+        # Invoke shard_fn if it is not None to shard parameters
+        if self._shard_fn is not None and isinstance(
+            self._shard_fn, ShardingStage3
+        ):
+            for param in self._inner_opt._parameter_list:
+                self._shard_fn._shard_parameter(param)
+
     def _shard_accumulator(self, param):
         # create the accumulators
         self._inner_opt._create_accumulators(self.target_block, [param])
@@ -733,6 +769,134 @@ def __getattr__(self, item):
         return getattr(self._inner_opt, item)
 
 
+class ShardingStage1:
+    """
+    A builtin shard_fn for shard_optimizer interface, users can pass it to shard_optimizer to implement sharding optimization with stage 1.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> batch = paddle.rand(shape=[8, 8])
+            >>> opt = paddle.optimizer.AdamW(parameters=layer.parameters())
+            >>> opt = dist.shard_optimizer(opt, dist.ShardingStage1())
+            >>> for _ in range(5):
+            >>>     loss = layer(batch)
+            >>>     loss.backward()
+            >>>     opt.step()
+            >>>     opt.clear_grad()
+            >>> # This case need to be executed in multi-card environment
+            >>> # python -m paddle.distributed.launch --gpus=0,1 {test_case}.py
+    """
+
+    def __call__(self, key, param, accumulator):
+        if param.is_dist():
+            # Only deal with momentum in optimizer, beta should be replicated cross param's mesh
+            if 'beta' not in key:
+                placements = get_placement_with_sharding(param)
+            else:
+                placements = [
+                    dist.Replicate()
+                    for _ in range(len(param.process_mesh.shape))
+                ]
+            return shard_tensor(
+                accumulator,
+                mesh=param.process_mesh,
+                placements=placements,
+            )
+        return accumulator
+
+
+class ShardingStage3:
+    """
+    A builtin shard_fn for shard_optimizer interface, users can pass it to shard_optimizer to implement sharding optimization with stage 3.
+
+    Args:
+        mesh(paddle.distributed.ProcessMesh): The `ProcessMesh` object describes the Cartesian topology of the used processes.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+            >>> class MLP(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.fc1 = paddle.nn.Linear(8, 8)
+            ...         self.fc2 = paddle.nn.Linear(8, 8)
+            ...
+            ...     def forward(self, input):
+            ...         return self.fc2(self.fc1(input))
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> layer = MLP()
+            >>> batch = paddle.rand(shape=[8, 8])
+            >>> opt = paddle.optimizer.AdamW(parameters=layer.parameters())
+            >>> opt = dist.shard_optimizer(opt, dist.ShardingStage3(mesh))
+            >>> for _ in range(5):
+            >>>     loss = layer(batch)
+            >>>     loss.backward()
+            >>>     opt.step()
+            >>>     opt.clear_grad()
+            >>> # This case need to be executed in multi-card environment
+            >>> # python -m paddle.distributed.launch --gpus=0,1 {test_case}.py
+    """
+
+    def __init__(self, mesh):
+        self._mesh = mesh
+
+    def _shard_parameter(self, param):
+        # TODO(liyurui): remove this trick dense to dist convert after adding
+        # dense_tensor.to_dist method.
+        if param.is_dense():
+            zero_dense = paddle.zeros(param.shape)
+            placements = []
+            for _ in range(len(self._mesh.shape)):
+                placements.append(dist.Replicate())
+            zero_dist = dist.shard_tensor(zero_dense, self._mesh, placements)
+            res = param + zero_dist
+
+        new_placements = get_placement_with_sharding(param)
+        shard_param = dist.reshard(param, param.process_mesh, new_placements)
+        # change the holder of param to new shard_param
+        param.get_tensor()._share_data_with(shard_param.get_tensor())
+
+    def __call__(self, key, param, accumulator):
+        if param.is_dist():
+            # Only deal with momentum in optimizer, beta should be replicated cross param's mesh
+            if 'beta' not in key:
+                placements = param.placements
+            else:
+                placements = [
+                    dist.Replicate()
+                    for _ in range(len(param.process_mesh.shape))
+                ]
+            return shard_tensor(
+                accumulator,
+                mesh=param.process_mesh,
+                placements=placements,
+            )
+        return accumulator
+
+
 def shard_optimizer(optimizer, shard_fn=None):
     """
 
 
@@ -166,6 +166,10 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
                   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_lazy_init MODULES
                   test_semi_auto_parallel_lazy_init)
   set_tests_properties(test_semi_auto_parallel_lazy_init
 
@@ -50,3 +50,11 @@ if((WITH_GPU) AND (LINUX))
   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()
@@ -0,0 +1,84 @@
+# 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
+#
+#     https://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 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.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_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())
+        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 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()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelShardingStage1().run_test_case()
@@ -0,0 +1,84 @@
+# 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
+#
+#     https://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 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(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 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()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelShardingStage3().run_test_case()