wehub-resource-sync
2004 lines
77 KiB
Python
2004 lines
77 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 logging
|
|
from functools import reduce
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from paddle.distributed.auto_parallel.static.operators.common import (
|
|
ParallelMode,
|
|
is_data_parallel_reduce_op,
|
|
is_parameter_related,
|
|
)
|
|
from paddle.distributed.auto_parallel.static.process_group import (
|
|
new_process_group,
|
|
)
|
|
from paddle.distributed.auto_parallel.static.utils import (
|
|
_get_comm_group,
|
|
get_logger,
|
|
get_var_numel,
|
|
insert_dependencies_for_vars,
|
|
is_backward_op,
|
|
is_dep_skip_op,
|
|
is_forward_op,
|
|
is_optimize_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,
|
|
)
|
|
from paddle.distributed.fleet.meta_optimizers.sharding.utils import get_var_size
|
|
from paddle.framework import core
|
|
from paddle.static import default_main_program, default_startup_program
|
|
from paddle.utils import unique_name
|
|
|
|
from .auto_parallel_master_grad import _is_master_grad_cast_op
|
|
from .pass_base import PassBase, register_pass
|
|
from .pass_utils import AutoParallelStreamType
|
|
|
|
OpRole = core.op_proto_and_checker_maker.OpRole
|
|
OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
|
|
_skip_ops = [
|
|
'create_py_reader',
|
|
'create_double_buffer_reader',
|
|
'read',
|
|
'slice',
|
|
'split',
|
|
'assign',
|
|
"send_v2",
|
|
]
|
|
# update here to support new optimizers
|
|
_supported_optimizer_type = [
|
|
"adam",
|
|
"adamax",
|
|
"adamw",
|
|
"decayed_adagrad",
|
|
"momentum",
|
|
"dgc_momentum",
|
|
"lars_momentum",
|
|
"merged_momentum",
|
|
"lamb",
|
|
"sgd",
|
|
]
|
|
|
|
_logger = get_logger(logging.INFO)
|
|
|
|
__amp_target_dtype__ = core.VarDesc.VarType.FP16
|
|
__amp_target_dtype_name__ = "float16"
|
|
|
|
|
|
def _is_reshard_op(op):
|
|
return op.desc.has_attr(
|
|
"op_namescope"
|
|
) and "/auto_parallel/reshard" in op.desc.attr('op_namescope')
|
|
|
|
|
|
# NOTE we add the "auto_parallel" prefix to the pass in order to
|
|
# indicate that this pass should obey some constrains by auto_parallel
|
|
# for example all ops and vars should has dist attr before and after pass
|
|
# should use dist op instead of custom comm op
|
|
@register_pass("auto_parallel_sharding")
|
|
class ShardingPass(PassBase):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.set_attr("dist_context", None)
|
|
self.set_attr("stage", None)
|
|
self.set_attr("sharding_degree", None) # for parallelizer
|
|
self.set_attr("degree", None) # for parallelizer_v2
|
|
self.set_attr("enable_overlap", None)
|
|
self.set_attr("param_comm_stream_num", None)
|
|
self.set_attr("grad_comm_stream_num", None)
|
|
self.set_attr("param_bucket_size_numel", None)
|
|
self.set_attr("grad_bucket_size_numel", None)
|
|
self.set_attr("partition_algor", None)
|
|
self.set_attr("enable_hierarchical_comm", None)
|
|
self.set_attr("params_grads", [])
|
|
self.set_attr("global_rank", -1)
|
|
self.set_attr("amp_dtype", "float16")
|
|
self.set_attr("gradient_sync_after_accumulate", False)
|
|
self.dp_groups = set()
|
|
self.sharding_infos = []
|
|
self.varname_to_sharding_info = {}
|
|
self.sharding_hybrid_dp = False
|
|
self.outer_dp_group = None
|
|
self.shared_params_grads = []
|
|
|
|
def _check_self(self):
|
|
if self.get_attr("dist_context") is None:
|
|
return False
|
|
|
|
if self.get_attr("stage") not in [1, 2, 3]:
|
|
return False
|
|
if self.get_attr("sharding_degree") is not None:
|
|
if (
|
|
not isinstance(self.get_attr("sharding_degree"), int)
|
|
) or self.get_attr("sharding_degree") <= 1:
|
|
return False
|
|
elif self.get_attr("degree") is not None:
|
|
if (not isinstance(self.get_attr("degree"), int)) or self.get_attr(
|
|
"degree"
|
|
) <= 1:
|
|
return False
|
|
else:
|
|
return False
|
|
if len(self.get_attr("params_grads")) <= 0:
|
|
return False
|
|
if (not isinstance(self.get_attr("global_rank"), int)) or self.get_attr(
|
|
"global_rank"
|
|
) < 0:
|
|
return False
|
|
if self.get_attr("enable_overlap") is None:
|
|
return False
|
|
if self.get_attr("param_comm_stream_num") is None:
|
|
return False
|
|
if self.get_attr("grad_comm_stream_num") is None:
|
|
return False
|
|
if self.get_attr("param_bucket_size_numel") is None:
|
|
return False
|
|
if self.get_attr("grad_bucket_size_numel") is None:
|
|
return False
|
|
if self.get_attr("partition_algor") is None:
|
|
return False
|
|
if self.get_attr("enable_hierarchical_comm") is None:
|
|
return False
|
|
return True
|
|
|
|
def _check_conflict(self, other_pass):
|
|
return True
|
|
|
|
def _apply_single_impl(self, main_program, startup_program, context):
|
|
self._dist_context = self.get_attr("dist_context")
|
|
self.sharding_world_size = int(
|
|
self.get_attr("sharding_degree") or self.get_attr("degree")
|
|
)
|
|
self.stage = int(self.get_attr("stage"))
|
|
self.global_rank = int(self.get_attr("global_rank"))
|
|
self.enable_overlap = self.get_attr("enable_overlap")
|
|
self.param_comm_stream_num = int(self.get_attr("param_comm_stream_num"))
|
|
self.grad_comm_stream_num = int(self.get_attr("grad_comm_stream_num"))
|
|
self.enable_hierarchical_comm = self.get_attr(
|
|
"enable_hierarchical_comm"
|
|
)
|
|
if self.param_comm_stream_num > 1 or self.grad_comm_stream_num > 1:
|
|
assert self.enable_overlap, (
|
|
"multiple comm stream need enable_overlap to be True"
|
|
)
|
|
self.param_bucket_size_numel = int(
|
|
self.get_attr("param_bucket_size_numel")
|
|
)
|
|
self.grad_bucket_size_numel = int(
|
|
self.get_attr("grad_bucket_size_numel")
|
|
)
|
|
self.partition_algor = self.get_attr("partition_algor")
|
|
|
|
params_grads = self.get_attr("params_grads")
|
|
main_block, startup_block = (
|
|
main_program.global_block(),
|
|
startup_program.global_block(),
|
|
)
|
|
|
|
self.amp_dtype = self.get_attr("amp_dtype")
|
|
if self.amp_dtype == "bfloat16":
|
|
__amp_target_dtype__ = core.VarDesc.VarType.BF16
|
|
__amp_target_dtype_name__ = "bfloat16"
|
|
|
|
# NOTE Multi / Sub-Block Support
|
|
# we assume that only parameter are present and partitioned in main_block,
|
|
# there is NO new param in sub_block, and all params in sub_block follows the same
|
|
# partition as main_block. the above constraint fulfill the 3 most common use-cases in Paddle sub_block:
|
|
# 1. subblock for lr scheduler
|
|
# 2. sub-block uses the same or partial network of main-block, e.g. GPT3 generation model
|
|
# 3. sub-block used for double backward
|
|
|
|
self._build_sharding_groups(main_block, params_grads)
|
|
for block in main_program.blocks:
|
|
self._shard_optimizer(block, startup_block)
|
|
self._shard_gradient_synchronization(block)
|
|
self._shard_parameter(block, startup_block)
|
|
|
|
context.set_attr("params_grads", self.shared_params_grads)
|
|
self._optimization_pass(main_program, startup_program)
|
|
|
|
def _build_sharding_groups(self, main_block, params_grads):
|
|
self._collective_data_parallel_groups(main_block)
|
|
self._build_sharding_infos(main_block, params_grads)
|
|
|
|
def _collective_data_parallel_groups(self, main_block):
|
|
for op in main_block.ops:
|
|
if not is_forward_op(op) or op.type in _skip_ops:
|
|
continue
|
|
# NOTE: there aren't dist_attr in the ops which reshard insert,
|
|
# and should be skip in sharding.
|
|
if _is_reshard_op(op):
|
|
continue
|
|
group = _inference_data_parallel_group_for_operator(
|
|
self.global_rank, op, self._dist_context, 0
|
|
)
|
|
if group is not None:
|
|
self.dp_groups.add(group)
|
|
|
|
# TODO(JZ-LIANG) allow more than one dp groups in network, support more general distribution
|
|
# generated by auto search
|
|
if len(self.dp_groups) != 1:
|
|
raise NotImplementedError(
|
|
f"So far Only and Exactly one data parallel group in network are supported, but got [{len(self.dp_groups)}] different data parallel groups"
|
|
)
|
|
|
|
def _build_sharding_infos(self, main_block, params_grads):
|
|
# order params
|
|
params_grads = re_order_program(
|
|
main_block, params_grads, self._dist_context
|
|
)
|
|
|
|
# partition
|
|
for dp_group in self.dp_groups:
|
|
assert dp_group.nranks >= self.sharding_world_size, (
|
|
f"sharding world size [{self.sharding_world_size}] should not larger than dp world size [{dp_group.nranks}]"
|
|
)
|
|
assert dp_group.nranks % self.sharding_world_size == 0, (
|
|
f"sharding world size [{self.sharding_world_size}] should be divisible by dp world size [{dp_group.nranks}]"
|
|
)
|
|
assert self.global_rank in dp_group.ranks, (
|
|
f"current ranks [{self.global_rank}] does NOT belong to the data parallel group [{dp_group.ranks}]"
|
|
)
|
|
assert len(params_grads) >= self.sharding_world_size, (
|
|
f"number of parameters [{len(params_grads)}] is not enough to be shard among [{self.sharding_world_size}] ranks"
|
|
)
|
|
|
|
# sharding hybrid data parallel: partial sharding param within
|
|
if dp_group.nranks > self.sharding_world_size:
|
|
self.sharding_hybrid_dp = True
|
|
assert self.param_comm_stream_num < 2
|
|
assert self.grad_comm_stream_num < 2
|
|
assert len(self.dp_groups) == 1, (
|
|
"hybrid sharding and data parallelism are supported only when there is exactly one data parallel group in the network"
|
|
)
|
|
outer_dp_group, sharding_group = _get_dp_and_sharding_groups(
|
|
dp_group.ranks, self.sharding_world_size, self.global_rank
|
|
)
|
|
sharding_group = new_process_group(sharding_group)
|
|
self.outer_dp_group = new_process_group(outer_dp_group)
|
|
else:
|
|
sharding_group = dp_group
|
|
|
|
self._dist_context._sharding_group = sharding_group
|
|
# TODO(JZ-LIANG) when support multiple dp groups in future, should group param and bind them to corresponding dp group
|
|
sharding_info = ShardingInfo(
|
|
sharding_group,
|
|
self.global_rank,
|
|
params_grads,
|
|
self.partition_algor,
|
|
)
|
|
self.sharding_infos.append(sharding_info)
|
|
for param in sharding_info.params:
|
|
self.varname_to_sharding_info[param.name] = sharding_info
|
|
|
|
def _shard_optimizer(self, main_block, startup_block):
|
|
"""
|
|
sharding all optimizer related ops and vars, include:
|
|
gradient clip ops & vars
|
|
weight decay ops & vars
|
|
optimizer ops and states
|
|
"""
|
|
self._shard_amp_related_op_and_vars(main_block)
|
|
self._shard_weight_decay(main_block)
|
|
# self._shard_gradient_clip(main_block)
|
|
self._shard_optimizer_ops_and_states(main_block, startup_block)
|
|
self._insert_optimizer_broadcasts(main_block, startup_block)
|
|
|
|
def _shard_amp_related_op_and_vars(self, main_block):
|
|
if self.stage < 1:
|
|
return
|
|
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
# shard amp related param_grad cast
|
|
if _is_param_grad_fp32_cast_op(main_block, op) and self.stage > 1:
|
|
output_name = op.output_arg_names[0]
|
|
param_name = output_name[: output_name.find("@")]
|
|
if not self._is_parameter_in_local_shard(param_name):
|
|
main_block._remove_op(idx, sync=False)
|
|
main_block._remove_var(output_name, sync=False)
|
|
|
|
# shard check nan inf
|
|
elif op.type in ["check_finite_and_unscale", "update_loss_scaling"]:
|
|
reversed_x = []
|
|
for input_name in op.desc.input('X'):
|
|
param_name = input_name[: input_name.find("@")]
|
|
|
|
if self._is_parameter_in_local_shard(param_name):
|
|
reversed_x.append(input_name)
|
|
|
|
# NOTE: When `reversed_x` is [], check_finite_and_unscale will be replaced by `fill_constant` op.
|
|
# The output of check_finite_and_unscale is be set False
|
|
if reversed_x:
|
|
op.desc.set_input('X', reversed_x)
|
|
op.desc.set_output('Out', reversed_x)
|
|
else:
|
|
if op.type == "check_finite_and_unscale":
|
|
op_role = op.attr('op_role')
|
|
out_name = op.output_arg_names[0]
|
|
out_var = main_block.vars[out_name]
|
|
main_block._remove_op(idx, sync=False)
|
|
main_block._insert_op_without_sync(
|
|
idx,
|
|
type="fill_constant",
|
|
outputs={"Out": out_var},
|
|
attrs={
|
|
"shape": out_var.shape,
|
|
"dtype": out_var.dtype,
|
|
"value": 0,
|
|
OP_ROLE_KEY: op_role,
|
|
},
|
|
)
|
|
dist_attr = (
|
|
self._dist_context.get_tensor_dist_attr_for_program(
|
|
out_var
|
|
)
|
|
)
|
|
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
main_block.ops[idx],
|
|
dist_attr.process_mesh,
|
|
dist_attr.dims_mapping,
|
|
self._dist_context,
|
|
chunk_id=dist_attr.chunk_id,
|
|
)
|
|
|
|
else:
|
|
main_block._remove_op(idx, sync=False)
|
|
|
|
main_block._sync_with_cpp()
|
|
|
|
def _shard_gradient_clip(self, main_block):
|
|
if self.stage < 2:
|
|
return
|
|
|
|
# TODO (JZ-LIANG) support calculate global norm with tensor parallelism
|
|
removed_op_type = ['elementwise_mul', 'squared_l2_norm', 'clip_by_norm']
|
|
removed_op_idx = set()
|
|
removed_tmp_var = set()
|
|
|
|
for idx, op in list(enumerate(main_block.ops)):
|
|
if not _is_gradient_clip_op(op):
|
|
continue
|
|
|
|
if op.type in removed_op_type:
|
|
input_name = op.input("X")[0]
|
|
param_name = input_name[: input_name.find("@GRAD")]
|
|
if not self._is_parameter_in_local_shard(param_name):
|
|
removed_op_idx.add(idx)
|
|
if op.type in ['squared_l2_norm', 'clip_by_norm']:
|
|
for output_name in op.output_arg_names:
|
|
removed_tmp_var.add(output_name)
|
|
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if not _is_gradient_clip_op(op):
|
|
continue
|
|
if idx in removed_op_idx:
|
|
main_block._remove_op(idx, sync=False)
|
|
|
|
for varname in removed_tmp_var:
|
|
main_block._remove_var(varname, sync=False)
|
|
|
|
for idx, op in list(enumerate(main_block.ops)):
|
|
if not _is_gradient_clip_op(op):
|
|
continue
|
|
if op.type == 'sum':
|
|
reserved_vars = []
|
|
for input_name in op.input_arg_names:
|
|
if input_name not in removed_tmp_var:
|
|
reserved_vars.append(input_name)
|
|
op.desc.set_input("X", reserved_vars)
|
|
|
|
sum_op_output = op.output_arg_names[0]
|
|
for i, sharding_info in enumerate(self.sharding_infos):
|
|
new_op = main_block._insert_op(
|
|
idx + i + 1,
|
|
type='all_reduce',
|
|
inputs={'x': [sum_op_output]},
|
|
outputs={'out': [sum_op_output]},
|
|
attrs={
|
|
'ring_id': sharding_info.group.id,
|
|
'op_namescope': "/gradient_clip_model_parallelism",
|
|
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
dist_attr = (
|
|
self._dist_context.get_tensor_dist_attr_for_program(
|
|
main_block.var(sum_op_output)
|
|
)
|
|
)
|
|
# assert dist_attr is not None
|
|
# naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
# new_op, dist_attr.process_mesh, dist_attr.dims_mapping,
|
|
# self._dist_context)
|
|
break
|
|
|
|
main_block._sync_with_cpp()
|
|
|
|
def _shard_weight_decay(self, main_block):
|
|
if self.stage < 2:
|
|
return
|
|
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if not _is_weight_decay_op(op):
|
|
continue
|
|
else:
|
|
raise NotImplementedError(
|
|
"weight decay is NOT supported by now"
|
|
)
|
|
main_block._sync_with_cpp()
|
|
|
|
def _shard_optimizer_ops_and_states(self, main_block, startup_block):
|
|
should_removed_optimizer_states = []
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if not is_optimize_op(op):
|
|
break
|
|
|
|
if op.type in _supported_optimizer_type:
|
|
assert "Param" in op.input_names
|
|
assert len(op.input("Param")) == 1
|
|
param_name = op.input("Param")[0]
|
|
if not self._is_parameter_in_local_shard(param_name):
|
|
should_removed_optimizer_states.extend(
|
|
[
|
|
varname
|
|
for varname in op.output_arg_names
|
|
if varname != param_name
|
|
]
|
|
)
|
|
main_block._remove_op(idx, sync=False)
|
|
else:
|
|
self.shared_params_grads.append(
|
|
self._get_param_grad(param_name)
|
|
)
|
|
|
|
for idx, op in reversed(list(enumerate(startup_block.ops))):
|
|
if (
|
|
len(op.output_arg_names) == 1
|
|
and op.output_arg_names[0] in should_removed_optimizer_states
|
|
):
|
|
startup_block._remove_op(idx, sync=False)
|
|
|
|
for varname in should_removed_optimizer_states:
|
|
if main_block.has_var(varname):
|
|
main_block._remove_var(varname, sync=False)
|
|
if startup_block.has_var(varname):
|
|
startup_block._remove_var(varname, sync=False)
|
|
|
|
main_block._sync_with_cpp()
|
|
startup_block._sync_with_cpp()
|
|
|
|
def _insert_optimizer_broadcasts(self, main_block, startup_block):
|
|
if self.stage > 2 or self.param_bucket_size_numel > 1:
|
|
return
|
|
|
|
for sharding_info in self.sharding_infos:
|
|
for param in sharding_info.params:
|
|
assert main_block.has_var(param.name)
|
|
assert startup_block.has_var(param.name)
|
|
|
|
new_op = main_block.append_op(
|
|
type='broadcast',
|
|
inputs={'x': param},
|
|
outputs={'out': param},
|
|
attrs={
|
|
'ring_id': sharding_info.group.id,
|
|
'root': sharding_info.get_var_rank(param.name),
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
new_op._set_attr(
|
|
'op_namescope', '/' + ParallelMode.DataParallel
|
|
)
|
|
param_dist_attr = (
|
|
self._dist_context.get_tensor_dist_attr_for_program(param)
|
|
)
|
|
assert param_dist_attr is not None
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
new_op,
|
|
param_dist_attr.process_mesh,
|
|
param_dist_attr.dims_mapping,
|
|
self._dist_context,
|
|
chunk_id=param_dist_attr.chunk_id,
|
|
)
|
|
main_block._sync_with_cpp()
|
|
|
|
def _is_parameter_in_local_shard(self, param_name):
|
|
assert param_name in self.varname_to_sharding_info
|
|
sharding_info = self.varname_to_sharding_info[param_name]
|
|
return sharding_info.is_in_local_shard(param_name)
|
|
|
|
def _get_param_grad(self, param_name):
|
|
assert param_name in self.varname_to_sharding_info
|
|
sharding_info = self.varname_to_sharding_info[param_name]
|
|
p_g = sharding_info.get_param_grad(param_name)
|
|
assert p_g is not None
|
|
return p_g
|
|
|
|
def _shard_gradient_synchronization(self, main_block):
|
|
if self.stage < 2:
|
|
return
|
|
|
|
dp_ring_ids = [group.id for group in self.dp_groups]
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if _is_param_grad_allreduce_op(op, main_block):
|
|
if (
|
|
op.type == "all_reduce"
|
|
and op.attr("reduce_type") == dist.ReduceOp.SUM
|
|
) or (
|
|
op.type == "reduce"
|
|
and op.attr("reduce_type") == dist.ReduceOp.SUM
|
|
):
|
|
reduce_op_type = "reduce"
|
|
reduce_type = dist.ReduceOp.SUM
|
|
else:
|
|
reduce_op_type = "reduce"
|
|
reduce_type = dist.ReduceOp.AVG
|
|
input_name = op.input_arg_names[0]
|
|
base_name = _get_base_name_from_grad_name(input_name)
|
|
sharding_info = self.varname_to_sharding_info[base_name]
|
|
reduce_op = _insert_reduce_op(
|
|
main_block,
|
|
reduce_op_type,
|
|
idx,
|
|
input_name,
|
|
sharding_info.group.id,
|
|
sharding_info.get_var_rank(base_name),
|
|
self._dist_context,
|
|
reduce_type,
|
|
)
|
|
if (
|
|
not self.sharding_hybrid_dp
|
|
or not sharding_info.is_in_local_shard(base_name)
|
|
):
|
|
main_block._remove_op(idx + 1, sync=False)
|
|
else:
|
|
op._set_attr("ring_id", self.outer_dp_group.id)
|
|
op._set_attr(
|
|
'op_namescope', '/' + ParallelMode.DataParallel
|
|
)
|
|
|
|
# NOTE:
|
|
# var@GRAD = sum(var@GRAD@RENAME@0, var@GRAD@RENAME@1)
|
|
# If the var is not in local rank and it is output of many ops, or the var is renamed in another words,
|
|
# the sum op should be removed.
|
|
if _is_param_grad_sum_op(op, main_block):
|
|
out_name = op.output_arg_names[0]
|
|
base_name = _get_base_name_from_grad_name(out_name)
|
|
sharding_info = self.varname_to_sharding_info[base_name]
|
|
if not sharding_info.is_in_local_shard(base_name):
|
|
main_block._remove_op(idx, sync=False)
|
|
|
|
main_block._sync_with_cpp()
|
|
|
|
def _shard_parameter(self, main_block, startup_block):
|
|
if self.stage < 3:
|
|
return
|
|
|
|
dp_ring_ids = [group.id for group in self.dp_groups]
|
|
for sharding_info in self.sharding_infos:
|
|
(
|
|
need_broadcast_vars,
|
|
param_usage,
|
|
) = sharding_info.get_broadcast_vars_and_param_usage(main_block)
|
|
not_used_param_name = []
|
|
for param_name in param_usage:
|
|
if (
|
|
param_usage[param_name] == 0
|
|
and sharding_info.get_var_rank(param_name)
|
|
!= sharding_info.local_rank
|
|
):
|
|
not_used_param_name.append(param_name)
|
|
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if is_optimize_op(op):
|
|
continue
|
|
|
|
for input_name in op.input_arg_names:
|
|
# NOTE hack for embedding op when AMP 02-3
|
|
# paddle amp force embedding (lookup table) to be run on fp32
|
|
if _is_param_fp16_cast_op(
|
|
main_block, op, sharding_info.param_names
|
|
):
|
|
# NOTE:
|
|
# param.cast_fp16 = cast(param)
|
|
# When param is not in current rank, the cast op need to be removed.
|
|
if not self._is_parameter_in_local_shard(input_name):
|
|
not_used_param_name.append(input_name)
|
|
continue
|
|
if input_name not in need_broadcast_vars:
|
|
continue
|
|
root_rank = sharding_info.get_var_rank(input_name)
|
|
if root_rank == sharding_info.local_rank:
|
|
broadcast_varname = input_name
|
|
else:
|
|
broadcast_varname = unique_name.generate(
|
|
input_name + "@BroadCast"
|
|
)
|
|
input_var = main_block.var(input_name)
|
|
new_var = main_block.create_var(
|
|
name=broadcast_varname,
|
|
shape=input_var.shape,
|
|
dtype=input_var.dtype,
|
|
persistable=False,
|
|
)
|
|
ref_dist_attr = (
|
|
self._dist_context.get_tensor_dist_attr_for_program(
|
|
input_var
|
|
)
|
|
)
|
|
set_var_dist_attr(
|
|
self._dist_context,
|
|
new_var,
|
|
ref_dist_attr.dims_mapping,
|
|
ref_dist_attr.process_mesh,
|
|
chunk_id=ref_dist_attr.chunk_id,
|
|
)
|
|
op_dist_attr = (
|
|
self._dist_context.get_op_dist_attr_for_program(op)
|
|
)
|
|
input_dist_attr = op_dist_attr.get_input_dist_attr(
|
|
input_name
|
|
)
|
|
op._rename_input(input_name, broadcast_varname)
|
|
op_dist_attr.set_input_dist_attr(
|
|
broadcast_varname, input_dist_attr
|
|
)
|
|
|
|
_insert_init_and_broadcast_op(
|
|
main_block,
|
|
idx,
|
|
broadcast_varname,
|
|
sharding_info.local_rank,
|
|
root_rank,
|
|
sharding_info.group.id,
|
|
op.attr('op_role'),
|
|
self._dist_context,
|
|
)
|
|
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if op.type != "cast":
|
|
continue
|
|
input_name = op.input_arg_names[0]
|
|
output_name = op.output_arg_names[0]
|
|
if input_name in not_used_param_name:
|
|
main_block._remove_op(idx, sync=False)
|
|
main_block._remove_var(output_name, sync=False)
|
|
|
|
for idx, op in reversed(list(enumerate(startup_block.ops))):
|
|
assert len(op.output_arg_names) == 1
|
|
output_name = op.output_arg_names[0]
|
|
|
|
if op.type == "broadcast":
|
|
if op.attr("ring_id") in dp_ring_ids:
|
|
if (
|
|
self.outer_dp_group
|
|
and sharding_info.get_var_rank(output_name)
|
|
== sharding_info.local_rank
|
|
):
|
|
op._set_attr("ring_id", self.outer_dp_group.id)
|
|
else:
|
|
startup_block._remove_op(idx, sync=False)
|
|
else: # We should remove the `broadcast` between `TensorParallel` mesh dim.
|
|
if (
|
|
sharding_info.get_var_rank(output_name)
|
|
!= sharding_info.local_rank
|
|
):
|
|
startup_block._remove_op(idx, sync=False)
|
|
continue
|
|
|
|
if (
|
|
op.type != "broadcast"
|
|
and output_name in param_usage
|
|
and sharding_info.get_var_rank(output_name)
|
|
!= sharding_info.local_rank
|
|
):
|
|
startup_block._remove_op(idx, sync=False)
|
|
|
|
for param_name in param_usage:
|
|
if (
|
|
sharding_info.get_var_rank(param_name)
|
|
!= sharding_info.local_rank
|
|
):
|
|
main_block._remove_var(param_name, sync=False)
|
|
startup_block._remove_var(param_name, sync=False)
|
|
|
|
main_block._sync_with_cpp()
|
|
startup_block._sync_with_cpp()
|
|
|
|
def _optimization_pass(self, main_program, startup_program):
|
|
if self.stage <= 1:
|
|
return
|
|
|
|
self.grad_coalesce_prefix = 'sharding_coalesce_grad_'
|
|
self.param_coalesce_prefix = 'sharding_coalesce_param_'
|
|
# NOTE PR#49275 for detail
|
|
self.comm_op_scheduling_priority = -1
|
|
|
|
# TODO support multiple sub_blocks
|
|
assert len(self.sharding_infos) == 1, (
|
|
f"gradient synchronization optimization only support one sharding group right now, but got [{len(self.sharding_infos)}]."
|
|
)
|
|
sharding_info = self.sharding_infos[0]
|
|
|
|
with paddle.static.program_guard(main_program, startup_program):
|
|
self._gradient_sync_optimization(sharding_info)
|
|
# TODO independent the logic of fuse and overlap
|
|
# support overlap when no fuse
|
|
if self.param_bucket_size_numel > 1:
|
|
if self.stage == 2:
|
|
self._fuse_overlap_parameter_comm_stage_two(sharding_info)
|
|
elif self.stage == 3:
|
|
self._fuse_overlap_parameter_comm_stage_three(sharding_info)
|
|
|
|
def _gradient_sync_optimization(self, sharding_info):
|
|
if self.grad_bucket_size_numel <= 1 and (not self.enable_overlap):
|
|
return
|
|
|
|
main_block = default_main_program().global_block()
|
|
startup_block = default_startup_program().global_block()
|
|
coalesce_to_group_map, grad_name_to_group_map = self._group_grads(
|
|
main_block,
|
|
sharding_info,
|
|
)
|
|
self._overlap_grad_comm(
|
|
main_block,
|
|
sharding_info,
|
|
coalesce_to_group_map,
|
|
grad_name_to_group_map,
|
|
)
|
|
|
|
def _fuse_overlap_parameter_comm_stage_two(self, sharding_info):
|
|
main_block = default_main_program().global_block()
|
|
startup_block = default_startup_program().global_block()
|
|
|
|
group_to_param_map, param_to_group_map = group_param(
|
|
sharding_info, self.param_bucket_size_numel
|
|
)
|
|
_logger.info("Sharding Stage2 Optimization:")
|
|
_logger.info(
|
|
f"Param Bucket size is [{self.param_bucket_size_numel}], [{len(param_to_group_map.keys())}] Parameters are fused into [{len(group_to_param_map.keys())}] Buckets"
|
|
)
|
|
broadcast_var_to_group_map = {}
|
|
|
|
if self.enable_overlap:
|
|
# if the communication is cross node, comm will be slow and calc will therefore
|
|
# wait for comm. enable multi-comm-stream
|
|
# TODO revise me in future
|
|
# 1. manager the comm and corresponding stream
|
|
# 2. allow more than two streams and open to be config
|
|
self.param_comm_group_stream_pairs = []
|
|
ranks = sharding_info.group.ranks
|
|
for i in range(self.param_comm_stream_num):
|
|
if i == 0:
|
|
group = sharding_info.group
|
|
else:
|
|
group = new_process_group(ranks, force_new_group=True)
|
|
|
|
self.param_comm_group_stream_pairs.append(
|
|
{
|
|
"comm_group": group,
|
|
"comm_stream": AutoParallelStreamType.SHARDING_STREAM.value,
|
|
}
|
|
)
|
|
_logger.info(
|
|
f"Parameter Communication would use [{self.param_comm_stream_num}] streams."
|
|
)
|
|
self.op_to_stream_idx = {}
|
|
|
|
for i, param_group in enumerate(group_to_param_map.keys()):
|
|
assert len(param_group) >= 1
|
|
if len(param_group) > 1:
|
|
coalesce_var_name = unique_name.generate(
|
|
self.param_coalesce_prefix + str(i)
|
|
)
|
|
startup_block.create_var(
|
|
name=coalesce_var_name,
|
|
dtype=param_group.dtype,
|
|
persistable=True,
|
|
stop_gradient=True,
|
|
)
|
|
param_group.coalesce_var = main_block.create_var(
|
|
name=coalesce_var_name,
|
|
dtype=param_group.dtype,
|
|
persistable=True,
|
|
stop_gradient=True,
|
|
)
|
|
startup_block.append_op(
|
|
type="coalesce_tensor",
|
|
inputs={"Input": param_group.vars},
|
|
outputs={
|
|
"Output": param_group.vars,
|
|
"FusedOutput": param_group.coalesce_var,
|
|
},
|
|
attrs={
|
|
"copy_data": True,
|
|
"use_align": True,
|
|
"dtype": param_group.dtype,
|
|
OP_ROLE_KEY: OpRole.Forward,
|
|
},
|
|
)
|
|
else:
|
|
param_group.coalesce_var = param_group.vars[0]
|
|
_logger.info(
|
|
f"Bucket[{i}] size [{sum([get_var_size(p) for p in param_group.vars])}]MB."
|
|
)
|
|
_logger.debug(
|
|
f"Bucket[{i}] parameters: {[p.name for p in param_group.vars]}."
|
|
)
|
|
|
|
broadcast_var_to_group_map[param_group.coalesce_var.name] = (
|
|
param_group
|
|
)
|
|
|
|
# TODO revise me to manager stream and comm
|
|
comm_stream_idx = i % self.param_comm_stream_num
|
|
comm_group = self.param_comm_group_stream_pairs[comm_stream_idx][
|
|
'comm_group'
|
|
]
|
|
comm_stream = self.param_comm_group_stream_pairs[comm_stream_idx][
|
|
'comm_stream'
|
|
]
|
|
new_op = main_block.append_op(
|
|
type='broadcast',
|
|
inputs={'x': param_group.coalesce_var},
|
|
outputs={'out': param_group.coalesce_var},
|
|
attrs={
|
|
'ring_id': comm_group.id,
|
|
'root': param_group.rank,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
self.op_to_stream_idx[new_op] = comm_stream_idx
|
|
new_op._set_attr('op_namescope', '/' + ParallelMode.DataParallel)
|
|
if self.enable_overlap:
|
|
new_op.dist_attr.execution_stream = comm_stream
|
|
new_op.dist_attr.scheduling_priority = (
|
|
self.comm_op_scheduling_priority
|
|
)
|
|
|
|
# NOTE the current dist context lack the presentation for bucket tensor which
|
|
# composes many tensor with different dims_mapping. we DO NOT assign dist attr
|
|
# for it currently.
|
|
|
|
# add dependencies:
|
|
# 1. all broadcast depend on its pre collective
|
|
# 2. coalesce broadcast add nop to resolute data flow dependencies
|
|
dep_map = {}
|
|
for i, op in enumerate(main_block.ops):
|
|
if is_sharding_param_broadcast_op(op):
|
|
broadcast_varname = op.output("Out")[0]
|
|
broadcast_var = main_block.vars[broadcast_varname]
|
|
param_group = broadcast_var_to_group_map[broadcast_varname]
|
|
comm_stream = None
|
|
if self.enable_overlap:
|
|
comm_stream = op.dist_attr.execution_stream
|
|
|
|
# FIXME remove me when upgrade to multi-comm version
|
|
if len(dep_map.keys()) < self.param_comm_stream_num:
|
|
op = _get_broadcast_first_depend_op(main_block)
|
|
prior_var = main_block.vars[op.output("ParamOut")[0]]
|
|
else:
|
|
pre_op = main_block.ops[i - self.param_comm_stream_num]
|
|
assert is_sharding_param_broadcast_op(pre_op), (
|
|
"Unexpected: sharding broadcast pre op should be broadcast."
|
|
)
|
|
prior_var = main_block.vars[pre_op.output("Out")[0]]
|
|
# broadcast order dependencies
|
|
dep_map[i] = [(i, [prior_var], [broadcast_var], comm_stream)]
|
|
|
|
if len(param_group.vars) > 1:
|
|
# in shard coalesce depend to optimizer
|
|
if param_group.is_in_local_shard:
|
|
last_grad = param_group.vars[-1]
|
|
dep_map[i].append(
|
|
(i, [last_grad], [broadcast_var], comm_stream)
|
|
)
|
|
# coalesce resolution post deps
|
|
dep_map[i].append(
|
|
(i + 1, [broadcast_var], param_group.vars, comm_stream)
|
|
)
|
|
|
|
# insert deps
|
|
indice = sorted(dep_map.keys(), reverse=True)
|
|
for i in indice:
|
|
for idx, prior_vars, post_vars, comm_stream in dep_map[i][::-1]:
|
|
depend_op = insert_dependencies_for_vars(
|
|
main_block,
|
|
idx,
|
|
prior_vars,
|
|
post_vars,
|
|
self._dist_context,
|
|
OpRole.Optimize,
|
|
process_mesh=[
|
|
-1
|
|
], # hack to avoid initialize the dist attr for coalesce var
|
|
is_recompute=False,
|
|
sync=False,
|
|
op_namescope="sharding_stage2_broadcast_dep",
|
|
)
|
|
if self.enable_overlap and depend_op is not None:
|
|
depend_op.dist_attr.execution_stream = comm_stream
|
|
depend_op.dist_attr.scheduling_priority = (
|
|
self.comm_op_scheduling_priority
|
|
)
|
|
|
|
main_block._sync_with_cpp()
|
|
|
|
def _fuse_overlap_parameter_comm_stage_three(self, sharding_info):
|
|
pass
|
|
|
|
def _group_grads(
|
|
self,
|
|
block,
|
|
sharding_info,
|
|
):
|
|
"""
|
|
conditions for gradients to be grouped:
|
|
1. group size < grad_bucket_size_numel
|
|
2. same dp group (TODO)
|
|
3. same src rank
|
|
4. same dtype
|
|
5. dependency: grad would NOT be used by other ops within group segment
|
|
|
|
main logic:
|
|
1. record coalesce group
|
|
2. record all dp allreduce/reduce op idx
|
|
|
|
3. insert coalesce op
|
|
4. insert coalesce dependency (avoid allocate memory too early)
|
|
5. modify and remove allreduce/reduce op
|
|
6. ensure sharding-dp hybrid parallel logic
|
|
|
|
gradients inside same group would be fuse into one coalesce tensor
|
|
"""
|
|
ops = block.ops
|
|
if self.grad_bucket_size_numel < 1:
|
|
# numel for transformer layer
|
|
# h = 4096 + 1
|
|
# ffn_numel = 2 * (4 * h) * h
|
|
# mha_numel = 3 * h * h + h * h
|
|
# max_fuse_numel = ffn_numel + mha_numel
|
|
self.grad_bucket_size_numel = 1
|
|
|
|
first_backward_op = None
|
|
for op in ops:
|
|
if is_backward_op(op):
|
|
first_backward_op = op
|
|
break
|
|
# not backward op, sharding for inference
|
|
if first_backward_op is None:
|
|
return
|
|
first_backward_varname = first_backward_op.output_arg_names[0]
|
|
|
|
cur_group = VarGroup(self.grad_bucket_size_numel)
|
|
grad_groups = []
|
|
grouped_grad_names = set()
|
|
|
|
def op_depend_on_group(op, group):
|
|
vars_ = set(op.input_arg_names + op.output_arg_names)
|
|
var_names = {var.name for var in group.vars}
|
|
return len(vars_.intersection(var_names)) > 0
|
|
|
|
# analyze groups
|
|
i = 0
|
|
while i < len(ops):
|
|
op = ops[i]
|
|
if is_data_parallel_reduce_op(op):
|
|
is_reduce = op.type == "reduce" and op.attr("reduce_type") in [
|
|
dist.ReduceOp.AVG,
|
|
dist.ReduceOp.SUM,
|
|
]
|
|
assert is_reduce, (
|
|
"Sharding should reduce grad first and than allreduce if Hybrid Sharding with Data-Parallel"
|
|
)
|
|
|
|
grad_name = op.output_arg_names[0]
|
|
param_name = _get_base_name_from_grad_name(grad_name)
|
|
rank = sharding_info.get_var_rank(param_name)
|
|
grad_var = block.var(grad_name)
|
|
|
|
if cur_group.acceptable(grad_var, rank):
|
|
assert grad_name not in grouped_grad_names
|
|
cur_group.collect(grad_var, rank)
|
|
else:
|
|
grad_groups.append(cur_group)
|
|
cur_group = VarGroup(self.grad_bucket_size_numel)
|
|
cur_group.collect(grad_var, rank)
|
|
|
|
if len(cur_group.vars) == 1:
|
|
cur_group.coalesce_op_idx = i - 1
|
|
# NOTE coalesce dependency: control when allocate memory for gradients
|
|
# too early would increase the peak memory requirement, too later would hurt the performance
|
|
j = 2
|
|
while is_dep_skip_op(ops[i - j]):
|
|
j += 1
|
|
dep_op = ops[i - j]
|
|
dep_varname = dep_op.output_arg_names[0]
|
|
cur_group.coalesce_dep_varname = dep_varname
|
|
|
|
grouped_grad_names.add(grad_name)
|
|
cur_group.reduce_op_indices.append(i)
|
|
|
|
if self.sharding_hybrid_dp and sharding_info.is_in_local_shard(
|
|
param_name
|
|
):
|
|
cur_group.is_in_local_shard = True
|
|
assert ops[i + 1].type == 'all_reduce' and ops[i + 1].attr(
|
|
'reduce_type'
|
|
) in [
|
|
paddle.distributed.ReduceOp.SUM,
|
|
], (
|
|
"Sharding should reduce grad first and than allreduce if Hybrid Sharding with Data-Parallel"
|
|
)
|
|
assert ops[i + 1].output_arg_names[0] == grad_name, (
|
|
"Hybrid Sharding with Data-Parallel should sync same gradient var"
|
|
)
|
|
cur_group.allreduce_op_indices.append(i + 1)
|
|
i += 1
|
|
elif op_depend_on_group(op, cur_group):
|
|
grad_groups.append(cur_group)
|
|
cur_group = VarGroup(self.grad_bucket_size_numel)
|
|
|
|
i += 1
|
|
# some grad not in this rank may not be used after dp reduced
|
|
if len(cur_group.vars) >= 1:
|
|
grad_groups.append(cur_group)
|
|
|
|
_logger.info("Sharding Gradient Communication Optimization:")
|
|
_logger.info(
|
|
f"Gradient Bucket size is [{self.grad_bucket_size_numel}], [{len(grouped_grad_names)}] Gradients are fused into [{len(grad_groups)}] Buckets."
|
|
)
|
|
|
|
# create coalesce tensor and record op idx
|
|
grad_name_to_group_map = {}
|
|
coalesce_to_group_map = {}
|
|
modify_reduce_op_map = {}
|
|
coalesce_op_map = {}
|
|
remove_reduce_op_indices = []
|
|
|
|
for i, group in enumerate(grad_groups):
|
|
if len(group.vars) > 1:
|
|
group.coalesce_var = block.create_var(
|
|
name=unique_name.generate(
|
|
self.grad_coalesce_prefix + str(i)
|
|
),
|
|
dtype=group.dtype,
|
|
persistable=False,
|
|
stop_gradient=True,
|
|
)
|
|
ref_dist_attr = (
|
|
self._dist_context.get_tensor_dist_attr_for_program(
|
|
group.vars[0]
|
|
)
|
|
)
|
|
set_var_dist_attr(
|
|
self._dist_context,
|
|
group.coalesce_var,
|
|
ref_dist_attr.dims_mapping,
|
|
ref_dist_attr.process_mesh,
|
|
chunk_id=ref_dist_attr.chunk_id,
|
|
)
|
|
coalesce_op_map[group.coalesce_op_idx] = group
|
|
last_reduce_op_idx = group.reduce_op_indices.pop()
|
|
modify_reduce_op_map[last_reduce_op_idx] = group
|
|
remove_reduce_op_indices.extend(group.reduce_op_indices)
|
|
if group.is_in_local_shard:
|
|
last_allreduce_op_idx = group.allreduce_op_indices.pop()
|
|
modify_reduce_op_map[last_allreduce_op_idx] = group
|
|
remove_reduce_op_indices.extend(group.allreduce_op_indices)
|
|
else:
|
|
group.coalesce_var = group.vars[0]
|
|
for grad in group.vars:
|
|
grad_name_to_group_map[grad.name] = group
|
|
coalesce_to_group_map[group.coalesce_var.name] = group
|
|
|
|
coalesce_op_set = set(coalesce_op_map.keys())
|
|
modify_op_set = set(modify_reduce_op_map.keys())
|
|
remove_op_set = set(remove_reduce_op_indices)
|
|
conflict = coalesce_op_set.intersection(modify_op_set)
|
|
|
|
assert len(conflict) == 0
|
|
conflict = coalesce_op_set.intersection(remove_op_set)
|
|
assert len(conflict) == 0
|
|
conflict = modify_op_set.intersection(remove_op_set)
|
|
assert len(conflict) == 0
|
|
|
|
# update block
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if idx in modify_reduce_op_map:
|
|
group = modify_reduce_op_map[idx]
|
|
grad_name = op.output_arg_names[0]
|
|
assert grad_name == group.vars[-1].name, (
|
|
f"Unexpected: it is supposed to sync [{group.vars[-1].name}] but got [{grad_name}]"
|
|
)
|
|
op._rename_input(grad_name, group.coalesce_var.name)
|
|
op._rename_output(grad_name, group.coalesce_var.name)
|
|
|
|
if idx in remove_reduce_op_indices:
|
|
block._remove_op(idx, sync=False)
|
|
|
|
if idx in coalesce_op_map:
|
|
group = coalesce_op_map[idx]
|
|
first_grad_name = group.vars[0].name
|
|
assert first_grad_name in op.output_arg_names, (
|
|
f"Unexpected: op is supposed to generate grad [{first_grad_name}] but got [{op}]"
|
|
)
|
|
grad_names = [grad.name for grad in group.vars]
|
|
|
|
concated_shapes = []
|
|
concated_ranks = []
|
|
for grad_ in group.vars:
|
|
shape = grad_.shape
|
|
concated_shapes.extend(shape)
|
|
concated_ranks.append(len(shape))
|
|
|
|
coalesce_op = block._insert_op_without_sync(
|
|
idx,
|
|
type="coalesce_tensor",
|
|
inputs={"Input": grad_names},
|
|
outputs={
|
|
"Output": grad_names,
|
|
"FusedOutput": group.coalesce_var,
|
|
},
|
|
attrs={
|
|
"copy_data": False,
|
|
"use_align": True,
|
|
"dtype": group.dtype,
|
|
"concated_shapes": concated_shapes,
|
|
"concated_ranks": concated_ranks,
|
|
OP_ROLE_KEY: OpRole.Backward,
|
|
},
|
|
)
|
|
|
|
ref_dist_attr = (
|
|
self._dist_context.get_tensor_dist_attr_for_program(
|
|
group.coalesce_var
|
|
)
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
coalesce_op,
|
|
ref_dist_attr.process_mesh,
|
|
ref_dist_attr.dims_mapping,
|
|
self._dist_context,
|
|
chunk_id=ref_dist_attr.chunk_id,
|
|
)
|
|
|
|
depend_op = insert_dependencies_for_vars(
|
|
block,
|
|
idx,
|
|
block.var(group.coalesce_dep_varname),
|
|
group.coalesce_var,
|
|
self._dist_context,
|
|
OpRole.Backward,
|
|
process_mesh=[
|
|
-1
|
|
], # hack to avoid initialize the dist attr for coalesce var
|
|
is_recompute=False,
|
|
sync=False,
|
|
op_namescope="sharding_grad_coalesce_dep",
|
|
)
|
|
block._sync_with_cpp()
|
|
|
|
return coalesce_to_group_map, grad_name_to_group_map
|
|
|
|
def _overlap_grad_comm(
|
|
self,
|
|
block,
|
|
sharding_info,
|
|
coalesce_to_group_map,
|
|
grad_name_to_group_map,
|
|
):
|
|
"""
|
|
overlap gradient communication with backward & optimizer computation.
|
|
|
|
1. assign gradient communications to grad comm stream
|
|
2. for coalesce gradient communication:
|
|
2.1 insert before communication dependencies
|
|
2.2 insert after communication dependencies only when need
|
|
3. there is not need to add explicit dependencies for non-coalesce gradient communication
|
|
|
|
P.S. this overlap pass is ONLY adapted for standalone executor (graph based) and stream award allocator.
|
|
"""
|
|
|
|
if not self.enable_overlap:
|
|
return
|
|
|
|
self.grad_comm_group_stream_pairs = []
|
|
ranks = sharding_info.group.ranks
|
|
# NOTE since the gradient synchronization has calculation, there would be computation
|
|
# competition between backward calculation. therefore should limit the number of stream used.
|
|
for i in range(self.grad_comm_stream_num):
|
|
if i == 0:
|
|
group = sharding_info.group
|
|
else:
|
|
group = new_process_group(ranks, force_new_group=True)
|
|
# NOTE here stream is just a presentation with different name,
|
|
# it is up to executor to create the exact streams given the name.
|
|
stream = f"sharding_grad_comm_stream{i}"
|
|
self.grad_comm_group_stream_pairs.append(
|
|
{
|
|
"comm_group": group,
|
|
"comm_stream": stream,
|
|
}
|
|
)
|
|
|
|
ops = block.ops
|
|
# analyze dependencies
|
|
dep_map = {}
|
|
reduce_op_count = 0
|
|
grad_comm_op_to_stream_idx = {}
|
|
for idx, op in enumerate(ops):
|
|
if is_data_parallel_reduce_op(op):
|
|
if op.type == 'all_reduce' and op.attr('reduce_type') in [
|
|
paddle.distributed.ReduceOp.SUM,
|
|
]:
|
|
continue
|
|
stream_idx = reduce_op_count % self.grad_comm_stream_num
|
|
grad_comm_op_to_stream_idx[op] = stream_idx
|
|
comm_group = self.grad_comm_group_stream_pairs[stream_idx][
|
|
"comm_group"
|
|
]
|
|
comm_stream = self.grad_comm_group_stream_pairs[stream_idx][
|
|
"comm_stream"
|
|
]
|
|
|
|
reduce_varname = op.output("Out")[0]
|
|
grad_group = coalesce_to_group_map[reduce_varname]
|
|
assert grad_group.coalesce_var.name == reduce_varname
|
|
|
|
# coalesce deps
|
|
if len(grad_group.vars) > 1:
|
|
# NOTE should prior vars to be all grads ?
|
|
# when the grad_ops' order is random
|
|
# prior dep
|
|
dep_map[idx] = [
|
|
(
|
|
idx,
|
|
grad_group.vars[-1],
|
|
grad_group.coalesce_var,
|
|
comm_stream,
|
|
"sharding_grad_comm_dep",
|
|
op.dist_attr,
|
|
)
|
|
]
|
|
# post dep
|
|
post_idx = idx + 1
|
|
if self.sharding_hybrid_dp and grad_group.is_in_local_shard:
|
|
post_idx += 1
|
|
dep_map[idx].append(
|
|
(
|
|
post_idx,
|
|
grad_group.coalesce_var,
|
|
grad_group.vars,
|
|
comm_stream,
|
|
"sharding_grad_comm_dep",
|
|
op.dist_attr,
|
|
)
|
|
)
|
|
|
|
# assign stream
|
|
op.dist_attr.execution_stream = comm_stream
|
|
op.dist_attr.scheduling_priority = (
|
|
self.comm_op_scheduling_priority
|
|
)
|
|
op._set_attr("ring_id", comm_group.id)
|
|
if self.sharding_hybrid_dp and grad_group.is_in_local_shard:
|
|
next_op = ops[idx + 1]
|
|
assert next_op.type == 'all_reduce' and next_op.attr(
|
|
'reduce_type'
|
|
) in [
|
|
paddle.distributed.ReduceOp.SUM,
|
|
]
|
|
assert next_op.output("Out")[0] == reduce_varname
|
|
# FIXME hybrid sharding-dp support multi comm & stream in feature
|
|
# next_op._set_attr("ring_id", comm_group.id)
|
|
next_op.dist_attr.execution_stream = comm_stream
|
|
next_op.dist_attr.scheduling_priority = (
|
|
self.comm_op_scheduling_priority
|
|
)
|
|
idx += 1
|
|
|
|
# NOTE(Ruibiao): Why add dependency here?
|
|
# It is hack to delay GC for coalesce_var, which significantly reduce memory usage.
|
|
# With the pattern of reduce_sum + scale, the coalesce_var is used by the reduce_sum
|
|
# op on the comm-stream, and then released by the scale op on the comp-stream. Since
|
|
# the generated and released op are both in comp-stream, the allocation of the
|
|
# coalesce_var can be fast-GC and reused by subsequent comp-op. However in reduce_avg
|
|
# parent, the coalesce_var is released on the reduce_avg op in comm-stream,
|
|
# triggering a cross-stream GC. In such case, an event is recorded on the underlying
|
|
# allocation, and the memory is unable to reused by other comp-ops, resulting in an
|
|
# increase in memory usage. For more details, see the code of StreamSafeCUDAAllocator.
|
|
# This issue should be fixed using CUDAMallocAsyncAllocator in the future.
|
|
if (
|
|
op.type == "reduce"
|
|
and op.attr("reduce_type") == dist.ReduceOp.AVG
|
|
and not grad_group.is_in_local_shard
|
|
and not self.get_attr("gradient_sync_after_accumulate")
|
|
):
|
|
if idx not in dep_map:
|
|
dep_map[idx] = []
|
|
dep_map[idx].append(
|
|
(
|
|
idx + 1,
|
|
grad_group.coalesce_var,
|
|
grad_group.coalesce_var,
|
|
None,
|
|
"sharding_reduce_avg_dep",
|
|
op.dist_attr,
|
|
)
|
|
)
|
|
|
|
reduce_op_count += 1
|
|
|
|
idx += 1
|
|
|
|
# insert deps
|
|
indice = sorted(dep_map.keys(), reverse=True)
|
|
for i in indice:
|
|
for (
|
|
idx,
|
|
prior_vars,
|
|
post_vars,
|
|
comm_stream,
|
|
op_namescope,
|
|
dist_attr,
|
|
) in dep_map[i][::-1]:
|
|
skip_insert_when_sequential_run = (
|
|
False if op_namescope == "sharding_reduce_avg_dep" else True
|
|
)
|
|
|
|
depend_op = insert_dependencies_for_vars(
|
|
block,
|
|
idx,
|
|
prior_vars,
|
|
post_vars,
|
|
self._dist_context,
|
|
OpRole.Backward,
|
|
process_mesh=[
|
|
-1
|
|
], # hack to avoid initialize the dist attr for coalesce var
|
|
is_recompute=False,
|
|
sync=False,
|
|
op_namescope=op_namescope,
|
|
skip_insert_when_sequential_run=skip_insert_when_sequential_run,
|
|
)
|
|
|
|
if depend_op is not None:
|
|
naive_set_dist_op_attr_for_program_by_mesh(
|
|
depend_op,
|
|
process_mesh=dist_attr.process_mesh,
|
|
ctx=self._dist_context,
|
|
chunk_id=dist_attr.chunk_id,
|
|
)
|
|
if comm_stream is not None:
|
|
depend_op.dist_attr.execution_stream = comm_stream
|
|
depend_op.dist_attr.scheduling_priority = (
|
|
self.comm_op_scheduling_priority
|
|
)
|
|
|
|
# hierarchical grad comm
|
|
if self.enable_hierarchical_comm:
|
|
# NOTE so far we only support Isomorphic cluster with 8 ranks per node
|
|
# TODO unify here create communicators
|
|
# create communicators
|
|
nranks_per_node = 8
|
|
assert self.sharding_world_size % nranks_per_node == 0
|
|
global_group = sharding_info.group
|
|
global_ranks = global_group.ranks
|
|
relative_idx_in_node = self.global_rank % nranks_per_node
|
|
node_idx = self.global_rank // nranks_per_node
|
|
inter_node_ranks = [
|
|
rank
|
|
for rank in global_ranks
|
|
if rank % nranks_per_node == relative_idx_in_node
|
|
]
|
|
_logger.info(
|
|
"Sharding Gradient Hierarchical Communication Optimization."
|
|
)
|
|
_logger.info(f"current global rank idx: {self.global_rank}.")
|
|
_logger.info(f"local inter node ranks idx: {inter_node_ranks}.")
|
|
assert (
|
|
len(inter_node_ranks)
|
|
== self.sharding_world_size // nranks_per_node
|
|
)
|
|
intra_node_ranks = [
|
|
rank
|
|
for rank in global_ranks
|
|
if rank // nranks_per_node == node_idx
|
|
]
|
|
assert len(intra_node_ranks) == nranks_per_node
|
|
_logger.info(f"local intra node ranks idx: {intra_node_ranks}.")
|
|
inter_node_groups = []
|
|
intra_node_groups = []
|
|
for _ in range(self.grad_comm_stream_num):
|
|
# TODO re-use one origin communicator
|
|
inter_node_groups.append(
|
|
new_process_group(inter_node_ranks, force_new_group=True)
|
|
)
|
|
intra_node_groups.append(
|
|
new_process_group(intra_node_ranks, force_new_group=True)
|
|
)
|
|
|
|
# update program
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if is_data_parallel_reduce_op(op):
|
|
assert (
|
|
op.type == "reduce"
|
|
and op.attr("reduce_type") == dist.ReduceOp.SUM
|
|
)
|
|
grad_comm_stream_idx = grad_comm_op_to_stream_idx[op]
|
|
inter_node_group = inter_node_groups[grad_comm_stream_idx]
|
|
intra_node_group = intra_node_groups[grad_comm_stream_idx]
|
|
|
|
reduce_varname = op.output("Out")[0]
|
|
if self.enable_overlap:
|
|
comm_stream = op.dist_attr.execution_stream
|
|
dst_rank = int(op.attr("root_id"))
|
|
|
|
in_peer = False
|
|
if dst_rank % nranks_per_node == relative_idx_in_node:
|
|
in_peer = True
|
|
intra_node_dst = dst_rank % nranks_per_node
|
|
|
|
op._set_attr('ring_id', intra_node_group.id)
|
|
op._set_attr('root_id', intra_node_dst)
|
|
|
|
if in_peer:
|
|
inter_node_dst = dst_rank // nranks_per_node
|
|
new_op = block._insert_op_without_sync(
|
|
idx + 1,
|
|
type='reduce',
|
|
inputs={"x": reduce_varname},
|
|
outputs={
|
|
"out": reduce_varname,
|
|
},
|
|
attrs={
|
|
'ring_id': inter_node_group.id,
|
|
'root_id': inter_node_dst,
|
|
'reduce_type': dist.ReduceOp.SUM,
|
|
OP_ROLE_KEY: OpRole.Backward,
|
|
},
|
|
)
|
|
new_op._set_attr(
|
|
'op_namescope', '/' + ParallelMode.DataParallel
|
|
)
|
|
|
|
if self.enable_overlap:
|
|
new_op.dist_attr.execution_stream = comm_stream
|
|
new_op.dist_attr.scheduling_priority = (
|
|
self.comm_op_scheduling_priority
|
|
)
|
|
|
|
block._sync_with_cpp()
|
|
|
|
|
|
def _get_broadcast_first_depend_op(block):
|
|
for op in block.ops:
|
|
if op.type in _supported_optimizer_type:
|
|
return op
|
|
|
|
raise Exception("Could not find optimizer op.")
|
|
|
|
|
|
def _insert_init_and_broadcast_op(
|
|
block,
|
|
insert_idx,
|
|
varname,
|
|
local_rank,
|
|
root_rank,
|
|
ring_id,
|
|
op_role,
|
|
dist_context,
|
|
):
|
|
"""
|
|
empty op for initialization
|
|
"""
|
|
broadcast_var = block.var(varname)
|
|
broadcast_var_dist_attr = dist_context.get_tensor_dist_attr_for_program(
|
|
broadcast_var
|
|
)
|
|
|
|
new_op = block._insert_op_without_sync(
|
|
insert_idx,
|
|
type='broadcast',
|
|
inputs={'x': varname},
|
|
outputs={'out': varname},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
'root': root_rank,
|
|
OP_ROLE_KEY: op_role,
|
|
},
|
|
)
|
|
new_op._set_attr('op_namescope', '/' + ParallelMode.DataParallel)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
new_op,
|
|
broadcast_var_dist_attr.process_mesh,
|
|
broadcast_var_dist_attr.dims_mapping,
|
|
dist_context,
|
|
chunk_id=broadcast_var_dist_attr.chunk_id,
|
|
)
|
|
if local_rank != root_rank:
|
|
new_op = block._insert_op_without_sync(
|
|
insert_idx,
|
|
type="empty",
|
|
outputs={"Out": broadcast_var.name},
|
|
attrs={
|
|
"shape": broadcast_var.shape,
|
|
"dtype": broadcast_var.dtype,
|
|
OP_ROLE_KEY: op_role,
|
|
},
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
new_op,
|
|
broadcast_var_dist_attr.process_mesh,
|
|
broadcast_var_dist_attr.dims_mapping,
|
|
dist_context,
|
|
chunk_id=broadcast_var_dist_attr.chunk_id,
|
|
)
|
|
|
|
|
|
def _insert_reduce_op(
|
|
block,
|
|
op_type,
|
|
insert_idx,
|
|
reduce_var,
|
|
ring_id,
|
|
root_id,
|
|
dist_context,
|
|
reduce_type,
|
|
op_role=OpRole.Backward,
|
|
):
|
|
assert root_id >= 0, (
|
|
f"root id should be a positive int, but now root id is {root_id}"
|
|
)
|
|
new_op = block._insert_op_without_sync(
|
|
insert_idx,
|
|
type=op_type,
|
|
inputs={'x': [reduce_var]},
|
|
outputs={'out': [reduce_var]},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
'root_id': root_id,
|
|
'reduce_type': reduce_type,
|
|
OP_ROLE_KEY: op_role,
|
|
},
|
|
)
|
|
|
|
dist_attr = dist_context.get_tensor_dist_attr_for_program(
|
|
block.var(reduce_var)
|
|
)
|
|
naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
|
|
new_op,
|
|
dist_attr.process_mesh,
|
|
dist_attr.dims_mapping,
|
|
dist_context,
|
|
chunk_id=dist_attr.chunk_id,
|
|
)
|
|
new_op._set_attr('op_namescope', '/' + ParallelMode.DataParallel)
|
|
return new_op
|
|
|
|
|
|
def _get_dp_and_sharding_groups(origin_group, sharding_group_size, rank):
|
|
dp_axis = 0
|
|
sharding_axis = 1
|
|
shape = [len(origin_group) // sharding_group_size, sharding_group_size]
|
|
|
|
dp_group = _get_comm_group(origin_group, shape, dp_axis, rank)
|
|
sharding_group = _get_comm_group(origin_group, shape, sharding_axis, rank)
|
|
|
|
return dp_group, sharding_group
|
|
|
|
|
|
def _is_gradient_clip_op(op):
|
|
return op.desc.has_attr("op_namescope") and op.desc.attr(
|
|
"op_namescope"
|
|
).startswith("/gradient_clip")
|
|
|
|
|
|
def _is_weight_decay_op(op):
|
|
return op.desc.has_attr("op_namescope") and op.desc.attr(
|
|
"op_namescope"
|
|
).startswith("/regularization")
|
|
|
|
|
|
def _is_param_grad_fp32_cast_op(block, op):
|
|
if not is_backward_op(op):
|
|
return False
|
|
if not _is_desired_cast_op(
|
|
block, op, __amp_target_dtype__, core.VarDesc.VarType.FP32
|
|
):
|
|
return False
|
|
if _is_master_grad_cast_op(block, op):
|
|
return False
|
|
output_name = op.output_arg_names[0]
|
|
base_name = output_name[: output_name.find("@")]
|
|
if not block.has_var(base_name):
|
|
return False
|
|
return block.var(base_name).is_parameter
|
|
|
|
|
|
def _is_param_fp16_cast_op(block, op, params):
|
|
if is_optimize_op(op):
|
|
return False
|
|
if not _is_desired_cast_op(block, op):
|
|
return False
|
|
input_name = op.input_arg_names[0]
|
|
if input_name not in params:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _is_desired_cast_op(
|
|
block,
|
|
op,
|
|
src_var_type=core.VarDesc.VarType.FP32,
|
|
dst_var_type=__amp_target_dtype__,
|
|
):
|
|
if op.type != "cast":
|
|
return False
|
|
assert len(op.input_arg_names) == 1
|
|
assert len(op.output_arg_names) == 1
|
|
input_var = block.var(op.input_arg_names[0])
|
|
output_var = block.var(op.output_arg_names[0])
|
|
|
|
if input_var.dtype != src_var_type or output_var.dtype != dst_var_type:
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def _get_base_name_from_grad_name(grad_name):
|
|
base_name = None
|
|
if ".cast_fp16@GRAD" in grad_name:
|
|
base_name = grad_name[: grad_name.find(".cast_fp16@GRAD")]
|
|
elif ".cast_bf16@GRAD" in grad_name:
|
|
base_name = grad_name[: grad_name.find(".cast_bf16@GRAD")]
|
|
elif "@GRAD" in grad_name:
|
|
base_name = grad_name[: grad_name.find("@GRAD")]
|
|
return base_name
|
|
|
|
|
|
def _is_param_grad_allreduce_op(op, block):
|
|
if not is_data_parallel_reduce_op(op):
|
|
return False
|
|
|
|
output_name = op.output_arg_names[0]
|
|
base_name = _get_base_name_from_grad_name(output_name)
|
|
|
|
if not block.has_var(base_name):
|
|
return False
|
|
|
|
return block.var(base_name).is_parameter
|
|
|
|
|
|
def _is_param_grad_sum_op(op, block):
|
|
if not is_backward_op(op):
|
|
return False
|
|
if op.type != "sum":
|
|
return False
|
|
|
|
output_name = op.output_arg_names[0]
|
|
base_name = _get_base_name_from_grad_name(output_name)
|
|
|
|
if not block.has_var(base_name):
|
|
return False
|
|
|
|
return block.var(base_name).is_parameter
|
|
|
|
|
|
def is_sharding_param_broadcast_op(op):
|
|
return (
|
|
op.type == "broadcast"
|
|
and op.desc.has_attr("op_namescope")
|
|
and ParallelMode.DataParallel in op.desc.attr("op_namescope")
|
|
)
|
|
|
|
|
|
def _inference_data_parallel_group_for_operator(
|
|
rank_id, op, dist_context, dp_axis=None
|
|
):
|
|
dp_group = None
|
|
for input_name in op.input_arg_names:
|
|
# TODO(zhaoyingli): maintain a dict in dist_context to record all variables which are renamed,
|
|
# to solve the param@RESHARD cannot be identified.
|
|
if not is_parameter_related(input_name, op.block, dist_context):
|
|
dist_attr = dist_context.get_op_dist_attr_for_program(op)
|
|
process_mesh = dist_attr.process_mesh
|
|
input_dim_mapping = dist_attr.get_input_dims_mapping(input_name)
|
|
mesh_shape = process_mesh.shape
|
|
# NOTE(zhaoyingli): OD-tensor's dims_mapping is empty list.
|
|
if len(input_dim_mapping) == 0:
|
|
continue
|
|
# TODO(JZ-LIANG) replace with specific batch size dimension
|
|
batch_size_axis = input_dim_mapping[0]
|
|
if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
|
|
if dp_axis is None or batch_size_axis == dp_axis:
|
|
group_ranks = _get_comm_group(
|
|
process_mesh.process_ids,
|
|
process_mesh.shape,
|
|
batch_size_axis,
|
|
rank_id,
|
|
)
|
|
dp_group = new_process_group(group_ranks)
|
|
break
|
|
|
|
return dp_group
|
|
|
|
|
|
def partition_by_use_order(params, group_size):
|
|
"""
|
|
shard the continuous param into same rank and divide the forward&backward computation into segment,
|
|
which will favor the fuse pass in later.
|
|
|
|
we assume that the params is already sorted by utilization order.
|
|
"""
|
|
mapping = {}
|
|
total_param_mem = 0.0
|
|
param2mem = []
|
|
for param in params:
|
|
mem = get_var_size(param)
|
|
total_param_mem += mem
|
|
param2mem.append((param, mem))
|
|
mapping = {x: [] for x in range(group_size)}
|
|
cur_rank = 0
|
|
mem_accu = 0.0
|
|
for param, mem in param2mem:
|
|
if mem_accu > total_param_mem * 1.0 * (cur_rank + 1) / group_size:
|
|
cur_rank += 1
|
|
mapping[cur_rank].append(param)
|
|
mem_accu += mem
|
|
|
|
return mapping
|
|
|
|
|
|
def partition_by_greedy_even(params, group_size):
|
|
"""
|
|
use greedy algorithm to partition parameter as even as possible.
|
|
"""
|
|
mapping = {}
|
|
for rank_ in range(group_size):
|
|
mapping[rank_] = []
|
|
sizes = [0] * group_size
|
|
for param in params:
|
|
rank = sizes.index(min(sizes))
|
|
mapping[rank].append(param)
|
|
numel = reduce(lambda x, y: x * y, param.shape, 1)
|
|
assert numel > 0, (
|
|
f"param [{param.name}] should larger than 0, but it is [{numel}]"
|
|
)
|
|
sizes[rank] += numel
|
|
|
|
return mapping
|
|
|
|
|
|
def partition_parameters(params, group_size, algor="greedy_even"):
|
|
if algor == "greedy_even":
|
|
rank_to_params = partition_by_greedy_even(params, group_size)
|
|
else:
|
|
rank_to_params = partition_by_use_order(params, group_size)
|
|
|
|
_logger.info("Sharding Parameter Partition:")
|
|
for k, v in rank_to_params.items():
|
|
_logger.info(
|
|
f"Rank:{k}, Parameter Size:{sum([get_var_size(var) for var in v])} MB."
|
|
)
|
|
_logger.info(f"Params in this rank: {[var.name for var in v]}.")
|
|
|
|
return rank_to_params
|
|
|
|
|
|
def re_order_program(block, param_grads, dist_context):
|
|
# record order
|
|
pname_to_pg_pairs = {}
|
|
for p, g in param_grads:
|
|
pname_to_pg_pairs[p.name] = (p, g)
|
|
|
|
use_order = []
|
|
for op in block.ops:
|
|
for input_name in op.input_arg_names:
|
|
if (input_name in pname_to_pg_pairs) and (
|
|
input_name not in use_order
|
|
):
|
|
use_order.append(input_name)
|
|
if len(use_order) == len(pname_to_pg_pairs):
|
|
break
|
|
|
|
# reorder optimizer
|
|
last_op = block.ops[-1]
|
|
pname_to_op = {}
|
|
num_ops = len(block.ops)
|
|
remove_op_indices = []
|
|
# TODO support case when optimizer is not the last op
|
|
if is_optimize_op(last_op) and last_op.type in _supported_optimizer_type:
|
|
# record optimizer
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if op.type in _supported_optimizer_type:
|
|
assert len(op.input("Param")) == 1
|
|
pname_to_op[op.input("Param")[0]] = op
|
|
remove_op_indices.append(idx)
|
|
assert len(use_order) == len(pname_to_op)
|
|
|
|
# append new opts
|
|
for pname in use_order:
|
|
new_op = block.append_op(type='nop')
|
|
new_op.desc.copy_from(pname_to_op[pname].desc)
|
|
dist_context.set_op_dist_attr_for_program(
|
|
new_op,
|
|
dist_context.get_op_dist_attr_for_program(pname_to_op[pname]),
|
|
)
|
|
|
|
# remove old opts
|
|
for idx in remove_op_indices:
|
|
block._remove_op(idx, sync=False)
|
|
|
|
block._sync_with_cpp()
|
|
assert len(block.ops) == num_ops
|
|
|
|
# TODO reorder gradient clip order
|
|
_logger.info(f"Sharding the Order of param being used: {use_order}.")
|
|
return [pname_to_pg_pairs[p] for p in use_order]
|
|
|
|
|
|
def group_param(sharding_info, fuse_size):
|
|
"""
|
|
param are group by:
|
|
rank id
|
|
fuse_size
|
|
dtype
|
|
"""
|
|
group_to_param_map = {}
|
|
param_to_group_map = {}
|
|
bucket = []
|
|
cur_group = VarGroup(fuse_size)
|
|
for param in sharding_info.params:
|
|
rank = sharding_info.get_var_rank(param.name)
|
|
|
|
if cur_group.acceptable(param, rank):
|
|
cur_group.collect(param, rank)
|
|
else:
|
|
cur_group = VarGroup(fuse_size)
|
|
cur_group.collect(param, rank)
|
|
|
|
cur_group.is_in_local_shard = sharding_info.is_in_local_shard(
|
|
param.name
|
|
)
|
|
|
|
if cur_group in group_to_param_map:
|
|
group_to_param_map[cur_group].append(param.name)
|
|
else:
|
|
group_to_param_map[cur_group] = [param.name]
|
|
|
|
param_to_group_map[param.name] = cur_group
|
|
|
|
return group_to_param_map, param_to_group_map
|
|
|
|
|
|
class ShardingInfo:
|
|
def __init__(self, group, rank, params_grads, partition_algor):
|
|
self.group = group
|
|
self.params_grads = {p.name: (p, g) for p, g in params_grads}
|
|
assert len(self.params_grads) == len(set(self.params_grads)), (
|
|
"found duplicated param in params_grads"
|
|
)
|
|
|
|
self.params = [p for p, _ in params_grads]
|
|
self.param_names = [p.name for p in self.params]
|
|
self.group_size = group.nranks
|
|
self.global_rank = rank
|
|
self.local_rank = group.ranks.index(self.global_rank)
|
|
self.partition_algor = partition_algor
|
|
# rank in below mapping are local rank in this sharding group
|
|
self.rank_to_params = partition_parameters(
|
|
self.params, self.group_size, self.partition_algor
|
|
)
|
|
# include fp32 and fp16 param
|
|
self.param_to_rank = {}
|
|
self._map_param_to_rank()
|
|
|
|
def _map_param_to_rank(self):
|
|
"""
|
|
mapping parameters to the rank which holds it.
|
|
"""
|
|
for rank, params in self.rank_to_params.items():
|
|
for param in params:
|
|
self.param_to_rank[param.name] = rank
|
|
|
|
def get_var_rank(self, varname):
|
|
if varname in self.param_to_rank:
|
|
return self.param_to_rank[varname]
|
|
return -1
|
|
|
|
# determine fp32 and fp16 (cast) param
|
|
def is_in_local_shard(self, param_name):
|
|
return self.get_var_rank(param_name) == self.local_rank
|
|
|
|
# NOTE the follow logic is designed for supporting AMP O1 when
|
|
# the param would be cast to fp16 before used for calculation.
|
|
# and sharding should only broadcast the casted fp16 param
|
|
# instead of the origin fp32 version param.
|
|
def get_broadcast_vars_and_param_usage(self, block):
|
|
broadcast_vars = set()
|
|
fp16_params = set()
|
|
fp16_to_fp32 = {}
|
|
|
|
param_usage = dict.fromkeys(self.param_names, 0)
|
|
for op in block.ops:
|
|
if is_optimize_op(op):
|
|
continue
|
|
for input_name in op.input_arg_names:
|
|
if input_name in self.param_names:
|
|
param_usage[input_name] += 1
|
|
|
|
for op in block.ops:
|
|
if not _is_param_fp16_cast_op(block, op, self.param_names):
|
|
continue
|
|
input_name = op.input_arg_names[0]
|
|
output_name = op.output_arg_names[0]
|
|
broadcast_vars.add(output_name)
|
|
fp16_params.add(output_name)
|
|
fp16_to_fp32[output_name] = input_name
|
|
param_usage[input_name] -= 1
|
|
self.param_to_rank[output_name] = self.param_to_rank[input_name]
|
|
|
|
for param, usage in param_usage.items():
|
|
if usage > 0:
|
|
broadcast_vars.add(param)
|
|
return broadcast_vars, param_usage
|
|
|
|
def get_param_grad(self, param_name):
|
|
if not self.is_in_local_shard(param_name):
|
|
raise ValueError(f"param[{param_name}] not in current rank.")
|
|
if param_name not in self.params_grads:
|
|
raise ValueError(f'param[{param_name}] not in params_grads')
|
|
return self.params_grads.get(param_name, None)
|
|
|
|
|
|
class VarGroup:
|
|
def __init__(self, max_size):
|
|
self.max_size = max_size
|
|
self.dtype = None
|
|
self.rank = -1
|
|
self.numel = 0
|
|
self.vars = []
|
|
self.coalesce_var = None
|
|
self.coalesce_dep_varname = None
|
|
self.coalesce_op_idx = None
|
|
self.reduce_op_indices = []
|
|
self.allreduce_op_indices = []
|
|
self.is_in_local_shard = False
|
|
|
|
def acceptable(self, param, rank):
|
|
if self.numel == 0:
|
|
return True
|
|
else:
|
|
if param.dtype != self.dtype:
|
|
return False
|
|
if rank != self.rank:
|
|
return False
|
|
if self.numel + get_var_numel(param) > self.max_size:
|
|
return False
|
|
return True
|
|
|
|
def collect(self, param, rank):
|
|
self.dtype = param.dtype
|
|
self.rank = rank
|
|
self.numel += get_var_numel(param)
|
|
self.vars.append(param)
|
|
|
|
def __len__(self):
|
|
return len(self.vars)
|