wehub-resource-sync
261 lines
9.3 KiB
Python
261 lines
9.3 KiB
Python
# Copyright (c) 2023 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.
|
|
|
|
|
|
from collections import defaultdict
|
|
from types import MethodType
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle import _legacy_C_ops, nn
|
|
from paddle.base import framework
|
|
from paddle.base.dygraph import (
|
|
base as imperative_base,
|
|
)
|
|
from paddle.distributed import fleet
|
|
from paddle.distributed.fleet.utils.hybrid_parallel_util import (
|
|
obtain_optimizer_parameters_list,
|
|
)
|
|
from paddle.framework import core
|
|
from paddle.utils import deprecated
|
|
|
|
|
|
class MixPrecisionLayer(nn.Layer):
|
|
def __init__(self, layers, dtype="float16"):
|
|
super().__init__(layers.full_name() + "_mix_precision")
|
|
|
|
self._layers = layers
|
|
self._dtype = dtype
|
|
|
|
assert self._dtype in ["float16", "bfloat16"]
|
|
|
|
for param in self._layers.parameters():
|
|
if not hasattr(param, "main_grad"):
|
|
param.main_grad = None
|
|
param._register_grad_hook(self._update_main_grad_hook(param))
|
|
|
|
def _update_main_grad_hook(self, param):
|
|
"""Create the update_main_grad hook for back-prop."""
|
|
|
|
# Hook used for back-prop and grad-merge.
|
|
@paddle.autograd.no_grad()
|
|
def param_hook(tmp_grad):
|
|
assert param.grad is None, (
|
|
f"In main_grad node, param.grad should be None, but find param[{param.name}] has grad."
|
|
)
|
|
if tmp_grad is not None and tmp_grad._is_initialized():
|
|
# Some previous pylayer may return None, should check grad validation.
|
|
if param.main_grad is None:
|
|
param.main_grad = core.eager.Tensor(
|
|
value=tmp_grad.cast(paddle.float32).value(),
|
|
place=tmp_grad.place,
|
|
name="main_grad@" + param.name,
|
|
)
|
|
else:
|
|
param.main_grad.add_(tmp_grad)
|
|
|
|
tmp_grad._clear_data()
|
|
|
|
return param_hook
|
|
|
|
def forward(self, *inputs, **kwargs):
|
|
outputs = self._layers(*inputs, **kwargs)
|
|
|
|
return outputs
|
|
|
|
def state_dict(
|
|
self,
|
|
destination=None,
|
|
include_sublayers=True,
|
|
structured_name_prefix="",
|
|
):
|
|
return self._layers.state_dict(
|
|
destination=destination,
|
|
include_sublayers=include_sublayers,
|
|
structured_name_prefix=structured_name_prefix,
|
|
)
|
|
|
|
@framework.deprecate_stat_dict
|
|
def set_state_dict(self, state_dict, use_structured_name=True):
|
|
self._layers.set_state_dict(
|
|
state_dict, use_structured_name=use_structured_name
|
|
)
|
|
|
|
|
|
class MixPrecisionOptimizer:
|
|
def __init__(self, optimizer):
|
|
self._inner_opt = optimizer
|
|
self._parameter_list = obtain_optimizer_parameters_list(optimizer)
|
|
|
|
@imperative_base.no_grad
|
|
@framework.dygraph_only
|
|
def step(self):
|
|
need_shard = any(
|
|
hasattr(p, '_need_shard') for p in self._parameter_list
|
|
)
|
|
if need_shard:
|
|
fleet.meta_parallel.sharding.group_sharded_fully_shard.FullyShardOptimizer(
|
|
self
|
|
)
|
|
self.step()
|
|
return
|
|
|
|
if not isinstance(self._parameter_list[0], dict):
|
|
params_grads = []
|
|
for param in self._parameter_list:
|
|
if param.stop_gradient:
|
|
continue
|
|
grad_var = param.main_grad
|
|
if grad_var is None:
|
|
continue
|
|
if paddle.in_dynamic_mode():
|
|
if (
|
|
hasattr(grad_var, "is_selected_rows")
|
|
and grad_var.is_selected_rows()
|
|
and self._inner_opt.regularization is not None
|
|
):
|
|
raise RuntimeError(
|
|
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
|
)
|
|
else:
|
|
if (
|
|
hasattr(grad_var, "_is_sparse")
|
|
and grad_var._is_sparse()
|
|
and self._inner_opt.regularization is not None
|
|
):
|
|
raise RuntimeError(
|
|
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
|
)
|
|
params_grads.append((param, grad_var))
|
|
|
|
optimize_ops = self._inner_opt._apply_optimize(
|
|
loss=None, startup_program=None, params_grads=params_grads
|
|
)
|
|
else:
|
|
# optimize parameters in groups
|
|
for param_group in self._inner_opt._param_groups:
|
|
params_grads = defaultdict(lambda: [])
|
|
for param in param_group['params']:
|
|
if param.stop_gradient:
|
|
continue
|
|
grad_var = param.main_grad
|
|
if grad_var is None:
|
|
continue
|
|
if paddle.in_dynamic_mode():
|
|
if (
|
|
hasattr(grad_var, "is_selected_rows")
|
|
and grad_var.is_selected_rows()
|
|
and self._inner_opt.regularization is not None
|
|
):
|
|
raise RuntimeError(
|
|
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
|
)
|
|
else:
|
|
if (
|
|
hasattr(grad_var, "_is_sparse")
|
|
and grad_var._is_sparse()
|
|
and self._inner_opt.regularization is not None
|
|
):
|
|
raise RuntimeError(
|
|
"AdamW don't support weight_decay with sparse parameters, please set it to None."
|
|
)
|
|
params_grads['params'].append((param, grad_var))
|
|
params_grads.update(
|
|
{k: v for k, v in param_group.items() if k != 'params'}
|
|
)
|
|
self._apply_optimize(
|
|
loss=None, startup_program=None, params_grads=params_grads
|
|
)
|
|
|
|
@framework.dygraph_only
|
|
def clear_grad(self, set_to_zero=True):
|
|
param_list = []
|
|
if self._parameter_list is None or not isinstance(
|
|
self._parameter_list[0], dict
|
|
):
|
|
for p in self._parameter_list:
|
|
if not p.stop_gradient:
|
|
param_list.append(p)
|
|
else:
|
|
for param_group in self._param_groups:
|
|
for p in param_group['params']:
|
|
if not p.stop_gradient:
|
|
param_list.append(p)
|
|
|
|
for p in param_list:
|
|
if hasattr(p, "main_grad") and p.main_grad is not None:
|
|
if set_to_zero:
|
|
p.main_grad.zero_()
|
|
else:
|
|
p.main_grad._clear()
|
|
p.main_grad = None
|
|
elif not hasattr(p, "main_grad"):
|
|
p.clear_gradient(set_to_zero)
|
|
|
|
def __getattr__(self, item):
|
|
return getattr(self._inner_opt, item)
|
|
|
|
|
|
def unscale_method(self, optimizer):
|
|
if not self._enable:
|
|
return
|
|
param_grads = []
|
|
if getattr(optimizer, '_param_groups', None) and isinstance(
|
|
optimizer._param_groups[0], dict
|
|
):
|
|
for group in optimizer._param_groups:
|
|
for param in group['params']:
|
|
if param.main_grad is not None:
|
|
assert param.main_grad.dtype == paddle.float32
|
|
param_grads.append(param.main_grad)
|
|
else:
|
|
for param in optimizer._parameter_list:
|
|
if param.main_grad is not None:
|
|
assert param.main_grad.dtype == paddle.float32
|
|
param_grads.append(param.main_grad)
|
|
|
|
temp_found_inf = paddle.to_tensor(np.array([0]).astype(np.bool_))
|
|
if len(param_grads):
|
|
_legacy_C_ops.check_finite_and_unscale(
|
|
param_grads,
|
|
self._scale,
|
|
param_grads,
|
|
temp_found_inf,
|
|
)
|
|
|
|
self._found_inf = 1 if temp_found_inf else 0
|
|
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
if hcg is not None and hcg.nranks > hcg.get_data_parallel_world_size():
|
|
is_found_inf = paddle.to_tensor([self._found_inf], dtype="int32")
|
|
paddle.distributed.all_reduce(
|
|
is_found_inf, op=paddle.distributed.ReduceOp.MAX, group=None
|
|
)
|
|
self._found_inf = int(is_found_inf)
|
|
|
|
|
|
@deprecated(
|
|
since="2.5.0",
|
|
update_to="paddle.distributed_scaler",
|
|
level=1,
|
|
)
|
|
class MixPrecisionScaler:
|
|
def __init__(self, scaler):
|
|
self._inner_scaler = scaler
|
|
self._inner_scaler._unscale = MethodType(unscale_method, scaler)
|
|
|
|
def __getattr__(self, item):
|
|
return getattr(self._inner_scaler, item)
|