chore: import upstream snapshot with attribution

python/paddle/distributed/fleet/meta_parallel/__init__.py

@@ -0,0 +1,57 @@
+# 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.
+
+# isort: skip_file
+from .meta_parallel_base import MetaParallelBase  # noqa: F401
+
+from .parallel_layers import (  # noqa: F401
+    ColumnParallelLinear,
+    LayerDesc,
+    LocalSharedLayerDesc,
+    ParallelCrossEntropy,
+    PipelineLayer,
+    RNGStatesTracker,
+    RowParallelLinear,
+    SharedLayerDesc,
+    VocabParallelEmbedding,
+    get_rng_state_tracker,
+    model_parallel_random_seed,
+    LayerSpec,
+    import_spec_layer,
+    get_spec_layer,
+    build_spec_layer,
+)
+from .pipeline_parallel import (  # noqa: F401
+    NoPipelineParallel,
+    PipelineParallel,
+    PipelineParallelMicroStepLocations,
+    PipelineParallelWithInterleave,
+    PipelineParallelWithInterleaveFthenB,
+    PipelineDatasetPreprocessor,
+    VPPFhenBInBalancedMemory,
+    register_global_pipeline_parallel_hook,
+)
+from .dualpipev import DualPipeVParallel  # noqa: F401
+from .segment_parallel import SegmentParallel  # noqa: F401
+from .sharding_parallel import ShardingParallel  # noqa: F401
+from .tensor_parallel import TensorParallel  # noqa: F401
+from .pp_utils.forward_backward_overlap_utils import (  # noqa: F401
+    ScheduleNode,
+    ScheduleChunk,
+)
+from .pp_utils.utils import (  # noqa: F401
+    dict_to_tuple_helper,
+)
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/dualpipev.py

@@ -0,0 +1,851 @@
+# 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.
+
+# The file has been adapted from DeepSeek DualPipe project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DualPipe/blob/main/LICENSE
+
+from __future__ import annotations
+
+import paddle
+from paddle import framework
+from paddle.distributed.communication.batch_isend_irecv import (
+    P2POp,
+    batch_isend_irecv,
+)
+
+try:
+    from paddle.distributed.communication import deep_ep
+except ImportError:
+    deep_ep = None
+
+from ..utils.log_util import logger
+from .pipeline_parallel import (
+    FakeMicroDataset,
+    HybridParallelOptimizer,
+    PipelineDatasetPreprocessor,
+    PipelineParallel,
+)
+from .pp_utils.batch_comm_helper import BatchCommHelper
+from .pp_utils.forward_backward_overlap_utils import ScheduleChunk
+from .zero_bubble_utils import EventStore, WeightGradStore
+
+__all__ = []
+
+
+def detach_and_requires_grad(x):
+    o = x.detach()
+    o.stop_gradient = False
+    return o
+
+
+class DualPipeVParallel(PipelineParallel):
+    """
+    An implementation of the DualPipeV, based on
+    https://github.com/deepseek-ai/DualPipe/blob/main/dualpipe/dualpipe.py.
+    """
+
+    def __init__(self, layers, hcg, strategy):
+        super().__init__(layers=layers, hcg=hcg, strategy=strategy)
+        self.overlapped_forward_backward = hasattr(
+            type(self._layers), "overlapped_forward_backward"
+        )
+        logger.info(
+            f"Using DualPipeVParallel with overlapping forward backward={self.overlapped_forward_backward}"
+        )
+
+        self.num_ranks = self.num_stages
+        self.group_rank = self.pp_group.rank
+        self.prev_rank = self.pp_group.ranks[
+            (self.group_rank - 1) % self.pp_group.world_size
+        ]
+        self.next_rank = self.pp_group.ranks[
+            (self.group_rank + 1) % self.pp_group.world_size
+        ]
+
+        # NOTE(zhangyuqin1998): The first rank has to broadcast the meta information
+        # of the P2P communication after the first forward.
+        self.need_broadcast_meta = self.is_pipeline_first_stage()
+        self.need_recv_meta = not self.is_pipeline_first_stage()
+        self._p2p_helper = BatchCommHelper(self._using_cache)
+
+    def is_pipeline_first_stage(self):
+        return self.group_rank == 0
+
+    def is_pipeline_last_stage(self):
+        return self.group_rank == self.num_ranks - 1
+
+    def _reset_states(self):
+        self.input_tensors = ([], [])
+        self.output_tensors = ([], [])
+        self.input_grad_tensors = ([], [])
+        self.output_grad_tensors = ([], [])
+        self.loss_tensors: list[paddle.Tensor] = []
+        self.schedule_chunks = ([], [])
+        self.loss_fn_chunks = []
+
+        # The first value in the list corresponds to phase 0, and the second value corresponds to phase 1.
+        self.current_f_acc_id = [0, 0]
+        self.current_b_acc_id = [0, 0]
+        self.current_send_f_acc_id = [0, 0]
+        self.current_send_b_acc_id = [0, 0]
+        self.current_recv_f_acc_id = [0, 0]
+        self.current_recv_b_acc_id = [0, 0]
+        self.comm_forward_ops: list[P2POp] = []
+        self.comm_backward_ops: list[P2POp] = []
+        self.to_free: list[paddle.Tensor] = []
+
+    def _get_forward_inputs(self, micro_datasets, phase, acc_id):
+        is_first_stage = self.is_pipeline_first_stage() and phase == 0
+        if is_first_stage:
+            assert micro_datasets is not None
+            self.input_tensors[phase].append(next(micro_datasets[phase])[0])
+        if self.forward_only:
+            self.input_tensors[phase][acc_id] = None
+        return self.input_tensors[phase][acc_id]
+
+    def _get_forward_labels(self, micro_datasets, phase, acc_id):
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+        if is_last_stage and self._compute_loss:
+            assert micro_datasets is not None
+            labels = next(micro_datasets[phase])[1]
+            self._check_micro_batch_data_valid(labels)
+            return labels
+        else:
+            return None
+
+    def _loss_compute(self, micro_datasets, phase, acc_id, logits):
+        labels = self._get_forward_labels(micro_datasets, phase, acc_id)
+        loss_fn_node = None
+        if not self.overlapped_forward_backward:
+            loss_tensor = self._layers._loss_fn[0](logits, labels)
+        else:
+            loss_fn_node = self._layers._loss_fn[0].build_schedule_node()
+            loss_fn_node.labels = labels
+            loss_tensor = loss_fn_node.forward(logits)
+        self._store_forward_loss(phase, loss_tensor, loss_fn_node)
+
+    def _store_forward_tensors(self, phase, outputs, schedule_chunk):
+        self.schedule_chunks[phase].append(schedule_chunk)
+        if self.is_pipeline_last_stage() and phase == 0:
+            self.input_tensors[1].append(
+                [detach_and_requires_grad(output) for output in outputs]
+            )
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+        if not is_last_stage:
+            self.output_tensors[phase].append(outputs)
+
+    def _forward_compute(self, phase: int, micro_datasets=None) -> None:
+        acc_id = self.current_f_acc_id[phase]
+        self.current_f_acc_id[phase] += 1
+
+        inputs = self._get_forward_inputs(micro_datasets, phase, acc_id)
+
+        if self.overlapped_forward_backward:
+            schedule_chunk = self._layers.get_schedule_chunk(chunk_id=phase)
+            outputs = schedule_chunk.forward(inputs)
+        else:
+            schedule_chunk = None
+            outputs = self._layers.forward(inputs, chunk_id=phase)
+        outputs = [outputs] if isinstance(outputs, paddle.Tensor) else outputs
+
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+        if is_last_stage and self._compute_loss:
+            self._loss_compute(micro_datasets, phase, acc_id, outputs)
+        self._store_forward_tensors(phase, outputs, schedule_chunk)
+
+    def _get_backward_inputs(self, phase, acc_id):
+        outputs = self.output_tensors[phase][acc_id]
+        self.output_tensors[phase][acc_id] = None
+        output_grads = self.output_grad_tensors[phase][acc_id]
+        self.output_grad_tensors[phase][acc_id] = None
+        non_empty = [
+            (t, g) for t, g in zip(outputs, output_grads) if g is not None
+        ]
+        outputs, output_grads = list(zip(*non_empty))
+        return outputs, output_grads
+
+    def _store_backward_tensors(self, phase, acc_id, input_grads=None):
+        if input_grads is None:
+            inputs = self.input_tensors[phase][acc_id]
+            input_grads = [
+                t.grad
+                for t in inputs
+                if (t is not None and not t.stop_gradient)
+            ]
+        self.input_tensors[phase][acc_id] = None
+
+        if isinstance(input_grads, paddle.Tensor):
+            input_grads = (input_grads,)
+        if self.is_pipeline_last_stage() and phase == 1:
+            self.output_grad_tensors[0].append(input_grads)
+        else:
+            self.input_grad_tensors[phase].append(input_grads)
+
+    def _store_forward_loss(self, phase, loss_tensor, loss_fn_node=None):
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+        if is_last_stage and self._compute_loss:
+            if isinstance(loss_tensor, (tuple, list)):
+                assert len(loss_tensor) == 1
+                loss_tensor = loss_tensor[0]
+            assert isinstance(loss_tensor, paddle.Tensor), (
+                "Currently, loss_fn should obtain Paddle.Tensor dtype"
+            )
+
+            self.loss_tensors.append(loss_tensor)
+            self.loss_fn_chunks.append(loss_fn_node)
+
+    def _backward_compute(self, phase: int, enable_zb: bool = False) -> None:
+        if self.forward_only:
+            return
+
+        acc_id = self.current_b_acc_id[phase]
+        self.current_b_acc_id[phase] += 1
+
+        is_last_stage = self.is_pipeline_first_stage() and phase == 1
+
+        WeightGradStore.enabled = enable_zb
+        input_grads = None
+        with paddle.amp.auto_cast(enable=False):
+            if is_last_stage:
+                loss = self.loss_tensors[acc_id]
+                if self.overlapped_forward_backward:
+                    loss_fn_node = self.loss_fn_chunks[acc_id]
+                    backward_chunk = self.schedule_chunks[phase][acc_id]
+                    _, _, input_grads = (
+                        self._layers.overlapped_forward_backward(
+                            ScheduleChunk([]),  # forward_chunk
+                            None,  # forward_inputs
+                            None,  # forward_loss_fn_node
+                            backward_chunk,
+                            loss_fn_node,
+                            None,  # input_grads
+                            self.scaler,
+                            combine_bw_event_to_wait=None,
+                            pp_stream=None,
+                        )
+                    )
+                    self.loss_fn_chunks[acc_id] = None
+                    self.schedule_chunks[phase][acc_id] = None
+                else:
+                    if self.scaler:
+                        paddle.autograd.backward(self.scaler.scale(loss))
+                    else:
+                        paddle.autograd.backward(loss)
+            else:
+                outputs, output_grads = self._get_backward_inputs(phase, acc_id)
+                if self.overlapped_forward_backward:
+                    backward_chunk = self.schedule_chunks[phase][acc_id]
+                    _, _, input_grads = (
+                        self._layers.overlapped_forward_backward(
+                            ScheduleChunk([]),  # forward_chunk
+                            None,  # forward_inputs
+                            None,  # forward_loss_fn_node
+                            backward_chunk,
+                            None,  # backward_loss_fn_node
+                            output_grads,
+                            None,  # scaler
+                            combine_bw_event_to_wait=None,
+                            pp_stream=None,
+                        )
+                    )
+                    self.schedule_chunks[phase][acc_id] = None
+                else:
+                    if len(outputs) > 0:
+                        outputs = [t for t in outputs if not t.stop_gradient]
+                        paddle.autograd.backward(
+                            tensors=outputs,
+                            grad_tensors=output_grads,
+                        )
+        WeightGradStore.enabled = False
+        if enable_zb:
+            WeightGradStore.flush()
+
+        self._store_backward_tensors(phase, acc_id, input_grads=input_grads)
+
+    def _forward_backward_compute(
+        self,
+        forward_phase: int,
+        backward_phase: int,
+        micro_datasets=None,
+        combine_backward_event_to_wait=None,
+        pass_pp_stream=False,
+    ) -> None:
+        if self.forward_only:
+            self._forward_compute(forward_phase, micro_datasets)
+            return
+
+        if not self.overlapped_forward_backward:
+            self._forward_compute(forward_phase, micro_datasets)
+            self._backward_compute(backward_phase)
+            return
+
+        # pre-forward
+        forward_acc_id = self.current_f_acc_id[forward_phase]
+        self.current_f_acc_id[forward_phase] += 1
+
+        forward_inputs = self._get_forward_inputs(
+            micro_datasets, forward_phase, forward_acc_id
+        )
+        forward_labels = self._get_forward_labels(
+            micro_datasets, forward_phase, forward_acc_id
+        )
+        if forward_labels is not None:
+            forward_loss_fn_node = self._layers._loss_fn[
+                0
+            ].build_schedule_node()
+            forward_loss_fn_node.labels = forward_labels
+        else:
+            forward_loss_fn_node = None
+
+        # pre-backward
+        backward_acc_id = self.current_b_acc_id[backward_phase]
+        self.current_b_acc_id[backward_phase] += 1
+
+        is_last_stage1 = self.is_pipeline_first_stage() and backward_phase == 1
+        if is_last_stage1:
+            backward_loss_fn_node = self.loss_fn_chunks[backward_acc_id]
+            backward_grads = None
+        else:
+            backward_loss_fn_node = None
+            _, backward_grads = self._get_backward_inputs(
+                backward_phase, backward_acc_id
+            )
+
+        # event_to_wait = deep_ep.get_event_from_custom_stream(paddle.device.current_stream().stream_base)
+
+        # forward & backward
+        forward_chunk = self._layers.get_schedule_chunk(chunk_id=forward_phase)
+        backward_chunk = self.schedule_chunks[backward_phase][backward_acc_id]
+        forward_outputs, forward_loss, backward_input_grads = (
+            self._layers.overlapped_forward_backward(
+                forward_chunk,
+                forward_inputs,
+                forward_loss_fn_node,
+                backward_chunk,
+                backward_loss_fn_node,
+                backward_grads,
+                self.scaler,
+                combine_bw_event_to_wait=combine_backward_event_to_wait,
+                pp_stream=(
+                    self.pp_group.process_group.get_stream(
+                        paddle.framework._current_expected_place_()
+                    )
+                    if pass_pp_stream
+                    else None
+                ),
+            )
+        )
+        self.schedule_chunks[backward_phase][backward_acc_id] = None
+
+        # post-forward
+        self._store_forward_tensors(
+            forward_phase, forward_outputs, forward_chunk
+        )
+        self._store_forward_loss(
+            forward_phase, forward_loss, forward_loss_fn_node
+        )
+
+        # post-backward
+        self._store_backward_tensors(
+            backward_phase, backward_acc_id, input_grads=backward_input_grads
+        )
+
+    def _commit_and_wait_comm(
+        self, p2p_overlap=False, use_outer_event_wait=False
+    ) -> None:
+        common_forward_ops_num = (
+            len(self.comm_forward_ops)
+            if self.comm_forward_ops is not None
+            else 0
+        )
+        common_backward_ops_num = (
+            len(self.comm_backward_ops)
+            if self.comm_backward_ops is not None
+            else 0
+        )
+        if common_forward_ops_num == 0 and common_backward_ops_num == 0:
+            if EventStore.event is not None:
+                e_t = EventStore.event
+                EventStore.event = None
+                return e_t
+            return deep_ep.get_event_from_custom_stream(
+                paddle.device.current_stream().stream_base
+            )
+
+        use_stream_wait_event = (
+            p2p_overlap and self._overlap_p2p_comm and deep_ep is not None
+        )
+
+        pp_raw_stream = self.pp_group.process_group.get_stream(
+            paddle.framework._current_expected_place_()
+        )
+        if use_outer_event_wait:
+            self.pp_group.process_group.set_outer_wait(True)
+
+        if common_forward_ops_num > 0:
+            fwd_reqs = batch_isend_irecv(self.comm_forward_ops)
+
+            if not use_stream_wait_event:
+                for req in fwd_reqs:
+                    req.wait()
+
+        if use_outer_event_wait:
+            self.pp_group.process_group.set_outer_wait(False)
+
+        if use_stream_wait_event:
+            forward_event_to_wait = deep_ep.get_event_from_custom_stream(
+                pp_raw_stream
+            )
+
+        backward_outer_event_wait = False
+        if EventStore.event is not None:
+            with paddle.device.stream_guard(
+                paddle.device.Stream(stream_base=pp_raw_stream)
+            ):
+                EventStore.event.current_stream_wait()
+
+            EventStore.set(None)
+            self.pp_group.process_group.set_outer_wait(True)
+
+            backward_outer_event_wait = True
+
+        if common_backward_ops_num > 0:
+            bwd_reqs = batch_isend_irecv(self.comm_backward_ops)
+
+            if not use_stream_wait_event:
+                for req in bwd_reqs:
+                    req.wait()
+
+        if backward_outer_event_wait:
+            self.pp_group.process_group.set_outer_wait(False)
+
+        if use_stream_wait_event:
+            forward_event_to_wait.current_stream_wait()
+
+            combine_bw_event_to_wait = deep_ep.get_event_from_custom_stream(
+                pp_raw_stream
+            )
+        else:
+            combine_bw_event_to_wait = deep_ep.get_event_from_custom_stream(
+                paddle.device.current_stream().stream_base
+            )
+
+        self.comm_forward_ops = []
+        self.comm_backward_ops = []
+
+        self._free_tensors()
+
+        return combine_bw_event_to_wait
+
+    def _weight_pass(self) -> None:
+        if self.forward_only:
+            return
+
+        self._commit_and_wait_comm()
+
+        # Assume FIFO
+        WeightGradStore.pop()
+
+    def _free_tensors(self) -> None:
+        self._release_output(self.to_free)
+        self.to_free = []
+
+    def _recv_forward(self, phase: int) -> None:
+        if (self.is_pipeline_first_stage() and phase == 0) or (
+            self.is_pipeline_last_stage() and phase == 1
+        ):
+            return
+
+        self.current_recv_f_acc_id[phase] += 1
+
+        tensors = self._p2p_helper.append_irecv(
+            self.comm_forward_ops,
+            self.prev_rank if phase == 0 else self.next_rank,
+            self.pp_group,
+            alloc_on_comm_stream=self._overlap_p2p_comm,
+        )
+        self.input_tensors[phase].append(tensors)
+
+    def _send_forward(self, phase: int) -> None:
+        if (self.is_pipeline_first_stage() and phase == 1) or (
+            self.is_pipeline_last_stage() and phase == 0
+        ):
+            return
+
+        acc_id = self.current_send_f_acc_id[phase]
+        self.current_send_f_acc_id[phase] += 1
+        tensors = self.output_tensors[phase][acc_id]
+
+        self._p2p_helper.append_isend(
+            self.comm_forward_ops,
+            tensors,
+            self.next_rank if phase == 0 else self.prev_rank,
+            self.pp_group,
+            self.need_broadcast_meta,
+        )
+        self.need_broadcast_meta = False
+
+        self.to_free.extend(tensors)
+
+    def _recv_backward(self, phase: int) -> None:
+        if self.forward_only:
+            return
+
+        if (self.is_pipeline_first_stage() and phase == 1) or (
+            self.is_pipeline_last_stage() and phase == 0
+        ):
+            return
+
+        self.current_recv_b_acc_id[phase] += 1
+        tensors = self._p2p_helper.append_irecv(
+            self.comm_backward_ops,
+            self.next_rank if phase == 0 else self.prev_rank,
+            self.pp_group,
+            alloc_on_comm_stream=self._overlap_p2p_comm,
+        )
+        self.output_grad_tensors[phase].append(tensors)
+
+    def _send_backward(self, phase: int) -> None:
+        if self.forward_only:
+            return
+
+        if (self.is_pipeline_first_stage() and phase == 0) or (
+            self.is_pipeline_last_stage() and phase == 1
+        ):
+            return
+
+        acc_id = self.current_send_b_acc_id[phase]
+        self.current_send_b_acc_id[phase] += 1
+        tensors = self.input_grad_tensors[phase][acc_id]
+        self.input_grad_tensors[phase][acc_id] = None
+
+        self._p2p_helper.append_isend(
+            self.comm_backward_ops,
+            tensors,
+            self.prev_rank if phase == 0 else self.next_rank,
+            self.pp_group,
+        )
+
+    def _forward_pass(
+        self,
+        phase: int,
+        micro_datasets=None,
+        recv: bool = True,
+        send: bool = True,
+    ) -> None:
+        if recv:
+            self._recv_forward(phase)
+        self._commit_and_wait_comm()
+
+        self._forward_compute(phase, micro_datasets)
+
+        if send:
+            self._send_forward(phase)
+
+    def _backward_pass(
+        self,
+        phase: int,
+        enable_zb: bool = False,
+        recv: bool = True,
+        send: bool = True,
+    ) -> None:
+        if recv:
+            self._recv_backward(phase)
+        self._commit_and_wait_comm()
+
+        self._backward_compute(phase, enable_zb)
+
+        if send:
+            self._send_backward(phase)
+
+    def _forward_backward_pass(
+        self,
+        forward_phase: int,
+        backward_phase: int,
+        micro_datasets=None,
+        recv0: bool = True,
+        first_chunk=False,
+        last_chunk=False,
+        main_stage=False,
+        last_stage_and_first_chunk=False,
+    ) -> None:
+        if recv0:
+            self._recv_forward(forward_phase)
+        self._recv_backward(backward_phase)
+
+        need_send_forward = not (
+            self.is_pipeline_first_stage() and forward_phase == 1
+        ) or (self.is_pipeline_last_stage() and forward_phase == 0)
+        need_send_backward = not (
+            self.is_pipeline_first_stage() and backward_phase == 0
+        ) or (self.is_pipeline_last_stage() and backward_phase == 1)
+
+        use_outer_event_wait = (
+            main_stage
+            and not first_chunk
+            and self._overlap_p2p_comm
+            and deep_ep is not None
+            and (need_send_forward and need_send_backward)
+        )
+
+        pass_pp_stream = (
+            main_stage
+            and not last_chunk
+            and self._overlap_p2p_comm
+            and deep_ep is not None
+            and (need_send_forward and need_send_backward)
+            and (not last_stage_and_first_chunk)
+        )
+
+        combine_bw_wait_event = self._commit_and_wait_comm(
+            not last_chunk, use_outer_event_wait
+        )
+
+        self._forward_backward_compute(
+            forward_phase,
+            backward_phase,
+            micro_datasets,
+            combine_backward_event_to_wait=combine_bw_wait_event,
+            pass_pp_stream=pass_pp_stream,
+        )
+
+        self._send_forward(forward_phase)
+        self._send_backward(backward_phase)
+
+    def _wrap_data(self, data, phase):
+        """
+        for backward compatibility, wrap data to Fake FakeMicroDataset if it is of type list or tuple
+        """
+        if isinstance(data, PipelineDatasetPreprocessor):
+            data = data()
+
+        if (not isinstance(data, tuple)) and (not isinstance(data, list)):
+            return data
+
+        micro_dataset = FakeMicroDataset(
+            data,
+            self.is_pipeline_first_stage() and phase == 0,
+            self.is_pipeline_first_stage() and phase == 1,
+            self.accumulate_steps,
+            self.micro_batch_size,
+        )
+        return micro_dataset
+
+    def _prepare_training(self, data, optimizer, lr_scheduler):
+        assert isinstance(optimizer, HybridParallelOptimizer), (
+            'optimizer should be HybridParallelOptimizer subclass.'
+        )
+
+        assert framework._dygraph_tracer()._has_grad, (
+            'Please enable the generation of gradients.'
+        )
+
+        if self.is_pipeline_first_stage():
+            assert data is not None, (
+                "For the first and the last stage, the data must be set."
+            )
+        else:
+            data = None
+
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+        self._layers.train()
+        self.register_sharding_comm_overlap_hook(optimizer)
+
+        return data
+
+    def _broadcast_final_loss(self):
+        loss_sum_tensor = paddle.zeros([1], "float32")
+        if self.is_pipeline_first_stage():
+            assert len(self.loss_tensors) > 0, (
+                "train_batch() in last stage should obtain valid loss"
+            )
+            for loss in self.loss_tensors:
+                loss_sum_tensor += loss.detach().astype("float32")
+            loss_sum_tensor /= self.accumulate_steps
+
+        paddle.distributed.all_reduce(
+            loss_sum_tensor, group=self.pp_group, sync_op=True
+        )
+        return loss_sum_tensor
+
+    def forward_backward_pipeline(
+        self,
+        data,
+        scaler,
+        forward_only=False,
+        compute_loss=True,
+    ):
+        self.scaler = scaler
+
+        rank = self.group_rank
+        num_ranks = self.num_ranks
+        assert (
+            self.accumulate_steps > 0 and self.accumulate_steps >= num_ranks * 2
+        ), f"{self.accumulate_steps=}, {num_ranks=}"
+        self.forward_only = forward_only
+
+        self._reset_states()
+
+        # NOTE(zhangyuqin1998): Tensors to be sent or received must have a
+        # consistent shape and data type throughout the entire pipeline. We
+        # broadcast the meta info in the first forward of the first rank.
+        self._p2p_helper.recv_meta_from_head(self.pp_group, self.need_recv_meta)
+        self.need_recv_meta = False
+
+        micro_dataset_phase0 = self._wrap_data(data, 0)
+        micro_dataset_phase1 = self._wrap_data(data, 1)
+        micro_datasets = [micro_dataset_phase0, micro_dataset_phase1]
+
+        # Step 1: nF0
+        step_1 = (num_ranks - rank - 1) * 2
+        for i in range(step_1):
+            self._forward_pass(0, micro_datasets)
+
+        # Step 2: nF0F1
+        step_2 = rank + 1
+        self._recv_forward(0)
+        for i in range(step_2):
+            self._forward_pass(0, micro_datasets, recv=False, send=False)
+            self._recv_forward(0)
+            self._forward_pass(
+                1,
+                micro_datasets,
+                send=(not self.is_pipeline_last_stage()) or (i < step_2 - 1),
+            )
+            self._send_forward(0)
+
+        # Step 3: nB1W1F1 (Use zero bubble)
+        step_3 = num_ranks - rank - 1
+        for i in range(step_3):
+            self._backward_pass(1, enable_zb=True)
+            self._recv_forward(1)
+            self._weight_pass()
+            self._forward_pass(1, micro_datasets, recv=False)
+
+        # Step 4 (Main step): nF0B1F1B0
+        step_4 = self.accumulate_steps - num_ranks * 2 + rank + 1
+        have_step5 = num_ranks - rank - 1 > 0
+        # Update code to support send/recv overlap
+        # Only support send/recv overlap in MainStep
+        for i in range(step_4):
+            is_last_chunk = i + 1 == step_4
+            if i == 0:
+                if self.is_pipeline_last_stage():
+                    # NOTE: We don't overlap these two passes to further reduce bubble size.
+                    self._forward_pass(
+                        0, micro_datasets, recv=False, send=False
+                    )
+                    self._send_forward(1)
+                    self._backward_pass(1, send=False)
+                    self._send_forward(0)
+                    self._send_backward(1)
+
+                    self._forward_backward_pass(
+                        1,
+                        0,
+                        micro_datasets,
+                        first_chunk=True,
+                        last_chunk=is_last_chunk,
+                        main_stage=True,
+                    )
+                else:
+                    self._forward_backward_pass(
+                        0,
+                        1,
+                        micro_datasets,
+                        recv0=False,
+                        first_chunk=True,
+                        main_stage=True,
+                    )
+
+                    self._forward_backward_pass(
+                        1,
+                        0,
+                        micro_datasets,
+                        last_chunk=is_last_chunk,
+                        main_stage=True,
+                    )
+            else:
+                self._forward_backward_pass(
+                    0,
+                    1,
+                    micro_datasets,
+                    main_stage=True,
+                    last_stage_and_first_chunk=self.is_pipeline_last_stage(),
+                )
+                self._forward_backward_pass(
+                    1,
+                    0,
+                    micro_datasets,
+                    last_chunk=is_last_chunk,
+                    main_stage=True,
+                )
+
+        # Step 5: nB1F1B0
+        step_5 = num_ranks - rank - 1
+        for i in range(step_5):
+            self._backward_pass(1)
+            self._forward_backward_pass(1, 0, micro_datasets)
+
+        # Step 6: nB1B0 (The second half of the passes use zero bubble)
+        step_6 = rank + 1
+        enable_zb = False
+        for i in range(step_6):
+            if i == step_6 // 2 and rank % 2 == 1:
+                enable_zb = True
+            self._backward_pass(1, enable_zb=enable_zb)
+            if i == step_6 // 2 and rank % 2 == 0:
+                enable_zb = True
+            self._backward_pass(0, enable_zb=enable_zb)
+
+        # Step 7: nWB0 (Use zero bubble)
+        step_7 = num_ranks - rank - 1
+        for i in range(step_7):
+            self._weight_pass()
+            self._backward_pass(0, enable_zb=True)
+
+        # Step 8: nW
+        step_8 = rank + 1
+        for i in range(step_8):
+            self._weight_pass()
+        assert WeightGradStore.funcs_queue.empty()
+
+        self._commit_and_wait_comm()
+
+        self._layers.allreduce_shared_weight_gradients()
+
+        with paddle.amp.auto_cast(enable=False):
+            train_loss = self._broadcast_final_loss()
+
+        self._reset_states()
+        return train_loss
+
+    def train_batch(
+        self,
+        data,
+        optimizer,
+        lr_scheduler=None,
+        scaler=None,
+    ):
+        data = self._prepare_training(data, optimizer, lr_scheduler)
+
+        train_loss = self.forward_backward_pipeline(data, scaler)
+
+        # optimizer
+        with paddle.amp.auto_cast(enable=False):
+            self._optimizer_step()
+
+        return train_loss

