wehub-resource-sync
3305 lines
130 KiB
Python
3305 lines
130 KiB
Python
# Copyright (c) 2021 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 operator
|
|
from collections import OrderedDict
|
|
from functools import reduce
|
|
|
|
import paddle
|
|
from paddle.distributed.fleet.meta_optimizers.common import OpRole
|
|
from paddle.distributed.utils.stream_utils import ExecutionStreamType
|
|
from paddle.framework import LayerHelper, OpProtoHolder, Program, core
|
|
from paddle.utils import unique_name
|
|
|
|
from .cost import (
|
|
AllgatherOpCost,
|
|
CommContext,
|
|
ConcatOpCost,
|
|
SendOpCost,
|
|
SliceOpCost,
|
|
SplitOpCost,
|
|
build_comm_desc,
|
|
)
|
|
from .dist_context import DistributedContext
|
|
from .process_group import new_process_group
|
|
from .utils import (
|
|
_g_gradient_clip_ops,
|
|
is_gradient_clip_op,
|
|
is_optimize_op,
|
|
is_reshard_op,
|
|
naive_set_dist_op_attr_for_program_by_mesh,
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping,
|
|
set_var_dist_attr,
|
|
)
|
|
|
|
# NOTE: If op in _g_special_ops or _g_gradient_clip_ops, it will not be resharded.
|
|
_g_special_ops = ['check_finite_and_unscale', 'update_loss_scaling']
|
|
_g_subblock_ops = ["while", "conditional_block"]
|
|
|
|
|
|
def get_var_with_recursion(var_name, block, program):
|
|
"""Get var in the parent block if not found in the current block"""
|
|
var = None
|
|
if var_name in block.vars:
|
|
var = block.vars[var_name]
|
|
else:
|
|
var = block._var_recursive(var_name)
|
|
assert var is not None, f"{var.name} is not found"
|
|
|
|
return var
|
|
|
|
|
|
class EndOpDesc:
|
|
"""
|
|
Describe to end reshard parse process.
|
|
It is supposed to contain a list of variables which are the outputs of one reshard process.
|
|
|
|
Args:
|
|
vars (list): a list of variables.
|
|
"""
|
|
|
|
def __init__(self, vars):
|
|
self._vars = vars
|
|
|
|
@property
|
|
def vars(self):
|
|
return self._vars
|
|
|
|
def __repr__(self):
|
|
return f"End vars : {self._vars}."
|
|
|
|
|
|
class AllGatherOpDesc:
|
|
"""
|
|
Describe the allgather op in the reshard phase.
|
|
|
|
Args:
|
|
group (list): Process group.
|
|
shape (list): The tensor shape.
|
|
is_bool (bool): Whether allgather bool data. Default: False.
|
|
"""
|
|
|
|
def __init__(self, group, shape, is_bool=False, need_split=True):
|
|
self._group = group
|
|
self._desc = "all_gather"
|
|
self._shape = shape
|
|
self._is_bool = is_bool
|
|
self._need_split = need_split
|
|
|
|
@property
|
|
def is_bool(self):
|
|
return self._is_bool
|
|
|
|
@property
|
|
def group(self):
|
|
return self._group
|
|
|
|
@property
|
|
def desc(self):
|
|
return self._desc
|
|
|
|
@property
|
|
def shape(self):
|
|
return self._shape
|
|
|
|
@property
|
|
def need_split(self):
|
|
return self._need_split
|
|
|
|
def __repr__(self):
|
|
return f"op: {self._desc}, group: {self._group}, shape: {self._shape}, is_bool: {self._is_bool}, need_split: {self._need_split}."
|
|
|
|
|
|
class AllGatherConcatOpDesc:
|
|
"""
|
|
Describe the c_concat op in the reshard phase.
|
|
|
|
Args:
|
|
group (list): Process group.
|
|
shape (list): The tensor shape.
|
|
is_bool (bool): Whether c_concat bool data. Default: False.
|
|
"""
|
|
|
|
def __init__(self, group, shape, is_bool=False):
|
|
self._group = group
|
|
self._desc = "c_concat"
|
|
self._shape = shape
|
|
self._is_bool = is_bool
|
|
|
|
@property
|
|
def is_bool(self):
|
|
return self._is_bool
|
|
|
|
@property
|
|
def group(self):
|
|
return self._group
|
|
|
|
@property
|
|
def desc(self):
|
|
return self._desc
|
|
|
|
@property
|
|
def shape(self):
|
|
return self._shape
|
|
|
|
def __repr__(self):
|
|
return f"op: {self._desc}, group: {self._group}, shape: {self._shape}, is_bool: {self._is_bool}."
|
|
|
|
|
|
class SendOpDesc:
|
|
"""
|
|
Describe the send op in the reshard phase.
|
|
|
|
Args:
|
|
partition_index (list): The index of partition in complete tensor.
|
|
src (int): The source process to send.
|
|
dst (int): The destination process to receive.
|
|
is_bool (bool): Whether send bool data. Default: False.
|
|
"""
|
|
|
|
def __init__(self, partition_index, src, dst, is_bool=False):
|
|
self._dst = dst
|
|
self._partition_index = partition_index
|
|
self._desc = "send"
|
|
self._shape = []
|
|
self._is_bool = is_bool
|
|
self._src = src
|
|
|
|
@property
|
|
def src(self):
|
|
return self._src
|
|
|
|
@property
|
|
def is_bool(self):
|
|
return self._is_bool
|
|
|
|
@property
|
|
def partition_index(self):
|
|
return self._partition_index
|
|
|
|
@property
|
|
def dst(self):
|
|
return self._dst
|
|
|
|
@property
|
|
def desc(self):
|
|
return self._desc
|
|
|
|
@property
|
|
def shape(self):
|
|
if not self._shape:
|
|
for item in self.partition_index:
|
|
self._shape.append(item[1] - item[0])
|
|
return self._shape
|
|
|
|
def __repr__(self):
|
|
return f"op: {self._desc}, partition_index: {self._partition_index}, dst: {self._dst}, shape: {self._shape}, is_bool: {self._is_bool}."
|
|
|
|
|
|
class RecvOpDesc:
|
|
"""
|
|
Describe the recv op in the reshard op.
|
|
|
|
Args:
|
|
partition_index (list): The index of partition in complete tensor.
|
|
src (int): The source process to send.
|
|
dst (int): The destination process to receive.
|
|
is_bool (bool): Whether receive bool data. Default: False.
|
|
"""
|
|
|
|
def __init__(self, partition_index, src, dst, is_bool=False):
|
|
self._src = src
|
|
self._partition_index = partition_index
|
|
self._desc = "recv"
|
|
self._shape = []
|
|
self._is_bool = is_bool
|
|
self._dst = dst
|
|
|
|
@property
|
|
def dst(self):
|
|
return self._dst
|
|
|
|
@property
|
|
def is_bool(self):
|
|
return self._is_bool
|
|
|
|
@property
|
|
def partition_index(self):
|
|
return self._partition_index
|
|
|
|
@property
|
|
def src(self):
|
|
return self._src
|
|
|
|
@property
|
|
def desc(self):
|
|
return self._desc
|
|
|
|
@property
|
|
def shape(self):
|
|
if not self._shape:
|
|
for item in self.partition_index:
|
|
self._shape.append(item[1] - item[0])
|
|
return self._shape
|
|
|
|
def __repr__(self):
|
|
return f"op: {self._desc}, partition_index: {self._partition_index}, dst: {self._dst}, shape: {self._shape}, is_bool: {self._is_bool}."
|
|
|
|
|
|
class SliceOpDesc:
|
|
"""
|
|
Describe the slice op in the reshard phase.
|
|
|
|
Args:
|
|
starts (list): It represents start indices of corresponding axis in ``axes``.
|
|
ends (list): It represents end indices of corresponding axis in ``axes``.
|
|
axes (list): Axes that `starts` and `ends` apply to.
|
|
shape (list): The shape of the tensor to be sliced.
|
|
"""
|
|
|
|
def __init__(self, starts, ends, axes, shape=None):
|
|
self._starts = starts
|
|
self._ends = ends
|
|
self._axes = axes
|
|
self._desc = "slice"
|
|
self._shape = shape
|
|
|
|
@property
|
|
def starts(self):
|
|
return self._starts
|
|
|
|
@property
|
|
def ends(self):
|
|
return self._ends
|
|
|
|
@property
|
|
def axes(self):
|
|
return self._axes
|
|
|
|
@property
|
|
def desc(self):
|
|
return self._desc
|
|
|
|
@property
|
|
def shape(self):
|
|
return self._shape
|
|
|
|
def __repr__(self):
|
|
if self._shape is not None:
|
|
return f"op: {self._desc}, starts: {self._starts}, ends: {self._ends}, axes: {self._axes}, shape: {self._shape}."
|
|
else:
|
|
return f"op: {self._desc}, starts: {self._starts}, ends: {self._ends}, axes: {self._axes}."
|
|
|
|
|
|
class ConcatOpDesc:
|
|
"""
|
|
Describe the concat op in the reshard phase.
|
|
|
|
Args:
|
|
partition_index_list (list): The list contains all partition index.
|
|
"""
|
|
|
|
def __init__(self, partition_index_list):
|
|
self._partition_index_list = partition_index_list
|
|
self._desc = "concat"
|
|
|
|
@property
|
|
def partition_index_list(self):
|
|
return self._partition_index_list
|
|
|
|
@property
|
|
def desc(self):
|
|
return self._desc
|
|
|
|
def __repr__(self):
|
|
return f"op: {self._desc}, partition_index_list: {self._partition_index_list}."
|
|
|
|
|
|
class Inserter:
|
|
"""Insert op required in the reshard process."""
|
|
|
|
@staticmethod
|
|
def insert_cast_op(block, idx, tensor, op_role, tensor_type, sync=True):
|
|
# to avoid name conflict with framework
|
|
new_var_name = paddle.utils.unique_name.generate_with_ignorable_key(
|
|
".".join(["cast@RESHARD", 'tmp'])
|
|
)
|
|
out = block.create_var(
|
|
name=new_var_name,
|
|
dtype=tensor_type,
|
|
type=tensor.type,
|
|
lod_level=tensor.lod_level,
|
|
)
|
|
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
cast_op = insert_operation(
|
|
idx,
|
|
type='cast',
|
|
inputs={'X': [tensor]},
|
|
outputs={'Out': [out]},
|
|
attrs={
|
|
'in_dtype': tensor.dtype,
|
|
'out_dtype': out.dtype,
|
|
'op_role': op_role,
|
|
},
|
|
)
|
|
cast_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return out
|
|
|
|
@staticmethod
|
|
def insert_send_op(block, idx, tensor, src, dst, op_role, sync=True):
|
|
"""Insert send op into block at the given index."""
|
|
op_type = 'send_v2'
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
# use pair comm group
|
|
process_group = new_process_group([src, dst], group_type='p2p')
|
|
send_op = insert_operation(
|
|
idx,
|
|
type=op_type,
|
|
inputs={'X': [tensor]},
|
|
attrs={
|
|
'ring_id': process_group.id,
|
|
'peer': process_group.ranks.index(dst),
|
|
'use_calc_stream': True,
|
|
'op_role': op_role,
|
|
'dynamic_shape': True,
|
|
},
|
|
)
|
|
send_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
|
|
@staticmethod
|
|
def insert_recv_op(block, idx, tensor, src, dst, op_role, sync=True):
|
|
"""Insert recv op into block at the given index."""
|
|
op_type = 'recv_v2'
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
# use pair group
|
|
process_group = new_process_group([src, dst], group_type='p2p')
|
|
recv_op = insert_operation(
|
|
idx,
|
|
type=op_type,
|
|
inputs={'X': [tensor]},
|
|
outputs={'Out': [tensor]},
|
|
attrs={
|
|
'ring_id': process_group.id,
|
|
'peer': process_group.ranks.index(src),
|
|
'out_shape': tensor.shape,
|
|
'dtype': tensor.dtype,
|
|
'use_calc_stream': True,
|
|
'op_role': op_role,
|
|
'dynamic_shape': True,
|
|
},
|
|
)
|
|
recv_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
|
|
@staticmethod
|
|
def insert_reset_lod_op(block, idx, X, Y, op_role, sync=True):
|
|
"""Insert reset_lod op into block at the given index."""
|
|
|
|
new_var_name = paddle.utils.unique_name.generate_with_ignorable_key(
|
|
".".join(["reset_lod@RESHARD", 'tmp'])
|
|
)
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
|
|
reset_lod_out = block.create_var(
|
|
name=new_var_name,
|
|
shape=X.shape,
|
|
type=X.type,
|
|
dtype=X.dtype,
|
|
lod_level=X.lod_level,
|
|
)
|
|
|
|
reset_op = insert_operation(
|
|
idx,
|
|
type="lod_reset",
|
|
inputs={'X': X, 'Y': Y},
|
|
outputs={'Out': reset_lod_out},
|
|
attrs={'op_role': op_role},
|
|
)
|
|
reset_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return reset_lod_out
|
|
|
|
@staticmethod
|
|
def insert_concat_op(block, idx, tensors, axis, op_role, sync=True):
|
|
"""Insert concat op into block at the given block."""
|
|
inputs = {'X': tensors}
|
|
attrs = {}
|
|
attrs['axis'] = axis
|
|
attrs['op_role'] = op_role
|
|
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
# to avoid name conflict with framework
|
|
helper = LayerHelper('concat@RESHARD', **locals())
|
|
with paddle.static.program_guard(block.program):
|
|
out = block.create_var(
|
|
name=paddle.utils.unique_name.generate_with_ignorable_key(
|
|
".".join([helper.name, 'tmp'])
|
|
),
|
|
dtype=tensors[0].dtype,
|
|
shape=None,
|
|
lod_level=tensors[0].lod_level,
|
|
type=tensors[0].type,
|
|
persistable=False,
|
|
stop_gradient=False,
|
|
)
|
|
concat_op = insert_operation(
|
|
idx,
|
|
type='concat',
|
|
inputs=inputs,
|
|
outputs={'Out': [out]},
|
|
attrs=attrs,
|
|
)
|
|
concat_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return out
|
|
|
|
@staticmethod
|
|
def insert_slice_op(
|
|
block, idx, tensor, starts, ends, axes, new_var_name, op_role, sync=True
|
|
):
|
|
"""Insert slice op into block at the given block."""
|
|
# This is a hack to insert split op to get slice tensor
|
|
# 1. [128, 128] => [64, 128]: split
|
|
# 2. [128, 128] => [128, 128]: assign
|
|
# 3. [128, 128] => [64, 64]: slice, it will replaced by multi split
|
|
global_shape = tensor.shape
|
|
slice_shape = [ends[i] - starts[i] for i in range(len(starts))]
|
|
diff_dims = []
|
|
for index, item in enumerate(slice_shape):
|
|
if item != global_shape[index]:
|
|
diff_dims.append(index)
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
|
|
# use assign
|
|
if len(diff_dims) == 0:
|
|
out = block.create_var(
|
|
name=new_var_name,
|
|
dtype=tensor.dtype,
|
|
type=tensor.type,
|
|
shape=slice_shape,
|
|
lod_level=tensor.lod_level,
|
|
)
|
|
inputs = {'X': [tensor]}
|
|
outputs = {"Out": [out]}
|
|
attrs = {"in_place": False, "op_role": op_role}
|
|
assign_op = insert_operation(
|
|
idx, type="assign", inputs=inputs, outputs=outputs, attrs=attrs
|
|
)
|
|
assign_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return out
|
|
|
|
# use split once
|
|
elif len(diff_dims) == 1:
|
|
diff_dim = diff_dims[0]
|
|
num_or_sections = global_shape[diff_dim] // slice_shape[diff_dim]
|
|
axis = diff_dim
|
|
cur_idx = starts[diff_dim] // slice_shape[diff_dim]
|
|
input_shape = global_shape
|
|
inputs = {'X': tensor}
|
|
attrs = {'num': num_or_sections, 'axis': axis, 'op_role': op_role}
|
|
new_shape = []
|
|
for index, item in enumerate(tensor.shape):
|
|
if index != axis:
|
|
new_shape.append(item)
|
|
else:
|
|
new_shape.append(item // num_or_sections)
|
|
with paddle.static.program_guard(block.program):
|
|
outs = [
|
|
block.create_var(
|
|
name=paddle.utils.unique_name.generate_with_ignorable_key(
|
|
".".join(['split@RESHARD', 'tmp'])
|
|
),
|
|
dtype=tensor.dtype,
|
|
shape=None,
|
|
type=tensor.type,
|
|
persistable=False,
|
|
lod_level=tensor.lod_level,
|
|
stop_gradient=False,
|
|
)
|
|
for i in range(num_or_sections)
|
|
]
|
|
out = outs[cur_idx]
|
|
split_op = insert_operation(
|
|
idx,
|
|
type="split",
|
|
inputs=inputs,
|
|
outputs={'Out': outs},
|
|
attrs=attrs,
|
|
)
|
|
split_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return out
|
|
|
|
# use slice
|
|
else:
|
|
inputs = {'Input': tensor}
|
|
infer_flags = [1 for i in range(len(axes))]
|
|
attrs = {
|
|
"axes": axes,
|
|
"starts": starts,
|
|
"ends": ends,
|
|
"infer_flags": infer_flags,
|
|
'op_role': op_role,
|
|
}
|
|
out = block.create_var(
|
|
name=new_var_name,
|
|
dtype=tensor.dtype,
|
|
type=tensor.type,
|
|
lod_level=tensor.lod_level,
|
|
)
|
|
slice_op = insert_operation(
|
|
idx,
|
|
type="slice",
|
|
inputs=inputs,
|
|
outputs={'Out': [out]},
|
|
attrs=attrs,
|
|
)
|
|
slice_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return out
|
|
|
|
@staticmethod
|
|
def insert_split_op(
|
|
block, idx, tensor, num_or_sections, op_role, axis=0, sync=True
|
|
):
|
|
"""Insert split op into block at the given index."""
|
|
helper = LayerHelper('split@RESHARD', **locals())
|
|
input_shape = tensor.shape
|
|
inputs = {'X': tensor}
|
|
attrs = {'num': num_or_sections, 'axis': axis, 'op_role': op_role}
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
|
|
new_shape = []
|
|
for index, item in enumerate(tensor.shape):
|
|
if index != axis:
|
|
new_shape.append(item)
|
|
else:
|
|
new_shape.append(item // num_or_sections)
|
|
with paddle.static.program_guard(block.program):
|
|
outs = [
|
|
block.create_var(
|
|
name=paddle.utils.unique_name.generate_with_ignorable_key(
|
|
".".join([helper.name, 'tmp'])
|
|
),
|
|
dtype=tensor.dtype,
|
|
shape=None,
|
|
lod_level=tensor.lod_level,
|
|
type=tensor.type,
|
|
persistable=False,
|
|
stop_gradient=False,
|
|
)
|
|
for i in range(num_or_sections)
|
|
]
|
|
split_op = insert_operation(
|
|
idx, type="split", inputs=inputs, outputs={'Out': outs}, attrs=attrs
|
|
)
|
|
split_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return outs
|
|
|
|
@staticmethod
|
|
def insert_fill_constant_op(block, idx, op_role, shape, sync=True):
|
|
"""Insert fill constant op into block at the given index."""
|
|
# to avoid name conflict with framework
|
|
helper = LayerHelper('fill_constant@RESHARD', **locals())
|
|
# use paddle.int64 as dtype
|
|
with paddle.static.program_guard(block.program):
|
|
out = block.create_var(
|
|
name=paddle.utils.unique_name.generate_with_ignorable_key(
|
|
".".join([helper.name, 'tmp'])
|
|
),
|
|
dtype=paddle.int64,
|
|
shape=None,
|
|
type=core.VarDesc.VarType.DENSE_TENSOR,
|
|
persistable=False,
|
|
stop_gradient=False,
|
|
)
|
|
inputs = {}
|
|
attrs = {'force_cpu': False}
|
|
attrs['str_value'] = str(int("1"))
|
|
attrs['value'] = int("1")
|
|
attrs['dtype'] = out.dtype
|
|
attrs['op_role'] = op_role
|
|
paddle.utils.get_shape_tensor_inputs(
|
|
inputs=inputs, attrs=attrs, shape=shape, op_type='fill_constant'
|
|
)
|
|
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
fillconstant_op = insert_operation(
|
|
idx,
|
|
type='fill_constant',
|
|
inputs=inputs,
|
|
outputs={'Out': [out]},
|
|
attrs=attrs,
|
|
)
|
|
out.stop_gradient = True
|
|
fillconstant_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return out
|
|
|
|
@staticmethod
|
|
def insert_allgather_op(
|
|
block, idx, tensor, ranks, op_role, need_split, sync=True
|
|
):
|
|
"""Insert allgather op into block at the given index."""
|
|
tensor_list = []
|
|
group = new_process_group(ranks)
|
|
idx_offset = 0
|
|
|
|
# insert all_gather op
|
|
op_type = 'all_gather'
|
|
# to avoid name conflict with framework
|
|
helper = LayerHelper(op_type + "@RESHARD", **locals())
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
|
|
with paddle.static.program_guard(block.program):
|
|
allgather_out = block.create_var(
|
|
name=paddle.utils.unique_name.generate_with_ignorable_key(
|
|
".".join([helper.name, 'tmp'])
|
|
),
|
|
dtype=tensor.dtype,
|
|
shape=None,
|
|
lod_level=tensor.lod_level,
|
|
type=tensor.type,
|
|
persistable=False,
|
|
stop_gradient=False,
|
|
)
|
|
allgather_op = insert_operation(
|
|
idx + idx_offset,
|
|
type=op_type,
|
|
inputs={'x': [tensor]},
|
|
outputs={'out': [allgather_out]},
|
|
attrs={
|
|
'ring_id': group.id,
|
|
'nranks': group.nranks,
|
|
'op_role': op_role,
|
|
},
|
|
)
|
|
allgather_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
allgather_op.dist_attr.execution_stream = (
|
|
ExecutionStreamType.DefaultStream.value
|
|
)
|
|
idx_offset += 1
|
|
|
|
# insert split op
|
|
if need_split:
|
|
split_out = Inserter.insert_split_op(
|
|
block,
|
|
idx + idx_offset,
|
|
allgather_out,
|
|
group.nranks,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
idx_offset += 1
|
|
tensor_list.extend(split_out)
|
|
else:
|
|
tensor_list.extend([allgather_out])
|
|
return tensor_list, idx_offset
|
|
|
|
@staticmethod
|
|
def insert_c_concat_op(block, idx, tensor, ranks, op_role, sync=True):
|
|
"""Insert c_concat op into block at the given index."""
|
|
group = new_process_group(ranks)
|
|
idx_offset = 0
|
|
insert_operation = (
|
|
block._insert_op if sync else block._insert_op_without_sync
|
|
)
|
|
|
|
# insert c_concat op
|
|
op_type = 'c_concat'
|
|
# to avoid name conflict with framework
|
|
helper = LayerHelper(op_type + "@RESHARD", **locals())
|
|
with paddle.static.program_guard(block.program):
|
|
c_concat_out = block.create_var(
|
|
name=paddle.utils.unique_name.generate_with_ignorable_key(
|
|
".".join([helper.name, 'tmp'])
|
|
),
|
|
dtype=tensor.dtype,
|
|
shape=None,
|
|
lod_level=tensor.lod_level,
|
|
type=tensor.type,
|
|
persistable=False,
|
|
stop_gradient=False,
|
|
)
|
|
cur_rank = paddle.distributed.get_rank()
|
|
c_concat_op = insert_operation(
|
|
idx + idx_offset,
|
|
type=op_type,
|
|
inputs={'X': [tensor]},
|
|
outputs={'Out': [c_concat_out]},
|
|
attrs={
|
|
'ring_id': group.id,
|
|
'use_calc_stream': True,
|
|
'use_model_parallel': True,
|
|
'nranks': group.nranks,
|
|
'op_role': op_role,
|
|
'rank': group.ranks.index(cur_rank) if cur_rank in ranks else 0,
|
|
},
|
|
)
|
|
c_concat_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
return c_concat_out
|
|
|
|
@staticmethod
|
|
def concat_partitions_with_op(
|
|
partition_tensor_list,
|
|
tensor,
|
|
partition_index,
|
|
block,
|
|
idx,
|
|
op_role,
|
|
sync=True,
|
|
):
|
|
"""Concat the tensors and insert concat op."""
|
|
if not partition_tensor_list:
|
|
partition_tensor_list.append((tensor, partition_index))
|
|
else:
|
|
i = 0
|
|
has_concat = False
|
|
while i < len(partition_tensor_list):
|
|
(
|
|
concat_axis,
|
|
first_order,
|
|
new_partition,
|
|
) = Resharder.compute_concat_info(
|
|
partition_tensor_list[i][1], partition_index
|
|
)
|
|
if concat_axis != -1:
|
|
has_concat = True
|
|
_ = (
|
|
Inserter.insert_concat_op(
|
|
block,
|
|
idx[0],
|
|
[partition_tensor_list[i][0], tensor],
|
|
concat_axis,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
if first_order == 0
|
|
else Inserter.insert_concat_op(
|
|
block,
|
|
idx[0],
|
|
[tensor, partition_tensor_list[i][0]],
|
|
concat_axis,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
)
|
|
partition_tensor_list.pop(i)
|
|
idx[0] += 1
|
|
Inserter.concat_partitions_with_op(
|
|
partition_tensor_list,
|
|
_,
|
|
new_partition,
|
|
block,
|
|
idx,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
break
|
|
i += 1
|
|
if not has_concat:
|
|
partition_tensor_list.append((tensor, partition_index))
|
|
|
|
|
|
class Remover:
|
|
"""Remove var and op in the reshard process."""
|
|
|
|
@staticmethod
|
|
def remove_no_need_ops(auto_parallel_main_prog, dist_context, rank_id):
|
|
"""Remove no need ops in the main program"""
|
|
not_remove_op_ref = [
|
|
"create_py_reader",
|
|
"create_double_buffer_reader",
|
|
"read",
|
|
]
|
|
|
|
# NOTE: The nested sub block is not be supported now.
|
|
remove_block_order = []
|
|
for block_idx in Resharder.while_block_info:
|
|
remove_block_order.append(block_idx)
|
|
|
|
for block_idx, block in enumerate(auto_parallel_main_prog.blocks):
|
|
if block_idx not in remove_block_order:
|
|
remove_block_order.append(block_idx)
|
|
|
|
# the sub block should be removed first
|
|
for block_idx in remove_block_order:
|
|
remove_op_idx = []
|
|
block = auto_parallel_main_prog.blocks[block_idx]
|
|
ops = block.ops
|
|
vars = block.vars
|
|
for idx, op in enumerate(ops):
|
|
if op.type == "read":
|
|
dim_list = []
|
|
for var_name in op.output_arg_names:
|
|
dim_list.extend(
|
|
get_var_with_recursion(
|
|
var_name, block, auto_parallel_main_prog
|
|
).shape
|
|
)
|
|
for i in range(idx, -1, -1):
|
|
if ops[i].type == "create_py_reader":
|
|
ops[i]._set_attr("shape_concat", dim_list)
|
|
break
|
|
continue
|
|
|
|
# replace the input and output of c_sync_comm_stream op when in pipeline scene.
|
|
if op.type == "c_sync_comm_stream":
|
|
need_save = []
|
|
for var_name in op.input_arg_names:
|
|
process_mesh = (
|
|
dist_context.get_tensor_dist_attr_for_program(
|
|
get_var_with_recursion(
|
|
var_name, block, auto_parallel_main_prog
|
|
)
|
|
).process_mesh
|
|
)
|
|
if rank_id in process_mesh.process_ids:
|
|
need_save.append(var_name)
|
|
if not need_save:
|
|
remove_op_idx.append(idx)
|
|
continue
|
|
|
|
proto = OpProtoHolder.instance().get_op_proto(op.type)
|
|
op.desc.set_input(proto.inputs[0].name, need_save)
|
|
op.desc.set_output(proto.outputs[0].name, need_save)
|
|
continue
|
|
|
|
# judge the other op whether should be removed.
|
|
op_dist_attr = dist_context.get_op_dist_attr_for_program(op)
|
|
if op_dist_attr is not None:
|
|
op_process_mesh = op_dist_attr.process_mesh
|
|
if (
|
|
rank_id not in op_process_mesh.process_ids
|
|
and op.type not in not_remove_op_ref
|
|
):
|
|
remove_op_idx.append(idx)
|
|
|
|
for idx in remove_op_idx[::-1]:
|
|
block._remove_op(idx, sync=False)
|
|
block._sync_with_cpp()
|
|
|
|
@staticmethod
|
|
def remove_no_need_vars(
|
|
auto_parallel_main_prog, dist_params_grads, feed_var_names
|
|
):
|
|
"""Remove no need vars in the main program"""
|
|
for block_idx, block in enumerate(auto_parallel_main_prog.blocks):
|
|
remove_vars = set()
|
|
ops = block.ops
|
|
vars = block.vars
|
|
need_vars = set()
|
|
for op in ops:
|
|
for var_name in op.input_arg_names:
|
|
if var_name in vars:
|
|
need_vars.add(var_name)
|
|
for var_name in op.output_arg_names:
|
|
if var_name in vars:
|
|
need_vars.add(var_name)
|
|
for var in vars:
|
|
if var not in need_vars:
|
|
remove_vars.add(var)
|
|
|
|
# change dist_params_grads, the optimize op just in block 0.
|
|
if block_idx == 0:
|
|
param_grad_map = {}
|
|
for op in ops:
|
|
if int(op.attr('op_role')) == int(OpRole.Optimize):
|
|
if (
|
|
"Param" in op.input_names
|
|
and "Grad" in op.input_names
|
|
):
|
|
param_name = op.input("Param")[0]
|
|
grad_name = op.input("Grad")[0]
|
|
param_grad_map[param_name] = grad_name
|
|
|
|
need_remove_idx = []
|
|
for idx, item in enumerate(dist_params_grads):
|
|
if item[0].name not in param_grad_map.keys():
|
|
need_remove_idx.append(idx)
|
|
|
|
for idx in need_remove_idx[::-1]:
|
|
dist_params_grads.pop(idx)
|
|
|
|
idx = 0
|
|
while idx < len(dist_params_grads):
|
|
param_name = dist_params_grads[idx][0].name
|
|
grad_name = dist_params_grads[idx][1].name
|
|
if grad_name != param_grad_map[param_name]:
|
|
dist_params_grads[idx] = (
|
|
vars[param_name],
|
|
vars[param_grad_map[param_name]],
|
|
)
|
|
idx += 1
|
|
|
|
for var in remove_vars:
|
|
if var in feed_var_names:
|
|
continue
|
|
block._remove_var(var, sync=False)
|
|
block._sync_with_cpp()
|
|
|
|
@staticmethod
|
|
def remove_no_need_in_main(
|
|
auto_parallel_main_prog, dist_context, rank_id, dist_params_grads
|
|
):
|
|
"""Remove no need vars and ops in the main program."""
|
|
Remover.remove_no_need_ops(
|
|
auto_parallel_main_prog, dist_context, rank_id
|
|
)
|
|
Resharder.change_while_op_input_and_output(
|
|
auto_parallel_main_prog, dist_context
|
|
)
|
|
# 'feed_var_names' cannot be removed from auto_parallel_main_prog
|
|
feed_var_names = []
|
|
for var in reduce(
|
|
operator.iadd, list(dist_context.serial_feed_vars.values()), []
|
|
):
|
|
feed_var_names.append(var.name)
|
|
Remover.remove_no_need_vars(
|
|
auto_parallel_main_prog, dist_params_grads, feed_var_names
|
|
)
|
|
|
|
@staticmethod
|
|
def remove_no_need_in_startup(
|
|
auto_parallel_main_prog, auto_parallel_startup_prog
|
|
):
|
|
"""Remove no need vars and ops in the startup program."""
|
|
main_input_vars = set()
|
|
main_ops = auto_parallel_main_prog.global_block().ops
|
|
for op in main_ops:
|
|
for var_name in op.input_arg_names:
|
|
main_input_vars.add(var_name)
|
|
|
|
startup_block = auto_parallel_startup_prog.global_block()
|
|
startup_output_vars = set()
|
|
startup_ops = startup_block.ops
|
|
for op in startup_ops:
|
|
# skip c_sync_comm_stream op
|
|
if op.type == "c_sync_comm_stream":
|
|
continue
|
|
for var_name in op.output_arg_names:
|
|
startup_output_vars.add(var_name)
|
|
|
|
need_vars = set()
|
|
for var_name in startup_output_vars:
|
|
if var_name in main_input_vars:
|
|
need_vars.add(var_name)
|
|
|
|
startup_ops = startup_block.ops
|
|
actual_need_vars = set()
|
|
for idx, op in enumerate(startup_ops):
|
|
is_need_op = False
|
|
if op.type == "c_sync_comm_stream":
|
|
continue
|
|
for var_name in op.output_arg_names:
|
|
if var_name in need_vars:
|
|
is_need_op = True
|
|
break
|
|
if is_need_op:
|
|
for var_name in op.output_arg_names:
|
|
actual_need_vars.add(var_name)
|
|
for var_name in op.input_arg_names:
|
|
actual_need_vars.add(var_name)
|
|
|
|
remove_vars = set()
|
|
for var_name in startup_block.vars:
|
|
if var_name not in actual_need_vars:
|
|
remove_vars.add(var_name)
|
|
for var in remove_vars:
|
|
startup_block._remove_var(var, sync=False)
|
|
startup_block._sync_with_cpp()
|
|
|
|
remove_op_idx = []
|
|
vars = startup_block.vars
|
|
for idx, op in enumerate(startup_block.ops):
|
|
is_no_need_op = False
|
|
if op.type == "c_sync_comm_stream":
|
|
var_names = []
|
|
for var_name in op.input_arg_names:
|
|
if var_name in vars:
|
|
var_names.append(var_name)
|
|
if not var_names:
|
|
remove_op_idx.append(idx)
|
|
else:
|
|
proto = OpProtoHolder.instance().get_op_proto(op.type)
|
|
op.desc.set_input(proto.inputs[0].name, var_names)
|
|
op.desc.set_output(proto.outputs[0].name, var_names)
|
|
continue
|
|
|
|
for var_name in op.output_arg_names:
|
|
if var_name not in vars:
|
|
is_no_need_op = True
|
|
break
|
|
if is_no_need_op:
|
|
remove_op_idx.append(idx)
|
|
for idx in remove_op_idx[::-1]:
|
|
startup_block._remove_op(idx, sync=False)
|
|
startup_block._sync_with_cpp()
|
|
|
|
|
|
class Resharder:
|
|
"""
|
|
Reshard tensor in the program according to its distributed attribute and corresponding op distributed attribute.
|
|
|
|
Args:
|
|
auto_parallel_main_prog (Program): An auto parallel main program.
|
|
auto_parallel_startup_prog (Program): An auto parallel startup program.
|
|
rank_id (int): The process id.
|
|
dist_context (DistributedContext): The distributed context of this rank.
|
|
dist_params_grads (list): The list contains the tuple of param and grad.
|
|
batch_size (int): The batch size. Default: None.
|
|
"""
|
|
|
|
while_block_info = {}
|
|
|
|
def __init__(
|
|
self,
|
|
auto_parallel_main_prog,
|
|
auto_parallel_startup_prog,
|
|
rank_id,
|
|
dist_context,
|
|
dist_params_grads,
|
|
batch_size=None,
|
|
):
|
|
assert isinstance(auto_parallel_main_prog, Program), (
|
|
"The type of auto_parallel_main_prog should be Program, "
|
|
f"but got {type(auto_parallel_main_prog)}."
|
|
)
|
|
if auto_parallel_startup_prog is not None:
|
|
assert isinstance(auto_parallel_main_prog, Program), (
|
|
"The type of auto_parallel_startup_prog should be Program or None, "
|
|
f"but got {type(auto_parallel_startup_prog)}."
|
|
)
|
|
assert isinstance(rank_id, int), (
|
|
f"The type of rank_id should be int, but got {type(rank_id)}."
|
|
)
|
|
assert isinstance(dist_context, DistributedContext), (
|
|
"The type of dist_context should be DistributedContext, "
|
|
f"but got {type(dist_context)}."
|
|
)
|
|
|
|
if batch_size is not None:
|
|
assert isinstance(batch_size, int), (
|
|
"The type of batch_size should be int, "
|
|
f"but got {type(batch_size)}."
|
|
)
|
|
|
|
self._auto_parallel_main_prog = auto_parallel_main_prog
|
|
self._auto_parallel_startup_prog = auto_parallel_startup_prog
|
|
self._rank_id = rank_id
|
|
self._dist_context = dist_context
|
|
self._dist_params_grads = dist_params_grads
|
|
self._batch_size = batch_size
|
|
self._has_sent = {}
|
|
self._has_recv = {}
|
|
self._has_allgather = {}
|
|
# to avoid reshard repeatedly
|
|
self._has_resharded = {}
|
|
|
|
@property
|
|
def auto_parallel_main_prog(self):
|
|
return self._auto_parallel_main_prog
|
|
|
|
@property
|
|
def auto_parallel_startup_prog(self):
|
|
return self._auto_parallel_startup_prog
|
|
|
|
@property
|
|
def rank_id(self):
|
|
return self._rank_id
|
|
|
|
@property
|
|
def dist_context(self):
|
|
return self._dist_context
|
|
|
|
@property
|
|
def dist_params_grads(self):
|
|
return self._dist_params_grads
|
|
|
|
@property
|
|
def batch_size(self):
|
|
return self._batch_size
|
|
|
|
@property
|
|
def has_sent(self):
|
|
return self._has_sent
|
|
|
|
@property
|
|
def has_recv(self):
|
|
return self._has_recv
|
|
|
|
@property
|
|
def has_allgather(self):
|
|
return self._has_allgather
|
|
|
|
@staticmethod
|
|
def compute_partition_shape(complete_shape, dims_mapping, process_shape):
|
|
"""Compute the shape of partition."""
|
|
partition_shape = []
|
|
for idx, item in enumerate(complete_shape):
|
|
if dims_mapping[idx] == -1:
|
|
partition_shape.append(item)
|
|
else:
|
|
partition_shape.append(item // process_shape[dims_mapping[idx]])
|
|
|
|
return partition_shape
|
|
|
|
@staticmethod
|
|
def compute_process_index(process, process_group, process_shape):
|
|
"""Compute the index of process_shape corresponding to the process."""
|
|
relative_process = process_group.index(process)
|
|
process_index = []
|
|
product = reduce(lambda x, y: x * y, process_shape, 1)
|
|
|
|
for i in range(len(process_shape)):
|
|
idx = relative_process // (product // process_shape[i])
|
|
product = product // process_shape[i]
|
|
relative_process = (
|
|
relative_process - relative_process // product * product
|
|
)
|
|
process_index.append(idx)
|
|
|
|
return process_index
|
|
|
|
@staticmethod
|
|
def compute_partition_index(
|
|
process, complete_shape, dims_mapping, process_shape, process_group
|
|
):
|
|
"""Compute the partition index in complete tensor."""
|
|
partition_shape = Resharder.compute_partition_shape(
|
|
complete_shape, dims_mapping, process_shape
|
|
)
|
|
process_index = Resharder.compute_process_index(
|
|
process, process_group, process_shape
|
|
)
|
|
partition_index = []
|
|
|
|
for i in range(len(complete_shape)):
|
|
if dims_mapping[i] == -1:
|
|
partition_index.append([0, partition_shape[i]])
|
|
else:
|
|
partition_index.append(
|
|
[
|
|
process_index[dims_mapping[i]] * partition_shape[i],
|
|
(process_index[dims_mapping[i]] + 1)
|
|
* partition_shape[i],
|
|
]
|
|
)
|
|
|
|
return partition_index
|
|
|
|
@staticmethod
|
|
def compute_concat_info(partition_index_x, partition_index_y):
|
|
"""Judge whether two partition can be concatenated and compute concatenated partition index."""
|
|
differ_count = 0
|
|
concat_axis = -1
|
|
first_order = 0
|
|
new_partition = []
|
|
|
|
for idx, item in enumerate(partition_index_x):
|
|
if item != partition_index_y[idx]:
|
|
differ_count += 1
|
|
if (
|
|
item[1] == partition_index_y[idx][0]
|
|
and item[0] < partition_index_y[idx][1]
|
|
):
|
|
concat_axis = idx
|
|
new_partition.append([item[0], partition_index_y[idx][1]])
|
|
elif (
|
|
item[0] == partition_index_y[idx][1]
|
|
and item[1] > partition_index_y[idx][0]
|
|
):
|
|
first_order = 1
|
|
concat_axis = idx
|
|
new_partition.append([partition_index_y[idx][0], item[1]])
|
|
else:
|
|
new_partition.append(item)
|
|
|
|
if differ_count == 1:
|
|
return concat_axis, first_order, new_partition
|
|
else:
|
|
return -1, first_order, new_partition
|
|
|
|
@staticmethod
|
|
def compute_complete_shape(slice_shape, process_shape, dims_mapping):
|
|
"""compute the complete shape of the slice tensor with its process mesh and dims mapping"""
|
|
complete_shape = []
|
|
for idx, item in enumerate(slice_shape):
|
|
if dims_mapping[idx] == -1:
|
|
complete_shape.append(item)
|
|
else:
|
|
complete_shape.append(item * process_shape[dims_mapping[idx]])
|
|
return complete_shape
|
|
|
|
@staticmethod
|
|
def concat_partitions(partition_index_list, partition_index):
|
|
"""Concat the given partitions without inserting concat op."""
|
|
if not partition_index_list:
|
|
partition_index_list.append(partition_index)
|
|
else:
|
|
i = 0
|
|
has_concat = False
|
|
while i < len(partition_index_list):
|
|
concat_axis, _, new_partition = Resharder.compute_concat_info(
|
|
partition_index_list[i], partition_index
|
|
)
|
|
if concat_axis != -1:
|
|
has_concat = True
|
|
partition_index_list.pop(i)
|
|
Resharder.concat_partitions(
|
|
partition_index_list, new_partition
|
|
)
|
|
break
|
|
i += 1
|
|
if not has_concat:
|
|
partition_index_list.append(partition_index)
|
|
|
|
@staticmethod
|
|
def change_while_op_input_and_output(auto_parallel_main_prog, dist_context):
|
|
"""Change while op input and output after the corresponding sub block ops removed"""
|
|
for sub_block_idx in Resharder.while_block_info:
|
|
sub_block = auto_parallel_main_prog.blocks[sub_block_idx]
|
|
parent_while_op_id = Resharder.while_block_info[sub_block_idx][
|
|
"op_id"
|
|
]
|
|
parent_block = auto_parallel_main_prog.blocks[sub_block.parent_idx]
|
|
|
|
sub_block_op_inputs = set()
|
|
sub_block_op_outputs = []
|
|
for op in sub_block.ops:
|
|
# skip the input and output of operators inserted in the reshard phase
|
|
dist_op = dist_context.get_dist_op_for_program(op)
|
|
if (
|
|
dist_op
|
|
or (op.type == "slice" and not dist_op)
|
|
or (op.type == "split" and not dist_op)
|
|
or (op.type == "assign" and not dist_op)
|
|
):
|
|
for var_name in op.output_arg_names:
|
|
if var_name not in sub_block_op_outputs:
|
|
sub_block_op_outputs.append(var_name)
|
|
for var_name in op.input_arg_names:
|
|
sub_block_op_inputs.add(var_name)
|
|
|
|
# find the while op
|
|
while_op = None
|
|
for op in parent_block.ops:
|
|
if op.desc.id() == parent_while_op_id and op.type == "while":
|
|
while_op = op
|
|
break
|
|
|
|
if while_op is None:
|
|
continue
|
|
|
|
# find the actual input and output of while op
|
|
proto = OpProtoHolder.instance().get_op_proto(while_op.type)
|
|
new_X = []
|
|
for var_name in while_op.input("X"):
|
|
if var_name in sub_block_op_inputs:
|
|
new_X.append(var_name)
|
|
assert new_X
|
|
new_X.sort()
|
|
while_op.desc.set_input(proto.inputs[0].name, new_X)
|
|
|
|
new_Out = []
|
|
for var_name in while_op.output("Out"):
|
|
for output_name in sub_block_op_outputs[::-1]:
|
|
if output_name.find(var_name) != -1 and (
|
|
len(var_name) == len(output_name)
|
|
or "@RESHARD" in output_name
|
|
):
|
|
if output_name not in new_Out:
|
|
new_Out.append(output_name)
|
|
assert new_Out
|
|
while_op.desc.set_output(proto.outputs[0].name, new_Out)
|
|
|
|
def is_overlapped(self, shape_x, shape_y):
|
|
"""Judge whether two partitions intersect on the specified dimension."""
|
|
overlapped = False
|
|
if (shape_y[0] <= shape_x[0] < shape_y[1]) or (
|
|
shape_x[0] <= shape_y[0] < shape_x[1]
|
|
):
|
|
overlapped = True
|
|
if shape_x == [0, 0] and shape_y == [0, 0]:
|
|
overlapped = True
|
|
return overlapped
|
|
|
|
def is_unshard(self, dims_mapping):
|
|
for dim in dims_mapping:
|
|
if dim != -1:
|
|
return False
|
|
return True
|
|
|
|
def is_special_op(self, op):
|
|
global _g_special_ops
|
|
if op.type in _g_special_ops:
|
|
return True
|
|
if is_gradient_clip_op(op) and op.type in _g_gradient_clip_ops:
|
|
return True
|
|
return False
|
|
|
|
def is_condition_replicative(self, op):
|
|
sub_block = self.auto_parallel_main_prog.blocks[op.attr("sub_block").id]
|
|
|
|
if op.type == "while":
|
|
input_cond = op.input("Condition")
|
|
elif op.type == "conditional_block":
|
|
input_cond = op.input("Cond")
|
|
|
|
# the dims mapping of condition tensor should be replicative
|
|
for var_name in input_cond:
|
|
var = get_var_with_recursion(
|
|
var_name, sub_block, self.auto_parallel_main_prog
|
|
)
|
|
dist_tensor = self.dist_context.get_dist_tensor_for_program(var)
|
|
tensor_dist_attr = dist_tensor.dist_attr
|
|
var_dims_mapping = tensor_dist_attr.dims_mapping
|
|
for dim in var_dims_mapping:
|
|
if dim != -1:
|
|
return False
|
|
|
|
return True
|
|
|
|
def need_reshard(self, dist_tensor, dist_attr, op_input=True, dist_op=None):
|
|
"""Judge the tensor whether needs to be resharded."""
|
|
is_reshard = False
|
|
tensor_dist_attr = dist_tensor.dist_attr
|
|
tensor_dims_mapping = tensor_dist_attr.dims_mapping
|
|
tensor_process_mesh = tensor_dist_attr.process_mesh
|
|
|
|
# dist_attr is [process_mesh, dims_mapping, chunk_id, op_role] and process_mesh is not a union
|
|
op_process_mesh = dist_attr[0]
|
|
|
|
if op_input:
|
|
op_input_dims_mapping = dist_attr[1]
|
|
if all(
|
|
x
|
|
for x in [
|
|
tensor_dims_mapping,
|
|
tensor_process_mesh,
|
|
op_input_dims_mapping,
|
|
op_process_mesh,
|
|
]
|
|
):
|
|
# judge whether need reshard by dims_mapping
|
|
if tensor_dims_mapping != op_input_dims_mapping:
|
|
if (
|
|
tensor_process_mesh
|
|
not in self.dist_context.process_meshes
|
|
):
|
|
# assert whether -1 when union.
|
|
for item in tensor_dims_mapping:
|
|
if item != -1:
|
|
raise ValueError(
|
|
"The dim must be -1 when tensor process mesh is a union."
|
|
)
|
|
is_reshard = True
|
|
|
|
# judge whether need reshard by process_mesh
|
|
if tensor_process_mesh != op_process_mesh:
|
|
is_reshard = True
|
|
# not reshard data in send/recv scene
|
|
if (
|
|
tensor_process_mesh != op_process_mesh
|
|
and len(tensor_process_mesh.process_ids)
|
|
== len(op_process_mesh.process_ids)
|
|
and dist_tensor.serial_tensor.is_data
|
|
):
|
|
is_reshard = False
|
|
else:
|
|
op_output_dims_mapping = dist_attr[1]
|
|
if all(
|
|
x
|
|
for x in [
|
|
tensor_dims_mapping,
|
|
tensor_process_mesh,
|
|
op_output_dims_mapping,
|
|
op_process_mesh,
|
|
]
|
|
):
|
|
if tensor_dims_mapping != op_output_dims_mapping:
|
|
raise ValueError(
|
|
"It is not supported that tensor dims mapping is different from op output dims mapping."
|
|
)
|
|
if tensor_process_mesh != op_process_mesh:
|
|
is_reshard = True
|
|
|
|
return is_reshard
|
|
|
|
def get_op_process_meshes(self, op):
|
|
"""Get sub process meshes of the given op if op process mesh is a union."""
|
|
process_meshes = []
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
op_process_mesh = dist_op.dist_attr.process_mesh
|
|
|
|
for process_mesh in self.dist_context.process_meshes:
|
|
if set(process_mesh.process_ids) & (
|
|
set(op_process_mesh.process_ids)
|
|
) and len(process_mesh.process_ids) < len(
|
|
op_process_mesh.process_ids
|
|
):
|
|
process_meshes.append(process_mesh)
|
|
|
|
# it means the process mesh is not a union when process meshes is null
|
|
if not process_meshes:
|
|
process_meshes.append(op_process_mesh)
|
|
|
|
return process_meshes
|
|
|
|
def find_op_desc_seq(
|
|
self,
|
|
dist_tensor,
|
|
dist_attr,
|
|
serial=False,
|
|
is_union_process_mesh_tensor=False,
|
|
):
|
|
"""
|
|
Find the op description sequence to reshard the source tensor for matching the op requirement.
|
|
|
|
Args:
|
|
dist_tensor (DistributedTensor): A distributed tensor.
|
|
dist_attr (list): A list contains process_mesh and dims_mapping such as [process_mesh, dims_mapping].
|
|
serial (bool): If serial is true, the dist tensor and dist op come from serial program. Otherwise, they come from auto program.
|
|
|
|
Returns:
|
|
Dict, the dict represents the required op description sequence corresponding to process, The key of dict is
|
|
process and value is a list containing op description.
|
|
"""
|
|
tensor_dist_attr = dist_tensor.dist_attr
|
|
source_tensor = dist_tensor.serial_tensor
|
|
|
|
source_dims_mapping = tensor_dist_attr.dims_mapping
|
|
source_process_mesh = tensor_dist_attr.process_mesh
|
|
source_process_group = source_process_mesh.process_ids
|
|
source_process_shape = source_process_mesh.shape
|
|
|
|
target_process_mesh = dist_attr[0]
|
|
target_dims_mapping = dist_attr[1]
|
|
target_process_group = target_process_mesh.process_ids
|
|
target_process_shape = target_process_mesh.shape
|
|
|
|
# NOTE(zhaoyingli):
|
|
# tensor's attr is process_mesh([0, 1, 2, 3]) dims_mapping([-1, -1]), which means the tensor is an union process_mesh tensor
|
|
# op input's attr is process_mesh([0, 1]) dims_mapping([0, -1])
|
|
# reshard will insert split op before the reshard_op
|
|
if is_union_process_mesh_tensor:
|
|
assert (
|
|
len(set(source_dims_mapping)) == 1
|
|
and next(iter(set(source_dims_mapping))) == -1
|
|
)
|
|
if set(target_process_group).intersection(
|
|
set(source_process_group)
|
|
):
|
|
source_process_group = target_process_group
|
|
source_process_shape = target_process_shape
|
|
|
|
if source_tensor.shape[0] < 0:
|
|
assert source_tensor.shape[0] == -1
|
|
new_shape = list(source_tensor.shape)
|
|
new_shape[0] = self.batch_size
|
|
source_tensor.desc.set_shape(new_shape)
|
|
|
|
complete_shape = (
|
|
Resharder.compute_complete_shape(
|
|
source_tensor.shape, source_process_shape, source_dims_mapping
|
|
)
|
|
if not serial
|
|
else source_tensor.shape
|
|
)
|
|
op_desc_seq = OrderedDict()
|
|
|
|
# TODO: if the target process group has the same process with source process group
|
|
if set(target_process_group).intersection(
|
|
set(source_process_group)
|
|
) and set(target_process_group).difference(set(source_process_group)):
|
|
pass
|
|
|
|
elif target_process_group != source_process_group:
|
|
partition_process_mapping_list = []
|
|
for source_process in source_process_group:
|
|
# get partition index of source process
|
|
source_partition_index = Resharder.compute_partition_index(
|
|
source_process,
|
|
complete_shape,
|
|
source_dims_mapping,
|
|
source_process_shape,
|
|
source_process_group,
|
|
)
|
|
if not partition_process_mapping_list:
|
|
# the item in partition_process_mapping_list is source_partition_index, which processes and whether has been used
|
|
partition_process_mapping_list.append(
|
|
[source_partition_index, [source_process], [False]]
|
|
)
|
|
else:
|
|
partition_list = [
|
|
item[0] for item in partition_process_mapping_list
|
|
]
|
|
process_list = [
|
|
item[1] for item in partition_process_mapping_list
|
|
]
|
|
has_used = [
|
|
item[2] for item in partition_process_mapping_list
|
|
]
|
|
|
|
if partition_list.count(source_partition_index) == 1:
|
|
index = partition_list.index(source_partition_index)
|
|
process_list[index].append(source_process)
|
|
has_used[index].append(False)
|
|
else:
|
|
partition_process_mapping_list.append(
|
|
[source_partition_index, [source_process], [False]]
|
|
)
|
|
|
|
for target_process in target_process_group:
|
|
# has_sent means the source_partition_index has been sent to target_process
|
|
has_sent = []
|
|
target_partition_index = Resharder.compute_partition_index(
|
|
target_process,
|
|
complete_shape,
|
|
target_dims_mapping,
|
|
target_process_shape,
|
|
target_process_group,
|
|
)
|
|
partition_index_list = []
|
|
all_partition_index_list = []
|
|
for source_process in source_process_group:
|
|
source_partition_index = Resharder.compute_partition_index(
|
|
source_process,
|
|
complete_shape,
|
|
source_dims_mapping,
|
|
source_process_shape,
|
|
source_process_group,
|
|
)
|
|
to_send_process = None
|
|
if (
|
|
all(
|
|
_
|
|
for _ in list(
|
|
map(
|
|
self.is_overlapped,
|
|
source_partition_index,
|
|
target_partition_index,
|
|
)
|
|
)
|
|
)
|
|
and source_partition_index not in has_sent
|
|
):
|
|
idx = [
|
|
item[0] for item in partition_process_mapping_list
|
|
].index(source_partition_index)
|
|
has_used = [
|
|
item[2] for item in partition_process_mapping_list
|
|
][idx]
|
|
process_list = [
|
|
item[1] for item in partition_process_mapping_list
|
|
][idx]
|
|
i = 0
|
|
while i < len(has_used):
|
|
if not has_used[i]:
|
|
to_send_process = process_list[i]
|
|
has_used[i] = True
|
|
break
|
|
i += 1
|
|
|
|
if i == len(has_used):
|
|
has_used = [False for x in has_used]
|
|
to_send_process = process_list[0]
|
|
has_used[0] = True
|
|
assert to_send_process is not None, (
|
|
"Failed to find the send process."
|
|
)
|
|
|
|
if to_send_process not in op_desc_seq.keys():
|
|
op_desc_seq[to_send_process] = []
|
|
if target_process not in op_desc_seq.keys():
|
|
op_desc_seq[target_process] = []
|
|
all_partition_index_list.append(source_partition_index)
|
|
|
|
# append send and recv op desc
|
|
is_bool = dist_tensor.serial_tensor.dtype == paddle.bool
|
|
send_op_desc = SendOpDesc(
|
|
source_partition_index,
|
|
to_send_process,
|
|
target_process,
|
|
is_bool=is_bool,
|
|
)
|
|
recv_op_desc = RecvOpDesc(
|
|
source_partition_index,
|
|
to_send_process,
|
|
target_process,
|
|
is_bool=is_bool,
|
|
)
|
|
op_desc_seq[to_send_process].append(send_op_desc)
|
|
op_desc_seq[target_process].append(recv_op_desc)
|
|
has_sent.append(source_partition_index)
|
|
Resharder.concat_partitions(
|
|
partition_index_list, source_partition_index
|
|
)
|
|
# TODO(zhaoyingli): Remove the method to a pass.
|
|
# Current method to get all pp_ranks' relationship must rely on reshard.
|
|
# When reshard insert send/recv pair, the process_group has the pp relationship.
|
|
# But the method to obtain pp_ranks' relationship is only supported in 'reshard_input',
|
|
# cause 'reshard_output' only has current process_group view instead of global view.
|
|
op_role = dist_attr[-1]
|
|
if int(op_role) == int(OpRole.Forward):
|
|
self.dist_context.up_down_streams.add_pair_stream(
|
|
to_send_process, target_process
|
|
)
|
|
|
|
# append concat op desc
|
|
op_desc_seq[target_process].append(
|
|
ConcatOpDesc(all_partition_index_list)
|
|
)
|
|
|
|
# append slice op desc
|
|
slice_starts = []
|
|
slice_ends = []
|
|
slices_axes = []
|
|
concatenated_partition_index = partition_index_list[0]
|
|
to_slice_tensor_shape = []
|
|
|
|
for idx, item in enumerate(concatenated_partition_index):
|
|
slice_starts.append(
|
|
target_partition_index[idx][0] - item[0]
|
|
)
|
|
slice_ends.append(target_partition_index[idx][1] - item[0])
|
|
slices_axes.append(idx)
|
|
to_slice_tensor_shape.append(item[1] - item[0])
|
|
|
|
op_desc_seq[target_process].append(
|
|
SliceOpDesc(
|
|
slice_starts,
|
|
slice_ends,
|
|
slices_axes,
|
|
shape=to_slice_tensor_shape,
|
|
)
|
|
)
|
|
|
|
# In the same process group, it will use allgather and slice op.
|
|
else:
|
|
# NOTE: It just supports even partition scene.
|
|
partition_index_list = []
|
|
all_partition_index_list = []
|
|
process_index = []
|
|
for source_process in source_process_group:
|
|
source_partition_index = Resharder.compute_partition_index(
|
|
source_process,
|
|
complete_shape,
|
|
source_dims_mapping,
|
|
source_process_shape,
|
|
source_process_group,
|
|
)
|
|
if source_partition_index not in partition_index_list:
|
|
partition_index_list.append(source_partition_index)
|
|
process_index.append(
|
|
[
|
|
[
|
|
source_process,
|
|
],
|
|
source_partition_index,
|
|
]
|
|
)
|
|
else:
|
|
process_index[
|
|
partition_index_list.index(source_partition_index)
|
|
][0].append(source_process)
|
|
|
|
for i in range(len(process_index[0][0])):
|
|
group = []
|
|
for j in range(len(process_index)):
|
|
group.append(process_index[j][0][i])
|
|
if i == 0:
|
|
all_partition_index_list.append(process_index[j][1])
|
|
for process in group:
|
|
min_comm_group = copy.deepcopy(group)
|
|
all_partition_index_list_copied = copy.deepcopy(
|
|
all_partition_index_list
|
|
)
|
|
target_partition_index = Resharder.compute_partition_index(
|
|
process,
|
|
complete_shape,
|
|
target_dims_mapping,
|
|
target_process_shape,
|
|
target_process_group,
|
|
)
|
|
for _process in group:
|
|
source_partition_index = (
|
|
Resharder.compute_partition_index(
|
|
_process,
|
|
complete_shape,
|
|
source_dims_mapping,
|
|
source_process_shape,
|
|
source_process_group,
|
|
)
|
|
)
|
|
if not all(
|
|
_
|
|
for _ in list(
|
|
map(
|
|
self.is_overlapped,
|
|
source_partition_index,
|
|
target_partition_index,
|
|
)
|
|
)
|
|
):
|
|
min_comm_group.remove(_process)
|
|
all_partition_index_list_copied.remove(
|
|
source_partition_index
|
|
)
|
|
|
|
concatenated_partition_index_list = []
|
|
for partition_index in all_partition_index_list_copied:
|
|
Resharder.concat_partitions(
|
|
concatenated_partition_index_list, partition_index
|
|
)
|
|
|
|
concatenated_partition_index = (
|
|
concatenated_partition_index_list[0]
|
|
)
|
|
|
|
slice_starts = []
|
|
slice_ends = []
|
|
slices_axes = []
|
|
to_slice_tensor_shape = []
|
|
for idx, item in enumerate(concatenated_partition_index):
|
|
slice_starts.append(
|
|
target_partition_index[idx][0] - item[0]
|
|
)
|
|
slice_ends.append(
|
|
target_partition_index[idx][1] - item[0]
|
|
)
|
|
slices_axes.append(idx)
|
|
to_slice_tensor_shape.append(item[1] - item[0])
|
|
slice_op_desc = SliceOpDesc(
|
|
starts=slice_starts,
|
|
ends=slice_ends,
|
|
axes=slices_axes,
|
|
shape=to_slice_tensor_shape,
|
|
)
|
|
allgather_shape = (
|
|
None
|
|
if not serial
|
|
else dist_tensor.local_sizes(rank=process)
|
|
)
|
|
# c_concat pass
|
|
if (
|
|
target_dims_mapping.count(-1)
|
|
== len(target_dims_mapping)
|
|
and source_dims_mapping[:-1].count(-1)
|
|
== len(source_dims_mapping[:-1])
|
|
and source_dims_mapping[-1] != -1
|
|
):
|
|
op_desc_seq[process] = [
|
|
AllGatherConcatOpDesc(
|
|
group=group, shape=allgather_shape
|
|
)
|
|
]
|
|
# optimization: [sharded, any x n] -> [unsharded, any x n], only need one allgather and no split or concat anymore.
|
|
elif (
|
|
target_dims_mapping[1:] == source_dims_mapping[1:]
|
|
and target_dims_mapping[0] == -1
|
|
and source_dims_mapping[0] != -1
|
|
):
|
|
op_desc_seq[process] = [
|
|
AllGatherOpDesc(
|
|
group=min_comm_group,
|
|
shape=allgather_shape,
|
|
is_bool=(source_tensor.dtype == paddle.bool),
|
|
need_split=False,
|
|
),
|
|
EndOpDesc(None),
|
|
]
|
|
else:
|
|
op_desc_seq[process] = (
|
|
[
|
|
AllGatherOpDesc(
|
|
group=min_comm_group,
|
|
shape=allgather_shape,
|
|
is_bool=(
|
|
source_tensor.dtype == paddle.bool
|
|
),
|
|
),
|
|
ConcatOpDesc(
|
|
partition_index_list=all_partition_index_list_copied
|
|
),
|
|
slice_op_desc,
|
|
]
|
|
if len(min_comm_group) > 1
|
|
else [slice_op_desc]
|
|
)
|
|
|
|
return op_desc_seq
|
|
|
|
def parse_op_desc(
|
|
self,
|
|
block,
|
|
op_desc_seq,
|
|
src_tensor,
|
|
reshard_op,
|
|
src_tensor_attr,
|
|
dst_input_attr,
|
|
sync=True,
|
|
):
|
|
"""
|
|
Parse op desc sequence and insert op in the block
|
|
|
|
src_tensor_attr(TensorDistAttr): tensor's dist_attr
|
|
dst_input_attr(list): input_var's dist_attrs of the op
|
|
"""
|
|
|
|
# Parse all communicator groups for all ranks
|
|
# Ensure every rank has a global view of communicator groups for entire cluster.
|
|
# When initialize communicators for pipeline parallel, every rank could
|
|
# conduct a correct global synchronization.
|
|
for rank_id in op_desc_seq:
|
|
op_desc_list = op_desc_seq[rank_id]
|
|
for op_desc in op_desc_list:
|
|
if isinstance(
|
|
op_desc, (AllGatherOpDesc, AllGatherConcatOpDesc)
|
|
):
|
|
new_process_group(op_desc.group)
|
|
elif isinstance(op_desc, SendOpDesc):
|
|
new_process_group(
|
|
[op_desc.src, op_desc.dst], group_type='p2p'
|
|
)
|
|
elif isinstance(op_desc, RecvOpDesc):
|
|
new_process_group(
|
|
[op_desc.src, op_desc.dst], group_type='p2p'
|
|
)
|
|
|
|
tensor_list = []
|
|
partition_tensor_list = []
|
|
if self.rank_id not in op_desc_seq.keys():
|
|
return
|
|
op_desc_list = op_desc_seq[self.rank_id]
|
|
|
|
idx = None
|
|
for index, op in list(enumerate(block.ops)):
|
|
if op.desc.id == reshard_op.desc.id:
|
|
idx = index
|
|
break
|
|
assert idx is not None, (
|
|
f"The op for reshard cannot be found in the rank {self.rank_id} program."
|
|
)
|
|
|
|
src_name = src_tensor.name
|
|
|
|
def is_grad(name):
|
|
return name.endswith('GRAD')
|
|
|
|
# all op that generate grad is marked as OpRole.Backward
|
|
op_role = (
|
|
OpRole.Backward
|
|
if is_optimize_op(reshard_op) and is_grad(src_name)
|
|
else reshard_op.attr('op_role')
|
|
)
|
|
|
|
# a Hack to send output vars from allgather_op to end_op
|
|
end_vars = None
|
|
for op_desc in op_desc_list:
|
|
if isinstance(op_desc, AllGatherOpDesc):
|
|
if src_name not in self.has_allgather.keys():
|
|
self.has_allgather[src_name] = []
|
|
if not self.has_allgather[src_name] or op_desc.group not in [
|
|
x[0] for x in self.has_allgather[src_name]
|
|
]:
|
|
if op_desc.is_bool:
|
|
# for bool data allgather, cast to int64 -> allgather -> cast bool
|
|
out_cast = Inserter.insert_cast_op(
|
|
block,
|
|
idx,
|
|
src_tensor,
|
|
op_role,
|
|
paddle.int64,
|
|
sync=sync,
|
|
)
|
|
tensor_list, idx_offset = Inserter.insert_allgather_op(
|
|
block,
|
|
idx + 1,
|
|
out_cast,
|
|
op_desc.group,
|
|
op_role,
|
|
need_split=op_desc.need_split,
|
|
sync=sync,
|
|
)
|
|
idx += idx_offset
|
|
tensor_name_list = []
|
|
for var in tensor_list:
|
|
out_cast = Inserter.insert_cast_op(
|
|
block,
|
|
idx,
|
|
var,
|
|
op_role,
|
|
paddle.bool,
|
|
sync=sync,
|
|
)
|
|
tensor_name_list.append(out_cast.name)
|
|
idx += 1
|
|
self.has_allgather[src_name].append(
|
|
[op_desc.group, tensor_name_list]
|
|
)
|
|
else:
|
|
tensor_list, idx_offset = Inserter.insert_allgather_op(
|
|
block,
|
|
idx,
|
|
src_tensor,
|
|
op_desc.group,
|
|
op_role,
|
|
need_split=op_desc.need_split,
|
|
sync=sync,
|
|
)
|
|
# NOTE(zhaoyingli): ONLY `process_mesh` and `chunk_id` are meaningful.
|
|
for offset in range(idx_offset):
|
|
op = block.ops[idx + offset]
|
|
for out_name in op.output_arg_names:
|
|
out_var = block.vars[out_name]
|
|
set_var_dist_attr(
|
|
self.dist_context,
|
|
out_var,
|
|
[-1] * len(out_var.shape),
|
|
src_tensor_attr.process_mesh,
|
|
chunk_id=src_tensor_attr.chunk_id,
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh(
|
|
op,
|
|
src_tensor_attr.process_mesh,
|
|
self.dist_context,
|
|
chunk_id=src_tensor_attr.chunk_id,
|
|
)
|
|
|
|
if idx_offset == 1:
|
|
end_vars = tensor_list
|
|
idx += idx_offset
|
|
tensor_name_list = [var.name for var in tensor_list]
|
|
self.has_allgather[src_name].append(
|
|
[op_desc.group, tensor_name_list]
|
|
)
|
|
else:
|
|
for item in self.has_allgather[src_name]:
|
|
if op_desc.group == item[0]:
|
|
tensor_list = [
|
|
get_var_with_recursion(
|
|
var_name,
|
|
block,
|
|
self.auto_parallel_main_prog,
|
|
)
|
|
for var_name in item[1]
|
|
]
|
|
break
|
|
assert tensor_list, (
|
|
"The result of parsing allgather op should not be None."
|
|
)
|
|
|
|
elif isinstance(op_desc, SendOpDesc):
|
|
if src_name not in self.has_sent.keys():
|
|
self.has_sent[src_name] = []
|
|
if op_desc.dst not in self.has_sent[src_name]:
|
|
if op_desc.is_bool:
|
|
out_cast = Inserter.insert_cast_op(
|
|
block,
|
|
idx,
|
|
src_tensor,
|
|
op_role,
|
|
paddle.int64,
|
|
sync=sync,
|
|
)
|
|
Inserter.insert_send_op(
|
|
block,
|
|
idx + 1,
|
|
out_cast,
|
|
op_desc.src,
|
|
op_desc.dst,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
idx += 2
|
|
else:
|
|
Inserter.insert_send_op(
|
|
block,
|
|
idx,
|
|
src_tensor,
|
|
op_desc.src,
|
|
op_desc.dst,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
block.ops[idx],
|
|
src_tensor_attr.process_mesh,
|
|
src_tensor_attr.dims_mapping,
|
|
self.dist_context,
|
|
chunk_id=src_tensor_attr.chunk_id,
|
|
)
|
|
idx += 1
|
|
self.has_sent[src_name].append(op_desc.dst)
|
|
|
|
elif isinstance(op_desc, RecvOpDesc):
|
|
if src_name not in self.has_recv.keys():
|
|
self.has_recv[src_name] = {}
|
|
if op_desc.src not in self.has_recv[src_name].keys():
|
|
partition_index = op_desc.partition_index
|
|
shape = []
|
|
for index in partition_index:
|
|
shape.append(index[1] - index[0])
|
|
if op_desc.is_bool:
|
|
# for bool data, recv int64 -> cast to bool
|
|
recv_tensor = block.create_var(
|
|
name=unique_name.generate(src_name + "@recv"),
|
|
shape=shape,
|
|
lod_level=src_tensor.lod_level,
|
|
dtype=paddle.int64,
|
|
type=src_tensor.type,
|
|
)
|
|
Inserter.insert_recv_op(
|
|
block,
|
|
idx,
|
|
recv_tensor,
|
|
op_desc.src,
|
|
op_desc.dst,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
out_cast = Inserter.insert_cast_op(
|
|
block,
|
|
idx + 1,
|
|
recv_tensor,
|
|
op_role,
|
|
paddle.bool,
|
|
sync=sync,
|
|
)
|
|
tensor_list.append(out_cast)
|
|
idx += 2
|
|
self.has_recv[src_name][op_desc.src] = out_cast
|
|
else:
|
|
recv_tensor = block.create_var(
|
|
name=unique_name.generate(src_name + "@recv"),
|
|
shape=shape,
|
|
lod_level=src_tensor.lod_level,
|
|
dtype=src_tensor.dtype,
|
|
type=src_tensor.type,
|
|
)
|
|
Inserter.insert_recv_op(
|
|
block,
|
|
idx,
|
|
recv_tensor,
|
|
op_desc.src,
|
|
op_desc.dst,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
set_var_dist_attr(
|
|
self.dist_context,
|
|
recv_tensor,
|
|
dst_input_attr[1], # dims_mapping
|
|
dst_input_attr[0], # process_mesh
|
|
chunk_id=dst_input_attr[2],
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
block.ops[idx],
|
|
dst_input_attr[0], # process_mesh
|
|
dst_input_attr[1], # dims_mapping
|
|
self.dist_context,
|
|
chunk_id=dst_input_attr[2],
|
|
)
|
|
|
|
# for lod tensor, need reset lod after received
|
|
if recv_tensor.lod_level != 0:
|
|
set_lod = False
|
|
# use data lod to reset tensor lod
|
|
for (
|
|
tmp_block
|
|
) in self.auto_parallel_main_prog.blocks:
|
|
for tmp_var_name in tmp_block.vars:
|
|
tmp_var = tmp_block.vars[tmp_var_name]
|
|
if (
|
|
tmp_var.is_data
|
|
and tmp_var.lod_level
|
|
== recv_tensor.lod_level
|
|
):
|
|
reset_lod_out = (
|
|
Inserter.insert_reset_lod_op(
|
|
block,
|
|
idx + 1,
|
|
recv_tensor,
|
|
tmp_var,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
)
|
|
tensor_list.append(reset_lod_out)
|
|
idx += 2
|
|
self.has_recv[src_name][op_desc.src] = (
|
|
reset_lod_out
|
|
)
|
|
set_lod = True
|
|
break
|
|
if set_lod:
|
|
break
|
|
assert set_lod is True
|
|
else:
|
|
tensor_list.append(recv_tensor)
|
|
idx += 1
|
|
self.has_recv[src_name][op_desc.src] = recv_tensor
|
|
else:
|
|
tensor_list.append(self.has_recv[src_name][op_desc.src])
|
|
|
|
elif isinstance(op_desc, ConcatOpDesc):
|
|
partition_index_list = op_desc.partition_index_list
|
|
pre_idx = idx
|
|
idx_list = [idx]
|
|
for index, tensor in enumerate(tensor_list):
|
|
Inserter.concat_partitions_with_op(
|
|
partition_tensor_list,
|
|
tensor,
|
|
partition_index_list[index],
|
|
block,
|
|
idx_list,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
idx = idx_list[0]
|
|
cur_idx = idx
|
|
# NOTE(zhaoyingli): ONLY `process_mesh` and `chunk_id` are meaningful.
|
|
for i in range(pre_idx, cur_idx):
|
|
op = block.ops[i]
|
|
for out_name in op.output_arg_names:
|
|
out_var = block.vars[out_name]
|
|
set_var_dist_attr(
|
|
self.dist_context,
|
|
out_var,
|
|
[-1] * len(out_var.shape),
|
|
dst_input_attr[0], # process_mesh
|
|
chunk_id=src_tensor_attr.chunk_id,
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh(
|
|
op,
|
|
dst_input_attr[0], # process_mesh
|
|
self.dist_context,
|
|
chunk_id=src_tensor_attr.chunk_id,
|
|
)
|
|
|
|
elif isinstance(
|
|
op_desc, (SliceOpDesc, AllGatherConcatOpDesc, EndOpDesc)
|
|
):
|
|
target_tensor = None
|
|
if isinstance(op_desc, SliceOpDesc):
|
|
assert (
|
|
len(partition_tensor_list) == 1
|
|
or not partition_tensor_list
|
|
)
|
|
to_slice_tensor = (
|
|
partition_tensor_list[0][0]
|
|
if len(partition_tensor_list) == 1
|
|
else src_tensor
|
|
)
|
|
new_name = unique_name.generate(src_name + "@RESHARD")
|
|
target_tensor = Inserter.insert_slice_op(
|
|
block,
|
|
idx,
|
|
to_slice_tensor,
|
|
starts=op_desc.starts,
|
|
ends=op_desc.ends,
|
|
axes=op_desc.axes,
|
|
new_var_name=new_name,
|
|
op_role=op_role,
|
|
sync=sync,
|
|
)
|
|
elif isinstance(op_desc, AllGatherConcatOpDesc):
|
|
target_tensor = Inserter.insert_c_concat_op(
|
|
block,
|
|
idx,
|
|
src_tensor,
|
|
op_desc.group,
|
|
op_role,
|
|
sync=sync,
|
|
)
|
|
else:
|
|
assert isinstance(op_desc, EndOpDesc)
|
|
assert len(end_vars) == 1
|
|
target_tensor = end_vars[0]
|
|
|
|
if not isinstance(op_desc, EndOpDesc):
|
|
assert target_tensor is not None
|
|
set_var_dist_attr(
|
|
self.dist_context,
|
|
target_tensor,
|
|
dst_input_attr[1], # dims_mapping
|
|
dst_input_attr[0], # process_mesh
|
|
chunk_id=dst_input_attr[2],
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
block.ops[idx],
|
|
dst_input_attr[0], # process_mesh
|
|
dst_input_attr[1], # dims_mapping
|
|
self.dist_context,
|
|
chunk_id=dst_input_attr[2],
|
|
)
|
|
|
|
if reshard_op.type == "while":
|
|
# var_reshard_mapping means the while op input need be changed to
|
|
if (
|
|
"var_reshard_mapping"
|
|
not in Resharder.while_block_info[
|
|
op.attr("sub_block").id
|
|
].keys()
|
|
):
|
|
Resharder.while_block_info[op.attr("sub_block").id][
|
|
"var_reshard_mapping"
|
|
] = {}
|
|
if (
|
|
src_name
|
|
not in Resharder.while_block_info[
|
|
op.attr("sub_block").id
|
|
]["var_reshard_mapping"].keys()
|
|
):
|
|
Resharder.while_block_info[op.attr("sub_block").id][
|
|
"var_reshard_mapping"
|
|
][src_name] = []
|
|
Resharder.while_block_info[op.attr("sub_block").id][
|
|
"var_reshard_mapping"
|
|
][src_name].append([dst_input_attr, target_tensor.name])
|
|
|
|
# rename op input from old name to new name and there is a scene that one var can be multi-ops' input
|
|
for op in block.ops[idx:]:
|
|
if is_reshard_op(op):
|
|
continue
|
|
while_op_X_append = [] # just for while op
|
|
for name in op.input_arg_names:
|
|
op_dist_attr = (
|
|
self.dist_context.get_op_dist_attr_for_program(op)
|
|
)
|
|
assert op_dist_attr is not None
|
|
if name == src_name:
|
|
op_input_dist_attr = (
|
|
op_dist_attr.get_input_dist_attr(src_name)
|
|
)
|
|
old_name = name
|
|
new_name = target_tensor.name
|
|
assert old_name != new_name
|
|
|
|
if op.desc.id() == reshard_op.desc.id():
|
|
op.desc._rename_input(name, new_name)
|
|
op_dist_attr.set_input_dist_attr(
|
|
new_name, op_input_dist_attr
|
|
)
|
|
self.dist_context.set_op_dist_attr_for_program(
|
|
op, op_dist_attr
|
|
)
|
|
self.dist_context.set_tensor_dist_attr_for_program(
|
|
target_tensor, op_input_dist_attr
|
|
)
|
|
if op.type == "while":
|
|
while_op_X_append.append(new_name)
|
|
continue
|
|
|
|
op_process_mesh = op_dist_attr.process_mesh
|
|
op_input_dims_mapping = (
|
|
op_dist_attr.get_input_dims_mapping(src_name)
|
|
)
|
|
# NOTE: For op whose process mesh is a union, its input will not be renamed by other op reshard result now which means that it will have more reshard operation.
|
|
if (
|
|
op_process_mesh == dst_input_attr[0]
|
|
and op_input_dims_mapping == dst_input_attr[1]
|
|
):
|
|
op.desc._rename_input(name, new_name)
|
|
op_dist_attr.set_input_dist_attr(
|
|
new_name, op_input_dist_attr
|
|
)
|
|
self.dist_context.set_op_dist_attr_for_program(
|
|
op, op_dist_attr
|
|
)
|
|
|
|
# for while op, the input X should reset
|
|
if while_op_X_append:
|
|
proto = OpProtoHolder.instance().get_op_proto(op.type)
|
|
op.desc.set_input(
|
|
proto.inputs[0].name,
|
|
op.input("X") + while_op_X_append,
|
|
)
|
|
|
|
def _get_subblock_input_attrs(self, op, var_name):
|
|
# NOTE: Multi while loop is not supported
|
|
assert op.type in _g_subblock_ops
|
|
sub_block = self.auto_parallel_main_prog.blocks[op.attr("sub_block").id]
|
|
ops = sub_block.ops
|
|
input_attrs = []
|
|
|
|
for op in ops:
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
if not dist_op:
|
|
continue
|
|
dist_attr = dist_op.dist_attr
|
|
for name in op.input_arg_names:
|
|
if name == var_name:
|
|
process_mesh = dist_attr.process_mesh
|
|
input_dims_mapping = dist_attr.get_input_dims_mapping(
|
|
var_name
|
|
)
|
|
chunk_id = dist_attr.chunk_id
|
|
has_exist = False
|
|
for input_attr in input_attrs:
|
|
if (
|
|
process_mesh == input_attr[0]
|
|
and input_dims_mapping == input_attr[1]
|
|
and chunk_id == input_attr[2]
|
|
):
|
|
has_exist = True
|
|
break
|
|
if not has_exist:
|
|
input_attrs.append(
|
|
[
|
|
process_mesh,
|
|
input_dims_mapping,
|
|
chunk_id,
|
|
op.attr('op_role'),
|
|
]
|
|
)
|
|
return input_attrs
|
|
|
|
def _get_subblock_output_attrs(self, op, var_name):
|
|
# NOTE: Multi while loop is not supported
|
|
assert op.type in _g_subblock_ops
|
|
sub_block = self.auto_parallel_main_prog.blocks[op.attr("sub_block").id]
|
|
ops = sub_block.ops
|
|
output_attrs = []
|
|
|
|
for op in ops:
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
if not dist_op:
|
|
continue
|
|
dist_attr = dist_op.dist_attr
|
|
for name in op.output_arg_names:
|
|
if name == var_name:
|
|
process_mesh = dist_attr.process_mesh
|
|
output_dims_mapping = dist_attr.get_output_dims_mapping(
|
|
var_name
|
|
)
|
|
chunk_id = dist_op.dist_attr.chunk_id
|
|
has_exist = False
|
|
for output_attr in output_attrs:
|
|
if (
|
|
process_mesh == output_attr[0]
|
|
and output_dims_mapping == output_attr[1]
|
|
and chunk_id == output_attr[2]
|
|
):
|
|
has_exist = True
|
|
break
|
|
if not has_exist:
|
|
output_attrs.append(
|
|
[
|
|
process_mesh,
|
|
output_dims_mapping,
|
|
chunk_id,
|
|
op.attr('op_role'),
|
|
]
|
|
)
|
|
return output_attrs
|
|
|
|
def _get_common_op_input_attrs(self, op, var_name):
|
|
process_meshes = []
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
dist_attr = dist_op.dist_attr
|
|
op_process_mesh = dist_attr.process_mesh
|
|
for process_mesh in self.dist_context.process_meshes:
|
|
if set(process_mesh.process_ids) & (
|
|
set(op_process_mesh.process_ids)
|
|
) and len(process_mesh.process_ids) < len(
|
|
op_process_mesh.process_ids
|
|
):
|
|
process_meshes.append(process_mesh)
|
|
|
|
# it means that the process mesh is not a union when process meshes is none
|
|
if not process_meshes:
|
|
process_meshes.append(op_process_mesh)
|
|
|
|
input_dims_mapping = dist_attr.get_input_dims_mapping(var_name)
|
|
chunk_id = dist_attr.chunk_id
|
|
input_attrs = []
|
|
for process_mesh in process_meshes:
|
|
input_attrs.append(
|
|
[process_mesh, input_dims_mapping, chunk_id, op.attr('op_role')]
|
|
)
|
|
|
|
return input_attrs
|
|
|
|
def get_op_input_attrs(self, op, var_name):
|
|
op_input_attrs = []
|
|
|
|
if op.type in _g_subblock_ops:
|
|
op_input_attrs = self._get_subblock_input_attrs(op, var_name)
|
|
if not op_input_attrs:
|
|
# NOTE: [hack method]
|
|
# Adapt to quantization pass, which persist_vars, including inputs and outputs, all are in global_block.
|
|
# Therefore, the while_op's inputs will contain the all persist_vars, which will be inputs or output of the quantization op in subblock.
|
|
op_input_attrs = self._get_subblock_output_attrs(op, var_name)
|
|
else:
|
|
op_input_attrs = self._get_common_op_input_attrs(op, var_name)
|
|
|
|
assert op_input_attrs, (
|
|
f"The input '{op.name}' of op '{var_name}' has no distributed attributes in subblock"
|
|
)
|
|
|
|
return op_input_attrs
|
|
|
|
def _remove_global_process_mesh(self):
|
|
"""Remove global process mesh from dist_context.process_meshes"""
|
|
process_ids = set()
|
|
process_mesh_count = len(self.dist_context.process_meshes)
|
|
if process_mesh_count > 1:
|
|
global_process_mesh_idx = []
|
|
has_sub_process_mesh = False
|
|
for process_mesh in self.dist_context.process_meshes:
|
|
for process_id in process_mesh.process_ids:
|
|
process_ids.add(process_id)
|
|
for idx, process_mesh in enumerate(
|
|
self.dist_context.process_meshes
|
|
):
|
|
if len(set(process_mesh.process_ids)) == len(process_ids):
|
|
global_process_mesh_idx.append(idx)
|
|
elif set(process_mesh.process_ids) < process_ids:
|
|
has_sub_process_mesh = True
|
|
|
|
if has_sub_process_mesh:
|
|
for idx in reversed(global_process_mesh_idx):
|
|
self.dist_context.process_meshes.pop(idx)
|
|
|
|
def _change_subblock_op_input_and_output(self, block_idx, block):
|
|
if "var_reshard_mapping" in Resharder.while_block_info[block_idx]:
|
|
var_reshard_mapping = Resharder.while_block_info[block_idx][
|
|
"var_reshard_mapping"
|
|
]
|
|
for op in block.ops:
|
|
for var_name in op.input_arg_names:
|
|
if var_name in var_reshard_mapping:
|
|
# in while sub block, the union process mesh is not split before reshard sub block
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
dist_attr = dist_op.dist_attr
|
|
target_name = None
|
|
for item in var_reshard_mapping[var_name]:
|
|
if (
|
|
dist_attr.process_mesh == item[0][0]
|
|
and dist_attr.get_input_dims_mapping(var_name)
|
|
== item[0][1]
|
|
):
|
|
target_name = item[1]
|
|
break
|
|
|
|
if target_name:
|
|
op.desc._rename_input(var_name, target_name)
|
|
op_input_dist_attr = dist_attr.get_input_dist_attr(
|
|
var_name
|
|
)
|
|
dist_attr.set_input_dist_attr(
|
|
target_name, op_input_dist_attr
|
|
)
|
|
|
|
# the outputs also need to be renamed when the output name is the same with input name in inplace op
|
|
for var_name in op.output_arg_names:
|
|
# if the tensor has been resharded multiply, it is not supported now.
|
|
if var_name in var_reshard_mapping:
|
|
if len(var_reshard_mapping[var_name]) > 1:
|
|
raise ValueError(
|
|
"The scene is not supported that the output is inplaced and the tensor has been resharded multiply when as input."
|
|
)
|
|
target_name = var_reshard_mapping[var_name][0][1]
|
|
|
|
op.desc._rename_output(var_name, target_name)
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
op_dist_attr = dist_op.dist_attr
|
|
op_output_dist_attr = op_dist_attr.get_output_dist_attr(
|
|
var_name
|
|
)
|
|
op_dist_attr.set_output_dist_attr(
|
|
target_name, op_output_dist_attr
|
|
)
|
|
|
|
def _reshard_input(self, block):
|
|
idx = 0
|
|
while idx < len(block.ops):
|
|
pre_op_count = len(block.ops)
|
|
op = block.ops[idx]
|
|
|
|
if self.is_special_op(op):
|
|
idx += 1
|
|
continue
|
|
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
if dist_op is not None:
|
|
if op.type in _g_subblock_ops:
|
|
if not self.is_condition_replicative(op):
|
|
raise ValueError(
|
|
"Please check the condition due to the dims mapping is not replicative."
|
|
)
|
|
if (
|
|
op.attr("sub_block").id
|
|
not in Resharder.while_block_info
|
|
):
|
|
Resharder.while_block_info[op.attr("sub_block").id] = {}
|
|
Resharder.while_block_info[op.attr("sub_block").id][
|
|
"op_id"
|
|
] = op.desc.id()
|
|
|
|
if op.type == "while":
|
|
# condition var process mesh is the same with op and dims_mapping is replicative, so it do not need reshard
|
|
input_var_names = op.input("X")
|
|
elif op.type == "conditional_block":
|
|
input_var_names = op.input("Input")
|
|
else:
|
|
input_var_names = op.input_arg_names
|
|
# to avoid while op X order different
|
|
input_var_names.sort()
|
|
|
|
idx_offset = 0
|
|
for var_name in input_var_names:
|
|
# skip lod_tensor_blocking_queue_? name
|
|
if "lod_tensor_blocking_queue" in var_name:
|
|
continue
|
|
var = get_var_with_recursion(
|
|
var_name, block, self.auto_parallel_main_prog
|
|
)
|
|
dist_tensor = self.dist_context.get_dist_tensor_for_program(
|
|
var
|
|
)
|
|
|
|
# judge whether union tensor dims_mapping all -1
|
|
is_union_process_mesh_tensor = False
|
|
if (
|
|
dist_tensor.dist_attr.process_mesh
|
|
not in self.dist_context.process_meshes
|
|
and self.dist_context.process_meshes
|
|
):
|
|
is_union_process_mesh_tensor = True
|
|
assert dist_tensor.dist_attr.dims_mapping.count(
|
|
-1
|
|
) == len(dist_tensor.dist_attr.dims_mapping)
|
|
|
|
op_input_attrs = self.get_op_input_attrs(op, var_name)
|
|
for input_attr in op_input_attrs:
|
|
# deal with union tensor
|
|
if is_union_process_mesh_tensor:
|
|
# if op process mesh is subset of union tensor process mesh
|
|
# and input's dims_mapping is equal to dist_tensor's dims_mapping,
|
|
# need no reshard
|
|
if (
|
|
set(input_attr[0].process_ids)
|
|
<= set(
|
|
dist_tensor.dist_attr.process_mesh.process_ids
|
|
)
|
|
and input_attr[1]
|
|
== dist_tensor.dist_attr.dims_mapping
|
|
):
|
|
continue
|
|
|
|
if dist_tensor is not None and self.need_reshard(
|
|
dist_tensor, input_attr
|
|
):
|
|
reshard_op_desc = self.find_op_desc_seq(
|
|
dist_tensor,
|
|
input_attr,
|
|
is_union_process_mesh_tensor=is_union_process_mesh_tensor,
|
|
)
|
|
self.parse_op_desc(
|
|
block,
|
|
reshard_op_desc,
|
|
var,
|
|
op,
|
|
dist_tensor.dist_attr,
|
|
input_attr,
|
|
)
|
|
cur_op_count = len(block.ops)
|
|
idx_offset = (
|
|
idx_offset + cur_op_count - pre_op_count
|
|
)
|
|
pre_op_count = cur_op_count
|
|
idx = idx + idx_offset + 1
|
|
else:
|
|
idx += 1
|
|
block._sync_with_cpp()
|
|
|
|
def _handle_recv(
|
|
self,
|
|
block,
|
|
idx,
|
|
var,
|
|
op,
|
|
send_rank,
|
|
recv_rank,
|
|
src_output_attr,
|
|
dst_tensor_attr,
|
|
):
|
|
if self.rank_id == recv_rank:
|
|
# if recv bool data, recv then cast
|
|
if var.dtype == paddle.bool:
|
|
recv_cast_out = block.create_var(
|
|
name=unique_name.generate(var.name + "@recv"),
|
|
shape=var.shape,
|
|
lod_level=var.lod_level,
|
|
dtype=paddle.int64,
|
|
type=var.type,
|
|
)
|
|
Inserter.insert_recv_op(
|
|
block,
|
|
idx + 1,
|
|
recv_cast_out,
|
|
send_rank,
|
|
recv_rank,
|
|
op.attr('op_role'),
|
|
)
|
|
reset_lod_out = None
|
|
if var.lod_level != 0:
|
|
set_lod = False
|
|
for tmp_block in self.auto_parallel_main_prog.blocks:
|
|
for tmp_var_name in tmp_block.vars:
|
|
tmp_var = tmp_block.vars[tmp_var_name]
|
|
if (
|
|
tmp_var.is_data
|
|
and tmp_var.lod_level == var.lod_level
|
|
):
|
|
reset_lod_out = block.create_var(
|
|
name=unique_name.generate(
|
|
var.name + "@RESETLOD"
|
|
),
|
|
shape=recv_cast_out.shape,
|
|
type=recv_cast_out.type,
|
|
dtype=recv_cast_out.dtype,
|
|
lod_level=recv_cast_out.lod_level,
|
|
)
|
|
idx += 1
|
|
block._insert_op(
|
|
idx,
|
|
type="lod_reset",
|
|
inputs={'X': recv_cast_out, 'Y': tmp_var},
|
|
outputs={'Out': reset_lod_out},
|
|
attrs={'op_role': op.attr("op_role")},
|
|
)
|
|
set_lod = True
|
|
break
|
|
if set_lod:
|
|
break
|
|
assert set_lod is True
|
|
|
|
# cast int64 to bool
|
|
cast_op = block._insert_op(
|
|
idx + 2,
|
|
type='cast',
|
|
inputs={
|
|
'X': (
|
|
[recv_cast_out]
|
|
if reset_lod_out is None
|
|
else [reset_lod_out]
|
|
)
|
|
},
|
|
outputs={'Out': [var]},
|
|
attrs={
|
|
'in_dtype': recv_cast_out.dtype,
|
|
'out_dtype': var.dtype,
|
|
'op_role': op.attr('op_role'),
|
|
},
|
|
)
|
|
cast_op._set_attr('op_namescope', "/auto_parallel/reshard")
|
|
else:
|
|
if var.lod_level != 0:
|
|
recv_out = block.create_var(
|
|
name=unique_name.generate(var.name + "@recv"),
|
|
shape=var.shape,
|
|
lod_level=var.lod_level,
|
|
dtype=var.int64,
|
|
type=var.type,
|
|
)
|
|
Inserter.insert_recv_op(
|
|
block,
|
|
idx + 1,
|
|
recv_out,
|
|
send_rank,
|
|
recv_rank,
|
|
op.attr('op_role'),
|
|
)
|
|
set_lod = False
|
|
for tmp_block in self.auto_parallel_main_prog.blocks:
|
|
for tmp_var_name in tmp_block.vars:
|
|
tmp_var = tmp_block.vars[tmp_var_name]
|
|
if (
|
|
tmp_var.is_data
|
|
and tmp_var.lod_level == var.lod_level
|
|
):
|
|
idx += 1
|
|
block._insert_op(
|
|
idx,
|
|
type="lod_reset",
|
|
inputs={'X': recv_out, 'Y': tmp_var},
|
|
outputs={'Out': var},
|
|
attrs={'op_role': op.attr("op_role")},
|
|
)
|
|
set_lod = True
|
|
break
|
|
if set_lod:
|
|
break
|
|
assert set_lod is True
|
|
else:
|
|
Inserter.insert_recv_op(
|
|
block,
|
|
idx + 1,
|
|
var,
|
|
send_rank,
|
|
recv_rank,
|
|
op.attr('op_role'),
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
block.ops[idx + 1],
|
|
dst_tensor_attr.process_mesh,
|
|
dst_tensor_attr.dims_mapping,
|
|
self.dist_context,
|
|
chunk_id=dst_tensor_attr.chunk_id,
|
|
)
|
|
|
|
def _handle_send(
|
|
self,
|
|
block,
|
|
idx,
|
|
var,
|
|
op,
|
|
send_rank,
|
|
recv_rank,
|
|
src_output_attr,
|
|
dst_tensor_attr,
|
|
):
|
|
if var.dtype == paddle.bool:
|
|
cast_out = Inserter.insert_cast_op(
|
|
block, idx + 1, var, op.attr('op_role'), paddle.int64
|
|
)
|
|
Inserter.insert_send_op(
|
|
block,
|
|
idx + 2,
|
|
cast_out,
|
|
send_rank,
|
|
recv_rank,
|
|
op.attr('op_role'),
|
|
)
|
|
else:
|
|
Inserter.insert_send_op(
|
|
block, idx + 1, var, send_rank, recv_rank, op.attr('op_role')
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
block.ops[idx + 1],
|
|
src_output_attr[0], # process_mesh
|
|
src_output_attr[1], # dims_mapping
|
|
self.dist_context,
|
|
chunk_id=src_output_attr[2],
|
|
)
|
|
|
|
def _reshard_output(self, block):
|
|
# insert send and recv op if output process mesh is different from tensor process mesh
|
|
idx = 0
|
|
|
|
# skip reader and ops whose process mesh is union
|
|
def _is_special_op(op):
|
|
skip_ops = [
|
|
"create_py_reader",
|
|
"create_double_buffer_reader",
|
|
"read",
|
|
"write_to_array",
|
|
"read_from_array",
|
|
"nop",
|
|
"depend",
|
|
]
|
|
global _g_special_ops
|
|
skip_ops += _g_special_ops
|
|
skip_ops += _g_subblock_ops
|
|
if op.type in skip_ops:
|
|
return True
|
|
if is_reshard_op(op):
|
|
return True
|
|
return False
|
|
|
|
while idx < len(block.ops):
|
|
pre_op_count = len(block.ops)
|
|
op = block.ops[idx]
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
if dist_op is not None and not _is_special_op(op):
|
|
idx_offset = 0
|
|
for var_name in op.output_arg_names:
|
|
var = get_var_with_recursion(
|
|
var_name, block, self.auto_parallel_main_prog
|
|
)
|
|
dist_tensor = self.dist_context.get_dist_tensor_for_program(
|
|
var
|
|
)
|
|
tensor_process_mesh = dist_tensor.dist_attr.process_mesh
|
|
output_attr = [
|
|
dist_op.dist_attr.process_mesh,
|
|
dist_op.dist_attr.get_output_dims_mapping(var_name),
|
|
dist_op.dist_attr.chunk_id,
|
|
op.attr("op_role"),
|
|
]
|
|
if dist_tensor is not None and self.need_reshard(
|
|
dist_tensor, output_attr, False
|
|
):
|
|
tensor_processes = set(
|
|
tensor_process_mesh.process_ids
|
|
) - (
|
|
set(tensor_process_mesh.process_ids)
|
|
& set(output_attr[0].process_ids)
|
|
)
|
|
if tensor_processes:
|
|
if len(tensor_processes) != len(
|
|
output_attr[0].process_ids
|
|
):
|
|
if dist_tensor.dist_attr.dims_mapping.count(
|
|
-1
|
|
) != len(
|
|
dist_tensor.dist_attr.dims_mapping
|
|
) or output_attr[1].count(-1) != len(
|
|
output_attr[1]
|
|
):
|
|
raise ValueError(
|
|
"The dims_mapping must be -1"
|
|
)
|
|
else:
|
|
for index, tensor_process in enumerate(
|
|
tensor_processes
|
|
):
|
|
recv_rank = tensor_process
|
|
actual_index = index
|
|
if index >= len(
|
|
output_attr[0].process_ids
|
|
):
|
|
actual_index = (
|
|
index
|
|
- len(
|
|
output_attr[0].process_ids
|
|
)
|
|
) % len(output_attr[0].process_ids)
|
|
item = output_attr[0].process_ids[
|
|
actual_index
|
|
]
|
|
if recv_rank == item:
|
|
continue
|
|
if var.shape[0] == -1:
|
|
new_shape = list(var.shape)
|
|
new_shape[0] = self.batch_size
|
|
var.desc.set_shape(new_shape)
|
|
if self.rank_id == item:
|
|
# if send bool data, cast then send
|
|
self._handle_send(
|
|
block,
|
|
idx,
|
|
var,
|
|
op,
|
|
item,
|
|
recv_rank,
|
|
output_attr,
|
|
dist_tensor.dist_attr,
|
|
)
|
|
elif self.rank_id == recv_rank:
|
|
# if recv bool data, recv then cast
|
|
self._handle_recv(
|
|
block,
|
|
idx,
|
|
var,
|
|
op,
|
|
item,
|
|
recv_rank,
|
|
output_attr,
|
|
dist_tensor.dist_attr,
|
|
)
|
|
else:
|
|
# Ensure every rank has a global view of communicator groups for entire cluster.
|
|
# When initialize communicators for pipeline parallel, every rank could
|
|
# conduct a correct global synchronization.
|
|
new_process_group(
|
|
[item, recv_rank],
|
|
group_type='p2p',
|
|
)
|
|
else:
|
|
for index, tensor_process in enumerate(
|
|
tensor_processes
|
|
):
|
|
recv_rank = tensor_process
|
|
item = output_attr[0].process_ids[index]
|
|
if recv_rank == item:
|
|
continue
|
|
if var.shape[0] == -1:
|
|
new_shape = list(var.shape)
|
|
new_shape[0] = self.batch_size
|
|
var.desc.set_shape(new_shape)
|
|
if self.rank_id == item:
|
|
# if send bool data, cast then send
|
|
self._handle_send(
|
|
block,
|
|
idx,
|
|
var,
|
|
op,
|
|
item,
|
|
recv_rank,
|
|
output_attr,
|
|
dist_tensor.dist_attr,
|
|
)
|
|
elif self.rank_id == recv_rank:
|
|
# if recv bool data, recv then cast
|
|
self._handle_recv(
|
|
block,
|
|
idx,
|
|
var,
|
|
op,
|
|
item,
|
|
recv_rank,
|
|
output_attr,
|
|
dist_tensor.dist_attr,
|
|
)
|
|
else:
|
|
# Ensure every rank has a global view of communicator groups for entire cluster.
|
|
# When initialize communicators for pipeline parallel, every rank could
|
|
# conduct a correct global synchronization.
|
|
new_process_group(
|
|
[item, recv_rank], group_type='p2p'
|
|
)
|
|
|
|
cur_op_count = len(block.ops)
|
|
idx_offset = (
|
|
idx_offset + cur_op_count - pre_op_count
|
|
)
|
|
pre_op_count = cur_op_count
|
|
|
|
idx = idx + idx_offset + 1
|
|
else:
|
|
idx += 1
|
|
|
|
def reshard(self):
|
|
self._remove_global_process_mesh()
|
|
for block_idx, block in enumerate(self.auto_parallel_main_prog.blocks):
|
|
# change the var_name before resharding sub block
|
|
if block_idx in Resharder.while_block_info:
|
|
self._change_subblock_op_input_and_output(block_idx, block)
|
|
|
|
# reshard input
|
|
self._reshard_input(block)
|
|
|
|
# reshard output
|
|
# NOTE: Only support that insert send and recv op if output process mesh is different from tensor process mesh
|
|
self._reshard_output(block)
|
|
|
|
# remove no need vars and ops in the main program
|
|
Remover.remove_no_need_in_main(
|
|
self.auto_parallel_main_prog,
|
|
self.dist_context,
|
|
self.rank_id,
|
|
self.dist_params_grads,
|
|
)
|
|
|
|
# remove no need vars and ops in the startup program
|
|
Remover.remove_no_need_in_startup(
|
|
self.auto_parallel_main_prog, self.auto_parallel_startup_prog
|
|
)
|
|
|
|
# reset some variable when remove operation ended
|
|
Resharder.while_block_info = {}
|
|
|
|
def get_cost(self, op, tensor, cluster):
|
|
# NOTE: The program should be the serial_program which is not been parted
|
|
global _g_special_ops
|
|
not_supported_op_type = [*_g_special_ops, 'while']
|
|
reshard_op_cost = None
|
|
if op.type in not_supported_op_type:
|
|
return reshard_op_cost
|
|
else:
|
|
tensor_name = tensor.name
|
|
if tensor_name == "lod_tensor_blocking_queue_0":
|
|
return reshard_op_cost
|
|
else:
|
|
dist_tensor = self.dist_context.get_dist_tensor_for_program(
|
|
tensor
|
|
)
|
|
# simplified processing: ignore union process mesh and output reshard
|
|
dist_op = self.dist_context.get_dist_op_for_program(op)
|
|
if not dist_tensor or not dist_op:
|
|
return reshard_op_cost
|
|
dims_mapping = dist_op.dist_attr.get_input_dims_mapping(
|
|
tensor.name
|
|
)
|
|
process_mesh = dist_op.dist_attr.process_mesh
|
|
dist_attr = [
|
|
process_mesh,
|
|
dims_mapping,
|
|
dist_op.dist_attr.chunk_id,
|
|
op.attr('op_role'),
|
|
]
|
|
if dist_tensor is not None and self.need_reshard(
|
|
dist_tensor, dist_attr
|
|
):
|
|
if tensor_name not in self._has_resharded:
|
|
self._has_resharded[tensor_name] = [dist_op]
|
|
else:
|
|
for item in self._has_resharded[tensor_name]:
|
|
item_dist_attr = item.dist_attr
|
|
item_dims_mapping = (
|
|
item_dist_attr.get_input_dims_mapping(
|
|
tensor_name
|
|
)
|
|
)
|
|
item_process_mesh = item_dist_attr.process_mesh
|
|
if (
|
|
dims_mapping == item_dims_mapping
|
|
and item_process_mesh == process_mesh
|
|
):
|
|
return reshard_op_cost
|
|
self._has_resharded[tensor_name].append(dist_op)
|
|
|
|
reshard_op_desc = self.find_op_desc_seq(
|
|
dist_tensor, dist_attr, serial=True
|
|
)
|
|
dtype = dist_tensor.serial_tensor.dtype
|
|
reshard_op_cost = self.parse_op_desc_for_cost(
|
|
reshard_op_desc, dtype, cluster
|
|
)
|
|
|
|
return reshard_op_cost
|
|
|
|
def _concat_partitions_for_cost(
|
|
self,
|
|
partition_tensor_list,
|
|
partition_index,
|
|
dtype,
|
|
rank_id,
|
|
local_rank_comp_cost,
|
|
cluster,
|
|
):
|
|
if not partition_tensor_list:
|
|
partition_tensor_list.append(partition_index)
|
|
else:
|
|
i = 0
|
|
has_concat = False
|
|
while i < len(partition_tensor_list):
|
|
(
|
|
concat_axis,
|
|
first_order,
|
|
new_partition,
|
|
) = Resharder.compute_concat_info(
|
|
partition_tensor_list[i], partition_index
|
|
)
|
|
if concat_axis != -1:
|
|
has_concat = True
|
|
concat_desc = {}
|
|
concat_desc["op"] = "concat"
|
|
concat_desc["attrs"] = {"axis": concat_axis}
|
|
if first_order == 0:
|
|
concat_desc["inputs"] = {
|
|
"X": [
|
|
(dtype, partition_tensor_list[i]),
|
|
(dtype, partition_index),
|
|
]
|
|
}
|
|
else:
|
|
concat_desc["inputs"] = {
|
|
"X": [
|
|
(dtype, partition_index),
|
|
(dtype, partition_tensor_list[i]),
|
|
]
|
|
}
|
|
partition_tensor_list.pop(i)
|
|
if rank_id not in local_rank_comp_cost:
|
|
local_rank_comp_cost[rank_id] = []
|
|
concat_desc["dtype"] = dtype
|
|
local_rank_comp_cost[rank_id].append(
|
|
ConcatOpCost(
|
|
op_desc=concat_desc, cluster=cluster, rank=rank_id
|
|
)
|
|
)
|
|
self._concat_partitions_for_cost(
|
|
partition_tensor_list,
|
|
new_partition,
|
|
dtype,
|
|
rank_id,
|
|
local_rank_comp_cost,
|
|
cluster,
|
|
)
|
|
break
|
|
i += 1
|
|
if not has_concat:
|
|
partition_tensor_list.append(partition_index)
|
|
|
|
def parse_op_desc_for_cost(self, reshard_op_desc, dtype, cluster):
|
|
def _get_idx(comm_ranks, group_ranks):
|
|
res, is_the_same = None, False
|
|
idx = 0
|
|
while idx < len(comm_ranks):
|
|
if comm_ranks[idx] == set(group_ranks):
|
|
is_the_same = True
|
|
|
|
for rank in group_ranks:
|
|
if rank in comm_ranks[idx]:
|
|
res = idx
|
|
comm_ranks[idx].add(rank)
|
|
if res is None:
|
|
idx += 1
|
|
else:
|
|
break
|
|
return res, is_the_same
|
|
|
|
comm_context = CommContext(cluster)
|
|
# run communication op before computation op
|
|
# TODO: Communication cost is not calculated when the var has been transferred by the same group in the past
|
|
comm_costs = []
|
|
comm_ranks = []
|
|
local_rank_comp_cost = {}
|
|
for key in reshard_op_desc:
|
|
partition_tensor_list = []
|
|
op_desc_list = reshard_op_desc[key]
|
|
for op_desc in op_desc_list:
|
|
if isinstance(op_desc, SendOpDesc):
|
|
group_ranks = [key, op_desc.dst]
|
|
shape = op_desc.shape
|
|
send_desc = build_comm_desc(
|
|
"send_v2", group_ranks, dtype, shape
|
|
)
|
|
idx, is_the_same = _get_idx(comm_ranks, group_ranks)
|
|
if idx is None:
|
|
comm_costs.append(
|
|
[
|
|
(
|
|
group_ranks,
|
|
SendOpCost(
|
|
op_desc=send_desc,
|
|
comm_context=comm_context,
|
|
),
|
|
)
|
|
]
|
|
)
|
|
comm_ranks.append(set(group_ranks))
|
|
else:
|
|
if not is_the_same:
|
|
comm_costs[idx].append(
|
|
(
|
|
group_ranks,
|
|
SendOpCost(
|
|
op_desc=send_desc,
|
|
comm_context=comm_context,
|
|
),
|
|
)
|
|
)
|
|
elif isinstance(op_desc, AllGatherOpDesc):
|
|
# NOTE: fill_const and other unnecessary op is not calculated because those cost is very small
|
|
group_ranks = op_desc.group
|
|
shape = op_desc.shape
|
|
allgather_desc = build_comm_desc(
|
|
"all_gather", group_ranks, dtype, shape
|
|
)
|
|
split_inputs_shape = []
|
|
for idx, dim in enumerate(shape):
|
|
if idx == 0:
|
|
split_inputs_shape.append(dim * len(group_ranks))
|
|
else:
|
|
split_inputs_shape.append(dim)
|
|
idx, is_the_same = _get_idx(comm_ranks, group_ranks)
|
|
if idx is None:
|
|
comm_costs.append(
|
|
[
|
|
(
|
|
group_ranks,
|
|
AllgatherOpCost(
|
|
op_desc=allgather_desc,
|
|
comm_context=comm_context,
|
|
),
|
|
)
|
|
]
|
|
)
|
|
comm_ranks.append(set(group_ranks))
|
|
else:
|
|
if not is_the_same:
|
|
comm_costs[idx].append(
|
|
(
|
|
group_ranks,
|
|
AllgatherOpCost(
|
|
op_desc=allgather_desc,
|
|
comm_context=comm_context,
|
|
),
|
|
)
|
|
)
|
|
# calc the split op cost
|
|
if key not in local_rank_comp_cost:
|
|
local_rank_comp_cost[key] = []
|
|
split_desc = {}
|
|
split_desc["op"] = "split"
|
|
split_desc["inputs"] = {
|
|
"inputs": [(dtype, split_inputs_shape)]
|
|
}
|
|
split_desc["attrs"] = {"num": len(group_ranks), "axis": 0}
|
|
split_desc["dtype"] = dtype
|
|
local_rank_comp_cost[key].append(
|
|
SplitOpCost(
|
|
op_desc=split_desc, cluster=cluster, rank=key
|
|
)
|
|
)
|
|
elif isinstance(op_desc, ConcatOpDesc):
|
|
partition_index_list = op_desc._partition_index_list
|
|
for idx, partition_idex in enumerate(partition_index_list):
|
|
self._concat_partitions_for_cost(
|
|
partition_tensor_list,
|
|
partition_idex,
|
|
dtype,
|
|
key,
|
|
local_rank_comp_cost,
|
|
cluster,
|
|
)
|
|
|
|
elif isinstance(op_desc, SliceOpDesc):
|
|
if key not in local_rank_comp_cost:
|
|
local_rank_comp_cost[key] = []
|
|
assert (
|
|
len(partition_tensor_list) == 1
|
|
or not partition_tensor_list
|
|
)
|
|
to_slice_tensor_shape = []
|
|
if len(partition_tensor_list) == 1:
|
|
for item in partition_tensor_list[0]:
|
|
to_slice_tensor_shape.append(item[1] - item[0])
|
|
else:
|
|
to_slice_tensor_shape = op_desc.shape
|
|
slice_desc = {}
|
|
slice_desc["op"] = "slice"
|
|
infer_flags = [1 for i in range(len(op_desc.axes))]
|
|
slice_desc["attrs"] = {
|
|
"axes": op_desc.axes,
|
|
"starts": op_desc.starts,
|
|
"ends": op_desc.ends,
|
|
"infer_flags": infer_flags,
|
|
}
|
|
slice_desc["inputs"] = {
|
|
"Input": [(dtype, to_slice_tensor_shape)]
|
|
}
|
|
slice_desc["dtype"] = dtype
|
|
local_rank_comp_cost[key].append(
|
|
SliceOpCost(
|
|
op_desc=slice_desc, cluster=cluster, rank=key
|
|
)
|
|
)
|
|
|
|
res = (comm_costs, local_rank_comp_cost)
|
|
|
|
return res
|