wehub-resource-sync
1065 lines
38 KiB
Python
1065 lines
38 KiB
Python
# Copyright (c) 2019 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 unique_name
|
|
from paddle.distributed.fleet.base.private_helper_function import (
|
|
wait_server_ready,
|
|
)
|
|
from paddle.framework import core
|
|
from paddle.static import default_main_program, default_startup_program
|
|
|
|
OpRole = core.op_proto_and_checker_maker.OpRole
|
|
|
|
|
|
class Collective:
|
|
''' '''
|
|
|
|
def __init__(self, nrings):
|
|
self.nrings = nrings
|
|
self.endpoints = None
|
|
self.current_endpoint = None
|
|
self.other_endpoints = None
|
|
self.nranks = None
|
|
self.rank = None
|
|
self.startup_program = None
|
|
self.main_program = None
|
|
op_maker = core.op_proto_and_checker_maker
|
|
self.op_role_key = op_maker.kOpRoleAttrName()
|
|
self.op_role_var_key = op_maker.kOpRoleVarAttrName()
|
|
|
|
def transpile(
|
|
self,
|
|
startup_program,
|
|
main_program,
|
|
rank,
|
|
endpoints,
|
|
current_endpoint,
|
|
wait_port,
|
|
):
|
|
# in case of '127.0.0.1:6700,127.0.0.1:6701,...'
|
|
if isinstance(endpoints, str):
|
|
endpoints = endpoints.split(',')
|
|
|
|
self.startup_program = startup_program
|
|
if startup_program is None:
|
|
self.startup_program = default_startup_program()
|
|
|
|
self.main_program = main_program
|
|
if main_program is None:
|
|
self.main_program = default_main_program()
|
|
|
|
self.nranks = len(endpoints)
|
|
if (
|
|
self.nranks == 1
|
|
and self.mode != "single_process_multi_thread"
|
|
and self.mode != "box"
|
|
):
|
|
raise ValueError('the number of endpoints must > 1')
|
|
|
|
if rank < 0:
|
|
raise ValueError('rank must >= 0')
|
|
self.rank = rank
|
|
|
|
if current_endpoint not in endpoints:
|
|
raise ValueError(
|
|
'current endpoint %s is not in %s',
|
|
current_endpoint,
|
|
str(endpoints),
|
|
)
|
|
|
|
self.endpoints = endpoints
|
|
self.current_endpoint = current_endpoint
|
|
|
|
if current_endpoint:
|
|
nranks = len(endpoints)
|
|
other_endpoints = endpoints[:]
|
|
other_endpoints.remove(current_endpoint)
|
|
self.other_endpoints = other_endpoints
|
|
|
|
self.wait_port = wait_port
|
|
|
|
self.startup_program._origin_program = self.startup_program.clone()
|
|
self._transpile_startup_program()
|
|
|
|
self.main_program._origin_program = self.main_program.clone()
|
|
self._transpile_main_program()
|
|
|
|
def _transpile_main_program(self):
|
|
raise NotImplementedError('call the inherited method of subclasses')
|
|
|
|
def _transpile_startup_program(self):
|
|
for ring_id in range(self.nrings):
|
|
self._init_communicator(
|
|
self.startup_program,
|
|
self.current_endpoint,
|
|
self.endpoints,
|
|
self.rank,
|
|
ring_id,
|
|
self.wait_port,
|
|
)
|
|
self._broadcast_params()
|
|
|
|
def _init_communicator(
|
|
self,
|
|
program,
|
|
current_endpoint,
|
|
endpoints,
|
|
rank,
|
|
ring_id,
|
|
wait_port,
|
|
has_multitrainer=False,
|
|
):
|
|
endpoints_str = ",".join(endpoints)
|
|
nranks = len(endpoints)
|
|
other_endpoints = endpoints[:]
|
|
other_endpoints.remove(current_endpoint)
|
|
block = program.global_block()
|
|
|
|
if rank == 0 and wait_port:
|
|
wait_server_ready(other_endpoints)
|
|
|
|
block = program.global_block()
|
|
if core.is_compiled_with_cuda():
|
|
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': rank,
|
|
'endpoint': current_endpoint,
|
|
'other_endpoints': other_endpoints,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
if not has_multitrainer:
|
|
# 'endpoints': endpoints_str,
|
|
block.append_op(
|
|
type='c_comm_init',
|
|
inputs={'X': nccl_id_var},
|
|
outputs={},
|
|
attrs={
|
|
'nranks': nranks,
|
|
'rank': rank,
|
|
'ring_id': ring_id,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
else:
|
|
block.append_op(
|
|
type='c_comm_init_multitrainer',
|
|
inputs={'X': nccl_id_var},
|
|
outputs={},
|
|
attrs={
|
|
'ntrainers': nranks,
|
|
'trainer_id': rank,
|
|
'ring_id': ring_id,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
elif core.is_compiled_with_xpu():
|
|
bkcl_id_var = block.create_var(
|
|
name=unique_name.generate('bkcl_id'),
|
|
persistable=True,
|
|
type=core.VarDesc.VarType.RAW,
|
|
)
|
|
block.append_op(
|
|
type='c_gen_bkcl_id',
|
|
inputs={},
|
|
outputs={'Out': bkcl_id_var},
|
|
attrs={
|
|
'rank': rank,
|
|
'endpoint': current_endpoint,
|
|
'other_endpoints': other_endpoints,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
block.append_op(
|
|
type='c_comm_init',
|
|
inputs={'X': bkcl_id_var},
|
|
outputs={},
|
|
attrs={
|
|
'nranks': nranks,
|
|
'rank': rank,
|
|
'ring_id': ring_id,
|
|
'endpoints': endpoints_str,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
elif (
|
|
paddle.distributed.ParallelEnv().device_type
|
|
in paddle.device.get_all_custom_device_type()
|
|
):
|
|
xccl_id_var = block.create_var(
|
|
name=unique_name.generate('xccl_id'),
|
|
persistable=True,
|
|
type=core.VarDesc.VarType.RAW,
|
|
)
|
|
block.append_op(
|
|
type='c_gen_xccl_id',
|
|
inputs={},
|
|
outputs={'Out': xccl_id_var},
|
|
attrs={
|
|
'rank': rank,
|
|
'endpoint': current_endpoint,
|
|
'other_endpoints': other_endpoints,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
block.append_op(
|
|
type='c_comm_init',
|
|
inputs={'X': xccl_id_var},
|
|
outputs={},
|
|
attrs={
|
|
'nranks': nranks,
|
|
'rank': rank,
|
|
'ring_id': ring_id,
|
|
'endpoints': endpoints_str,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
|
|
def _broadcast_params(self):
|
|
block = self.startup_program.global_block()
|
|
ring_id = -1
|
|
for param in block.iter_parameters():
|
|
if param.is_distributed:
|
|
continue
|
|
|
|
ring_id = (ring_id + 1) % self.nrings
|
|
block.append_op(
|
|
type='broadcast',
|
|
inputs={'x': param},
|
|
outputs={'out': param},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
'root': 0,
|
|
self.op_role_key: OpRole.Forward,
|
|
},
|
|
)
|
|
|
|
for ring_id in range(self.nrings):
|
|
block.append_op(
|
|
type='c_sync_comm_stream',
|
|
inputs={'X': param},
|
|
outputs={'Out': param},
|
|
attrs={'ring_id': ring_id, self.op_role_key: OpRole.Forward},
|
|
)
|
|
|
|
def _is_loss_grad_op(self, op):
|
|
if self.op_role_key not in op.attr_names:
|
|
return False
|
|
op_role = int(op.all_attrs()[self.op_role_key])
|
|
return op_role & int(OpRole.Backward) and op_role & int(OpRole.Loss)
|
|
|
|
def _is_backward_op(self, op):
|
|
return self.op_role_key in op.attr_names and int(
|
|
op.all_attrs()[self.op_role_key]
|
|
) & int(OpRole.Backward)
|
|
|
|
def _is_update_op(self, op):
|
|
return (
|
|
'Param' in op.input_names
|
|
and 'Grad' in op.input_names
|
|
and "LearningRate" in op.input_names
|
|
)
|
|
|
|
def _is_optimizer_op(self, op):
|
|
return self.op_role_key in op.attr_names and int(
|
|
op.all_attrs()[self.op_role_key]
|
|
) & int(OpRole.Optimize)
|
|
|
|
|
|
class GradAllReduce(Collective):
|
|
''' '''
|
|
|
|
def __init__(self, nrings=2):
|
|
Collective.__init__(self, nrings)
|
|
self.mode = "grad_allreduce"
|
|
|
|
def _transpile_main_program(self):
|
|
self._insert_scale_loss_grad_ops()
|
|
self._insert_allreduce_ops()
|
|
|
|
def _insert_scale_loss_grad_ops(self):
|
|
'''
|
|
In order to keep the learning rate consistent in different numbers of
|
|
training workers, we scale the loss grad by the number of workers
|
|
'''
|
|
block = self.main_program.global_block()
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if self._is_loss_grad_op(op):
|
|
loss_grad_var = block.vars[op.output_arg_names[0]]
|
|
block._insert_op(
|
|
idx + 1,
|
|
type='scale',
|
|
inputs={'X': loss_grad_var},
|
|
outputs={'Out': loss_grad_var},
|
|
attrs={
|
|
'scale': 1.0 / self.nranks,
|
|
self.op_role_key: OpRole.Backward,
|
|
},
|
|
)
|
|
|
|
def _insert_allreduce_ops(self):
|
|
block = self.main_program.global_block()
|
|
ring_id = -1
|
|
grad = None
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if (
|
|
self._is_backward_op(op)
|
|
and self.op_role_var_key in op.attr_names
|
|
):
|
|
op_role_var = op.all_attrs()[self.op_role_var_key]
|
|
|
|
if len(op_role_var) == 0:
|
|
continue
|
|
assert len(op_role_var) % 2 == 0
|
|
|
|
offset = idx
|
|
for i in range(0, len(op_role_var), 2):
|
|
param = block.vars[op_role_var[i]]
|
|
grad = block.vars[op_role_var[i + 1]]
|
|
if param.is_distributed:
|
|
continue
|
|
|
|
if offset == idx:
|
|
offset += 1
|
|
block._insert_op(
|
|
offset,
|
|
type='c_sync_calc_stream',
|
|
inputs={'X': grad},
|
|
outputs={'Out': grad},
|
|
attrs={self.op_role_key: OpRole.Backward},
|
|
)
|
|
offset += 1
|
|
|
|
# As we search ops reversely, we should insert all_reduce sum
|
|
# op in the same way to keep the ring_id alternate
|
|
ring_id = (ring_id + 1) % self.nrings
|
|
block._insert_op(
|
|
offset,
|
|
type='all_reduce',
|
|
inputs={'x': grad},
|
|
outputs={'out': grad},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
|
self.op_role_key: OpRole.Backward,
|
|
},
|
|
)
|
|
|
|
if grad is None:
|
|
return
|
|
|
|
for idx, op in enumerate(block.ops):
|
|
if self._is_optimizer_op(op):
|
|
for ring_id in range(self.nrings):
|
|
block._insert_op(
|
|
idx + ring_id,
|
|
type='c_sync_comm_stream',
|
|
inputs={'X': grad},
|
|
outputs={'Out': grad},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
self.op_role_key: OpRole.Backward,
|
|
},
|
|
)
|
|
break
|
|
|
|
|
|
class LocalSGD(Collective):
|
|
''' '''
|
|
|
|
def __init__(self, nrings=2):
|
|
Collective.__init__(self, nrings)
|
|
self.snapshot_key = '@SNAPSHOT'
|
|
self.mode = "local_sgd"
|
|
|
|
def _transpile_startup_program(self):
|
|
Collective._transpile_startup_program(self)
|
|
|
|
block = self.startup_program.global_block()
|
|
non_dist_params = []
|
|
for param in block.iter_parameters():
|
|
if not param.is_distributed:
|
|
non_dist_params.append(param)
|
|
|
|
for param in non_dist_params:
|
|
snapshot = block.create_var(
|
|
name=self.snapshot_name(param.name),
|
|
shape=param.shape,
|
|
persistable=True,
|
|
stop_gradient=True,
|
|
)
|
|
block.append_op(
|
|
type='assign',
|
|
inputs={'X': [param]},
|
|
outputs={'Out': [snapshot]},
|
|
attrs={self.op_role_key: OpRole.Forward},
|
|
)
|
|
|
|
def snapshot_name(self, param_name):
|
|
return param_name + self.snapshot_key
|
|
|
|
def _transpile_main_program(self):
|
|
block = self.main_program.global_block()
|
|
ordered_param_snapshot = []
|
|
ring_id = -1
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if self._is_update_op(op):
|
|
param = block.vars[op.input('Param')[0]]
|
|
if param.is_distributed:
|
|
continue
|
|
|
|
snapshot = block.create_var(
|
|
name=self.snapshot_name(param.name),
|
|
shape=param.shape,
|
|
persistable=True,
|
|
stop_gradient=True,
|
|
dtype=param.dtype,
|
|
)
|
|
|
|
block._insert_op(
|
|
idx + 1,
|
|
type='elementwise_sub',
|
|
inputs={'X': [snapshot], 'Y': [param]},
|
|
outputs={'Out': [param]},
|
|
attrs={self.op_role_key: OpRole.Optimize},
|
|
)
|
|
block._insert_op(
|
|
idx + 2,
|
|
type='c_sync_calc_stream',
|
|
inputs={'X': param},
|
|
outputs={'Out': param},
|
|
attrs={self.op_role_key: OpRole.Optimize},
|
|
)
|
|
ring_id = (ring_id + 1) % self.nrings
|
|
block._insert_op(
|
|
idx + 3,
|
|
type='all_reduce',
|
|
inputs={'x': [param]},
|
|
outputs={'out': [param]},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
|
self.op_role_key: OpRole.Optimize,
|
|
},
|
|
)
|
|
|
|
ordered_param_snapshot.append((param, snapshot))
|
|
|
|
for ring_id in range(self.nrings):
|
|
block.append_op(
|
|
type='c_sync_comm_stream',
|
|
inputs={'X': param},
|
|
outputs={'Out': param},
|
|
attrs={'ring_id': ring_id, self.op_role_key: OpRole.Optimize},
|
|
)
|
|
|
|
for param_snapshot in reversed(ordered_param_snapshot):
|
|
param = param_snapshot[0]
|
|
snapshot = param_snapshot[1]
|
|
block.append_op(
|
|
type='scale',
|
|
inputs={'X': [param]},
|
|
outputs={'Out': [param]},
|
|
attrs={
|
|
'scale': 1.0 / self.nranks,
|
|
self.op_role_key: OpRole.Optimize,
|
|
},
|
|
)
|
|
block.append_op(
|
|
type='elementwise_sub',
|
|
inputs={'X': [snapshot], 'Y': [param]},
|
|
outputs={'Out': [param]},
|
|
attrs={self.op_role_key: OpRole.Optimize},
|
|
)
|
|
block.append_op(
|
|
type='assign',
|
|
inputs={'X': [param]},
|
|
outputs={'Out': [snapshot]},
|
|
attrs={self.op_role_key: OpRole.Optimize},
|
|
)
|
|
|
|
|
|
class SingleProcessMultiThread(GradAllReduce):
|
|
"""
|
|
single process multi thread mode
|
|
"""
|
|
|
|
def __init__(self):
|
|
GradAllReduce.__init__(self, 1)
|
|
self.mode = "single_process_multi_thread"
|
|
self.fuse_allreduce = int(os.getenv("PADDLE_FUSE_ALLREDUCE", "1"))
|
|
self.loss_scale = int(os.getenv("PADDLE_LOSS_SCALE", "1"))
|
|
self.gpu_nums = len(
|
|
os.getenv("FLAGS_selected_gpus", "0,1,2,3,4,5,6,7").split(",")
|
|
)
|
|
|
|
def _transpile_startup_program(self):
|
|
nodes_num = 0
|
|
if len(self.endpoints) > 1:
|
|
nodes_num = len({x.split(':')[0] for x in self.endpoints})
|
|
# different ip num is multi node
|
|
if nodes_num > 1:
|
|
self.nranks = nodes_num
|
|
print("begin to _transpile_startup_program for multi-node")
|
|
print("current_endpoint: ", self.current_endpoint)
|
|
print("total endpoints: ", self.endpoints)
|
|
print(f"rank: {self.rank}, ring_id: {self.nrings}")
|
|
for ring_id in range(self.nrings):
|
|
self._init_communicator(
|
|
self.startup_program,
|
|
self.current_endpoint,
|
|
self.endpoints,
|
|
self.rank,
|
|
ring_id,
|
|
self.wait_port,
|
|
True,
|
|
)
|
|
else:
|
|
self.nranks = 1
|
|
print("begin to _transpile_startup_program for single-node")
|
|
block = self.startup_program.global_block()
|
|
block.append_op(type='comm_init_all', attrs={'ring_id': 0})
|
|
|
|
def _transpile_main_program(self):
|
|
# not need loss scale and no dense param
|
|
param_cnt = self._get_update_param_count()
|
|
if self.loss_scale == 0 and param_cnt == 0:
|
|
return
|
|
# scale loss
|
|
if self.loss_scale:
|
|
self._insert_scale_loss_grad_ops(param_cnt)
|
|
# no param
|
|
if param_cnt == 0:
|
|
return
|
|
# fuse allreduce
|
|
if self.fuse_allreduce > 0:
|
|
print(f"begin used fuse_allreduce param count = {param_cnt}")
|
|
# use fuse allreduce
|
|
self._insert_fuse_allreduce_ops()
|
|
else:
|
|
self._insert_allreduce_ops()
|
|
|
|
def _get_update_param_count(self):
|
|
"""
|
|
get need update param count
|
|
"""
|
|
param_count = 0
|
|
block = self.main_program.global_block()
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if not self._is_backward_op(op):
|
|
continue
|
|
if self.op_role_var_key not in op.attr_names:
|
|
continue
|
|
op_role_var = op.all_attrs()[self.op_role_var_key]
|
|
if len(op_role_var) == 0:
|
|
continue
|
|
|
|
assert len(op_role_var) % 2 == 0
|
|
for i in range(0, len(op_role_var), 2):
|
|
param = block.vars[op_role_var[i]]
|
|
if param.is_distributed:
|
|
continue
|
|
param_count = param_count + 1
|
|
|
|
return param_count
|
|
|
|
def _insert_scale_loss_grad_ops(self, param_cnt):
|
|
'''
|
|
In order to keep the learning rate consistent in different numbers of
|
|
training workers, we scale the loss grad by the number of workers
|
|
'''
|
|
if param_cnt > 0:
|
|
scale = 1.0 / self.nranks / self.gpu_nums
|
|
else:
|
|
scale = 1.0 / self.gpu_nums
|
|
print(f"begin _insert_scale_loss_grad_ops scale = {scale}")
|
|
block = self.main_program.global_block()
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if not self._is_loss_grad_op(op):
|
|
continue
|
|
loss_grad_var = block.vars[op.output_arg_names[0]]
|
|
block._insert_op(
|
|
idx + 1,
|
|
type='scale',
|
|
inputs={'X': loss_grad_var},
|
|
outputs={'Out': loss_grad_var},
|
|
attrs={'scale': scale, self.op_role_key: OpRole.Backward},
|
|
)
|
|
|
|
def _insert_fuse_allreduce_ops(self):
|
|
"""
|
|
insert coalesce_tensor and all reduce ops
|
|
"""
|
|
block = self.main_program.global_block()
|
|
ring_id = -1
|
|
grad = None
|
|
input_grads = []
|
|
global_offset = 0 # find insert offset of fuse tensor, after the max dense grad offset
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if (
|
|
self._is_backward_op(op)
|
|
and self.op_role_var_key in op.attr_names
|
|
):
|
|
op_role_var = op.all_attrs()[self.op_role_var_key]
|
|
if len(op_role_var) == 0:
|
|
continue
|
|
assert len(op_role_var) % 2 == 0
|
|
offset = idx
|
|
for i in range(0, len(op_role_var), 2):
|
|
param = block.vars[op_role_var[i]]
|
|
grad = block.vars[op_role_var[i + 1]]
|
|
if param.is_distributed:
|
|
continue
|
|
if offset == idx:
|
|
input_grads.append(grad)
|
|
global_offset = max(global_offset, offset + 1)
|
|
if grad is None:
|
|
return
|
|
|
|
if self.fuse_allreduce == 2:
|
|
# grads aggregation of multi-gpus
|
|
block._insert_op(
|
|
global_offset,
|
|
type='c_sync_calc_stream',
|
|
inputs={'X': input_grads[0]},
|
|
outputs={'Out': input_grads[0]},
|
|
attrs={self.op_role_key: OpRole.Backward},
|
|
)
|
|
global_offset += 1
|
|
ring_id = (ring_id + 1) % self.nrings
|
|
block._insert_op(
|
|
global_offset,
|
|
type='c_allreduce_xsum',
|
|
inputs={'X': input_grads},
|
|
outputs={'Out': input_grads},
|
|
attrs={'ring_id': ring_id, self.op_role_key: OpRole.Backward},
|
|
)
|
|
global_offset += 1
|
|
# sync before adam
|
|
block._insert_op(
|
|
global_offset,
|
|
type='c_sync_comm_stream',
|
|
inputs={'X': input_grads[0]},
|
|
outputs={'Out': input_grads[0]},
|
|
attrs={'ring_id': ring_id, self.op_role_key: OpRole.Backward},
|
|
)
|
|
global_offset += 1
|
|
else:
|
|
# init output_grads
|
|
output_grads = input_grads
|
|
# init fused_output with temp shape, it will calculate real shape depend on inputs
|
|
fused_output = block.create_var(
|
|
name="fused_output",
|
|
shape=[1],
|
|
persistable=False,
|
|
dtype=core.VarDesc.VarType.FP32,
|
|
stop_gradient=True,
|
|
)
|
|
# fuse all grad tensors
|
|
coalesce_tensor_attrs = {
|
|
"copy_data": True,
|
|
"set_constant": False,
|
|
"dtype": core.VarDesc.VarType.FP32,
|
|
}
|
|
block._insert_op(
|
|
global_offset,
|
|
type='coalesce_tensor',
|
|
inputs={'Input': input_grads},
|
|
outputs={'Output': output_grads, 'FusedOutput': fused_output},
|
|
attrs=coalesce_tensor_attrs,
|
|
)
|
|
global_offset += 1
|
|
# grads aggregation of multi-gpus
|
|
block._insert_op(
|
|
global_offset,
|
|
type='c_sync_calc_stream',
|
|
inputs={'X': fused_output},
|
|
outputs={'Out': fused_output},
|
|
attrs={self.op_role_key: OpRole.Backward},
|
|
)
|
|
global_offset += 1
|
|
ring_id = (ring_id + 1) % self.nrings
|
|
block._insert_op(
|
|
global_offset,
|
|
type='all_reduce',
|
|
inputs={'x': fused_output},
|
|
outputs={'out': fused_output},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
self.op_role_key: OpRole.Backward,
|
|
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
|
},
|
|
)
|
|
global_offset += 1
|
|
|
|
# sync before adam
|
|
block._insert_op(
|
|
global_offset,
|
|
type='c_sync_comm_stream',
|
|
inputs={'X': fused_output},
|
|
outputs={'Out': fused_output},
|
|
attrs={'ring_id': ring_id, self.op_role_key: OpRole.Backward},
|
|
)
|
|
global_offset += 1
|
|
|
|
|
|
class MultiThread(GradAllReduce):
|
|
''' '''
|
|
|
|
def __init__(self, nrings=1, trans_mode="fuse_all_reduce"):
|
|
GradAllReduce.__init__(self, nrings)
|
|
self.mode = "box"
|
|
self.trans_mode = trans_mode
|
|
self.fuse_grad_size_in_num = 128
|
|
gpu_nums = os.getenv("FLAGS_selected_gpus", "0,1,2,3,4,5,6,7,8").split(
|
|
","
|
|
)
|
|
self.gpu_num = len(gpu_nums)
|
|
|
|
def _transpile_startup_program(self):
|
|
if len(self.endpoints) > 1:
|
|
print("begin to _transpile_startup_program for multi-node")
|
|
print("current_endpoint: ", self.current_endpoint)
|
|
print("total endpoints: ", self.endpoints)
|
|
print(f"rank: {self.rank}, ring_id: {self.nrings}")
|
|
for ring_id in range(self.nrings):
|
|
self._init_communicator(
|
|
self.startup_program,
|
|
self.current_endpoint,
|
|
self.endpoints,
|
|
self.rank,
|
|
ring_id,
|
|
self.wait_port,
|
|
True,
|
|
)
|
|
|
|
else:
|
|
if "xpu" in self.trans_mode:
|
|
print(
|
|
"begin to _transpile_startup_program for single-node in XPU"
|
|
)
|
|
block = self.startup_program.global_block()
|
|
block.append_op(
|
|
type='comm_init_all',
|
|
attrs={
|
|
'devices': list(
|
|
map(
|
|
int, os.getenv("FLAGS_selected_gpus").split(",")
|
|
)
|
|
),
|
|
'ring_id': 0,
|
|
},
|
|
)
|
|
else:
|
|
print("begin to _transpile_startup_program for single-node")
|
|
block = self.startup_program.global_block()
|
|
block.append_op(type='comm_init_all', attrs={'ring_id': 0})
|
|
|
|
def _transpile_main_program(self):
|
|
self._insert_scale_loss_grad_ops()
|
|
if self.trans_mode == "all_gather":
|
|
print("begin to transpile in all-gather mode")
|
|
self.allgather_ranks = self.nranks * self.gpu_num
|
|
self._insert_allgather_ops()
|
|
self._update_adam_ops()
|
|
elif self.trans_mode == "fuse_all_reduce":
|
|
print("begin to transpile in fuse all-reduce mode")
|
|
self._insert_fuse_allreduce_ops()
|
|
elif (
|
|
self.trans_mode == "all_reduce_xpu"
|
|
and len(os.getenv("FLAGS_selected_gpus").split(",")) == 1
|
|
):
|
|
print(
|
|
"skip transpile in all-reduce-xpu mode when number of devices is only one"
|
|
)
|
|
else:
|
|
print("begin to transpile in all-reduce mode")
|
|
self._insert_allreduce_ops()
|
|
|
|
def _insert_allgather_ops(self):
|
|
"""
|
|
insert allgather op to the main_program
|
|
"""
|
|
block = self.main_program.global_block()
|
|
ring_id = -1
|
|
grad = None
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if (
|
|
self._is_backward_op(op)
|
|
and self.op_role_var_key in op.attr_names
|
|
):
|
|
op_role_var = op.all_attrs()[self.op_role_var_key]
|
|
if len(op_role_var) == 0:
|
|
continue
|
|
assert len(op_role_var) % 2 == 0
|
|
|
|
offset = idx
|
|
for i in range(0, len(op_role_var), 2):
|
|
param = block.vars[op_role_var[i]]
|
|
new_grad_var = block.create_var(
|
|
name=op_role_var[i] + "_allgather",
|
|
shape=[self.allgather_ranks, *list(param.shape)],
|
|
persistable=False,
|
|
dtype=core.VarDesc.VarType.FP32,
|
|
stop_gradient=True,
|
|
)
|
|
grad = block.vars[op_role_var[i + 1]]
|
|
if param.is_distributed: # no need to care: used in PLSC
|
|
continue
|
|
|
|
if offset == idx:
|
|
offset += 1
|
|
block._insert_op(
|
|
offset,
|
|
type='c_sync_calc_stream',
|
|
inputs={'X': grad},
|
|
outputs={'Out': grad},
|
|
attrs={self.op_role_key: OpRole.Backward},
|
|
)
|
|
offset += 1
|
|
|
|
# As we search ops reversely, we should insert all_gather
|
|
# op in the same way to keep the ring_id alternate
|
|
ring_id = (ring_id + 1) % self.nrings
|
|
block._insert_op(
|
|
offset,
|
|
type='all_gather',
|
|
inputs={'x': grad},
|
|
outputs={'out': new_grad_var},
|
|
attrs={
|
|
'nranks': self.allgather_ranks,
|
|
'ring_id': ring_id,
|
|
self.op_role_key: OpRole.Backward,
|
|
},
|
|
)
|
|
|
|
if grad is None:
|
|
return
|
|
|
|
for idx, op in enumerate(block.ops):
|
|
if self._is_optimizer_op(op):
|
|
for ring_id in range(self.nrings):
|
|
block._insert_op(
|
|
idx + ring_id,
|
|
type='c_sync_comm_stream',
|
|
inputs={'X': grad},
|
|
outputs={'Out': grad},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
self.op_role_key: OpRole.Backward,
|
|
},
|
|
)
|
|
break
|
|
|
|
def _update_adam_ops(self):
|
|
"""
|
|
remove the original adam op, and add new adam ops
|
|
"""
|
|
block = self.main_program.global_block()
|
|
|
|
for idx, op in reversed(list(enumerate(block.ops))):
|
|
if self._is_optimizer_op(op):
|
|
offset = idx
|
|
if (
|
|
op.type != 'adam' and op.type != 'lamb'
|
|
): # filter out scale op
|
|
continue
|
|
param_name = op.input("Param")[0]
|
|
inputs = {
|
|
"Param": block.vars[op.input("Param")[0]],
|
|
"LearningRate": block.vars[op.input("LearningRate")[0]],
|
|
"Moment1": block.vars[op.input("Moment1")[0]],
|
|
"Moment2": block.vars[op.input("Moment2")[0]],
|
|
"Beta1Pow": block.vars[op.input("Beta1Pow")[0]],
|
|
"Beta2Pow": block.vars[op.input("Beta2Pow")[0]],
|
|
}
|
|
outputs = {
|
|
"ParamOut": block.vars[op.output("ParamOut")[0]],
|
|
"Moment1Out": block.vars[op.output("Moment1Out")[0]],
|
|
"Moment2Out": block.vars[op.output("Moment2Out")[0]],
|
|
"Beta1PowOut": block.vars[op.output("Beta1PowOut")[0]],
|
|
"Beta2PowOut": block.vars[op.output("Beta2PowOut")[0]],
|
|
}
|
|
attrs = {
|
|
"epsilon": op.attr('epsilon'),
|
|
"beta1": op.attr('beta1'),
|
|
"beta2": op.attr('beta2'),
|
|
"lazy_mode": op.attr('lazy_mode'),
|
|
"min_row_size_to_use_multithread": op.attr(
|
|
'min_row_size_to_use_multithread'
|
|
),
|
|
}
|
|
split_vars = [
|
|
block.create_var(
|
|
name=param_name + "_" + str(i),
|
|
shape=block.vars[op.input("Param")[0]].shape,
|
|
persistable=False,
|
|
dtype=core.VarDesc.VarType.FP32,
|
|
stop_gradient=True,
|
|
)
|
|
for i in range(self.allgather_ranks)
|
|
]
|
|
block._insert_op(
|
|
offset,
|
|
type="split",
|
|
inputs={
|
|
'X': block.vars[op.input("Param")[0] + "_allgather"]
|
|
},
|
|
outputs={'Out': split_vars},
|
|
attrs={'num': self.allgather_ranks, 'axis': 0},
|
|
)
|
|
offset += 1
|
|
|
|
for i in range(self.allgather_ranks):
|
|
inputs["Grad"] = split_vars[i]
|
|
block._insert_op(
|
|
offset,
|
|
type=op.type,
|
|
inputs=inputs,
|
|
outputs=outputs,
|
|
attrs=attrs,
|
|
)
|
|
offset += 1
|
|
# remove the original adam op
|
|
block._remove_op(offset)
|
|
|
|
def _insert_fuse_allreduce_ops(self):
|
|
"""
|
|
insert coalesce_tensor and all reduce ops
|
|
"""
|
|
block = self.main_program.global_block()
|
|
ring_id = 0 % self.nrings
|
|
grad = None
|
|
param_grads = []
|
|
# find all grad params
|
|
for op in reversed(block.ops):
|
|
if (
|
|
self._is_backward_op(op)
|
|
and self.op_role_var_key in op.attr_names
|
|
):
|
|
op_role_var = op.all_attrs()[self.op_role_var_key]
|
|
if len(op_role_var) == 0:
|
|
continue
|
|
assert len(op_role_var) % 2 == 0, (
|
|
"vars need to be one param var followed by one grad var, "
|
|
"but got odd number of vars"
|
|
)
|
|
for i in range(0, len(op_role_var), 2):
|
|
param_name = op_role_var[i]
|
|
param = block.var(param_name)
|
|
grad_name = op_role_var[i + 1]
|
|
grad = block.var(grad_name)
|
|
if param.is_distributed:
|
|
continue
|
|
param_grads.append(grad)
|
|
if grad is None:
|
|
return
|
|
|
|
segments = []
|
|
last_dtype = None
|
|
# split the grad based on dtype and fused size
|
|
for var in param_grads:
|
|
if (
|
|
len(segments) == 0
|
|
or len(segments[-1]) == self.fuse_grad_size_in_num
|
|
or var.dtype != last_dtype
|
|
):
|
|
segments.append([var])
|
|
last_dtype = var.dtype
|
|
else:
|
|
segments[-1].append(var)
|
|
|
|
fused_vars = []
|
|
for idx, op in enumerate(block.ops):
|
|
if self._is_optimizer_op(op):
|
|
for segment in segments:
|
|
# insert coalesce tensor
|
|
tmp_var = block.create_var(
|
|
name=unique_name.generate(
|
|
f'FusedOutput_{segment[0].name}'
|
|
),
|
|
dtype=segment[0].dtype,
|
|
persistable=False,
|
|
stop_gradient=True,
|
|
)
|
|
fused_vars.append(tmp_var)
|
|
block._insert_op(
|
|
idx,
|
|
type="coalesce_tensor",
|
|
inputs={"Input": segment},
|
|
outputs={"Output": segment, "FusedOutput": tmp_var},
|
|
attrs={
|
|
"copy_data": True,
|
|
"use_align": True,
|
|
"dtype": segment[0].dtype,
|
|
self.op_role_key: OpRole.Backward,
|
|
},
|
|
)
|
|
break
|
|
|
|
# insert the allreduce_sum op
|
|
for idx, op in enumerate(block.ops):
|
|
if self._is_optimizer_op(op):
|
|
for fused_var in fused_vars:
|
|
block._insert_op(
|
|
idx,
|
|
type='all_reduce',
|
|
inputs={'x': fused_var},
|
|
outputs={'out': fused_var},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
'reduce_type': paddle.distributed.ReduceOp.SUM,
|
|
self.op_role_key: OpRole.Backward,
|
|
},
|
|
)
|
|
block._insert_op(
|
|
idx,
|
|
type='c_sync_calc_stream',
|
|
inputs={'X': fused_var},
|
|
outputs={'Out': fused_var},
|
|
attrs={self.op_role_key: OpRole.Backward},
|
|
)
|
|
break
|
|
|
|
if len(fused_vars) == 0:
|
|
block._sync_with_cpp()
|
|
return
|
|
|
|
# insert the sync comm op
|
|
for idx, op in enumerate(block.ops):
|
|
if self._is_optimizer_op(op):
|
|
block._insert_op(
|
|
idx,
|
|
type='c_sync_comm_stream',
|
|
inputs={'X': fused_vars[0]},
|
|
outputs={'Out': fused_vars[0]},
|
|
attrs={
|
|
'ring_id': ring_id,
|
|
self.op_role_key: OpRole.Backward,
|
|
},
|
|
)
|
|
break
|
|
block._sync_with_cpp()
|