python/paddle/distributed/fleet/meta_parallel/meta_parallel_base.py

@@ -0,0 +1,44 @@
+#   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.
+
+from paddle import nn
+
+__all__ = []
+
+
+class MetaParallelBase(nn.Layer):
+    def __init__(self, layers, hcg, strategy):
+        super().__init__(layers.full_name() + "_meta_parallel_base")
+        self._layers = layers
+        self._hcg = hcg
+        self._strategy = strategy
+        self._prepare_for_model()
+
+    def _prepare_for_model(self):
+        pass
+
+    def _pre_forward(self, *inputs, **kwargs):
+        pass
+
+    def forward(self, *inputs, **kwargs):
+        self._pre_forward(*inputs, **kwargs)
+
+        output = self._layers(*inputs, **kwargs)
+
+        self._post_forward(output)
+
+        return output
+
+    def _post_forward(self, output):
+        pass

python/paddle/distributed/fleet/meta_parallel/parallel_layers/__init__.py

@@ -0,0 +1,39 @@
+# 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.
+
+from .mp_layers import (  # noqa: F401
+    ColumnParallelLinear,
+    ParallelCrossEntropy,
+    RowParallelLinear,
+    VocabParallelEmbedding,
+)
+from .pp_layers import (  # noqa: F401
+    LayerDesc,
+    LocalSharedLayerDesc,
+    PipelineLayer,
+    SharedLayerDesc,
+)
+from .random import (  # noqa: F401
+    RNGStatesTracker,
+    get_rng_state_tracker,
+    model_parallel_random_seed,
+)
+from .spec_utils import (
+    LayerSpec as LayerSpec,
+    build_spec_layer as build_spec_layer,
+    get_spec_layer as get_spec_layer,
+    import_spec_layer as import_spec_layer,
+)
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/parallel_layers/mp_layers.py

@@ -0,0 +1,22 @@
+#   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.
+
+from ...layers.mpu.mp_layers import (  # noqa: F401
+    ColumnParallelLinear,
+    ParallelCrossEntropy,
+    RowParallelLinear,
+    VocabParallelEmbedding,
+)
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/parallel_layers/random.py

@@ -0,0 +1,22 @@
+#   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.
+
+from ...layers.mpu.random import (  # noqa: F401
+    RNGStatesTracker,
+    dropout,
+    get_rng_state_tracker,
+    model_parallel_random_seed,
+)
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/parallel_layers/spec_utils.py

@@ -0,0 +1,141 @@
+# 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.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+from __future__ import annotations
+
+import types
+import warnings
+from dataclasses import dataclass, field
+
+
+@dataclass
+class LayerSpec:
+    """This is a Layer Specification dataclass.
+
+    Specification defines the location of the layer (to import dynamically)
+    or the imported layer itself. It also defines the extra_kwargs that need to be
+    passed to initialize the layer.
+
+    Args:
+        layer (tuple | type): A tuple describing the location of the
+            layer class e.g. `(layer.location, LayerClass)` or the imported
+            layer class itself e.g. `LayerClass` (which is already imported
+            using `from layer.location import LayerClass`).
+        extra_kwargs (dict): A dictionary of extra_kwargs that need to be passed while init.
+
+    """
+
+    layer: tuple | type
+    extra_kwargs: dict = field(default_factory=lambda: {})
+    sublayers_spec: type = None
+
+    def __repr__(self):
+        rst = ""
+        if isinstance(self.layer, tuple):
+            for sub_layer in self.layer:
+                rst = rst + repr(sub_layer) + ","
+        else:
+            rst = repr(self.layer) + repr(self.extra_kwargs)
+        return rst
+
+
+def import_spec_layer(layer_path: tuple[str]):
+    """Import a named object from a layer in the context of this function."""
+    base_path, name = layer_path
+    try:
+        layer = __import__(base_path, globals(), locals(), [name])
+    except ImportError as e:
+        print(f"couldn't import layer due to {e}")
+        return None
+    return vars(layer)[name]
+
+
+def get_spec_layer(spec_or_layer: LayerSpec | type, **additional_kwargs):
+    # If a layer class is already provided return it as is
+    if isinstance(spec_or_layer, (type, types.FunctionType)):
+        return spec_or_layer
+
+    # If the layer is provided instead of layer path, then return it as is
+    if isinstance(spec_or_layer.layer, (type, types.FunctionType)):
+        return spec_or_layer.layer
+
+    # Otherwise, return the dynamically imported layer from the layer path
+    return import_spec_layer(spec_or_layer.layer)
+
+
+def build_spec_layer(spec_or_layer: LayerSpec | type, *args, **kwargs):
+    # If the passed `spec_or_layer` is
+    # a `Function`, then return it as it is
+    # NOTE: to support an already initialized layer add the following condition
+    # `or isinstance(spec_or_layer, paddle.nn.Layer)` to the following if check
+    if isinstance(spec_or_layer, types.FunctionType):
+        return spec_or_layer
+
+    # If the passed `spec_or_layer` is actually a spec (instance of
+    # `LayerSpec`) and it specifies a `Function` using its `layer`
+    # field, return the `Function` as it is
+    if isinstance(spec_or_layer, LayerSpec) and isinstance(
+        spec_or_layer.layer, types.FunctionType
+    ):
+        return spec_or_layer.layer
+
+    # Check if a layer class is provided as a spec or if the layer path
+    # itself is a class
+    if isinstance(spec_or_layer, type):
+        layer = spec_or_layer
+    elif hasattr(spec_or_layer, "layer") and isinstance(
+        spec_or_layer.layer, type
+    ):
+        layer = spec_or_layer.layer
+    else:
+        # Otherwise, dynamically import the layer from the layer path
+        layer = import_spec_layer(spec_or_layer.layer)
+
+    # If the imported layer is actually a `Function` return it as it is
+    if isinstance(layer, types.FunctionType):
+        return layer
+
+    # Finally return the initialized layer with extra_kwargs from the spec as well
+    # as those passed as **kwargs from the code
+
+    # Add the `sublayers_spec` argument to the layer init call if it exists in the
+    # spec.
+    if (
+        hasattr(spec_or_layer, "sublayers_spec")
+        and spec_or_layer.sublayers_spec is not None
+    ):
+        kwargs["sublayers_spec"] = spec_or_layer.sublayers_spec
+    if hasattr(spec_or_layer, "extra_kwargs"):
+        for key in spec_or_layer.extra_kwargs.keys():
+            if key in kwargs:
+                warnings.warn(
+                    f"Got same key {key} in extra_kwargs and kwargs during init {layer.__name__}. Will keep the value ing extra_kwargs."
+                )
+                kwargs.pop(key)
+    try:
+        return layer(
+            *args,
+            **spec_or_layer.extra_kwargs
+            if hasattr(spec_or_layer, "extra_kwargs")
+            else {},
+            **kwargs,
+        )
+    except Exception as e:
+        # improve the error message since we hide the layer name in the line above
+        import sys
+
+        raise type(e)(
+            f"{e!s} when instantiating {layer.__name__}"
+        ).with_traceback(sys.exc_info()[2])

python/paddle/distributed/fleet/meta_parallel/pipeline_hooks.py

@@ -0,0 +1,56 @@
+# 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.
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class PipelineHook:
+    def __init__(self):
+        self.hooks: dict[int, list[Callable]] = defaultdict(list)
+        self._hooks_capacity = 0
+        self.reset_current_id()
+
+    def reset_current_id(self):
+        self._current_id = 0
+
+    def set_hooks_capacity(self, capacity: int):
+        self._hooks_capacity = capacity
+
+    def register_hook(self, hook_id: int, hook: Callable):
+        assert hook_id < self._hooks_capacity, (
+            f"hook_id {hook_id} is out of range, maximum capacity is {self._hooks_capacity}."
+        )
+        self.hooks[hook_id].append(hook)
+
+    def run_hook(self):
+        assert self._current_id < self._hooks_capacity, (
+            f"hook_id {self._current_id} is out of range, maximum capacity is {self._hooks_capacity}."
+        )
+        for hook in self.hooks[self._current_id]:
+            hook(self._current_id)
+        self._current_id += 1
+
+    @property
+    def current_id(self):
+        return self._current_id
+
+    @property
+    def hooks_capacity(self):
+        return self._hooks_capacity

python/paddle/distributed/fleet/meta_parallel/pp_utils/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+__all__ = []

python/paddle/distributed/fleet/meta_parallel/pp_utils/batch_comm_helper.py

@@ -0,0 +1,133 @@
+# 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 paddle
+from paddle.distributed.communication.batch_isend_irecv import (
+    P2POp,
+)
+
+from .p2p_communication import SendRecvMeta
+from .utils import (
+    number_2_dtype,
+    paddle_2_number,
+)
+
+
+class BatchCommHelper:
+    # NOTE(zhangyuqin1998): Tensors to be sent or received must have a
+    # consistent shape and data type throughout the entire pipeline.
+    def __init__(self, use_cache=True):
+        self._send_recv_meta = SendRecvMeta()
+        self._use_cache = use_cache
+
+    def clear_meta_cache(self):
+        self._send_recv_meta.init_or_erase_meta()
+
+    def _send_meta(self, tensors, group, broadcast=False):
+        self._send_recv_meta.set_send_message(tensors)
+        self._send_recv_meta.send_meta(tensors, group, broadcast=broadcast)
+        self._send_recv_meta.recv_shape_message = (
+            self._send_recv_meta.send_shape_message
+        )
+        self._send_recv_meta.recv_dtype_message = (
+            self._send_recv_meta.send_dtype_message
+        )
+
+    def _recv_meta(self, group, broadcast=False):
+        self._send_recv_meta.recv_meta(group, broadcast=broadcast)
+
+    def _build_from_meta(self):
+        shape_message = self._send_recv_meta.recv_shape_message
+        dtype_message = self._send_recv_meta.recv_dtype_message
+        stop_gradient = self._send_recv_meta.recv_stop_gradient
+        assert (shape_message is not None) and (dtype_message is not None), (
+            "Failed to build from meta."
+        )
+
+        res = []
+        if isinstance(shape_message, tuple):
+            for idx, shape in enumerate(shape_message):
+                tmp = paddle.empty(
+                    shape=shape, dtype=number_2_dtype(dtype_message[idx])
+                )
+                tmp.stop_gradient = (
+                    stop_gradient[idx] if stop_gradient is not None else False
+                )
+                res.append(tmp)
+        else:
+            tmp = paddle.empty(
+                shape=shape_message, dtype=number_2_dtype(dtype_message)
+            )
+            tmp.stop_gradient = stop_gradient
+            res.append(tmp)
+        return res
+
+    def _check_valid(self, tensors):
+        shape_message = self._send_recv_meta.recv_shape_message
+        dtype_message = self._send_recv_meta.recv_dtype_message
+
+        assert (shape_message is not None) and (dtype_message is not None), (
+            "Failed to build from meta."
+        )
+
+        if isinstance(shape_message, tuple):
+            assert isinstance(tensors, (list, tuple))
+            assert len(tensors) == len(shape_message)
+            for idx, (shape, dtype, tensor) in enumerate(
+                zip(shape_message, dtype_message, tensors)
+            ):
+                assert tensor.shape == shape, "Invalid shape."
+                assert number_2_dtype(
+                    paddle_2_number(tensor.dtype)
+                ) == number_2_dtype(dtype), "Invalid dtype."
+        else:
+            if isinstance(tensors, (list, tuple)):
+                assert len(tensors) == 1
+                tensors = tensors[0]
+
+            assert tensors.shape == shape_message, "Invalid shape."
+            assert number_2_dtype(
+                paddle_2_number(tensors.dtype)
+            ) == number_2_dtype(dtype_message), "Invalid dtype."
+
+    def recv_meta_from_head(self, group, need_recv_meta):
+        if not need_recv_meta:
+            return
+        self._recv_meta(group, broadcast=True)
+
+    def append_irecv(self, ops, src, group, alloc_on_comm_stream=False):
+        if alloc_on_comm_stream:
+            send_recv_stream = paddle.device.Stream(
+                stream_base=group.process_group.get_stream(
+                    paddle.framework._current_expected_place_()
+                )
+            )
+            with paddle.device.stream_guard(send_recv_stream):
+                tensors = self._build_from_meta()
+        else:
+            tensors = self._build_from_meta()
+        for tensor in tensors:
+            if tensor is not None:
+                ops.append(P2POp(paddle.distributed.irecv, tensor, src, group))
+        return tensors
+
+    def append_isend(self, ops, tensors, dst, group, need_broadcast_meta=False):
+        if need_broadcast_meta:
+            self._send_meta(tensors, group, broadcast=True)
+        self._check_valid(tensors)
+        for tensor in tensors:
+            if tensor is not None:
+                ops.append(P2POp(paddle.distributed.isend, tensor, dst, group))

python/paddle/distributed/fleet/meta_parallel/pp_utils/forward_backward_overlap_utils.py

