chore: import upstream snapshot with attribution

python/paddle/distributed/auto_parallel/pipelining/__init__.py

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

python/paddle/distributed/auto_parallel/pipelining/_backward.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+import paddle
+
+from .utils import _map_debug_info
+
+logger = logging.getLogger(__name__)
+
+
+def stage_backward_input(
+    stage_outputs_or_loss: list[paddle.Tensor],
+    output_grads: list[paddle.Tensor] | None,
+    input_values: list[paddle.Tensor],
+    weights: Iterator[paddle.Tensor],
+) -> tuple[tuple[paddle.Tensor | None, ...], list[dict[str, Any]]]:
+    raise NotImplementedError("stage_backward_input is not implemented yet")
+
+
+def stage_backward_weight(
+    weights: Iterator[paddle.Tensor],
+    param_groups: list[dict[str, Any]],
+    retain_graph=False,
+) -> tuple[paddle.Tensor | None, ...]:
+    raise NotImplementedError("stage_backward_weight is not implemented yet")
+
+
+def stage_backward(
+    stage_output,
+    output_grads,
+    input_values,
+) -> tuple[paddle.Tensor | None, ...]:
+    """
+    This is a helper function to:
+    1. compute the gradients for the stage inputs, and
+    2. accumulate gradients for the stage module's parameters.
+
+    Given the input value(s) and the corresponding gradient for the output
+    value(s), compute and accumulate gradients for all parameter values (leaves
+    in the autograd trace) as well as return a list of the gradients for the
+    input values
+
+    """
+
+    try:
+        # stage_output may be a composite datatype like dict. Extract all individual
+        # tensor values here
+        stage_output_tensors: list[paddle.Tensor] = []
+        output_grad_tensors: list[paddle.Tensor | None] = []
+
+        def extract_tensors_with_grads(
+            output_val,
+            grad_val,
+            extract_tensors_with_grads,
+        ):
+            if isinstance(output_val, paddle.Tensor):
+                if output_val.stop_gradient and output_val.grad_fn is None:
+                    return
+                assert isinstance(grad_val, (paddle.Tensor, type(None))), (
+                    f"Expected Tensor or None gradient but got {type(grad_val)}"
+                )
+                stage_output_tensors.append(output_val)
+                output_grad_tensors.append(grad_val)
+            elif isinstance(output_val, (tuple, list)):
+                if grad_val is None:
+                    return
+                assert isinstance(grad_val, (tuple, list)), (
+                    f"grad_value expected to have type {type(output_val)} but got {type(grad_val)}"
+                )
+                assert len(output_val) == len(grad_val)
+                for ov, gv in zip(output_val, grad_val):
+                    extract_tensors_with_grads(
+                        ov,
+                        gv,
+                        extract_tensors_with_grads,
+                    )
+            elif isinstance(output_val, dict):
+                if grad_val is None:
+                    return
+                assert isinstance(grad_val, dict)
+                assert set(output_val.keys()) == set(grad_val.keys())
+                for k in output_val.keys():
+                    extract_tensors_with_grads(
+                        output_val[k], grad_val[k], extract_tensors_with_grads
+                    )
+            else:
+                # Output is a non-tensor type; just ignore it
+                pass
+
+        # Note: ref cycle
+        # break a ref cycle that would keep tensors alive until GC runs
+        # 1. extract_tensors_with_grads refers to a cell that holds refs to any vars defined in stage_backward
+        #    and used in extract_tensors_with_grads
+        # 2. extract_tensors_with_grads referred to both stage_output_tensors, output_grad_tensors,
+        #    and to itself (extract_tensors_with_grads) since it makes a recursive call
+        # 3. stage_output_tensors was kept alive by the above refcycle, and it holds activation tensors, which is bad
+        # fix -> explicitly pass in the ref to the fn, so there is no gc cycle anymore
+        extract_tensors_with_grads(
+            stage_output, output_grads, extract_tensors_with_grads
+        )
+        # Deactivate auto mixed precision context in the backward phase
+        with paddle.amp.auto_cast(enable=False):
+            paddle.autograd.backward(
+                stage_output_tensors,
+                grad_tensors=output_grad_tensors,
+            )
+
+        # Extract gradients wrt the input values
+        grad_inputs: list[paddle.Tensor | None] = []
+        for val in input_values:
+            if isinstance(val, paddle.Tensor):
+                grad_inputs.append(val.grad)
+            else:
+                grad_inputs.append(None)
+
+    except Exception as e:
+        exc_msg = f"""
+        Failed to run stage backward:
+        Stage output: {_map_debug_info(stage_output)}
+        Output gradient: {_map_debug_info(output_grads)}
+        Input: {_map_debug_info(input_values)}
+        """
+        raise RuntimeError(exc_msg) from e
+
+    return tuple(grad_inputs)

