wehub-resource-sync
128 lines
4.4 KiB
Python
128 lines
4.4 KiB
Python
# Copyright 2020-present the HuggingFace Inc. team.
|
|
# 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 types
|
|
|
|
import numpy as np
|
|
import paddle
|
|
from paddle.common_ops_import import LayerHelper
|
|
|
|
from ...utils.log import logger
|
|
|
|
|
|
def npu_accelerate_plugin(optimizer):
|
|
"""npu_accelerate_plugin uses the flatten_param_grads method to speed up the performance of the model on NPU devices.
|
|
flatten_param_grads method will be added to `step` function of optimizer.
|
|
|
|
Args:
|
|
optimizer (`paddle.optimizer.Optimizer`):
|
|
The Optimizer whose `step` method will be modified.
|
|
"""
|
|
optimizer.step = types.MethodType(_optimizer_step_with_flatten_param_grads, optimizer)
|
|
|
|
|
|
def _optimizer_step_with_flatten_param_grads(optimizer):
|
|
if not isinstance(optimizer._param_groups[0], dict):
|
|
params_grads = []
|
|
for param in optimizer._param_groups:
|
|
if param.stop_gradient:
|
|
continue
|
|
if param._grad_ivar() is not None:
|
|
grad_var = param._grad_ivar()
|
|
params_grads.append((param, grad_var))
|
|
|
|
# currently, only support ClipGradByGlobalNorm and without regularization.
|
|
if isinstance(params_grads, list) and optimizer.regularization is None:
|
|
if optimizer._grad_clip is None or isinstance(optimizer._grad_clip, paddle.nn.ClipGradByGlobalNorm):
|
|
params_grads = _flatten_param_grads(optimizer, params_grads)
|
|
|
|
optimizer._apply_optimize(
|
|
loss=None,
|
|
startup_program=None,
|
|
params_grads=params_grads,
|
|
param_group_idx=0,
|
|
)
|
|
else:
|
|
raise RuntimeError("flatten_param_grads is not supported when _param_groups[0] is dict.")
|
|
|
|
|
|
def _flatten_param_grads(optimizer, params_grads):
|
|
optimizer.helper = LayerHelper(optimizer.__class__.__name__)
|
|
need_flatten_params = []
|
|
need_flatten_grads = []
|
|
for p, g in params_grads:
|
|
if g is None:
|
|
continue
|
|
g.persistable = True
|
|
if getattr(p, "need_clip", True) is False or getattr(p, "regularizer", None) is not None:
|
|
logger.warning(
|
|
f"flatten_param_grads=True will be discarded since parameter {p.name}'s need_clip is False or "
|
|
"the regularizer is set."
|
|
)
|
|
return params_grads
|
|
|
|
need_flatten_params.append(p)
|
|
need_flatten_grads.append(g)
|
|
|
|
shape = [np.prod(p.shape) for p in need_flatten_params]
|
|
|
|
flatten_param = optimizer.helper.create_global_variable(
|
|
name="flatten_param",
|
|
persistable=True,
|
|
dtype=need_flatten_params[0].dtype,
|
|
shape=[np.sum(shape)],
|
|
belong_to_optimizer=True,
|
|
)
|
|
|
|
flatten_grad = optimizer.helper.create_global_variable(
|
|
name="flatten_grad",
|
|
persistable=True,
|
|
dtype=need_flatten_grads[0].dtype,
|
|
shape=[np.sum(shape)],
|
|
belong_to_optimizer=True,
|
|
)
|
|
|
|
flatten_param.stop_gradient = False
|
|
# In the final state of the dynamic graph, the `coalesce_tensor` op
|
|
# does not support passing the output as an input into the op in
|
|
# temporary, so _legacy_C_ops is temporarily used here.
|
|
# `use_align` is set to false, which is different from the behavior
|
|
# under static graphs. `use_align` can be set to true after calling
|
|
# the coalesce_tensor op of the final state (_C_ops).
|
|
paddle._legacy_C_ops.coalesce_tensor(
|
|
need_flatten_params,
|
|
need_flatten_params,
|
|
flatten_param,
|
|
"copy_data",
|
|
True,
|
|
"use_align",
|
|
False,
|
|
"dtype",
|
|
need_flatten_params[0].dtype,
|
|
)
|
|
|
|
paddle._legacy_C_ops.coalesce_tensor(
|
|
need_flatten_grads,
|
|
need_flatten_grads,
|
|
flatten_grad,
|
|
"copy_data",
|
|
True,
|
|
"use_align",
|
|
False,
|
|
"dtype",
|
|
need_flatten_grads[0].dtype,
|
|
)
|
|
return [(flatten_param, flatten_grad)]
|