chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,730 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from types import MethodType
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle import nn
|
||||
from paddle.distributed import collective, fleet
|
||||
from paddle.framework import core
|
||||
from paddle.nn import ClipGradByGlobalNorm
|
||||
|
||||
from .group_sharded_stage3 import (
|
||||
ForwardPostHooks,
|
||||
ForwardPreHooks,
|
||||
OrderedSet,
|
||||
TaskFlow,
|
||||
_current_layer_params,
|
||||
_PartitionParam,
|
||||
_TensorWrapper,
|
||||
_UnsliceParam,
|
||||
align,
|
||||
alignment,
|
||||
)
|
||||
from .group_sharded_storage import GradStorage
|
||||
from .group_sharded_utils import GroupShardedClipGrad, Type, device_guard
|
||||
|
||||
|
||||
def _OptimizerWrapper(optimizer, offload, group, update_params_slice):
|
||||
if not hasattr(optimizer, "_optim"):
|
||||
optimizer._optim = optimizer
|
||||
optimizer.offload = offload
|
||||
optimizer._group = group
|
||||
optimizer.update_scaler = None
|
||||
optimizer.update_slice = update_params_slice
|
||||
return optimizer
|
||||
|
||||
|
||||
class FullyShardOptimizer:
|
||||
def __init__(
|
||||
self,
|
||||
optimizer,
|
||||
group=None,
|
||||
sync_buffers=False,
|
||||
device="xpu" if core.is_compiled_with_xpu() else "gpu",
|
||||
segment_size=2**20,
|
||||
pretrain_sync_models=True,
|
||||
offload=False,
|
||||
sync_comm=False,
|
||||
dp_group=None,
|
||||
exclude_layer=None,
|
||||
):
|
||||
self._default_device = device
|
||||
self.__sync_buffers = sync_buffers
|
||||
self._offload = offload
|
||||
self._sync_comm = sync_comm
|
||||
|
||||
# stage3 support some layer set by users to be unslice
|
||||
# _exclude_layer=[layer_name or id(layer)]
|
||||
self._exclude_layer = [] if exclude_layer is None else exclude_layer
|
||||
assert isinstance(self._exclude_layer, (list, tuple)), (
|
||||
"the exclude_layers must be a list with layers' name or layers' id"
|
||||
)
|
||||
|
||||
# segmentation size
|
||||
assert segment_size >= 0, "segment_size must be GE than 0."
|
||||
self._segment_size = segment_size
|
||||
|
||||
global param2dtype
|
||||
param2dtype = {}
|
||||
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
group = hcg.get_sharding_parallel_group()
|
||||
# Communication group establishment
|
||||
self._group = (
|
||||
collective.new_group(collective._get_global_group().ranks)
|
||||
if group is None
|
||||
else group
|
||||
)
|
||||
self._dp_group = dp_group
|
||||
self._world_size_scaling = 1.0 / self._group.nranks
|
||||
assert self._group.nranks > 1, (
|
||||
"Training must be distributed, ranks must be greater than 1."
|
||||
)
|
||||
self._rank = self._group.rank
|
||||
self._global_root_rank = self._group.ranks[
|
||||
0
|
||||
] # picking ranks index 0 as the reference
|
||||
|
||||
# Parameter segmentation for global ranks
|
||||
self._unslice_params = OrderedSet() # param's numel <= segment_size
|
||||
self._unslice_params2align = {} # {param.name: param's align}
|
||||
self._grad_storages = {} # {param.dtype: GradStorage}
|
||||
|
||||
assert not isinstance(optimizer, list), (
|
||||
"Multiple optimizers are not supported now."
|
||||
)
|
||||
self._optim = _OptimizerWrapper(
|
||||
optimizer,
|
||||
self._offload,
|
||||
self._group,
|
||||
self._update_params_slice,
|
||||
)
|
||||
self._ori_parameter_list = self._optim._parameter_list
|
||||
self._ori_param_groups = self._optim._param_groups
|
||||
|
||||
for p in self._ori_parameter_list:
|
||||
del p._need_shard
|
||||
if p._numel() > self._segment_size:
|
||||
pass
|
||||
elif p.trainable:
|
||||
self._unslice_params.add(_UnsliceParam(p))
|
||||
|
||||
# check main_grad
|
||||
self._check_main_grad()
|
||||
|
||||
# Replace optimizer's _grad_clip
|
||||
if isinstance(self._optim._grad_clip, ClipGradByGlobalNorm):
|
||||
logging.warning(
|
||||
"While using ClipGradByGlobalNorm in GroupShardedStage3, the grad clip of original optimizer will be changed."
|
||||
)
|
||||
if self.use_main_grad:
|
||||
self._optim._inner_opt._grad_clip = GroupShardedClipGrad(
|
||||
self._optim._inner_opt._grad_clip,
|
||||
paddle.get_device(),
|
||||
self._group,
|
||||
)
|
||||
else:
|
||||
self._optim._grad_clip = GroupShardedClipGrad(
|
||||
self._optim._grad_clip, paddle.get_device(), self._group
|
||||
)
|
||||
if self._optim._parameter_list and isinstance(
|
||||
self._optim._parameter_list[0], dict
|
||||
):
|
||||
for item in self._optim._param_groups:
|
||||
if "grad_clip" in item.keys():
|
||||
item["grad_clip"] = self._optim._grad_clip
|
||||
|
||||
# Add unslice params to master_weight in fp16
|
||||
self._setup_master_weights_for_unslice()
|
||||
|
||||
# Redefine optimizer step and clear function
|
||||
self._redefine_opt_step()
|
||||
self._redefine_opt_clear()
|
||||
|
||||
def _check_main_grad(self):
|
||||
self.use_main_grad = None
|
||||
for param in self._ori_parameter_list:
|
||||
if self.use_main_grad is None and hasattr(param, "main_grad"):
|
||||
self.use_main_grad = True
|
||||
if self.use_main_grad:
|
||||
assert hasattr(param, "main_grad"), (
|
||||
"Params have different main grad attributes."
|
||||
)
|
||||
|
||||
def _setup_master_weights_for_unslice(self):
|
||||
for param in self._unslice_params:
|
||||
# Update optimizer master weights
|
||||
if (
|
||||
param.dtype == Type.fp16.value or param.dtype == Type.bf16.value
|
||||
) and not self._offload:
|
||||
master_tensor = paddle.cast(param, Type.fp32.value)
|
||||
master_tensor.name = param.name
|
||||
self._optim._master_weights[param.name] = master_tensor
|
||||
|
||||
def _clear_gradients(self):
|
||||
current_layer_params = self._ori_parameter_list
|
||||
# 1.Handle param's slice
|
||||
trainable_params = list(
|
||||
filter(
|
||||
lambda p: p.trainable and p not in self._unslice_params,
|
||||
current_layer_params,
|
||||
)
|
||||
)
|
||||
for param in trainable_params:
|
||||
if not hasattr(param, "fw_storage"):
|
||||
continue
|
||||
assert hasattr(param, "fw_storage"), (
|
||||
f"Find {param.name} don't have fw_storage attribute."
|
||||
)
|
||||
if self.use_main_grad:
|
||||
param.fw_storage.main_grad._clear()
|
||||
param.fw_storage.main_grad = None
|
||||
else:
|
||||
param.fw_storage.clear_gradient(False)
|
||||
param.bw_storage._clear()
|
||||
param.bw_storage = None
|
||||
|
||||
# Update param memory slice
|
||||
def _update_params_slice(self):
|
||||
update_list = self._update_params()
|
||||
|
||||
if not isinstance(self._optim._param_groups[0], dict):
|
||||
slice_params = [param.fw_storage for param in update_list]
|
||||
self._optim._parameter_list = slice_params + list(
|
||||
self._unslice_params
|
||||
)
|
||||
self._optim._param_groups = slice_params + list(
|
||||
self._unslice_params
|
||||
)
|
||||
else:
|
||||
for param_group in self._optim._param_groups:
|
||||
p_group = []
|
||||
for p in param_group['params']:
|
||||
if hasattr(p, "fw_storage"):
|
||||
p_group.append(p.fw_storage)
|
||||
else:
|
||||
p_group.append(p)
|
||||
|
||||
param_group['params'] = p_group
|
||||
|
||||
def _update_params(self):
|
||||
"""
|
||||
Update parameters to optimizer memory slice.
|
||||
"""
|
||||
update_list = []
|
||||
current_layer_params = self._ori_parameter_list
|
||||
trainable_params = list(
|
||||
filter(
|
||||
lambda p: p.trainable and p not in self._unslice_params,
|
||||
current_layer_params,
|
||||
)
|
||||
)
|
||||
# 1.Handle param's slice
|
||||
for param in trainable_params:
|
||||
assert hasattr(param, "fw_storage"), (
|
||||
f"Find {param.name} don't have fw_storage attribute"
|
||||
)
|
||||
|
||||
param.fw_storage = _TensorWrapper(param)
|
||||
if self.use_main_grad:
|
||||
param.fw_storage.main_grad = param.bw_storage
|
||||
else:
|
||||
assert param.fw_storage.grad is None
|
||||
param.fw_storage._copy_gradient_from(param.bw_storage)
|
||||
update_list.append(param)
|
||||
|
||||
return update_list
|
||||
|
||||
def _redefine_opt_step(self):
|
||||
params_slice_func = self._update_params_slice
|
||||
opt_step = self._optim.step
|
||||
|
||||
def _opt_step(self):
|
||||
if not self.update_scaler:
|
||||
params_slice_func()
|
||||
opt_step()
|
||||
|
||||
self._optim.step = MethodType(_opt_step, self._optim)
|
||||
|
||||
def _redefine_opt_clear(self):
|
||||
clear_func = self._clear_gradients
|
||||
|
||||
def _opt_clear(self):
|
||||
clear_func()
|
||||
|
||||
self._optim.clear_grad = MethodType(_opt_clear, self._optim)
|
||||
|
||||
|
||||
class FullyShard(nn.Layer):
|
||||
"""
|
||||
A wrapper for Sharding Stage3 Layer in Dygraph.
|
||||
|
||||
.. warning: GroupShardedStage3 encapsulates the layer strategy and integrates it into the nn.Layer.
|
||||
|
||||
.. ZeRO: https://arxiv.org/pdf/1910.02054.pdf.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
layer,
|
||||
group=None,
|
||||
sync_buffers=False,
|
||||
device="xpu" if core.is_compiled_with_xpu() else "gpu",
|
||||
segment_size=2**20,
|
||||
pretrain_sync_models=True,
|
||||
offload=False,
|
||||
sync_comm=False,
|
||||
dp_group=None,
|
||||
exclude_layer=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Default configs
|
||||
assert (
|
||||
core.is_compiled_with_cuda()
|
||||
or core.is_compiled_with_xpu()
|
||||
or (device in core.get_all_custom_device_type())
|
||||
), "Only support CUDA / XPU / CustomDevice."
|
||||
|
||||
self._layer = layer
|
||||
self._default_device = device
|
||||
self.__sync_buffers = sync_buffers
|
||||
self._offload = offload
|
||||
self._sync_comm = sync_comm
|
||||
|
||||
# stage3 support some layer set by users to be unslice
|
||||
self._exclude_layer = [] if exclude_layer is None else exclude_layer
|
||||
assert isinstance(self._exclude_layer, (list, tuple)), (
|
||||
"the exclude_layers must be a list with layers' name or layers' id"
|
||||
)
|
||||
|
||||
# segmentation size
|
||||
assert segment_size >= 0, "segment_size must be GE than 0."
|
||||
self._segment_size = segment_size
|
||||
|
||||
global DEV
|
||||
DEV = (
|
||||
"cpu"
|
||||
if paddle.get_device() == "cpu"
|
||||
else paddle.get_device().split(":")[0]
|
||||
)
|
||||
global DEV_ID
|
||||
DEV_ID = (
|
||||
0
|
||||
if paddle.get_device() == "cpu"
|
||||
else int(paddle.get_device().split(":")[1])
|
||||
)
|
||||
global param2dtype
|
||||
param2dtype = {}
|
||||
|
||||
hcg = fleet.get_hybrid_communicate_group()
|
||||
group = hcg.get_sharding_parallel_group()
|
||||
|
||||
self._group = (
|
||||
collective.new_group(collective._get_global_group().ranks)
|
||||
if group is None
|
||||
else group
|
||||
)
|
||||
self._dp_group = dp_group
|
||||
self._world_size_scaling = 1.0 / self._group.nranks
|
||||
assert self._group.nranks > 1, (
|
||||
"Training must be distributed, ranks must be greater than 1."
|
||||
)
|
||||
self._rank = self._group.rank
|
||||
self._global_root_rank = self._group.ranks[
|
||||
0
|
||||
] # picking ranks index 0 as the reference
|
||||
|
||||
# Parameter segmentation for global ranks
|
||||
# After flatten -> self._param2buffer_size, self._param2buffer, self._trainable_params
|
||||
self._param2buffer_size = {} # {param.name: size}
|
||||
self._param2buffer = {} # {param.name: [(start0, end0),(start1, end1), ...]}
|
||||
self._trainable_params = {} # {id(layer): [trainable_params]}
|
||||
self._unslice_params = OrderedSet() # param's numel <= segment_size
|
||||
self._unslice_params2align = {} # {param.name: param's align}
|
||||
self._grad_storages = {} # {param.dtype: GradStorage}
|
||||
|
||||
self._ori_parameter_list = self._layer.parameters()
|
||||
for param in self._ori_parameter_list:
|
||||
param._need_shard = True
|
||||
# check main_grad
|
||||
self._check_main_grad()
|
||||
|
||||
# Synchronous all ranks models
|
||||
if pretrain_sync_models:
|
||||
self._sync_params_and_buffers()
|
||||
|
||||
self._segment_rank_params(self._layer)
|
||||
|
||||
# In the first step, record the execution order of the layer
|
||||
self._order_tracer = OrderedDict()
|
||||
self._order_tracer["order"] = 0
|
||||
self._order_tracer["layer"] = []
|
||||
|
||||
# Add unslice params GradStorage
|
||||
self._handle_unslice_params()
|
||||
|
||||
# Register task flow
|
||||
self._task_flow = TaskFlow()
|
||||
|
||||
# Register forward hooks
|
||||
self._register_forward_hooks(self._layer)
|
||||
|
||||
# Register backward parameter hooks
|
||||
self._register_backward_hooks()
|
||||
|
||||
def _handle_unslice_params(self):
|
||||
buffer_size = {}
|
||||
buffer_size[Type.bf16.value] = 0
|
||||
buffer_size[Type.fp32.value] = 0
|
||||
buffer_size[Type.fp16.value] = 0
|
||||
for param in self._unslice_params:
|
||||
param2dtype[param.name] = param.dtype
|
||||
p_align = self._param2align(param)
|
||||
self._unslice_params2align[param.name] = p_align
|
||||
buffer_size[param.dtype] += param._numel() + p_align
|
||||
|
||||
# Create unslice_params'grad
|
||||
for param in sorted(self._unslice_params, key=lambda p: p.name):
|
||||
if param.dtype not in self._grad_storages.keys():
|
||||
self._grad_storages[param.dtype] = GradStorage(
|
||||
buffer_size[param.dtype],
|
||||
dtype=(
|
||||
param.dtype
|
||||
if not self.use_main_grad
|
||||
else paddle.float32
|
||||
),
|
||||
device=self._default_device,
|
||||
destination=self._rank,
|
||||
param2align=self._unslice_params2align,
|
||||
)
|
||||
self._grad_storages[param.dtype].add_grad(
|
||||
param, self._unslice_params2align[param.name]
|
||||
)
|
||||
|
||||
def _check_main_grad(self):
|
||||
self.use_main_grad = None
|
||||
for param in self._layer.parameters():
|
||||
if self.use_main_grad is None and hasattr(param, "main_grad"):
|
||||
self.use_main_grad = True
|
||||
if self.use_main_grad:
|
||||
assert hasattr(param, "main_grad"), (
|
||||
"Params have different main grad attributes."
|
||||
)
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _sync_params_and_buffers(self):
|
||||
"""
|
||||
Sync all model states for all ranks
|
||||
"""
|
||||
|
||||
for p in self._layer.parameters():
|
||||
dist.broadcast(
|
||||
p, src=self._global_root_rank, group=self._group, sync_op=True
|
||||
)
|
||||
if self._dp_group is not None and self._dp_group.nranks > 1:
|
||||
dist.broadcast(
|
||||
p,
|
||||
src=self._dp_group.ranks[0],
|
||||
group=self._dp_group,
|
||||
sync_op=True,
|
||||
)
|
||||
|
||||
def _sync_grad_storages_hook(self):
|
||||
for grad_storage in self._grad_storages.values():
|
||||
grad_storage.buffer.scale_(scale=self._world_size_scaling)
|
||||
dist.all_reduce(tensor=grad_storage.buffer, group=self._group)
|
||||
if self._dp_group is not None and self._dp_group.nranks > 1:
|
||||
grad_storage.buffer.scale_(scale=(1.0 / self._dp_group.nranks))
|
||||
dist.all_reduce(
|
||||
tensor=grad_storage.buffer, group=self._dp_group
|
||||
)
|
||||
|
||||
def forward(self, *inputs, **kwargs):
|
||||
"""
|
||||
A wrapper for Sharding Stage3 layer.
|
||||
"""
|
||||
# add hook to sync grad storage
|
||||
for grad_storage in self._grad_storages.values():
|
||||
grad_storage.buffer.zero_()
|
||||
grad_storage.manual_release()
|
||||
grad_storage.rebuild()
|
||||
core.eager._add_backward_final_hook(self._sync_grad_storages_hook)
|
||||
|
||||
# 1.Sync layer's buffers state
|
||||
if self.__sync_buffers:
|
||||
self._sync_buffers()
|
||||
|
||||
# 2.Normal FW on the base model
|
||||
fw = self._layer(*inputs, **kwargs)
|
||||
|
||||
return fw
|
||||
|
||||
def _segment_rank_params(self, layer, name="last_layer"):
|
||||
"""
|
||||
Flatten parameters according to layer.
|
||||
"""
|
||||
current_layer_params = _current_layer_params(layer)
|
||||
if current_layer_params:
|
||||
self._flatten_layer_params(layer, current_layer_params)
|
||||
|
||||
for name, sub_layer in layer.named_children():
|
||||
self._segment_rank_params(sub_layer, name)
|
||||
|
||||
def _flatten_layer_params(self, layer, current_layer_params):
|
||||
"""
|
||||
Parameter segmentation and memory integration.
|
||||
"""
|
||||
|
||||
if id(layer) in self._trainable_params.keys():
|
||||
return
|
||||
|
||||
def _add_manage_info(trainable_param):
|
||||
return _PartitionParam(trainable_param)
|
||||
|
||||
current_params = []
|
||||
for p in current_layer_params:
|
||||
if p._numel() > self._segment_size:
|
||||
current_params.append(_add_manage_info(p))
|
||||
elif p.trainable:
|
||||
self._unslice_params.add(_UnsliceParam(p))
|
||||
|
||||
self._trainable_params[id(layer)] = current_params
|
||||
|
||||
for param in self._trainable_params[id(layer)]:
|
||||
if param.name in self._param2buffer.keys():
|
||||
continue
|
||||
self._param2buffer[param.name] = []
|
||||
# 1.Params alignment
|
||||
align_ = self._param2align(param)
|
||||
|
||||
offset = align_ + param._numel()
|
||||
buffer_size = (
|
||||
offset
|
||||
if offset % self._group.nranks == 0
|
||||
else offset + self._group.nranks - (offset % self._group.nranks)
|
||||
)
|
||||
self._param2buffer_size[param.name] = buffer_size
|
||||
|
||||
# 2.Combination param buffer
|
||||
assert buffer_size % self._group.nranks == 0
|
||||
pre_buffer = buffer_size // self._group.nranks
|
||||
|
||||
for rank_ in range(self._group.nranks):
|
||||
self._param2buffer[param.name].append(
|
||||
(rank_ * pre_buffer, (rank_ + 1) * pre_buffer)
|
||||
)
|
||||
|
||||
# Record param's dtype
|
||||
param2dtype[param.name] = param.dtype
|
||||
# 3.Flatten layer params and release other rank buffer
|
||||
self._param_storage(param, buffer_size)
|
||||
|
||||
def _param_storage(self, param, buffer_size):
|
||||
"""
|
||||
This is a function to simplify the handling of parameter InternalStorages.
|
||||
"""
|
||||
assert isinstance(buffer_size, int)
|
||||
value = (
|
||||
np.zeros(buffer_size, dtype=np.float16)
|
||||
if (
|
||||
Type.fp16.value == param.dtype or Type.bf16.value == param.dtype
|
||||
)
|
||||
else np.zeros(buffer_size, dtype=np.float32)
|
||||
)
|
||||
buffer = core.eager.Tensor(value=value, place=core.CPUPlace())
|
||||
if Type.bf16.value == param.dtype:
|
||||
buffer = buffer.cast(Type.bf16.value)
|
||||
|
||||
param_shape = param.shape
|
||||
origin_state = param.stop_gradient
|
||||
param.stop_gradient = True
|
||||
param.flatten_()
|
||||
param.stop_gradient = origin_state
|
||||
start, end = self._param2buffer[param.name][self._rank]
|
||||
|
||||
# Copy the current param value
|
||||
with device_guard():
|
||||
tmp_var = buffer._slice(0, param._numel())
|
||||
param_cpu = param.cpu()
|
||||
tmp_var.get_tensor().set(param_cpu.get_tensor(), core.CPUPlace())
|
||||
del tmp_var
|
||||
param.get_tensor()._set_dims(param_shape)
|
||||
|
||||
# Current rank param_storage
|
||||
param.fw_storage = core.eager.Tensor(
|
||||
value=buffer._slice(start, end), name="slice@" + param.name
|
||||
)
|
||||
param.status = "part"
|
||||
param._clear_data()
|
||||
|
||||
def _register_forward_hooks(self, layer):
|
||||
"""
|
||||
Register PyLayer to manage memory slices.
|
||||
There are four stages:
|
||||
FW
|
||||
1. Before the forward layers, synchronize the full parameters.
|
||||
2. After the forward layers, release the full parameter and keep the parameter slice.
|
||||
BW
|
||||
3. Before the backward layers, synchronize the full parameters and create param's grad.
|
||||
4. After the gradient accumulation, release the full parameter and keep the parameter slice.
|
||||
"""
|
||||
current_layer_params = _current_layer_params(layer)
|
||||
if current_layer_params:
|
||||
# the layer in self._exclude_layer will be added hooks.
|
||||
if not (
|
||||
id(layer) in self._exclude_layer
|
||||
or layer.__class__.__name__ in self._exclude_layer
|
||||
):
|
||||
self._register_forward_all_hooks(layer, self._task_flow)
|
||||
|
||||
for _, sub_layer in layer.named_children():
|
||||
self._register_forward_hooks(sub_layer)
|
||||
|
||||
def _register_forward_all_hooks(self, sub_layer, task_flow):
|
||||
def _forward_pre_hook(layer, inputs):
|
||||
return ForwardPreHooks(
|
||||
layer,
|
||||
self._order_tracer,
|
||||
self._trainable_params,
|
||||
self._param2buffer_size,
|
||||
self._group,
|
||||
self._sync_comm,
|
||||
self._offload,
|
||||
task_flow,
|
||||
)
|
||||
|
||||
def _forward_post_hook(layer, inputs, outputs):
|
||||
if isinstance(outputs, paddle.Tensor):
|
||||
outputs = (outputs,)
|
||||
return ForwardPostHooks.apply(
|
||||
*outputs,
|
||||
layer=layer,
|
||||
order_tracer=self._order_tracer,
|
||||
trainable_params=self._trainable_params,
|
||||
param2buffer=self._param2buffer,
|
||||
param2buffer_size=self._param2buffer_size,
|
||||
rank=self._rank,
|
||||
group=self._group,
|
||||
sync_comm=self._sync_comm,
|
||||
offload=self._offload,
|
||||
task_flow=task_flow,
|
||||
)
|
||||
|
||||
# register previous forward hooks
|
||||
sub_layer.register_forward_pre_hook(_forward_pre_hook)
|
||||
|
||||
# register post forward hooks
|
||||
sub_layer.register_forward_post_hook(_forward_post_hook)
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _sync_buffers(self):
|
||||
"""
|
||||
Sync all the param buffers from all ranks (exp: batch norm statistics).
|
||||
"""
|
||||
|
||||
for buffer in self._layer.buffers(include_sublayers=True):
|
||||
dist.broadcast(
|
||||
buffer, self._global_root_rank, self._group, sync_op=True
|
||||
)
|
||||
if self._dp_group is not None and self._dp_group.nranks > 1:
|
||||
dist.broadcast(
|
||||
buffer,
|
||||
self._dp_group.ranks[0],
|
||||
self._dp_group,
|
||||
sync_op=True,
|
||||
)
|
||||
|
||||
def __getattr__(self, name):
|
||||
"""Forward missing attributes to wrapped layer."""
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
return getattr(self._layer, name)
|
||||
|
||||
def _register_backward_hooks(self):
|
||||
current_layer_params = self._layer.parameters(include_sublayers=True)
|
||||
trainable_params = list(
|
||||
filter(
|
||||
lambda p: p.trainable and p not in self._unslice_params,
|
||||
current_layer_params,
|
||||
)
|
||||
)
|
||||
|
||||
for param in trainable_params:
|
||||
allreduce_function = self._get_allreduce_fn(param)
|
||||
param._register_backward_hook(allreduce_function)
|
||||
|
||||
def _get_allreduce_fn(self, param):
|
||||
@paddle.autograd.no_grad()
|
||||
def allreduce_(*_):
|
||||
assert param.trainable, (
|
||||
"the param must be trainable for grad allreduced"
|
||||
)
|
||||
if param.name in self._task_flow.full_grad.keys():
|
||||
full_grad = self._task_flow.full_grad[param.name]
|
||||
# Only support sync allreduce current rank's layer now
|
||||
full_grad.scale_(scale=self._world_size_scaling)
|
||||
dist.all_reduce(tensor=full_grad, group=self._group)
|
||||
if self._dp_group is not None and self._dp_group.nranks > 1:
|
||||
full_grad.scale_(scale=1.0 / self._dp_group.nranks)
|
||||
dist.all_reduce(tensor=full_grad, group=self._dp_group)
|
||||
|
||||
start, end = self._param2buffer[param.name][self._rank]
|
||||
if param.bw_storage is None:
|
||||
param.bw_storage = (
|
||||
full_grad._slice(start, end).detach().clone()
|
||||
)
|
||||
else:
|
||||
param.bw_storage = paddle.add(
|
||||
param.bw_storage,
|
||||
full_grad._slice(start, end).detach().clone(),
|
||||
)
|
||||
|
||||
if self.use_main_grad:
|
||||
param.main_grad = None
|
||||
else:
|
||||
param.clear_gradient(False)
|
||||
del self._task_flow.full_grad[param.name]
|
||||
|
||||
if param.name in self._task_flow.full_param.keys():
|
||||
if param.status == "all":
|
||||
param.use_count = 0
|
||||
param._clear_data()
|
||||
start, end = self._param2buffer[param.name][self._rank]
|
||||
param.fw_storage = (
|
||||
self._task_flow.full_param[param.name][0]
|
||||
._slice(start, end)
|
||||
.detach()
|
||||
.clone()
|
||||
)
|
||||
param.status = "part"
|
||||
del self._task_flow.full_param[param.name]
|
||||
|
||||
return allreduce_
|
||||
|
||||
def _param2align(self, param):
|
||||
# CUDA alignment 256 bytes
|
||||
size = param._numel() * align[param.dtype]
|
||||
device_alignment = alignment[self._default_device]
|
||||
remaining = size % device_alignment
|
||||
ali = 0 if remaining == 0 else device_alignment - remaining
|
||||
align_ = ali // align[param.dtype]
|
||||
return align_
|
||||
+792
@@ -0,0 +1,792 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
# The file has been adapted from fairscale file:
|
||||
# https://github.com/facebookresearch/fairscale/blob/main/fairscale/optim/oss.py
|
||||
# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
|
||||
# We retain the following license from the original files:
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
import warnings
|
||||
from collections import OrderedDict
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed import ParallelMode, fleet
|
||||
from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
|
||||
ShardedStateDict,
|
||||
ShardedWeight,
|
||||
create_sharded_weight_with_new_local,
|
||||
)
|
||||
from paddle.framework import core
|
||||
from paddle.nn import ClipGradByGlobalNorm
|
||||
from paddle.optimizer import Optimizer
|
||||
|
||||
HybridParallelClipGrad = fleet.meta_optimizers.dygraph_optimizer.hybrid_parallel_optimizer.HybridParallelClipGrad
|
||||
from paddle.distributed.collective import _get_global_group, new_group
|
||||
|
||||
from .group_sharded_storage import GradStorage, ParamStorage
|
||||
from .group_sharded_utils import GroupShardedClipGrad, Type, device_guard
|
||||
|
||||
# CUDA alignment 256 bytes, cpu alignment 4096 bytes
|
||||
alignment = {"gpu": 256, "cpu": 4096, "xpu": 256}
|
||||
align = {
|
||||
Type.fp16.value: 2,
|
||||
Type.bf16.value: 2,
|
||||
Type.fp32.value: 4,
|
||||
}
|
||||
|
||||
|
||||
class GroupShardedOptimizerStage2(Optimizer):
|
||||
"""
|
||||
A wrapper for Sharding Stage2 Optimizer in Dygraph.
|
||||
|
||||
.. warning: ShardingOptimizer encapsulates the optimization strategy and integrates it into the optimizer.
|
||||
|
||||
.. ZeRO: 1.https://arxiv.org/pdf/1910.02054.pdf 2.https://arxiv.org/pdf/1910.02054.pdf.
|
||||
|
||||
"""
|
||||
|
||||
# TODO (Baibaifan)
|
||||
# Feature Notes:
|
||||
# 1. Unified memory for parameters and parameters.grad to InternalStorage.
|
||||
# 2. Support the segmentation of optimizer parameters and partial updating of parameters.
|
||||
# 3. Dynamically adjust training parameters and models.
|
||||
# 4. Support offload function.
|
||||
# 5. Support the establishment of independent communication groups.
|
||||
# 6. Broadcast_fp16 is not supported now.
|
||||
def __init__(
|
||||
self,
|
||||
params,
|
||||
optim,
|
||||
group=None,
|
||||
offload=False,
|
||||
device="xpu" if core.is_compiled_with_xpu() else "gpu",
|
||||
pretrain_sync_models=True,
|
||||
dp_group=None,
|
||||
**kw,
|
||||
):
|
||||
super().__init__(learning_rate=optim._learning_rate, parameters=params)
|
||||
assert (
|
||||
core.is_compiled_with_cuda()
|
||||
or core.is_compiled_with_xpu()
|
||||
or (device in core.get_all_custom_device_type())
|
||||
), "Only GPU and XPU and CustomDevice is supported now"
|
||||
|
||||
# Segmentation information
|
||||
self._dtype_rank_params = (
|
||||
OrderedDict()
|
||||
) # {dtype:[param1,param2]} device, rank, params
|
||||
self._param2rank = {}
|
||||
self.__segment_params = []
|
||||
self._rank_buffer_size = {} # {dtype: {rank: numel+alignment}}
|
||||
self._param2align = {} # {param.name: align}
|
||||
|
||||
# Default information
|
||||
self._optim = optim
|
||||
|
||||
# sharing stage 2 comm overlap flag
|
||||
self._reduce_overlap = False
|
||||
# record the last task used for comm overlap for sharding stage 2
|
||||
self._comm_task = None
|
||||
|
||||
assert hasattr(self._optim, "_master_weights"), (
|
||||
"Must use optimizer with _master_weights attribute"
|
||||
)
|
||||
|
||||
# Support parameter group and parameter list
|
||||
self._local_params = []
|
||||
if isinstance(params[0], dict):
|
||||
for param_group in params:
|
||||
self._local_params.extend(list(param_group["params"]))
|
||||
else:
|
||||
self._local_params.extend(list(params))
|
||||
|
||||
self.use_main_grad = None
|
||||
for param in self._local_params:
|
||||
if self.use_main_grad is None and hasattr(param, "main_grad"):
|
||||
self.use_main_grad = True
|
||||
if self.use_main_grad:
|
||||
assert hasattr(param, "main_grad"), (
|
||||
"Params have different main grad attributes."
|
||||
)
|
||||
if self.use_main_grad:
|
||||
assert not offload, "offload not support main_grad for now"
|
||||
|
||||
self._default_device = device
|
||||
self._pfp16 = (
|
||||
len(
|
||||
list(
|
||||
filter(
|
||||
lambda x: x.trainable and x.dtype == Type.fp16.value,
|
||||
self._local_params,
|
||||
)
|
||||
)
|
||||
)
|
||||
> 0
|
||||
)
|
||||
self._pbf16 = (
|
||||
len(
|
||||
list(
|
||||
filter(
|
||||
lambda x: x.trainable and x.dtype == Type.bf16.value,
|
||||
self._local_params,
|
||||
)
|
||||
)
|
||||
)
|
||||
> 0
|
||||
)
|
||||
|
||||
self._broadcast_overlap = False
|
||||
self._forward_pre_hook_remove_helper = []
|
||||
try:
|
||||
# The fp32 params such as layer_norm_0.w_0 will be at the end of param_list.
|
||||
# Have to sort the params to make sure all params are in the forward using order.
|
||||
self._broadcast_order_params = sorted(
|
||||
self.local_params,
|
||||
key=lambda x: int(x.name.split('.')[0].split('_')[-1]),
|
||||
)
|
||||
except ValueError:
|
||||
self._broadcast_order_params = None
|
||||
|
||||
self._group = (
|
||||
new_group(_get_global_group().ranks) if group is None else group
|
||||
)
|
||||
|
||||
# only support to combine stage2 and dp hybrid parallel now.
|
||||
self._dp_group = dp_group
|
||||
self.world_size = self._group.nranks
|
||||
self._rank = self._group.rank
|
||||
self._global_root_rank = self._group.ranks[0]
|
||||
|
||||
if self._dp_group is not None and self._dp_group.nranks > 1:
|
||||
assert not offload, (
|
||||
"Not support! when using offload with sharding stage2, please use pure sharding stage2, exclude data parallel."
|
||||
)
|
||||
|
||||
# Synchronous all ranks models
|
||||
if pretrain_sync_models:
|
||||
self._sync_params_and_buffers()
|
||||
|
||||
self.param_storages = {} # {dtype: {rank: InternalStorage}}
|
||||
|
||||
if isinstance(self._optim._grad_clip, ClipGradByGlobalNorm):
|
||||
logging.warning(
|
||||
"While using ClipGradByGlobalNorm in GroupShardedOptimizerStage2, the grad clip of original optimizer will be changed."
|
||||
)
|
||||
|
||||
hcg = fleet.fleet._hcg if hasattr(fleet.fleet, "_hcg") else None
|
||||
if (
|
||||
hcg
|
||||
and hcg.get_parallel_mode() is not ParallelMode.DATA_PARALLEL
|
||||
and not offload
|
||||
):
|
||||
if self.use_main_grad:
|
||||
self._optim._inner_opt._grad_clip = HybridParallelClipGrad(
|
||||
self._optim._inner_opt._grad_clip, hcg
|
||||
)
|
||||
else:
|
||||
self._optim._grad_clip = HybridParallelClipGrad(
|
||||
self._optim._grad_clip, hcg
|
||||
)
|
||||
else:
|
||||
if self.use_main_grad:
|
||||
self._optim._inner_opt._grad_clip = GroupShardedClipGrad(
|
||||
self._optim._inner_opt._grad_clip,
|
||||
paddle.get_device(),
|
||||
self._group,
|
||||
)
|
||||
else:
|
||||
self._optim._grad_clip = GroupShardedClipGrad(
|
||||
self._optim._grad_clip, paddle.get_device(), self._group
|
||||
)
|
||||
|
||||
if self._optim._parameter_list and isinstance(
|
||||
self._optim._parameter_list[0], dict
|
||||
):
|
||||
for item in self._optim._param_groups:
|
||||
if "grad_clip" in item.keys():
|
||||
item["grad_clip"] = self._optim._grad_clip
|
||||
|
||||
if offload:
|
||||
assert self._pfp16, (
|
||||
"Only support offload strategy while using 'Adam', 'AdamW' and 'Momentum' optimizer with AMP/Pure FP16"
|
||||
)
|
||||
|
||||
self.offload = offload # Using for offload
|
||||
self.offload_device = "cpu"
|
||||
self.offload_buffer_size = 0
|
||||
self.offload_param2align = {}
|
||||
self.offload_params = None
|
||||
self.offload_grads = None
|
||||
self.dev_id = int(paddle.get_device().split(":")[1])
|
||||
|
||||
self._master_params = {}
|
||||
|
||||
# Update optimizer parameters and adjust parameter storage and use according to rank.
|
||||
self._update_opt_status()
|
||||
|
||||
def _set_auxiliary_var(self, key, val):
|
||||
super()._set_auxiliary_var(key, val)
|
||||
self._optim._set_auxiliary_var(key, val)
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _sync_params_and_buffers(self):
|
||||
"""
|
||||
Sync all model states for all ranks
|
||||
"""
|
||||
|
||||
for p in self._local_params:
|
||||
dist.broadcast(
|
||||
p, src=self._global_root_rank, group=self._group, sync_op=True
|
||||
)
|
||||
|
||||
if self._dp_group:
|
||||
dist.broadcast(
|
||||
p,
|
||||
src=self._dp_group.ranks[0],
|
||||
group=self._dp_group,
|
||||
sync_op=True,
|
||||
)
|
||||
|
||||
def _update_task(self, task):
|
||||
if self._reduce_overlap:
|
||||
assert task is not None
|
||||
# Only track of the last reduce task.
|
||||
# Since all tasks are on the same stream, only need to wait the last one.
|
||||
# After waiting for the last reduce task, all reduce tasks before have already finished.
|
||||
self._comm_task = task
|
||||
|
||||
def _set_reduce_overlap(self, reduce_overlap):
|
||||
# Enable gradients' reduces overlap with backward calculation.
|
||||
self._reduce_overlap = reduce_overlap
|
||||
|
||||
def _set_broadcast_overlap(
|
||||
self, broadcast_overlap, layers=None, num_groups=None
|
||||
):
|
||||
# Enable post optimizer broadcasts overlap with the forward calculation of next batch.
|
||||
self._broadcast_overlap = broadcast_overlap
|
||||
if self._broadcast_overlap:
|
||||
assert layers is not None, (
|
||||
"To enable broadcast overlap forward, please pass the module to the function."
|
||||
)
|
||||
self._layers = layers
|
||||
warnings.warn(
|
||||
"Setting overlap broadcast means the `paddle.device.cuda.synchronize()` "
|
||||
"must be called manually before calling `paddle.save()` and before and inference."
|
||||
)
|
||||
if self._broadcast_order_params is None:
|
||||
# Params' names should be like column_linear_32.w_0 pattern to get the best performance.
|
||||
warnings.warn(
|
||||
r"The param name passed to the optimizer doesn't follow .+_[0-9]+\..+ pattern, "
|
||||
"overlap broadcast may harm the performance."
|
||||
)
|
||||
self._broadcast_order_params = self._local_params
|
||||
|
||||
if num_groups is None or num_groups > len(self._broadcast_order_params):
|
||||
warnings.warn(
|
||||
"The num_groups for broadcast is larger than the number of params to be broadcast. "
|
||||
"It will set to default value: 1 (use the default sharding group)."
|
||||
)
|
||||
num_groups = 1
|
||||
|
||||
assert isinstance(num_groups, int) and num_groups > 0, (
|
||||
"num_groups should be a positive integer"
|
||||
)
|
||||
|
||||
self._number_of_broadcast_groups = num_groups
|
||||
self._broadcast_groups = [
|
||||
None for _ in range(self._number_of_broadcast_groups)
|
||||
]
|
||||
self._broadcast_groups[0] = self._group
|
||||
|
||||
ranks = self._group.ranks
|
||||
for i in range(1, self._number_of_broadcast_groups):
|
||||
self._broadcast_groups[i] = new_group(ranks)
|
||||
|
||||
def _generate_master_params(self, trainable_params):
|
||||
if self.offload:
|
||||
for param in trainable_params:
|
||||
if param.name not in self._master_params.keys():
|
||||
self._master_params[param.name] = core.eager.Tensor(
|
||||
name=param.name,
|
||||
value=param.cast(dtype=Type.fp32.value).numpy(),
|
||||
place=core.CPUPlace(),
|
||||
stop_gradient=param.stop_gradient,
|
||||
)
|
||||
else:
|
||||
for param in trainable_params:
|
||||
if (
|
||||
param.dtype == Type.fp16.value
|
||||
or param.dtype == Type.bf16.value
|
||||
):
|
||||
master_tensor = paddle.cast(param, Type.fp32.value)
|
||||
master_tensor.name = param.name
|
||||
self._optim._master_weights[param.name] = master_tensor
|
||||
|
||||
def _update_opt_status(self):
|
||||
"""Update optimizer status and parameter storage information, and special functions to be developed."""
|
||||
# func 1
|
||||
self._integration_params()
|
||||
|
||||
# Segment helpers
|
||||
|
||||
def _segment_params(self):
|
||||
"""
|
||||
Divide all optimizer parameters equally into rank.
|
||||
"""
|
||||
if len(self.__segment_params) == 0:
|
||||
self.__segment_params, param_lists = (
|
||||
[[] for _ in range(self.world_size)],
|
||||
[[] for _ in range(self.world_size)],
|
||||
)
|
||||
sizes = [0] * self.world_size
|
||||
for param in self._local_params:
|
||||
# Add this param to rank with smallest size.
|
||||
rank = sizes.index(min(sizes))
|
||||
param_lists[rank].append(param)
|
||||
|
||||
# Statistical real numels
|
||||
sizes[rank] += param._numel() if param.trainable else 0
|
||||
|
||||
for rank, params in enumerate(param_lists):
|
||||
self.__segment_params[rank].extend(params)
|
||||
return self.__segment_params
|
||||
|
||||
@property
|
||||
def local_params(self):
|
||||
return self._local_params
|
||||
|
||||
@property
|
||||
def param2rank(self):
|
||||
"""Map the params to the rank which owns them"""
|
||||
if len(self._param2rank) == 0:
|
||||
for rank, params in enumerate(self._segment_params()):
|
||||
for param in params:
|
||||
self._param2rank[param.name] = rank
|
||||
return self._param2rank
|
||||
|
||||
@property
|
||||
def dtype_rank_params(self):
|
||||
"""
|
||||
Divide the parameters into groups according to rank and dtype.
|
||||
"""
|
||||
if len(self._dtype_rank_params) == 0:
|
||||
# Assign the parameters of each rank according to the type
|
||||
trainable_params = list(
|
||||
filter(lambda x: x.trainable, self._local_params)
|
||||
)
|
||||
for param in trainable_params:
|
||||
if param.dtype not in self._dtype_rank_params.keys():
|
||||
self._dtype_rank_params[param.dtype] = [
|
||||
[] for _ in range(self.world_size)
|
||||
]
|
||||
self._dtype_rank_params[param.dtype][
|
||||
self.param2rank[param.name]
|
||||
].append(param)
|
||||
|
||||
# Sort per rank params by size
|
||||
for dtype in self._dtype_rank_params.keys():
|
||||
for rank_params in self._dtype_rank_params[dtype]:
|
||||
rank_params.sort(key=lambda x: x._numel())
|
||||
|
||||
return self._dtype_rank_params
|
||||
|
||||
@property
|
||||
def rank_buffer_size(self):
|
||||
"""
|
||||
Count the memory size of the parameters corresponding to rank under the corresponding dtype.
|
||||
"""
|
||||
# CUDA alignment 256 bytes
|
||||
if self._default_device in core.get_all_custom_device_type():
|
||||
device_alignment = core.libpaddle._get_device_min_chunk_size(
|
||||
self._default_device
|
||||
)
|
||||
else:
|
||||
device_alignment = alignment[self._default_device]
|
||||
|
||||
if len(self._rank_buffer_size) == 0:
|
||||
for dtype in self.dtype_rank_params.keys():
|
||||
if dtype not in self._rank_buffer_size.keys():
|
||||
self._rank_buffer_size[dtype] = {}
|
||||
for dst_rank, per_rank_params in enumerate(
|
||||
self.dtype_rank_params[dtype]
|
||||
):
|
||||
if dst_rank not in self._rank_buffer_size[dtype].keys():
|
||||
self._rank_buffer_size[dtype][dst_rank] = 0
|
||||
for param in per_rank_params:
|
||||
if not param.trainable:
|
||||
continue
|
||||
size = param._numel() * align[dtype]
|
||||
remaining = size % device_alignment
|
||||
ali = (
|
||||
0
|
||||
if remaining == 0
|
||||
else device_alignment - remaining
|
||||
)
|
||||
align_ = ali // align[dtype]
|
||||
self._rank_buffer_size[dtype][dst_rank] += (
|
||||
param._numel() + align_
|
||||
)
|
||||
self._param2align[param.name] = align_
|
||||
|
||||
return self._rank_buffer_size
|
||||
|
||||
def _integration_params(self):
|
||||
"""
|
||||
Integrate the parameters into a continuous memory according to rank, and support the update of training parameters.
|
||||
"""
|
||||
|
||||
for dtype, per_rank_params in self.dtype_rank_params.items():
|
||||
if dtype not in self.param_storages.keys():
|
||||
self.param_storages[dtype] = {}
|
||||
|
||||
for dst_rank, params in enumerate(per_rank_params):
|
||||
if len(params) > 0:
|
||||
# Merge all the trainable params in a single InternalStorage
|
||||
trainable_params = list(
|
||||
filter(lambda x: x.trainable, params)
|
||||
)
|
||||
if (self._pfp16 or self._pbf16) and dst_rank == self._rank:
|
||||
self._generate_master_params(trainable_params)
|
||||
if trainable_params:
|
||||
param_storage = ParamStorage(
|
||||
size=self.rank_buffer_size[dtype][dst_rank],
|
||||
dtype=dtype,
|
||||
device=self._default_device,
|
||||
)
|
||||
|
||||
param_storage.add_rank_params(
|
||||
trainable_params, self._param2align
|
||||
)
|
||||
self.param_storages[dtype][dst_rank] = param_storage
|
||||
|
||||
# Clear the InternalStorage keys which are not in use anymore
|
||||
dtype_in_use = list(self.dtype_rank_params.keys())
|
||||
dtype_to_pop = list(
|
||||
filter(lambda x: x not in dtype_in_use, self.param_storages.keys())
|
||||
)
|
||||
for d in dtype_to_pop:
|
||||
self.param_storages.pop(d)
|
||||
|
||||
if self.offload:
|
||||
self._optim._master_weights = self._master_params
|
||||
cpu_master_params = list(self._master_params.values())
|
||||
if self._default_device in core.get_all_custom_device_type():
|
||||
device_alignment = core.libpaddle._get_device_min_chunk_size(
|
||||
self._default_device
|
||||
)
|
||||
else:
|
||||
device_alignment = alignment[self._default_device]
|
||||
|
||||
for param in cpu_master_params:
|
||||
size = param._numel() * align[Type.fp32.value]
|
||||
remaining = size % device_alignment
|
||||
ali = 0 if remaining == 0 else device_alignment - remaining
|
||||
align_ = ali // align[Type.fp32.value]
|
||||
self.offload_buffer_size += param._numel() + align_
|
||||
self.offload_param2align[param.name] = align_
|
||||
|
||||
if cpu_master_params:
|
||||
with device_guard(self._rank, self.offload_device):
|
||||
self.offload_params = ParamStorage(
|
||||
size=self.offload_buffer_size,
|
||||
dtype=Type.fp32.value,
|
||||
device=self.offload_device,
|
||||
)
|
||||
self.offload_params.buffer.name = "offload_buffer"
|
||||
self.offload_params.add_rank_params(
|
||||
cpu_master_params, self.offload_param2align, False
|
||||
)
|
||||
self.offload_params.buffer.stop_gradient = False
|
||||
|
||||
self.offload_grads = GradStorage(
|
||||
size=self.offload_buffer_size,
|
||||
dtype=Type.fp32.value,
|
||||
device=self.offload_device,
|
||||
destination=self._rank,
|
||||
param2align=self.offload_param2align,
|
||||
convert_cpu=True,
|
||||
)
|
||||
for p in cpu_master_params:
|
||||
self.offload_grads.add_grad(
|
||||
p, self.offload_param2align[p.name]
|
||||
)
|
||||
|
||||
self._optim._master_weights[
|
||||
self.offload_params.buffer.name
|
||||
] = self.offload_params.buffer
|
||||
|
||||
def _offload_acc_grad(self, param_name, grad_fp32_cpu):
|
||||
"""accumulate grads with offload strategy"""
|
||||
with device_guard(self._rank, self.offload_device):
|
||||
if param_name in self._master_params.keys():
|
||||
if self._master_params[param_name].grad is None:
|
||||
self._master_params[param_name]._copy_gradient_from(
|
||||
grad_fp32_cpu
|
||||
)
|
||||
else:
|
||||
self._master_params[param_name].grad.add_(grad_fp32_cpu)
|
||||
|
||||
self.offload_params.buffer._copy_gradient_from(
|
||||
self.offload_grads.buffer
|
||||
)
|
||||
|
||||
def _offload_scale_grad(self, scale_size):
|
||||
"""scale grads with offload strategy"""
|
||||
with device_guard(self._rank, self.offload_device):
|
||||
self.offload_grads.buffer.scale_(scale=scale_size)
|
||||
|
||||
def _offload_clear_grad(self):
|
||||
"""clear grads with offload strategy"""
|
||||
with device_guard(self._rank, self.offload_device):
|
||||
self.offload_grads.buffer.zero_()
|
||||
|
||||
def _step(self):
|
||||
if self._broadcast_overlap:
|
||||
# Clear the pre forward hook in the optimizer step.
|
||||
for hook_remove in self._forward_pre_hook_remove_helper:
|
||||
hook_remove.remove()
|
||||
self._forward_pre_hook_remove_helper = []
|
||||
|
||||
if self.offload:
|
||||
params_list = [self.offload_params.buffer]
|
||||
|
||||
# TODO(Baibaifan): Offload will support param_groups later
|
||||
if not isinstance(self._optim._param_groups[0], dict):
|
||||
self._optim._parameter_list = params_list
|
||||
self._optim._param_groups = params_list
|
||||
|
||||
# Run the optimizer of the current rank step
|
||||
if self.offload:
|
||||
with device_guard(device=self.offload_device):
|
||||
self._optim.step()
|
||||
|
||||
for param in self._local_params:
|
||||
if param.name in self._master_params.keys():
|
||||
if (
|
||||
self._default_device
|
||||
in core.get_all_custom_device_type()
|
||||
):
|
||||
param.set_value(
|
||||
self._master_params[param.name]
|
||||
._copy_to(
|
||||
paddle.CustomPlace(
|
||||
self._default_device, self.dev_id
|
||||
),
|
||||
True,
|
||||
)
|
||||
.cast(dtype=param.dtype)
|
||||
)
|
||||
elif self._default_device == "xpu":
|
||||
param.set_value(
|
||||
self._master_params[param.name]
|
||||
.to("xpu:" + str(self.dev_id))
|
||||
.cast(dtype=param.dtype)
|
||||
)
|
||||
else:
|
||||
param.set_value(
|
||||
self._master_params[param.name]
|
||||
.cuda(self.dev_id)
|
||||
.cast(dtype=param.dtype)
|
||||
)
|
||||
else:
|
||||
self._optim.step()
|
||||
|
||||
# Synchronize all the updated shards in between the ranks
|
||||
self._broadcast_params()
|
||||
|
||||
def step(self):
|
||||
"""
|
||||
A wrapper for Optimizer's step function to finish the update operation of the optimizer.
|
||||
"""
|
||||
# This method won't be called directly by opt.step()!
|
||||
# The _redefine_opt_step() in class GroupShardedStage2 will wrap this function.
|
||||
self._step()
|
||||
|
||||
def minimize(self):
|
||||
raise RuntimeError(
|
||||
"optimizer.minimize() not support now, please use optimizer.step()"
|
||||
)
|
||||
|
||||
def set_state_dict(self, state_dict):
|
||||
self._optim.set_state_dict(state_dict)
|
||||
|
||||
def state_dict(self):
|
||||
return self._optim.state_dict()
|
||||
|
||||
def _clear_cache(self):
|
||||
self.__segment_params.clear()
|
||||
self._dtype_rank_params.clear()
|
||||
self._param2rank.clear()
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _broadcast_params(self):
|
||||
"""Broadcast the parameters of the current rank to each rank"""
|
||||
|
||||
# Exchange all the shards with the other ranks
|
||||
if self._broadcast_overlap:
|
||||
self._broadcast_params_overlap_forward()
|
||||
else:
|
||||
for dtype_per_rank in self.param_storages.values():
|
||||
for dst_rank, internal_storage in dtype_per_rank.items():
|
||||
dist.broadcast(
|
||||
tensor=internal_storage.buffer,
|
||||
src=self._group.ranks[dst_rank],
|
||||
group=self._group,
|
||||
sync_op=True,
|
||||
)
|
||||
|
||||
def _forward_pre_hook_function(self, tasks):
|
||||
# Since the layers will call pre hook by `forward_pre_hook(self, inputs)`,
|
||||
# the helper functions needs the x and y to take those params.
|
||||
def __impl__(x, y):
|
||||
for task in tasks:
|
||||
# Wait for broadcast task before using the result of the broadcast.
|
||||
task.wait()
|
||||
|
||||
return __impl__
|
||||
|
||||
def set_lr(self, lr):
|
||||
super().set_lr(lr)
|
||||
self._optim.set_lr(lr)
|
||||
|
||||
def get_lr(self):
|
||||
return self._optim.get_lr()
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _broadcast_params_overlap_forward(self):
|
||||
# Exchange all the shards with the other ranks,
|
||||
# but overlap the broadcast with next batch's calculation.
|
||||
group_idx = 0
|
||||
|
||||
param2task = {}
|
||||
for x in self._broadcast_order_params:
|
||||
if x.trainable:
|
||||
group = self._broadcast_groups[group_idx]
|
||||
group_idx = (group_idx + 1) % self._number_of_broadcast_groups
|
||||
task = dist.broadcast(
|
||||
tensor=x,
|
||||
src=group.ranks[self._param2rank[x.name]],
|
||||
group=group,
|
||||
sync_op=False,
|
||||
)
|
||||
assert x.name not in param2task
|
||||
param2task[x.name] = task
|
||||
|
||||
for layer in self._layers.sublayers():
|
||||
if len(layer.sublayers()) == 0:
|
||||
# Register forward pre hood for leaf layers. This will get the best performance.
|
||||
tasks = []
|
||||
for param in layer.parameters():
|
||||
if param.trainable:
|
||||
if param.name in param2task:
|
||||
tasks.append(param2task[param.name])
|
||||
self._forward_pre_hook_remove_helper.append(
|
||||
layer.register_forward_pre_hook(
|
||||
self._forward_pre_hook_function(tasks)
|
||||
)
|
||||
)
|
||||
|
||||
def sharded_state_dict(
|
||||
self,
|
||||
model_sharded_state_dict: ShardedStateDict,
|
||||
) -> ShardedStateDict:
|
||||
"""
|
||||
Convert optimizer state dict to a sharded state dict based on model sharding information.
|
||||
|
||||
Args:
|
||||
model_sharded_state_dict (dict): Sharded state dict of the model, containing tensor metadata.
|
||||
|
||||
Returns:
|
||||
dict: A new optimizer state dict where weights are wrapped as ShardedWeight.
|
||||
"""
|
||||
|
||||
_FP32_MASTER = "fp32_master_0"
|
||||
_MOMENT_NAME = "moment"
|
||||
_optimizer_scalar_name = [
|
||||
"beta1_pow_acc_0",
|
||||
"beta2_pow_acc_0",
|
||||
]
|
||||
_optimizer_non_scaler_name = [
|
||||
"moment1_0",
|
||||
"moment2_0",
|
||||
"velocity_0",
|
||||
]
|
||||
|
||||
def _generate_base_static_name(vname):
|
||||
if _FP32_MASTER in vname:
|
||||
return tuple(vname.split("_" + _FP32_MASTER + "_", 1))
|
||||
for name in _optimizer_scalar_name + _optimizer_non_scaler_name:
|
||||
if vname.endswith(name):
|
||||
return vname[: -(len(name) + 1)], name
|
||||
raise ValueError(f"Cannot split variable name: {vname}.")
|
||||
|
||||
optimizer_sharded_state_dict = {}
|
||||
optimizer_state_dict = self.state_dict()
|
||||
# Build name mapping and remove non-tensor entries from optimizer state
|
||||
static_to_struct_mapping = {}
|
||||
model_sharded_state_dict = dict(
|
||||
sorted(model_sharded_state_dict.items())
|
||||
)
|
||||
for k, v in model_sharded_state_dict.items():
|
||||
# When shared weights exist, the v.local_tensor.name of shared parameters are identical, but only the first parameter has optimizer states. Therefore, only the key-value pairs of the first occurrence in the shared parameter group need to be retained.
|
||||
if v.local_tensor.name not in static_to_struct_mapping:
|
||||
static_to_struct_mapping[v.local_tensor.name] = k
|
||||
|
||||
master_weights = optimizer_state_dict.pop("master_weights", None)
|
||||
optimizer_state_dict.pop("LR_Scheduler", None)
|
||||
|
||||
# Process main optimizer states
|
||||
for key, tensor in optimizer_state_dict.items():
|
||||
static_name, optim_state_type = _generate_base_static_name(key)
|
||||
struct_name = static_to_struct_mapping[static_name]
|
||||
sharded_weight = model_sharded_state_dict[struct_name]
|
||||
|
||||
unified_name = f"{struct_name}.{optim_state_type}"
|
||||
|
||||
# Determine tensor partitioning scheme
|
||||
if _MOMENT_NAME in optim_state_type:
|
||||
optimizer_sharded_state_dict[unified_name] = (
|
||||
create_sharded_weight_with_new_local(
|
||||
unified_name, tensor, sharded_weight
|
||||
)
|
||||
)
|
||||
else: # Non-momentum parameters
|
||||
optimizer_sharded_state_dict[unified_name] = ShardedWeight(
|
||||
key=unified_name,
|
||||
local_tensor=tensor,
|
||||
local_shape=(1,),
|
||||
global_shape=(1,),
|
||||
global_offset=(0,),
|
||||
)
|
||||
|
||||
# Process master weights if using mixed precision
|
||||
if master_weights is not None:
|
||||
for key, tensor in master_weights.items():
|
||||
struct_name = static_to_struct_mapping[key]
|
||||
sharded_weight = model_sharded_state_dict[struct_name]
|
||||
unified_name = f"{struct_name}.w_0"
|
||||
optimizer_sharded_state_dict[unified_name] = (
|
||||
create_sharded_weight_with_new_local(
|
||||
unified_name, tensor, sharded_weight
|
||||
)
|
||||
)
|
||||
|
||||
return optimizer_sharded_state_dict
|
||||
@@ -0,0 +1,720 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
# The file has been adapted from fairscale file:
|
||||
# https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/data_parallel/sharded_ddp.py
|
||||
# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
|
||||
# We retain the following license from the original files:
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
from functools import reduce
|
||||
from types import MethodType
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
from paddle import nn
|
||||
from paddle.distributed import collective
|
||||
from paddle.distributed.utils.log_utils import get_logger
|
||||
from paddle.framework import core
|
||||
|
||||
from .group_sharded_optimizer_stage2 import GroupShardedOptimizerStage2
|
||||
from .group_sharded_storage import GradStorage
|
||||
from .group_sharded_utils import Type, device_guard
|
||||
|
||||
logger_ = get_logger(logging.WARNING)
|
||||
|
||||
|
||||
def _trainable(param):
|
||||
return param.trainable
|
||||
|
||||
|
||||
class GroupShardedStage2(nn.Layer):
|
||||
"""
|
||||
A wrapper for Sharding Stage2 Layer in Dygraph.
|
||||
.. warning: GroupShardedStage2 encapsulates the layer strategy and integrates it into the nn.Layer.
|
||||
.. ZeRO: https://arxiv.org/pdf/1910.02054.pdf.
|
||||
"""
|
||||
|
||||
# TODO (Baibaifan)
|
||||
# Feature Notes::
|
||||
# 1. Unified memory for param and param.grad to InternalStorage.
|
||||
# 2. Divide param.grad according to rank to centrally apply for and release GPU memory.
|
||||
# 3. Dynamically adjust training parameters and models.
|
||||
# 4. Support offload function.
|
||||
# 5. Support the establishment of independent communication groups.
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
layer,
|
||||
sharding_optimizer,
|
||||
group=None,
|
||||
sync_buffers=False,
|
||||
buffer_max_size=2**23, # 8MB
|
||||
auto_refresh_trainable=True,
|
||||
device="xpu" if core.is_compiled_with_xpu() else "gpu",
|
||||
dp_group=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# training options
|
||||
self._layer = layer
|
||||
self._sharding_optimizers = (
|
||||
[sharding_optimizer]
|
||||
if not isinstance(sharding_optimizer, list)
|
||||
else sharding_optimizer
|
||||
)
|
||||
assert all(
|
||||
isinstance(opt, GroupShardedOptimizerStage2)
|
||||
for opt in self._sharding_optimizers
|
||||
), "Please use GroupShardedOptimizerStage2 optimizer"
|
||||
self._sync_buffers = sync_buffers
|
||||
self._auto_refresh_trainable = auto_refresh_trainable
|
||||
|
||||
# Communication related attributes
|
||||
self._group = (
|
||||
collective.new_group(collective._get_global_group().ranks)
|
||||
if group is None
|
||||
else group
|
||||
)
|
||||
self._world_size_scaling = 1.0 / self._group.nranks
|
||||
assert self._group.nranks > 1, (
|
||||
"Training must be distributed, ranks must be greater than 1"
|
||||
)
|
||||
self._rank = self._group.rank
|
||||
self._global_root_rank = self._group.ranks[
|
||||
0
|
||||
] # picking ranks index 0 as the reference
|
||||
self._default_device = device
|
||||
|
||||
self._dp_group = dp_group
|
||||
|
||||
# Global statistical parameters
|
||||
self._all_params = []
|
||||
for optim in self._sharding_optimizers:
|
||||
self._all_params.extend(list(optim.local_params))
|
||||
self.use_main_grad = None
|
||||
for param in self._all_params:
|
||||
if self.use_main_grad is None and hasattr(param, "main_grad"):
|
||||
self.use_main_grad = True
|
||||
if self.use_main_grad:
|
||||
assert hasattr(param, "main_grad"), (
|
||||
"Params have different main grad attributes."
|
||||
)
|
||||
|
||||
# sharing stage 2 comm overlap flag
|
||||
self._reduce_overlap = False
|
||||
|
||||
self._grad_reduced = []
|
||||
self._trainable_param2rank = {}
|
||||
self._trainable_param2align = {}
|
||||
self._trainable_params = list(
|
||||
filter(lambda x: x.trainable, self._all_params)
|
||||
)
|
||||
self._trainable_mask = list(map(_trainable, self._trainable_params))
|
||||
self._param_grads = []
|
||||
|
||||
# Set grad storage size & Display param sizes and model sizes
|
||||
model_size = sum([p._numel() for p in self._layer.parameters()])
|
||||
assert buffer_max_size >= 0, "buffer_max_size must be GE than 0."
|
||||
self._buffer_max_size = self._rank_buffer_size(
|
||||
buffer_max_size, model_size
|
||||
)
|
||||
self._use_grad_storage = buffer_max_size > 0
|
||||
self._grad_storages = {} # {dtype: {rank: GradStorage}}
|
||||
self._has_grad_storage = []
|
||||
self._grad_storage_list = []
|
||||
|
||||
# Offload
|
||||
# TODO(haohongxiang): Now it's not be supported for multi-optimizers using Offload strategy
|
||||
self._offload_optims = list(
|
||||
filter(lambda optim: optim.offload, self._sharding_optimizers)
|
||||
)
|
||||
if len(self._offload_optims) > 0:
|
||||
assert len(self._sharding_optimizers) == 1, (
|
||||
"Only support offload strategy for single optimizer"
|
||||
)
|
||||
|
||||
self._offload = len(self._offload_optims) > 0
|
||||
self._offload_device = "cpu"
|
||||
|
||||
# Set backward pass hooks
|
||||
self._bw_hooks = []
|
||||
|
||||
self.scale_in_opt = False
|
||||
|
||||
# TODO (Baibaifan) Set tasks flow support asynchronous communicate
|
||||
# self._tasks_flow = deque()
|
||||
|
||||
# Define optimizer step and clear_grad
|
||||
self._redefine_opt_step()
|
||||
self._redefine_opt_clear()
|
||||
|
||||
def forward(self, *inputs, **kwargs):
|
||||
"""
|
||||
A wrapper for Sharding Stage2 layer.
|
||||
- Fresh trainable params or rebuild grad storage
|
||||
- Sync layer's buffer params
|
||||
- Clear all flags states
|
||||
- Forward for origin layers
|
||||
"""
|
||||
|
||||
# Whether to need to reset trainable parameters
|
||||
needs_fresh = len(self._bw_hooks) == 0 and self.training
|
||||
|
||||
if self._auto_refresh_trainable:
|
||||
needs_fresh |= self._detect_train_change()
|
||||
|
||||
# Front hook
|
||||
self._init_internal_storage(needs_fresh)
|
||||
|
||||
# Sync layer's buffers state
|
||||
if self._sync_buffers:
|
||||
self.__sync_buffers()
|
||||
|
||||
# Normal FW on the base model
|
||||
fw = self._layer(*inputs, **kwargs)
|
||||
|
||||
return fw
|
||||
|
||||
def set_state_dict(self, state_dict, use_structured_name=True):
|
||||
self._layer.set_state_dict(
|
||||
state_dict, use_structured_name=use_structured_name
|
||||
)
|
||||
|
||||
def state_dict(
|
||||
self,
|
||||
destination=None,
|
||||
include_sublayers=True,
|
||||
structured_name_prefix="",
|
||||
):
|
||||
return self._layer.state_dict(
|
||||
destination=destination,
|
||||
include_sublayers=include_sublayers,
|
||||
structured_name_prefix=structured_name_prefix,
|
||||
)
|
||||
|
||||
def _clear_gradients(self):
|
||||
"""
|
||||
Set zero to the gradient of the optimizer's current rank trainable parameters.
|
||||
"""
|
||||
# Release grad storages
|
||||
for dtype in self._grad_storages.keys():
|
||||
if (
|
||||
not self._offload
|
||||
and self._rank in self._grad_storages[dtype].keys()
|
||||
):
|
||||
self._grad_storages[dtype][self._rank].buffer.zero_()
|
||||
|
||||
# Release grads of params
|
||||
for param in self._trainable_params:
|
||||
if param.name in self._param_grads:
|
||||
if self.use_main_grad and param.main_grad is not None:
|
||||
param.main_grad.zero_()
|
||||
elif param.grad is not None:
|
||||
param._zero_grads()
|
||||
|
||||
# Release grads of master params with offload strategy
|
||||
if self._offload:
|
||||
self._sharding_optimizers[0]._offload_clear_grad()
|
||||
|
||||
def _grad_scale(self):
|
||||
"""
|
||||
this function will do 2 things:
|
||||
1. Before the optimization, scale main_grad to support gradient merge if param has main_grad, or to support fused_linear_param_grad_add gradient merge.
|
||||
2. Before the optimization, scale the gradients before allreduce of dp_group.
|
||||
"""
|
||||
|
||||
need_dp_scale = self._dp_group is not None and self._dp_group.nranks > 1
|
||||
if self.scale_in_opt:
|
||||
scale_factor = self._world_size_scaling
|
||||
else:
|
||||
scale_factor = 1.0
|
||||
|
||||
if need_dp_scale:
|
||||
dp_scale_factor = 1.0 / (self._dp_group.nranks)
|
||||
scale_factor = scale_factor * dp_scale_factor
|
||||
|
||||
# Scale grad storages
|
||||
for dtype in self._grad_storages.keys():
|
||||
if (
|
||||
not self._offload
|
||||
and self._rank in self._grad_storages[dtype].keys()
|
||||
):
|
||||
self._grad_storages[dtype][self._rank].buffer.scale_(
|
||||
scale=scale_factor
|
||||
)
|
||||
# Scale grads of params
|
||||
with paddle.no_grad():
|
||||
for param in self._trainable_params:
|
||||
if param.name in self._param_grads:
|
||||
if self.use_main_grad and param.main_grad is not None:
|
||||
param.main_grad.scale_(scale=scale_factor)
|
||||
elif param.grad is not None:
|
||||
param.grad.scale_(scale=scale_factor)
|
||||
|
||||
# Scale grads of master params with offload strategy
|
||||
if self._offload:
|
||||
if need_dp_scale is False:
|
||||
return
|
||||
self._sharding_optimizers[0]._offload_scale_grad(
|
||||
scale=1.0 / (self._dp_group.nranks)
|
||||
)
|
||||
|
||||
def _init_internal_storage(self, needs_fresh):
|
||||
"""
|
||||
Judge Fresh trainable params or rebuild grad storage.
|
||||
"""
|
||||
if needs_fresh:
|
||||
self._fresh_trainable()
|
||||
else:
|
||||
self._build_grad_storages()
|
||||
|
||||
# Clear all flags state
|
||||
self._clear_counters()
|
||||
|
||||
def to(self, device=None, dtype=None, blocking=True):
|
||||
"""
|
||||
Synchronously or asynchronously convert the data type of the layer, the device is not supported now.
|
||||
"""
|
||||
assert isinstance(device, str), "Device must be type str"
|
||||
assert device == self._default_device, (
|
||||
"New devices are not supported, because of the optimizer state is not sync"
|
||||
)
|
||||
|
||||
self._layer.to(device=device, dtype=dtype, blocking=blocking)
|
||||
|
||||
# Re-build the buckets, hooks, etc..
|
||||
self._fresh_trainable()
|
||||
|
||||
def _fresh_trainable(self):
|
||||
"""Whether to update training parameters."""
|
||||
|
||||
# Make sure that this is not done while gradients are waiting to be reduced (if no_sync context for instance)
|
||||
if reduce(lambda x, y: x or y, self._grad_reduced, False):
|
||||
logging.warning("Grads waiting to be reduced.")
|
||||
|
||||
self._trainable_params = list(
|
||||
filter(lambda x: x.trainable, self._all_params)
|
||||
)
|
||||
self._trainable_params.sort(key=lambda x: x._numel())
|
||||
|
||||
self._trainable_param2rank = {}
|
||||
for optim in self._sharding_optimizers:
|
||||
# Need to be wrapped for Sharding Stage2 Optimizer
|
||||
if len(optim.param_storages.keys()) == 0:
|
||||
optim._update_opt_status()
|
||||
|
||||
# Get the parameters split by the optimizer according to rank
|
||||
for per_rank_params in (
|
||||
optim.dtype_rank_params.values()
|
||||
): # all the params from all ranks
|
||||
for params in per_rank_params:
|
||||
for param in filter(lambda x: x.trainable, params):
|
||||
self._trainable_param2rank[param.name] = (
|
||||
optim.param2rank[param.name]
|
||||
)
|
||||
self._trainable_param2align[param.name] = (
|
||||
optim._param2align[param.name]
|
||||
)
|
||||
|
||||
# Create grad_storage
|
||||
self._setup_use_grad_storage()
|
||||
# setup backward hooks
|
||||
self._setup_backward_hooks()
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def __sync_buffers(self):
|
||||
"""
|
||||
Sync all the param buffers from all ranks (exp: batch norm statistics).
|
||||
"""
|
||||
|
||||
for buffer in self._layer.buffers(include_sublayers=True):
|
||||
dist.broadcast(
|
||||
buffer, self._global_root_rank, self._group, sync_op=True
|
||||
)
|
||||
|
||||
if self._dp_group and self._dp_group.nranks > 1:
|
||||
dist.broadcast(
|
||||
buffer,
|
||||
self._dp_group.ranks[0],
|
||||
self._dp_group,
|
||||
sync_op=True,
|
||||
)
|
||||
|
||||
def __getattr__(self, name):
|
||||
"""Forward missing attributes to wrapped layer."""
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
return getattr(self._layer, name)
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _clear_counters(self):
|
||||
"""Reset all the grad reduce and call counters."""
|
||||
if self.training:
|
||||
self._grad_reduced = [True for _ in self._trainable_params]
|
||||
|
||||
if self._use_grad_storage:
|
||||
for grad_storage in self._grad_storage_list:
|
||||
grad_storage.reset_checked_in()
|
||||
|
||||
def _set_reduce_overlap(self, reduce_overlap):
|
||||
# Hacky way to not add an extra parameter to the `group_sharded_parallel` funct.
|
||||
# User should use this like:
|
||||
# model, optimizer, scaler = group_sharded_parallel(...)
|
||||
# model._set_reduce_overlap(True)
|
||||
self._reduce_overlap = reduce_overlap
|
||||
if self._reduce_overlap:
|
||||
assert len(self._sharding_optimizers) == 1, (
|
||||
"Only support comm overlap strategy for single optimizer"
|
||||
)
|
||||
self._sharding_optimizers[0]._set_reduce_overlap(reduce_overlap)
|
||||
|
||||
def _get_scaled_grad_fn(self, param):
|
||||
@paddle.autograd.no_grad()
|
||||
def scale(grad):
|
||||
# do gradient scale separately
|
||||
# For grad scale, we need to do it in the backward hook due to fp16 may overflow if we first add grad and then scale
|
||||
# For main_grad scale and fused_linear_param_grad_add, we do scale in the optimizer.
|
||||
if not self.scale_in_opt:
|
||||
if (
|
||||
not hasattr(param, "main_grad")
|
||||
and grad is not None
|
||||
and grad.dtype == Type.fp16.value
|
||||
):
|
||||
assert grad._is_initialized(), (
|
||||
"grad should be initialized in stage2"
|
||||
)
|
||||
grad.scale_(self._world_size_scaling)
|
||||
else:
|
||||
self.scale_in_opt = True
|
||||
|
||||
return scale
|
||||
|
||||
def _get_reduce_fn(self, index, param, dst_rank):
|
||||
"""
|
||||
There are two ways to reduce gradient.
|
||||
- 1. Do not use self._use_grad_storage or exceeded buffer_max_size will be reduced separately.
|
||||
- 2. Use grad_storage Reduce the storage to get the full gradient from different ranks.
|
||||
"""
|
||||
|
||||
if not self._use_grad_storage or not self._has_grad_storage[index]:
|
||||
# Direct reduction
|
||||
@paddle.autograd.no_grad()
|
||||
def reduce(*_):
|
||||
# Skip gradient reduction, do not change status information
|
||||
if self._grad_reduced[index]:
|
||||
assert (
|
||||
param.grad is not None or param.main_grad is not None
|
||||
), "Parameter should have grad or main grad"
|
||||
|
||||
# Change reduce information
|
||||
self._grad_reduced[index] = False
|
||||
|
||||
# Clear the gradient that does not belong to the current rank through the callback function
|
||||
def cleanup():
|
||||
if dst_rank != self._rank:
|
||||
if self.use_main_grad:
|
||||
param.main_grad._clear_data()
|
||||
param.main_grad = None
|
||||
else:
|
||||
param.clear_gradient(False)
|
||||
elif self._offload:
|
||||
tmp_grad = param.grad.cast(
|
||||
dtype=Type.fp32.value
|
||||
).cpu()
|
||||
|
||||
self._sharding_optimizers[0]._offload_acc_grad(
|
||||
param.name, tmp_grad
|
||||
)
|
||||
del tmp_grad
|
||||
param.clear_gradient(False)
|
||||
|
||||
# Synchronize the reduce parameter gradient asynchronize
|
||||
self._sharding_optimizers[0]._update_task(
|
||||
dist.reduce(
|
||||
tensor=(
|
||||
param.grad
|
||||
if not self.use_main_grad
|
||||
else param.main_grad
|
||||
),
|
||||
dst=self._group.ranks[dst_rank],
|
||||
group=self._group,
|
||||
sync_op=not self._reduce_overlap,
|
||||
)
|
||||
)
|
||||
|
||||
# Clear the task flow and trigger callback to clear the redundant gradient
|
||||
# self._clear_task_flow()
|
||||
|
||||
cleanup()
|
||||
|
||||
else:
|
||||
# Buffer reduction
|
||||
@paddle.autograd.no_grad()
|
||||
def reduce(*_):
|
||||
# Skip gradient reduction, do not change status information
|
||||
if self._grad_reduced[index]:
|
||||
assert (
|
||||
param.grad is not None or param.main_grad is not None
|
||||
), "Parameter should have grad or main grad"
|
||||
|
||||
# Change reduce information
|
||||
self._grad_reduced[index] = False
|
||||
grad_storage = self._grad_storages[param.dtype][dst_rank]
|
||||
grad_storage.params_checked_in += 1
|
||||
|
||||
if grad_storage.all_checked_in:
|
||||
assert grad_storage.buffer is not None
|
||||
|
||||
# Clearing up the grad_storage buffer
|
||||
def cleanup():
|
||||
if dst_rank != self._rank:
|
||||
for p in grad_storage._params:
|
||||
if self.use_main_grad:
|
||||
p.main_grad._clear_data()
|
||||
p.main_grad = None
|
||||
else:
|
||||
p.clear_gradient(False)
|
||||
|
||||
grad_storage.buffer._clear_data()
|
||||
elif self._offload:
|
||||
grad_storage.to(device=self._offload_device)
|
||||
for p in grad_storage._params:
|
||||
with device_guard():
|
||||
tmp_grad = p.grad.cast(
|
||||
dtype=Type.fp32.value
|
||||
)
|
||||
self._sharding_optimizers[
|
||||
0
|
||||
]._offload_acc_grad(p.name, tmp_grad)
|
||||
p.clear_gradient(False)
|
||||
grad_storage._device = self._default_device
|
||||
grad_storage.buffer._clear_data()
|
||||
|
||||
# Reduce the bucket
|
||||
grad_storage.sent = True
|
||||
# Synchronize the reduce parameter gradient asynchronize
|
||||
self._sharding_optimizers[0]._update_task(
|
||||
dist.reduce(
|
||||
tensor=grad_storage.buffer,
|
||||
dst=self._group.ranks[grad_storage.destination],
|
||||
group=self._group,
|
||||
sync_op=not self._reduce_overlap,
|
||||
)
|
||||
)
|
||||
|
||||
cleanup()
|
||||
|
||||
# Clear the task flow and trigger callback to clear the redundant gradient
|
||||
# self._clear_task_flow()
|
||||
|
||||
return reduce
|
||||
|
||||
def _setup_backward_hooks(self):
|
||||
"""
|
||||
Set the backward hook to synchronize the gradients of all rank by reduce group ranks.
|
||||
"""
|
||||
|
||||
# Remove previous backward hooks
|
||||
while len(self._bw_hooks) > 0:
|
||||
self._bw_hooks.pop().remove()
|
||||
|
||||
# Go through the parameters, attach the hook
|
||||
if not self.training:
|
||||
return
|
||||
|
||||
for index, param in enumerate(self._trainable_params):
|
||||
param._register_grad_hook(self._get_scaled_grad_fn(param))
|
||||
|
||||
dst_rank = self._trainable_param2rank[param.name]
|
||||
|
||||
reduce_function = self._get_reduce_fn(index, param, dst_rank)
|
||||
|
||||
self._bw_hooks.append(
|
||||
param._register_backward_hook(reduce_function)
|
||||
)
|
||||
|
||||
def _setup_use_grad_storage(self):
|
||||
"""
|
||||
Integrate the parameters gradient into a continuous memory according to rank, and support the update of training parameters.
|
||||
"""
|
||||
|
||||
# According to parameters's numel sort, allocate memory of parameter gradient to continuous memory according to rank
|
||||
self._grad_storages = {}
|
||||
self._has_grad_storage = [False for _ in self._trainable_params]
|
||||
|
||||
for index, param in enumerate(self._trainable_params):
|
||||
dst_rank = self._trainable_param2rank[param.name]
|
||||
|
||||
if param.dtype not in self._grad_storages.keys():
|
||||
self._grad_storages[param.dtype] = {}
|
||||
|
||||
if dst_rank not in self._grad_storages[param.dtype].keys():
|
||||
self._grad_storages[param.dtype][dst_rank] = GradStorage(
|
||||
self._buffer_max_size[param.dtype],
|
||||
dtype=(
|
||||
param.dtype
|
||||
if not self.use_main_grad
|
||||
else paddle.float32
|
||||
),
|
||||
device=self._default_device,
|
||||
destination=dst_rank,
|
||||
param2align=self._trainable_param2align,
|
||||
)
|
||||
|
||||
# Criteria to decide whether this parameter is to be put in GradStorage
|
||||
if self._grad_storages[param.dtype][dst_rank].can_add_grad_view(
|
||||
param, self._trainable_param2align[param.name]
|
||||
):
|
||||
self._grad_storages[param.dtype][dst_rank].add_grad(
|
||||
param, self._trainable_param2align[param.name]
|
||||
)
|
||||
self._has_grad_storage[index] = True
|
||||
else:
|
||||
self._param_grads.append(param.name)
|
||||
|
||||
for dtype in self._grad_storages.keys():
|
||||
self._grad_storage_list.extend(
|
||||
list(self._grad_storages[dtype].values())
|
||||
)
|
||||
|
||||
# def _clear_task_flow(self):
|
||||
# """Try to consume the previous tasks."""
|
||||
# while len(self._tasks_flow) > 0:
|
||||
# task = self._tasks_flow.popleft()
|
||||
# task.wait()
|
||||
# if task.callback is not None:
|
||||
# task.callback()
|
||||
|
||||
def _detect_train_change(self):
|
||||
# Current trainable parameters
|
||||
trainable_mask = list(map(_trainable, self._trainable_params))
|
||||
|
||||
# Whether parameters trainability changed
|
||||
trainability_changed = trainable_mask != self._trainable_mask
|
||||
|
||||
if trainability_changed:
|
||||
logging.warning(
|
||||
"Trainable params changed, because of eval/train mode or parameter freezing/unfreeze."
|
||||
)
|
||||
self._trainable_mask = trainable_mask
|
||||
|
||||
return trainability_changed
|
||||
|
||||
def _build_grad_storages(self):
|
||||
"""
|
||||
Rebuild grad storages.
|
||||
"""
|
||||
# Rebuild fp16/fp32 grad storages
|
||||
for dtype in self._grad_storages.keys():
|
||||
for dst_rank, grad_storage in self._grad_storages[dtype].items():
|
||||
if self._offload or dst_rank != self._rank:
|
||||
grad_storage.manual_release()
|
||||
grad_storage.rebuild()
|
||||
|
||||
def _rank_buffer_size(self, buffer_max_size, model_size):
|
||||
"""
|
||||
Generate the minimum buffer size for each rank & Display param sizes and model sizes.
|
||||
"""
|
||||
|
||||
# Initialize buffer size
|
||||
rank_buffer_size = {}
|
||||
for shard_opt in self._sharding_optimizers:
|
||||
if shard_opt.rank_buffer_size:
|
||||
for dtype in shard_opt.rank_buffer_size.keys():
|
||||
sizes = max(shard_opt.rank_buffer_size[dtype].values())
|
||||
rank_buffer_size[dtype] = min(sizes, buffer_max_size)
|
||||
|
||||
if Type.fp16.value in rank_buffer_size.keys():
|
||||
# FP16 GradStorage and model size
|
||||
logger_.info(
|
||||
f"====== FP16 GradStorage size: {rank_buffer_size[Type.fp16.value] / 2**19:.2f}M parameters, Model size {model_size / 2**19:.2f}M parameters ======"
|
||||
)
|
||||
if Type.bf16.value in rank_buffer_size.keys():
|
||||
# FP16 GradStorage and model size
|
||||
logger_.info(
|
||||
f"====== BF16 GradStorage size: {rank_buffer_size[Type.bf16.value] / 2**19:.2f}M parameters, Model size {model_size / 2**19:.2f}M parameters ======"
|
||||
)
|
||||
if Type.fp32.value in rank_buffer_size.keys():
|
||||
# FP32 GradStorage and model size
|
||||
logger_.info(
|
||||
f"====== FP32 GradStorage size: {rank_buffer_size[Type.fp32.value] / 2**18:.2f}M parameters, Model size {model_size / 2**18:.2f}M parameters ======"
|
||||
)
|
||||
return rank_buffer_size
|
||||
|
||||
def _dp_allreduce(self):
|
||||
# do dp allreduce here for gradient merge.
|
||||
if self._dp_group and self._dp_group.nranks > 1:
|
||||
for dtype in self._grad_storages.keys():
|
||||
for rank, g in sorted(
|
||||
self._grad_storages[dtype].items(), key=lambda x: x[0]
|
||||
):
|
||||
if g.destination == self._rank:
|
||||
assert g.buffer._is_initialized()
|
||||
dist.all_reduce(
|
||||
tensor=g.buffer,
|
||||
group=self._dp_group,
|
||||
sync_op=True,
|
||||
)
|
||||
for param in self._trainable_params:
|
||||
if param.name in self._param_grads:
|
||||
if self.use_main_grad and param.main_grad is None:
|
||||
continue
|
||||
elif param.grad is None:
|
||||
continue
|
||||
dst_rank = self._trainable_param2rank[param.name]
|
||||
if dst_rank == self._rank:
|
||||
dist.all_reduce(
|
||||
tensor=(
|
||||
param.grad
|
||||
if not self.use_main_grad
|
||||
else param.main_grad
|
||||
),
|
||||
group=self._dp_group,
|
||||
sync_op=True,
|
||||
)
|
||||
|
||||
def _redefine_opt_step(self):
|
||||
grad_func = self._grad_scale
|
||||
dp_allreduce_func = self._dp_allreduce
|
||||
|
||||
for opt in self._sharding_optimizers:
|
||||
opt_step = opt.step
|
||||
|
||||
def _opt_step(self):
|
||||
if self._reduce_overlap:
|
||||
# Wait for the last reduce task. This wait must before grad scale function.
|
||||
assert self._comm_task is not None
|
||||
self._comm_task.wait()
|
||||
|
||||
grad_func()
|
||||
dp_allreduce_func()
|
||||
opt_step()
|
||||
|
||||
opt.step = MethodType(_opt_step, opt)
|
||||
|
||||
def _redefine_opt_clear(self):
|
||||
clear_func = self._clear_gradients
|
||||
|
||||
def _opt_clear(self):
|
||||
clear_func()
|
||||
|
||||
for opt in self._sharding_optimizers:
|
||||
opt.clear_grad = MethodType(_opt_clear, opt)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,363 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
# The file has been adapted from fairscale file:
|
||||
# https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/misc/param_bucket.py
|
||||
# Git commit hash: 8acbec718f3c70a6b9785470bb9e05cd84fc3f8e
|
||||
# We retain the following license from the original files:
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.framework import core
|
||||
|
||||
from .group_sharded_utils import Type, cvt_to_device, device_guard
|
||||
|
||||
|
||||
class BufferWarper(core.eager.Tensor):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.need_clip = True
|
||||
self.is_distributed = False
|
||||
self.trainable = True
|
||||
|
||||
|
||||
class InternalStorage:
|
||||
"""
|
||||
This is a basic class, which is responsible for consolidating the basic storage tensor.
|
||||
|
||||
"""
|
||||
|
||||
# Support integration parameter tensor
|
||||
def __init__(self, size, dtype, device, convert_cpu=False):
|
||||
self._params = []
|
||||
self._param_ids = []
|
||||
self._fill = 0
|
||||
self._device = device
|
||||
self._dtype = dtype
|
||||
|
||||
# The flatten tensor
|
||||
size = [size] if isinstance(size, int) else size
|
||||
if convert_cpu:
|
||||
value = (
|
||||
np.zeros(size, dtype=np.float16)
|
||||
if Type.fp16.value == dtype
|
||||
else np.zeros(size, dtype=np.float32)
|
||||
)
|
||||
self.buffer = core.eager.Tensor(value=value, place=core.CPUPlace())
|
||||
if dtype == Type.bf16.value:
|
||||
self.buffer = paddle.cast(self.buffer, dtype=paddle.bfloat16)
|
||||
else:
|
||||
self.buffer = paddle.zeros(size, dtype=dtype)
|
||||
|
||||
self.dev_id = (
|
||||
0
|
||||
if paddle.get_device() == "cpu"
|
||||
else int(paddle.get_device().split(":")[1])
|
||||
)
|
||||
|
||||
def to(self, device, dtype=None, keep_alignment=True):
|
||||
"""
|
||||
Move the underlying buffer
|
||||
"""
|
||||
assert self.buffer is not None, (
|
||||
"Cannot move a collapsed bucket, please rebuild it"
|
||||
)
|
||||
assert dtype == Type.fp32.value or Type.fp16.value, (
|
||||
"Conversion type is not supported now"
|
||||
)
|
||||
|
||||
if self._device != device:
|
||||
if device in paddle.device.get_all_custom_device_type():
|
||||
tmp_buffer = self.buffer._copy_to(
|
||||
paddle.CustomPlace(device, self.dev_id), True
|
||||
)
|
||||
else:
|
||||
tmp_buffer = (
|
||||
cvt_to_device(self.buffer, self.dev_id)
|
||||
if device in ["gpu", "xpu"]
|
||||
else self.buffer.cpu()
|
||||
)
|
||||
for param in self._params:
|
||||
param.clear_gradient(False)
|
||||
|
||||
del self.buffer
|
||||
self.buffer = tmp_buffer
|
||||
self._device = device
|
||||
|
||||
if dtype is not None:
|
||||
self.buffer = self.buffer.cast(dtype=dtype)
|
||||
self._dtype = dtype
|
||||
|
||||
def warp_buffer(self):
|
||||
tmp_buffer = BufferWarper()
|
||||
self._buffer = self.buffer
|
||||
tmp_buffer.get_tensor()._share_data_with(self.buffer.get_tensor())
|
||||
self.buffer = tmp_buffer
|
||||
|
||||
|
||||
class ParamStorage(InternalStorage):
|
||||
"""
|
||||
This is a basic class to simplify the handling of parameter InternalStorages.
|
||||
"""
|
||||
|
||||
def __init__(self, size, dtype, device):
|
||||
super().__init__(size, dtype, device, convert_cpu=True)
|
||||
self.param2align = None
|
||||
|
||||
def to(self, device, dtype=None, keep_alignment=True):
|
||||
"""
|
||||
Move the underlying buffer
|
||||
"""
|
||||
|
||||
super().to(device, dtype)
|
||||
|
||||
if keep_alignment:
|
||||
self._array_params()
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def add_rank_params(self, trainable_params, param2align, convert_gpu=True):
|
||||
"""
|
||||
Add new parameters to the InternalStorage. Params becomes a view of this InternalStorage buffer.
|
||||
"""
|
||||
|
||||
assert all(
|
||||
id(param) not in self._param_ids for param in trainable_params
|
||||
), "The same param cannot be checked in twice"
|
||||
assert self.buffer is not None
|
||||
|
||||
self.param2align = param2align
|
||||
|
||||
cpu_param_shape = []
|
||||
for param in trainable_params:
|
||||
p_shape = self._add_param_as_view(
|
||||
param, param2align[param.name], convert_gpu
|
||||
)
|
||||
cpu_param_shape.append(p_shape)
|
||||
|
||||
if convert_gpu:
|
||||
if self._device in paddle.device.get_all_custom_device_type():
|
||||
self.buffer = self.buffer._copy_to(
|
||||
paddle.CustomPlace(self._device, self.dev_id), True
|
||||
)
|
||||
else:
|
||||
# buffer convert from cpu to cuda
|
||||
self.buffer = cvt_to_device(self.buffer, self.dev_id)
|
||||
|
||||
self._fill = 0
|
||||
|
||||
for idx, param in enumerate(trainable_params):
|
||||
self._convert_buffer(
|
||||
param, cpu_param_shape[idx], param2align[param.name]
|
||||
)
|
||||
self._params.append(param)
|
||||
self._param_ids.append(id(param))
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _add_param_as_view(self, param, align, convert_gpu=True):
|
||||
assert param.dtype == self.buffer.dtype, (
|
||||
f"Different types for the InternalStorage and the param, cannot proceed: {param.dtype} - {self.buffer.dtype}"
|
||||
)
|
||||
|
||||
var_end = self._fill + param._numel()
|
||||
offset = var_end + align
|
||||
assert offset <= self.buffer._numel()
|
||||
|
||||
p_shape = param.shape
|
||||
|
||||
origin_state = param.stop_gradient
|
||||
param.stop_gradient = True
|
||||
param.flatten_()
|
||||
param.stop_gradient = origin_state
|
||||
|
||||
# Copy the current param value
|
||||
|
||||
with device_guard(self.dev_id, "cpu"):
|
||||
tmp_var = self.buffer._slice(self._fill, var_end)
|
||||
if convert_gpu:
|
||||
param_cpu = param.cpu()
|
||||
param._clear_data()
|
||||
tmp_var.set_value(param_cpu)
|
||||
else:
|
||||
tmp_var.set_value(param)
|
||||
del tmp_var
|
||||
|
||||
self._fill = offset
|
||||
return p_shape
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _convert_buffer(self, param, p_shape, align):
|
||||
var_end = self._fill + np.prod(p_shape).tolist()
|
||||
offset = var_end + align
|
||||
assert offset <= self.buffer._numel()
|
||||
|
||||
# Convert the param value
|
||||
with device_guard(self.dev_id, self._device):
|
||||
tmp_tensor = self.buffer._slice(self._fill, var_end)
|
||||
tmp_tensor._share_buffer_to(param)
|
||||
param.get_tensor()._set_dims(p_shape)
|
||||
|
||||
self._fill = offset
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _array_params(self):
|
||||
"""
|
||||
Given the parameters which have been registered previously, rebuild the whole InternalStorage.
|
||||
"""
|
||||
assert len(self._params) > 0
|
||||
assert self.param2align is not None
|
||||
|
||||
self._fill = 0
|
||||
for p in self._params:
|
||||
self._convert_buffer(p, p.shape, self.param2align[p.name]) # modify
|
||||
|
||||
|
||||
class GradStorage(InternalStorage):
|
||||
"""
|
||||
This is a basic class to simplify the handling of gradient InternalStorages
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, size, dtype, device, destination, param2align, convert_cpu=False
|
||||
):
|
||||
if isinstance(size, np.int64):
|
||||
size = size.tolist()
|
||||
super().__init__(size, dtype, device, convert_cpu)
|
||||
|
||||
self._max_size = size
|
||||
self._release = False
|
||||
|
||||
self.params_checked_in = 0
|
||||
self.destination = destination
|
||||
self._param2align = param2align
|
||||
self.sent = False
|
||||
|
||||
def reset_checked_in(self):
|
||||
"""Reset the counter of the parameter grads which have been checked in"""
|
||||
self.params_checked_in = 0
|
||||
self.sent = False
|
||||
|
||||
@property
|
||||
def all_checked_in(self):
|
||||
"""Judge all the expected gradient check-in happened"""
|
||||
return len(self._params) == self.params_checked_in
|
||||
|
||||
def can_add_grad_view(self, param, align):
|
||||
"""Is there enough InternalStorage to add this parameter gradient, and whether this param have already checked in."""
|
||||
return (
|
||||
self._fill + param._numel() + align <= self._max_size
|
||||
and id(param) not in self._param_ids
|
||||
)
|
||||
|
||||
def to(self, device, dtype=None, keep_alignment=True):
|
||||
"""
|
||||
Move the underlying buffer
|
||||
"""
|
||||
if self._release:
|
||||
self.rebuild()
|
||||
|
||||
super().to(device, dtype)
|
||||
|
||||
if keep_alignment:
|
||||
self._array_grads()
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def add_grad(self, param, align):
|
||||
"""
|
||||
Add a new parameter gradient to the InternalStorage. Param.grad becomes a view of this InternalStorage buffer.
|
||||
"""
|
||||
|
||||
assert id(param) not in self._param_ids, (
|
||||
"The same gradients cannot be checked in twice"
|
||||
)
|
||||
|
||||
self._add_grad_as_view(param, align)
|
||||
self._params.append(param)
|
||||
self._param_ids.append(id(param))
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def manual_release(self):
|
||||
"""
|
||||
Release the buffer from InternalStorage. The InternalStorage will need to be rebuilt before use.
|
||||
"""
|
||||
if not self._release:
|
||||
for p in self._params:
|
||||
use_main_grad = hasattr(p, "main_grad")
|
||||
if use_main_grad and p.main_grad is not None:
|
||||
p.main_grad._clear_data()
|
||||
p.main_grad = None
|
||||
elif p.grad is not None:
|
||||
p.clear_gradient(False)
|
||||
|
||||
self.buffer = None
|
||||
self._fill = 0
|
||||
self.params_checked_in = 0
|
||||
self._release = True
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def rebuild(self):
|
||||
"""
|
||||
Given the parameter gradients which have been registered previously, rebuild the whole InternalStorage.
|
||||
"""
|
||||
|
||||
if self._release:
|
||||
self.buffer = paddle.zeros([self._max_size], dtype=self._dtype)
|
||||
|
||||
for p in self._params:
|
||||
self._add_grad_as_view(p, self._param2align[p.name])
|
||||
|
||||
self._release = False
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _array_grads(self):
|
||||
"""
|
||||
Given the parameters gradients which have been registered previously, rebuild the whole InternalStorage.
|
||||
"""
|
||||
if len(self._params) > 0:
|
||||
self._fill = 0
|
||||
for p in self._params:
|
||||
self._add_grad_as_view(p, self._param2align[p.name])
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _add_grad_as_view(self, param, align):
|
||||
assert param._numel() > 0, (
|
||||
"Cannot add a gradient to a released InternalStorage, please rebuild"
|
||||
)
|
||||
|
||||
use_main_grad = hasattr(param, "main_grad")
|
||||
if use_main_grad:
|
||||
assert self.buffer.dtype == paddle.float32
|
||||
else:
|
||||
assert param.dtype == self.buffer.dtype
|
||||
|
||||
grad_end = self._fill + param._numel()
|
||||
offset = grad_end + align
|
||||
assert offset <= self.buffer._numel()
|
||||
|
||||
# Copy the current grad value to InternalStorage
|
||||
with device_guard(self.dev_id, self._device):
|
||||
tmp_var = self.buffer._slice(self._fill, grad_end)
|
||||
tmp_var.get_tensor()._set_dims(param.shape)
|
||||
if not use_main_grad:
|
||||
param._copy_gradient_from(tmp_var)
|
||||
else:
|
||||
param.main_grad = tmp_var
|
||||
del tmp_var
|
||||
|
||||
self._fill = offset
|
||||
@@ -0,0 +1,352 @@
|
||||
# Copyright (c) 2022 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 contextlib
|
||||
from enum import Enum
|
||||
from types import MethodType
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle import _C_ops, _legacy_C_ops
|
||||
from paddle.base import core
|
||||
from paddle.common_ops_import import dygraph_only
|
||||
from paddle.nn import clip
|
||||
|
||||
|
||||
class Taskflow:
|
||||
"""
|
||||
Task flows, one way linked list for task acquisition.
|
||||
"""
|
||||
|
||||
def __init__(self, task, callback):
|
||||
self.task = task
|
||||
self.callback = callback
|
||||
|
||||
|
||||
class Type(Enum):
|
||||
"""
|
||||
Type of trainable parameters
|
||||
"""
|
||||
|
||||
fp16 = paddle.float16
|
||||
bf16 = paddle.bfloat16
|
||||
fp32 = paddle.float32
|
||||
|
||||
|
||||
class GroupShardedClipGrad:
|
||||
def __init__(self, clip, device, group):
|
||||
self._clip = clip
|
||||
self._device = device
|
||||
self._group = group
|
||||
|
||||
@paddle.autograd.no_grad()
|
||||
def _dygraph_clip(self, params_grads):
|
||||
sum_square_fp32, sum_square_fp16, sum_square_bfp16 = [], [], []
|
||||
unslice_params_fp32, unslice_params_fp16, unslice_params_bfp16 = (
|
||||
[],
|
||||
[],
|
||||
[],
|
||||
)
|
||||
|
||||
for p, g in params_grads:
|
||||
p_slice = True # using for slice parameter in sharding stage3
|
||||
if g is None or getattr(p, 'need_clip', True) is False:
|
||||
continue
|
||||
if hasattr(p, "unslice"):
|
||||
p_slice = False
|
||||
|
||||
merge_grad = g
|
||||
if g.type == core.VarDesc.VarType.SELECTED_ROWS:
|
||||
merge_grad = clip.get_tensor_from_selected_rows(
|
||||
clip.merge_selected_rows(g)
|
||||
)
|
||||
square = paddle.square(merge_grad)
|
||||
sum_square = paddle.sum(square)
|
||||
|
||||
if p.dtype == paddle.float16:
|
||||
if p_slice:
|
||||
sum_square_fp16.append(sum_square)
|
||||
else:
|
||||
unslice_params_fp16.append(sum_square)
|
||||
elif p.dtype == paddle.float32:
|
||||
if p_slice:
|
||||
sum_square_fp32.append(sum_square)
|
||||
else:
|
||||
unslice_params_fp32.append(sum_square)
|
||||
elif p.dtype == paddle.bfloat16:
|
||||
if p_slice:
|
||||
sum_square_bfp16.append(sum_square)
|
||||
else:
|
||||
unslice_params_bfp16.append(sum_square)
|
||||
|
||||
# global norm of non-distributed FP16 params_and_grads
|
||||
if len(sum_square_fp16) == 0:
|
||||
global_norm_fp16 = paddle.to_tensor(
|
||||
np.array(0.0), dtype=paddle.float32
|
||||
)
|
||||
else:
|
||||
global_norm_fp16 = paddle.add_n(sum_square_fp16)
|
||||
global_norm_fp16 = paddle.cast(
|
||||
global_norm_fp16, dtype=paddle.float32
|
||||
)
|
||||
|
||||
# global norm of non-distributed BFP16 params_and_grads
|
||||
if len(sum_square_bfp16) == 0:
|
||||
global_norm_bfp16 = paddle.to_tensor(
|
||||
np.array(0.0), dtype=paddle.float32
|
||||
)
|
||||
else:
|
||||
global_norm_bfp16 = paddle.add_n(sum_square_bfp16)
|
||||
global_norm_bfp16 = paddle.cast(
|
||||
global_norm_bfp16, dtype=paddle.float32
|
||||
)
|
||||
|
||||
# global norm of non-distributed FP16 params_and_grads for unslice parameters
|
||||
if len(unslice_params_fp16) == 0:
|
||||
global_unslice_fp16 = paddle.to_tensor(
|
||||
np.array(0.0), dtype=paddle.float32
|
||||
)
|
||||
else:
|
||||
global_unslice_fp16 = paddle.add_n(unslice_params_fp16)
|
||||
global_unslice_fp16 = paddle.cast(
|
||||
global_unslice_fp16, dtype=paddle.float32
|
||||
)
|
||||
|
||||
# global norm of non-distributed BFP16 params_and_grads for unslice parameters
|
||||
if len(unslice_params_bfp16) == 0:
|
||||
global_unslice_bfp16 = paddle.to_tensor(
|
||||
np.array(0.0), dtype=paddle.float32
|
||||
)
|
||||
else:
|
||||
global_unslice_bfp16 = paddle.add_n(unslice_params_bfp16)
|
||||
global_unslice_bfp16 = paddle.cast(
|
||||
global_unslice_bfp16, dtype=paddle.float32
|
||||
)
|
||||
|
||||
# global norm of non-distributed FP32 params_and_grads
|
||||
if len(sum_square_fp32) == 0:
|
||||
global_norm_fp32 = paddle.to_tensor(
|
||||
np.array(0.0), dtype=paddle.float32
|
||||
)
|
||||
else:
|
||||
global_norm_fp32 = paddle.add_n(sum_square_fp32)
|
||||
|
||||
# global norm of non-distributed FP32 params_and_grads for unslice parameters
|
||||
if len(unslice_params_fp32) == 0:
|
||||
global_unslice_fp32 = paddle.to_tensor(
|
||||
np.array(0.0), dtype=paddle.float32
|
||||
)
|
||||
else:
|
||||
global_unslice_fp32 = paddle.add_n(unslice_params_fp32)
|
||||
|
||||
global_unslice_var = (
|
||||
global_unslice_fp16 + global_unslice_fp32 + global_unslice_bfp16
|
||||
)
|
||||
|
||||
global_norm_var = (
|
||||
global_norm_fp16 + global_norm_fp32 + global_norm_bfp16
|
||||
)
|
||||
|
||||
# add all reduce to get global norm of distributed params_and_grads
|
||||
dev_id = int(self._device.split(":")[1])
|
||||
dev_type = self._device.split(':')[0]
|
||||
if paddle.device.get_device() == "cpu":
|
||||
if dev_type in paddle.device.get_all_custom_device_type():
|
||||
global_norm_var = global_norm_var._copy_to(
|
||||
paddle.CustomPlace(dev_type, dev_id), True
|
||||
)
|
||||
elif dev_type == "xpu":
|
||||
global_norm_var = global_norm_var.to(self._device)
|
||||
else:
|
||||
global_norm_var = global_norm_var.cuda(dev_id)
|
||||
|
||||
with device_guard(dev_id, self._device.split(":")[0]):
|
||||
paddle.distributed.all_reduce(global_norm_var, group=self._group)
|
||||
|
||||
global_norm_var = paddle.sqrt(global_norm_var + global_unslice_var)
|
||||
max_global_norm = paddle.full(
|
||||
shape=[], dtype=global_norm_var.dtype, fill_value=self.clip_norm
|
||||
)
|
||||
|
||||
clip_var = paddle.divide(
|
||||
x=max_global_norm,
|
||||
y=paddle.maximum(x=global_norm_var, y=max_global_norm),
|
||||
)
|
||||
clip_var_fp16 = paddle.cast(clip_var, paddle.float16)
|
||||
|
||||
for p, g in params_grads:
|
||||
if getattr(p, 'need_clip', True) is False or g is None:
|
||||
continue
|
||||
origin_state = g.stop_gradient
|
||||
g.stop_gradient = True
|
||||
if p.dtype == paddle.float16:
|
||||
g.scale_(clip_var_fp16)
|
||||
else:
|
||||
g.scale_(clip_var)
|
||||
g.stop_gradient = origin_state
|
||||
# p._reset_grad_inplace_version(True)
|
||||
|
||||
return params_grads
|
||||
|
||||
def __getattr__(self, item):
|
||||
return getattr(self._clip, item)
|
||||
|
||||
def __call__(self, params_grads):
|
||||
return self._dygraph_clip(params_grads)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def device_guard(dev_id=0, device="cpu"):
|
||||
origin_device = paddle.device.get_device()
|
||||
if device == "cpu":
|
||||
paddle.set_device(device)
|
||||
elif device in ["gpu", "xpu"]:
|
||||
paddle.set_device(f"{device}:{dev_id}")
|
||||
elif device in paddle.device.get_all_custom_device_type():
|
||||
paddle.set_device(f"{device}:{dev_id}")
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
paddle.set_device(origin_device)
|
||||
|
||||
|
||||
@dygraph_only
|
||||
def GroupShardedScaler(scaler):
|
||||
def unscale_method(self, optimizer):
|
||||
if not self._enable:
|
||||
return
|
||||
param_grads = []
|
||||
param_grads_bfp16 = []
|
||||
param_grads_fp16 = []
|
||||
param_grads_fp32 = []
|
||||
if hasattr(optimizer, "update_slice"):
|
||||
optimizer.update_slice()
|
||||
optimizer.update_scaler = True
|
||||
|
||||
if getattr(optimizer._optim, '_param_groups', None) and isinstance(
|
||||
optimizer._optim._param_groups[0], dict
|
||||
):
|
||||
for group in optimizer._optim._param_groups:
|
||||
for param in group['params']:
|
||||
tgt_grad = None
|
||||
if (
|
||||
hasattr(param, "main_grad")
|
||||
and param.main_grad is not None
|
||||
):
|
||||
tgt_grad = param.main_grad
|
||||
elif param.grad is not None:
|
||||
tgt_grad = param.grad
|
||||
if tgt_grad is not None:
|
||||
param_grads.append(tgt_grad)
|
||||
if tgt_grad.dtype in [
|
||||
core.VarDesc.VarType.FP16,
|
||||
paddle.float16,
|
||||
]:
|
||||
param_grads_fp16.append(tgt_grad)
|
||||
elif tgt_grad.dtype in [paddle.bfloat16]:
|
||||
param_grads_bfp16.append(tgt_grad)
|
||||
else:
|
||||
param_grads_fp32.append(tgt_grad)
|
||||
else:
|
||||
for param in optimizer._optim._parameter_list:
|
||||
tgt_grad = None
|
||||
if hasattr(param, "main_grad") and param.main_grad is not None:
|
||||
tgt_grad = param.main_grad
|
||||
elif param.grad is not None:
|
||||
tgt_grad = param.grad
|
||||
if tgt_grad is not None:
|
||||
param_grads.append(tgt_grad)
|
||||
if tgt_grad.dtype in [
|
||||
core.VarDesc.VarType.FP16,
|
||||
paddle.float16,
|
||||
]:
|
||||
param_grads_fp16.append(tgt_grad)
|
||||
elif tgt_grad.dtype in [paddle.bfloat16]:
|
||||
param_grads_bfp16.append(tgt_grad)
|
||||
else:
|
||||
param_grads_fp32.append(tgt_grad)
|
||||
|
||||
temp_found_inf_fp16 = paddle.to_tensor(np.array([0]).astype(np.bool_))
|
||||
temp_found_inf_bfp16 = paddle.to_tensor(np.array([0]).astype(np.bool_))
|
||||
temp_found_inf_fp32 = paddle.to_tensor(np.array([0]).astype(np.bool_))
|
||||
|
||||
device = paddle.get_device().split(":")[0]
|
||||
device = "cpu" if optimizer.offload else device
|
||||
dev_id = (
|
||||
0 if device == "cpu" else int(paddle.get_device().split(":")[1])
|
||||
)
|
||||
|
||||
self._found_inf = self._temp_found_inf_value_false
|
||||
with device_guard(dev_id, device):
|
||||
if len(param_grads_bfp16):
|
||||
_legacy_C_ops.check_finite_and_unscale(
|
||||
param_grads_bfp16,
|
||||
self._scale,
|
||||
param_grads_bfp16,
|
||||
temp_found_inf_bfp16,
|
||||
)
|
||||
self._found_inf = _C_ops.bitwise_or(
|
||||
self._found_inf, temp_found_inf_bfp16
|
||||
)
|
||||
if len(param_grads_fp16):
|
||||
_legacy_C_ops.check_finite_and_unscale(
|
||||
param_grads_fp16,
|
||||
self._scale,
|
||||
param_grads_fp16,
|
||||
temp_found_inf_fp16,
|
||||
)
|
||||
self._found_inf = _C_ops.bitwise_or(
|
||||
self._found_inf, temp_found_inf_fp16
|
||||
)
|
||||
if len(param_grads_fp32):
|
||||
_legacy_C_ops.check_finite_and_unscale(
|
||||
param_grads_fp32,
|
||||
self._scale,
|
||||
param_grads_fp32,
|
||||
temp_found_inf_fp32,
|
||||
)
|
||||
self._found_inf = _C_ops.bitwise_or(
|
||||
self._found_inf, temp_found_inf_fp32
|
||||
)
|
||||
|
||||
self._found_inf = self._found_inf.cast("int32")
|
||||
|
||||
paddle.distributed.all_reduce(
|
||||
self._found_inf, op=paddle.distributed.ReduceOp.MAX, group=None
|
||||
)
|
||||
|
||||
self._found_inf = self._found_inf.cast("bool")
|
||||
|
||||
scaler._unscale = MethodType(unscale_method, scaler)
|
||||
return scaler
|
||||
|
||||
|
||||
def cvt_to_device(x, dev_id, blocking=True):
|
||||
"""
|
||||
Copy data in x from cpu memory to supported device
|
||||
"""
|
||||
if paddle.is_compiled_with_cuda():
|
||||
place = paddle.CUDAPlace(dev_id)
|
||||
elif paddle.is_compiled_with_xpu():
|
||||
place = paddle.XPUPlace(dev_id)
|
||||
else:
|
||||
supported_custom_devices = ["npu"]
|
||||
place = paddle.framework._current_expected_place()
|
||||
if place.get_device_type() not in supported_custom_devices:
|
||||
raise OSError(
|
||||
"Only supported compiled paddle with gpu/rocm and xpu, but current version is compiled with cpu."
|
||||
)
|
||||
return x._copy_to(place, blocking)
|
||||
Reference in New Issue
Block a user