python/paddle/distributed/auto_parallel/pipelining/microbatch.py

@@ -0,0 +1,326 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import Replicate, Shard
+from paddle.distributed.auto_parallel.api import (
+    dtensor_from_local,
+    dtensor_to_local,
+)
+from paddle.utils import flatten, map_structure, pack_sequence_as
+
+logger = logging.getLogger(__name__)
+
+# Default chunking dimension is 0. This is used for the case where the user did
+# not specify a chunking dimension.
+DEFAULT_CHUNK_DIM = 0
+
+
+def _split_tensor(x, num_chunks, split_axis=0):
+    if not x.is_dist():
+        chunk_tensors = paddle.tensor_split(x, num_chunks, split_axis)
+    # dp_degree > 1 , placements of model input is [S(0), R, ...]
+    else:
+        if dist.in_auto_parallel_align_mode():
+
+            def _reorder_data_for_align():
+                nonlocal x
+                assert x.placements[0] == dist.Shard(0), (
+                    "inputs should be placed on S(0)."
+                )
+
+                shardings = x.process_mesh.shape[0]
+
+                rows_per_shard = x.shape[0] // shardings
+                new_indices = []
+                for s_id in range(shardings):
+                    for row_in_shard in range(rows_per_shard):
+                        new_indices.append(s_id + row_in_shard * shardings)
+                tmp = x[new_indices]
+                x = dist.reshard(tmp, x.process_mesh, x.placements)
+
+            _reorder_data_for_align()
+        mesh = x.process_mesh
+        placements = x.placements
+        dense_x = dtensor_to_local(x, mesh, placements)
+        chunk_tensors = paddle.tensor_split(dense_x, num_chunks, split_axis)
+        for i in range(num_chunks):
+            chunk_tensors[i] = dtensor_from_local(
+                chunk_tensors[i], mesh, placements
+            )
+    return chunk_tensors
+
+
+def _concat_tensor(chunk_tensors, axis=0):
+    chunk0 = chunk_tensors[0]
+    if not chunk0.is_dist():
+        out = paddle.concat(chunk_tensors, axis)
+
+    else:
+        # loss_fun(out, labels), placements of labels is [S(0), R, ...]
+        mesh = chunk0.process_mesh
+        placements = [Replicate() for _ in range(mesh.ndim)]
+        dp_index = mesh.dim_names.index("dp") if "dp" in mesh.dim_names else 0
+        placements[dp_index] = Shard(0)
+
+        for i in range(len(chunk_tensors)):
+            chunk_tensors[i] = dist.reshard(chunk_tensors[i], mesh, placements)
+
+            chunk_tensors[i] = dtensor_to_local(
+                chunk_tensors[i], mesh, placements
+            )
+        out = paddle.concat(chunk_tensors, axis)
+        out = dtensor_from_local(out, mesh, placements)
+    return out
+
+
+class TensorChunkSpec:
+    """
+    Class used to specify chunking of inputs
+    """
+
+    def __init__(self, split_axis):
+        self.split_axis = split_axis
+
+    split_axis: int
+
+    def __repr__(self):
+        return f"{self.__class__.__module__}.{self.__class__.__name__}({self.split_axis})"
+
+    def __str__(self):
+        return f"TensorChunkSpec({self.split_axis})"
+
+
+def _split_args_helper(
+    args_dict,
+    args_chunk_spec,
+    num_chunks,
+):
+    """
+    A helper function of split_args_kwargs_into_chunks.
+    """
+    assert len(args_dict) == len(args_chunk_spec), (
+        f"args_dict.keys() = {list(args_dict.keys())} args_chunk_spec.keys() = {list(args_chunk_spec.keys())}"
+    )
+
+    shared_args_dict_flat = {}
+    # handle args one by one
+    for arg_key, arg in args_dict.items():
+        arg_flat = flatten(arg)
+
+        chunk_spec = args_chunk_spec[arg_key]
+        assert chunk_spec is not None
+
+        chunk_spec_flat = flatten(chunk_spec)
+        assert len(chunk_spec_flat) == len(arg_flat), (
+            f"{arg_key} {len(arg_flat)} != {len(chunk_spec_flat)}"
+        )
+
+        shard_arg_flat = []
+
+        for v, chunk_v in zip(arg_flat, chunk_spec_flat):
+            if not isinstance(v, paddle.Tensor):
+                shard_arg_flat.append([v] * num_chunks)
+            elif isinstance(chunk_v, TensorChunkSpec):
+                v_split_axis_size = v.shape[chunk_v.split_axis]
+
+                if v_split_axis_size < num_chunks:
+                    raise ValueError(
+                        f"Arg {arg_key} on chunking dimension has a size of {v_split_axis_size}, "
+                        f"smaller than the number of chunks {num_chunks}. "
+                        "Please adjust your num_chunks setting."
+                    )
+                # split tensor v
+                chunk_tensors = _split_tensor(v, num_chunks, chunk_v.split_axis)
+
+                shard_arg_flat.append(chunk_tensors)
+            else:
+                raise TypeError(f"Unrecognized chunk spec: {chunk_v}")
+
+        shared_args_dict_flat[arg_key] = shard_arg_flat
+
+    # the structure of each element in args_split is the same as the original args_dict
+    args_split = []
+    for idx in range(num_chunks):
+        chunk_args = {}
+        for key, arg in shared_args_dict_flat.items():
+            last_arg = None if not arg else arg[0][idx]
+            arg_of_curr_chunk = (
+                [v[idx] for v in arg] if len(arg) > 1 else last_arg
+            )
+            chunk_args[key] = arg_of_curr_chunk
+
+        # flatten chunk_args first, and then pack chunk_args as the origin args_dict
+        flatten_chunk_args = [x for x in flatten(chunk_args) if x is not None]
+        chunk_args = pack_sequence_as(args_dict, flatten_chunk_args)
+        args_split.append(chunk_args)
+    return args_split
+
+
+def split_args_kwargs_into_chunks(
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None,
+    chunks: int,
+    args_chunk_spec: (
+        tuple[
+            tuple[TensorChunkSpec, ...]
+            | list[TensorChunkSpec, ...]
+            | TensorChunkSpec,
+            ...,
+        ]
+        | None
+    ) = None,
+    kwargs_chunk_spec: (
+        dict[
+            str,
+            tuple[TensorChunkSpec, ...]
+            | list[TensorChunkSpec, ...]
+            | TensorChunkSpec,
+        ]
+        | None
+    ) = None,
+) -> tuple[list[tuple], list[dict]]:
+    """
+    Given a sequence of args and kwargs, split them into a number of chunks
+    according to  their respective chunking specs.
+
+    Args:
+        args: tuple of args
+        kwargs: dict of kwargs
+        chunks: Number of chunks to split the args and kwargs into
+        args_chunk_spec: chunking specs for args, in same shape as args
+        kwargs_chunk_spec: chunking specs for kwargs, in same shape as kwargs
+
+    Returns:
+        args_split: list of sharded args
+        kwargs_split: list of sharded kwargs
+    """
+
+    if kwargs is None:
+        kwargs = {}
+
+    if args_chunk_spec is None:
+        args_chunk_spec = map_structure(
+            lambda _: TensorChunkSpec(DEFAULT_CHUNK_DIM), args
+        )
+
+    if kwargs_chunk_spec is None:
+        kwargs_chunk_spec = map_structure(
+            lambda _: TensorChunkSpec(DEFAULT_CHUNK_DIM), kwargs
+        )
+
+    args_split_dict = _split_args_helper(
+        dict(enumerate(args)),
+        dict(enumerate(args_chunk_spec)),
+        chunks,
+    )
+    kwargs_split = _split_args_helper(
+        kwargs,
+        kwargs_chunk_spec,
+        chunks,
+    )
+
+    assert len(args_split_dict) == len(kwargs_split), (
+        "args and kwargs are split into difference number of chunks: "
+        f"{len(args_split_dict)}, {len(kwargs_split)}"
+    )
+
+    # the form of each args_chunk should be tuple
+    args_split = [
+        tuple(args_chunk[i] for i in range(len(args_chunk)))
+        for args_chunk in args_split_dict
+    ]
+
+    return args_split, kwargs_split
+
+
+def merge_chunks(
+    chunks: list[Any],
+    chunk_spec,
+):
+    """
+    Given a list of chunks, merge them into a single chunk according to
+    the chunk spec.
+
+    Args:
+        chunks: list of chunks
+        chunk_spec: Chunking spec for the chunks
+
+    Returns:
+        chunk: chunks merged value
+    """
+    if len(chunks) == 0:
+        logger.warning("No chunks to merge.")
+        return chunks
+
+    if chunk_spec is None:
+        chunk_spec = map_structure(
+            lambda _: TensorChunkSpec(DEFAULT_CHUNK_DIM), chunks[0]
+        )
+
+    chunks_flat = []
+    # flatten chunk_spec first
+    chunk_spec = flatten(chunk_spec)
+    for chunk in chunks:
+        chunk_flat = flatten(chunk)
+        assert len(chunk_flat) == len(chunk_spec), (
+            f"Chunk {chunk} did not match chunk spec {chunk_spec}"
+        )
+        chunks_flat.append(chunk_flat)
+
+    def _merge_non_tensor_type_arg(chunks, idx, chunk_spec_of_arg=None):
+        # use the first chunk's value as the merged result
+        arg_0 = chunks[0][idx]
+        for chunk_idx in range(1, len(chunks)):
+            assert chunks[chunk_idx][idx] == arg_0, (
+                f"Cannot merge chunks with index 0 and {idx} with different values,"
+                f"When the arg's TensorChunkSpec is {chunk_spec_of_arg}"
+            )
+        return arg_0
+
+    args_flat = []
+    for arg_idx, chunk_spec_of_arg in enumerate(chunk_spec):
+        if isinstance(chunk_spec_of_arg, TensorChunkSpec):
+            if isinstance(chunks_flat[0][arg_idx], paddle.Tensor):
+                arg_chunks_to_merge = [
+                    chunks_flat[chunk_idx][arg_idx]
+                    for chunk_idx in range(len(chunks_flat))
+                ]
+                merged_arg = _concat_tensor(
+                    arg_chunks_to_merge, axis=chunk_spec_of_arg.split_axis
+                )
+            else:
+                logger.warning(
+                    f"Cannot merge chunks with TensorChunkSpec {chunk_spec_of_arg}."
+                    "The TensorChunkSpec only supports paddle.Tensor type."
+                )
+
+                merged_arg = _merge_non_tensor_type_arg(
+                    chunks_flat, arg_idx, chunk_spec_of_arg
+                )
+        else:
+            merged_arg = _merge_non_tensor_type_arg(
+                chunks_flat, arg_idx, chunk_spec_of_arg
+            )
+
+        args_flat.append(merged_arg)
+
+    # pack args_flat as the input chunks[0]
+    return pack_sequence_as(chunks[0], args_flat)

