chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,131 @@
|
||||
# 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.
|
||||
|
||||
import paddle
|
||||
|
||||
# all registered reshard functions
|
||||
_g_reshard_func_list = []
|
||||
|
||||
|
||||
class ReshardFunction:
|
||||
def is_suitable(self, dist_tensor, dist_attr):
|
||||
raise NotImplementedError
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def choose_reshard_func(src_dist_attr, dst_dist_attr):
|
||||
global _g_reshard_func_list
|
||||
for reshard_func in _g_reshard_func_list:
|
||||
if reshard_func.is_suitable(src_dist_attr, dst_dist_attr):
|
||||
return reshard_func
|
||||
return None
|
||||
|
||||
|
||||
def register_reshard_func(reshard_func):
|
||||
global _g_reshard_func_list
|
||||
_g_reshard_func_list.append(reshard_func)
|
||||
|
||||
|
||||
def clean_reshard_funcs():
|
||||
global _g_reshard_func_list
|
||||
_g_reshard_func_list.clear()
|
||||
|
||||
|
||||
def is_shard(dist_attr):
|
||||
for v in dist_attr.dims_mapping:
|
||||
if v != -1:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_partial(dist_attr):
|
||||
if len(dist_attr.partial_status) > 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_replicated(dist_attr):
|
||||
dims_mapping_set = set(dist_attr.dims_mapping)
|
||||
if len(dist_attr.partial_status) == 0 and (
|
||||
len(dims_mapping_set) == 0
|
||||
or (len(dims_mapping_set) == 1 and -1 in dims_mapping_set)
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def copy_dist_attr_with_new_member(
|
||||
src_dist_attr,
|
||||
new_process_mesh=None,
|
||||
new_dims_mapping=None,
|
||||
new_partial_status=None,
|
||||
):
|
||||
if new_process_mesh is None:
|
||||
new_process_mesh = src_dist_attr.process_mesh
|
||||
if new_dims_mapping is None:
|
||||
new_dims_mapping = src_dist_attr.dims_mapping
|
||||
if new_partial_status is None:
|
||||
new_partial_status = src_dist_attr.partial_status
|
||||
|
||||
return paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
new_process_mesh,
|
||||
new_dims_mapping,
|
||||
new_partial_status,
|
||||
)
|
||||
|
||||
|
||||
def copy_op_attr_with_new_member(
|
||||
src_dist_attr,
|
||||
new_process_mesh=None,
|
||||
new_operands=None,
|
||||
new_results=None,
|
||||
new_chunk_id=None,
|
||||
):
|
||||
if new_process_mesh is None:
|
||||
new_process_mesh = src_dist_attr.process_mesh
|
||||
if new_operands is None:
|
||||
new_operands = src_dist_attr.operands()
|
||||
if new_results is None:
|
||||
new_results = src_dist_attr.results()
|
||||
if new_chunk_id is None:
|
||||
new_chunk_id = src_dist_attr.chunk_id
|
||||
|
||||
return paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
new_process_mesh,
|
||||
new_operands,
|
||||
new_results,
|
||||
new_chunk_id,
|
||||
)
|
||||
|
||||
|
||||
def copy_process_mesh_with_new_member(
|
||||
src_process_mesh,
|
||||
new_shape=None,
|
||||
new_process_ids=None,
|
||||
new_dim_names=None,
|
||||
):
|
||||
if new_shape is None:
|
||||
new_shape = src_process_mesh.shape
|
||||
if new_process_ids is None:
|
||||
new_process_ids = src_process_mesh.process_ids
|
||||
if new_dim_names is None:
|
||||
new_dim_names = src_process_mesh.dim_names
|
||||
|
||||
return paddle.base.libpaddle.pir.create_process_mesh(
|
||||
new_shape,
|
||||
new_process_ids,
|
||||
new_dim_names,
|
||||
)
|
||||
+94
@@ -0,0 +1,94 @@
|
||||
# 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.
|
||||
|
||||
import paddle
|
||||
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
is_replicated,
|
||||
)
|
||||
from .nd_mesh_reshard_func import NdMeshReshardFunction
|
||||
|
||||
|
||||
class GlobalToSubMeshFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
# NOTE we could allow the src_dist_attr is not replicated and reshard it as replicated before go through the global_to_sub logic
|
||||
# but the dst_dist_attr should be replicated otherwise there will be un-defined result when change the mesh.
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
if in_mesh.ndim > out_mesh.ndim + 1:
|
||||
return False
|
||||
if in_mesh.ndim == out_mesh.ndim:
|
||||
return set(out_mesh.process_ids) < set(in_mesh.process_ids)
|
||||
else:
|
||||
sub_meshes = paddle.base.libpaddle.pir.get_sub_meshes(in_mesh)
|
||||
return out_mesh in sub_meshes
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
# reshard operand as replicated before change the mesh.
|
||||
if not is_replicated(src_dist_attr):
|
||||
tmp_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[-1] * len(src_dist_attr.dims_mapping),
|
||||
{},
|
||||
)
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
|
||||
pre_reshard_func = NdMeshReshardFunction()
|
||||
src_value = pre_reshard_func.reshard(
|
||||
src_dist_attr,
|
||||
tmp_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dist_attr
|
||||
|
||||
if src_value.has_one_use():
|
||||
src_value.update_dist_attr(dst_dist_attr)
|
||||
prev_op = src_value.get_defining_op()
|
||||
op_dist_attr = prev_op.dist_attr
|
||||
op_mesh = op_dist_attr.process_mesh
|
||||
operands = op_dist_attr.operands()
|
||||
results = op_dist_attr.results()
|
||||
chunk_id = op_dist_attr.chunk_id
|
||||
results[src_value.index()] = dst_dist_attr
|
||||
prev_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
op_mesh, operands, results, chunk_id
|
||||
)
|
||||
)
|
||||
return src_value
|
||||
else:
|
||||
dst_value = paddle._C_ops.share_data_(src_value)
|
||||
share_data_op = dst_value.get_defining_op()
|
||||
# set dist type and dist attr
|
||||
dst_value.set_type(dst_type)
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
share_data_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
return dst_value
|
||||
@@ -0,0 +1,365 @@
|
||||
# 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.
|
||||
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed.auto_parallel.static.utils import split_mesh
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
copy_dist_attr_with_new_member,
|
||||
is_partial,
|
||||
)
|
||||
from .p_to_r_reshard_func import PToRReshardFunction
|
||||
from .p_to_s_reshard_func import PToSReshardFunction
|
||||
from .r_to_p_reshard_func import RToPReshardFunction
|
||||
from .r_to_s_reshard_func import RToSReshardFunction
|
||||
from .s_to_r_reshard_func import SToRReshardFunction
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
|
||||
|
||||
def find_first_diff_shard_axis(src_dist_attr, dst_dist_attr):
|
||||
src_dims_mapping = src_dist_attr.dims_mapping
|
||||
dst_dims_mapping = dst_dist_attr.dims_mapping
|
||||
ndim = len(src_dims_mapping)
|
||||
for i in range(ndim - 1, -1, -1):
|
||||
if src_dims_mapping[i] != dst_dims_mapping[i]:
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
def get_1D_sub_process_mesh(process_mesh, mesh_dim):
|
||||
"""
|
||||
Get the 1-D sub process mesh on specific mesh_dim which:
|
||||
1) where the reshard should be performed.
|
||||
2) contains current process.
|
||||
|
||||
Args:
|
||||
process_mesh (ProcessMesh): the global process mesh.
|
||||
mesh_dim (int): the mesh dimension where the dist_tensor is
|
||||
sharded or partial.
|
||||
|
||||
e.g.
|
||||
1) process_mesh = [[0, 1, 2], [3, 4, 5]], axis = 0:
|
||||
process rank id returned sub mesh
|
||||
0 or 3 [0, 3]
|
||||
1 or 4 [1, 4]
|
||||
2 or 5 [2, 5]
|
||||
2) process_mesh = [[0, 1, 2], [3, 4, 5]], axis = 1:
|
||||
process rank id returned sub mesh
|
||||
0 or 1 or 2 [0, 1, 2]
|
||||
3 or 4 or 5 [3, 4, 5]
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
mesh_shape = process_mesh.shape
|
||||
dim_names = process_mesh.dim_names
|
||||
process_ids = np.array(process_mesh.process_ids).reshape(mesh_shape)
|
||||
|
||||
rank_id = dist.get_rank()
|
||||
# FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
|
||||
if rank_id not in process_mesh.process_ids:
|
||||
rank_id = process_mesh.process_ids[0]
|
||||
coord = list(np.where(process_ids == rank_id))
|
||||
coord[mesh_dim] = range(mesh_shape[mesh_dim])
|
||||
sub_process_ids = process_ids[tuple(coord)].flatten()
|
||||
sub_mesh_name = dim_names[mesh_dim]
|
||||
|
||||
return dist.ProcessMesh(sub_process_ids, [sub_mesh_name])
|
||||
|
||||
|
||||
class NdMeshReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
if out_mesh.ndim <= 1:
|
||||
return False
|
||||
|
||||
# check dims_mapping and partial_status
|
||||
if src_dist_attr == dst_dist_attr:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
"""
|
||||
Reshard on N-d mesh:
|
||||
1. Find the tensor dimensions where the dims_mapping values
|
||||
differ between src_dist_attr and dst_dist_attr.
|
||||
2. From higher to lower, convert the non-replicated dimensions
|
||||
in step1 to replicated using corresponding 1-D mesh functions.
|
||||
3. Convert the replicated dimensions in step2 to the status in
|
||||
dst_dist_attr with corresponding 1-D mesh functions.
|
||||
"""
|
||||
# Step1. find first dimension with different shard status in src_dist_attr
|
||||
# and dst_dist_attr.
|
||||
first_diff_axis = find_first_diff_shard_axis(
|
||||
src_dist_attr, dst_dist_attr
|
||||
)
|
||||
# out_value = src_value # intermediate result
|
||||
# src_type = src_value.type()
|
||||
tensor_ndim = len(src_value.shape)
|
||||
process_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
# Step2. Convert the non-replicated dimensions to replicated.
|
||||
|
||||
# Step2.1 convert shard status to replicated
|
||||
for i in range(first_diff_axis, -1, -1):
|
||||
in_mesh_axis = src_dist_attr.dims_mapping[i]
|
||||
out_mesh_axis = dst_dist_attr.dims_mapping[i]
|
||||
if in_mesh_axis == -1 or in_mesh_axis == out_mesh_axis:
|
||||
continue
|
||||
|
||||
# calculate the dist_attr after converting
|
||||
tmp_dims_mapping = src_dist_attr.dims_mapping
|
||||
tmp_dims_mapping[i] = -1
|
||||
tmp_dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
src_dist_attr, new_dims_mapping=tmp_dims_mapping
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dst_dist_attr
|
||||
)
|
||||
sub_mesh_list = split_mesh(process_mesh, in_mesh_axis)
|
||||
for sub_mesh in sub_mesh_list:
|
||||
new_process_group(sorted(sub_mesh.process_ids))
|
||||
# get the process_mesh on specific axis
|
||||
sub_mesh = get_1D_sub_process_mesh(process_mesh, in_mesh_axis)
|
||||
|
||||
# calculate corresponding 1-D dist_attr of src_dst_attr
|
||||
in_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
in_one_dim_dims_mapping[i] = 0
|
||||
in_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, in_one_dim_dims_mapping, {}
|
||||
)
|
||||
)
|
||||
|
||||
# calculate corresponding 1-D dist_attr of dst_dst_attr
|
||||
out_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
out_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, out_one_dim_dims_mapping, {}
|
||||
)
|
||||
)
|
||||
|
||||
one_dim_func = SToRReshardFunction()
|
||||
src_value = one_dim_func.reshard(
|
||||
in_one_dim_dist_attr,
|
||||
out_one_dim_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dst_dist_attr
|
||||
|
||||
# Step2.2. convert partial status to replicated
|
||||
if is_partial(src_dist_attr):
|
||||
in_partial_status = src_dist_attr.partial_status
|
||||
out_partial_status = dst_dist_attr.partial_status # read-only
|
||||
# convert each partial dim to replicated with corresponding
|
||||
# 1-D mesh function
|
||||
for partial_dim, partial_type in in_partial_status.items():
|
||||
if partial_dim in out_partial_status:
|
||||
if out_partial_status[partial_dim] != partial_type:
|
||||
raise NotImplementedError(
|
||||
f"Reshard tensor from one partial type {partial_type} to another partial type {out_partial_status[partial_dim]} is not supported yet."
|
||||
)
|
||||
continue
|
||||
|
||||
p_to_s = False
|
||||
if partial_dim in dst_dist_attr.dims_mapping:
|
||||
p_to_s = True
|
||||
shard_index = dst_dist_attr.dims_mapping.index(partial_dim)
|
||||
# get the partial status after converting
|
||||
tmp_partial_status = src_dist_attr.partial_status
|
||||
tmp_partial_status.pop(partial_dim)
|
||||
|
||||
tmp_dims_mapping = src_dist_attr.dims_mapping
|
||||
if p_to_s:
|
||||
tmp_dims_mapping[shard_index] = partial_dim
|
||||
|
||||
tmp_dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
src_dist_attr,
|
||||
new_dims_mapping=tmp_dims_mapping,
|
||||
new_partial_status=tmp_partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dst_dist_attr
|
||||
)
|
||||
sub_mesh_list = split_mesh(process_mesh, partial_dim)
|
||||
for sub_mesh in sub_mesh_list:
|
||||
new_process_group(sorted(sub_mesh.process_ids))
|
||||
# get the process_mesh on specific axis
|
||||
sub_mesh = get_1D_sub_process_mesh(process_mesh, partial_dim)
|
||||
|
||||
# calculate corresponding 1-D dist_attr of src_dst_attr
|
||||
in_one_dim_partial_status = {0: partial_type}
|
||||
in_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh,
|
||||
[-1] * tensor_ndim,
|
||||
in_one_dim_partial_status,
|
||||
)
|
||||
)
|
||||
out_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
one_dim_func = PToRReshardFunction()
|
||||
if p_to_s:
|
||||
out_one_dim_dims_mapping[shard_index] = 0
|
||||
one_dim_func = PToSReshardFunction()
|
||||
|
||||
# calculate corresponding 1-D dist_attr of dst_dst_attr
|
||||
out_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh,
|
||||
out_one_dim_dims_mapping,
|
||||
{},
|
||||
)
|
||||
)
|
||||
|
||||
src_value = one_dim_func.reshard(
|
||||
in_one_dim_dist_attr,
|
||||
out_one_dim_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dst_dist_attr
|
||||
|
||||
# Step3. Convert the replicated status to the status in dst_dist_attr
|
||||
# Step3.1 convert replicated to partial
|
||||
if is_partial(dst_dist_attr):
|
||||
in_partial_status = src_dist_attr.partial_status
|
||||
out_partial_status = dst_dist_attr.partial_status
|
||||
for partial_dim, partial_type in out_partial_status.items():
|
||||
if partial_dim in in_partial_status:
|
||||
continue
|
||||
|
||||
sub_mesh = get_1D_sub_process_mesh(process_mesh, partial_dim)
|
||||
|
||||
in_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh,
|
||||
[-1] * tensor_ndim,
|
||||
{},
|
||||
)
|
||||
)
|
||||
out_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, [-1] * tensor_ndim, {0: partial_type}
|
||||
)
|
||||
)
|
||||
|
||||
tmp_dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
dst_dist_attr,
|
||||
new_partial_status={partial_dim: partial_type},
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dst_dist_attr
|
||||
)
|
||||
|
||||
src_value = RToPReshardFunction().reshard(
|
||||
in_one_dim_dist_attr,
|
||||
out_one_dim_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dst_dist_attr
|
||||
|
||||
# Step3.2 convert replicated to shard
|
||||
for i in range(first_diff_axis, -1, -1):
|
||||
in_mesh_axis = src_dist_attr.dims_mapping[i]
|
||||
out_mesh_axis = dst_dist_attr.dims_mapping[i]
|
||||
if in_mesh_axis == out_mesh_axis:
|
||||
continue
|
||||
|
||||
# calculate the dist_attr after converting
|
||||
tmp_dims_mapping = src_dist_attr.dims_mapping
|
||||
tmp_dims_mapping[i] = out_mesh_axis
|
||||
tmp_dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
src_dist_attr, new_dims_mapping=tmp_dims_mapping
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dst_dist_attr
|
||||
)
|
||||
|
||||
# get the process_mesh on specific axis
|
||||
sub_mesh = get_1D_sub_process_mesh(process_mesh, out_mesh_axis)
|
||||
|
||||
# calculate the corresponding 1-D input dist attr
|
||||
in_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
in_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, in_one_dim_dims_mapping, {}
|
||||
)
|
||||
)
|
||||
|
||||
# calculate the corresponding 1-D output dist attr
|
||||
out_one_dim_dims_mapping = [-1] * tensor_ndim
|
||||
out_one_dim_dims_mapping[i] = 0
|
||||
out_one_dim_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
sub_mesh, out_one_dim_dims_mapping, {}
|
||||
)
|
||||
)
|
||||
one_dim_func = RToSReshardFunction()
|
||||
src_value = one_dim_func.reshard(
|
||||
in_one_dim_dist_attr,
|
||||
out_one_dim_dist_attr,
|
||||
src_value,
|
||||
tmp_dst_type,
|
||||
)
|
||||
src_dist_attr = tmp_dst_dist_attr
|
||||
return src_value
|
||||
|
||||
|
||||
class NdMeshReshardFunctionCrossMesh(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
if in_mesh.shape != out_mesh.shape:
|
||||
return False
|
||||
if out_mesh.ndim <= 1:
|
||||
return False
|
||||
if src_dist_attr == dst_dist_attr:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
same_status_func = SameStatusReshardFunction()
|
||||
tmp_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
src_value = same_status_func.reshard(
|
||||
src_dist_attr, tmp_dist_attr, src_value, tmp_dst_type
|
||||
)
|
||||
|
||||
nd_mesh_func = NdMeshReshardFunction()
|
||||
assert nd_mesh_func.is_suitable(tmp_dist_attr, dst_dist_attr), (
|
||||
f"Invoke the p to r reshard function is not valid from {tmp_dist_attr} to {dst_dist_attr}"
|
||||
)
|
||||
return nd_mesh_func.reshard(
|
||||
tmp_dist_attr, dst_dist_attr, src_value, dst_type
|
||||
)
|
||||
@@ -0,0 +1,113 @@
|
||||
# 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.
|
||||
|
||||
import paddle
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import ReshardFunction, is_partial, is_replicated
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
|
||||
|
||||
class PToRReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_partial(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
src_reduce_type = src_dist_attr.partial_status[0]
|
||||
# reduce_mean = False
|
||||
# if src_reduce_type == paddle.base.core.ReduceType.kRedAvg:
|
||||
# src_reduce_type = paddle.base.core.ReduceType.kRedSum
|
||||
# reduce_mean = True
|
||||
|
||||
group = new_process_group(sorted(src_mesh.process_ids))
|
||||
reduced_value = paddle._C_ops.all_reduce(
|
||||
src_value, group.id, int(src_reduce_type)
|
||||
)
|
||||
# set dist type and dist attr
|
||||
reduced_value.set_type(dst_type)
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
|
||||
reduced_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
return reduced_value
|
||||
|
||||
|
||||
class PToRReshardFunctionCrossMesh(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_partial(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if (
|
||||
in_mesh.ndim != 1
|
||||
or out_mesh.ndim != 1
|
||||
or in_mesh.shape != out_mesh.shape
|
||||
):
|
||||
return False
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
same_status_func = SameStatusReshardFunction()
|
||||
tmp_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
src_value = same_status_func.reshard(
|
||||
src_dist_attr, tmp_dist_attr, src_value, tmp_dst_type
|
||||
)
|
||||
|
||||
p_to_r_func = PToRReshardFunction()
|
||||
assert p_to_r_func.is_suitable(tmp_dist_attr, dst_dist_attr), (
|
||||
f"Invoke the p to r reshard function is not valid from {tmp_dist_attr} to {dst_dist_attr}"
|
||||
)
|
||||
return p_to_r_func.reshard(
|
||||
tmp_dist_attr, dst_dist_attr, src_value, dst_type
|
||||
)
|
||||
@@ -0,0 +1,245 @@
|
||||
# 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.
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed.utils.stream_utils import ExecutionStreamType
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
copy_dist_attr_with_new_member,
|
||||
is_partial,
|
||||
is_shard,
|
||||
)
|
||||
|
||||
|
||||
class PToSReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_partial(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
src_reduce_type = src_dist_attr.partial_status[0]
|
||||
assert src_reduce_type == paddle.base.core.ReduceType.kRedSum, (
|
||||
f"The p to s reshard func only support sum op, but received {src_reduce_type}"
|
||||
)
|
||||
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
|
||||
split_axis = dst_dist_attr.dims_mapping.index(0)
|
||||
num_of_process = len(src_dist_attr.process_mesh.process_ids)
|
||||
remainder_of_padding = src_value.shape[split_axis] % num_of_process
|
||||
is_balanced_split = remainder_of_padding == 0
|
||||
|
||||
permute = False
|
||||
if split_axis != 0:
|
||||
perm = list(range(0, len(src_value.shape)))
|
||||
perm[0] = split_axis
|
||||
perm[split_axis] = 0
|
||||
src_value = paddle._C_ops.transpose(src_value, perm)
|
||||
permute = True
|
||||
tmp_dims_mapping = dst_dist_attr.dims_mapping
|
||||
tmp_dims_mapping[split_axis] = -1
|
||||
tmp_dims_mapping[0] = 0
|
||||
dst_dist_attr = copy_dist_attr_with_new_member(
|
||||
dst_dist_attr, new_dims_mapping=tmp_dims_mapping
|
||||
)
|
||||
|
||||
if is_balanced_split:
|
||||
global_dst_attr = dst_type.as_dist_type().dist_attr()
|
||||
global_dims_mapping = global_dst_attr.dims_mapping
|
||||
axis = global_dims_mapping[0]
|
||||
global_dims_mapping[0] = global_dims_mapping[split_axis]
|
||||
global_dims_mapping[split_axis] = axis
|
||||
global_dist_attr = copy_dist_attr_with_new_member(
|
||||
global_dst_attr, new_dims_mapping=global_dims_mapping
|
||||
)
|
||||
dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), global_dist_attr
|
||||
)
|
||||
group = new_process_group(sorted(src_mesh.process_ids))
|
||||
dst_value = paddle._C_ops.reduce_scatter(
|
||||
src_value, group.id, num_of_process
|
||||
)
|
||||
dst_value.get_defining_op().set_execution_stream(
|
||||
ExecutionStreamType.DefaultStream.value
|
||||
)
|
||||
|
||||
# set dist type and dist attr
|
||||
dst_value.set_type(dst_type)
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [dst_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
|
||||
if split_axis != 0:
|
||||
dst_value = paddle._C_ops.transpose(dst_value, perm)
|
||||
return dst_value
|
||||
else:
|
||||
dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), dst_dist_attr
|
||||
)
|
||||
original_dims_mapping = dst_dist_attr.dims_mapping.copy()
|
||||
original_split_axis = split_axis
|
||||
split_axis = 0
|
||||
avg_size_on_split_axis = int(
|
||||
(src_value.shape[split_axis] + num_of_process - 1)
|
||||
/ num_of_process
|
||||
)
|
||||
padding_num = (
|
||||
avg_size_on_split_axis * num_of_process
|
||||
- src_value.shape[split_axis]
|
||||
)
|
||||
padding_shape = src_value._local_shape
|
||||
padding_shape[split_axis] = padding_num
|
||||
padding_tensor = paddle.full(
|
||||
padding_shape,
|
||||
0.0,
|
||||
src_value.dtype,
|
||||
)
|
||||
tmp_src_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
padding_tensor.type(), src_dist_attr
|
||||
)
|
||||
padding_tensor.set_type(tmp_src_type)
|
||||
padding_tensor.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh, [], [src_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
concat_value = paddle._C_ops.concat(
|
||||
[src_value, padding_tensor], split_axis
|
||||
)
|
||||
axis_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[-1],
|
||||
{0: paddle.base.core.ReduceType.kRedSum},
|
||||
)
|
||||
)
|
||||
concat_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[
|
||||
paddle.base.libpaddle.pir.create_array_attribute(
|
||||
[src_dist_attr, src_dist_attr]
|
||||
),
|
||||
axis_dist_attr,
|
||||
],
|
||||
[src_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
|
||||
concat_global_shape = list(src_value.shape)
|
||||
concat_global_shape[split_axis] = (
|
||||
avg_size_on_split_axis * num_of_process
|
||||
)
|
||||
concat_type = paddle.pir.create_shaped_type(
|
||||
src_value.type(), concat_global_shape
|
||||
)
|
||||
concat_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
concat_type, src_dist_attr
|
||||
)
|
||||
concat_value.set_type(concat_type)
|
||||
|
||||
dst_value = self.reshard_p_to_s_with_padding(
|
||||
concat_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num,
|
||||
)
|
||||
if permute:
|
||||
dst_value = paddle._C_ops.transpose(dst_value, perm)
|
||||
split_axis = original_split_axis
|
||||
return dst_value
|
||||
|
||||
def reshard_p_to_s_with_padding(
|
||||
self,
|
||||
src_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num=0,
|
||||
):
|
||||
group = new_process_group(
|
||||
sorted(src_dist_attr.process_mesh.process_ids)
|
||||
)
|
||||
dst_value = paddle._C_ops.reduce_scatter(
|
||||
src_value, group.id, len(src_dist_attr.process_mesh.process_ids)
|
||||
)
|
||||
out_global_shape = dst_type.shape
|
||||
out_global_shape[split_axis] = (
|
||||
padding_num + out_global_shape[split_axis]
|
||||
)
|
||||
dst_tmp_type = paddle.pir.create_shaped_type(
|
||||
dst_value.type(), out_global_shape
|
||||
)
|
||||
dst_tmp_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
dst_tmp_type, dst_dist_attr
|
||||
)
|
||||
dst_value.set_type(dst_tmp_type)
|
||||
dst_value.get_defining_op().set_execution_stream(
|
||||
ExecutionStreamType.DefaultStream.value
|
||||
)
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
src_value.get_defining_op().dist_attr.chunk_id,
|
||||
)
|
||||
)
|
||||
if padding_num != 0:
|
||||
if dist.get_rank() == dst_dist_attr.process_mesh.process_ids[-1]:
|
||||
dst_value = paddle._C_ops.split(
|
||||
dst_value,
|
||||
[
|
||||
dst_value.shape[split_axis] - padding_num,
|
||||
padding_num,
|
||||
],
|
||||
0,
|
||||
)[0]
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
[dst_dist_attr],
|
||||
[dst_dist_attr],
|
||||
src_value.get_defining_op().dist_attr.chunk_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
dst_value.set_type(dst_type)
|
||||
return dst_value
|
||||
@@ -0,0 +1,73 @@
|
||||
# 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.
|
||||
|
||||
import paddle
|
||||
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
is_partial,
|
||||
is_replicated,
|
||||
is_shard,
|
||||
)
|
||||
|
||||
|
||||
class RToPReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_replicated(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_partial(dst_dist_attr) or is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
dst_mesh = dst_dist_attr.process_mesh
|
||||
dst_reduce_type = dst_dist_attr.partial_status[0]
|
||||
local_rank = paddle.distributed.get_rank()
|
||||
|
||||
assert dst_reduce_type in [
|
||||
paddle.base.core.ReduceType.kRedSum,
|
||||
paddle.distributed.ReduceType.kRedAvg,
|
||||
paddle.distributed.ReduceType.kRedMax,
|
||||
], f"Unsupported reduce type {dst_reduce_type}"
|
||||
|
||||
if (
|
||||
dst_reduce_type == paddle.distributed.ReduceType.kRedSum
|
||||
and local_rank != 0
|
||||
):
|
||||
dst_value = paddle.full(src_value.shape, 0, dtype=src_value.dtype)
|
||||
else:
|
||||
dst_value = paddle.assign(src_value)
|
||||
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
dst_value.set_type(dst_type)
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_mesh, [src_dist_attr], [dst_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
|
||||
return dst_value
|
||||
@@ -0,0 +1,142 @@
|
||||
# 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.
|
||||
|
||||
import paddle
|
||||
|
||||
from .base_reshard_func import ReshardFunction, is_replicated, is_shard
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
|
||||
|
||||
class RToSReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_replicated(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
split_axis = -1
|
||||
mesh_axis = -1
|
||||
for idx, v in enumerate(dst_dist_attr.dims_mapping):
|
||||
if v != -1:
|
||||
split_axis = idx
|
||||
mesh_axis = v
|
||||
|
||||
mesh = src_dist_attr.process_mesh
|
||||
curr_global_rank = paddle.distributed.get_rank()
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
|
||||
if curr_global_rank in mesh.process_ids:
|
||||
total_nums = src_value.shape[split_axis]
|
||||
num_of_pieces = mesh.shape[mesh_axis]
|
||||
if num_of_pieces == 1:
|
||||
dst_value = paddle._C_ops.share_data_(src_value)
|
||||
share_data_op = dst_value.get_defining_op()
|
||||
# set dist type and dist attr
|
||||
dst_value.set_type(dst_type)
|
||||
share_data_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
return dst_value
|
||||
piece_len = (total_nums + num_of_pieces - 1) // num_of_pieces
|
||||
rank_relative = mesh.process_ids.index(curr_global_rank)
|
||||
start = rank_relative * piece_len
|
||||
end = start + piece_len
|
||||
if curr_global_rank == mesh.process_ids[-1]:
|
||||
end = total_nums
|
||||
|
||||
out_value = paddle.slice(src_value, [split_axis], [start], [end])
|
||||
|
||||
out_value.set_type(dst_type)
|
||||
out_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
mesh, [src_dist_attr], [dst_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
return out_value
|
||||
# fake var will be removed in remove_other_rank_op_pass.
|
||||
fake_var = paddle._C_ops.reshard_v2(src_value, dst_dist_attr)
|
||||
fake_var.set_type(dst_type)
|
||||
return fake_var
|
||||
|
||||
|
||||
class RToSReshardFunctionCrossMesh(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_replicated(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if (
|
||||
in_mesh.ndim != 1
|
||||
or out_mesh.ndim != 1
|
||||
or in_mesh.shape != out_mesh.shape
|
||||
):
|
||||
return False
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
same_status_func = SameStatusReshardFunction()
|
||||
tmp_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
out_value = same_status_func.reshard(
|
||||
src_dist_attr, tmp_dist_attr, src_value, tmp_dst_type
|
||||
)
|
||||
|
||||
if out_value is None:
|
||||
return None
|
||||
|
||||
curr_global_rank = paddle.distributed.get_rank()
|
||||
if curr_global_rank in dst_dist_attr.process_mesh.process_ids:
|
||||
r_to_s_func = RToSReshardFunction()
|
||||
assert r_to_s_func.is_suitable(tmp_dist_attr, dst_dist_attr), (
|
||||
f"Invoke the r to s reshard function is not valid from {tmp_dist_attr} to {dst_dist_attr}"
|
||||
)
|
||||
return r_to_s_func.reshard(
|
||||
tmp_dist_attr, dst_dist_attr, out_value, dst_type
|
||||
)
|
||||
return None
|
||||
@@ -0,0 +1,59 @@
|
||||
# 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 .base_reshard_func import register_reshard_func
|
||||
from .global_to_sub_mesh_func import GlobalToSubMeshFunction
|
||||
from .nd_mesh_reshard_func import (
|
||||
NdMeshReshardFunction,
|
||||
NdMeshReshardFunctionCrossMesh,
|
||||
)
|
||||
from .p_to_r_reshard_func import (
|
||||
PToRReshardFunction,
|
||||
PToRReshardFunctionCrossMesh,
|
||||
)
|
||||
from .p_to_s_reshard_func import (
|
||||
PToSReshardFunction,
|
||||
)
|
||||
from .r_to_p_reshard_func import RToPReshardFunction
|
||||
from .r_to_s_reshard_func import (
|
||||
RToSReshardFunction,
|
||||
RToSReshardFunctionCrossMesh,
|
||||
)
|
||||
from .s_to_r_reshard_func import (
|
||||
SToRReshardFunction,
|
||||
SToRReshardFunctionCrossMesh,
|
||||
)
|
||||
from .s_to_s_reshard_func import SToSReshardFunction
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
from .sub_to_global_mesh_func import SubToGlobalMeshFunction
|
||||
|
||||
|
||||
def register_reshard_funcs():
|
||||
register_reshard_func(PToRReshardFunction())
|
||||
register_reshard_func(PToRReshardFunctionCrossMesh())
|
||||
register_reshard_func(PToSReshardFunction())
|
||||
register_reshard_func(RToSReshardFunction())
|
||||
register_reshard_func(RToSReshardFunctionCrossMesh())
|
||||
register_reshard_func(RToPReshardFunction())
|
||||
register_reshard_func(SameStatusReshardFunction())
|
||||
register_reshard_func(SToRReshardFunction())
|
||||
register_reshard_func(SToRReshardFunctionCrossMesh())
|
||||
register_reshard_func(NdMeshReshardFunction())
|
||||
register_reshard_func(NdMeshReshardFunctionCrossMesh())
|
||||
register_reshard_func(GlobalToSubMeshFunction())
|
||||
register_reshard_func(SubToGlobalMeshFunction())
|
||||
register_reshard_func(SToSReshardFunction())
|
||||
|
||||
|
||||
register_reshard_funcs()
|
||||
@@ -0,0 +1,363 @@
|
||||
# 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.
|
||||
|
||||
|
||||
import paddle
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
copy_op_attr_with_new_member,
|
||||
is_replicated,
|
||||
is_shard,
|
||||
)
|
||||
from .same_status_reshard_func import SameStatusReshardFunction
|
||||
|
||||
|
||||
class SToRReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_shard(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def infer_allgather_dist_type(self, in_value, split_axis):
|
||||
tensor_ndim = len(in_value.shape)
|
||||
in_dist_attr = in_value.dist_attr()
|
||||
split_mesh_dim = in_dist_attr.dims_mapping[split_axis]
|
||||
mesh = in_dist_attr.process_mesh
|
||||
|
||||
# Calculate local shape. In nd_mesh_reshard, multiple tensor axis
|
||||
# may be shard and it will call this 1-D s_to_r function on each
|
||||
# axis. In this case, we should recompute the local and global shape.
|
||||
out_local_shape = list(in_value.shape)
|
||||
out_local_shape[split_axis] = int(
|
||||
(in_value.shape[split_axis] + mesh.shape[split_mesh_dim] - 1)
|
||||
/ mesh.shape[split_mesh_dim]
|
||||
)
|
||||
out_global_shape = list(out_local_shape)
|
||||
out_global_shape[0] *= mesh.shape[split_mesh_dim]
|
||||
out_type = paddle.pir.create_shaped_type(
|
||||
in_value.type(), out_global_shape
|
||||
)
|
||||
|
||||
out_dims_mapping = list(in_dist_attr.dims_mapping)
|
||||
out_dims_mapping[split_axis] = -1
|
||||
out_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
mesh, out_dims_mapping, in_dist_attr.partial_status
|
||||
)
|
||||
out_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
out_type, out_dist_attr
|
||||
)
|
||||
return out_type
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
if src_dist_attr.process_mesh.size == 1:
|
||||
dst_value = paddle._C_ops.share_data_(src_value)
|
||||
share_data_op = dst_value.get_defining_op()
|
||||
# set dist type and dist attr
|
||||
dst_value.set_type(dst_type)
|
||||
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
share_data_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[src_dist_attr],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
return dst_value
|
||||
|
||||
def get_split_axis_with_dims_mapping(dims_mapping):
|
||||
split_axis = {}
|
||||
for idx, v in enumerate(dims_mapping):
|
||||
if v != -1:
|
||||
split_axis[idx] = v
|
||||
return split_axis
|
||||
|
||||
split_axis_map = get_split_axis_with_dims_mapping(
|
||||
src_dist_attr.dims_mapping
|
||||
)
|
||||
|
||||
split_axis = -1
|
||||
for k, v in split_axis_map.items():
|
||||
split_axis = k
|
||||
break
|
||||
num_of_process = src_dist_attr.process_mesh.size
|
||||
num_of_padding = src_value.shape[split_axis] % num_of_process
|
||||
is_balanced_split = num_of_padding == 0
|
||||
|
||||
if is_balanced_split:
|
||||
new_value = self.reshard_s_to_r_with_padding(
|
||||
src_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
num_of_padding,
|
||||
)
|
||||
return new_value
|
||||
else:
|
||||
# find the last one
|
||||
need_padding = (
|
||||
paddle.distributed.get_rank()
|
||||
== src_dist_attr.process_mesh.process_ids[-1]
|
||||
)
|
||||
|
||||
# get padding_num
|
||||
avg_size_on_split_axis = int(
|
||||
(src_value.shape[split_axis] + num_of_process - 1)
|
||||
/ num_of_process
|
||||
)
|
||||
padding_num = (
|
||||
avg_size_on_split_axis * num_of_process
|
||||
- src_value.shape[split_axis]
|
||||
)
|
||||
if need_padding:
|
||||
# set right _local_shape
|
||||
local_shape_at_split_axis = src_value.shape[
|
||||
split_axis
|
||||
] - avg_size_on_split_axis * (num_of_process - 1)
|
||||
local_shape = src_value._local_shape
|
||||
local_shape[split_axis] = local_shape_at_split_axis
|
||||
tmp_src_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), src_dist_attr, list(local_shape)
|
||||
)
|
||||
src_value.set_type(tmp_src_type)
|
||||
padding_shape = src_value._local_shape
|
||||
padding_shape[split_axis] = padding_num
|
||||
padding_tensor = paddle.full(
|
||||
padding_shape,
|
||||
0.0,
|
||||
src_value.dtype,
|
||||
)
|
||||
tmp_src_type1 = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
padding_tensor.type(), dst_dist_attr
|
||||
)
|
||||
padding_tensor.set_type(tmp_src_type1)
|
||||
padding_tensor.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_dist_attr.process_mesh, [], [dst_dist_attr]
|
||||
)
|
||||
)
|
||||
|
||||
concat_value = paddle._C_ops.concat(
|
||||
[src_value, padding_tensor], split_axis
|
||||
)
|
||||
# set concat dist_attr
|
||||
axis_dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
src_dist_attr.process_mesh, [-1], {}
|
||||
)
|
||||
)
|
||||
concat_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_dist_attr.process_mesh,
|
||||
[
|
||||
paddle.base.libpaddle.pir.create_array_attribute(
|
||||
[src_dist_attr, dst_dist_attr]
|
||||
),
|
||||
axis_dist_attr,
|
||||
],
|
||||
[src_dist_attr],
|
||||
)
|
||||
)
|
||||
# set concat_value type
|
||||
concat_global_shape = list(src_value.shape)
|
||||
concat_global_shape[split_axis] = (
|
||||
avg_size_on_split_axis * num_of_process
|
||||
)
|
||||
concat_type = paddle.pir.create_shaped_type(
|
||||
src_value.type(), concat_global_shape
|
||||
)
|
||||
concat_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
concat_type, src_dist_attr
|
||||
)
|
||||
concat_value.set_type(concat_type)
|
||||
|
||||
new_value = self.reshard_s_to_r_with_padding(
|
||||
concat_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num,
|
||||
)
|
||||
return new_value
|
||||
else:
|
||||
new_value = self.reshard_s_to_r_with_padding(
|
||||
src_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num,
|
||||
)
|
||||
return new_value
|
||||
|
||||
def reshard_s_to_r_with_padding(
|
||||
self,
|
||||
src_value,
|
||||
split_axis,
|
||||
src_dist_attr,
|
||||
dst_dist_attr,
|
||||
dst_type,
|
||||
padding_num=0,
|
||||
):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
num_of_process = len(src_mesh.process_ids)
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
|
||||
group = new_process_group(sorted(src_mesh.process_ids))
|
||||
allgather_value = paddle._C_ops.all_gather(
|
||||
src_value, group.id, num_of_process
|
||||
)
|
||||
allgather_type = self.infer_allgather_dist_type(src_value, split_axis)
|
||||
allgather_value.set_type(allgather_type)
|
||||
|
||||
# set op_dist_attr
|
||||
new_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
[-1] * len(dst_dist_attr.dims_mapping),
|
||||
dst_dist_attr.partial_status,
|
||||
)
|
||||
allgather_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [new_dist_attr], chunk_id
|
||||
)
|
||||
)
|
||||
|
||||
if split_axis != 0 or padding_num != 0:
|
||||
allgather_op = allgather_value.get_defining_op()
|
||||
split_values = paddle._C_ops.split_with_num(
|
||||
allgather_op.result(0), num_of_process, 0
|
||||
)
|
||||
builtin_split_op = split_values[0].get_defining_op()
|
||||
pd_split_op = builtin_split_op.operand_source(0).get_defining_op()
|
||||
pd_split_op.dist_attr = copy_op_attr_with_new_member(
|
||||
pd_split_op.dist_attr, new_chunk_id=chunk_id
|
||||
)
|
||||
|
||||
# fix the split_with_num dist attribute.
|
||||
new_inner_types = []
|
||||
for sub_value in split_values:
|
||||
new_inner_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
sub_value.type(), allgather_value.dist_attr()
|
||||
)
|
||||
new_inner_types.append(new_inner_type)
|
||||
sub_value.set_type(new_inner_type)
|
||||
vec_type = paddle.base.libpaddle.pir.create_vec_type(
|
||||
new_inner_types
|
||||
)
|
||||
pd_split_op.result(0).set_type(vec_type)
|
||||
|
||||
if padding_num != 0:
|
||||
tmp_split_values = paddle._C_ops.split(
|
||||
split_values[-1],
|
||||
[
|
||||
split_values[-1].shape[split_axis] - padding_num,
|
||||
padding_num,
|
||||
],
|
||||
split_axis,
|
||||
)
|
||||
split_op = tmp_split_values.get_defining_op()
|
||||
split_op.dist_attr = copy_op_attr_with_new_member(
|
||||
split_op.dist_attr, new_chunk_id=chunk_id
|
||||
)
|
||||
split_values[-1] = tmp_split_values[0]
|
||||
concat_value = paddle._C_ops.concat(split_values, split_axis)
|
||||
concat_op = concat_value.get_defining_op()
|
||||
concat_op.dist_attr = copy_op_attr_with_new_member(
|
||||
concat_op.dist_attr, new_chunk_id=chunk_id
|
||||
)
|
||||
return concat_value
|
||||
else:
|
||||
concat_value = paddle._C_ops.concat(split_values, split_axis)
|
||||
# fold builtin.split op and builtin.combine op
|
||||
concat_op = concat_value.get_defining_op()
|
||||
concat_op.dist_attr = copy_op_attr_with_new_member(
|
||||
concat_op.dist_attr, new_chunk_id=chunk_id
|
||||
)
|
||||
builtin_combine_op = concat_op.operand_source(
|
||||
0
|
||||
).get_defining_op()
|
||||
concat_op.operand(0).set_source(pd_split_op.result(0))
|
||||
builtin_combine_op.erase()
|
||||
builtin_split_op.erase()
|
||||
return concat_value
|
||||
|
||||
return allgather_value
|
||||
|
||||
|
||||
class SToRReshardFunctionCrossMesh(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_shard(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_replicated(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if (
|
||||
in_mesh.ndim != 1
|
||||
or out_mesh.ndim != 1
|
||||
or in_mesh.shape != out_mesh.shape
|
||||
):
|
||||
return False
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
same_status_func = SameStatusReshardFunction()
|
||||
tmp_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
|
||||
dst_dist_attr.process_mesh,
|
||||
src_dist_attr.dims_mapping,
|
||||
src_dist_attr.partial_status,
|
||||
)
|
||||
tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), tmp_dist_attr
|
||||
)
|
||||
out_value = same_status_func.reshard(
|
||||
src_dist_attr, tmp_dist_attr, src_value, tmp_dst_type
|
||||
)
|
||||
|
||||
s_to_r_func = SToRReshardFunction()
|
||||
assert s_to_r_func.is_suitable(tmp_dist_attr, dst_dist_attr), (
|
||||
f"Invoke the p to r reshard function is not valid from {tmp_dist_attr} to {dst_dist_attr}"
|
||||
)
|
||||
return s_to_r_func.reshard(
|
||||
tmp_dist_attr, dst_dist_attr, out_value, dst_type
|
||||
)
|
||||
@@ -0,0 +1,124 @@
|
||||
# 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.
|
||||
|
||||
|
||||
import copy
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.utils.stream_utils import ExecutionStreamType
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import (
|
||||
ReshardFunction,
|
||||
is_shard,
|
||||
)
|
||||
|
||||
|
||||
class SToSReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if not is_shard(src_dist_attr):
|
||||
return False
|
||||
|
||||
if not is_shard(dst_dist_attr):
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh.ndim != 1:
|
||||
return False
|
||||
if out_mesh.ndim != 1:
|
||||
return False
|
||||
if in_mesh != out_mesh:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
"""
|
||||
Reshard from shard to shard status on 1D mesh.
|
||||
E.g. tensor shape: [B, S, H], mesh = [0, 1]
|
||||
1. [Shard(0)] --> [Shard(1)], N ranks:
|
||||
1). reshape from [B, S, H] -> [B, N, S/N, H]
|
||||
2). transpose from [B, N, S/N, H] -> [N, B, S/N, H]
|
||||
3). reshape from [N, B, S/N, H] -> [N*B, S/N, H]
|
||||
4). all to all communicate
|
||||
2. [Shard(1)] --> [Shard(0)], N ranks:
|
||||
1). all to all communicate
|
||||
2). reshape from [B, S, H] -> [N, B/N, S, H]
|
||||
3). transpose from [N, B/N, S, H] -> [B/N, N, S/N, H]
|
||||
4). reshape from [B/N, N, S/N, H] -> [B, S, H]
|
||||
"""
|
||||
in_split_axis = src_dist_attr.dims_mapping.index(0)
|
||||
out_split_axis = dst_dist_attr.dims_mapping.index(0)
|
||||
nranks = len(src_dist_attr.process_mesh.process_ids)
|
||||
|
||||
if out_split_axis != 0:
|
||||
pre_shape = copy.copy(src_value.shape)
|
||||
if pre_shape[out_split_axis] != -1:
|
||||
pre_shape[out_split_axis] = pre_shape[out_split_axis] // nranks
|
||||
pre_shape.insert(out_split_axis, nranks)
|
||||
out_reshape1 = paddle._C_ops.reshape(src_value, pre_shape)
|
||||
|
||||
axes = [out_split_axis]
|
||||
for i in range(len(pre_shape)):
|
||||
if i != out_split_axis:
|
||||
axes.append(i)
|
||||
out_transpose = paddle._C_ops.transpose(out_reshape1, axes)
|
||||
|
||||
pre_shape.pop(out_split_axis)
|
||||
if pre_shape[in_split_axis] != -1:
|
||||
pre_shape[in_split_axis] *= nranks
|
||||
in_all2all = paddle._C_ops.reshape(out_transpose, pre_shape)
|
||||
in_all2all_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), dst_dist_attr
|
||||
)
|
||||
in_all2all.set_type(in_all2all_type)
|
||||
else:
|
||||
in_all2all = paddle._C_ops.share_data_(src_value)
|
||||
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
group = new_process_group(sorted(src_mesh.process_ids))
|
||||
dst_value = paddle._C_ops.all_to_all(in_all2all, group.id)
|
||||
dst_value.get_defining_op().set_execution_stream(
|
||||
ExecutionStreamType.DefaultStream.value
|
||||
)
|
||||
out_all2all_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
in_all2all.type(), src_dist_attr
|
||||
)
|
||||
dst_value.set_type(out_all2all_type)
|
||||
dst_value.get_defining_op().dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [dst_dist_attr], -1
|
||||
)
|
||||
)
|
||||
|
||||
if in_split_axis != 0:
|
||||
post_shape = copy.copy(src_value.shape)
|
||||
if post_shape[0] != -1:
|
||||
post_shape[0] = post_shape[0] // nranks
|
||||
post_shape.insert(0, nranks)
|
||||
dst_value = paddle.reshape(dst_value, post_shape)
|
||||
|
||||
axes = list(range(1, len(post_shape)))
|
||||
axes.insert(in_split_axis, 0)
|
||||
dst_value = paddle._C_ops.transpose(dst_value, axes)
|
||||
|
||||
post_shape.pop(0)
|
||||
if post_shape[in_split_axis] != -1:
|
||||
post_shape[in_split_axis] *= nranks
|
||||
dst_value = paddle._C_ops.reshape(dst_value, post_shape)
|
||||
|
||||
dst_value.set_type(dst_type)
|
||||
|
||||
return dst_value
|
||||
+158
@@ -0,0 +1,158 @@
|
||||
# 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.
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.passes.pass_utils import find_var_used_op_chunk_id
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from .base_reshard_func import ReshardFunction
|
||||
|
||||
|
||||
class SameStatusReshardFunction(ReshardFunction):
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
if src_dist_attr.dims_mapping != dst_dist_attr.dims_mapping:
|
||||
return False
|
||||
if src_dist_attr.partial_dims != dst_dist_attr.partial_dims:
|
||||
return False
|
||||
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
if in_mesh == out_mesh:
|
||||
return False
|
||||
if in_mesh.shape != out_mesh.shape:
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
dst_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
all_process_ids = list(
|
||||
set(src_mesh.process_ids) | set(dst_mesh.process_ids)
|
||||
)
|
||||
all_process_ids = sorted(all_process_ids)
|
||||
|
||||
cur_global_rank = paddle.distributed.get_rank()
|
||||
|
||||
for src, dst in zip(src_mesh.process_ids, dst_mesh.process_ids):
|
||||
if src != dst:
|
||||
new_process_group([src, dst], group_type="p2p")
|
||||
new_process_group([dst, src], group_type="p2p")
|
||||
|
||||
is_send = True
|
||||
for src, dst in zip(src_mesh.process_ids, dst_mesh.process_ids):
|
||||
if src == cur_global_rank:
|
||||
chunk_id = -1
|
||||
if (
|
||||
src_value.get_defining_op().name() == "pd_op.add_n"
|
||||
and src_value.get_defining_op()
|
||||
.operand_source(0)
|
||||
.get_defining_op()
|
||||
.name()
|
||||
== "builtin.combine"
|
||||
):
|
||||
add_n_op = src_value.get_defining_op()
|
||||
combine_op = add_n_op.operand_source(0).get_defining_op()
|
||||
combine_op_chunk_id_list = []
|
||||
for input in combine_op.operands():
|
||||
if input.source().get_defining_op().dist_attr:
|
||||
combine_op_chunk_id_list.append(
|
||||
input.source()
|
||||
.get_defining_op()
|
||||
.dist_attr.chunk_id
|
||||
)
|
||||
else:
|
||||
combine_op_chunk_id_list.append(-1)
|
||||
# check combine_op operands chunk_id equal
|
||||
assert all(
|
||||
x == combine_op_chunk_id_list[0]
|
||||
for x in combine_op_chunk_id_list
|
||||
), "combine_op's operands has different chunk_id."
|
||||
chunk_id = combine_op_chunk_id_list[0]
|
||||
# reset add_n chunk_id
|
||||
add_n_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
add_n_op.dist_attr.process_mesh,
|
||||
add_n_op.dist_attr.operands(),
|
||||
add_n_op.dist_attr.results(),
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = (
|
||||
src_value.get_defining_op().dist_attr.chunk_id
|
||||
)
|
||||
|
||||
comm_group = new_process_group([src, dst], group_type="p2p")
|
||||
paddle._C_ops.send_v2(
|
||||
src_value,
|
||||
comm_group.id,
|
||||
comm_group.ranks.index(dst),
|
||||
True,
|
||||
False,
|
||||
)
|
||||
point = paddle.base.libpaddle.pir.get_current_insertion_point()
|
||||
point.prev()
|
||||
new_op = point.get_operation()
|
||||
assert new_op.name() == "pd_op.send_v2"
|
||||
new_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [], chunk_id
|
||||
)
|
||||
)
|
||||
break
|
||||
|
||||
elif dst == cur_global_rank:
|
||||
all_used_ops = src_value.all_used_ops()
|
||||
chunk_id = -1
|
||||
for used_op in all_used_ops:
|
||||
var = used_op.result(0)
|
||||
if var.dist_attr().process_mesh == dst_mesh:
|
||||
chunk_id = find_var_used_op_chunk_id(var)
|
||||
|
||||
assert -1 not in dst_type.shape, (
|
||||
"dynamic shape is not supported by pir-auto parallel yet."
|
||||
)
|
||||
|
||||
comm_group = new_process_group([src, dst], group_type="p2p")
|
||||
recv_value = paddle._C_ops.recv_v2(
|
||||
dst_type._local_shape,
|
||||
dst_type.dtype,
|
||||
comm_group.ranks.index(src),
|
||||
comm_group.id,
|
||||
True,
|
||||
False,
|
||||
)
|
||||
new_op = recv_value.get_defining_op()
|
||||
new_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_mesh,
|
||||
[],
|
||||
[dst_dist_attr],
|
||||
chunk_id,
|
||||
)
|
||||
)
|
||||
recv_value.set_type(dst_type)
|
||||
is_send = False
|
||||
break
|
||||
|
||||
if is_send:
|
||||
# fake var will be removed in remove_other_rank_op_pass.
|
||||
fake_var = paddle._C_ops.reshard_v2(src_value, dst_dist_attr)
|
||||
fake_var.set_type(dst_type)
|
||||
return fake_var
|
||||
else:
|
||||
return recv_value
|
||||
+171
@@ -0,0 +1,171 @@
|
||||
# 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.
|
||||
|
||||
import copy
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed.auto_parallel.placement_type import (
|
||||
check_placements_equal,
|
||||
)
|
||||
|
||||
from ..process_group import new_process_group
|
||||
from ..utils import split_mesh
|
||||
from .base_reshard_func import ReshardFunction, copy_dist_attr_with_new_member
|
||||
|
||||
|
||||
def _mesh_equal_ignore_shape_one(mesh1, mesh2, dim: int):
|
||||
"""
|
||||
Check if two process meshes are equal, ignoring the shape value `1`
|
||||
in the specified dimension. This is used when mesh1 is a sub-mesh
|
||||
split from a global mesh, in this case, the shape of mesh1 is `1`
|
||||
in the split dim.
|
||||
E.g, the following two meshes are equal:
|
||||
mesh1: shape = [1,2,2], process_ids = [0,1,2,3]
|
||||
mesh2: shape = [2,2], process_ids = [0,1,2,3]
|
||||
"""
|
||||
assert dim >= 0 and dim < len(mesh1.shape), "invalid dim arg"
|
||||
if mesh1 == mesh2:
|
||||
return True
|
||||
|
||||
if mesh1.process_ids != mesh2.process_ids:
|
||||
return False
|
||||
|
||||
a_shape = copy.copy(mesh1.shape)
|
||||
b_shape = copy.copy(mesh2.shape)
|
||||
|
||||
if a_shape[dim] != 1:
|
||||
return False
|
||||
a_shape.pop(dim)
|
||||
|
||||
return a_shape == b_shape
|
||||
|
||||
|
||||
class SubToGlobalMeshFunction(ReshardFunction):
|
||||
"""
|
||||
Reshard from sub-mesh to global mesh, now only supports
|
||||
both input and output values are replicated, e.g.
|
||||
1. input: mesh:[0], placements:[Replicate()]
|
||||
output: mesh:[0,1], placements:[Replicate()]
|
||||
2. input: mesh:[0,1], placements:[Replicate()]
|
||||
output: mesh:[[0,1],[2,3]], placements:[Replicate(), Replicate()]
|
||||
"""
|
||||
|
||||
def is_suitable(self, src_dist_attr, dst_dist_attr):
|
||||
in_mesh = src_dist_attr.process_mesh
|
||||
out_mesh = dst_dist_attr.process_mesh
|
||||
sub_mesh_dim = paddle.base.core.sub_mesh_dim(out_mesh, in_mesh)
|
||||
if sub_mesh_dim == -1:
|
||||
return False
|
||||
sub_meshes, sub_placements = (
|
||||
dist.auto_parallel.api._get_sub_meshes_and_local_placements(
|
||||
out_mesh, dst_dist_attr.placements_attr, sub_mesh_dim
|
||||
)
|
||||
)
|
||||
if not check_placements_equal(
|
||||
src_dist_attr.placements_attr, sub_placements
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
|
||||
src_mesh = src_dist_attr.process_mesh
|
||||
dst_mesh = dst_dist_attr.process_mesh
|
||||
|
||||
sub_mesh_dim = paddle.base.core.sub_mesh_dim(dst_mesh, src_mesh)
|
||||
sub_meshes = split_mesh(dst_mesh, sub_mesh_dim)
|
||||
dst_meshes = [
|
||||
mesh
|
||||
for mesh in sub_meshes
|
||||
if not _mesh_equal_ignore_shape_one(mesh, src_mesh, sub_mesh_dim)
|
||||
]
|
||||
|
||||
comm_group_ids = []
|
||||
root_ranks = []
|
||||
for p_id in src_mesh.process_ids:
|
||||
comm_group_ids.append([p_id])
|
||||
root_ranks.append(p_id)
|
||||
for i, group_ids in enumerate(comm_group_ids):
|
||||
for mesh in dst_meshes:
|
||||
group_ids.append(mesh.process_ids[i])
|
||||
|
||||
other_ranks = copy.copy(dst_mesh.process_ids)
|
||||
for rank in other_ranks:
|
||||
if rank in src_mesh.process_ids:
|
||||
other_ranks.remove(rank)
|
||||
|
||||
cur_rank = paddle.distributed.get_rank()
|
||||
|
||||
if cur_rank in src_mesh.process_ids:
|
||||
# the root rank will broadcast the src_value to other ranks
|
||||
chunk_id = -1
|
||||
if src_value.get_defining_op().dist_attr:
|
||||
chunk_id = src_value.get_defining_op().dist_attr.chunk_id
|
||||
tmp_value = paddle._C_ops.share_data_(src_value)
|
||||
value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
src_value.type(), src_value.dist_attr()
|
||||
)
|
||||
tmp_value.set_type(value_type)
|
||||
op = tmp_value.get_defining_op()
|
||||
op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
src_mesh, [src_dist_attr], [src_dist_attr], chunk_id
|
||||
)
|
||||
elif cur_rank in other_ranks:
|
||||
# create the buffer on other ranks for receiving the data
|
||||
tmp_value = paddle.zeros(dst_type.shape, dst_type.dtype)
|
||||
op = tmp_value.get_defining_op()
|
||||
mesh = paddle.distributed.ProcessMesh(other_ranks)
|
||||
value_dist_attr = copy_dist_attr_with_new_member(
|
||||
dst_dist_attr, new_process_mesh=mesh
|
||||
)
|
||||
value_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
|
||||
dst_type, value_dist_attr
|
||||
)
|
||||
tmp_value.set_type(value_type)
|
||||
op.dist_attr = paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
mesh, [], [value_dist_attr]
|
||||
)
|
||||
else:
|
||||
# do nothing if the current rank is not in src_mesh and dst_mesh.
|
||||
# use reshard_op to create and return a fake value, and the fake
|
||||
# value will be removed 'remove_other_rank_op_pass'.
|
||||
fake_var = paddle._C_ops.reshard_v2(src_value, dst_dist_attr)
|
||||
return fake_var
|
||||
|
||||
# create communication groups
|
||||
for i, group_ids in enumerate(comm_group_ids):
|
||||
comm_group_ids[i] = sorted(group_ids)
|
||||
# the root arg in broadcast is the local index
|
||||
# of the rank in the communication group
|
||||
root_ranks[i] = comm_group_ids[i].index(root_ranks[i])
|
||||
|
||||
comm_groups = []
|
||||
for i, group_ids in enumerate(comm_group_ids):
|
||||
comm_groups.append(new_process_group(group_ids))
|
||||
if cur_rank in group_ids:
|
||||
cur_group_id = i
|
||||
|
||||
broadcast_value = paddle._C_ops.broadcast(
|
||||
tmp_value, comm_groups[cur_group_id].id, root_ranks[cur_group_id]
|
||||
)
|
||||
broadcast_value.set_type(dst_type)
|
||||
|
||||
broadcast_op = broadcast_value.get_defining_op()
|
||||
broadcast_op.dist_attr = (
|
||||
paddle.base.libpaddle.pir.create_op_dist_attribute(
|
||||
dst_mesh, [src_dist_attr], [dst_dist_attr]
|
||||
)
|
||||
)
|
||||
|
||||
return broadcast_value
|
||||
Reference in New Issue
Block a user