@@ -0,0 +1,201 @@
+# 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 .utils import dict_to_tuple_helper
+
+
+class ScheduleChunk:
+    # NOTE(zhangyuqin): ScheduleChunk is the atomic unit of pipeline scheduling.
+    # A ScheduleChunk can contain several ScheduleNodes
+    def __init__(self, nodes):
+        self.nodes = nodes
+        self._check_nodes_valid()
+
+    def forward(self, inputs):
+        for n in self.nodes:
+            inputs = n.forward(inputs)
+        return inputs
+
+    def backward(self, output_grad):
+        for n in reversed(self.nodes):
+            output_grad = n.backward(output_grad)
+        return output_grad
+
+    def _check_nodes_valid(self):
+        for n in self.nodes:
+            assert isinstance(n, (ScheduleNode, ScheduleChunk))
+
+
+def detach_and_requires_grad(inputs):
+    if isinstance(inputs, (tuple, list)):
+        is_tuple = isinstance(inputs, tuple)
+        ret = []
+        for input in inputs:
+            if isinstance(input, (tuple, list)):
+                ret.append(detach_and_requires_grad(input))
+            elif isinstance(input, paddle.Tensor):
+                tmp = input.detach() if input is not None else None
+                if tmp is not None:
+                    tmp.stop_gradient = input.stop_gradient
+                ret.append(tmp)
+            else:
+                ret.append(input)
+        if is_tuple:
+            ret = tuple(ret)
+        return ret
+    elif isinstance(inputs, dict):
+        ret = {}
+        for key in inputs.keys():
+            input = inputs[key]
+            tmp = input.detach() if input is not None else None
+            if tmp is not None:
+                tmp.stop_gradient = input.stop_gradient
+            ret[key] = tmp
+        return ret
+    else:
+        tmp = inputs.detach()
+        tmp.stop_gradient = inputs.stop_gradient
+        return tmp
+
+
+def clone_and_clear_dataptr(outputs, clear_dataptr=False):
+    if isinstance(outputs, (tuple, list)):
+        is_tuple = isinstance(outputs, tuple)
+        ret = [
+            FakeClone.apply(o)
+            for o in outputs
+            if o is not None and isinstance(o, paddle.Tensor)
+        ]
+
+        if clear_dataptr:
+            for o in ret:
+                o._clear_dataptr()
+        if is_tuple:
+            ret = tuple(ret)
+        return ret
+    elif isinstance(outputs, dict):
+        ret = {}
+        for key in outputs.keys():
+            o = outputs[key]
+            if o is not None and isinstance(o, paddle.Tensor):
+                ret[key] = FakeClone.apply(o)
+        if clear_dataptr:
+            for key in ret:
+                ret[key]._clear_dataptr()
+        return ret
+    else:
+        ret = FakeClone.apply(outputs)
+        if clear_dataptr:
+            ret._clear_dataptr()
+        return ret
+
+
+class FakeClone(paddle.autograd.PyLayer):
+    # NOTE(zhangyuqin): Some input tensors may not be used in the forward function, but their gradients
+    # need to be retained. Therefore, we need a clone here. To avoid the DtoD copy, we need a FakeClone
+    @staticmethod
+    def forward(ctx, input):
+        return paddle.empty_like(input)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+class ScheduleNode:
+    # NOTE(zhangyuqin): ScheduleNode is a subgraph of the pipeline, capable of independently calling
+    # forward and backward. Users should not use paddle.autograd.backward on the results of ScheduleNode.forward.
+    # Instead, they should use ScheduleNode.backward. Otherwise, resource leakage may occur.
+    def __init__(self, fwd_func, name=""):
+        self.name = name
+        self.fwd_func = fwd_func
+        self.inputs = None
+        self.outputs = None
+
+        self.labels = None
+        self.scale_loss_factor = None
+
+    def forward(self, inputs=(), **kwargs):
+        detached_inputs = detach_and_requires_grad(inputs)
+        self.inputs = detached_inputs
+        if self.labels is not None:
+            outputs = self.fwd_func(self.inputs, self.labels, **kwargs)
+        else:
+            outputs = self.fwd_func(self.inputs, **kwargs)
+        if self.scale_loss_factor is not None:
+            outputs /= self.scale_loss_factor
+
+        # Do not release the loss tensor.
+        clear_dataptr = self.labels is None
+        self.outputs = clone_and_clear_dataptr(outputs, clear_dataptr)
+        return outputs
+
+    def backward(self, output_grad=None, scaler=None):
+        if output_grad is None:
+            if isinstance(self.outputs, (tuple, list)):
+                assert len(self.outputs) == 1
+                outputs = self.outputs[0]
+            else:
+                outputs = self.outputs
+            assert isinstance(outputs, paddle.Tensor)
+            if scaler is not None:
+                paddle.autograd.backward(scaler.scale(outputs))
+            else:
+                paddle.autograd.backward(outputs)
+        else:
+            # Record the original type (tuple or list) to preserve it after filtering
+            is_output_grad_tuple = isinstance(output_grad, tuple)
+            if not isinstance(output_grad, (tuple, list)):
+                is_output_grad_tuple = True  # Single value becomes tuple
+                output_grad = (output_grad,)
+
+            outputs = dict_to_tuple_helper(self.outputs)
+            if not isinstance(outputs, (tuple, list)):
+                outputs = (outputs,)
+            outputs = [t for t in outputs if not t.stop_gradient]
+
+            # Filter None values from output_grad
+            output_grad = [grad for grad in output_grad if grad is not None]
+            # Preserve original type (tuple or list)
+            output_grad = (
+                tuple(output_grad)
+                if is_output_grad_tuple
+                else list(output_grad)
+            )
+
+            assert len(outputs) == len(output_grad), (
+                f"{len(outputs)} of {type(outputs[0])} vs {len(output_grad)} of {type(output_grad[0])}"
+            )
+
+            paddle.autograd.backward(outputs, output_grad)
+
+        inputs = dict_to_tuple_helper(self.inputs)
+        if not isinstance(inputs, (tuple, list)):
+            inputs = (inputs,)
+        grad = tuple([e.grad if e is not None else None for e in inputs])
+        # grad = tuple([e.grad if e is not None and not e.stop_gradient else None for e in inputs])
+        self._reset_states()
+
+        # if len(grad) == 1:
+        #     grad = grad[0]
+        return grad
+
+    def _reset_states(self):
+        self.inputs = None
+        self.outputs = None
+        self.labels = None
+        self.scale_loss_factor = None

python/paddle/distributed/fleet/meta_parallel/pp_utils/four_directions_p2p_communication.py

