wehub-resource-sync
2208 lines
84 KiB
Python
Executable File
2208 lines
84 KiB
Python
Executable File
# Copyright (c) 2020 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 os
|
|
|
|
import paddle
|
|
from paddle.base import core
|
|
from paddle.incubate.optimizer import PipelineOptimizer
|
|
from paddle.static import (
|
|
create_global_var,
|
|
default_startup_program,
|
|
device_guard,
|
|
)
|
|
from paddle.utils import unique_name
|
|
|
|
from ..utils.log_util import logger
|
|
from .common import (
|
|
OP_ROLE_KEY,
|
|
OP_ROLE_VAR_KEY,
|
|
CollectiveHelper,
|
|
OpRole,
|
|
is_backward_op,
|
|
is_optimizer_op,
|
|
is_update_op,
|
|
)
|
|
from .meta_optimizer_base import MetaOptimizerBase
|
|
from .sharding import utils
|
|
from .sharding.fp16_helper import FP16Utils
|
|
from .sharding.gradient_clip_helper import GradientClipHelper
|
|
from .sharding.offload_helper import OffloadHelper
|
|
from .sharding.prune import ProgramDeps
|
|
from .sharding.shard import ProgramSegment, Shard
|
|
from .sharding.utils import (
|
|
get_first_optimize_op_idx,
|
|
get_grad_device,
|
|
get_var_size,
|
|
insert_allreduce_ops,
|
|
insert_broadcast_ops,
|
|
insert_cast_ops,
|
|
insert_fill_constant_ops,
|
|
insert_reduce_ops,
|
|
insert_scale_loss_grad_ops,
|
|
insert_sync_calc_op,
|
|
insert_sync_comm_ops,
|
|
)
|
|
from .sharding.weight_decay_helper import WeightDecayHelper
|
|
|
|
__all__ = []
|
|
|
|
|
|
class ShardingOptimizer(MetaOptimizerBase):
|
|
"""Sharding Optimizer."""
|
|
|
|
def __init__(self, optimizer):
|
|
super().__init__(optimizer)
|
|
self.inner_opt = optimizer
|
|
self.meta_optimizers_white_list = [
|
|
"RecomputeOptimizer",
|
|
"AMPOptimizer",
|
|
"LarsOptimizer",
|
|
"LambOptimizer",
|
|
"ASPOptimizer",
|
|
# "ModelParallelOptimizer",
|
|
# "PipelineOptimizer",
|
|
]
|
|
self.meta_optimizers_black_list = []
|
|
self._main_program = None
|
|
self._startup_program = None
|
|
self._segments = []
|
|
# params and fp16 params is for broadcast
|
|
self._params = set()
|
|
self._broadcast_vars = set()
|
|
# reduced grads to param name
|
|
self._reduced_grads_to_param = {}
|
|
self._shard = Shard()
|
|
self._verbose = False
|
|
self._thread_mode = False
|
|
self._use_calc_stream = False
|
|
|
|
# use sharding as outer parallelism (e.g. inner:Megatron & outer sharding)
|
|
self.mp_degree = 1
|
|
|
|
def _can_apply(self):
|
|
if not self.role_maker._is_collective:
|
|
return False
|
|
if self.role_maker._worker_num() <= 1:
|
|
return False
|
|
return self.user_defined_strategy.sharding
|
|
|
|
def _disable_strategy(self, dist_strategy):
|
|
dist_strategy.sharding = False
|
|
dist_strategy.sharding_configs = {}
|
|
|
|
def _enable_strategy(self, dist_strategy, context):
|
|
dist_strategy.sharding = True
|
|
dist_strategy.sharding_configs = {"segment_broadcast_MB": 32}
|
|
|
|
def _get_sharding_segment_strategy(self):
|
|
"""get
|
|
self._sharding_segment_strategy
|
|
1. if by_size: self._broadcast_MB
|
|
2. if by_anchors: self._sharding_segment_anchors
|
|
self._backward_remain_anchors
|
|
self._forward_remain_anchors
|
|
"""
|
|
strategy = self.user_defined_strategy
|
|
sharding_configs = strategy.sharding_configs
|
|
segment_strategy = str(sharding_configs["sharding_segment_strategy"])
|
|
|
|
if segment_strategy == "segment_broadcast_MB":
|
|
self._broadcast_MB = sharding_configs["segment_broadcast_MB"]
|
|
assert self._broadcast_MB > 0, (
|
|
"segment size should larger than zero !"
|
|
)
|
|
elif segment_strategy == "segment_anchors":
|
|
self._sharding_segment_anchors = sharding_configs["segment_anchors"]
|
|
assert len(self._sharding_segment_anchors) > 0, (
|
|
"you should set the sharding segment anchors !"
|
|
)
|
|
self._backward_remain_anchors = self._sharding_segment_anchors[:]
|
|
self._forward_remain_anchors = []
|
|
else:
|
|
raise NotImplementedError(
|
|
f"the sharding segment strategy [{segment_strategy}] is not implemented"
|
|
)
|
|
self._sharding_segment_strategy = segment_strategy
|
|
|
|
def _get_hybrid_degree(self):
|
|
"""get
|
|
self.hybrid_dp
|
|
self.sharding_degree
|
|
self.mp_degree
|
|
self.pp_degree
|
|
self.dp_degree
|
|
"""
|
|
strategy = self.user_defined_strategy
|
|
sharding_configs = strategy.sharding_configs
|
|
|
|
# parallelism
|
|
sharding_degree = int(sharding_configs["sharding_degree"])
|
|
mp_degree = int(sharding_configs["mp_degree"])
|
|
pp_degree = int(sharding_configs["pp_degree"])
|
|
dp_degree = int(sharding_configs['dp_degree'])
|
|
global_world_size = self.role_maker._worker_num()
|
|
|
|
assert sharding_degree > 0, "sharding degree must be larger than zero"
|
|
# pipeline setting
|
|
# TODO (JZ-LIANG) should revise here for support mix parallelism with pipeline
|
|
if pp_degree > 1:
|
|
assert strategy.pipeline is True
|
|
|
|
if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None):
|
|
assert pp_degree == 2, (
|
|
"For manually set pipeline, only pp_degree = 2 is supported."
|
|
)
|
|
assert (
|
|
global_world_size == mp_degree * sharding_degree * dp_degree
|
|
), (
|
|
f"global work size [{global_world_size}], mp_degree [{mp_degree}], sharding_degree [{sharding_degree}], dp_degree [{dp_degree}]."
|
|
)
|
|
else:
|
|
assert (
|
|
global_world_size
|
|
== mp_degree * sharding_degree * pp_degree * dp_degree
|
|
), (
|
|
f"global work size [{global_world_size}], mp_degree [{mp_degree}], sharding_degree [{sharding_degree}], pp_degree [{pp_degree}], dp_degree [{dp_degree}]."
|
|
)
|
|
|
|
# FIXME (JZ-LIANG) deprecated hybrid_dp
|
|
if sharding_configs["hybrid_dp"]:
|
|
logger.warning(
|
|
"[hybrid_dp] API setting is deprecated. Now when "
|
|
"dp_degree >= 2, its will be in hybrid dp mode automatically"
|
|
)
|
|
assert dp_degree >= 1
|
|
|
|
self.hybrid_dp = True if dp_degree > 1 else False
|
|
self.sharding_degree = sharding_degree
|
|
self.mp_degree = mp_degree
|
|
self.pp_degree = pp_degree
|
|
self.dp_degree = dp_degree
|
|
|
|
def _get_hybrid_dp_mode(self):
|
|
"""get
|
|
self.hybrid_dp_mode = 'pp_hybrid_dp' or 'sharding_hybrid_dp'
|
|
self.gradient_merge_mode = 'pp_gm' or 'sharding_gm'
|
|
self._gradient_merge_acc_step
|
|
self.pp_allreduce_in_optimize
|
|
self._optimizer_sharding
|
|
"""
|
|
strategy = self.user_defined_strategy
|
|
sharding_configs = strategy.sharding_configs
|
|
|
|
# NOTE (JZ-LIANG)
|
|
# There 2 kind of modes for gradient-merge and hybrid-dp in mixed parallelism [sharding] and [pipeline].
|
|
# We distinguish this two modes since the gm/hybrid-dp related allreduce should be insert in different place
|
|
# according different mode to have best performance:
|
|
# sharding: communication within node, and therefore should insert within backward segment
|
|
# to overlap with bw calc, conduct every micro step.
|
|
# pipeline: communication across nodes, and therefore should insert in update segment,
|
|
# conduct just once per global step.
|
|
dp_mode = None
|
|
# dp here is the pure dp as the outermost parallelism
|
|
if self.hybrid_dp:
|
|
if self.pp_degree > 1:
|
|
dp_mode = "pp_hybrid_dp"
|
|
else:
|
|
assert self.sharding_degree > 1, (
|
|
"by now we only support five kind of hybrid dp: sharding_hybrid_dp, "
|
|
"mp_sharding_hybrid_dp, pp_hybrid_dp, mp_sharding_pp_hybrid_dp, sharding_pp_hybrid_dp."
|
|
)
|
|
dp_mode = "sharding_hybrid_dp"
|
|
|
|
# gradient merge
|
|
gm_mode = None
|
|
gm_acc_step = int(sharding_configs["gradient_merge_acc_step"])
|
|
if self.pp_degree <= 1:
|
|
gm_mode = "sharding_gm"
|
|
self._grad2merged_grad = {}
|
|
else:
|
|
gm_mode = "pp_gm"
|
|
gm_acc_step = strategy.pipeline_configs['accumulate_steps']
|
|
gradient_scale_configs = strategy.gradient_scale_configs
|
|
assert gradient_scale_configs['scale_strategy'] == 'avg', (
|
|
'For pipeline mode, the '
|
|
'gradient scale mode should '
|
|
'be "avg", but got {}'.format(
|
|
gradient_scale_configs['scale_strategy']
|
|
)
|
|
)
|
|
# Note (Yuang Liu): this avg_loss flag determines where to do the average op for grad merge.
|
|
# If True, will do sum firstly for gradient merge, then do scale by gm_acc_step.
|
|
# If False, will scale loss by gm_acc_step first, then do sum for gradient merge.
|
|
self.scale_gradient = gradient_scale_configs['scale_gradient']
|
|
if gm_acc_step > 1:
|
|
logger.info(
|
|
f"Gradient merge in [{gm_mode}], acc step = [{gm_acc_step}]"
|
|
)
|
|
|
|
optimizer_sharding = False
|
|
# TODO(wangxi): need support dp_as_opt_sharding with sharding
|
|
# need support without pp in future
|
|
if (
|
|
self.sharding_degree == 1
|
|
and self.dp_degree > 1
|
|
and sharding_configs['_dp_as_optimizer_sharding']
|
|
and self.pp_degree > 1
|
|
):
|
|
optimizer_sharding = True
|
|
|
|
self.hybrid_dp_mode = dp_mode
|
|
self.gradient_merge_mode = gm_mode
|
|
self._gradient_merge_acc_step = gm_acc_step
|
|
self._optimizer_sharding = optimizer_sharding
|
|
|
|
# this feature is design for ascend, and should NOT be used in GPU training
|
|
self.pp_allreduce_in_optimize = sharding_configs[
|
|
"pp_allreduce_in_optimize"
|
|
]
|
|
|
|
def _inner_opt_minimize(
|
|
self, loss, startup_program, parameter_list, no_grad_set
|
|
):
|
|
pipeline_configs = self.user_defined_strategy.pipeline_configs
|
|
|
|
if self.inner_opt is None:
|
|
raise ValueError(
|
|
"self.inner_opt of ShardingOptimizer should not be None."
|
|
)
|
|
|
|
if self.pp_degree > 1:
|
|
pp_optimizer = PipelineOptimizer(
|
|
self.inner_opt, self._gradient_merge_acc_step
|
|
)
|
|
self._pp_optimizer = pp_optimizer
|
|
|
|
global_rank = self.role_maker._worker_index()
|
|
schedule_mode = pipeline_configs['schedule_mode']
|
|
|
|
pipeline_opt = {
|
|
'schedule_mode': schedule_mode,
|
|
'micro_batch_size': pipeline_configs['micro_batch_size'],
|
|
'local_rank': self.pp_rank,
|
|
'global_rank': global_rank,
|
|
'use_sharding': True,
|
|
# TODO (JZ-LIANG) should revise here for support mix parallelism with pipeline
|
|
'ring_id': 20,
|
|
'global_ring_id': 3,
|
|
'mp_degree': self.mp_degree,
|
|
'mp_rank': global_rank % self.mp_degree,
|
|
'scale_gradient': self.scale_gradient,
|
|
}
|
|
main_program = loss.block.program
|
|
main_program._pipeline_opt = pipeline_opt
|
|
|
|
(
|
|
optimize_ops,
|
|
params_grads,
|
|
program_list,
|
|
self.pipeline_pair,
|
|
self.pp_ring_map,
|
|
) = pp_optimizer.minimize(
|
|
loss, startup_program, parameter_list, no_grad_set
|
|
)
|
|
assert self.pp_degree == len(program_list)
|
|
else:
|
|
optimize_ops, params_grads = self.inner_opt.minimize(
|
|
loss, startup_program, parameter_list, no_grad_set
|
|
)
|
|
|
|
if startup_program is None:
|
|
startup_program = default_startup_program()
|
|
|
|
if self.pp_degree > 1:
|
|
startup_program = startup_program._pipeline_opt['startup_program']
|
|
print("pp_rank:", self.pp_rank)
|
|
if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None):
|
|
main_program = program_list[
|
|
int(os.getenv("PADDLE_MANUAL_PIPELINE_STAGE"))
|
|
]
|
|
else:
|
|
main_program = program_list[self.pp_rank]
|
|
with open(f"main_{self.role_maker._worker_index()}", 'w') as f:
|
|
f.writelines(str(main_program))
|
|
main_block = main_program.global_block()
|
|
new_params_grads = []
|
|
for param, grad in params_grads:
|
|
if main_block.has_var(param.name):
|
|
new_params_grads.append((param, grad))
|
|
params_grads = new_params_grads
|
|
else:
|
|
main_block = loss.block
|
|
|
|
startup_block = startup_program.global_block()
|
|
self._main_program = main_block.program
|
|
self._startup_program = startup_program
|
|
|
|
if self.pp_degree > 1:
|
|
pp_optimizer._rename_gradient_var_name(main_block)
|
|
with open(f"main_{self.role_maker._worker_index()}", 'w') as f:
|
|
f.writelines(str(main_program))
|
|
|
|
return optimize_ops, params_grads
|
|
|
|
def _apply_sharding_pass(self, params_grads):
|
|
if self.sharding_degree == 1:
|
|
return
|
|
|
|
main_block = self._main_program.global_block()
|
|
startup_block = self._startup_program.global_block()
|
|
|
|
# step1: build shard
|
|
self._build_shard(
|
|
params_grads, self.sharding_rank, self.sharding_degree
|
|
)
|
|
|
|
# step2: split_program
|
|
self._split_program(main_block)
|
|
|
|
# step3: add broadcast and reduce ops
|
|
self._add_broadcast_allreduce(main_block)
|
|
main_block._sync_with_cpp()
|
|
startup_block._sync_with_cpp()
|
|
|
|
# step4: remove unneeded ops and vars from block
|
|
self._prune_main_program(
|
|
main_block,
|
|
self._shard,
|
|
[self.mp_ring_id, self.sharding_ring_id, self.pp_ring_id],
|
|
)
|
|
self._prune_startup_program(startup_block, self._shard)
|
|
|
|
def _apply_opt_sharding_pass(self, params_grads):
|
|
"""outer dp as optimizer sharding"""
|
|
if self._optimizer_sharding is False:
|
|
return
|
|
|
|
main_block = self._main_program.global_block()
|
|
startup_block = self._startup_program.global_block()
|
|
|
|
# step1: build shard
|
|
self._build_shard(params_grads, self.dp_rank, self.dp_degree)
|
|
|
|
# NOTE(wangxi): prune_main_program will prune cast if not add this
|
|
for param, grad in params_grads:
|
|
self._reduced_grads_to_param[grad.name] = param.name
|
|
|
|
# step4: remove unneeded ops and vars from block
|
|
self._prune_main_program(
|
|
main_block,
|
|
self._shard,
|
|
[self.mp_ring_id, self.pp_ring_id, self.dp_ring_id],
|
|
)
|
|
self._prune_startup_program(startup_block, self._shard)
|
|
|
|
def _insert_allreduce_for_pp(self, params_grads):
|
|
if self.pp_degree == 1:
|
|
return
|
|
|
|
strategy = self.user_defined_strategy
|
|
sharding_configs = strategy.sharding_configs
|
|
|
|
main_block = self._main_program.global_block()
|
|
startup_block = self._startup_program.global_block()
|
|
|
|
# sharding-pp related logic
|
|
# pp_optimizer._rename_gradient_var_name(main_block)
|
|
# crop ops
|
|
if self.sharding_degree > 1:
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if is_update_op(op):
|
|
op_role_var = op.attr('op_role_var')
|
|
param_name = op_role_var[0]
|
|
if not self._shard.has_param(param_name):
|
|
main_block._remove_op(idx)
|
|
|
|
for idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if op.type != 'cast':
|
|
continue
|
|
in_name = op.input_arg_names[0]
|
|
if in_name not in self._params:
|
|
continue
|
|
# if self._shard.has_param(param_name): continue
|
|
if in_name not in main_block.vars:
|
|
main_block._remove_op(idx)
|
|
|
|
if self._optimizer_sharding:
|
|
# TODO(wangxi): support fp16_allreduce with optimizer sharding
|
|
strategy.fp16_allreduce = False
|
|
|
|
shard = self._shard if self._optimizer_sharding else None
|
|
accumulated_grad_names = self._pp_optimizer._accumulate_gradients(
|
|
main_block, strategy=strategy, shard=shard
|
|
)
|
|
|
|
len_of_ops = len(main_block.ops)
|
|
if self.scale_gradient:
|
|
self._avg_grad_merge_after_sum(main_block, accumulated_grad_names)
|
|
first_optimize_op_index = get_first_optimize_op_idx(main_block)
|
|
|
|
if self.pp_allreduce_in_optimize:
|
|
logger.info(
|
|
f"Pipeline Persistable grad is {accumulated_grad_names}"
|
|
)
|
|
# FIXME(wangxi): accumulated_grad get from pipeline is not
|
|
# include sharding's param@BroadCast grad when
|
|
# pp_allreduce_in_optimize
|
|
accumulated_grad_names = insert_reduce_ops(
|
|
main_block,
|
|
first_optimize_op_index,
|
|
self.sharding_ring_id,
|
|
accumulated_grad_names,
|
|
self._shard,
|
|
core.op_proto_and_checker_maker.OpRole.Optimize,
|
|
use_calc_stream=True,
|
|
rank=self.sharding_rank,
|
|
)
|
|
|
|
logger.info(f"PP-Sharding grad is {accumulated_grad_names}")
|
|
first_optimize_op_index += len(main_block.ops) - len_of_ops
|
|
len_of_ops = len(main_block.ops)
|
|
|
|
if self._optimizer_sharding:
|
|
accumulated_grad_names = utils.insert_reduce_ops(
|
|
main_block,
|
|
first_optimize_op_index,
|
|
self.dp_ring_id,
|
|
accumulated_grad_names,
|
|
self._shard,
|
|
OpRole.Optimize,
|
|
use_calc_stream=True,
|
|
rank=self.dp_rank,
|
|
strategy=strategy,
|
|
)
|
|
logger.info(f"Optimizer grad in this rank {accumulated_grad_names}")
|
|
first_optimize_op_index += len(main_block.ops) - len_of_ops
|
|
len_of_ops = len(main_block.ops)
|
|
|
|
# NOTE(wangxi): we fused after optimize_cast
|
|
optimize_cast = sharding_configs['optimize_cast']
|
|
optimizer_param = utils.insert_broadcast_param_ops(
|
|
main_block,
|
|
len_of_ops,
|
|
self.dp_ring_id,
|
|
[x[0].name for x in params_grads],
|
|
self._shard,
|
|
OpRole.Optimize,
|
|
use_calc_stream=True,
|
|
rank=self.dp_rank,
|
|
strategy=None if optimize_cast else strategy,
|
|
)
|
|
logger.info(f"Optimizer param in this rank {optimizer_param}")
|
|
if not strategy.fuse_grad_merge and not optimize_cast:
|
|
assert len(accumulated_grad_names) == len(optimizer_param)
|
|
elif self.hybrid_dp and self.hybrid_dp_mode == "pp_hybrid_dp":
|
|
insert_allreduce_ops(
|
|
main_block,
|
|
first_optimize_op_index,
|
|
self.dp_ring_id,
|
|
accumulated_grad_names,
|
|
core.op_proto_and_checker_maker.OpRole.Optimize,
|
|
use_calc_stream=True,
|
|
user_defined_strategy=strategy,
|
|
)
|
|
first_optimize_op_index += len(main_block.ops) - len_of_ops
|
|
len_of_ops = len(main_block.ops)
|
|
|
|
# FIXME(wangxi): if fp16_allreduce, put cast fp16->fp32 to there?
|
|
|
|
def _avg_grad_merge_after_sum(self, main_block, accumulated_grad_names):
|
|
if (
|
|
self.user_defined_strategy.amp
|
|
and self.user_defined_strategy.amp_configs[
|
|
'use_dynamic_loss_scaling'
|
|
]
|
|
):
|
|
# For AMP, if using dynamic loss scaling the avg
|
|
# operation can be simple done by modify the LossScaling op.
|
|
for idx, op in enumerate(main_block.ops):
|
|
if op.type == 'check_finite_and_unscale':
|
|
loss_scale_name = op.input('Scale')[0]
|
|
loss_scaling_var = main_block.var(loss_scale_name)
|
|
loss_scale_tmp_var_name = loss_scale_name + '@TMP'
|
|
loss_scale_tmp_var = main_block.create_var(
|
|
name=loss_scale_tmp_var_name,
|
|
shape=loss_scaling_var.shape,
|
|
dtype=loss_scaling_var.dtype,
|
|
)
|
|
main_block._insert_op_without_sync(
|
|
idx,
|
|
type='scale',
|
|
inputs={'X': loss_scaling_var},
|
|
outputs={'Out': loss_scale_tmp_var},
|
|
attrs={
|
|
'scale': self._gradient_merge_acc_step,
|
|
'bias': 0.0,
|
|
'bias_after_scale': False,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
op._rename_input(loss_scale_name, loss_scale_tmp_var_name)
|
|
break
|
|
else:
|
|
# For pp, do the avg operation for gradient merge after merging
|
|
# the gradient to meet the logic for gradient merge under pure dp.
|
|
tmp_first_opt_idx = None
|
|
for idx, op in enumerate(main_block.ops):
|
|
if is_optimizer_op(op) and op.type != 'c_sync_comm_stream':
|
|
tmp_first_opt_idx = idx
|
|
break
|
|
assert tmp_first_opt_idx is not None, (
|
|
'Occurs some errors, no optimize ops'
|
|
)
|
|
for grad in accumulated_grad_names:
|
|
main_block._insert_op_without_sync(
|
|
tmp_first_opt_idx,
|
|
type='scale',
|
|
inputs={'X': grad},
|
|
outputs={'Out': grad},
|
|
attrs={
|
|
'scale': 1.0 / self._gradient_merge_acc_step,
|
|
'bias': 0.0,
|
|
'bias_after_scale': False,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
|
|
def _adapt_amp_clip_without_sharding(self):
|
|
# if not use sharding, adapt amp/clip, for remain parallelism.
|
|
# cast --> amp --> clip --> opt
|
|
if self.sharding_degree > 1:
|
|
return
|
|
if self._optimizer_sharding:
|
|
return
|
|
|
|
main_block = self._main_program.global_block()
|
|
startup_block = self._startup_program.global_block()
|
|
|
|
# amp inf_var & clip global_norm_var
|
|
|
|
rings = [self.mp_ring_id, self.pp_ring_id]
|
|
FP16Utils.sync_amp_check_nan_inf(main_block, rings)
|
|
|
|
gradient_clip_helper = GradientClipHelper(None)
|
|
gradient_clip_helper.sync_global_norm(
|
|
main_block, [self.mp_ring_id, self.pp_ring_id], self.mp_rank
|
|
)
|
|
|
|
def _insert_loss_grad_scale_op(self):
|
|
main_block = self._main_program.global_block()
|
|
|
|
# step6: loss div dp_degree
|
|
global_dp_degree = self.sharding_degree * self.dp_degree
|
|
assert int(global_dp_degree) == global_dp_degree
|
|
if global_dp_degree > 1:
|
|
insert_scale_loss_grad_ops(main_block, scale=global_dp_degree)
|
|
|
|
main_block._sync_with_cpp()
|
|
|
|
def _apply_optimize_offload_pass(self, params_grads):
|
|
strategy = self.user_defined_strategy
|
|
sharding_configs = strategy.sharding_configs
|
|
main_block = self._main_program.global_block()
|
|
startup_block = self._startup_program.global_block()
|
|
|
|
mp_ring_id = self.mp_ring_id if self.mp_degree > 1 else None
|
|
dp_ring_id = self.dp_ring_id if self.dp_degree > 1 else None
|
|
offload_helper = OffloadHelper(
|
|
mp_ring_id=mp_ring_id, dp_ring_id=dp_ring_id
|
|
)
|
|
|
|
# optimize offload should be enable while gradient merge is enable and
|
|
# acc_step is quite large (e.g. >> 100). Since its memcpy could not be
|
|
# overlap with calc, otherwise it will slower down training severely.
|
|
if sharding_configs["optimize_offload"]:
|
|
logger.info("Sharding with optimize offload !")
|
|
offload_helper.offload(main_block, startup_block)
|
|
# The optimize_cast is already included in offload_fp32param
|
|
offload_helper.offload_fp32param(main_block, startup_block)
|
|
elif sharding_configs['optimize_cast']:
|
|
logger.info("Sharding with optimize cast !")
|
|
# NOTE(wangxi): optimize_cast will persist fp16 param, it
|
|
# will take more memory, but will be faster. Trade space for time.
|
|
if self._optimizer_sharding:
|
|
offload_helper.opt_sharding_cast_fp32param(
|
|
main_block, startup_block, [x[0].name for x in params_grads]
|
|
)
|
|
# NOTE(wangxi): fused after optimize_cast
|
|
utils.fuse_opt_broadcast_param_ops(
|
|
main_block, dp_ring_id, self._shard, strategy=strategy
|
|
)
|
|
else:
|
|
offload_helper.cast_fp32param_in_optimize(
|
|
main_block, startup_block
|
|
)
|
|
|
|
def _dump_program_for_debug(self):
|
|
main_block = self._main_program.global_block()
|
|
startup_block = self._startup_program.global_block()
|
|
with open(
|
|
f"start_sharding_{self.role_maker._worker_index()}", 'w'
|
|
) as f:
|
|
f.writelines(str(startup_block.program))
|
|
with open(f"main_sharding_{self.role_maker._worker_index()}", 'w') as f:
|
|
f.writelines(str(main_block.program))
|
|
|
|
def minimize_impl(
|
|
self, loss, startup_program=None, parameter_list=None, no_grad_set=None
|
|
):
|
|
# TODO: (JZ-LIANG) support multiple comm in future
|
|
# self._nrings = self.user_defined_strategy.nccl_comm_num
|
|
self._nrings_sharding = 1
|
|
self._nrings_dp = 1
|
|
|
|
self._get_sharding_segment_strategy()
|
|
self._get_hybrid_degree()
|
|
self._get_hybrid_dp_mode()
|
|
|
|
# config sharding & dp groups
|
|
self._build_groups()
|
|
|
|
# inner optimize minimize
|
|
optimize_ops, params_grads = self._inner_opt_minimize(
|
|
loss, startup_program, parameter_list, no_grad_set
|
|
)
|
|
|
|
self._init_comm()
|
|
|
|
self._apply_sharding_pass(params_grads)
|
|
|
|
self._apply_opt_sharding_pass(params_grads)
|
|
|
|
self._insert_allreduce_for_pp(params_grads)
|
|
|
|
self._adapt_amp_clip_without_sharding()
|
|
|
|
# loss div dp_degree
|
|
self._insert_loss_grad_scale_op()
|
|
|
|
# apply optimize offload or optimize cast
|
|
self._apply_optimize_offload_pass(params_grads)
|
|
|
|
# step6: (optional) sharding gradient merge
|
|
self._sharding_gradient_merge()
|
|
|
|
# # check op dependency
|
|
# FIXME (JZ-LIANG) enable checking in future.
|
|
# check_broadcast(main_block)
|
|
# check_allreduce_sum(main_block, self._shard, self.sharding_ring_id,
|
|
# self.dp_ring_id)
|
|
|
|
# NOTE(JZ-LIANG) ensure in both sharding_hybrid_dp & pp_hybrid_dp
|
|
# init param broadcast should be called after startup pruning
|
|
self._initialization_broadcast()
|
|
|
|
# NOTE(wangxi): if param is not persistable, program.clone will
|
|
# failed, so we remove no persistable param, recreate param as a var
|
|
self._recreate_not_persist_param_as_var()
|
|
|
|
self._dump_program_for_debug()
|
|
return optimize_ops, params_grads
|
|
|
|
def _init_pair_comm(self, pair, ring_id):
|
|
pp_group_endpoints = [
|
|
self.pp_group_endpoints[pair[0]],
|
|
self.pp_group_endpoints[pair[1]],
|
|
]
|
|
pp_rank = 0 if self.pp_rank == pair[0] else 1
|
|
if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None) is None:
|
|
self._collective_helper._init_communicator(
|
|
self._startup_program,
|
|
self.current_endpoint,
|
|
pp_group_endpoints,
|
|
pp_rank,
|
|
ring_id,
|
|
False,
|
|
sync=False,
|
|
)
|
|
|
|
def _init_pipeline_comm(self, startup_block):
|
|
# TODO (JZ-LIANG) to unify pp_rank_ and pp_rank
|
|
if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None) is None:
|
|
self._collective_helper._init_communicator(
|
|
self._startup_program,
|
|
self.current_endpoint,
|
|
self.pp_group_endpoints,
|
|
self.pp_rank,
|
|
self.pp_ring_id,
|
|
False,
|
|
sync=False,
|
|
)
|
|
|
|
# GPU
|
|
for pair in self.pipeline_pair:
|
|
pair_key = pair[0] * 1000 + pair[1]
|
|
ring_id = self.pp_ring_map[pair_key]
|
|
logger.info(f"pp pair:{pair}, ring_id: {ring_id}")
|
|
if self.pp_rank in pair:
|
|
self._init_pair_comm(pair, ring_id)
|
|
|
|
def _init_comm(self):
|
|
# sync var
|
|
startup_block = self._startup_program.global_block()
|
|
|
|
# mp ring
|
|
if self.mp_degree > 1:
|
|
self._collective_helper._init_communicator(
|
|
self._startup_program,
|
|
self.current_endpoint,
|
|
self.mp_group_endpoints,
|
|
self.mp_rank,
|
|
self.mp_ring_id,
|
|
False,
|
|
sync=False,
|
|
)
|
|
|
|
# sharding ring
|
|
if self.sharding_degree > 1:
|
|
self._collective_helper._init_communicator(
|
|
self._startup_program,
|
|
self.current_endpoint,
|
|
self.sharding_group_endpoints,
|
|
self.sharding_rank,
|
|
self.sharding_ring_id,
|
|
False,
|
|
sync=False,
|
|
)
|
|
|
|
# pp ring
|
|
if self.pp_degree > 1:
|
|
self._init_pipeline_comm(startup_block)
|
|
|
|
# pure dp ring
|
|
if self.dp_degree > 1:
|
|
self._collective_helper._init_communicator(
|
|
self._startup_program,
|
|
self.current_endpoint,
|
|
self.dp_group_endpoints,
|
|
self.dp_rank,
|
|
self.dp_ring_id,
|
|
False,
|
|
sync=False,
|
|
)
|
|
|
|
startup_block._sync_with_cpp()
|
|
|
|
def _build_shard(self, params_grads, shard_rank, shard_size):
|
|
# step 2: split params
|
|
self._params = {x[0].name for x in params_grads}
|
|
self._shard.setup(params_grads, shard_rank, shard_size)
|
|
|
|
# step 3: get broadcast vars
|
|
self._broadcast_vars = self._shard.find_broadcast_params(
|
|
self._main_program.global_block()
|
|
)
|
|
|
|
def _wait(
|
|
self,
|
|
):
|
|
endpoints = self.global_endpoints[:]
|
|
current_endpoint = endpoints[self.global_rank]
|
|
if self.global_rank == 0:
|
|
self._collective_helper._wait(current_endpoint, endpoints)
|
|
|
|
def collect_segment(self, segment, op_idx, block):
|
|
segment._start_idx = op_idx + 1
|
|
self._segments.insert(0, segment)
|
|
new_segment = ProgramSegment(block)
|
|
new_segment._end_idx = op_idx + 1
|
|
|
|
return new_segment
|
|
|
|
def _split_program(self, block):
|
|
for op_idx, op in reversed(list(enumerate(block.ops))):
|
|
if int(op.attr('op_role')) != int(OpRole.Optimize):
|
|
last_backward_op_idx = op_idx + 1
|
|
break
|
|
|
|
var2broadcast_time = {}
|
|
segment = ProgramSegment(block)
|
|
segment._end_idx = last_backward_op_idx
|
|
for op_idx in reversed(range(last_backward_op_idx)):
|
|
op = block.ops[op_idx]
|
|
assert int(op.attr('op_role')) != int(OpRole.Optimize)
|
|
if self._sharding_segment_strategy == "segment_broadcast_MB":
|
|
if segment._param_mem >= self._broadcast_MB:
|
|
segment = self.collect_segment(segment, op_idx, block)
|
|
|
|
elif self._sharding_segment_strategy == "segment_anchors":
|
|
if int(op.attr('op_role')) == int(OpRole.Backward):
|
|
for input_name in op.desc.input_arg_names():
|
|
# NOTE (JZ-LIANG) naive rule to support amp, if amp change, should modify here accordingly
|
|
if self.user_defined_strategy.amp:
|
|
if ".cast_fp16@GRAD" not in input_name:
|
|
continue
|
|
else:
|
|
input_name = input_name[
|
|
: input_name.find(".cast_fp16@GRAD")
|
|
]
|
|
|
|
if input_name in self._backward_remain_anchors:
|
|
segment = self.collect_segment(
|
|
segment, op_idx, block
|
|
)
|
|
assert (
|
|
input_name not in self._forward_remain_anchors
|
|
), f"segment anchor [{input_name}] met twice !"
|
|
self._backward_remain_anchors.remove(input_name)
|
|
self._forward_remain_anchors.append(input_name)
|
|
elif int(op.attr('op_role')) == int(OpRole.Forward):
|
|
for output_name in op.desc.output_arg_names():
|
|
if output_name in self._forward_remain_anchors:
|
|
segment = self.collect_segment(
|
|
segment, op_idx, block
|
|
)
|
|
self._forward_remain_anchors.remove(output_name)
|
|
|
|
# find broadcast vars
|
|
for input_name in op.desc.input_arg_names():
|
|
if input_name not in self._broadcast_vars:
|
|
continue
|
|
if input_name in segment._param2broadcast:
|
|
# skip broadcast because it reuse the old broadcast var
|
|
broadcast_name = segment._param2broadcast[input_name]
|
|
if input_name != broadcast_name:
|
|
op._rename_input(input_name, broadcast_name)
|
|
continue
|
|
if self._shard.has_param(input_name):
|
|
broadcast_var_name = input_name
|
|
else:
|
|
broadcast_var_name = unique_name.generate(
|
|
input_name + "@BroadCast"
|
|
)
|
|
if not self._thread_mode:
|
|
segment._fill_constant_vars.append(broadcast_var_name)
|
|
|
|
# (JZ-LIANG) should use Param base name ?
|
|
broadcast_var_base_name = input_name
|
|
if "subprog" in broadcast_var_base_name:
|
|
# remove suffix
|
|
broadcast_var_base_name = broadcast_var_base_name[
|
|
: broadcast_var_base_name.find(".subprog")
|
|
]
|
|
|
|
var2broadcast_time[broadcast_var_base_name] = (
|
|
var2broadcast_time.get(broadcast_var_base_name, 0) + 1
|
|
)
|
|
|
|
segment._param2broadcast[input_name] = broadcast_var_name
|
|
segment._broadcast_vars.append(
|
|
(broadcast_var_name, self._shard.device(input_name))
|
|
)
|
|
segment._param_mem += get_var_size(
|
|
self._main_program.global_block().var(input_name)
|
|
)
|
|
|
|
# find reduce vars
|
|
if self.pp_degree > 1 and self.pp_allreduce_in_optimize:
|
|
# place pipeline gradient allreduce in optimize
|
|
pass
|
|
else:
|
|
if is_backward_op(op) and OP_ROLE_VAR_KEY in op.attr_names:
|
|
op_role_var = op.all_attrs()[OP_ROLE_VAR_KEY]
|
|
if len(op_role_var) != 0:
|
|
assert len(op_role_var) % 2 == 0
|
|
for i in range(0, len(op_role_var), 2):
|
|
param, reduced_grad = (
|
|
op_role_var[i],
|
|
op_role_var[i + 1],
|
|
)
|
|
segment._allreduce_vars.append(reduced_grad)
|
|
assert (
|
|
reduced_grad not in self._reduced_grads_to_param
|
|
)
|
|
self._reduced_grads_to_param[reduced_grad] = param
|
|
|
|
# find cast op
|
|
if FP16Utils.is_fp16_cast_op(block, op, self._params):
|
|
fp32_param = op.desc.input_arg_names()[0]
|
|
fp16_param = op.desc.output_arg_names()[0]
|
|
if self._shard.has_param(fp32_param):
|
|
segment._cast_ops[fp16_param] = fp32_param
|
|
|
|
if segment._param_mem > 0:
|
|
segment._start_idx = 0
|
|
self._segments.insert(0, segment)
|
|
|
|
if self._sharding_segment_strategy == "segment_anchors":
|
|
assert len(self._forward_remain_anchors) == 0, (
|
|
f"remain anchors {self._forward_remain_anchors}"
|
|
)
|
|
assert len(self._backward_remain_anchors) == 0, (
|
|
f"remain anchors {self._backward_remain_anchors}"
|
|
)
|
|
|
|
if self._verbose:
|
|
for varname in sorted(
|
|
var2broadcast_time, key=var2broadcast_time.get, reverse=True
|
|
):
|
|
logger.info(
|
|
f"Sharding broadcast: [{var2broadcast_time[varname]}] times [{varname}]"
|
|
)
|
|
for idx_ in range(len(self._segments)):
|
|
logger.info(f"segment [{idx_}] :")
|
|
logger.info(
|
|
"start op: [{}] [{}]".format(
|
|
block.ops[self._segments[idx_]._start_idx].desc.type(),
|
|
block.ops[
|
|
self._segments[idx_]._start_idx
|
|
].desc.input_arg_names(),
|
|
)
|
|
)
|
|
logger.info(
|
|
"end op: [{}] [{}]".format(
|
|
block.ops[self._segments[idx_]._end_idx].desc.type(),
|
|
block.ops[
|
|
self._segments[idx_]._end_idx
|
|
].desc.input_arg_names(),
|
|
)
|
|
)
|
|
|
|
def _prune_main_program(self, block, shard, rings):
|
|
"""
|
|
calculate deps from allreduce op to optimize op,
|
|
remove ops and vars not needed in this worker
|
|
|
|
1. prune regularization (weight decay)
|
|
2. prune cast_fp32_to_fp16; update amp_inline_checking
|
|
3. prune gradient_clip related; update global_norm_sum
|
|
4. prune optimizer op + param + gradient
|
|
|
|
"""
|
|
weight_decay_helper = WeightDecayHelper()
|
|
weight_decay_helper.prune_weight_decay(block, shard)
|
|
|
|
# FIXME(wangxi): mp should prune duplicated param_grads
|
|
# NOTE (JZ-LIANG) the sync of FoundInfinite should among one entire Model Parallelism
|
|
# group. and each Data Parallelism group should have its own sync of FoundInfinite
|
|
# amp could use global group for sync
|
|
FP16Utils.prune_fp16(block, shard, self._reduced_grads_to_param, rings)
|
|
|
|
# clipbyglobalnorm should only use the Model parallelism group (mp-sharding-pp)
|
|
gradient_clip_helper = GradientClipHelper(None)
|
|
gradient_clip_helper.prune_gradient_clip(block, shard, rings)
|
|
|
|
# build prog deps
|
|
reduced_grads = []
|
|
for idx, op in enumerate(block.ops):
|
|
input_names = op.desc.input_arg_names()
|
|
output_names = op.desc.output_arg_names()
|
|
# FIXME(wangxi): need use grads, pipeline grad is @GRAD@MERGE
|
|
if (
|
|
op.type == "all_reduce"
|
|
and op.attr('reduce_type') == paddle.distributed.ReduceOp.SUM
|
|
):
|
|
assert len(output_names) == 1
|
|
output_name = output_names[0]
|
|
reduced_grads.append(output_name)
|
|
|
|
# prune optimizer state and param
|
|
pruned_opti_vars = []
|
|
for var_name in list(block.vars.keys()):
|
|
if shard.is_opti_var(var_name) and not shard.has_opt_var(var_name):
|
|
pruned_opti_vars.append(var_name)
|
|
program_deps = ProgramDeps(block, reduced_grads, pruned_opti_vars)
|
|
|
|
# Init
|
|
for var_name in program_deps._end_vars:
|
|
program_deps._should_removed_var.add(var_name)
|
|
|
|
# Prune
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if op.type in [
|
|
"all_reduce",
|
|
"c_sync_comm_stream",
|
|
"c_calc_comm_stream",
|
|
"c_gen_nccl_id",
|
|
"c_gen_bkcl_id",
|
|
"c_gen_xccl_id",
|
|
"c_comm_init",
|
|
'send_v2',
|
|
'recv_v2',
|
|
]:
|
|
pass
|
|
elif op.type == "conditional_block":
|
|
assert op.desc.has_attr("sub_block")
|
|
subblock_idx = op.desc.attr("sub_block").id
|
|
subblock_deps = program_deps.get_sub_block_deps(subblock_idx)
|
|
# only prune amp subblock
|
|
if subblock_deps is None or not self._is_amp_subblock(op):
|
|
continue
|
|
# init
|
|
reversed_output_vars = []
|
|
for output_name in op.desc.output("Out"):
|
|
if output_name in program_deps._should_removed_var:
|
|
subblock_deps._should_removed_var.add(output_name)
|
|
program_deps.crop_output_var_from_op(idx, output_name)
|
|
else:
|
|
reversed_output_vars.append(output_name)
|
|
# prune
|
|
for sub_op_idx, _ in reversed(
|
|
list(enumerate(subblock_deps._block.ops))
|
|
):
|
|
if subblock_deps.should_remove_op(sub_op_idx):
|
|
subblock_deps.remove_op(sub_op_idx)
|
|
reversed_input_vars = []
|
|
for input_name in op.desc.input('Input'):
|
|
if input_name not in subblock_deps._should_removed_var:
|
|
reversed_input_vars.append(input_name)
|
|
else:
|
|
program_deps.crop_input_var_from_op(idx, input_name)
|
|
op.desc.set_input('Input', reversed_input_vars)
|
|
op.desc.set_output('Out', reversed_output_vars)
|
|
else:
|
|
# if all outputs of this op are in _should_removed_var
|
|
# _should_removed_var: opt state not cur shard
|
|
if program_deps.should_remove_op(idx):
|
|
# NOTE(wangxi): need reserve all param in optimizer_sharding
|
|
reserved_vars = (
|
|
self._params if self._optimizer_sharding else None
|
|
)
|
|
program_deps.remove_op(idx, reserved_vars)
|
|
|
|
# NOTE (JZ-LIANG) revise and unify logic here
|
|
# sharding support fp16_allreduce logic
|
|
block._sync_with_cpp()
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if op.type == 'concat' and is_optimizer_op(op):
|
|
# remove inputs that not on this card
|
|
reserved_x = []
|
|
for var_name in op.desc.input("X"):
|
|
if block.has_var(var_name):
|
|
reserved_x.append(var_name)
|
|
op.desc.set_input('X', reserved_x)
|
|
block._sync_with_cpp()
|
|
|
|
def _add_broadcast_allreduce(self, block):
|
|
"""
|
|
add broadcast allreduce op
|
|
if enable gradient_merge, insert related ops
|
|
|
|
if combined with pipeline(grad accumulate),
|
|
the grad allreduce should be done in optimize role
|
|
"""
|
|
if len(self._segments) < 1:
|
|
return
|
|
# sharding
|
|
if self.pp_degree > 1 and self.pp_allreduce_in_optimize:
|
|
for idx in range(len(self._segments)):
|
|
assert len(self._segments[idx]._allreduce_vars) == 0
|
|
|
|
# NOTE (JZ-LIANG) revise and unify logic here
|
|
# fix the _end_idx for segments[-1] if pp is used.
|
|
new_end_idx = self._segments[-1]._end_idx
|
|
for idx in range(
|
|
self._segments[-1]._end_idx - 1,
|
|
self._segments[-1]._start_idx - 1,
|
|
-1,
|
|
):
|
|
op = block.ops[idx]
|
|
if op.type == "fill_constant" or op.type == "sum":
|
|
if "MERGED" in op.output_arg_names[0]:
|
|
new_end_idx = idx + 1
|
|
elif op.type == "cast":
|
|
if "@TMP" in op.output_arg_names[0]:
|
|
new_end_idx = idx + 1
|
|
self._segments[-1]._end_idx = new_end_idx
|
|
|
|
if self._segments[-1]._allreduce_vars:
|
|
shard_allreduce_vars = self._shard.filter_grads(
|
|
self._segments[-1]._allreduce_vars
|
|
)
|
|
if (
|
|
self.gradient_merge_mode != "sharding_gm"
|
|
or self._gradient_merge_acc_step <= 1
|
|
):
|
|
if (
|
|
self.hybrid_dp
|
|
and self.hybrid_dp_mode == "sharding_hybrid_dp"
|
|
and len(shard_allreduce_vars) >= 1
|
|
):
|
|
if not self._use_calc_stream:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
self._segments[-1]._end_idx,
|
|
self.dp_ring_id,
|
|
shard_allreduce_vars,
|
|
)
|
|
insert_allreduce_ops(
|
|
block,
|
|
self._segments[-1]._end_idx,
|
|
self.dp_ring_id,
|
|
shard_allreduce_vars,
|
|
user_defined_strategy=self.user_defined_strategy,
|
|
use_calc_stream=self._use_calc_stream,
|
|
)
|
|
# gradient merge
|
|
elif (
|
|
self.gradient_merge_mode == "sharding_gm"
|
|
and self._gradient_merge_acc_step > 1
|
|
):
|
|
self.create_persistable_gradients_and_insert_merge_ops(
|
|
block,
|
|
self._startup_program.global_block(),
|
|
self._segments[-1]._end_idx,
|
|
shard_allreduce_vars,
|
|
self._shard,
|
|
)
|
|
if not self._use_calc_stream:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
self._segments[-1]._end_idx,
|
|
self.sharding_ring_id,
|
|
self._segments[-1]._allreduce_vars,
|
|
)
|
|
# allreduce --> reduce
|
|
insert_reduce_ops(
|
|
block,
|
|
self._segments[-1]._end_idx,
|
|
self.sharding_ring_id,
|
|
self._segments[-1]._allreduce_vars,
|
|
self._shard,
|
|
op_role=OpRole.Backward,
|
|
use_calc_stream=self._use_calc_stream,
|
|
)
|
|
|
|
for idx, segment in reversed(list(enumerate(self._segments))):
|
|
allreduce_vars = (
|
|
self._segments[idx - 1]._allreduce_vars if idx > 0 else []
|
|
)
|
|
broadcast_vars = (
|
|
self._segments[idx + 1]._broadcast_vars
|
|
if idx < len(self._segments) - 1
|
|
else []
|
|
)
|
|
fill_constant_vars = (
|
|
self._segments[idx + 2]._fill_constant_vars
|
|
if idx < len(self._segments) - 2
|
|
else []
|
|
)
|
|
cast_ops = (
|
|
self._segments[idx + 2]._cast_ops
|
|
if idx < len(self._segments) - 2
|
|
else {}
|
|
)
|
|
|
|
for op_idx in reversed(range(segment._start_idx, segment._end_idx)):
|
|
op = block.ops[op_idx]
|
|
for input_name in op.desc.input_arg_names():
|
|
if (
|
|
input_name in segment._param2broadcast
|
|
and input_name != segment._param2broadcast[input_name]
|
|
):
|
|
op._rename_input(
|
|
input_name, segment._param2broadcast[input_name]
|
|
)
|
|
|
|
for param_name, broadcast_name in segment._param2broadcast.items():
|
|
if param_name != broadcast_name:
|
|
block.create_var(
|
|
name=broadcast_name,
|
|
shape=self._main_program.global_block()
|
|
.var(param_name)
|
|
.shape,
|
|
dtype=self._main_program.global_block()
|
|
.var(param_name)
|
|
.dtype,
|
|
persistable=False,
|
|
)
|
|
|
|
# step1: remove cast ops
|
|
block._sync_with_cpp()
|
|
segment._end_idx += FP16Utils.remove_cast_op(
|
|
block, self._params, segment, 0
|
|
)
|
|
|
|
# step2: add Sync ops
|
|
shard_allreduce_vars = self._shard.filter_grads(allreduce_vars)
|
|
|
|
if (
|
|
self.gradient_merge_mode != "sharding_gm"
|
|
or self._gradient_merge_acc_step <= 1
|
|
):
|
|
if (
|
|
self.hybrid_dp
|
|
and self.hybrid_dp_mode == "sharding_hybrid_dp"
|
|
and len(shard_allreduce_vars) >= 1
|
|
):
|
|
if not self._use_calc_stream:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
segment._end_idx,
|
|
self.dp_ring_id,
|
|
shard_allreduce_vars,
|
|
)
|
|
|
|
broad_cast_vars = [x[0] for x in broadcast_vars]
|
|
if not self._use_calc_stream and len(broad_cast_vars) > 0:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
segment._end_idx,
|
|
self.sharding_ring_id,
|
|
broad_cast_vars,
|
|
)
|
|
else:
|
|
comm_dep_vars = allreduce_vars + [
|
|
x[0] for x in broadcast_vars
|
|
]
|
|
if not self._use_calc_stream and len(comm_dep_vars) > 0:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
segment._end_idx,
|
|
self.sharding_ring_id,
|
|
comm_dep_vars,
|
|
)
|
|
# gradient merge
|
|
elif (
|
|
self.gradient_merge_mode == "sharding_gm"
|
|
and self._gradient_merge_acc_step > 1
|
|
):
|
|
broad_cast_vars = [x[0] for x in broadcast_vars]
|
|
if not self._use_calc_stream and len(broad_cast_vars) > 0:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
segment._end_idx,
|
|
self.sharding_ring_id,
|
|
broad_cast_vars,
|
|
)
|
|
|
|
calc_dep_vars = (
|
|
fill_constant_vars
|
|
+ [k for k, v in cast_ops.items()]
|
|
+ self._segments[idx]._allreduce_vars
|
|
)
|
|
|
|
if not self._use_calc_stream and len(calc_dep_vars) > 0:
|
|
insert_sync_calc_op(
|
|
block, segment._end_idx, [calc_dep_vars[-1]]
|
|
)
|
|
|
|
# step3: insert `fill_constant` ops
|
|
insert_fill_constant_ops(
|
|
block, segment._end_idx, fill_constant_vars
|
|
)
|
|
|
|
# step4: add `cast` ops
|
|
insert_cast_ops(block, segment._end_idx, cast_ops)
|
|
|
|
# step5: add broadcast ops
|
|
# gradient merge
|
|
if (
|
|
self.gradient_merge_mode == "sharding_gm"
|
|
and self._gradient_merge_acc_step > 1
|
|
):
|
|
self.create_persistable_gradients_and_insert_merge_ops(
|
|
block,
|
|
self._startup_program.global_block(),
|
|
segment._start_idx,
|
|
shard_allreduce_vars,
|
|
self._shard,
|
|
)
|
|
|
|
insert_broadcast_ops(
|
|
block,
|
|
segment._start_idx,
|
|
self.sharding_ring_id,
|
|
broadcast_vars,
|
|
self._use_calc_stream,
|
|
)
|
|
|
|
# step6: add all_reduce ops
|
|
# dp
|
|
if (
|
|
self.gradient_merge_mode != "sharding_gm"
|
|
or self._gradient_merge_acc_step <= 1
|
|
):
|
|
if (
|
|
self.hybrid_dp
|
|
and self.hybrid_dp_mode == "sharding_hybrid_dp"
|
|
and len(shard_allreduce_vars) >= 1
|
|
):
|
|
insert_allreduce_ops(
|
|
block,
|
|
segment._start_idx,
|
|
self.dp_ring_id,
|
|
shard_allreduce_vars,
|
|
user_defined_strategy=self.user_defined_strategy,
|
|
use_calc_stream=self._use_calc_stream,
|
|
)
|
|
if not self._use_calc_stream:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
segment._start_idx,
|
|
self.sharding_ring_id,
|
|
allreduce_vars,
|
|
)
|
|
# gradient merge
|
|
elif (
|
|
self.gradient_merge_mode == "sharding_gm"
|
|
and self._gradient_merge_acc_step > 1
|
|
):
|
|
if not self._use_calc_stream:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
segment._start_idx,
|
|
self.sharding_ring_id,
|
|
allreduce_vars,
|
|
)
|
|
# sharding
|
|
# allreduce --> reduce
|
|
# TODO temp change
|
|
if len(allreduce_vars) > 0:
|
|
insert_reduce_ops(
|
|
block,
|
|
segment._start_idx,
|
|
self.sharding_ring_id,
|
|
allreduce_vars,
|
|
self._shard,
|
|
op_role=OpRole.Backward,
|
|
use_calc_stream=self._use_calc_stream,
|
|
)
|
|
|
|
block._sync_with_cpp()
|
|
|
|
if self._segments[0]._broadcast_vars:
|
|
broadcast_vars = [x[0] for x in self._segments[0]._broadcast_vars]
|
|
if not self._use_calc_stream:
|
|
insert_sync_comm_ops(
|
|
block,
|
|
self._segments[0]._start_idx,
|
|
self.sharding_ring_id,
|
|
broadcast_vars,
|
|
)
|
|
insert_broadcast_ops(
|
|
block,
|
|
self._segments[0]._start_idx,
|
|
self.sharding_ring_id,
|
|
self._segments[0]._broadcast_vars,
|
|
self._use_calc_stream,
|
|
)
|
|
|
|
fill_constant_vars = []
|
|
for x in self._segments[:2]:
|
|
fill_constant_vars += x._fill_constant_vars
|
|
|
|
# Join
|
|
cast_ops = {}
|
|
for x in self._segments[:2]:
|
|
for k, v in x._cast_ops.items():
|
|
cast_ops[k] = v
|
|
|
|
calc_deps_vars = fill_constant_vars + [k for k, v in cast_ops.items()]
|
|
if not self._use_calc_stream and (fill_constant_vars or cast_ops):
|
|
insert_sync_calc_op(
|
|
block, self._segments[0]._start_idx, [calc_deps_vars[-1]]
|
|
)
|
|
|
|
if fill_constant_vars:
|
|
insert_fill_constant_ops(
|
|
block, self._segments[0]._start_idx, fill_constant_vars
|
|
)
|
|
|
|
if cast_ops:
|
|
insert_cast_ops(block, self._segments[0]._start_idx, cast_ops)
|
|
|
|
return
|
|
|
|
def _prune_startup_program(self, block, shard):
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
for output_name in op.desc.output_arg_names():
|
|
if shard.has_var(output_name):
|
|
continue
|
|
if self._optimizer_sharding and shard.is_param(output_name):
|
|
continue
|
|
# TODO why do we remove op, when only one var is removed
|
|
block._remove_op(idx, sync=False)
|
|
break
|
|
for var_name in list(block.vars.keys()):
|
|
if shard.has_var(var_name):
|
|
continue
|
|
if self._optimizer_sharding and shard.is_param(var_name):
|
|
continue
|
|
block._remove_var(var_name, sync=False)
|
|
block._sync_with_cpp()
|
|
|
|
def _build_groups(self):
|
|
"""
|
|
pre-assign ring ids
|
|
mp: 0
|
|
sharding: 1
|
|
pure-dp: 2
|
|
global: 3
|
|
pp: 4
|
|
pp-pair: >= 20
|
|
if one parallelism is not enable: -1
|
|
and only support parallelism hierarchy: mp --> sharding --> pp --> dp
|
|
"""
|
|
# step 1: initialize nccl
|
|
self.global_word_size = self.role_maker._worker_num()
|
|
self.global_rank = self.role_maker._worker_index()
|
|
self.global_endpoints = self.role_maker._get_trainer_endpoints()
|
|
if self._thread_mode:
|
|
self.current_endpoint = self.global_endpoints[
|
|
self.role_maker._role_id()
|
|
]
|
|
else:
|
|
self.current_endpoint = self.global_endpoints[self.global_rank]
|
|
self._collective_helper = CollectiveHelper(
|
|
self.role_maker, nrings=self._nrings_sharding
|
|
)
|
|
assert self.global_word_size % self.mp_degree == 0, (
|
|
f"global_word_size: {self.global_word_size} should be divisible to the mp_degree: {self.mp_degree}"
|
|
)
|
|
assert self.global_word_size % self.sharding_degree == 0, (
|
|
f"global_word_size: {self.global_word_size} should be divisible to the sharding_degree: {self.sharding_degree}"
|
|
)
|
|
assert self.global_word_size % self.pp_degree == 0, (
|
|
f"global_word_size: {self.global_word_size} should be divisible to the pp_degree: {self.pp_degree}"
|
|
)
|
|
assert self.global_word_size % self.dp_degree == 0, (
|
|
f"global_word_size: {self.global_word_size} should be divisible to the dp_degree: {self.dp_degree}"
|
|
)
|
|
|
|
# mp group
|
|
if self.mp_degree > 1:
|
|
self.mp_ring_id = 0
|
|
self.mp_rank = self.global_rank % self.mp_degree
|
|
self.mp_group_id = self.global_rank // self.mp_degree
|
|
self.mp_group_endpoints = [
|
|
ep
|
|
for idx, ep in enumerate(self.global_endpoints)
|
|
if idx // self.mp_degree == self.mp_group_id
|
|
]
|
|
assert self.current_endpoint in self.mp_group_endpoints
|
|
assert len(self.mp_group_endpoints) == self.mp_degree, (
|
|
f"num of mp worker in group is [{len(self.mp_group_endpoints)}], but mp group size is [{self.mp_degree}]"
|
|
)
|
|
else:
|
|
self.mp_degree = 1
|
|
self.mp_ring_id = -1
|
|
self.mp_rank = -1
|
|
self.mp_group_id = -1
|
|
self.mp_group_endpoints = []
|
|
|
|
# sharding
|
|
if self.sharding_degree > 1:
|
|
self.sharding_ring_id = 1
|
|
self.sharding_rank = (
|
|
self.global_rank // self.mp_degree
|
|
) % self.sharding_degree
|
|
self.sharding_group_id = self.global_rank // (
|
|
self.mp_degree * self.sharding_degree
|
|
)
|
|
# mp + sharding + ...
|
|
if self.mp_degree > 1:
|
|
self.sharding_group_endpoints = [
|
|
ep
|
|
for idx, ep in enumerate(self.global_endpoints)
|
|
if (idx // (self.mp_degree * self.sharding_degree))
|
|
== self.sharding_group_id
|
|
and idx % self.mp_degree == self.mp_rank
|
|
]
|
|
# sharding + ...
|
|
else:
|
|
self.sharding_group_endpoints = [
|
|
ep
|
|
for idx, ep in enumerate(self.global_endpoints)
|
|
if (idx // (self.mp_degree * self.sharding_degree))
|
|
== self.sharding_group_id
|
|
]
|
|
assert self.current_endpoint in self.sharding_group_endpoints
|
|
else:
|
|
self.sharding_degree = 1
|
|
self.sharding_ring_id = -1
|
|
self.sharding_rank = -1
|
|
self.sharding_group_id = -1
|
|
self.sharding_group_endpoints = []
|
|
|
|
# pp
|
|
if self.pp_degree > 1:
|
|
self.pp_pair_ring_id = 20
|
|
# pipeline global ring_id set to 4 for sharding0, mp1, dp2, global3
|
|
self.pp_ring_id = 4
|
|
self.pp_rank = (
|
|
self.global_rank
|
|
// (self.sharding_degree * self.mp_degree)
|
|
% self.pp_degree
|
|
)
|
|
# (NOTE): Already adjust for (outer-pure) dp
|
|
self.pp_group_id = self.global_rank // (
|
|
self.mp_degree * self.sharding_degree * self.pp_degree
|
|
)
|
|
pp_first_stage_idx = self.global_rank % (
|
|
self.sharding_degree * self.mp_degree
|
|
) + self.pp_group_id * (
|
|
self.mp_degree * self.sharding_degree * self.pp_degree
|
|
)
|
|
pp_stage_offset = self.sharding_degree * self.mp_degree
|
|
self.pp_group_endpoints = []
|
|
for i in range(self.pp_degree):
|
|
self.pp_group_endpoints.append(
|
|
self.global_endpoints[
|
|
pp_first_stage_idx + pp_stage_offset * i
|
|
]
|
|
)
|
|
assert self.current_endpoint in self.pp_group_endpoints
|
|
else:
|
|
self.pp_ring_id = -1
|
|
self.pp_degree = 1
|
|
self.pp_pair_ring_id = -1
|
|
self.pp_rank = -1
|
|
self.pp_group_id = -1
|
|
self.pp_group_endpoints = []
|
|
|
|
# outer-pure-dp group
|
|
# NOTE (JZ-LIANG) support outer-pure-dp to scale the throughput in 3D parallelism
|
|
# e.g. mp-sharding-pp-dp
|
|
# sharding-hybrid-dp as one scenario of outer-pure-dp
|
|
local_pp_degree = self.pp_degree
|
|
if os.getenv("PADDLE_MANUAL_PIPELINE_STAGE", None):
|
|
assert self.pp_degree == 2, (
|
|
"For manually set pipeline, only pp_degree = 2 is supported."
|
|
)
|
|
assert (
|
|
self.global_word_size
|
|
== self.mp_degree * self.sharding_degree * self.dp_degree
|
|
), (
|
|
f"global work size [{self.global_word_size}], mp_degree [{self.mp_degree}], sharding_degree [{self.sharding_degree}], dp_degree [{self.dp_degree}]."
|
|
)
|
|
local_pp_degree = 1
|
|
else:
|
|
assert (
|
|
self.global_word_size
|
|
== self.mp_degree
|
|
* self.sharding_degree
|
|
* self.pp_degree
|
|
* self.dp_degree
|
|
), (
|
|
f"mp_degree: [{self.mp_degree}], sharding_degree: [{self.sharding_degree}], pp_degree: [{self.pp_degree}], dp_degree: [{self.dp_degree}]; BUT global nrank: [{self.global_word_size}]"
|
|
)
|
|
|
|
if self.dp_degree > 1:
|
|
self.dp_ring_id = 2
|
|
self.dp_rank = self.global_rank // (
|
|
self.sharding_degree * self.mp_degree * local_pp_degree
|
|
)
|
|
dp_first_rank_idx = self.global_rank % (
|
|
self.sharding_degree * self.mp_degree * local_pp_degree
|
|
)
|
|
dp_offset = self.sharding_degree * self.mp_degree * local_pp_degree
|
|
self.dp_group_endpoints = []
|
|
for i in range(self.dp_degree):
|
|
self.dp_group_endpoints.append(
|
|
self.global_endpoints[dp_first_rank_idx + dp_offset * i]
|
|
)
|
|
assert self.current_endpoint in self.dp_group_endpoints
|
|
logger.info("Hybrid DP mode turn on !")
|
|
else:
|
|
self.dp_ring_id = -1
|
|
self.dp_rank = -1
|
|
self.dp_group_endpoints = []
|
|
|
|
# global group
|
|
# use for gen_nccl_comm_sync, amp check nan inf, clip by global norm
|
|
# NOTE (JZ-LIANG) when use global ring for calc global norm and dp_degree > 1, the allreduce result should be divided by dp_degree
|
|
self.global_ring_id = 3
|
|
|
|
logger.info(f"global word size: {self.global_word_size}")
|
|
logger.info(f"global rank: {self.global_rank}")
|
|
logger.info(f"global endpoints: {self.global_endpoints}")
|
|
logger.info(f"global ring id: {self.global_ring_id}")
|
|
logger.info("#####" * 6)
|
|
|
|
logger.info(f"mp group size: {self.mp_degree}")
|
|
logger.info(f"mp rank: {self.mp_rank}")
|
|
logger.info(f"mp group id: {self.mp_group_id}")
|
|
logger.info(f"mp group endpoints: {self.mp_group_endpoints}")
|
|
logger.info(f"mp ring id: {self.mp_ring_id}")
|
|
logger.info("#####" * 6)
|
|
|
|
logger.info(f"sharding group size: {self.sharding_degree}")
|
|
logger.info(f"sharding rank: {self.sharding_rank}")
|
|
logger.info(f"sharding group id: {self.sharding_group_id}")
|
|
logger.info(
|
|
f"sharding group endpoints: {self.sharding_group_endpoints}"
|
|
)
|
|
logger.info(f"sharding ring id: {self.sharding_ring_id}")
|
|
logger.info("#####" * 6)
|
|
|
|
logger.info(f"pp group size: {self.pp_degree}")
|
|
logger.info(f"pp rank: {self.pp_rank}")
|
|
logger.info(f"pp group id: {self.pp_group_id}")
|
|
logger.info(f"pp group endpoints: {self.pp_group_endpoints}")
|
|
logger.info(f"pp ring id: {self.pp_ring_id}")
|
|
logger.info("#####" * 6)
|
|
|
|
logger.info(f"pure dp group size: {self.dp_degree}")
|
|
logger.info(f"pure dp rank: {self.dp_rank}")
|
|
logger.info(f"pure dp group endpoints: {self.dp_group_endpoints}")
|
|
logger.info(f"pure dp ring id: {self.dp_ring_id}")
|
|
logger.info("#####" * 6)
|
|
|
|
def _recreate_not_persist_param_as_var(self):
|
|
def recreate_not_persist_param_as_var(program):
|
|
block = program.global_block()
|
|
params = block.all_parameters()
|
|
for param in params:
|
|
if param.persistable:
|
|
continue
|
|
|
|
name = param.name
|
|
shape = param.shape
|
|
dtype = param.dtype
|
|
type = param.type
|
|
lod_level = param.lod_level
|
|
stop_gradient = param.stop_gradient
|
|
trainable = param.trainable
|
|
optimize_attr = param.optimize_attr
|
|
regularizer = param.regularizer
|
|
have_dist_attr = False
|
|
is_distributed = False
|
|
if hasattr(param, 'is_distributed'):
|
|
have_dist_attr = True
|
|
is_distributed = param.is_distributed
|
|
|
|
block._remove_var(name, sync=False)
|
|
var = block.create_var(
|
|
name=name,
|
|
shape=shape,
|
|
dtype=dtype,
|
|
type=type,
|
|
lod_level=lod_level,
|
|
stop_gradient=stop_gradient,
|
|
trainable=trainable,
|
|
persistable=False,
|
|
)
|
|
if have_dist_attr:
|
|
var.is_distributed = is_distributed
|
|
|
|
block._sync_with_cpp()
|
|
|
|
recreate_not_persist_param_as_var(self._startup_program)
|
|
recreate_not_persist_param_as_var(self._main_program)
|
|
|
|
def _initialization_broadcast(self):
|
|
"""
|
|
this function is to ensure the initialization between dp group to be
|
|
identical when hybrid-dp is used, and the initialization of
|
|
not distributed param between mp group to be identical.
|
|
"""
|
|
if self.dp_degree <= 1 and self.mp_degree <= 1:
|
|
return
|
|
|
|
startup_block = self._startup_program.global_block()
|
|
|
|
params = startup_block.all_parameters()
|
|
params_name = []
|
|
not_dist_param_name = set()
|
|
|
|
for param in params:
|
|
params_name.append(param.name)
|
|
if not hasattr(param, 'is_distributed') or not param.is_distributed:
|
|
not_dist_param_name.add(param.name)
|
|
|
|
# offload and optimize_cast will insert broadcast op
|
|
broadcast_params = set()
|
|
for op in startup_block.ops:
|
|
if op.type == 'broadcast':
|
|
broadcast_params.add(op.desc.output_arg_names()[0])
|
|
|
|
for param in params_name:
|
|
if param in broadcast_params:
|
|
continue
|
|
|
|
rings = []
|
|
# need sync not distributed param in mp group
|
|
if self.mp_degree > 1 and param in not_dist_param_name:
|
|
rings.append(self.mp_ring_id)
|
|
if self.dp_degree > 1:
|
|
rings.append(self.dp_ring_id)
|
|
|
|
for ring in rings:
|
|
startup_block.append_op(
|
|
type='broadcast',
|
|
inputs={'x': param},
|
|
outputs={'out': param},
|
|
attrs={
|
|
'ring_id': ring,
|
|
'root': 0,
|
|
OP_ROLE_KEY: OpRole.Forward,
|
|
},
|
|
)
|
|
|
|
startup_block._sync_with_cpp()
|
|
|
|
# sharding gradient merge
|
|
def create_persistable_gradients_and_insert_merge_ops(
|
|
self, main_block, startup_block, insert_idx, grad_names, shard
|
|
):
|
|
for grad_name in grad_names:
|
|
assert get_grad_device(grad_name, shard) == shard.worker_idx, (
|
|
f"try to merge gradient not belong to current shard: [{grad_name}]"
|
|
)
|
|
persistable_grad_name = grad_name + '@GradientMerge'
|
|
assert grad_name not in self._grad2merged_grad, (
|
|
f"grad [{grad_name}] already in grad2merged_grad, maybe you meet sharing weight case !"
|
|
)
|
|
self._grad2merged_grad[grad_name] = persistable_grad_name
|
|
grad_var = main_block.var(grad_name)
|
|
# create var
|
|
gradient_merge_var = main_block.create_var(
|
|
name=persistable_grad_name,
|
|
shape=grad_var.shape,
|
|
dtype=grad_var.dtype,
|
|
persistable=True,
|
|
)
|
|
startup_gradient_merge_var = startup_block.create_var(
|
|
name=persistable_grad_name,
|
|
shape=grad_var.shape,
|
|
dtype=grad_var.dtype,
|
|
persistable=True,
|
|
)
|
|
|
|
# merge gradient
|
|
main_block._insert_op_without_sync(
|
|
insert_idx,
|
|
type="elementwise_add",
|
|
inputs={'X': grad_name, 'Y': gradient_merge_var},
|
|
outputs={'Out': gradient_merge_var},
|
|
attrs={
|
|
'axis': -1,
|
|
OP_ROLE_KEY: OpRole.Backward,
|
|
},
|
|
)
|
|
|
|
# startup initialization
|
|
startup_block.append_op(
|
|
type="fill_constant",
|
|
outputs={"Out": startup_gradient_merge_var},
|
|
attrs={
|
|
"shape": grad_var.shape,
|
|
"dtype": grad_var.dtype,
|
|
"value": float(0),
|
|
},
|
|
)
|
|
|
|
main_block._sync_with_cpp()
|
|
startup_block._sync_with_cpp()
|
|
|
|
def _create_gm_cond(self, main_block):
|
|
# Add const var
|
|
acc_step_var = create_global_var(
|
|
name="gradient_merge_acc_step",
|
|
shape=[1],
|
|
value=int(self._gradient_merge_acc_step),
|
|
dtype='int32',
|
|
persistable=True,
|
|
force_cpu=True,
|
|
)
|
|
|
|
zero_var = create_global_var(
|
|
name="gradient_merge_zero",
|
|
shape=[1],
|
|
value=0,
|
|
dtype='int32',
|
|
persistable=True,
|
|
force_cpu=True,
|
|
)
|
|
|
|
# Add step var & cond var
|
|
current_step_var = create_global_var(
|
|
name="gradient_merge_current_step",
|
|
shape=[1],
|
|
value=0,
|
|
dtype='int32',
|
|
persistable=True,
|
|
force_cpu=True,
|
|
)
|
|
|
|
cond_var = main_block.create_var(
|
|
name="gradient_merge_cond", shape=[1], dtype='bool'
|
|
)
|
|
|
|
with device_guard("cpu"):
|
|
# step_var = (step_var + 1) % k_step
|
|
main_block.append_op(
|
|
type='increment',
|
|
inputs={'X': [current_step_var]},
|
|
outputs={'Out': [current_step_var]},
|
|
attrs={'step': float(1), OP_ROLE_KEY: OpRole.Optimize},
|
|
)
|
|
|
|
main_block.append_op(
|
|
type='elementwise_mod',
|
|
inputs={'X': current_step_var, 'Y': acc_step_var},
|
|
outputs={'Out': current_step_var},
|
|
attrs={
|
|
'axis': -1,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
|
|
# cond_var = (step_var == 0)
|
|
main_block.append_op(
|
|
type='equal',
|
|
inputs={'X': current_step_var, 'Y': zero_var},
|
|
outputs={'Out': cond_var},
|
|
attrs={OP_ROLE_KEY: OpRole.Optimize},
|
|
)
|
|
# paddle.static.Print(current_step_var, message="in FWBW last conditional")
|
|
return cond_var
|
|
|
|
def _true_apply_gradient(self):
|
|
"""
|
|
allreduce grad@gradientmerge in dp group
|
|
grad@gradientmerge / acc_step
|
|
re-create all optimize ops of origin main block and rename them
|
|
cast(backward)
|
|
amp
|
|
clip
|
|
opt
|
|
# fill constant grad@gradientmerge
|
|
|
|
"""
|
|
# current conditional block
|
|
main_block = self._main_program.global_block()
|
|
cur_block_idx = self._main_program.current_block_idx
|
|
cur_block = self._main_program.current_block()
|
|
self.cond_block = self._main_program.current_block()
|
|
|
|
# cur_block's forward_block & backward_block is itself
|
|
cur_block._set_forward_block_idx(cur_block_idx)
|
|
|
|
# allreduce grad@gradientmerge
|
|
if self.hybrid_dp:
|
|
assert self.dp_ring_id >= 0, (
|
|
"dp_ring_id should larger than 0 when in sharding&DP mode"
|
|
)
|
|
for grad, merged_grad in self._grad2merged_grad.items():
|
|
merged_grad_var = main_block.var(merged_grad)
|
|
cur_block.append_op(
|
|
type='all_reduce',
|
|
inputs={'X': merged_grad_var},
|
|
outputs={'Out': merged_grad_var},
|
|
attrs={
|
|
'ring_id': self.dp_ring_id,
|
|
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
|
|
# grad@gradientmerge / acc_step
|
|
for grad, merged_grad in self._grad2merged_grad.items():
|
|
# grad /= k_steps
|
|
merged_grad_var = main_block.var(merged_grad)
|
|
cur_block.append_op(
|
|
type='scale',
|
|
inputs={'X': merged_grad_var},
|
|
outputs={'Out': merged_grad_var},
|
|
attrs={
|
|
'scale': 1.0 / float(self._gradient_merge_acc_step),
|
|
'bias': 0.0,
|
|
'bias_after_scale': False,
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
|
|
# re-create optimize ops
|
|
already_moved_var_names = []
|
|
for op_desc in self.original_optimize_ops_desc:
|
|
new_op_desc = cur_block.desc.append_op()
|
|
new_op_desc.copy_from(op_desc)
|
|
|
|
for input_name in new_op_desc.input_arg_names():
|
|
if input_name in self._grad2merged_grad:
|
|
new_op_desc._rename_input(
|
|
input_name, self._grad2merged_grad[input_name]
|
|
)
|
|
|
|
for output_name in new_op_desc.output_arg_names():
|
|
if output_name in self._grad2merged_grad:
|
|
new_op_desc._rename_output(
|
|
output_name, self._grad2merged_grad[output_name]
|
|
)
|
|
|
|
# move non temp optimize vars from block0 to cond block
|
|
if (
|
|
output_name not in already_moved_var_names
|
|
and output_name not in self._grad2merged_grad.keys()
|
|
):
|
|
var_ = self._main_program.global_block().var(output_name)
|
|
if not var_.persistable:
|
|
# move
|
|
name_ = var_.name
|
|
shape_ = var_.shape
|
|
type_ = var_.dtype
|
|
self._main_program.global_block()._remove_var(
|
|
var_.name, sync=False
|
|
)
|
|
self.cond_block.create_var(
|
|
name=name_,
|
|
shape=shape_,
|
|
dtype=type_,
|
|
persistable=False,
|
|
)
|
|
already_moved_var_names.append(name_)
|
|
|
|
self._main_program.global_block()._sync_with_cpp()
|
|
cur_block._sync_with_cpp()
|
|
|
|
# fill zero to grad@gradientmerge
|
|
for grad, merged_grad in self._grad2merged_grad.items():
|
|
merged_grad_var = main_block.var(merged_grad)
|
|
cur_block.append_op(
|
|
type='fill_constant',
|
|
outputs={'Out': merged_grad_var},
|
|
attrs={
|
|
"shape": merged_grad_var.shape,
|
|
"dtype": merged_grad_var.dtype,
|
|
"value": float(0),
|
|
OP_ROLE_KEY: OpRole.Optimize,
|
|
},
|
|
)
|
|
|
|
# lr_var = main_block.var("gradient_merge_current_step")
|
|
# paddle.static.Print(lr_var, message="in OPTIMIZE last conditional")
|
|
|
|
def _sharding_gradient_merge(self):
|
|
"""
|
|
copy all optimize ops in origin main block
|
|
remove all optimize ops in origin main block
|
|
create cond block
|
|
|
|
"""
|
|
if (
|
|
self.gradient_merge_mode != "sharding_gm"
|
|
or self._gradient_merge_acc_step <= 1
|
|
):
|
|
return
|
|
|
|
main_block = self._main_program.global_block()
|
|
# copy original optimize ops to temp ops desc list
|
|
# remove them from block 0
|
|
tmp_copy_block = self._main_program._create_block()
|
|
|
|
self.original_optimize_ops_desc = []
|
|
for op_idx, op in reversed(list(enumerate(main_block.ops))):
|
|
if int(op.attr('op_role')) != int(OpRole.Optimize):
|
|
continue
|
|
else:
|
|
tmp_op_desc = tmp_copy_block.desc.append_op()
|
|
tmp_op_desc.copy_from(op.desc)
|
|
self.original_optimize_ops_desc.append(tmp_op_desc)
|
|
main_block._remove_op(op_idx, sync=False)
|
|
tmp_copy_block._sync_with_cpp()
|
|
self.original_optimize_ops_desc = list(
|
|
reversed(self.original_optimize_ops_desc)
|
|
)
|
|
|
|
# back to block 0
|
|
self._main_program._rollback()
|
|
|
|
# create cond vars and ops at the end of block 0
|
|
cond = self._create_gm_cond(main_block)
|
|
|
|
# create cond block
|
|
cond_block = self._main_program._create_block()
|
|
self._true_apply_gradient()
|
|
|
|
# back to block 0
|
|
self._main_program._rollback()
|
|
|
|
# cond op
|
|
step_scope = self._main_program.global_block().create_var(
|
|
type=core.VarDesc.VarType.STEP_SCOPES
|
|
)
|
|
conditional_block_op = self._main_program.global_block().append_op(
|
|
type='conditional_block',
|
|
inputs={
|
|
'Cond': cond,
|
|
'Input': [],
|
|
},
|
|
outputs={'Out': [], 'Scope': [step_scope]},
|
|
attrs={
|
|
'sub_block': cond_block,
|
|
'is_scalar_condition': True,
|
|
},
|
|
)
|
|
|
|
|
|
class ThreadShardingOptimizer(ShardingOptimizer):
|
|
"""Sharding Optimizer."""
|
|
|
|
def __init__(self, optimizer):
|
|
super().__init__(optimizer)
|
|
self.inner_opt = optimizer
|
|
self.meta_optimizers_white_list = [
|
|
"ParameterServerOptimizer",
|
|
"RecomputeOptimizer",
|
|
"AMPOptimizer",
|
|
"LarsOptimizer",
|
|
"LambOptimizer",
|
|
"ASPOptimizer",
|
|
# "ModelParallelOptimizer",
|
|
# "PipelineOptimizer",
|
|
]
|
|
self._thread_mode = True
|
|
self._use_calc_stream = False
|
|
op_maker = core.op_proto_and_checker_maker
|
|
self.op_role_key = op_maker.kOpRoleAttrName()
|
|
|
|
def _prune_main_program(self, block, shard, rings):
|
|
"""
|
|
rename BroadCast param
|
|
|
|
"""
|
|
var_names = set()
|
|
for idx, op in enumerate(block.ops):
|
|
for input_name in op.desc.input_arg_names():
|
|
pos = input_name.find("@BroadCast")
|
|
if pos <= 0:
|
|
continue
|
|
new_name = input_name[0:pos]
|
|
op.desc._rename_input(input_name, new_name)
|
|
var_names.add(input_name)
|
|
for output_name in op.desc.output_arg_names():
|
|
pos = output_name.find("@BroadCast")
|
|
if pos <= 0:
|
|
continue
|
|
new_name = output_name[0:pos]
|
|
op.desc._rename_output(output_name, new_name)
|
|
var_names.add(output_name)
|
|
|
|
for var_name in var_names:
|
|
block._remove_var(var_name, sync=False)
|
|
|
|
print("remove broadcast param count=", len(var_names))
|
|
block._sync_with_cpp()
|
|
|
|
def _prune_startup_program(self, block, shard):
|
|
"""
|
|
not need process
|
|
"""
|
|
block._sync_with_cpp()
|
|
|
|
def minimize_impl(
|
|
self, loss, startup_program=None, parameter_list=None, no_grad_set=None
|
|
):
|
|
"""
|
|
reset start program and main program
|
|
"""
|
|
sharding_configs = self.user_defined_strategy.sharding_configs
|
|
if "use_calc_stream" in sharding_configs:
|
|
self._use_calc_stream = sharding_configs["use_calc_stream"]
|
|
optimize_ops, params_grads = super().minimize_impl(
|
|
loss, startup_program, parameter_list, no_grad_set
|
|
)
|
|
# main_block = self._main_program.global_block()
|
|
# startup_block = self._startup_program.global_block()
|
|
loss.block.program = self._main_program
|
|
from paddle import fluid
|
|
|
|
fluid.framework.switch_startup_program(self._startup_program)
|
|
return optimize_ops, params_grads
|
|
|
|
def _init_comm(self):
|
|
# sync var
|
|
self.role_id = self.role_maker._role_id()
|
|
self.node_nums = self.role_maker._node_num()
|
|
startup_block = self._startup_program.global_block()
|
|
node_nums = len(self.global_endpoints)
|
|
assert self.node_nums == node_nums, "end points not equal node nums"
|
|
self.current_endpoint = self.global_endpoints[self.role_id]
|
|
|
|
# mp ring
|
|
if self.mp_degree > 1:
|
|
self._init_communicator(
|
|
self._startup_program,
|
|
self.current_endpoint,
|
|
self.mp_group_endpoints,
|
|
self.role_id,
|
|
self.mp_ring_id,
|
|
)
|
|
|
|
# sharding ring
|
|
if self.sharding_degree > 1:
|
|
self._init_communicator(
|
|
self._startup_program,
|
|
self.current_endpoint,
|
|
self.sharding_group_endpoints,
|
|
self.role_id,
|
|
self.sharding_ring_id,
|
|
)
|
|
|
|
# pure dp ring
|
|
if self.dp_degree > 1:
|
|
self._init_communicator(
|
|
self._startup_program,
|
|
self.current_endpoint,
|
|
self.dp_group_endpoints,
|
|
self.role_id,
|
|
self.dp_ring_id,
|
|
)
|
|
|
|
startup_block._sync_with_cpp()
|
|
|
|
def _wait(self):
|
|
if len(self.global_endpoints) <= 1:
|
|
return
|
|
endpoints = self.global_endpoints[:]
|
|
current_endpoint = endpoints[self.role_maker._role_id()]
|
|
if self.global_rank == 0:
|
|
from paddle.fluid.transpiler.details import wait_server_ready
|
|
|
|
endpoints.remove(current_endpoint)
|
|
wait_server_ready(endpoints)
|
|
|
|
def _init_communicator(
|
|
self, program, current_endpoint, endpoints, role_id, ring_id
|
|
):
|
|
nranks = len(endpoints)
|
|
block = program.global_block()
|
|
# init multi node nccl
|
|
if nranks > 1:
|
|
other_endpoints = endpoints[:]
|
|
other_endpoints.remove(current_endpoint)
|
|
|
|
nccl_id_var = block.create_var(
|
|
name=unique_name.generate('nccl_id'),
|
|
persistable=True,
|
|
type=core.VarDesc.VarType.RAW,
|
|
)
|
|
block.append_op(
|
|
type='c_gen_nccl_id',
|
|
inputs={},
|
|
outputs={'Out': nccl_id_var},
|
|
attrs={
|
|
'rank': role_id,
|
|
'endpoint': current_endpoint,
|
|
'other_endpoints': other_endpoints,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
block.append_op(
|
|
type='c_comm_init_multitrainer',
|
|
inputs={'X': nccl_id_var},
|
|
outputs={},
|
|
attrs={
|
|
'ntrainers': nranks,
|
|
'trainer_id': role_id,
|
|
'ring_id': ring_id,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
else:
|
|
block.append_op(type='comm_init_all', attrs={'ring_id': ring_id})
|