chore: import upstream snapshot with attribution

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)