python/paddle/distributed/auto_parallel/pipelining/utils.py

@@ -0,0 +1,177 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+import paddle
+from paddle.distributed import fleet
+from paddle.distributed.auto_parallel.api import (
+    dtensor_from_local,
+)
+from paddle.utils import map_structure
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+logger = logging.getLogger(__name__)
+
+
+def _detach_and_requires_grad(x):
+    o = x.detach()
+    o.stop_gradient = False
+    return o
+
+
+def _detach_and_keep_grad(x):
+    o = x.detach_()
+    o.stop_gradient = x.stop_gradient
+    return o
+
+
+def _zero_initialize_with_meta(meta, mesh):
+    assert isinstance(meta, TensorMeta)
+    x = paddle.zeros(
+        meta._local_shape if meta._local_shape else meta.shape, dtype=meta.dtype
+    )
+    if meta.placements:
+        x = dtensor_from_local(x, mesh, meta.placements)
+    return x
+
+
+def _flatten_args(args):
+    """
+    Flatten the args into a list form.
+    """
+    flat_args = []
+
+    def extract_tensor_args(a):
+        nonlocal flat_args
+        if isinstance(a, paddle.Tensor):
+            flat_args.append(a)
+        return a
+
+    paddle.utils.map_structure(
+        extract_tensor_args,
+        args,
+    )
+
+    return flat_args
+
+
+class PipeliningShapeError(RuntimeError):
+    """Shape mismatch between configured and runtime values."""
+
+
+def _validate_tensor_metadata(desc, expected, given):
+    if not expected.shape == given.shape:
+        raise PipeliningShapeError(
+            f"{desc} has a shape mismatch: expected {expected.shape} actual {given.shape}"
+        )
+    if not expected.dtype == given.dtype:
+        raise PipeliningShapeError(
+            f"{desc} has a dtype mismatch: expected {expected.dtype} actual {given.dtype}"
+        )
+
+
+def _validate_tensors_metadata(
+    desc,
+    expected_tensors: list[paddle.Tensor] | tuple[paddle.Tensor, ...],
+    actual_tensors: list[paddle.Tensor] | tuple[paddle.Tensor, ...],
+):
+    if len(expected_tensors) != len(actual_tensors):
+        raise PipeliningShapeError(
+            f"{desc}: Number of values ({len(actual_tensors)}) does not match expected number ({len(expected_tensors)})"
+        )
+    for i in range(len(expected_tensors)):
+        _validate_tensor_metadata(
+            f"{desc}: value {i}", expected_tensors[i], actual_tensors[i]
+        )
+
+
+NestedStruct = list[Any] | tuple[Any, ...] | dict[Any, Any]
+
+
+def _map_structure_only(
+    type_: Any, fn: Callable[[Any], Any], structure: NestedStruct
+) -> NestedStruct:
+    """
+    Apply `fn` to each entry which matches `type_` in `structure` and return a new structure with the same shape.
+    """
+    return map_structure(
+        lambda x: fn(x) if isinstance(x, type_) else x, structure
+    )
+
+
+class TensorMeta:
+    def __init__(self, tensor: paddle.Tensor):
+        if tensor.is_dist():
+            self.shape = tensor.shape
+            self._local_shape = tensor._local_shape
+        else:
+            self.shape = tensor.shape
+            self._local_shape = None
+        self.dtype = tensor.dtype
+        self.placements = None if not tensor.is_dist() else tensor.placements
+        self.stop_gradient = tensor.stop_gradient
+
+    def __repr__(self):
+        return f"TensorMeta(global_shape={self.shape},local_shape={self._local_shape}, dtype={self.dtype}, placements={self.placements})"
+
+
+def _get_pp_mesh(pp_idx=0, pp_dim_names="pp"):
+    """
+    Get the mesh of the {pp_idx}th PipelineStage.
+    """
+    mesh = fleet.auto.get_mesh()
+    assert mesh is not None, (
+        "the mesh is None, please call fleet.auto.set_mesh first."
+    )
+    if "pp" in mesh.dim_names:
+        mesh = mesh.get_mesh_with_dim("pp", pp_idx)
+    else:
+        logger.warning(
+            f"The dim name of pp {pp_dim_names} not exist in global mesh {mesh}"
+        )
+    return mesh
+
+
+def _get_stage_mesh(stage_index, pp_group_size, style=None):
+    if style == "v":
+        raise NotImplementedError
+    if style is not None:
+        raise ValueError(f"Unknown style: {style}, style can be None, v.")
+    else:
+        pp_idx = stage_index % pp_group_size
+        return _get_pp_mesh(pp_idx)
+
+
+def _friendly_debug_info(v):
+    """
+    Helper function to print out debug info in a friendly way.
+    """
+    if isinstance(v, paddle.Tensor):
+        return f"Tensor({v.shape}, stop_gradient={v.stop_gradient}, dtype={v.dtype})"
+    else:
+        return str(v)
+
+
+def _map_debug_info(a):
+    """
+    Helper function to apply `friendly_debug_info` to items in `a`.
+    `a` may be a list, tuple, or dict.
+    """
+    return map_structure(_friendly_debug_info, a)