wehub-resource-sync
341 lines
11 KiB
Python
341 lines
11 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import paddle
|
|
from paddle import framework
|
|
|
|
# (TODO: GhostScreaming) It will be removed later.
|
|
from paddle.base import core
|
|
from paddle.distributed.parallel import (
|
|
_split_tensors,
|
|
build_groups,
|
|
in_dynamic_mode,
|
|
sync_params_buffers,
|
|
)
|
|
|
|
from .log_util import logger
|
|
|
|
__all__ = []
|
|
|
|
|
|
def obtain_optimizer_parameters_list(optimizer):
|
|
if getattr(optimizer, '_param_groups', None) and isinstance(
|
|
optimizer._param_groups[0], dict
|
|
):
|
|
parameters_list = []
|
|
for group in optimizer._param_groups:
|
|
for param in group['params']:
|
|
parameters_list.append(param)
|
|
else:
|
|
parameters_list = list(optimizer._parameter_list)
|
|
|
|
return parameters_list
|
|
|
|
|
|
def _apply_collective_grads(parameters, comm_group, bucket_size, scale=None):
|
|
grad_var_set = set()
|
|
grad_vars = []
|
|
sparse_grad_vars = []
|
|
|
|
for param in parameters:
|
|
if param.trainable and (param._grad_ivar() is not None):
|
|
g_var = param._grad_ivar()
|
|
assert not g_var._is_sparse(), (
|
|
"Now, it doesn't support sparse parameters"
|
|
)
|
|
grad_vars.append(g_var)
|
|
assert g_var not in grad_var_set
|
|
grad_var_set.add(g_var)
|
|
|
|
coalesced_grads_and_vars = build_groups(grad_vars, bucket_size)
|
|
|
|
nranks = (
|
|
paddle.distributed.get_world_size()
|
|
if comm_group is None
|
|
else comm_group.nranks
|
|
)
|
|
|
|
scale = nranks if scale is None else 1.0 / scale
|
|
scale = None if scale == 1.0 else scale
|
|
|
|
for coalesced_grad, _, _ in coalesced_grads_and_vars:
|
|
# need to div nranks
|
|
if scale is not None:
|
|
div_factor = paddle.to_tensor(scale, dtype=coalesced_grad.dtype)
|
|
paddle.base.framework._dygraph_tracer().trace_op(
|
|
type="elementwise_div",
|
|
inputs={'X': coalesced_grad, 'Y': div_factor},
|
|
outputs={'Out': coalesced_grad},
|
|
attrs={'axis': -1},
|
|
)
|
|
paddle.distributed.all_reduce(coalesced_grad, group=comm_group)
|
|
|
|
_split_tensors(coalesced_grads_and_vars)
|
|
|
|
|
|
def _apply_collective_grads_eager(
|
|
parameters, comm_group, bucket_size, scale=None
|
|
):
|
|
grad_var_set = set()
|
|
grad_vars = []
|
|
|
|
for param in parameters:
|
|
g_var = None
|
|
if param.trainable and (param._grad_ivar() is not None):
|
|
g_var = param._grad_ivar()
|
|
if param.trainable and hasattr(param, "main_grad"):
|
|
assert param._grad_ivar() is None, "param.grad is not None"
|
|
g_var = param.main_grad
|
|
if g_var is not None:
|
|
assert not g_var.is_sparse(), (
|
|
"Now, it doesn't support sparse parameters"
|
|
)
|
|
grad_vars.append(g_var)
|
|
assert g_var not in grad_var_set
|
|
grad_var_set.add(g_var)
|
|
|
|
if len(grad_vars) == 0:
|
|
return
|
|
|
|
coalesced_grads_and_vars = build_groups(grad_vars, bucket_size)
|
|
|
|
nranks = (
|
|
paddle.distributed.get_world_size()
|
|
if comm_group is None
|
|
else comm_group.nranks
|
|
)
|
|
|
|
scale = 1.0 / nranks if scale is None else scale
|
|
scale = None if scale == 1.0 else scale
|
|
|
|
for coalesced_grad, _, _ in coalesced_grads_and_vars:
|
|
# need to div nranks
|
|
if scale is not None:
|
|
coalesced_grad.scale_(scale)
|
|
paddle.distributed.all_reduce(coalesced_grad, group=comm_group)
|
|
|
|
_split_tensors(coalesced_grads_and_vars)
|
|
|
|
|
|
def _broadcast_data_help(data, shape, dtype, hcg):
|
|
model_parallel_group = hcg.get_model_parallel_group()
|
|
src_rank = hcg.get_model_parallel_group_src_rank()
|
|
mp_rank = hcg.get_model_parallel_rank()
|
|
|
|
shape_gpu = paddle.to_tensor(shape, dtype="int32")
|
|
paddle.distributed.broadcast(
|
|
shape_gpu, src=src_rank, group=model_parallel_group, sync_op=True
|
|
)
|
|
|
|
if mp_rank != 0:
|
|
input_data = paddle.zeros(shape_gpu, dtype=dtype)
|
|
else:
|
|
input_data = data
|
|
|
|
paddle.distributed.broadcast(
|
|
input_data, src=src_rank, group=model_parallel_group, sync_op=True
|
|
)
|
|
|
|
if mp_rank != 0:
|
|
if in_dynamic_mode():
|
|
data._clear_data()
|
|
input_data._share_buffer_to(data)
|
|
else:
|
|
data.value().get_tensor()._clear()
|
|
data.value().get_tensor()._share_data_with(
|
|
input_data.value().get_tensor()
|
|
)
|
|
|
|
|
|
def _broadcast_object_list_help(object_list, hcg):
|
|
model_parallel_group = hcg.get_model_parallel_group()
|
|
src_rank = hcg.get_model_parallel_group_src_rank()
|
|
mp_rank = hcg.get_model_parallel_rank()
|
|
|
|
paddle.distributed.broadcast_object_list(
|
|
object_list, src=src_rank, group=model_parallel_group
|
|
)
|
|
|
|
|
|
def _process_element(hcg, dev, place, element):
|
|
if isinstance(element, core.eager.Tensor):
|
|
with framework.no_grad():
|
|
if (
|
|
in_dynamic_mode()
|
|
and not eval(f"element.place.is_{dev}_place")()
|
|
):
|
|
element_gpu = element._copy_to(place, True)
|
|
element._clear_data()
|
|
element_gpu._share_buffer_to(element)
|
|
_broadcast_data_help(element, element.shape, element.dtype, hcg)
|
|
elif isinstance(element, (dict, list, tuple)):
|
|
return _broadcast_nested_data(hcg, dev, place, element)
|
|
else:
|
|
_broadcast_object_list_help([element], hcg)
|
|
|
|
|
|
def _broadcast_nested_data(hcg, dev, place, data):
|
|
if isinstance(data, dict):
|
|
return {
|
|
key: _process_element(hcg, dev, place, value)
|
|
for key, value in data.items()
|
|
}
|
|
elif isinstance(data, list):
|
|
return [_process_element(hcg, dev, place, item) for item in data]
|
|
elif isinstance(data, tuple):
|
|
return tuple(_process_element(hcg, dev, place, item) for item in data)
|
|
else:
|
|
raise TypeError(f"Unsupported data type: {type(data)}")
|
|
|
|
|
|
def broadcast_input_data(hcg, *inputs, **kwargs):
|
|
cur_device = paddle.get_device()
|
|
dev = cur_device.split(":")[0]
|
|
assert (
|
|
dev
|
|
in [
|
|
"xpu",
|
|
"gpu",
|
|
]
|
|
or dev in paddle.device.get_all_custom_device_type()
|
|
), f"Only support xpu, gpu and custom_device now, but this is {dev}"
|
|
dev_idx = int(cur_device.split(':')[1])
|
|
if dev == "gpu":
|
|
place = paddle.CUDAPlace(dev_idx)
|
|
elif dev in paddle.device.get_all_custom_device_type():
|
|
place = paddle.CustomPlace(dev, dev_idx)
|
|
dev = 'custom'
|
|
else:
|
|
place = eval(f"paddle.{dev.upper()}Place")(dev_idx)
|
|
|
|
if len(inputs) > 0:
|
|
inputs = _broadcast_nested_data(hcg, dev, place, inputs)
|
|
if len(kwargs) > 0:
|
|
kwargs = _broadcast_nested_data(hcg, dev, place, kwargs)
|
|
return inputs, kwargs
|
|
|
|
|
|
def broadcast_mp_parameters(model, hcg, fuse_params=True):
|
|
model_parallel_group = hcg.get_model_parallel_group()
|
|
src_rank = hcg.get_model_parallel_group_src_rank()
|
|
sync_params_buffers(
|
|
model,
|
|
model_parallel_group,
|
|
src_rank,
|
|
is_model_parallel=True,
|
|
fuse_params=fuse_params,
|
|
)
|
|
|
|
|
|
def broadcast_dp_parameters(model, hcg, fuse_params=True):
|
|
data_parallel_group = hcg.get_data_parallel_group()
|
|
src_rank = hcg.get_data_parallel_group_src_rank()
|
|
sync_params_buffers(
|
|
model,
|
|
data_parallel_group,
|
|
src_rank,
|
|
is_model_parallel=False,
|
|
fuse_params=fuse_params,
|
|
)
|
|
|
|
|
|
def fused_allreduce_gradients_with_group(
|
|
parameter_list, group, bucket_size=128 * 1024 * 1024, scale=None
|
|
):
|
|
apply_func = (
|
|
_apply_collective_grads_eager
|
|
if in_dynamic_mode()
|
|
else _apply_collective_grads
|
|
)
|
|
with framework.no_grad():
|
|
apply_func(parameter_list, group, bucket_size, scale)
|
|
|
|
|
|
def fused_allreduce_gradients(parameter_list, hcg):
|
|
group = None
|
|
scale = None
|
|
if hcg is not None:
|
|
dp_enabled = hcg.get_data_parallel_world_size() > 1
|
|
sep_enabled = hcg.get_sep_parallel_world_size() > 1
|
|
assert dp_enabled or sep_enabled, (
|
|
f"dp_enabled {dp_enabled}; sep_enabled {sep_enabled}"
|
|
)
|
|
group = None
|
|
# sep all reduce is not scaled
|
|
scale = 1.0
|
|
if dp_enabled:
|
|
group = hcg.get_data_parallel_group()
|
|
scale = scale / group.nranks
|
|
if sep_enabled:
|
|
sep_group = hcg.get_sep_parallel_group()
|
|
dp_sep_group = hcg.get_dp_sep_parallel_group()
|
|
group = sep_group if group is None else dp_sep_group
|
|
|
|
logger.debug("dp or sep start fuse allreduce gradients")
|
|
from paddle.distributed import in_auto_parallel_align_mode
|
|
|
|
if in_auto_parallel_align_mode():
|
|
scale = 1.0
|
|
fused_allreduce_gradients_with_group(parameter_list, group, scale=scale)
|
|
|
|
|
|
def broadcast_sharding_parameters(model, hcg, fuse_params=True):
|
|
# TODO TO save memory, use un-fused broadcast to avoid potential OOM
|
|
logger.debug("sharding start init parameters sync")
|
|
sharding_parallel_group = hcg.get_sharding_parallel_group()
|
|
src_rank = hcg.get_sharding_parallel_group_src_rank()
|
|
sync_params_buffers(
|
|
model,
|
|
sharding_parallel_group,
|
|
src_rank,
|
|
is_model_parallel=False,
|
|
fuse_params=fuse_params,
|
|
)
|
|
|
|
|
|
def broadcast_sep_parameters(model, hcg, fuse_params=True):
|
|
# TODO TO save memory, use un-fused broadcast to avoid potential OOM
|
|
logger.debug("sep start init parameters sync")
|
|
sep_group = hcg.get_sep_parallel_group()
|
|
src_rank = hcg.get_sep_parallel_group_src_rank()
|
|
sync_params_buffers(
|
|
model,
|
|
sep_group,
|
|
src_rank,
|
|
is_model_parallel=False,
|
|
fuse_params=fuse_params,
|
|
)
|
|
|
|
|
|
def broadcast_moe_sharding_parameters(model, hcg, fuse_params=True):
|
|
# TODO TO save memory, use un-fused broadcast to avoid potential OOM
|
|
logger.debug("moe sharding start init parameters sync")
|
|
moe_sharding_group = hcg.get_moe_sharding_parallel_group()
|
|
src_rank = hcg.get_moe_sharding_parallel_group_src_rank()
|
|
sync_params_buffers(
|
|
model,
|
|
moe_sharding_group,
|
|
src_rank,
|
|
is_model_parallel=False,
|
|
fuse_params=fuse_params,
|
|
is_moe_sharding_parallel=True,
|
|
)
|
|
|
|
|
|
def unwrap_optimizer(optimizer, optimizer_instances=()):
|
|
_inner_opt = optimizer
|
|
while isinstance(_inner_opt, optimizer_instances):
|
|
_inner_opt = _inner_opt._inner_opt
|
|
return _inner_opt
|