chore: import upstream snapshot with attribution

python/paddle/distributed/fsdp/__init__.py

@@ -0,0 +1,17 @@
+# 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 .fully_shard import fully_shard
+
+__all__ = ["fully_shard"]

python/paddle/distributed/fsdp/fully_shard.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+    import paddle
+    import paddle.distributed as dist
+import paddle
+from paddle.distributed.auto_parallel.fully_shard import FullyShardAuto
+from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_fully_shard import (
+    FullyShard,
+)
+
+
+def in_auto_parallel_mode() -> bool:
+    return getattr(
+        paddle.base.framework.global_var, '_in_auto_parallel_', False
+    )
+
+
+# @dataclass
+class MixedPrecisionPolicy:
+    param_dtype: paddle.dtype | None = None
+    reduce_dtype: paddle.dtype | None = None
+    output_dtype: paddle.dtype | None = None
+    cast_forward_inputs: bool = True
+
+
+# @dataclass
+class OffloadPolicy:
+    pin_memory: bool = True
+
+
+def _fully_shard_manual_parallel(
+    module,
+    mesh,
+    reshard_after_forward,
+    shard_placement_fn,
+    mp_policy,
+    offload_policy,
+    ignored_params,
+    enable_tensor_fusion_and_overlap,
+):
+    return FullyShard(module)
+
+
+def _fully_shard_auto_parallel(
+    module,
+    mesh,
+    reshard_after_forward,
+    shard_placement_fn,
+    mp_policy,
+    offload_policy,
+    ignored_params,
+    enable_tensor_fusion_and_overlap,
+):
+    FullyShardAuto(
+        module,
+        mesh,
+        enable_tensor_fusion_and_overlap=enable_tensor_fusion_and_overlap,
+    )
+    return module
+
+
+def fully_shard(
+    module: paddle.nn.Layer,
+    *,
+    mesh: dist.ProcessMesh = None,
+    reshard_after_forward: bool | int | None = None,
+    shard_placement_fn: Callable[[paddle.Tensor], dist.Shard | None]
+    | None = None,
+    mp_policy: MixedPrecisionPolicy | None = None,
+    offload_policy: OffloadPolicy | None = None,
+    ignored_params: Iterable[paddle.Tensor] | None = None,
+    enable_tensor_fusion_and_overlap: bool = True,
+) -> paddle.nn.Layer:
+    """
+    Apply fully sharded data parallel (FSDP) to the given module.
+
+    This function wraps the input module with fully sharded data parallelism, which shards
+    model parameters, gradients, and optimizer states across multiple devices. It supports
+    both auto_parallel mode and manual_parallel mode.
+
+    Args:
+        module (Layer): The neural network module to be wrapped with fully sharded data parallelism.
+        mesh (dist.ProcessMesh, optional): The process mesh defining the device arrangement for sharding.
+            Defaults to None, which uses the default mesh.
+        reshard_after_forward (bool | int | None, optional): Controls when to reshard the parameters after forward pass.
+            If True or 1, reshard after each forward pass. If False or 0, keep sharded.
+            If None, use default strategy. Defaults to None.
+        shard_placement_fn (Callable[[paddle.Tensor], dist.Shard | None] | None, optional):
+            A function that determines how each tensor should be sharded. Takes a tensor as input
+            and returns a Shard placement or None. If None, uses default sharding strategy.
+            Defaults to None.
+        mp_policy (MixedPrecisionPolicy | None, optional): Mixed precision policy configuration.
+            If None, creates a default MixedPrecisionPolicy. Defaults to None.
+        offload_policy (OffloadPolicy | None, optional): Offload policy configuration for CPU offloading.
+            If None, creates a default OffloadPolicy. Defaults to None.
+        ignored_params (Iterable[paddle.Tensor] | None, optional): Parameters that should not be sharded.
+            These parameters will be kept in full precision and not distributed. Defaults to None.
+        enable_tensor_fusion_and_overlap (bool, optional): Whether to enable tensor fusion and
+            communication/computation overlap optimization. When True, parameters are fused into
+            contiguous buffers for more efficient communication and memory access, and enables
+            prefetching and asynchronous communication to overlap with computation for better
+            performance. Defaults to True.
+
+    Returns:
+        module: A wrapper module that applies FSDP to the input module.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # type: ignore
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> # python -m paddle.distributed.launch --device=0,1 train.py
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> from paddle.distributed.fsdp import fully_shard
+
+            >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+            >>> model = paddle.nn.Linear(10, 10)
+            >>> inputs = paddle.rand(shape=[10, 10])
+            >>> inputs = dist.shard_tensor(inputs, mesh, [dist.Shard(0)])
+            >>> opt = paddle.optimizer.AdamW(parameters=model.parameters())
+            >>> model = fully_shard(model, mesh)
+            >>> tr_loss = model(inputs)
+            >>> tr_loss.backward()
+            >>> opt.step()
+            >>> opt.clear_grad()
+    """
+    if mp_policy is None:
+        mp_policy = MixedPrecisionPolicy()
+    if offload_policy is None:
+        offload_policy = OffloadPolicy()
+    ignored_params_set: set[paddle.Tensor] = (
+        set(ignored_params) if ignored_params else set()
+    )
+
+    args = (
+        module,
+        mesh,
+        reshard_after_forward,
+        shard_placement_fn,
+        mp_policy,
+        offload_policy,
+        ignored_params_set,
+        enable_tensor_fusion_and_overlap,
+    )
+
+    if in_auto_parallel_mode():
+        return _fully_shard_auto_parallel(*args)
+    else:
+        return _fully_shard_manual_parallel(*args)