@@ -0,0 +1,870 @@
+# 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
+
+import numpy as np
+
+import paddle
+from paddle import framework
+
+from ...utils import timer_helper as timer
+from ...utils.log_util import logger
+from .utils import number_2_dtype, paddle_2_number
+
+_hcg = None
+_enable_partial_send_recv = True
+_timers = None
+
+_xpu_comm_group_started = False
+
+_sync_send = os.environ.get("PADDLE_P2P_SYNC_SEND", "0")
+_sync_send = _sync_send.lower() in ['1', 'true']
+
+
+def _xpu_comm_group_start():
+    if not paddle.is_compiled_with_xpu():
+        return
+    global _xpu_comm_group_started
+    assert not _xpu_comm_group_started
+    framework.core.ProcessGroupBKCL.group_start()
+    _xpu_comm_group_started = True
+
+
+def _xpu_comm_group_end():
+    if not paddle.is_compiled_with_xpu():
+        return
+    global _xpu_comm_group_started
+    if _xpu_comm_group_started:
+        framework.core.ProcessGroupBKCL.group_end()
+        _xpu_comm_group_started = False
+
+
+def initialize_p2p_groups(
+    hcg, enable_partial_send_recv=True, enable_timer=False
+):
+    global _hcg, _enable_partial_send_recv, _timers
+    _hcg = hcg
+    _enable_partial_send_recv = enable_partial_send_recv
+    if enable_timer:
+        _timers = timer.get_timers()
+    (
+        send_next_group,
+        send_prev_group,
+        recv_next_group,
+        recv_prev_group,
+    ) = _hcg.get_p2p_groups()
+
+    debug_str = (
+        f"P2pInfo: send_next_group: {send_next_group!r}, send_prev_group: {send_prev_group!r}, "
+        f"recv_next_group: {recv_next_group!r}, recv_prev_group: {recv_prev_group!r}"
+    )
+    logger.info(debug_str)
+
+
+class SendRecvMeta:
+    """Mainly used to help p2p communication context information"""
+
+    def __init__(self):
+        self.send_shape_message = None
+        self.send_dtype_message = None
+
+        self.recv_shape_message = None
+        self.recv_dtype_message = None
+        self.recv_stop_gradient = None
+
+        self.has_send_meta = False
+        self.has_recv_meta = False
+
+    def _recv_shape_dtype(self, group):
+        # recv len(shape)
+        dims = paddle.to_tensor([0])
+        src_rank = _hcg._get_p2p_prev_rank()
+
+        paddle.distributed.recv(dims, src=src_rank, group=group)
+        dims = dims.item()
+
+        # recv shape
+        shape = paddle.to_tensor([0] * dims)
+        paddle.distributed.recv(shape, src=src_rank, group=group)
+
+        # recv dtype
+        dtype = paddle.to_tensor([0])
+        paddle.distributed.recv(dtype, src=src_rank, group=group)
+
+        # recv stop_gradient
+        stop_grad = paddle.to_tensor([0])
+        paddle.distributed.recv(stop_grad, src=src_rank, group=group)
+        return shape.tolist(), dtype.item(), stop_grad.item()
+
+    def recv_meta(self, group):
+        tensor_type = paddle.to_tensor([0])
+        src_rank = _hcg._get_p2p_prev_rank()
+
+        paddle.distributed.recv(tensor_type, src=src_rank, group=group)
+        tensor_type = tensor_type.item()
+
+        if tensor_type == 0:
+            shape, dtype, stop_grad = self._recv_shape_dtype(group)
+            self.recv_shape_message = shape
+            self.recv_dtype_message = dtype
+            self.recv_stop_gradient = bool(stop_grad)
+
+        elif tensor_type == 1:
+            num = paddle.to_tensor([0])
+            paddle.distributed.recv(num, src=src_rank, group=group)
+            num = num.item()
+            shapes = []
+            dtypes = []
+            stop_grads = []
+            for i in range(num):
+                shape, dtype, stop_grad = self._recv_shape_dtype(group)
+                shapes.append(shape)
+                dtypes.append(dtype)
+                stop_grads.append(bool(stop_grad))
+
+            self.recv_shape_message = tuple(shapes)
+            self.recv_dtype_message = tuple(dtypes)
+            self.recv_stop_gradient = tuple(stop_grads)
+
+    def _send_dims_shape_dtype(self, tensor, group):
+        # send len(shape)
+        dims = paddle.to_tensor([len(tensor.shape)])
+        dst_rank = _hcg._get_p2p_next_rank()
+
+        paddle.distributed.send(dims, dst=dst_rank, group=group)
+
+        # send shape
+        shape = paddle.to_tensor(tensor.shape)
+        paddle.distributed.send(shape, dst=dst_rank, group=group)
+
+        # send dtype
+        dtype = paddle.to_tensor([paddle_2_number(tensor.dtype)])
+        paddle.distributed.send(dtype, dst=dst_rank, group=group)
+
+        # send trainable
+        stop_grad = paddle.to_tensor([int(tensor.stop_gradient)])
+        paddle.distributed.send(stop_grad, dst=dst_rank, group=group)
+
+    def send_meta(self, tensor, group):
+        dst_rank = _hcg._get_p2p_next_rank()
+
+        if isinstance(tensor, paddle.Tensor):
+            tensor_type = paddle.to_tensor([0])
+            # send tensor type
+            paddle.distributed.send(tensor_type, dst=dst_rank, group=group)
+
+            self._send_dims_shape_dtype(tensor, group)
+        elif isinstance(tensor, tuple):
+            tensor_type = paddle.to_tensor([1])
+            # send tensor type
+            paddle.distributed.send(tensor_type, dst=dst_rank, group=group)
+
+            nums = paddle.to_tensor([len(tensor)])
+            paddle.distributed.send(nums, dst=dst_rank, group=group)
+
+            for d in tensor:
+                assert isinstance(d, paddle.Tensor)
+                self._send_dims_shape_dtype(d, group=group)
+
+    def set_send_message(self, tensor):
+        if isinstance(tensor, paddle.Tensor):
+            self.send_shape_message = tensor.shape
+            self.send_dtype_message = paddle_2_number(tensor.dtype)
+        elif isinstance(tensor, tuple):
+            self.send_shape_message = tuple(
+                [d.shape for d in tensor if not d.stop_gradient]
+            )
+            self.send_dtype_message = tuple(
+                [
+                    paddle_2_number(d.dtype)
+                    for d in tensor
+                    if not d.stop_gradient
+                ]
+            )
+
+
+def _is_valid_send_recv_partial(tensor, mp_degree):
+    if not _enable_partial_send_recv:
+        return False
+    tensor_numel = np.prod(tensor.shape)
+    assert tensor_numel != 0, "can't send/recv zero element"
+    return mp_degree > 1 and tensor_numel % mp_degree == 0
+
+
+def _partial_send_op(
+    tensor, group, use_calc_stream, ring_id, dst, nranks, rank_id
+):
+    dst_rank_in_group = dst if group is None else group.get_group_rank(dst)
+    if framework.in_dynamic_mode():
+        group = (
+            paddle.distributed.collective._get_default_group()
+            if group is None
+            else group
+        )
+        comm_op = (
+            group.process_group.send_partial_on_calc_stream
+            if use_calc_stream
+            else group.process_group.send_partial
+        )
+        return comm_op(tensor, dst_rank_in_group, nranks, rank_id)
+
+
+def send_partial(
+    tensor, dst=0, nranks=1, rank_id=0, group=None, use_calc_stream=True
+):
+    # dst: local rank in group
+    if group is not None and not group.is_member():
+        return
+    ring_id = 0 if group is None else group.id
+
+    dst_rank = (
+        _hcg._get_p2p_next_rank() if dst == 1 else _hcg._get_p2p_prev_rank()
+    )
+
+    if _is_valid_send_recv_partial(tensor, nranks):
+        return _partial_send_op(
+            tensor, group, use_calc_stream, ring_id, dst_rank, nranks, rank_id
+        )
+    else:
+        send_op = paddle.distributed.isend
+        return send_op(tensor.detach(), dst=dst_rank, group=group)
+
+
+def _partial_recv_op(
+    tensor, group, use_calc_stream, ring_id, src, nranks, rank_id
+):
+    src_rank_in_group = src if group is None else group.get_group_rank(src)
+    group = (
+        paddle.distributed.collective._get_default_group()
+        if group is None
+        else group
+    )
+    comm_op = (
+        group.process_group.recv_partial_on_calc_stream
+        if use_calc_stream
+        else group.process_group.recv_partial
+    )
+    return comm_op(tensor, src_rank_in_group, nranks, rank_id)
+
+
+def recv_partial(
+    tensor, src=0, nranks=1, rank_id=0, group=None, use_calc_stream=True
+):
+    # src: local rank in group
+    if group is not None and not group.is_member():
+        return
+    ring_id = 0 if group is None else group.id
+
+    src_rank = (
+        _hcg._get_p2p_prev_rank() if src == 0 else _hcg._get_p2p_next_rank()
+    )
+
+    if _is_valid_send_recv_partial(tensor, nranks):
+        return _partial_recv_op(
+            tensor, group, use_calc_stream, ring_id, src_rank, nranks, rank_id
+        )
+    else:
+        if use_calc_stream:
+            recv_op = paddle.distributed.recv
+        elif framework.in_dynamic_mode():
+            recv_op = paddle.distributed.irecv
+        return recv_op(tensor.detach(), src=src_rank, group=group)
+
+
+def _partial_allgather_op(
+    tensor, group, use_calc_stream, ring_id, nranks, rank_id
+):
+    group = (
+        paddle.distributed.collective._get_default_group()
+        if group is None
+        else group
+    )
+    comm_op = (
+        group.process_group.all_gather_partial_on_calc_stream
+        if use_calc_stream
+        else group.process_group.all_gather_partial
+    )
+    return comm_op(tensor, tensor, nranks, rank_id)
+
+
+def allgather_partial(
+    tensor, nranks=1, rank_id=0, group=None, use_calc_stream=True
+):
+    if not _is_valid_send_recv_partial(tensor, nranks):
+        return tensor
+    if group is not None and not group.is_member():
+        return
+    ring_id = 0 if group is None else group.id
+
+    return _partial_allgather_op(
+        tensor, group, use_calc_stream, ring_id, nranks, rank_id
+    )
+
+
+def _p2p_helper(
+    tensor_send_next,
+    tensor_send_prev,
+    recv_prev,
+    recv_next,
+    sync_recv=True,
+    send_recv_meta=None,
+):
+    global _hcg
+
+    tensor_recv_prev = None
+    tensor_recv_next = None
+
+    # send / recv message
+    assert send_recv_meta is not None, "send_recv_meta should not be None"
+    recv_shape_msg = send_recv_meta.recv_shape_message
+    recv_dtype_msg = send_recv_meta.recv_dtype_message
+    recv_stop_gradient = send_recv_meta.recv_stop_gradient
+
+    send_shape_msg = send_recv_meta.send_shape_message
+    send_dtype_msg = send_recv_meta.send_dtype_message
+
+    # model parallel message
+    mp_group = _hcg.get_model_parallel_group()
+    mp_degree = _hcg.get_model_parallel_world_size()
+    mp_rank = _hcg.get_model_parallel_rank()
+
+    if recv_prev:
+        if isinstance(recv_shape_msg, tuple):
+            tensor_recv_prev = []
+            for idx, shape in enumerate(recv_shape_msg):
+                tmp = paddle.empty(
+                    shape=shape, dtype=number_2_dtype(recv_dtype_msg[idx])
+                )
+                tmp.stop_gradient = recv_stop_gradient[idx]
+                tensor_recv_prev.append(tmp)
+            tensor_recv_prev = tuple(tensor_recv_prev)
+        else:
+            tensor_recv_prev = paddle.empty(
+                shape=recv_shape_msg, dtype=number_2_dtype(recv_dtype_msg)
+            )
+            tensor_recv_prev.stop_gradient = recv_stop_gradient
+
+    if recv_next:
+        if isinstance(send_shape_msg, tuple):
+            tensor_recv_next = []
+            for idx, shape in enumerate(send_shape_msg):
+                tensor_recv_next.append(
+                    paddle.empty(
+                        shape=shape, dtype=number_2_dtype(send_dtype_msg[idx])
+                    )
+                )
+            tensor_recv_next = tuple(tensor_recv_next)
+        else:
+            tensor_recv_next = paddle.empty(
+                shape=send_shape_msg, dtype=number_2_dtype(send_dtype_msg)
+            )
+
+    # TODO(Yuang Liu): use batch_isend_irecv replace all these comm ops
+    tasks = []
+    # start to p2p communicate
+
+    if _sync_send:
+        # Some devices(NPU for example) do not support asynchronized send op, So the order is
+        # recv_prev -> send_next -> recv_next -> send_prev
+        # When using this order, the environment variable
+        # 'PADDLE_P2P_SYNC_SEND' should be set True
+        if tensor_recv_prev is not None:
+            if isinstance(tensor_recv_prev, tuple):
+                for d in tensor_recv_prev:
+                    task = recv_partial(
+                        d,
+                        src=0,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.recv_prev_group,
+                        use_calc_stream=sync_recv,
+                    )
+                    if sync_recv:
+                        allgather_partial(
+                            d,
+                            nranks=mp_degree,
+                            rank_id=mp_rank,
+                            group=mp_group,
+                            use_calc_stream=True,
+                        )
+                    else:
+                        tasks.append(task)
+            else:
+                task = recv_partial(
+                    tensor_recv_prev,
+                    src=0,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.recv_prev_group,
+                    use_calc_stream=sync_recv,
+                )
+
+                if sync_recv:
+                    allgather_partial(
+                        tensor_recv_prev,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=mp_group,
+                        use_calc_stream=True,
+                    )
+                else:
+                    tasks.append(task)
+
+        if tensor_send_next is not None:
+            if isinstance(tensor_send_next, tuple):
+                for d in tensor_send_next:
+                    paddle.distributed.wait(d, use_calc_stream=True)
+                    send_partial(
+                        d,
+                        dst=1,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.send_next_group,
+                        use_calc_stream=False,
+                    )
+            else:
+                paddle.distributed.wait(tensor_send_next, use_calc_stream=True)
+                send_partial(
+                    tensor_send_next,
+                    dst=1,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.send_next_group,
+                    use_calc_stream=False,
+                )
+
+        if tensor_recv_next is not None:
+            if isinstance(tensor_recv_next, tuple):
+                for d in tensor_recv_next:
+                    task = recv_partial(
+                        d,
+                        src=1,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.recv_next_group,
+                        use_calc_stream=sync_recv,
+                    )
+
+                    if sync_recv:
+                        allgather_partial(
+                            d,
+                            nranks=mp_degree,
+                            rank_id=mp_rank,
+                            group=mp_group,
+                            use_calc_stream=True,
+                        )
+                    else:
+                        tasks.append(task)
+
+            else:
+                task = recv_partial(
+                    tensor_recv_next,
+                    src=1,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.recv_next_group,
+                    use_calc_stream=sync_recv,
+                )
+                if sync_recv:
+                    allgather_partial(
+                        tensor_recv_next,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=mp_group,
+                        use_calc_stream=True,
+                    )
+                else:
+                    tasks.append(task)
+
+        if tensor_send_prev is not None:
+            if isinstance(tensor_send_prev, tuple):
+                for d in tensor_send_prev:
+                    paddle.distributed.wait(d, use_calc_stream=True)
+                    send_partial(
+                        d,
+                        dst=0,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.send_prev_group,
+                        use_calc_stream=False,
+                    )
+            else:
+                paddle.distributed.wait(tensor_send_prev, use_calc_stream=True)
+                send_partial(
+                    tensor_send_prev,
+                    dst=0,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.send_prev_group,
+                    use_calc_stream=False,
+                )
+    else:
+        _xpu_comm_group_start()
+        if tensor_send_prev is not None:
+            if isinstance(tensor_send_prev, tuple):
+                for d in tensor_send_prev:
+                    paddle.distributed.wait(d, use_calc_stream=True)
+                    send_partial(
+                        d,
+                        dst=0,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.send_prev_group,
+                        use_calc_stream=False,
+                    )
+            else:
+                paddle.distributed.wait(tensor_send_prev, use_calc_stream=True)
+                send_partial(
+                    tensor_send_prev,
+                    dst=0,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.send_prev_group,
+                    use_calc_stream=False,
+                )
+
+        if tensor_recv_prev is not None:
+            if isinstance(tensor_recv_prev, tuple):
+                for d in tensor_recv_prev:
+                    task = recv_partial(
+                        d,
+                        src=0,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.recv_prev_group,
+                        use_calc_stream=sync_recv,
+                    )
+                    if sync_recv:
+                        _xpu_comm_group_end()
+                        allgather_partial(
+                            d,
+                            nranks=mp_degree,
+                            rank_id=mp_rank,
+                            group=mp_group,
+                            use_calc_stream=True,
+                        )
+                    else:
+                        tasks.append(task)
+            else:
+                task = recv_partial(
+                    tensor_recv_prev,
+                    src=0,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.recv_prev_group,
+                    use_calc_stream=sync_recv,
+                )
+
+                if sync_recv:
+                    _xpu_comm_group_end()
+                    allgather_partial(
+                        tensor_recv_prev,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=mp_group,
+                        use_calc_stream=True,
+                    )
+                else:
+                    tasks.append(task)
+
+        if tensor_send_next is not None:
+            if isinstance(tensor_send_next, tuple):
+                for d in tensor_send_next:
+                    paddle.distributed.wait(d, use_calc_stream=True)
+                    send_partial(
+                        d,
+                        dst=1,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.send_next_group,
+                        use_calc_stream=False,
+                    )
+            else:
+                paddle.distributed.wait(tensor_send_next, use_calc_stream=True)
+                send_partial(
+                    tensor_send_next,
+                    dst=1,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.send_next_group,
+                    use_calc_stream=False,
+                )
+
+        if tensor_recv_next is not None:
+            if isinstance(tensor_recv_next, tuple):
+                for d in tensor_recv_next:
+                    task = recv_partial(
+                        d,
+                        src=1,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=_hcg.recv_next_group,
+                        use_calc_stream=sync_recv,
+                    )
+
+                    if sync_recv:
+                        _xpu_comm_group_end()
+                        allgather_partial(
+                            d,
+                            nranks=mp_degree,
+                            rank_id=mp_rank,
+                            group=mp_group,
+                            use_calc_stream=True,
+                        )
+                    else:
+                        tasks.append(task)
+
+            else:
+                task = recv_partial(
+                    tensor_recv_next,
+                    src=1,
+                    nranks=mp_degree,
+                    rank_id=mp_rank,
+                    group=_hcg.recv_next_group,
+                    use_calc_stream=sync_recv,
+                )
+                if sync_recv:
+                    _xpu_comm_group_end()
+                    allgather_partial(
+                        tensor_recv_next,
+                        nranks=mp_degree,
+                        rank_id=mp_rank,
+                        group=mp_group,
+                        use_calc_stream=True,
+                    )
+                else:
+                    tasks.append(task)
+        _xpu_comm_group_end()
+    if not sync_recv:
+        if framework.in_dynamic_mode():
+            # wait irecv tasks in eager dygraph mode with new comm library
+            for task in tasks:
+                assert task is not None
+                task.wait()
+
+        tensors_for_all_gather = []
+        if tensor_recv_prev is not None:
+            if isinstance(tensor_recv_prev, tuple):
+                for d in tensor_recv_prev:
+                    tensors_for_all_gather.append(d)
+            else:
+                tensors_for_all_gather.append(tensor_recv_prev)
+        if tensor_recv_next is not None:
+            if isinstance(tensor_recv_next, tuple):
+                for d in tensor_recv_next:
+                    tensors_for_all_gather.append(d)
+            else:
+                tensors_for_all_gather.append(tensor_recv_next)
+
+        for tensor in tensors_for_all_gather:
+            allgather_partial(
+                tensor,
+                nranks=mp_degree,
+                rank_id=mp_rank,
+                group=mp_group,
+                use_calc_stream=True,
+            )
+
+    return tensor_recv_prev, tensor_recv_next
+
+
+class P2pHelper:
+    def __init__(self, use_cache=True):
+        self._send_recv_meta = SendRecvMeta()
+        self._use_cache = use_cache
+
+    def _send_meta(self, output_tensor, skip_check_meta=False):
+        if not self._send_recv_meta.has_send_meta:
+            self._send_recv_meta.set_send_message(output_tensor)
+            self._send_recv_meta.send_meta(
+                output_tensor, _hcg.get_pipe_parallel_group()
+            )
+            self._send_recv_meta.has_send_meta = self._use_cache
+
+    def _recv_meta(self):
+        if not self._send_recv_meta.has_recv_meta:
+            self._send_recv_meta.recv_meta(_hcg.get_pipe_parallel_group())
+            self._send_recv_meta.has_recv_meta = self._use_cache
+
+    def recv_forward(self, pp_first_stage, sync_recv=True):
+        global _timers
+        if _timers is not None:
+            _timers("recv_forward").start()
+        if pp_first_stage:
+            input_tensor = None
+        else:
+            self._recv_meta()
+
+            input_tensor, _ = _p2p_helper(
+                tensor_send_next=None,
+                tensor_send_prev=None,
+                recv_prev=True,
+                recv_next=False,
+                sync_recv=sync_recv,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("recv_forward").stop()
+        return input_tensor
+
+    def recv_backward(self, pp_last_stage, sync_recv=True):
+        global _timers
+        if _timers is not None:
+            _timers("recv_backward").start()
+        if pp_last_stage:
+            output_tensor_grad = None
+        else:
+            _, output_tensor_grad = _p2p_helper(
+                tensor_send_next=None,
+                tensor_send_prev=None,
+                recv_prev=False,
+                recv_next=True,
+                sync_recv=sync_recv,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("recv_backward").stop()
+        return output_tensor_grad
+
+    def send_forward(self, output_tensor, pp_last_stage, skip_check_meta=False):
+        global _timers
+        if _timers is not None:
+            _timers("send_forward").start()
+        if not pp_last_stage:
+            self._send_meta(output_tensor, skip_check_meta=skip_check_meta)
+
+            _p2p_helper(
+                tensor_send_next=output_tensor,
+                tensor_send_prev=None,
+                recv_prev=False,
+                recv_next=False,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("send_forward").stop()
+
+    def send_backward(self, input_tensor_grad, pp_first_stage):
+        global _timers
+        if _timers is not None:
+            _timers("send_backward").start()
+        if not pp_first_stage:
+            _p2p_helper(
+                tensor_send_next=None,
+                tensor_send_prev=input_tensor_grad,
+                recv_prev=False,
+                recv_next=False,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("send_backward").stop()
+
+    def send_forward_recv_backward(self, output_tensor, pp_last_stage):
+        global _timers
+        if _timers is not None:
+            _timers("send_forward_recv_backward").start()
+        if pp_last_stage:
+            output_tensor_grad = None
+        else:
+            _, output_tensor_grad = _p2p_helper(
+                tensor_send_next=output_tensor,
+                tensor_send_prev=None,
+                recv_prev=False,
+                recv_next=True,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("send_forward_recv_backward").stop()
+        return output_tensor_grad
+
+    def send_backward_recv_forward(self, input_tensor_grad, pp_first_stage):
+        global _timers
+        if _timers is not None:
+            _timers("send_backward_recv_forward").start()
+        if pp_first_stage:
+            input_tensor = None
+        else:
+            input_tensor, _ = _p2p_helper(
+                tensor_send_next=None,
+                tensor_send_prev=input_tensor_grad,
+                recv_prev=True,
+                recv_next=False,
+                send_recv_meta=self._send_recv_meta,
+            )
+        if _timers is not None:
+            _timers("send_backward_recv_forward").stop()
+        return input_tensor
+
+    def send_forward_backward_recv_forward_backward(
+        self, output_tensor, input_tensor_grad, recv_prev, recv_next
+    ):
+        # always have to send dtype info to downstream
+        global _timers
+        if _timers is not None:
+            _timers("send_forward_backward_recv_forward_backward").start()
+
+        self._send_meta(output_tensor)
+        if recv_prev:
+            self._recv_meta()
+        input_tensor, output_tensor_grad = _p2p_helper(
+            tensor_send_next=output_tensor,
+            tensor_send_prev=input_tensor_grad,
+            recv_prev=recv_prev,
+            recv_next=recv_next,
+            sync_recv=False,
+            send_recv_meta=self._send_recv_meta,
+        )
+        if _timers is not None:
+            _timers("send_forward_backward_recv_forward_backward").stop()
+        return input_tensor, output_tensor_grad
+
+    def send_forward_recv_forward(self, output_tensor, recv_prev):
+        # always have to send dtype info to downstream
+        global _timers
+        if _timers is not None:
+            _timers("send_forward_recv_forward").start()
+
+        if output_tensor is not None:
+            self._send_meta(output_tensor)
+
+        if recv_prev:
+            self._recv_meta()
+
+        input_tensor, _ = _p2p_helper(
+            tensor_send_next=output_tensor,
+            tensor_send_prev=None,
+            recv_prev=recv_prev,
+            recv_next=False,
+            sync_recv=False,
+            send_recv_meta=self._send_recv_meta,
+        )
+        if _timers is not None:
+            _timers("send_forward_recv_forward").stop()
+        return input_tensor
+
+    def send_backward_recv_backward(self, input_tensor_grad, recv_next):
+        global _timers
+        if _timers is not None:
+            _timers("send_backward_recv_backward").start()
+        _, output_tensor_grad = _p2p_helper(
+            tensor_send_next=None,
+            tensor_send_prev=input_tensor_grad,
+            recv_prev=False,
+            recv_next=recv_next,
+            sync_recv=False,
+            send_recv_meta=self._send_recv_meta,
+        )
+        if _timers is not None:
+            _timers("send_backward_recv_backward").stop()
+        return output_tensor_grad

python/paddle/distributed/fleet/meta_parallel/pp_utils/profiler_helper.py

@@ -0,0 +1,41 @@
+# 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
+
+
+def main():
+    all_record = []
+    all_files = os.listdir('./')
+    all_files = sorted(
+        filter(
+            lambda file: file.startswith("profile_record_tmp_file_for_rank_"),
+            all_files,
+        )
+    )
+
+    for files in all_files:
+        with open(files, 'r') as f:
+            for line in f:
+                all_record.append(line.strip())
+
+    with open('pipeline_profile.json', 'w') as f:
+        f.write('[ ')
+        f.writelines(all_record[i] + ',\n' for i in range(len(all_record) - 1))
+        f.write(all_record[-1])
+        f.write(' ]\n')
+
+
+if __name__ == "__main__":
+    main()

python/paddle/distributed/fleet/meta_parallel/pp_utils/utils.py

@@ -0,0 +1,160 @@
+# 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 paddle
+from paddle import _C_ops
+
+__all__ = []
+
+
+FLOAT_TYPE_DICT = {
+    paddle.float16: "float16",
+    paddle.float32: "float32",
+    paddle.float64: "float64",
+    paddle.bfloat16: "bfloat16",
+    paddle.bool: "bool",
+}
+
+PADDLE_TO_NUMBER = {
+    paddle.float16: 0,
+    paddle.float32: 1,
+    paddle.float64: 2,
+    paddle.int32: 3,
+    paddle.int64: 4,
+    paddle.bfloat16: 5,
+    paddle.bool: 6,
+}
+
+NUMBER_TO_DTYPE = {
+    0: "float16",
+    1: "float32",
+    2: "float64",
+    3: "int32",
+    4: "int64",
+    5: "bfloat16",
+    6: "bool",
+}
+
+
+def is_float_tensor(tensor):
+    """Is a float tensor"""
+    return tensor.dtype in FLOAT_TYPE_DICT.keys()
+
+
+def get_tensor_dtype(dtype):
+    assert dtype in FLOAT_TYPE_DICT.keys()
+    return FLOAT_TYPE_DICT[dtype]
+
+
+def paddle_2_number(dtype):
+    assert dtype in PADDLE_TO_NUMBER.keys()
+    return PADDLE_TO_NUMBER[dtype]
+
+
+def number_2_dtype(number):
+    assert number in NUMBER_TO_DTYPE.keys()
+    return NUMBER_TO_DTYPE[number]
+
+
+def get_tensor_bytes(tensor):
+    """Get the bytes a tensor occupied."""
+    elem_size = None
+    if tensor.dtype == paddle.float32:
+        elem_size = 4
+    elif tensor.dtype == paddle.float64:
+        elem_size = 8
+    elif tensor.dtype == paddle.int64:
+        elem_size = 8
+    elif tensor.dtype == paddle.int32:
+        elem_size = 4
+    elif tensor.dtype == paddle.float16:
+        elem_size = 2
+    elif tensor.dtype == paddle.int8:
+        elem_size = 1
+    else:
+        raise ValueError(f"unknown data type: {tensor.dtype}")
+    return tensor.numel() * elem_size
+
+
+def _all_gather(tensor, group=None, use_calc_stream=True):
+    """
+    The main difference with paddle.distributed.all_gather:
+    no need to pass in tensor_list, the returned tensor is spliced
+    """
+    if group is not None and not group.is_member():
+        return
+    ring_id = 0 if group is None else group.id
+    nranks = (
+        paddle.distributed.collective._get_global_group().nranks
+        if group is None
+        else group.nranks
+    )
+    return _C_ops.all_gather(
+        tensor,
+        ring_id,
+        nranks,
+    )
+
+
+def tuple_to_dict_helper(input_tensor):
+    # recv tuple -> fwd input dict
+    use_dict = False
+    if isinstance(input_tensor, tuple):
+        use_dict = hasattr(input_tensor[0], "key")
+    else:  # single tensor
+        use_dict = hasattr(input_tensor, "key")
+    if use_dict:
+        input_tensor = convert_tensor_tuple_to_dict(input_tensor)
+    return input_tensor, use_dict
+
+
+def dict_to_tuple_helper(output_tensor):
+    if isinstance(output_tensor, dict):
+        output_tensor_tuple = convert_tensor_dict_to_tuple(
+            output_tensor_dict=output_tensor
+        )
+    else:  # single tensor or tensor tuple
+        output_tensor_tuple = output_tensor
+    return output_tensor_tuple
+
+
+def convert_tensor_dict_to_tuple(output_tensor_dict):
+    output_tensor = []
+    for key, tensor in output_tensor_dict.items():
+        if isinstance(tensor, (list, tuple)):
+            for idx, t in enumerate(tensor):
+                t.key = key + " " + str(idx)
+                output_tensor.append(t)
+        else:  # single tensor
+            tensor.key = key
+            output_tensor.append(tensor)
+
+    return tuple(output_tensor)
+
+
+def convert_tensor_tuple_to_dict(input_tensor_tuple):
+    input_tensor_dict = {}
+    for tensor in input_tensor_tuple:
+        key = tensor.key
+        if " " in key:
+            real_key, _ = key.split(" ")
+            if real_key in input_tensor_dict.keys():
+                input_tensor_dict[real_key].append(tensor)
+            else:
+                input_tensor_dict[real_key] = [tensor]
+        else:
+            input_tensor_dict[key] = tensor
+        delattr(tensor, "key")
+    return input_tensor_dict

python/paddle/distributed/fleet/meta_parallel/segment_parallel.py

@@ -0,0 +1,40 @@
+#   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.hybrid_parallel_util import (
+    broadcast_dp_parameters,
+    broadcast_sep_parameters,
+    broadcast_sharding_parameters,
+)
+from ..utils.log_util import logger
+from .meta_parallel_base import MetaParallelBase
+
+__all__ = []
+
+
+class SegmentParallel(MetaParallelBase):
+    def __init__(self, layers, hcg, **kwargs):
+        super().__init__(layers, hcg, **kwargs)
+
+    def _prepare_for_model(self):
+        logger.info("start broadcast sep parameters")
+        broadcast_sep_parameters(self._layers, self._hcg)
+
+        if self._hcg.get_sharding_parallel_world_size() > 1:
+            logger.info("start broadcast sharding parameters")
+            broadcast_sharding_parameters(self._layers, self._hcg)
+
+        if self._hcg.get_data_parallel_world_size() > 1:
+            logger.info("start broadcast dp parameters")
+            broadcast_dp_parameters(self._layers, self._hcg)

python/paddle/distributed/fleet/meta_parallel/sharding/__init__.py

@@ -0,0 +1,13 @@
+#   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.

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_fully_shard.py

@@ -0,0 +1,730 @@
+# 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 logging
+from collections import OrderedDict
+from types import MethodType
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import collective, fleet
+from paddle.framework import core
+from paddle.nn import ClipGradByGlobalNorm
+
+from .group_sharded_stage3 import (
+    ForwardPostHooks,
+    ForwardPreHooks,
+    OrderedSet,
+    TaskFlow,
+    _current_layer_params,
+    _PartitionParam,
+    _TensorWrapper,
+    _UnsliceParam,
+    align,
+    alignment,
+)
+from .group_sharded_storage import GradStorage
+from .group_sharded_utils import GroupShardedClipGrad, Type, device_guard
+
+
+def _OptimizerWrapper(optimizer, offload, group, update_params_slice):
+    if not hasattr(optimizer, "_optim"):
+        optimizer._optim = optimizer
+        optimizer.offload = offload
+        optimizer._group = group
+        optimizer.update_scaler = None
+        optimizer.update_slice = update_params_slice
+    return optimizer
+
+
+class FullyShardOptimizer:
+    def __init__(
+        self,
+        optimizer,
+        group=None,
+        sync_buffers=False,
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
+        segment_size=2**20,
+        pretrain_sync_models=True,
+        offload=False,
+        sync_comm=False,
+        dp_group=None,
+        exclude_layer=None,
+    ):
+        self._default_device = device
+        self.__sync_buffers = sync_buffers
+        self._offload = offload
+        self._sync_comm = sync_comm
+
+        # stage3 support some layer set by users to be unslice
+        # _exclude_layer=[layer_name or id(layer)]
+        self._exclude_layer = [] if exclude_layer is None else exclude_layer
+        assert isinstance(self._exclude_layer, (list, tuple)), (
+            "the exclude_layers must be a list with layers' name or layers' id"
+        )
+
+        # segmentation size
+        assert segment_size >= 0, "segment_size must be GE than 0."
+        self._segment_size = segment_size
+
+        global param2dtype
+        param2dtype = {}
+
+        hcg = fleet.get_hybrid_communicate_group()
+        group = hcg.get_sharding_parallel_group()
+        # Communication group establishment
+        self._group = (
+            collective.new_group(collective._get_global_group().ranks)
+            if group is None
+            else group
+        )
+        self._dp_group = dp_group
+        self._world_size_scaling = 1.0 / self._group.nranks
+        assert self._group.nranks > 1, (
+            "Training must be distributed, ranks must be greater than 1."
+        )
+        self._rank = self._group.rank
+        self._global_root_rank = self._group.ranks[
+            0
+        ]  # picking ranks index 0 as the reference
+
+        # Parameter segmentation for global ranks
+        self._unslice_params = OrderedSet()  # param's numel <= segment_size
+        self._unslice_params2align = {}  # {param.name: param's align}
+        self._grad_storages = {}  # {param.dtype: GradStorage}
+
+        assert not isinstance(optimizer, list), (
+            "Multiple optimizers are not supported now."
+        )
+        self._optim = _OptimizerWrapper(
+            optimizer,
+            self._offload,
+            self._group,
+            self._update_params_slice,
+        )
+        self._ori_parameter_list = self._optim._parameter_list
+        self._ori_param_groups = self._optim._param_groups
+
+        for p in self._ori_parameter_list:
+            del p._need_shard
+            if p._numel() > self._segment_size:
+                pass
+            elif p.trainable:
+                self._unslice_params.add(_UnsliceParam(p))
+
+        # check main_grad
+        self._check_main_grad()
+
+        # Replace optimizer's _grad_clip
+        if isinstance(self._optim._grad_clip, ClipGradByGlobalNorm):
+            logging.warning(
+                "While using ClipGradByGlobalNorm in GroupShardedStage3, the grad clip of original optimizer will be changed."
+            )
+            if self.use_main_grad:
+                self._optim._inner_opt._grad_clip = GroupShardedClipGrad(
+                    self._optim._inner_opt._grad_clip,
+                    paddle.get_device(),
+                    self._group,
+                )
+            else:
+                self._optim._grad_clip = GroupShardedClipGrad(
+                    self._optim._grad_clip, paddle.get_device(), self._group
+                )
+            if self._optim._parameter_list and isinstance(
+                self._optim._parameter_list[0], dict
+            ):
+                for item in self._optim._param_groups:
+                    if "grad_clip" in item.keys():
+                        item["grad_clip"] = self._optim._grad_clip
+
+        # Add unslice params to master_weight in fp16
+        self._setup_master_weights_for_unslice()
+
+        # Redefine optimizer step and clear function
+        self._redefine_opt_step()
+        self._redefine_opt_clear()
+
+    def _check_main_grad(self):
+        self.use_main_grad = None
+        for param in self._ori_parameter_list:
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(param, "main_grad"), (
+                    "Params have different main grad attributes."
+                )
+
+    def _setup_master_weights_for_unslice(self):
+        for param in self._unslice_params:
+            # Update optimizer master weights
+            if (
+                param.dtype == Type.fp16.value or param.dtype == Type.bf16.value
+            ) and not self._offload:
+                master_tensor = paddle.cast(param, Type.fp32.value)
+                master_tensor.name = param.name
+                self._optim._master_weights[param.name] = master_tensor
+
+    def _clear_gradients(self):
+        current_layer_params = self._ori_parameter_list
+        # 1.Handle param's slice
+        trainable_params = list(
+            filter(
+                lambda p: p.trainable and p not in self._unslice_params,
+                current_layer_params,
+            )
+        )
+        for param in trainable_params:
+            if not hasattr(param, "fw_storage"):
+                continue
+            assert hasattr(param, "fw_storage"), (
+                f"Find {param.name} don't have fw_storage attribute."
+            )
+            if self.use_main_grad:
+                param.fw_storage.main_grad._clear()
+                param.fw_storage.main_grad = None
+            else:
+                param.fw_storage.clear_gradient(False)
+            param.bw_storage._clear()
+            param.bw_storage = None
+
+    # Update param memory slice
+    def _update_params_slice(self):
+        update_list = self._update_params()
+
+        if not isinstance(self._optim._param_groups[0], dict):
+            slice_params = [param.fw_storage for param in update_list]
+            self._optim._parameter_list = slice_params + list(
+                self._unslice_params
+            )
+            self._optim._param_groups = slice_params + list(
+                self._unslice_params
+            )
+        else:
+            for param_group in self._optim._param_groups:
+                p_group = []
+                for p in param_group['params']:
+                    if hasattr(p, "fw_storage"):
+                        p_group.append(p.fw_storage)
+                    else:
+                        p_group.append(p)
+
+                param_group['params'] = p_group
+
+    def _update_params(self):
+        """
+        Update parameters to optimizer memory slice.
+        """
+        update_list = []
+        current_layer_params = self._ori_parameter_list
+        trainable_params = list(
+            filter(
+                lambda p: p.trainable and p not in self._unslice_params,
+                current_layer_params,
+            )
+        )
+        # 1.Handle param's slice
+        for param in trainable_params:
+            assert hasattr(param, "fw_storage"), (
+                f"Find {param.name} don't have fw_storage attribute"
+            )
+
+            param.fw_storage = _TensorWrapper(param)
+            if self.use_main_grad:
+                param.fw_storage.main_grad = param.bw_storage
+            else:
+                assert param.fw_storage.grad is None
+                param.fw_storage._copy_gradient_from(param.bw_storage)
+            update_list.append(param)
+
+        return update_list
+
+    def _redefine_opt_step(self):
+        params_slice_func = self._update_params_slice
+        opt_step = self._optim.step
+
+        def _opt_step(self):
+            if not self.update_scaler:
+                params_slice_func()
+            opt_step()
+
+        self._optim.step = MethodType(_opt_step, self._optim)
+
+    def _redefine_opt_clear(self):
+        clear_func = self._clear_gradients
+
+        def _opt_clear(self):
+            clear_func()
+
+        self._optim.clear_grad = MethodType(_opt_clear, self._optim)
+
+
+class FullyShard(nn.Layer):
+    """
+    A wrapper for Sharding Stage3 Layer in Dygraph.
+
+    .. warning: GroupShardedStage3 encapsulates the layer strategy and integrates it into the nn.Layer.
+
+    .. ZeRO: https://arxiv.org/pdf/1910.02054.pdf.
+    """
+
+    def __init__(
+        self,
+        layer,
+        group=None,
+        sync_buffers=False,
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
+        segment_size=2**20,
+        pretrain_sync_models=True,
+        offload=False,
+        sync_comm=False,
+        dp_group=None,
+        exclude_layer=None,
+    ):
+        super().__init__()
+
+        # Default configs
+        assert (
+            core.is_compiled_with_cuda()
+            or core.is_compiled_with_xpu()
+            or (device in core.get_all_custom_device_type())
+        ), "Only support CUDA / XPU / CustomDevice."
+
+        self._layer = layer
+        self._default_device = device
+        self.__sync_buffers = sync_buffers
+        self._offload = offload
+        self._sync_comm = sync_comm
+
+        # stage3 support some layer set by users to be unslice
+        self._exclude_layer = [] if exclude_layer is None else exclude_layer
+        assert isinstance(self._exclude_layer, (list, tuple)), (
+            "the exclude_layers must be a list with layers' name or layers' id"
+        )
+
+        # segmentation size
+        assert segment_size >= 0, "segment_size must be GE than 0."
+        self._segment_size = segment_size
+
+        global DEV
+        DEV = (
+            "cpu"
+            if paddle.get_device() == "cpu"
+            else paddle.get_device().split(":")[0]
+        )
+        global DEV_ID
+        DEV_ID = (
+            0
+            if paddle.get_device() == "cpu"
+            else int(paddle.get_device().split(":")[1])
+        )
+        global param2dtype
+        param2dtype = {}
+
+        hcg = fleet.get_hybrid_communicate_group()
+        group = hcg.get_sharding_parallel_group()
+
+        self._group = (
+            collective.new_group(collective._get_global_group().ranks)
+            if group is None
+            else group
+        )
+        self._dp_group = dp_group
+        self._world_size_scaling = 1.0 / self._group.nranks
+        assert self._group.nranks > 1, (
+            "Training must be distributed, ranks must be greater than 1."
+        )
+        self._rank = self._group.rank
+        self._global_root_rank = self._group.ranks[
+            0
+        ]  # picking ranks index 0 as the reference
+
+        # Parameter segmentation for global ranks
+        # After flatten -> self._param2buffer_size, self._param2buffer, self._trainable_params
+        self._param2buffer_size = {}  # {param.name: size}
+        self._param2buffer = {}  # {param.name: [(start0, end0),(start1, end1), ...]}
+        self._trainable_params = {}  # {id(layer): [trainable_params]}
+        self._unslice_params = OrderedSet()  # param's numel <= segment_size
+        self._unslice_params2align = {}  # {param.name: param's align}
+        self._grad_storages = {}  # {param.dtype: GradStorage}
+
+        self._ori_parameter_list = self._layer.parameters()
+        for param in self._ori_parameter_list:
+            param._need_shard = True
+        # check main_grad
+        self._check_main_grad()
+
+        # Synchronous all ranks models
+        if pretrain_sync_models:
+            self._sync_params_and_buffers()
+
+        self._segment_rank_params(self._layer)
+
+        # In the first step, record the execution order of the layer
+        self._order_tracer = OrderedDict()
+        self._order_tracer["order"] = 0
+        self._order_tracer["layer"] = []
+
+        # Add unslice params GradStorage
+        self._handle_unslice_params()
+
+        # Register task flow
+        self._task_flow = TaskFlow()
+
+        # Register forward hooks
+        self._register_forward_hooks(self._layer)
+
+        # Register backward parameter hooks
+        self._register_backward_hooks()
+
+    def _handle_unslice_params(self):
+        buffer_size = {}
+        buffer_size[Type.bf16.value] = 0
+        buffer_size[Type.fp32.value] = 0
+        buffer_size[Type.fp16.value] = 0
+        for param in self._unslice_params:
+            param2dtype[param.name] = param.dtype
+            p_align = self._param2align(param)
+            self._unslice_params2align[param.name] = p_align
+            buffer_size[param.dtype] += param._numel() + p_align
+
+        # Create unslice_params'grad
+        for param in sorted(self._unslice_params, key=lambda p: p.name):
+            if param.dtype not in self._grad_storages.keys():
+                self._grad_storages[param.dtype] = GradStorage(
+                    buffer_size[param.dtype],
+                    dtype=(
+                        param.dtype
+                        if not self.use_main_grad
+                        else paddle.float32
+                    ),
+                    device=self._default_device,
+                    destination=self._rank,
+                    param2align=self._unslice_params2align,
+                )
+            self._grad_storages[param.dtype].add_grad(
+                param, self._unslice_params2align[param.name]
+            )
+
+    def _check_main_grad(self):
+        self.use_main_grad = None
+        for param in self._layer.parameters():
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(param, "main_grad"), (
+                    "Params have different main grad attributes."
+                )
+
+    @paddle.autograd.no_grad()
+    def _sync_params_and_buffers(self):
+        """
+        Sync all model states for all ranks
+        """
+
+        for p in self._layer.parameters():
+            dist.broadcast(
+                p, src=self._global_root_rank, group=self._group, sync_op=True
+            )
+            if self._dp_group is not None and self._dp_group.nranks > 1:
+                dist.broadcast(
+                    p,
+                    src=self._dp_group.ranks[0],
+                    group=self._dp_group,
+                    sync_op=True,
+                )
+
+    def _sync_grad_storages_hook(self):
+        for grad_storage in self._grad_storages.values():
+            grad_storage.buffer.scale_(scale=self._world_size_scaling)
+            dist.all_reduce(tensor=grad_storage.buffer, group=self._group)
+            if self._dp_group is not None and self._dp_group.nranks > 1:
+                grad_storage.buffer.scale_(scale=(1.0 / self._dp_group.nranks))
+                dist.all_reduce(
+                    tensor=grad_storage.buffer, group=self._dp_group
+                )
+
+    def forward(self, *inputs, **kwargs):
+        """
+        A wrapper for Sharding Stage3 layer.
+        """
+        # add hook to sync grad storage
+        for grad_storage in self._grad_storages.values():
+            grad_storage.buffer.zero_()
+            grad_storage.manual_release()
+            grad_storage.rebuild()
+        core.eager._add_backward_final_hook(self._sync_grad_storages_hook)
+
+        # 1.Sync layer's buffers state
+        if self.__sync_buffers:
+            self._sync_buffers()
+
+        # 2.Normal FW on the base model
+        fw = self._layer(*inputs, **kwargs)
+
+        return fw
+
+    def _segment_rank_params(self, layer, name="last_layer"):
+        """
+        Flatten parameters according to layer.
+        """
+        current_layer_params = _current_layer_params(layer)
+        if current_layer_params:
+            self._flatten_layer_params(layer, current_layer_params)
+
+        for name, sub_layer in layer.named_children():
+            self._segment_rank_params(sub_layer, name)
+
+    def _flatten_layer_params(self, layer, current_layer_params):
+        """
+        Parameter segmentation and memory integration.
+        """
+
+        if id(layer) in self._trainable_params.keys():
+            return
+
+        def _add_manage_info(trainable_param):
+            return _PartitionParam(trainable_param)
+
+        current_params = []
+        for p in current_layer_params:
+            if p._numel() > self._segment_size:
+                current_params.append(_add_manage_info(p))
+            elif p.trainable:
+                self._unslice_params.add(_UnsliceParam(p))
+
+        self._trainable_params[id(layer)] = current_params
+
+        for param in self._trainable_params[id(layer)]:
+            if param.name in self._param2buffer.keys():
+                continue
+            self._param2buffer[param.name] = []
+            # 1.Params alignment
+            align_ = self._param2align(param)
+
+            offset = align_ + param._numel()
+            buffer_size = (
+                offset
+                if offset % self._group.nranks == 0
+                else offset + self._group.nranks - (offset % self._group.nranks)
+            )
+            self._param2buffer_size[param.name] = buffer_size
+
+            # 2.Combination param buffer
+            assert buffer_size % self._group.nranks == 0
+            pre_buffer = buffer_size // self._group.nranks
+
+            for rank_ in range(self._group.nranks):
+                self._param2buffer[param.name].append(
+                    (rank_ * pre_buffer, (rank_ + 1) * pre_buffer)
+                )
+
+            # Record param's dtype
+            param2dtype[param.name] = param.dtype
+            # 3.Flatten layer params and release other rank buffer
+            self._param_storage(param, buffer_size)
+
+    def _param_storage(self, param, buffer_size):
+        """
+        This is a function to simplify the handling of parameter InternalStorages.
+        """
+        assert isinstance(buffer_size, int)
+        value = (
+            np.zeros(buffer_size, dtype=np.float16)
+            if (
+                Type.fp16.value == param.dtype or Type.bf16.value == param.dtype
+            )
+            else np.zeros(buffer_size, dtype=np.float32)
+        )
+        buffer = core.eager.Tensor(value=value, place=core.CPUPlace())
+        if Type.bf16.value == param.dtype:
+            buffer = buffer.cast(Type.bf16.value)
+
+        param_shape = param.shape
+        origin_state = param.stop_gradient
+        param.stop_gradient = True
+        param.flatten_()
+        param.stop_gradient = origin_state
+        start, end = self._param2buffer[param.name][self._rank]
+
+        # Copy the current param value
+        with device_guard():
+            tmp_var = buffer._slice(0, param._numel())
+        param_cpu = param.cpu()
+        tmp_var.get_tensor().set(param_cpu.get_tensor(), core.CPUPlace())
+        del tmp_var
+        param.get_tensor()._set_dims(param_shape)
+
+        # Current rank param_storage
+        param.fw_storage = core.eager.Tensor(
+            value=buffer._slice(start, end), name="slice@" + param.name
+        )
+        param.status = "part"
+        param._clear_data()
+
+    def _register_forward_hooks(self, layer):
+        """
+        Register PyLayer to manage memory slices.
+        There are four stages:
+        FW
+        1. Before the forward layers, synchronize the full parameters.
+        2. After the forward layers, release the full parameter and keep the parameter slice.
+        BW
+        3. Before the backward layers, synchronize the full parameters and create param's grad.
+        4. After the gradient accumulation, release the full parameter and keep the parameter slice.
+        """
+        current_layer_params = _current_layer_params(layer)
+        if current_layer_params:
+            # the layer in self._exclude_layer will be added hooks.
+            if not (
+                id(layer) in self._exclude_layer
+                or layer.__class__.__name__ in self._exclude_layer
+            ):
+                self._register_forward_all_hooks(layer, self._task_flow)
+
+        for _, sub_layer in layer.named_children():
+            self._register_forward_hooks(sub_layer)
+
+    def _register_forward_all_hooks(self, sub_layer, task_flow):
+        def _forward_pre_hook(layer, inputs):
+            return ForwardPreHooks(
+                layer,
+                self._order_tracer,
+                self._trainable_params,
+                self._param2buffer_size,
+                self._group,
+                self._sync_comm,
+                self._offload,
+                task_flow,
+            )
+
+        def _forward_post_hook(layer, inputs, outputs):
+            if isinstance(outputs, paddle.Tensor):
+                outputs = (outputs,)
+            return ForwardPostHooks.apply(
+                *outputs,
+                layer=layer,
+                order_tracer=self._order_tracer,
+                trainable_params=self._trainable_params,
+                param2buffer=self._param2buffer,
+                param2buffer_size=self._param2buffer_size,
+                rank=self._rank,
+                group=self._group,
+                sync_comm=self._sync_comm,
+                offload=self._offload,
+                task_flow=task_flow,
+            )
+
+        # register previous forward hooks
+        sub_layer.register_forward_pre_hook(_forward_pre_hook)
+
+        # register post forward hooks
+        sub_layer.register_forward_post_hook(_forward_post_hook)
+
+    @paddle.autograd.no_grad()
+    def _sync_buffers(self):
+        """
+        Sync all the param buffers from all ranks (exp: batch norm statistics).
+        """
+
+        for buffer in self._layer.buffers(include_sublayers=True):
+            dist.broadcast(
+                buffer, self._global_root_rank, self._group, sync_op=True
+            )
+            if self._dp_group is not None and self._dp_group.nranks > 1:
+                dist.broadcast(
+                    buffer,
+                    self._dp_group.ranks[0],
+                    self._dp_group,
+                    sync_op=True,
+                )
+
+    def __getattr__(self, name):
+        """Forward missing attributes to wrapped layer."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self._layer, name)
+
+    def _register_backward_hooks(self):
+        current_layer_params = self._layer.parameters(include_sublayers=True)
+        trainable_params = list(
+            filter(
+                lambda p: p.trainable and p not in self._unslice_params,
+                current_layer_params,
+            )
+        )
+
+        for param in trainable_params:
+            allreduce_function = self._get_allreduce_fn(param)
+            param._register_backward_hook(allreduce_function)
+
+    def _get_allreduce_fn(self, param):
+        @paddle.autograd.no_grad()
+        def allreduce_(*_):
+            assert param.trainable, (
+                "the param must be trainable for grad allreduced"
+            )
+            if param.name in self._task_flow.full_grad.keys():
+                full_grad = self._task_flow.full_grad[param.name]
+                # Only support sync allreduce current rank's layer now
+                full_grad.scale_(scale=self._world_size_scaling)
+                dist.all_reduce(tensor=full_grad, group=self._group)
+                if self._dp_group is not None and self._dp_group.nranks > 1:
+                    full_grad.scale_(scale=1.0 / self._dp_group.nranks)
+                    dist.all_reduce(tensor=full_grad, group=self._dp_group)
+
+                start, end = self._param2buffer[param.name][self._rank]
+                if param.bw_storage is None:
+                    param.bw_storage = (
+                        full_grad._slice(start, end).detach().clone()
+                    )
+                else:
+                    param.bw_storage = paddle.add(
+                        param.bw_storage,
+                        full_grad._slice(start, end).detach().clone(),
+                    )
+
+                if self.use_main_grad:
+                    param.main_grad = None
+                else:
+                    param.clear_gradient(False)
+                del self._task_flow.full_grad[param.name]
+
+            if param.name in self._task_flow.full_param.keys():
+                if param.status == "all":
+                    param.use_count = 0
+                    param._clear_data()
+                    start, end = self._param2buffer[param.name][self._rank]
+                    param.fw_storage = (
+                        self._task_flow.full_param[param.name][0]
+                        ._slice(start, end)
+                        .detach()
+                        .clone()
+                    )
+                    param.status = "part"
+                    del self._task_flow.full_param[param.name]
+
+        return allreduce_
+
+    def _param2align(self, param):
+        # CUDA alignment 256 bytes
+        size = param._numel() * align[param.dtype]
+        device_alignment = alignment[self._default_device]
+        remaining = size % device_alignment
+        ali = 0 if remaining == 0 else device_alignment - remaining
+        align_ = ali // align[param.dtype]
+        return align_

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_optimizer_stage2.py

@@ -0,0 +1,792 @@
+#   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.
+
+# The file has been adapted from fairscale file:
+# https://github.com/facebookresearch/fairscale/blob/main/fairscale/optim/oss.py
+# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
+# We retain the following license from the original files:
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import warnings
+from collections import OrderedDict
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import ParallelMode, fleet
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
+    ShardedStateDict,
+    ShardedWeight,
+    create_sharded_weight_with_new_local,
+)
+from paddle.framework import core
+from paddle.nn import ClipGradByGlobalNorm
+from paddle.optimizer import Optimizer
+
+HybridParallelClipGrad = fleet.meta_optimizers.dygraph_optimizer.hybrid_parallel_optimizer.HybridParallelClipGrad
+from paddle.distributed.collective import _get_global_group, new_group
+
+from .group_sharded_storage import GradStorage, ParamStorage
+from .group_sharded_utils import GroupShardedClipGrad, Type, device_guard
+
+# CUDA alignment 256 bytes, cpu alignment 4096 bytes
+alignment = {"gpu": 256, "cpu": 4096, "xpu": 256}
+align = {
+    Type.fp16.value: 2,
+    Type.bf16.value: 2,
+    Type.fp32.value: 4,
+}
+
+
+class GroupShardedOptimizerStage2(Optimizer):
+    """
+    A wrapper for Sharding Stage2 Optimizer in Dygraph.
+
+    .. warning: ShardingOptimizer encapsulates the optimization strategy and integrates it into the optimizer.
+
+    .. ZeRO: 1.https://arxiv.org/pdf/1910.02054.pdf 2.https://arxiv.org/pdf/1910.02054.pdf.
+
+    """
+
+    # TODO (Baibaifan)
+    # Feature Notes:
+    # 1. Unified memory for parameters and parameters.grad to InternalStorage.
+    # 2. Support the segmentation of optimizer parameters and partial updating of parameters.
+    # 3. Dynamically adjust training parameters and models.
+    # 4. Support offload function.
+    # 5. Support the establishment of independent communication groups.
+    # 6. Broadcast_fp16 is not supported now.
+    def __init__(
+        self,
+        params,
+        optim,
+        group=None,
+        offload=False,
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
+        pretrain_sync_models=True,
+        dp_group=None,
+        **kw,
+    ):
+        super().__init__(learning_rate=optim._learning_rate, parameters=params)
+        assert (
+            core.is_compiled_with_cuda()
+            or core.is_compiled_with_xpu()
+            or (device in core.get_all_custom_device_type())
+        ), "Only GPU and XPU and CustomDevice is supported now"
+
+        # Segmentation information
+        self._dtype_rank_params = (
+            OrderedDict()
+        )  # {dtype:[param1,param2]} device, rank, params
+        self._param2rank = {}
+        self.__segment_params = []
+        self._rank_buffer_size = {}  # {dtype: {rank: numel+alignment}}
+        self._param2align = {}  # {param.name: align}
+
+        # Default information
+        self._optim = optim
+
+        # sharing stage 2 comm overlap flag
+        self._reduce_overlap = False
+        # record the last task used for comm overlap for sharding stage 2
+        self._comm_task = None
+
+        assert hasattr(self._optim, "_master_weights"), (
+            "Must use optimizer with _master_weights attribute"
+        )
+
+        # Support parameter group and parameter list
+        self._local_params = []
+        if isinstance(params[0], dict):
+            for param_group in params:
+                self._local_params.extend(list(param_group["params"]))
+        else:
+            self._local_params.extend(list(params))
+
+        self.use_main_grad = None
+        for param in self._local_params:
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(param, "main_grad"), (
+                    "Params have different main grad attributes."
+                )
+        if self.use_main_grad:
+            assert not offload, "offload not support main_grad for now"
+
+        self._default_device = device
+        self._pfp16 = (
+            len(
+                list(
+                    filter(
+                        lambda x: x.trainable and x.dtype == Type.fp16.value,
+                        self._local_params,
+                    )
+                )
+            )
+            > 0
+        )
+        self._pbf16 = (
+            len(
+                list(
+                    filter(
+                        lambda x: x.trainable and x.dtype == Type.bf16.value,
+                        self._local_params,
+                    )
+                )
+            )
+            > 0
+        )
+
+        self._broadcast_overlap = False
+        self._forward_pre_hook_remove_helper = []
+        try:
+            # The fp32 params such as layer_norm_0.w_0 will be at the end of param_list.
+            # Have to sort the params to make sure all params are in the forward using order.
+            self._broadcast_order_params = sorted(
+                self.local_params,
+                key=lambda x: int(x.name.split('.')[0].split('_')[-1]),
+            )
+        except ValueError:
+            self._broadcast_order_params = None
+
+        self._group = (
+            new_group(_get_global_group().ranks) if group is None else group
+        )
+
+        # only support to combine stage2 and dp hybrid parallel now.
+        self._dp_group = dp_group
+        self.world_size = self._group.nranks
+        self._rank = self._group.rank
+        self._global_root_rank = self._group.ranks[0]
+
+        if self._dp_group is not None and self._dp_group.nranks > 1:
+            assert not offload, (
+                "Not support! when using offload with sharding stage2, please use pure sharding stage2, exclude data parallel."
+            )
+
+        # Synchronous all ranks models
+        if pretrain_sync_models:
+            self._sync_params_and_buffers()
+
+        self.param_storages = {}  # {dtype: {rank: InternalStorage}}
+
+        if isinstance(self._optim._grad_clip, ClipGradByGlobalNorm):
+            logging.warning(
+                "While using ClipGradByGlobalNorm in GroupShardedOptimizerStage2, the grad clip of original optimizer will be changed."
+            )
+
+            hcg = fleet.fleet._hcg if hasattr(fleet.fleet, "_hcg") else None
+            if (
+                hcg
+                and hcg.get_parallel_mode() is not ParallelMode.DATA_PARALLEL
+                and not offload
+            ):
+                if self.use_main_grad:
+                    self._optim._inner_opt._grad_clip = HybridParallelClipGrad(
+                        self._optim._inner_opt._grad_clip, hcg
+                    )
+                else:
+                    self._optim._grad_clip = HybridParallelClipGrad(
+                        self._optim._grad_clip, hcg
+                    )
+            else:
+                if self.use_main_grad:
+                    self._optim._inner_opt._grad_clip = GroupShardedClipGrad(
+                        self._optim._inner_opt._grad_clip,
+                        paddle.get_device(),
+                        self._group,
+                    )
+                else:
+                    self._optim._grad_clip = GroupShardedClipGrad(
+                        self._optim._grad_clip, paddle.get_device(), self._group
+                    )
+
+            if self._optim._parameter_list and isinstance(
+                self._optim._parameter_list[0], dict
+            ):
+                for item in self._optim._param_groups:
+                    if "grad_clip" in item.keys():
+                        item["grad_clip"] = self._optim._grad_clip
+
+        if offload:
+            assert self._pfp16, (
+                "Only support offload strategy while using 'Adam', 'AdamW' and 'Momentum' optimizer with AMP/Pure FP16"
+            )
+
+        self.offload = offload  # Using for offload
+        self.offload_device = "cpu"
+        self.offload_buffer_size = 0
+        self.offload_param2align = {}
+        self.offload_params = None
+        self.offload_grads = None
+        self.dev_id = int(paddle.get_device().split(":")[1])
+
+        self._master_params = {}
+
+        # Update optimizer parameters and adjust parameter storage and use according to rank.
+        self._update_opt_status()
+
+    def _set_auxiliary_var(self, key, val):
+        super()._set_auxiliary_var(key, val)
+        self._optim._set_auxiliary_var(key, val)
+
+    @paddle.autograd.no_grad()
+    def _sync_params_and_buffers(self):
+        """
+        Sync all model states for all ranks
+        """
+
+        for p in self._local_params:
+            dist.broadcast(
+                p, src=self._global_root_rank, group=self._group, sync_op=True
+            )
+
+            if self._dp_group:
+                dist.broadcast(
+                    p,
+                    src=self._dp_group.ranks[0],
+                    group=self._dp_group,
+                    sync_op=True,
+                )
+
+    def _update_task(self, task):
+        if self._reduce_overlap:
+            assert task is not None
+        # Only track of the last reduce task.
+        # Since all tasks are on the same stream, only need to wait the last one.
+        # After waiting for the last reduce task, all reduce tasks before have already finished.
+        self._comm_task = task
+
+    def _set_reduce_overlap(self, reduce_overlap):
+        # Enable gradients' reduces overlap with backward calculation.
+        self._reduce_overlap = reduce_overlap
+
+    def _set_broadcast_overlap(
+        self, broadcast_overlap, layers=None, num_groups=None
+    ):
+        # Enable post optimizer broadcasts overlap with the forward calculation of next batch.
+        self._broadcast_overlap = broadcast_overlap
+        if self._broadcast_overlap:
+            assert layers is not None, (
+                "To enable broadcast overlap forward, please pass the module to the function."
+            )
+            self._layers = layers
+            warnings.warn(
+                "Setting overlap broadcast means the `paddle.device.cuda.synchronize()` "
+                "must be called manually before calling `paddle.save()` and before and inference."
+            )
+            if self._broadcast_order_params is None:
+                # Params' names should be like column_linear_32.w_0 pattern to get the best performance.
+                warnings.warn(
+                    r"The param name passed to the optimizer doesn't follow .+_[0-9]+\..+ pattern, "
+                    "overlap broadcast may harm the performance."
+                )
+                self._broadcast_order_params = self._local_params
+
+        if num_groups is None or num_groups > len(self._broadcast_order_params):
+            warnings.warn(
+                "The num_groups for broadcast is larger than the number of params to be broadcast. "
+                "It will set to default value: 1 (use the default sharding group)."
+            )
+            num_groups = 1
+
+        assert isinstance(num_groups, int) and num_groups > 0, (
+            "num_groups should be a positive integer"
+        )
+
+        self._number_of_broadcast_groups = num_groups
+        self._broadcast_groups = [
+            None for _ in range(self._number_of_broadcast_groups)
+        ]
+        self._broadcast_groups[0] = self._group
+
+        ranks = self._group.ranks
+        for i in range(1, self._number_of_broadcast_groups):
+            self._broadcast_groups[i] = new_group(ranks)
+
+    def _generate_master_params(self, trainable_params):
+        if self.offload:
+            for param in trainable_params:
+                if param.name not in self._master_params.keys():
+                    self._master_params[param.name] = core.eager.Tensor(
+                        name=param.name,
+                        value=param.cast(dtype=Type.fp32.value).numpy(),
+                        place=core.CPUPlace(),
+                        stop_gradient=param.stop_gradient,
+                    )
+        else:
+            for param in trainable_params:
+                if (
+                    param.dtype == Type.fp16.value
+                    or param.dtype == Type.bf16.value
+                ):
+                    master_tensor = paddle.cast(param, Type.fp32.value)
+                    master_tensor.name = param.name
+                    self._optim._master_weights[param.name] = master_tensor
+
+    def _update_opt_status(self):
+        """Update optimizer status and parameter storage information, and special functions to be developed."""
+        # func 1
+        self._integration_params()
+
+    # Segment helpers
+
+    def _segment_params(self):
+        """
+        Divide all optimizer parameters equally into rank.
+        """
+        if len(self.__segment_params) == 0:
+            self.__segment_params, param_lists = (
+                [[] for _ in range(self.world_size)],
+                [[] for _ in range(self.world_size)],
+            )
+            sizes = [0] * self.world_size
+            for param in self._local_params:
+                # Add this param to rank with smallest size.
+                rank = sizes.index(min(sizes))
+                param_lists[rank].append(param)
+
+                # Statistical real numels
+                sizes[rank] += param._numel() if param.trainable else 0
+
+            for rank, params in enumerate(param_lists):
+                self.__segment_params[rank].extend(params)
+        return self.__segment_params
+
+    @property
+    def local_params(self):
+        return self._local_params
+
+    @property
+    def param2rank(self):
+        """Map the params to the rank which owns them"""
+        if len(self._param2rank) == 0:
+            for rank, params in enumerate(self._segment_params()):
+                for param in params:
+                    self._param2rank[param.name] = rank
+        return self._param2rank
+
+    @property
+    def dtype_rank_params(self):
+        """
+        Divide the parameters into groups according to rank and dtype.
+        """
+        if len(self._dtype_rank_params) == 0:
+            # Assign the parameters of each rank according to the type
+            trainable_params = list(
+                filter(lambda x: x.trainable, self._local_params)
+            )
+            for param in trainable_params:
+                if param.dtype not in self._dtype_rank_params.keys():
+                    self._dtype_rank_params[param.dtype] = [
+                        [] for _ in range(self.world_size)
+                    ]
+                self._dtype_rank_params[param.dtype][
+                    self.param2rank[param.name]
+                ].append(param)
+
+            # Sort per rank params by size
+            for dtype in self._dtype_rank_params.keys():
+                for rank_params in self._dtype_rank_params[dtype]:
+                    rank_params.sort(key=lambda x: x._numel())
+
+        return self._dtype_rank_params
+
+    @property
+    def rank_buffer_size(self):
+        """
+        Count the memory size of the parameters corresponding to rank under the corresponding dtype.
+        """
+        # CUDA alignment 256 bytes
+        if self._default_device in core.get_all_custom_device_type():
+            device_alignment = core.libpaddle._get_device_min_chunk_size(
+                self._default_device
+            )
+        else:
+            device_alignment = alignment[self._default_device]
+
+        if len(self._rank_buffer_size) == 0:
+            for dtype in self.dtype_rank_params.keys():
+                if dtype not in self._rank_buffer_size.keys():
+                    self._rank_buffer_size[dtype] = {}
+                for dst_rank, per_rank_params in enumerate(
+                    self.dtype_rank_params[dtype]
+                ):
+                    if dst_rank not in self._rank_buffer_size[dtype].keys():
+                        self._rank_buffer_size[dtype][dst_rank] = 0
+                    for param in per_rank_params:
+                        if not param.trainable:
+                            continue
+                        size = param._numel() * align[dtype]
+                        remaining = size % device_alignment
+                        ali = (
+                            0
+                            if remaining == 0
+                            else device_alignment - remaining
+                        )
+                        align_ = ali // align[dtype]
+                        self._rank_buffer_size[dtype][dst_rank] += (
+                            param._numel() + align_
+                        )
+                        self._param2align[param.name] = align_
+
+        return self._rank_buffer_size
+
+    def _integration_params(self):
+        """
+        Integrate the parameters into a continuous memory according to rank, and support the update of training parameters.
+        """
+
+        for dtype, per_rank_params in self.dtype_rank_params.items():
+            if dtype not in self.param_storages.keys():
+                self.param_storages[dtype] = {}
+
+            for dst_rank, params in enumerate(per_rank_params):
+                if len(params) > 0:
+                    # Merge all the trainable params in a single InternalStorage
+                    trainable_params = list(
+                        filter(lambda x: x.trainable, params)
+                    )
+                    if (self._pfp16 or self._pbf16) and dst_rank == self._rank:
+                        self._generate_master_params(trainable_params)
+                    if trainable_params:
+                        param_storage = ParamStorage(
+                            size=self.rank_buffer_size[dtype][dst_rank],
+                            dtype=dtype,
+                            device=self._default_device,
+                        )
+
+                        param_storage.add_rank_params(
+                            trainable_params, self._param2align
+                        )
+                        self.param_storages[dtype][dst_rank] = param_storage
+
+        # Clear the InternalStorage keys which are not in use anymore
+        dtype_in_use = list(self.dtype_rank_params.keys())
+        dtype_to_pop = list(
+            filter(lambda x: x not in dtype_in_use, self.param_storages.keys())
+        )
+        for d in dtype_to_pop:
+            self.param_storages.pop(d)
+
+        if self.offload:
+            self._optim._master_weights = self._master_params
+            cpu_master_params = list(self._master_params.values())
+            if self._default_device in core.get_all_custom_device_type():
+                device_alignment = core.libpaddle._get_device_min_chunk_size(
+                    self._default_device
+                )
+            else:
+                device_alignment = alignment[self._default_device]
+
+            for param in cpu_master_params:
+                size = param._numel() * align[Type.fp32.value]
+                remaining = size % device_alignment
+                ali = 0 if remaining == 0 else device_alignment - remaining
+                align_ = ali // align[Type.fp32.value]
+                self.offload_buffer_size += param._numel() + align_
+                self.offload_param2align[param.name] = align_
+
+            if cpu_master_params:
+                with device_guard(self._rank, self.offload_device):
+                    self.offload_params = ParamStorage(
+                        size=self.offload_buffer_size,
+                        dtype=Type.fp32.value,
+                        device=self.offload_device,
+                    )
+                    self.offload_params.buffer.name = "offload_buffer"
+                    self.offload_params.add_rank_params(
+                        cpu_master_params, self.offload_param2align, False
+                    )
+                    self.offload_params.buffer.stop_gradient = False
+
+                    self.offload_grads = GradStorage(
+                        size=self.offload_buffer_size,
+                        dtype=Type.fp32.value,
+                        device=self.offload_device,
+                        destination=self._rank,
+                        param2align=self.offload_param2align,
+                        convert_cpu=True,
+                    )
+                    for p in cpu_master_params:
+                        self.offload_grads.add_grad(
+                            p, self.offload_param2align[p.name]
+                        )
+
+                    self._optim._master_weights[
+                        self.offload_params.buffer.name
+                    ] = self.offload_params.buffer
+
+    def _offload_acc_grad(self, param_name, grad_fp32_cpu):
+        """accumulate grads with offload strategy"""
+        with device_guard(self._rank, self.offload_device):
+            if param_name in self._master_params.keys():
+                if self._master_params[param_name].grad is None:
+                    self._master_params[param_name]._copy_gradient_from(
+                        grad_fp32_cpu
+                    )
+                else:
+                    self._master_params[param_name].grad.add_(grad_fp32_cpu)
+
+        self.offload_params.buffer._copy_gradient_from(
+            self.offload_grads.buffer
+        )
+
+    def _offload_scale_grad(self, scale_size):
+        """scale grads with offload strategy"""
+        with device_guard(self._rank, self.offload_device):
+            self.offload_grads.buffer.scale_(scale=scale_size)
+
+    def _offload_clear_grad(self):
+        """clear grads with offload strategy"""
+        with device_guard(self._rank, self.offload_device):
+            self.offload_grads.buffer.zero_()
+
+    def _step(self):
+        if self._broadcast_overlap:
+            # Clear the pre forward hook in the optimizer step.
+            for hook_remove in self._forward_pre_hook_remove_helper:
+                hook_remove.remove()
+            self._forward_pre_hook_remove_helper = []
+
+        if self.offload:
+            params_list = [self.offload_params.buffer]
+
+            # TODO(Baibaifan): Offload will support param_groups later
+            if not isinstance(self._optim._param_groups[0], dict):
+                self._optim._parameter_list = params_list
+                self._optim._param_groups = params_list
+
+        # Run the optimizer of the current rank step
+        if self.offload:
+            with device_guard(device=self.offload_device):
+                self._optim.step()
+
+            for param in self._local_params:
+                if param.name in self._master_params.keys():
+                    if (
+                        self._default_device
+                        in core.get_all_custom_device_type()
+                    ):
+                        param.set_value(
+                            self._master_params[param.name]
+                            ._copy_to(
+                                paddle.CustomPlace(
+                                    self._default_device, self.dev_id
+                                ),
+                                True,
+                            )
+                            .cast(dtype=param.dtype)
+                        )
+                    elif self._default_device == "xpu":
+                        param.set_value(
+                            self._master_params[param.name]
+                            .to("xpu:" + str(self.dev_id))
+                            .cast(dtype=param.dtype)
+                        )
+                    else:
+                        param.set_value(
+                            self._master_params[param.name]
+                            .cuda(self.dev_id)
+                            .cast(dtype=param.dtype)
+                        )
+        else:
+            self._optim.step()
+
+        # Synchronize all the updated shards in between the ranks
+        self._broadcast_params()
+
+    def step(self):
+        """
+        A wrapper for Optimizer's step function to finish the update operation of the optimizer.
+        """
+        # This method won't be called directly by opt.step()!
+        # The _redefine_opt_step() in class GroupShardedStage2 will wrap this function.
+        self._step()
+
+    def minimize(self):
+        raise RuntimeError(
+            "optimizer.minimize() not support now, please use optimizer.step()"
+        )
+
+    def set_state_dict(self, state_dict):
+        self._optim.set_state_dict(state_dict)
+
+    def state_dict(self):
+        return self._optim.state_dict()
+
+    def _clear_cache(self):
+        self.__segment_params.clear()
+        self._dtype_rank_params.clear()
+        self._param2rank.clear()
+
+    @paddle.autograd.no_grad()
+    def _broadcast_params(self):
+        """Broadcast the parameters of the current rank to each rank"""
+
+        # Exchange all the shards with the other ranks
+        if self._broadcast_overlap:
+            self._broadcast_params_overlap_forward()
+        else:
+            for dtype_per_rank in self.param_storages.values():
+                for dst_rank, internal_storage in dtype_per_rank.items():
+                    dist.broadcast(
+                        tensor=internal_storage.buffer,
+                        src=self._group.ranks[dst_rank],
+                        group=self._group,
+                        sync_op=True,
+                    )
+
+    def _forward_pre_hook_function(self, tasks):
+        # Since the layers will call pre hook by `forward_pre_hook(self, inputs)`,
+        # the helper functions needs the x and y to take those params.
+        def __impl__(x, y):
+            for task in tasks:
+                # Wait for broadcast task before using the result of the broadcast.
+                task.wait()
+
+        return __impl__
+
+    def set_lr(self, lr):
+        super().set_lr(lr)
+        self._optim.set_lr(lr)
+
+    def get_lr(self):
+        return self._optim.get_lr()
+
+    @paddle.autograd.no_grad()
+    def _broadcast_params_overlap_forward(self):
+        # Exchange all the shards with the other ranks,
+        # but overlap the broadcast with next batch's calculation.
+        group_idx = 0
+
+        param2task = {}
+        for x in self._broadcast_order_params:
+            if x.trainable:
+                group = self._broadcast_groups[group_idx]
+                group_idx = (group_idx + 1) % self._number_of_broadcast_groups
+                task = dist.broadcast(
+                    tensor=x,
+                    src=group.ranks[self._param2rank[x.name]],
+                    group=group,
+                    sync_op=False,
+                )
+                assert x.name not in param2task
+                param2task[x.name] = task
+
+        for layer in self._layers.sublayers():
+            if len(layer.sublayers()) == 0:
+                # Register forward pre hood for leaf layers. This will get the best performance.
+                tasks = []
+                for param in layer.parameters():
+                    if param.trainable:
+                        if param.name in param2task:
+                            tasks.append(param2task[param.name])
+                self._forward_pre_hook_remove_helper.append(
+                    layer.register_forward_pre_hook(
+                        self._forward_pre_hook_function(tasks)
+                    )
+                )
+
+    def sharded_state_dict(
+        self,
+        model_sharded_state_dict: ShardedStateDict,
+    ) -> ShardedStateDict:
+        """
+        Convert optimizer state dict to a sharded state dict based on model sharding information.
+
+        Args:
+            model_sharded_state_dict (dict): Sharded state dict of the model, containing tensor metadata.
+
+        Returns:
+            dict: A new optimizer state dict where weights are wrapped as ShardedWeight.
+        """
+
+        _FP32_MASTER = "fp32_master_0"
+        _MOMENT_NAME = "moment"
+        _optimizer_scalar_name = [
+            "beta1_pow_acc_0",
+            "beta2_pow_acc_0",
+        ]
+        _optimizer_non_scaler_name = [
+            "moment1_0",
+            "moment2_0",
+            "velocity_0",
+        ]
+
+        def _generate_base_static_name(vname):
+            if _FP32_MASTER in vname:
+                return tuple(vname.split("_" + _FP32_MASTER + "_", 1))
+            for name in _optimizer_scalar_name + _optimizer_non_scaler_name:
+                if vname.endswith(name):
+                    return vname[: -(len(name) + 1)], name
+            raise ValueError(f"Cannot split variable name: {vname}.")
+
+        optimizer_sharded_state_dict = {}
+        optimizer_state_dict = self.state_dict()
+        # Build name mapping and remove non-tensor entries from optimizer state
+        static_to_struct_mapping = {}
+        model_sharded_state_dict = dict(
+            sorted(model_sharded_state_dict.items())
+        )
+        for k, v in model_sharded_state_dict.items():
+            # When shared weights exist, the v.local_tensor.name of shared parameters are identical, but only the first parameter has optimizer states. Therefore, only the key-value pairs of the first occurrence in the shared parameter group need to be retained.
+            if v.local_tensor.name not in static_to_struct_mapping:
+                static_to_struct_mapping[v.local_tensor.name] = k
+
+        master_weights = optimizer_state_dict.pop("master_weights", None)
+        optimizer_state_dict.pop("LR_Scheduler", None)
+
+        # Process main optimizer states
+        for key, tensor in optimizer_state_dict.items():
+            static_name, optim_state_type = _generate_base_static_name(key)
+            struct_name = static_to_struct_mapping[static_name]
+            sharded_weight = model_sharded_state_dict[struct_name]
+
+            unified_name = f"{struct_name}.{optim_state_type}"
+
+            # Determine tensor partitioning scheme
+            if _MOMENT_NAME in optim_state_type:
+                optimizer_sharded_state_dict[unified_name] = (
+                    create_sharded_weight_with_new_local(
+                        unified_name, tensor, sharded_weight
+                    )
+                )
+            else:  # Non-momentum parameters
+                optimizer_sharded_state_dict[unified_name] = ShardedWeight(
+                    key=unified_name,
+                    local_tensor=tensor,
+                    local_shape=(1,),
+                    global_shape=(1,),
+                    global_offset=(0,),
+                )
+
+        # Process master weights if using mixed precision
+        if master_weights is not None:
+            for key, tensor in master_weights.items():
+                struct_name = static_to_struct_mapping[key]
+                sharded_weight = model_sharded_state_dict[struct_name]
+                unified_name = f"{struct_name}.w_0"
+                optimizer_sharded_state_dict[unified_name] = (
+                    create_sharded_weight_with_new_local(
+                        unified_name, tensor, sharded_weight
+                    )
+                )
+
+        return optimizer_sharded_state_dict

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_stage2.py

@@ -0,0 +1,720 @@
+#   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.
+
+# The file has been adapted from fairscale file:
+# https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/data_parallel/sharded_ddp.py
+# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
+# We retain the following license from the original files:
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from functools import reduce
+from types import MethodType
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import collective
+from paddle.distributed.utils.log_utils import get_logger
+from paddle.framework import core
+
+from .group_sharded_optimizer_stage2 import GroupShardedOptimizerStage2
+from .group_sharded_storage import GradStorage
+from .group_sharded_utils import Type, device_guard
+
+logger_ = get_logger(logging.WARNING)
+
+
+def _trainable(param):
+    return param.trainable
+
+
+class GroupShardedStage2(nn.Layer):
+    """
+    A wrapper for Sharding Stage2 Layer in Dygraph.
+    .. warning: GroupShardedStage2 encapsulates the layer strategy and integrates it into the nn.Layer.
+    .. ZeRO: https://arxiv.org/pdf/1910.02054.pdf.
+    """
+
+    # TODO (Baibaifan)
+    # Feature Notes::
+    # 1. Unified memory for param and param.grad to InternalStorage.
+    # 2. Divide param.grad according to rank to centrally apply for and release GPU memory.
+    # 3. Dynamically adjust training parameters and models.
+    # 4. Support offload function.
+    # 5. Support the establishment of independent communication groups.
+
+    def __init__(
+        self,
+        layer,
+        sharding_optimizer,
+        group=None,
+        sync_buffers=False,
+        buffer_max_size=2**23,  # 8MB
+        auto_refresh_trainable=True,
+        device="xpu" if core.is_compiled_with_xpu() else "gpu",
+        dp_group=None,
+    ):
+        super().__init__()
+
+        # training options
+        self._layer = layer
+        self._sharding_optimizers = (
+            [sharding_optimizer]
+            if not isinstance(sharding_optimizer, list)
+            else sharding_optimizer
+        )
+        assert all(
+            isinstance(opt, GroupShardedOptimizerStage2)
+            for opt in self._sharding_optimizers
+        ), "Please use GroupShardedOptimizerStage2 optimizer"
+        self._sync_buffers = sync_buffers
+        self._auto_refresh_trainable = auto_refresh_trainable
+
+        # Communication related attributes
+        self._group = (
+            collective.new_group(collective._get_global_group().ranks)
+            if group is None
+            else group
+        )
+        self._world_size_scaling = 1.0 / self._group.nranks
+        assert self._group.nranks > 1, (
+            "Training must be distributed, ranks must be greater than 1"
+        )
+        self._rank = self._group.rank
+        self._global_root_rank = self._group.ranks[
+            0
+        ]  # picking ranks index 0 as the reference
+        self._default_device = device
+
+        self._dp_group = dp_group
+
+        # Global statistical parameters
+        self._all_params = []
+        for optim in self._sharding_optimizers:
+            self._all_params.extend(list(optim.local_params))
+        self.use_main_grad = None
+        for param in self._all_params:
+            if self.use_main_grad is None and hasattr(param, "main_grad"):
+                self.use_main_grad = True
+            if self.use_main_grad:
+                assert hasattr(param, "main_grad"), (
+                    "Params have different main grad attributes."
+                )
+
+        # sharing stage 2 comm overlap flag
+        self._reduce_overlap = False
+
+        self._grad_reduced = []
+        self._trainable_param2rank = {}
+        self._trainable_param2align = {}
+        self._trainable_params = list(
+            filter(lambda x: x.trainable, self._all_params)
+        )
+        self._trainable_mask = list(map(_trainable, self._trainable_params))
+        self._param_grads = []
+
+        # Set grad storage size & Display param sizes and model sizes
+        model_size = sum([p._numel() for p in self._layer.parameters()])
+        assert buffer_max_size >= 0, "buffer_max_size must be GE than 0."
+        self._buffer_max_size = self._rank_buffer_size(
+            buffer_max_size, model_size
+        )
+        self._use_grad_storage = buffer_max_size > 0
+        self._grad_storages = {}  # {dtype: {rank: GradStorage}}
+        self._has_grad_storage = []
+        self._grad_storage_list = []
+
+        # Offload
+        # TODO(haohongxiang): Now it's not be supported for multi-optimizers using Offload strategy
+        self._offload_optims = list(
+            filter(lambda optim: optim.offload, self._sharding_optimizers)
+        )
+        if len(self._offload_optims) > 0:
+            assert len(self._sharding_optimizers) == 1, (
+                "Only support offload strategy for single optimizer"
+            )
+
+        self._offload = len(self._offload_optims) > 0
+        self._offload_device = "cpu"
+
+        # Set backward pass hooks
+        self._bw_hooks = []
+
+        self.scale_in_opt = False
+
+        # TODO (Baibaifan) Set tasks flow support asynchronous communicate
+        # self._tasks_flow = deque()
+
+        # Define optimizer step and clear_grad
+        self._redefine_opt_step()
+        self._redefine_opt_clear()
+
+    def forward(self, *inputs, **kwargs):
+        """
+        A wrapper for Sharding Stage2 layer.
+        - Fresh trainable params or rebuild grad storage
+        - Sync layer's buffer params
+        - Clear all flags states
+        - Forward for origin layers
+        """
+
+        # Whether to need to reset trainable parameters
+        needs_fresh = len(self._bw_hooks) == 0 and self.training
+
+        if self._auto_refresh_trainable:
+            needs_fresh |= self._detect_train_change()
+
+        # Front hook
+        self._init_internal_storage(needs_fresh)
+
+        # Sync layer's buffers state
+        if self._sync_buffers:
+            self.__sync_buffers()
+
+        # Normal FW on the base model
+        fw = self._layer(*inputs, **kwargs)
+
+        return fw
+
+    def set_state_dict(self, state_dict, use_structured_name=True):
+        self._layer.set_state_dict(
+            state_dict, use_structured_name=use_structured_name
+        )
+
+    def state_dict(
+        self,
+        destination=None,
+        include_sublayers=True,
+        structured_name_prefix="",
+    ):
+        return self._layer.state_dict(
+            destination=destination,
+            include_sublayers=include_sublayers,
+            structured_name_prefix=structured_name_prefix,
+        )
+
+    def _clear_gradients(self):
+        """
+        Set zero to the gradient of the optimizer's current rank trainable parameters.
+        """
+        # Release grad storages
+        for dtype in self._grad_storages.keys():
+            if (
+                not self._offload
+                and self._rank in self._grad_storages[dtype].keys()
+            ):
+                self._grad_storages[dtype][self._rank].buffer.zero_()
+
+        # Release grads of params
+        for param in self._trainable_params:
+            if param.name in self._param_grads:
+                if self.use_main_grad and param.main_grad is not None:
+                    param.main_grad.zero_()
+                elif param.grad is not None:
+                    param._zero_grads()
+
+        # Release grads of master params with offload strategy
+        if self._offload:
+            self._sharding_optimizers[0]._offload_clear_grad()
+
+    def _grad_scale(self):
+        """
+        this function will do 2 things:
+        1.  Before the optimization, scale main_grad to support gradient merge if param has main_grad, or to support fused_linear_param_grad_add gradient merge.
+        2.  Before the optimization, scale the gradients before allreduce of dp_group.
+        """
+
+        need_dp_scale = self._dp_group is not None and self._dp_group.nranks > 1
+        if self.scale_in_opt:
+            scale_factor = self._world_size_scaling
+        else:
+            scale_factor = 1.0
+
+        if need_dp_scale:
+            dp_scale_factor = 1.0 / (self._dp_group.nranks)
+            scale_factor = scale_factor * dp_scale_factor
+
+        # Scale grad storages
+        for dtype in self._grad_storages.keys():
+            if (
+                not self._offload
+                and self._rank in self._grad_storages[dtype].keys()
+            ):
+                self._grad_storages[dtype][self._rank].buffer.scale_(
+                    scale=scale_factor
+                )
+        # Scale grads of params
+        with paddle.no_grad():
+            for param in self._trainable_params:
+                if param.name in self._param_grads:
+                    if self.use_main_grad and param.main_grad is not None:
+                        param.main_grad.scale_(scale=scale_factor)
+                    elif param.grad is not None:
+                        param.grad.scale_(scale=scale_factor)
+
+        # Scale grads of master params with offload strategy
+        if self._offload:
+            if need_dp_scale is False:
+                return
+            self._sharding_optimizers[0]._offload_scale_grad(
+                scale=1.0 / (self._dp_group.nranks)
+            )
+
+    def _init_internal_storage(self, needs_fresh):
+        """
+        Judge Fresh trainable params or rebuild grad storage.
+        """
+        if needs_fresh:
+            self._fresh_trainable()
+        else:
+            self._build_grad_storages()
+
+        # Clear all flags state
+        self._clear_counters()
+
+    def to(self, device=None, dtype=None, blocking=True):
+        """
+        Synchronously or asynchronously convert the data type of the layer, the device is not supported now.
+        """
+        assert isinstance(device, str), "Device must be type str"
+        assert device == self._default_device, (
+            "New devices are not supported, because of the optimizer state is not sync"
+        )
+
+        self._layer.to(device=device, dtype=dtype, blocking=blocking)
+
+        # Re-build the buckets, hooks, etc..
+        self._fresh_trainable()
+
+    def _fresh_trainable(self):
+        """Whether to update training parameters."""
+
+        # Make sure that this is not done while gradients are waiting to be reduced (if no_sync context for instance)
+        if reduce(lambda x, y: x or y, self._grad_reduced, False):
+            logging.warning("Grads waiting to be reduced.")
+
+        self._trainable_params = list(
+            filter(lambda x: x.trainable, self._all_params)
+        )
+        self._trainable_params.sort(key=lambda x: x._numel())
+
+        self._trainable_param2rank = {}
+        for optim in self._sharding_optimizers:
+            # Need to be wrapped for Sharding Stage2 Optimizer
+            if len(optim.param_storages.keys()) == 0:
+                optim._update_opt_status()
+
+            # Get the parameters split by the optimizer according to rank
+            for per_rank_params in (
+                optim.dtype_rank_params.values()
+            ):  # all the params from all ranks
+                for params in per_rank_params:
+                    for param in filter(lambda x: x.trainable, params):
+                        self._trainable_param2rank[param.name] = (
+                            optim.param2rank[param.name]
+                        )
+                        self._trainable_param2align[param.name] = (
+                            optim._param2align[param.name]
+                        )
+
+        # Create grad_storage
+        self._setup_use_grad_storage()
+        # setup backward hooks
+        self._setup_backward_hooks()
+
+    @paddle.autograd.no_grad()
+    def __sync_buffers(self):
+        """
+        Sync all the param buffers from all ranks (exp: batch norm statistics).
+        """
+
+        for buffer in self._layer.buffers(include_sublayers=True):
+            dist.broadcast(
+                buffer, self._global_root_rank, self._group, sync_op=True
+            )
+
+            if self._dp_group and self._dp_group.nranks > 1:
+                dist.broadcast(
+                    buffer,
+                    self._dp_group.ranks[0],
+                    self._dp_group,
+                    sync_op=True,
+                )
+
+    def __getattr__(self, name):
+        """Forward missing attributes to wrapped layer."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self._layer, name)
+
+    @paddle.autograd.no_grad()
+    def _clear_counters(self):
+        """Reset all the grad reduce and call counters."""
+        if self.training:
+            self._grad_reduced = [True for _ in self._trainable_params]
+
+        if self._use_grad_storage:
+            for grad_storage in self._grad_storage_list:
+                grad_storage.reset_checked_in()
+
+    def _set_reduce_overlap(self, reduce_overlap):
+        # Hacky way to not add an extra parameter to the `group_sharded_parallel` funct.
+        # User should use this like:
+        # model, optimizer, scaler = group_sharded_parallel(...)
+        # model._set_reduce_overlap(True)
+        self._reduce_overlap = reduce_overlap
+        if self._reduce_overlap:
+            assert len(self._sharding_optimizers) == 1, (
+                "Only support comm overlap strategy for single optimizer"
+            )
+        self._sharding_optimizers[0]._set_reduce_overlap(reduce_overlap)
+
+    def _get_scaled_grad_fn(self, param):
+        @paddle.autograd.no_grad()
+        def scale(grad):
+            # do gradient scale separately
+            # For grad scale, we need to do it in the backward hook due to fp16 may overflow if we first add grad and then scale
+            # For main_grad scale and fused_linear_param_grad_add, we do scale in the optimizer.
+            if not self.scale_in_opt:
+                if (
+                    not hasattr(param, "main_grad")
+                    and grad is not None
+                    and grad.dtype == Type.fp16.value
+                ):
+                    assert grad._is_initialized(), (
+                        "grad should be initialized in stage2"
+                    )
+                    grad.scale_(self._world_size_scaling)
+                else:
+                    self.scale_in_opt = True
+
+        return scale
+
+    def _get_reduce_fn(self, index, param, dst_rank):
+        """
+        There are two ways to reduce gradient.
+        - 1. Do not use self._use_grad_storage or exceeded buffer_max_size will be reduced separately.
+        - 2. Use grad_storage Reduce the storage to get the full gradient from different ranks.
+        """
+
+        if not self._use_grad_storage or not self._has_grad_storage[index]:
+            # Direct reduction
+            @paddle.autograd.no_grad()
+            def reduce(*_):
+                # Skip gradient reduction, do not change status information
+                if self._grad_reduced[index]:
+                    assert (
+                        param.grad is not None or param.main_grad is not None
+                    ), "Parameter should have grad or main grad"
+
+                    # Change reduce information
+                    self._grad_reduced[index] = False
+
+                    # Clear the gradient that does not belong to the current rank through the callback function
+                    def cleanup():
+                        if dst_rank != self._rank:
+                            if self.use_main_grad:
+                                param.main_grad._clear_data()
+                                param.main_grad = None
+                            else:
+                                param.clear_gradient(False)
+                        elif self._offload:
+                            tmp_grad = param.grad.cast(
+                                dtype=Type.fp32.value
+                            ).cpu()
+
+                            self._sharding_optimizers[0]._offload_acc_grad(
+                                param.name, tmp_grad
+                            )
+                            del tmp_grad
+                            param.clear_gradient(False)
+
+                    # Synchronize the reduce parameter gradient asynchronize
+                    self._sharding_optimizers[0]._update_task(
+                        dist.reduce(
+                            tensor=(
+                                param.grad
+                                if not self.use_main_grad
+                                else param.main_grad
+                            ),
+                            dst=self._group.ranks[dst_rank],
+                            group=self._group,
+                            sync_op=not self._reduce_overlap,
+                        )
+                    )
+
+                    # Clear the task flow and trigger callback to clear the redundant gradient
+                    # self._clear_task_flow()
+
+                    cleanup()
+
+        else:
+            # Buffer reduction
+            @paddle.autograd.no_grad()
+            def reduce(*_):
+                # Skip gradient reduction, do not change status information
+                if self._grad_reduced[index]:
+                    assert (
+                        param.grad is not None or param.main_grad is not None
+                    ), "Parameter should have grad or main grad"
+
+                    # Change reduce information
+                    self._grad_reduced[index] = False
+                    grad_storage = self._grad_storages[param.dtype][dst_rank]
+                    grad_storage.params_checked_in += 1
+
+                    if grad_storage.all_checked_in:
+                        assert grad_storage.buffer is not None
+
+                        # Clearing up the grad_storage buffer
+                        def cleanup():
+                            if dst_rank != self._rank:
+                                for p in grad_storage._params:
+                                    if self.use_main_grad:
+                                        p.main_grad._clear_data()
+                                        p.main_grad = None
+                                    else:
+                                        p.clear_gradient(False)
+
+                                grad_storage.buffer._clear_data()
+                            elif self._offload:
+                                grad_storage.to(device=self._offload_device)
+                                for p in grad_storage._params:
+                                    with device_guard():
+                                        tmp_grad = p.grad.cast(
+                                            dtype=Type.fp32.value
+                                        )
+                                    self._sharding_optimizers[
+                                        0
+                                    ]._offload_acc_grad(p.name, tmp_grad)
+                                    p.clear_gradient(False)
+                                grad_storage._device = self._default_device
+                                grad_storage.buffer._clear_data()
+
+                        # Reduce the bucket
+                        grad_storage.sent = True
+                        # Synchronize the reduce parameter gradient asynchronize
+                        self._sharding_optimizers[0]._update_task(
+                            dist.reduce(
+                                tensor=grad_storage.buffer,
+                                dst=self._group.ranks[grad_storage.destination],
+                                group=self._group,
+                                sync_op=not self._reduce_overlap,
+                            )
+                        )
+
+                        cleanup()
+
+                    # Clear the task flow and trigger callback to clear the redundant gradient
+                    # self._clear_task_flow()
+
+        return reduce
+
+    def _setup_backward_hooks(self):
+        """
+        Set the backward hook to synchronize the gradients of all rank by reduce group ranks.
+        """
+
+        # Remove previous backward hooks
+        while len(self._bw_hooks) > 0:
+            self._bw_hooks.pop().remove()
+
+        # Go through the parameters, attach the hook
+        if not self.training:
+            return
+
+        for index, param in enumerate(self._trainable_params):
+            param._register_grad_hook(self._get_scaled_grad_fn(param))
+
+            dst_rank = self._trainable_param2rank[param.name]
+
+            reduce_function = self._get_reduce_fn(index, param, dst_rank)
+
+            self._bw_hooks.append(
+                param._register_backward_hook(reduce_function)
+            )
+
+    def _setup_use_grad_storage(self):
+        """
+        Integrate the parameters gradient into a continuous memory according to rank, and support the update of training parameters.
+        """
+
+        # According to parameters's numel sort, allocate memory of parameter gradient to continuous memory according to rank
+        self._grad_storages = {}
+        self._has_grad_storage = [False for _ in self._trainable_params]
+
+        for index, param in enumerate(self._trainable_params):
+            dst_rank = self._trainable_param2rank[param.name]
+
+            if param.dtype not in self._grad_storages.keys():
+                self._grad_storages[param.dtype] = {}
+
+            if dst_rank not in self._grad_storages[param.dtype].keys():
+                self._grad_storages[param.dtype][dst_rank] = GradStorage(
+                    self._buffer_max_size[param.dtype],
+                    dtype=(
+                        param.dtype
+                        if not self.use_main_grad
+                        else paddle.float32
+                    ),
+                    device=self._default_device,
+                    destination=dst_rank,
+                    param2align=self._trainable_param2align,
+                )
+
+            # Criteria to decide whether this parameter is to be put in GradStorage
+            if self._grad_storages[param.dtype][dst_rank].can_add_grad_view(
+                param, self._trainable_param2align[param.name]
+            ):
+                self._grad_storages[param.dtype][dst_rank].add_grad(
+                    param, self._trainable_param2align[param.name]
+                )
+                self._has_grad_storage[index] = True
+            else:
+                self._param_grads.append(param.name)
+
+        for dtype in self._grad_storages.keys():
+            self._grad_storage_list.extend(
+                list(self._grad_storages[dtype].values())
+            )
+
+    # def _clear_task_flow(self):
+    #     """Try to consume the previous tasks."""
+    #     while len(self._tasks_flow) > 0:
+    #         task = self._tasks_flow.popleft()
+    #         task.wait()
+    #         if task.callback is not None:
+    #             task.callback()
+
+    def _detect_train_change(self):
+        # Current trainable parameters
+        trainable_mask = list(map(_trainable, self._trainable_params))
+
+        # Whether parameters trainability changed
+        trainability_changed = trainable_mask != self._trainable_mask
+
+        if trainability_changed:
+            logging.warning(
+                "Trainable params changed, because of eval/train mode or parameter freezing/unfreeze."
+            )
+            self._trainable_mask = trainable_mask
+
+        return trainability_changed
+
+    def _build_grad_storages(self):
+        """
+        Rebuild grad storages.
+        """
+        # Rebuild fp16/fp32 grad storages
+        for dtype in self._grad_storages.keys():
+            for dst_rank, grad_storage in self._grad_storages[dtype].items():
+                if self._offload or dst_rank != self._rank:
+                    grad_storage.manual_release()
+                    grad_storage.rebuild()
+
+    def _rank_buffer_size(self, buffer_max_size, model_size):
+        """
+        Generate the minimum buffer size for each rank & Display param sizes and model sizes.
+        """
+
+        # Initialize buffer size
+        rank_buffer_size = {}
+        for shard_opt in self._sharding_optimizers:
+            if shard_opt.rank_buffer_size:
+                for dtype in shard_opt.rank_buffer_size.keys():
+                    sizes = max(shard_opt.rank_buffer_size[dtype].values())
+                    rank_buffer_size[dtype] = min(sizes, buffer_max_size)
+
+        if Type.fp16.value in rank_buffer_size.keys():
+            # FP16 GradStorage and model size
+            logger_.info(
+                f"====== FP16 GradStorage size: {rank_buffer_size[Type.fp16.value] / 2**19:.2f}M parameters, Model size {model_size / 2**19:.2f}M parameters ======"
+            )
+        if Type.bf16.value in rank_buffer_size.keys():
+            # FP16 GradStorage and model size
+            logger_.info(
+                f"====== BF16 GradStorage size: {rank_buffer_size[Type.bf16.value] / 2**19:.2f}M parameters, Model size {model_size / 2**19:.2f}M parameters ======"
+            )
+        if Type.fp32.value in rank_buffer_size.keys():
+            # FP32 GradStorage and model size
+            logger_.info(
+                f"====== FP32 GradStorage size: {rank_buffer_size[Type.fp32.value] / 2**18:.2f}M parameters, Model size {model_size / 2**18:.2f}M parameters ======"
+            )
+        return rank_buffer_size
+
+    def _dp_allreduce(self):
+        # do dp allreduce here for gradient merge.
+        if self._dp_group and self._dp_group.nranks > 1:
+            for dtype in self._grad_storages.keys():
+                for rank, g in sorted(
+                    self._grad_storages[dtype].items(), key=lambda x: x[0]
+                ):
+                    if g.destination == self._rank:
+                        assert g.buffer._is_initialized()
+                        dist.all_reduce(
+                            tensor=g.buffer,
+                            group=self._dp_group,
+                            sync_op=True,
+                        )
+            for param in self._trainable_params:
+                if param.name in self._param_grads:
+                    if self.use_main_grad and param.main_grad is None:
+                        continue
+                    elif param.grad is None:
+                        continue
+                    dst_rank = self._trainable_param2rank[param.name]
+                    if dst_rank == self._rank:
+                        dist.all_reduce(
+                            tensor=(
+                                param.grad
+                                if not self.use_main_grad
+                                else param.main_grad
+                            ),
+                            group=self._dp_group,
+                            sync_op=True,
+                        )
+
+    def _redefine_opt_step(self):
+        grad_func = self._grad_scale
+        dp_allreduce_func = self._dp_allreduce
+
+        for opt in self._sharding_optimizers:
+            opt_step = opt.step
+
+            def _opt_step(self):
+                if self._reduce_overlap:
+                    # Wait for the last reduce task. This wait must before grad scale function.
+                    assert self._comm_task is not None
+                    self._comm_task.wait()
+
+                grad_func()
+                dp_allreduce_func()
+                opt_step()
+
+            opt.step = MethodType(_opt_step, opt)
+
+    def _redefine_opt_clear(self):
+        clear_func = self._clear_gradients
+
+        def _opt_clear(self):
+            clear_func()
+
+        for opt in self._sharding_optimizers:
+            opt.clear_grad = MethodType(_opt_clear, opt)

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_storage.py

@@ -0,0 +1,363 @@
+#   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.
+
+# The file has been adapted from fairscale file:
+# https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/misc/param_bucket.py
+# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
+# We retain the following license from the original files:
+
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+
+from .group_sharded_utils import Type, cvt_to_device, device_guard
+
+
+class BufferWarper(core.eager.Tensor):
+    def __init__(self):
+        super().__init__()
+        self.need_clip = True
+        self.is_distributed = False
+        self.trainable = True
+
+
+class InternalStorage:
+    """
+    This is a basic class, which is responsible for consolidating the basic storage tensor.
+
+    """
+
+    # Support integration parameter tensor
+    def __init__(self, size, dtype, device, convert_cpu=False):
+        self._params = []
+        self._param_ids = []
+        self._fill = 0
+        self._device = device
+        self._dtype = dtype
+
+        # The flatten tensor
+        size = [size] if isinstance(size, int) else size
+        if convert_cpu:
+            value = (
+                np.zeros(size, dtype=np.float16)
+                if Type.fp16.value == dtype
+                else np.zeros(size, dtype=np.float32)
+            )
+            self.buffer = core.eager.Tensor(value=value, place=core.CPUPlace())
+            if dtype == Type.bf16.value:
+                self.buffer = paddle.cast(self.buffer, dtype=paddle.bfloat16)
+        else:
+            self.buffer = paddle.zeros(size, dtype=dtype)
+
+        self.dev_id = (
+            0
+            if paddle.get_device() == "cpu"
+            else int(paddle.get_device().split(":")[1])
+        )
+
+    def to(self, device, dtype=None, keep_alignment=True):
+        """
+        Move the underlying buffer
+        """
+        assert self.buffer is not None, (
+            "Cannot move a collapsed bucket, please rebuild it"
+        )
+        assert dtype == Type.fp32.value or Type.fp16.value, (
+            "Conversion type is not supported now"
+        )
+
+        if self._device != device:
+            if device in paddle.device.get_all_custom_device_type():
+                tmp_buffer = self.buffer._copy_to(
+                    paddle.CustomPlace(device, self.dev_id), True
+                )
+            else:
+                tmp_buffer = (
+                    cvt_to_device(self.buffer, self.dev_id)
+                    if device in ["gpu", "xpu"]
+                    else self.buffer.cpu()
+                )
+            for param in self._params:
+                param.clear_gradient(False)
+
+            del self.buffer
+            self.buffer = tmp_buffer
+            self._device = device
+
+        if dtype is not None:
+            self.buffer = self.buffer.cast(dtype=dtype)
+            self._dtype = dtype
+
+    def warp_buffer(self):
+        tmp_buffer = BufferWarper()
+        self._buffer = self.buffer
+        tmp_buffer.get_tensor()._share_data_with(self.buffer.get_tensor())
+        self.buffer = tmp_buffer
+
+
+class ParamStorage(InternalStorage):
+    """
+    This is a basic class to simplify the handling of parameter InternalStorages.
+    """
+
+    def __init__(self, size, dtype, device):
+        super().__init__(size, dtype, device, convert_cpu=True)
+        self.param2align = None
+
+    def to(self, device, dtype=None, keep_alignment=True):
+        """
+        Move the underlying buffer
+        """
+
+        super().to(device, dtype)
+
+        if keep_alignment:
+            self._array_params()
+
+    @paddle.autograd.no_grad()
+    def add_rank_params(self, trainable_params, param2align, convert_gpu=True):
+        """
+        Add new parameters to the InternalStorage. Params becomes a view of this InternalStorage buffer.
+        """
+
+        assert all(
+            id(param) not in self._param_ids for param in trainable_params
+        ), "The same param cannot be checked in twice"
+        assert self.buffer is not None
+
+        self.param2align = param2align
+
+        cpu_param_shape = []
+        for param in trainable_params:
+            p_shape = self._add_param_as_view(
+                param, param2align[param.name], convert_gpu
+            )
+            cpu_param_shape.append(p_shape)
+
+        if convert_gpu:
+            if self._device in paddle.device.get_all_custom_device_type():
+                self.buffer = self.buffer._copy_to(
+                    paddle.CustomPlace(self._device, self.dev_id), True
+                )
+            else:
+                # buffer convert from cpu to cuda
+                self.buffer = cvt_to_device(self.buffer, self.dev_id)
+
+        self._fill = 0
+
+        for idx, param in enumerate(trainable_params):
+            self._convert_buffer(
+                param, cpu_param_shape[idx], param2align[param.name]
+            )
+            self._params.append(param)
+            self._param_ids.append(id(param))
+
+    @paddle.autograd.no_grad()
+    def _add_param_as_view(self, param, align, convert_gpu=True):
+        assert param.dtype == self.buffer.dtype, (
+            f"Different types for the InternalStorage and the param, cannot proceed: {param.dtype} - {self.buffer.dtype}"
+        )
+
+        var_end = self._fill + param._numel()
+        offset = var_end + align
+        assert offset <= self.buffer._numel()
+
+        p_shape = param.shape
+
+        origin_state = param.stop_gradient
+        param.stop_gradient = True
+        param.flatten_()
+        param.stop_gradient = origin_state
+
+        # Copy the current param value
+
+        with device_guard(self.dev_id, "cpu"):
+            tmp_var = self.buffer._slice(self._fill, var_end)
+            if convert_gpu:
+                param_cpu = param.cpu()
+                param._clear_data()
+                tmp_var.set_value(param_cpu)
+            else:
+                tmp_var.set_value(param)
+            del tmp_var
+
+        self._fill = offset
+        return p_shape
+
+    @paddle.autograd.no_grad()
+    def _convert_buffer(self, param, p_shape, align):
+        var_end = self._fill + np.prod(p_shape).tolist()
+        offset = var_end + align
+        assert offset <= self.buffer._numel()
+
+        # Convert the param value
+        with device_guard(self.dev_id, self._device):
+            tmp_tensor = self.buffer._slice(self._fill, var_end)
+            tmp_tensor._share_buffer_to(param)
+            param.get_tensor()._set_dims(p_shape)
+
+        self._fill = offset
+
+    @paddle.autograd.no_grad()
+    def _array_params(self):
+        """
+        Given the parameters which have been registered previously, rebuild the whole InternalStorage.
+        """
+        assert len(self._params) > 0
+        assert self.param2align is not None
+
+        self._fill = 0
+        for p in self._params:
+            self._convert_buffer(p, p.shape, self.param2align[p.name])  # modify
+
+
+class GradStorage(InternalStorage):
+    """
+    This is a basic class to simplify the handling of gradient InternalStorages
+    """
+
+    def __init__(
+        self, size, dtype, device, destination, param2align, convert_cpu=False
+    ):
+        if isinstance(size, np.int64):
+            size = size.tolist()
+        super().__init__(size, dtype, device, convert_cpu)
+
+        self._max_size = size
+        self._release = False
+
+        self.params_checked_in = 0
+        self.destination = destination
+        self._param2align = param2align
+        self.sent = False
+
+    def reset_checked_in(self):
+        """Reset the counter of the parameter grads which have been checked in"""
+        self.params_checked_in = 0
+        self.sent = False
+
+    @property
+    def all_checked_in(self):
+        """Judge all the expected gradient check-in happened"""
+        return len(self._params) == self.params_checked_in
+
+    def can_add_grad_view(self, param, align):
+        """Is there enough InternalStorage to add this parameter gradient, and whether this param have already checked in."""
+        return (
+            self._fill + param._numel() + align <= self._max_size
+            and id(param) not in self._param_ids
+        )
+
+    def to(self, device, dtype=None, keep_alignment=True):
+        """
+        Move the underlying buffer
+        """
+        if self._release:
+            self.rebuild()
+
+        super().to(device, dtype)
+
+        if keep_alignment:
+            self._array_grads()
+
+    @paddle.autograd.no_grad()
+    def add_grad(self, param, align):
+        """
+        Add a new parameter gradient to the InternalStorage. Param.grad becomes a view of this InternalStorage buffer.
+        """
+
+        assert id(param) not in self._param_ids, (
+            "The same gradients cannot be checked in twice"
+        )
+
+        self._add_grad_as_view(param, align)
+        self._params.append(param)
+        self._param_ids.append(id(param))
+
+    @paddle.autograd.no_grad()
+    def manual_release(self):
+        """
+        Release the buffer from InternalStorage. The InternalStorage will need to be rebuilt before use.
+        """
+        if not self._release:
+            for p in self._params:
+                use_main_grad = hasattr(p, "main_grad")
+                if use_main_grad and p.main_grad is not None:
+                    p.main_grad._clear_data()
+                    p.main_grad = None
+                elif p.grad is not None:
+                    p.clear_gradient(False)
+
+            self.buffer = None
+            self._fill = 0
+            self.params_checked_in = 0
+            self._release = True
+
+    @paddle.autograd.no_grad()
+    def rebuild(self):
+        """
+        Given the parameter gradients which have been registered previously, rebuild the whole InternalStorage.
+        """
+
+        if self._release:
+            self.buffer = paddle.zeros([self._max_size], dtype=self._dtype)
+
+            for p in self._params:
+                self._add_grad_as_view(p, self._param2align[p.name])
+
+            self._release = False
+
+    @paddle.autograd.no_grad()
+    def _array_grads(self):
+        """
+        Given the parameters gradients which have been registered previously, rebuild the whole InternalStorage.
+        """
+        if len(self._params) > 0:
+            self._fill = 0
+            for p in self._params:
+                self._add_grad_as_view(p, self._param2align[p.name])
+
+    @paddle.autograd.no_grad()
+    def _add_grad_as_view(self, param, align):
+        assert param._numel() > 0, (
+            "Cannot add a gradient to a released InternalStorage, please rebuild"
+        )
+
+        use_main_grad = hasattr(param, "main_grad")
+        if use_main_grad:
+            assert self.buffer.dtype == paddle.float32
+        else:
+            assert param.dtype == self.buffer.dtype
+
+        grad_end = self._fill + param._numel()
+        offset = grad_end + align
+        assert offset <= self.buffer._numel()
+
+        # Copy the current grad value to InternalStorage
+        with device_guard(self.dev_id, self._device):
+            tmp_var = self.buffer._slice(self._fill, grad_end)
+            tmp_var.get_tensor()._set_dims(param.shape)
+            if not use_main_grad:
+                param._copy_gradient_from(tmp_var)
+            else:
+                param.main_grad = tmp_var
+            del tmp_var
+
+        self._fill = offset

python/paddle/distributed/fleet/meta_parallel/sharding/group_sharded_utils.py

@@ -0,0 +1,352 @@
+#   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 contextlib
+from enum import Enum
+from types import MethodType
+
+import numpy as np
+
+import paddle
+from paddle import _C_ops, _legacy_C_ops
+from paddle.base import core
+from paddle.common_ops_import import dygraph_only
+from paddle.nn import clip
+
+
+class Taskflow:
+    """
+    Task flows, one way linked list for task acquisition.
+    """
+
+    def __init__(self, task, callback):
+        self.task = task
+        self.callback = callback
+
+
+class Type(Enum):
+    """
+    Type of trainable parameters
+    """
+
+    fp16 = paddle.float16
+    bf16 = paddle.bfloat16
+    fp32 = paddle.float32
+
+
+class GroupShardedClipGrad:
+    def __init__(self, clip, device, group):
+        self._clip = clip
+        self._device = device
+        self._group = group
+
+    @paddle.autograd.no_grad()
+    def _dygraph_clip(self, params_grads):
+        sum_square_fp32, sum_square_fp16, sum_square_bfp16 = [], [], []
+        unslice_params_fp32, unslice_params_fp16, unslice_params_bfp16 = (
+            [],
+            [],
+            [],
+        )
+
+        for p, g in params_grads:
+            p_slice = True  # using for slice parameter in sharding stage3
+            if g is None or getattr(p, 'need_clip', True) is False:
+                continue
+            if hasattr(p, "unslice"):
+                p_slice = False
+
+            merge_grad = g
+            if g.type == core.VarDesc.VarType.SELECTED_ROWS:
+                merge_grad = clip.get_tensor_from_selected_rows(
+                    clip.merge_selected_rows(g)
+                )
+            square = paddle.square(merge_grad)
+            sum_square = paddle.sum(square)
+
+            if p.dtype == paddle.float16:
+                if p_slice:
+                    sum_square_fp16.append(sum_square)
+                else:
+                    unslice_params_fp16.append(sum_square)
+            elif p.dtype == paddle.float32:
+                if p_slice:
+                    sum_square_fp32.append(sum_square)
+                else:
+                    unslice_params_fp32.append(sum_square)
+            elif p.dtype == paddle.bfloat16:
+                if p_slice:
+                    sum_square_bfp16.append(sum_square)
+                else:
+                    unslice_params_bfp16.append(sum_square)
+
+        # global norm of non-distributed FP16 params_and_grads
+        if len(sum_square_fp16) == 0:
+            global_norm_fp16 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_norm_fp16 = paddle.add_n(sum_square_fp16)
+            global_norm_fp16 = paddle.cast(
+                global_norm_fp16, dtype=paddle.float32
+            )
+
+        # global norm of non-distributed BFP16 params_and_grads
+        if len(sum_square_bfp16) == 0:
+            global_norm_bfp16 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_norm_bfp16 = paddle.add_n(sum_square_bfp16)
+            global_norm_bfp16 = paddle.cast(
+                global_norm_bfp16, dtype=paddle.float32
+            )
+
+        # global norm of non-distributed FP16 params_and_grads for unslice parameters
+        if len(unslice_params_fp16) == 0:
+            global_unslice_fp16 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_unslice_fp16 = paddle.add_n(unslice_params_fp16)
+            global_unslice_fp16 = paddle.cast(
+                global_unslice_fp16, dtype=paddle.float32
+            )
+
+        # global norm of non-distributed BFP16 params_and_grads for unslice parameters
+        if len(unslice_params_bfp16) == 0:
+            global_unslice_bfp16 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_unslice_bfp16 = paddle.add_n(unslice_params_bfp16)
+            global_unslice_bfp16 = paddle.cast(
+                global_unslice_bfp16, dtype=paddle.float32
+            )
+
+        # global norm of non-distributed FP32 params_and_grads
+        if len(sum_square_fp32) == 0:
+            global_norm_fp32 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_norm_fp32 = paddle.add_n(sum_square_fp32)
+
+        # global norm of non-distributed FP32 params_and_grads for unslice parameters
+        if len(unslice_params_fp32) == 0:
+            global_unslice_fp32 = paddle.to_tensor(
+                np.array(0.0), dtype=paddle.float32
+            )
+        else:
+            global_unslice_fp32 = paddle.add_n(unslice_params_fp32)
+
+        global_unslice_var = (
+            global_unslice_fp16 + global_unslice_fp32 + global_unslice_bfp16
+        )
+
+        global_norm_var = (
+            global_norm_fp16 + global_norm_fp32 + global_norm_bfp16
+        )
+
+        # add all reduce to get global norm of distributed params_and_grads
+        dev_id = int(self._device.split(":")[1])
+        dev_type = self._device.split(':')[0]
+        if paddle.device.get_device() == "cpu":
+            if dev_type in paddle.device.get_all_custom_device_type():
+                global_norm_var = global_norm_var._copy_to(
+                    paddle.CustomPlace(dev_type, dev_id), True
+                )
+            elif dev_type == "xpu":
+                global_norm_var = global_norm_var.to(self._device)
+            else:
+                global_norm_var = global_norm_var.cuda(dev_id)
+
+        with device_guard(dev_id, self._device.split(":")[0]):
+            paddle.distributed.all_reduce(global_norm_var, group=self._group)
+
+        global_norm_var = paddle.sqrt(global_norm_var + global_unslice_var)
+        max_global_norm = paddle.full(
+            shape=[], dtype=global_norm_var.dtype, fill_value=self.clip_norm
+        )
+
+        clip_var = paddle.divide(
+            x=max_global_norm,
+            y=paddle.maximum(x=global_norm_var, y=max_global_norm),
+        )
+        clip_var_fp16 = paddle.cast(clip_var, paddle.float16)
+
+        for p, g in params_grads:
+            if getattr(p, 'need_clip', True) is False or g is None:
+                continue
+            origin_state = g.stop_gradient
+            g.stop_gradient = True
+            if p.dtype == paddle.float16:
+                g.scale_(clip_var_fp16)
+            else:
+                g.scale_(clip_var)
+            g.stop_gradient = origin_state
+            # p._reset_grad_inplace_version(True)
+
+        return params_grads
+
+    def __getattr__(self, item):
+        return getattr(self._clip, item)
+
+    def __call__(self, params_grads):
+        return self._dygraph_clip(params_grads)
+
+
+@contextlib.contextmanager
+def device_guard(dev_id=0, device="cpu"):
+    origin_device = paddle.device.get_device()
+    if device == "cpu":
+        paddle.set_device(device)
+    elif device in ["gpu", "xpu"]:
+        paddle.set_device(f"{device}:{dev_id}")
+    elif device in paddle.device.get_all_custom_device_type():
+        paddle.set_device(f"{device}:{dev_id}")
+
+    try:
+        yield
+    finally:
+        paddle.set_device(origin_device)
+
+
+@dygraph_only
+def GroupShardedScaler(scaler):
+    def unscale_method(self, optimizer):
+        if not self._enable:
+            return
+        param_grads = []
+        param_grads_bfp16 = []
+        param_grads_fp16 = []
+        param_grads_fp32 = []
+        if hasattr(optimizer, "update_slice"):
+            optimizer.update_slice()
+            optimizer.update_scaler = True
+
+        if getattr(optimizer._optim, '_param_groups', None) and isinstance(
+            optimizer._optim._param_groups[0], dict
+        ):
+            for group in optimizer._optim._param_groups:
+                for param in group['params']:
+                    tgt_grad = None
+                    if (
+                        hasattr(param, "main_grad")
+                        and param.main_grad is not None
+                    ):
+                        tgt_grad = param.main_grad
+                    elif param.grad is not None:
+                        tgt_grad = param.grad
+                    if tgt_grad is not None:
+                        param_grads.append(tgt_grad)
+                        if tgt_grad.dtype in [
+                            core.VarDesc.VarType.FP16,
+                            paddle.float16,
+                        ]:
+                            param_grads_fp16.append(tgt_grad)
+                        elif tgt_grad.dtype in [paddle.bfloat16]:
+                            param_grads_bfp16.append(tgt_grad)
+                        else:
+                            param_grads_fp32.append(tgt_grad)
+        else:
+            for param in optimizer._optim._parameter_list:
+                tgt_grad = None
+                if hasattr(param, "main_grad") and param.main_grad is not None:
+                    tgt_grad = param.main_grad
+                elif param.grad is not None:
+                    tgt_grad = param.grad
+                if tgt_grad is not None:
+                    param_grads.append(tgt_grad)
+                    if tgt_grad.dtype in [
+                        core.VarDesc.VarType.FP16,
+                        paddle.float16,
+                    ]:
+                        param_grads_fp16.append(tgt_grad)
+                    elif tgt_grad.dtype in [paddle.bfloat16]:
+                        param_grads_bfp16.append(tgt_grad)
+                    else:
+                        param_grads_fp32.append(tgt_grad)
+
+        temp_found_inf_fp16 = paddle.to_tensor(np.array([0]).astype(np.bool_))
+        temp_found_inf_bfp16 = paddle.to_tensor(np.array([0]).astype(np.bool_))
+        temp_found_inf_fp32 = paddle.to_tensor(np.array([0]).astype(np.bool_))
+
+        device = paddle.get_device().split(":")[0]
+        device = "cpu" if optimizer.offload else device
+        dev_id = (
+            0 if device == "cpu" else int(paddle.get_device().split(":")[1])
+        )
+
+        self._found_inf = self._temp_found_inf_value_false
+        with device_guard(dev_id, device):
+            if len(param_grads_bfp16):
+                _legacy_C_ops.check_finite_and_unscale(
+                    param_grads_bfp16,
+                    self._scale,
+                    param_grads_bfp16,
+                    temp_found_inf_bfp16,
+                )
+                self._found_inf = _C_ops.bitwise_or(
+                    self._found_inf, temp_found_inf_bfp16
+                )
+            if len(param_grads_fp16):
+                _legacy_C_ops.check_finite_and_unscale(
+                    param_grads_fp16,
+                    self._scale,
+                    param_grads_fp16,
+                    temp_found_inf_fp16,
+                )
+                self._found_inf = _C_ops.bitwise_or(
+                    self._found_inf, temp_found_inf_fp16
+                )
+            if len(param_grads_fp32):
+                _legacy_C_ops.check_finite_and_unscale(
+                    param_grads_fp32,
+                    self._scale,
+                    param_grads_fp32,
+                    temp_found_inf_fp32,
+                )
+                self._found_inf = _C_ops.bitwise_or(
+                    self._found_inf, temp_found_inf_fp32
+                )
+
+        self._found_inf = self._found_inf.cast("int32")
+
+        paddle.distributed.all_reduce(
+            self._found_inf, op=paddle.distributed.ReduceOp.MAX, group=None
+        )
+
+        self._found_inf = self._found_inf.cast("bool")
+
+    scaler._unscale = MethodType(unscale_method, scaler)
+    return scaler
+
+
+def cvt_to_device(x, dev_id, blocking=True):
+    """
+    Copy data in x from cpu memory to supported device
+    """
+    if paddle.is_compiled_with_cuda():
+        place = paddle.CUDAPlace(dev_id)
+    elif paddle.is_compiled_with_xpu():
+        place = paddle.XPUPlace(dev_id)
+    else:
+        supported_custom_devices = ["npu"]
+        place = paddle.framework._current_expected_place()
+        if place.get_device_type() not in supported_custom_devices:
+            raise OSError(
+                "Only supported compiled paddle with gpu/rocm and xpu, but current version is compiled with cpu."
+            )
+    return x._copy_to(place, blocking)

python/paddle/distributed/fleet/meta_parallel/sharding_parallel.py

@@ -0,0 +1,37 @@
+#   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.
+
+from ..utils.hybrid_parallel_util import (
+    broadcast_dp_parameters,
+    broadcast_sharding_parameters,
+)
+from ..utils.log_util import logger
+from .meta_parallel_base import MetaParallelBase
+
+__all__ = []
+
+
+class ShardingParallel(MetaParallelBase):
+    def __init__(self, layers, hcg, **kwargs):
+        super().__init__(layers, hcg, **kwargs)
+
+    def _prepare_for_model(self):
+        logger.info("start broadcast sharding parameters")
+        broadcast_sharding_parameters(self._layers, self._hcg)
+
+        if self._hcg.get_data_parallel_world_size() > 1:
+            logger.info("start broadcast dp parameters")
+            broadcast_dp_parameters(self._layers, self._hcg)
+
+        logger.info("sharding's parameters is ready")

python/paddle/distributed/fleet/meta_parallel/tensor_parallel.py

@@ -0,0 +1,66 @@
+#   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.
+
+from ..utils.hybrid_parallel_util import (
+    broadcast_dp_parameters,
+    broadcast_input_data,
+    broadcast_moe_sharding_parameters,
+    broadcast_mp_parameters,
+    broadcast_sep_parameters,
+    broadcast_sharding_parameters,
+)
+from ..utils.log_util import logger
+from .meta_parallel_base import MetaParallelBase
+
+__all__ = []
+
+
+class TensorParallel(MetaParallelBase):
+    def __init__(self, layers, hcg, **kwargs):
+        super().__init__(layers, hcg, **kwargs)
+
+    def _prepare_for_model(self):
+        logger.info("start broadcast mp parameters")
+        broadcast_mp_parameters(self._layers, self._hcg)
+
+        if self._hcg.get_sep_parallel_world_size() > 1:
+            logger.info("start broadcast sep parameters")
+            broadcast_sep_parameters(self._layers, self._hcg, fuse_params=False)
+
+        if self._hcg.get_sharding_parallel_world_size() > 1:
+            logger.info("start broadcast sharding parameters")
+            broadcast_sharding_parameters(
+                self._layers, self._hcg, fuse_params=False
+            )
+
+        if self._hcg.get_data_parallel_world_size() > 1:
+            logger.info("start broadcast dp parameters")
+            broadcast_dp_parameters(self._layers, self._hcg, fuse_params=False)
+
+        if self._hcg.get_moe_sharding_parallel_world_size() > 1:
+            logger.info("start broadcast moe sharding parameters")
+            broadcast_moe_sharding_parameters(
+                self._layers, self._hcg, fuse_params=False
+            )
+
+        logger.info("mp's parameters is ready")
+
+    def _pre_forward(self, *inputs, **kwargs):
+        need_broadcast_data = True
+        if self._strategy is not None:
+            mp_configs = self._strategy.hybrid_configs["mp_configs"]
+            need_broadcast_data = mp_configs.need_broadcast_data
+        if need_broadcast_data:
+            logger.debug("mp start broadcast input data")
+            return broadcast_input_data(self._hcg, *inputs, **kwargs)

python/paddle/distributed/fleet/meta_parallel/zero_bubble_utils.py

@@ -0,0 +1,133 @@
+# 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.
+
+# The file has been adapted from DeepSeek DualPipe project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DualPipe/blob/main/LICENSE
+
+
+import queue
+from functools import partial
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.autograd import PyLayer
+
+
+class WeightGradStore:
+    enabled = False
+    cache = []
+    funcs_queue = queue.Queue()
+
+    @classmethod
+    def put(cls, func) -> None:
+        cls.cache.append(func)
+
+    @classmethod
+    def flush(cls) -> None:
+        cls.funcs_queue.put(cls.cache)
+        cls.cache = []
+
+    @classmethod
+    def pop(cls) -> None:
+        assert not cls.funcs_queue.empty(), "Pop empty queue."
+        funcs = cls.funcs_queue.get()
+        for func in funcs:
+            func()
+
+    @classmethod
+    def clear(cls) -> None:
+        cls.cache = []
+        cls.funcs_queue = queue.Queue()
+
+
+class EventStore:
+    event = None
+
+    @classmethod
+    def set(cls, event) -> None:
+        cls.event = event
+
+
+def fold_init_dims(tensor):
+    # NOTE(zhangyuqin1998): Reshape a rank-3 tensor from P x M x N to (P * M) x N,
+    # to keep weight_grad in a correct rank. See phi::FoldInitDims.
+    if tensor.ndim == 3:
+        tensor = paddle.reshape(tensor, [-1, tensor.shape[-1]])
+    return tensor
+
+
+def grad_weight_fn(input, weight, out_grad, inplace_update_grad=True):
+    if weight.stop_gradient:
+        return
+    with paddle.no_grad():
+        weight_grad = paddle.matmul(
+            x=fold_init_dims(input),
+            y=fold_init_dims(out_grad),
+            transpose_x=True,
+            transpose_y=False,
+        )
+
+        if hasattr(weight, "main_grad"):
+            if weight.main_grad is None:
+                weight.main_grad = paddle.base.framework.core.eager.Tensor(
+                    value=weight_grad.cast(paddle.float32).value(),
+                    place=weight_grad.place,
+                    name="main_grad@" + weight.name,
+                )
+            else:
+                weight.main_grad.add_(weight_grad)
+            weight_grad._clear_data()
+        else:
+            if weight.grad is None:
+                weight.grad = paddle.zeros_like(weight, dtype=weight.dtype)
+            weight.grad = paddle.add(weight.grad, weight_grad)
+
+
+class SplitBWMatmul(PyLayer):
+    @staticmethod
+    def forward(ctx, input, weight, bias):
+        ctx.save_for_backward(input, weight, bias)
+        out = F.linear(x=input, weight=weight, bias=bias)
+        return out
+
+    @staticmethod
+    def backward(ctx, out_grad):
+        input, weight, bias = ctx.saved_tensor()
+
+        if WeightGradStore.enabled:
+            WeightGradStore.put(
+                partial(grad_weight_fn, input, weight, out_grad)
+            )
+        else:
+            grad_weight_fn(input, weight, out_grad)
+
+        input_grad = None
+        if not input.stop_gradient:
+            input_grad = paddle.matmul(
+                x=out_grad, y=weight, transpose_x=False, transpose_y=True
+            )
+        if bias is not None:
+            bias_grad = None
+            if not bias.stop_gradient:
+                bias_grad = paddle.sum(fold_init_dims(out_grad), axis=0)
+            return input_grad, None, bias_grad
+        else:
+            return input_grad, None
+
+
+class SplitBWLinear(nn.Linear):
+    def forward(self, input):
+        return SplitBWMatmul.apply(input, self.weight, bias=self.bias)