chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,256 @@
|
||||
# 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 functools
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.utils import flatten, pack_sequence_as
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from paddle.distributed import ProcessMesh
|
||||
|
||||
|
||||
def local_map(
|
||||
func: Callable[..., Any],
|
||||
out_placements: list[list[dist.Placement]],
|
||||
in_placements: list[list[dist.Placement]] | None = None,
|
||||
process_mesh: ProcessMesh | None = None,
|
||||
reshard_inputs: bool = False,
|
||||
) -> Callable[..., Any]:
|
||||
"""
|
||||
The `local_map` API allows users to pass dist_tensors to a function that is written
|
||||
to be applied on ``paddle.Tensor`` s. It works by extracting the local components
|
||||
of dist_tensors, calling the function, and wrapping the outputs as dist_tensors
|
||||
according to the ``out_placements``.
|
||||
|
||||
Args:
|
||||
func (Callable): The function to be applied on each local shard of dist_tensors.
|
||||
|
||||
out_placements (list[list[dist.Placement]]):
|
||||
The desired placements for each output tensor. Must be a list where each element
|
||||
is a list of Placement objects specifying the distribution strategy for that
|
||||
output tensor. The length of the outer list must match the number of outputs
|
||||
from ``func``. For non-tensor outputs, the corresponding placement must be None.
|
||||
When there are no dist_tensor inputs, process_mesh must be specified to use
|
||||
non-None placements.
|
||||
|
||||
in_placements (Optional[list[list[dist.Placement]]], optional):
|
||||
The required placements for each input tensor. If specified, must be a list
|
||||
where each element is a list of Placement objects defining the distribution
|
||||
strategy for that input tensor. The length of the outer list must match the
|
||||
number of input tensors.
|
||||
Default: None
|
||||
|
||||
process_mesh (ProcessMesh, optional):
|
||||
The process mesh that all dist_tensors are placed on. If not specified,
|
||||
this will be inferred from the input dist_tensors' process mesh.
|
||||
local_map requires all dist_tensors to be placed on the same process mesh.
|
||||
Must be specified when there are no dist_tensor inputs but out_placements
|
||||
contains non-None values.
|
||||
Default: None
|
||||
|
||||
reshard_inputs (bool, optional):
|
||||
the bool value indicating whether to reshard the input :dist_tensors when
|
||||
their placements are different from the required input placements. If this
|
||||
value is ``False`` and some :dist_tensor input has a different placement,
|
||||
an exception will be raised. Default: False.
|
||||
|
||||
Returns:
|
||||
Callable: A function that applies func to local shards of input dist_tensors and returns dist_tensors or original values.
|
||||
|
||||
Example:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> from __future__ import annotations
|
||||
>>> import paddle
|
||||
>>> import paddle.distributed as dist
|
||||
>>> from paddle import Tensor
|
||||
>>> from paddle.distributed import ProcessMesh
|
||||
|
||||
>>> def custom_function(x):
|
||||
... mask = paddle.zeros_like(x)
|
||||
... if dist.get_rank() == 0:
|
||||
... mask[1:3] = 1
|
||||
... else:
|
||||
... mask[4:7] = 1
|
||||
... x = x * mask
|
||||
... mask_sum = paddle.sum(x)
|
||||
... mask_sum = mask_sum / mask.sum()
|
||||
... return mask_sum
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
||||
>>> dist.init_parallel_env()
|
||||
>>> mesh = ProcessMesh([0, 1], dim_names=["x"])
|
||||
>>> local_input = paddle.arange(0, 10, dtype="float32")
|
||||
>>> local_input = local_input + dist.get_rank()
|
||||
>>> input_dist = dist.auto_parallel.api.dtensor_from_local(local_input, mesh, [dist.Shard(0)])
|
||||
>>> wrapped_func = dist.local_map(
|
||||
... custom_function,
|
||||
... out_placements=[[dist.Partial(dist.ReduceType.kRedSum)]],
|
||||
... in_placements=[[dist.Shard(0)]],
|
||||
... process_mesh=mesh,
|
||||
... )
|
||||
>>> output_dist = wrapped_func(input_dist)
|
||||
|
||||
>>> local_value = output_dist._local_value()
|
||||
>>> gathered_values: list[Tensor] = []
|
||||
>>> dist.all_gather(gathered_values, local_value)
|
||||
|
||||
>>> print(f"[Rank 0] local_value={gathered_values[0].item()}")
|
||||
[Rank 0] local_value=1.5
|
||||
>>> print(f"[Rank 1] local_value={gathered_values[1].item()}")
|
||||
[Rank 1] local_value=6.0
|
||||
>>> print(f"global_value (distributed)={output_dist.item()}")
|
||||
global_value (distributed)=7.5
|
||||
|
||||
>>> # This case needs to be executed in a multi-card environment
|
||||
>>> # export CUDA_VISIBLE_DEVICES=0,1
|
||||
>>> # python -m paddle.distributed.launch {test_case}.py
|
||||
"""
|
||||
|
||||
def wrapped(process_mesh: ProcessMesh | None, *args, **kwargs):
|
||||
# Process input arguments
|
||||
flat_dist_args = flatten(args)
|
||||
if in_placements is not None:
|
||||
assert len(in_placements) == len(flat_dist_args), (
|
||||
f"in_placements length {len(in_placements)} does not match "
|
||||
f"number of input args {len(flat_dist_args)}!"
|
||||
)
|
||||
|
||||
flat_local_args = []
|
||||
seen_dist_tensor = False
|
||||
|
||||
for idx, arg in enumerate(flat_dist_args):
|
||||
if dist.auto_parallel.api.is_dist_tensor(arg):
|
||||
dist_tensor = arg
|
||||
if process_mesh is None:
|
||||
if paddle.in_dynamic_mode():
|
||||
process_mesh = dist_tensor.process_mesh
|
||||
else:
|
||||
process_mesh = dist_tensor.dist_attr().process_mesh
|
||||
|
||||
seen_dist_tensor = True
|
||||
|
||||
if in_placements is not None:
|
||||
in_placement = in_placements[idx]
|
||||
if in_placement is None:
|
||||
if paddle.in_dynamic_mode():
|
||||
in_placement = dist_tensor.placements
|
||||
else:
|
||||
in_placement = dist_tensor.dist_attr().placements
|
||||
else:
|
||||
if paddle.in_dynamic_mode():
|
||||
if in_placement != dist_tensor.placements:
|
||||
if reshard_inputs:
|
||||
dist_tensor = dist.reshard(
|
||||
dist_tensor, process_mesh, in_placement
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"in_placement {in_placement} does not match dist_tensor.placements {dist_tensor.placements}"
|
||||
)
|
||||
|
||||
else:
|
||||
if (
|
||||
in_placement
|
||||
!= dist_tensor.dist_attr().placements
|
||||
):
|
||||
if reshard_inputs:
|
||||
dist_tensor = dist.reshard(
|
||||
dist_tensor, process_mesh, in_placement
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"in_placement {in_placement} does not match dist_tensor.dist_attr().placements {dist_tensor.dist_attr().placements}"
|
||||
"If reshard_inputs is wanted, set "
|
||||
"reshard_inputs=True to local_map."
|
||||
)
|
||||
local_tensor = dist.auto_parallel.api.dtensor_to_local(
|
||||
dist_tensor, process_mesh, in_placement
|
||||
)
|
||||
flat_local_args.append(local_tensor)
|
||||
else:
|
||||
flat_local_args.append(arg)
|
||||
|
||||
local_args = pack_sequence_as(args, flat_local_args)
|
||||
out = func(*local_args, **kwargs)
|
||||
original_out = out
|
||||
if seen_dist_tensor:
|
||||
flat_out = flatten(out)
|
||||
assert len(flat_out) == len(out_placements), (
|
||||
"local_map requires one PlacementType for each output value, "
|
||||
f"got {len(out_placements)} placements but expected "
|
||||
f"{len(flat_out)}!"
|
||||
)
|
||||
|
||||
flat_dist_and_arg_out = []
|
||||
for out, out_placement in zip(flat_out, out_placements):
|
||||
if paddle.in_dynamic_mode():
|
||||
if isinstance(out, paddle.Tensor):
|
||||
assert not dist.auto_parallel.api.is_dist_tensor(out), (
|
||||
f"Expected dense tensor output but got {type(out)}: {out}"
|
||||
)
|
||||
|
||||
flat_dist_and_arg_out.append(
|
||||
dist.auto_parallel.api.dtensor_from_local(
|
||||
out, process_mesh, out_placement
|
||||
)
|
||||
)
|
||||
else:
|
||||
assert out_placement is None, (
|
||||
f"Expected None placements for non-tensor output {out} "
|
||||
f"but got {out_placement}!"
|
||||
)
|
||||
flat_dist_and_arg_out.append(out)
|
||||
else:
|
||||
if isinstance(out, paddle.base.libpaddle.pir.Value):
|
||||
assert not dist.auto_parallel.api.is_dist_tensor(out), (
|
||||
f"Expected dense tensor output but got {type(out)}: {out}"
|
||||
)
|
||||
|
||||
flat_dist_and_arg_out.append(
|
||||
dist.auto_parallel.api.dtensor_from_local(
|
||||
out, process_mesh, out_placement
|
||||
)
|
||||
)
|
||||
else:
|
||||
assert out_placement is None, (
|
||||
f"Expected None placements for non-tensor output {out} "
|
||||
f"but got {out_placement}!"
|
||||
)
|
||||
flat_dist_and_arg_out.append(out)
|
||||
return pack_sequence_as(original_out, flat_dist_and_arg_out)
|
||||
else:
|
||||
flat_out = flatten(out)
|
||||
flat_dist_and_arg_out = []
|
||||
for out, out_placement in zip(flat_out, out_placements):
|
||||
if out_placement is not None:
|
||||
assert process_mesh is not None, (
|
||||
"process_mesh must be specified when out_placements is not None"
|
||||
)
|
||||
flat_dist_and_arg_out.append(
|
||||
dist.auto_parallel.api.dtensor_from_local(
|
||||
out, process_mesh, out_placement
|
||||
)
|
||||
)
|
||||
else:
|
||||
flat_dist_and_arg_out.append(out)
|
||||
return pack_sequence_as(original_out, flat_dist_and_arg_out)
|
||||
|
||||
return functools.partial(wrapped, process_mesh)
|
||||
Reference in New Issue
Block a user