wehub-resource-sync
1035 lines
43 KiB
Python
1035 lines
43 KiB
Python
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import types
|
|
import warnings
|
|
|
|
import paddle
|
|
from paddle.base import (
|
|
core,
|
|
default_main_program,
|
|
default_startup_program,
|
|
program_guard,
|
|
unique_name,
|
|
)
|
|
from paddle.base.framework import auto_complete_op_role, in_pir_mode
|
|
|
|
from .amp_nn import check_finite_and_unscale, update_loss_scaling
|
|
from .fp16_lists import AutoMixedPrecisionLists, check_amp_dtype
|
|
from .fp16_utils import (
|
|
cast_model_to_fp16,
|
|
cast_parameters_to_fp16,
|
|
update_role_var_grad,
|
|
)
|
|
from .function_overload import FunctionType, overload
|
|
|
|
OpRole = core.op_proto_and_checker_maker.OpRole
|
|
|
|
|
|
def _set_multi_precision(optimizer, multi_precision):
|
|
if not isinstance(
|
|
optimizer,
|
|
(paddle.optimizer.Optimizer),
|
|
):
|
|
raise RuntimeError(
|
|
f"Current AMP training level is O2, optimizer is expected to be paddle.optimizer.Optimizer, but receive {type(optimizer)}."
|
|
)
|
|
|
|
if multi_precision and hasattr(optimizer, "_multi_precision"):
|
|
optimizer._multi_precision = multi_precision
|
|
|
|
|
|
class OptimizerWithMixedPrecision:
|
|
"""
|
|
Optimizer with mixed-precision (MP) training. This is a wrapper of a common
|
|
optimizer, plus the support of mixed-precision pre-training. The object
|
|
of this class almost has the same behavior as the common optimizer, with the
|
|
methods `minimize()`, `backward()`, `apply_gradients()` implemented.
|
|
Additionally, it enables the MP training automatically, i.e, the creation
|
|
and maintenance of master parameters, scaling of loss, etc.
|
|
|
|
Args:
|
|
optimizer (Optimizer): A common Optimizer object.
|
|
amp_lists (AutoMixedPrecisionLists): An AutoMixedPrecisionLists object.
|
|
level(str): Auto mixed precision level. Accepted values are "O1", "O2" and "OD": At the O1 level, operators in the white list
|
|
will use float16/bfloat16 inputs for calculations, and operators in the black list will use float32 inputs for calculations. At the O2
|
|
level, model's parameters will be casted to float16/bfloat16 by using `decorator`, and operators that have all float16/bfloat16 inputs
|
|
will be converted to float16/bfloat16, and that have any float32 input will be converted to float32. For the OD level, operators in
|
|
default white list will compute in float16/bfloat16.
|
|
dtype(str): Whether to use 'float16' or 'bfloat16'.
|
|
init_loss_scaling (float): The initial loss scaling factor.
|
|
use_dynamic_loss_scaling (bool): Whether to use dynamic loss scaling.
|
|
incr_every_n_steps(int): Increases loss scaling every n consecutive
|
|
steps with finite gradients.
|
|
decr_every_n_nan_or_inf(int): Decreases loss scaling every n
|
|
accumulated steps with nan or
|
|
inf gradients.
|
|
incr_ratio(float): The multiplier to use when increasing the loss
|
|
scaling.
|
|
decr_ratio(float): The less-than-one-multiplier to use when decreasing
|
|
the loss scaling.
|
|
use_amp_guard(bool): Whether to use `fp16_guard` when constructing the program.
|
|
Default None, which means that its value is equal to `use_pure_fp16`.
|
|
use_master_grad(bool): Whether to use fp32 master gradients during optimizer. Default is False.
|
|
use_promote(bool): Whether to promotes to fp32 when op has any float32 inputs. Default is False.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer,
|
|
amp_lists,
|
|
level,
|
|
dtype,
|
|
init_loss_scaling,
|
|
use_dynamic_loss_scaling,
|
|
incr_every_n_steps,
|
|
decr_every_n_nan_or_inf,
|
|
incr_ratio,
|
|
decr_ratio,
|
|
use_amp_guard=None,
|
|
use_master_grad=False,
|
|
use_promote=False,
|
|
):
|
|
self._optimizer = optimizer
|
|
self._amp_lists = amp_lists
|
|
self._param_grads = None
|
|
self._train_program = None
|
|
|
|
self._is_distributed = False
|
|
self._use_master_grad = False
|
|
self._scaled_loss = None
|
|
self._loss_scaling = None
|
|
self._init_loss_scaling = init_loss_scaling
|
|
self._use_dynamic_loss_scaling = use_dynamic_loss_scaling
|
|
if dtype == "bfloat16":
|
|
if use_dynamic_loss_scaling:
|
|
self._use_dynamic_loss_scaling = False
|
|
self._init_loss_scaling = 1.0
|
|
warnings.warn(
|
|
"Dynamic loss scaling for bfloat16 amp training is disabled, and the init_loss_scaling is changed to 1.0 automatically by PaddlePaddle."
|
|
)
|
|
if in_pir_mode():
|
|
self._amp_vartype = core.DataType.BFLOAT16
|
|
else:
|
|
self._amp_vartype = core.VarDesc.VarType.BF16
|
|
else:
|
|
if in_pir_mode():
|
|
self._amp_vartype = core.DataType.FLOAT16
|
|
else:
|
|
self._amp_vartype = core.VarDesc.VarType.FP16
|
|
|
|
self._learning_rate = optimizer._learning_rate
|
|
self._learning_rate_map = optimizer._learning_rate_map
|
|
self._use_pure_fp16 = level == "O2"
|
|
if self._use_pure_fp16 and (dtype == "bfloat16" or dtype == "float16"):
|
|
self._use_master_grad = use_master_grad
|
|
self._optimizer._master_grad = use_master_grad
|
|
self._amp_level = level
|
|
self._use_fp16_guard = use_amp_guard
|
|
self._to_fp16_var_names = None
|
|
if self._use_dynamic_loss_scaling:
|
|
self._incr_every_n_steps = incr_every_n_steps
|
|
self._decr_every_n_nan_or_inf = decr_every_n_nan_or_inf
|
|
self._incr_ratio = incr_ratio
|
|
self._decr_ratio = decr_ratio
|
|
self._num_good_steps = None
|
|
self._num_bad_steps = None
|
|
self.use_promote = use_promote
|
|
|
|
def _set_distributed(self, flag):
|
|
# if distributed, all cards will communication with each other,
|
|
# overlap communication and computation by split the
|
|
# check_finite_and_unscale op.
|
|
self._is_distributed = flag
|
|
|
|
def get_loss_scaling(self):
|
|
"""Return the real-time loss scaling factor."""
|
|
assert self._loss_scaling is not None, (
|
|
'Please call minimize() before calling get_loss_scaling().'
|
|
)
|
|
return self._loss_scaling
|
|
|
|
def get_scaled_loss(self):
|
|
"""Return the scaled loss.
|
|
It's useful when you feed customed loss into executor.
|
|
"""
|
|
return self._scaled_loss
|
|
|
|
def _supports_check_nan_inf(self):
|
|
return getattr(self._optimizer, "_supports_check_nan_inf", False)
|
|
|
|
def _init_amp_var(self):
|
|
if in_pir_mode():
|
|
if self._use_dynamic_loss_scaling:
|
|
self._num_good_steps = paddle.pir.core.create_persistable_value(
|
|
dtype='int32',
|
|
shape=[1],
|
|
name=unique_name.generate("num_good_steps"),
|
|
initializer=paddle.nn.initializer.ConstantInitializer(
|
|
value=0
|
|
),
|
|
)
|
|
self._num_bad_steps = paddle.pir.core.create_persistable_value(
|
|
dtype='int32',
|
|
shape=[1],
|
|
name=unique_name.generate("num_bad_steps"),
|
|
initializer=paddle.nn.initializer.ConstantInitializer(
|
|
value=0
|
|
),
|
|
)
|
|
|
|
if isinstance(self._optimizer._learning_rate, float):
|
|
self._optimizer._learning_rate_map[
|
|
paddle.static.default_main_program()
|
|
] = paddle.pir.core.create_persistable_value(
|
|
dtype=self._optimizer.get_lr_dtype(),
|
|
shape=[1],
|
|
name=unique_name.generate("learning_rate"),
|
|
initializer=paddle.nn.initializer.ConstantInitializer(
|
|
value=float(self._optimizer._learning_rate)
|
|
),
|
|
)
|
|
|
|
return
|
|
|
|
self._loss_scaling = paddle.static.create_global_var(
|
|
name=unique_name.generate("loss_scaling"),
|
|
shape=[1],
|
|
value=self._init_loss_scaling,
|
|
dtype='float32',
|
|
persistable=True,
|
|
)
|
|
|
|
if self._use_dynamic_loss_scaling:
|
|
self._num_good_steps = paddle.static.create_global_var(
|
|
name=unique_name.generate("num_good_steps"),
|
|
shape=[1],
|
|
value=0,
|
|
dtype='int32',
|
|
persistable=True,
|
|
)
|
|
self._num_bad_steps = paddle.static.create_global_var(
|
|
name=unique_name.generate("num_bad_steps"),
|
|
shape=[1],
|
|
value=0,
|
|
dtype='int32',
|
|
persistable=True,
|
|
)
|
|
|
|
# Ensure the data type of learning rate vars matches the optimizer's
|
|
# preferred dtype (e.g. float64 for AdamW, float32 for others).
|
|
if isinstance(self._optimizer._learning_rate, float):
|
|
_lr_dtype = self._optimizer.get_lr_dtype()
|
|
self._optimizer._learning_rate_map[default_main_program()] = (
|
|
paddle.static.create_global_var(
|
|
name=unique_name.generate("learning_rate"),
|
|
shape=[1],
|
|
value=float(self._optimizer._learning_rate),
|
|
dtype=_lr_dtype,
|
|
persistable=True,
|
|
)
|
|
)
|
|
|
|
def backward(
|
|
self,
|
|
loss,
|
|
startup_program=None,
|
|
parameter_list=None,
|
|
no_grad_set=None,
|
|
callbacks=None,
|
|
):
|
|
"""
|
|
Backward propagation or auto differentiation for gradients' computation.
|
|
|
|
Args:
|
|
loss (Variable): The loss Variable to minimize.
|
|
startup_program (Program|None): The startup Program for initializing
|
|
parameters in `parameter_list`.
|
|
parameter_list (list|None): A list of Variables to update.
|
|
no_grad_set (set|None): A set of Variables should be ignored.
|
|
callbacks (list|None): A list of callable objects to run when appending
|
|
backward operator for one parameter.
|
|
|
|
Returns:
|
|
A list of (param, grad), which is a tuple of a parameter and its
|
|
gradient respectively, and the scaled loss.
|
|
"""
|
|
train_program = loss.block.program
|
|
self._train_program = train_program
|
|
self._float_status = None
|
|
|
|
if in_pir_mode():
|
|
with paddle.static.program_guard(
|
|
self._train_program, startup_program
|
|
):
|
|
self._init_amp_var()
|
|
if self._scaled_loss is None:
|
|
self._scaled_loss = loss
|
|
params_grads = self._optimizer.backward(
|
|
self._scaled_loss,
|
|
startup_program,
|
|
parameter_list,
|
|
no_grad_set,
|
|
callbacks,
|
|
)
|
|
return params_grads
|
|
|
|
with program_guard(self._train_program, startup_program):
|
|
self._init_amp_var()
|
|
|
|
if self._use_pure_fp16:
|
|
self._to_fp16_var_names = cast_model_to_fp16(
|
|
self._train_program,
|
|
self._amp_lists,
|
|
self._use_fp16_guard,
|
|
self._amp_vartype,
|
|
level='O2',
|
|
use_promote=self.use_promote,
|
|
)
|
|
else:
|
|
# use_fp16_guard is not support amp-o1.
|
|
cast_model_to_fp16(
|
|
self._train_program,
|
|
self._amp_lists,
|
|
use_fp16_guard=False,
|
|
dest_type=self._amp_vartype,
|
|
level=self._amp_level,
|
|
use_promote=self.use_promote,
|
|
)
|
|
|
|
if loss.dtype != core.VarDesc.VarType.FP32:
|
|
loss = loss.astype('float32')
|
|
# When not using dynamic loss scaling and the init loss scaling value is equal to 1.0,
|
|
# the model can be optimized.
|
|
if self._use_dynamic_loss_scaling or self._init_loss_scaling != 1.0:
|
|
self._scaled_loss = loss * self._loss_scaling
|
|
else:
|
|
self._scaled_loss = loss
|
|
|
|
params_grads = self._optimizer.backward(
|
|
self._scaled_loss,
|
|
startup_program,
|
|
parameter_list,
|
|
no_grad_set,
|
|
callbacks,
|
|
)
|
|
if self._supports_check_nan_inf():
|
|
self._add_cast_ops_to_startup_program(startup_program)
|
|
return params_grads
|
|
|
|
def _add_cast_ops_to_startup_program(self, startup_program):
|
|
names = list(self._to_fp16_var_names) if self._to_fp16_var_names else []
|
|
names.sort()
|
|
startup_program = (
|
|
default_startup_program()
|
|
if startup_program is None
|
|
else startup_program
|
|
)
|
|
block = startup_program.global_block()
|
|
param_names = [p.name for p in block.all_parameters()]
|
|
for name in names:
|
|
if name not in param_names:
|
|
continue
|
|
|
|
tmp = block.create_var(dtype=core.VarDesc.VarType.FP32)
|
|
block.append_op(
|
|
type='assign', inputs={'X': [name]}, outputs={'Out': [tmp]}
|
|
)
|
|
block.append_op(
|
|
type='cast',
|
|
inputs={'X': [tmp]},
|
|
outputs={'Out': [name]},
|
|
attrs={
|
|
'in_dtype': core.VarDesc.VarType.FP32,
|
|
'out_dtype': self._amp_vartype,
|
|
},
|
|
)
|
|
self._to_fp16_var_names = None
|
|
|
|
def amp_init(
|
|
self,
|
|
place,
|
|
scope=None,
|
|
test_program=None,
|
|
use_fp16_test=False,
|
|
rewrite_master_weight=False,
|
|
):
|
|
"""
|
|
Init the amp training, such as cast fp32 parameters to fp16 type.
|
|
|
|
Args:
|
|
place(CUDAPlace): place is used to initialize
|
|
fp16 parameters with fp32 values.
|
|
scope(Scope): The scope is used to find fp32 parameters.
|
|
test_program(Program): The program is used for testing.
|
|
use_fp16_test(bool): Whether to use fp16 testing.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
>>> import paddle.nn.functional as F
|
|
>>> paddle.enable_static()
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> def run_example_code():
|
|
... place = paddle.CUDAPlace(0)
|
|
... exe = paddle.static.Executor(place)
|
|
... data = paddle.static.data(name='X', shape=[None, 1, 28, 28], dtype='float32')
|
|
... conv2d = paddle.static.nn.conv2d(input=data, num_filters=6, filter_size=3)
|
|
... # 1) Use fp16_guard to control the range of fp16 kernels used.
|
|
... with paddle.static.amp.fp16_guard():
|
|
... bn = paddle.static.nn.batch_norm(input=conv2d, act="relu")
|
|
... pool = F.max_pool2d(bn, kernel_size=2, stride=2)
|
|
... hidden = paddle.static.nn.fc(pool, size=10)
|
|
... loss = paddle.mean(hidden)
|
|
... # 2) Create the optimizer and set `multi_precision` to True.
|
|
... # Setting `multi_precision` to True can avoid the poor accuracy
|
|
... # or the slow convergence in a way.
|
|
... optimizer = paddle.optimizer.Momentum(learning_rate=0.01, multi_precision=True)
|
|
... # 3) These ops in `custom_black_list` will keep in the float32 computation type.
|
|
... amp_list = paddle.static.amp.CustomOpLists(custom_black_list=['pool2d'])
|
|
... # 4) The entry of Paddle AMP.
|
|
... # Enable pure fp16 training by setting `use_pure_fp16` to True.
|
|
... optimizer = paddle.static.amp.decorate(
|
|
... optimizer,
|
|
... amp_list,
|
|
... init_loss_scaling=128.0,
|
|
... use_dynamic_loss_scaling=True,
|
|
... use_pure_fp16=True,
|
|
... )
|
|
... # If you don't use the default_startup_program(), you should pass
|
|
... # your defined `startup_program` into `minimize`.
|
|
... optimizer.minimize(loss)
|
|
... exe.run(paddle.static.default_startup_program())
|
|
... # 5) Use `amp_init` after FP32 parameters initialization(such as `exe.run(startup_program)`).
|
|
... # If you want to perform the testing process, you should pass `test_program` into `amp_init`.
|
|
... optimizer.amp_init(place, scope=paddle.static.global_scope())
|
|
|
|
>>> if paddle.is_compiled_with_cuda() and len(paddle.static.cuda_places()) > 0:
|
|
... run_example_code()
|
|
"""
|
|
assert self._train_program is not None, (
|
|
"Please call the minimize method first."
|
|
)
|
|
if self._use_pure_fp16:
|
|
cast_parameters_to_fp16(
|
|
place,
|
|
self._train_program,
|
|
scope,
|
|
self._to_fp16_var_names,
|
|
self._amp_vartype,
|
|
rewrite_master_weight,
|
|
self._optimizer._master_weights,
|
|
)
|
|
if test_program is not None:
|
|
if self._use_pure_fp16:
|
|
cast_model_to_fp16(
|
|
test_program,
|
|
self._amp_lists,
|
|
self._use_fp16_guard,
|
|
self._amp_vartype,
|
|
level='O2',
|
|
use_promote=self.use_promote,
|
|
)
|
|
elif use_fp16_test:
|
|
# use_fp16_guard is not support amp-o1.
|
|
cast_model_to_fp16(
|
|
test_program,
|
|
self._amp_lists,
|
|
use_fp16_guard=False,
|
|
dest_type=self._amp_vartype,
|
|
level=self._amp_level,
|
|
use_promote=self.use_promote,
|
|
)
|
|
|
|
def _append_cast_to_master_grad_op(self, param_grads):
|
|
"""
|
|
Create master gradient vars and add cast gradient to master gradient op in main program
|
|
|
|
Args:
|
|
param_grads(list(tuple(Tensor, Tensor))): A list of (parameter, gradient) pair to update.
|
|
|
|
Returns:
|
|
list: A list of (parameter, master_gradient) pair. In the following grad clip step and optimizer step, params can be updated by master gradient. main_prog will also append cast ops before grad clip ops.
|
|
|
|
"""
|
|
|
|
if not self._use_master_grad:
|
|
return param_grads
|
|
|
|
global_block = self._train_program.global_block()
|
|
target_block = global_block
|
|
if not in_pir_mode():
|
|
current_block = self._train_program.current_block()
|
|
if current_block.idx != global_block.idx:
|
|
target_block = self._train_program.blocks[
|
|
current_block.backward_block_idx
|
|
]
|
|
params_master_grads = []
|
|
|
|
assert isinstance(
|
|
target_block, (paddle.base.framework.Block, paddle.pir.Block)
|
|
)
|
|
|
|
if in_pir_mode():
|
|
for p, g in param_grads:
|
|
if g not in self._optimizer._master_grads:
|
|
if self._optimizer._is_dtype_fp16_or_bf16(g.dtype):
|
|
master_g = self._optimizer._create_master_grad(g)
|
|
params_master_grads.append((p, master_g))
|
|
else:
|
|
params_master_grads.append((p, g))
|
|
else:
|
|
# create
|
|
for p, g in param_grads:
|
|
if g.name not in self._optimizer._master_grads.keys():
|
|
if self._optimizer._is_dtype_fp16_or_bf16(g.dtype):
|
|
master_g = self._optimizer._create_master_grad(g)
|
|
params_master_grads.append((p, master_g))
|
|
target_block.append_op(
|
|
type="cast",
|
|
inputs={"X": [g]},
|
|
outputs={"Out": [master_g]},
|
|
attrs={
|
|
"in_dtype": g.dtype,
|
|
"out_dtype": master_g.dtype,
|
|
},
|
|
)
|
|
else:
|
|
params_master_grads.append((p, g))
|
|
|
|
return params_master_grads
|
|
|
|
def apply_gradients(self, params_grads):
|
|
"""
|
|
Check scaled gradients to determine whether to update loss scaling and update
|
|
parameters by their scaled gradients.
|
|
|
|
Args:
|
|
params_grads (list): A list of params and scaled grads.
|
|
|
|
Returns:
|
|
A list of optimize operators.
|
|
"""
|
|
|
|
if not in_pir_mode():
|
|
# Change the op_role_var attr for some ops, so that gradients
|
|
# transferred across GPUs can be FP16.
|
|
update_role_var_grad(self._train_program, params_grads)
|
|
|
|
# Create master grad and add cast op into program
|
|
params_grads = self._append_cast_to_master_grad_op(params_grads)
|
|
|
|
# When not using dynamic loss scaling and the init loss scaling value is equal to 1.0,
|
|
# the model can be optimized.
|
|
if (
|
|
not self._use_dynamic_loss_scaling
|
|
and self._init_loss_scaling == 1.0
|
|
):
|
|
return self._optimizer.apply_gradients(params_grads)
|
|
|
|
if self._supports_check_nan_inf():
|
|
self._optimizer._set_scale(self._loss_scaling)
|
|
optimize_ops = self._optimizer.apply_gradients(params_grads)
|
|
found_inf = self._optimizer._found_inf
|
|
self._add_dynamic_loss_scaling(params_grads, found_inf)
|
|
return optimize_ops
|
|
|
|
found_inf = self._check_finite_and_unscale(params_grads)
|
|
if self._use_dynamic_loss_scaling and (
|
|
self._amp_vartype == paddle.float16
|
|
or self._amp_vartype == core.DataType.FLOAT16
|
|
):
|
|
self._add_dynamic_loss_scaling(params_grads, found_inf)
|
|
|
|
# Pass found_inf to adam, to skip update for not only param, but also momentum and beta_pow
|
|
# With fleet, optimizers are nested and the real optimizer set by user is the inner most one.
|
|
real_optimizer = self._optimizer
|
|
while hasattr(real_optimizer, "inner_opt"):
|
|
real_optimizer = real_optimizer.inner_opt
|
|
if isinstance(
|
|
real_optimizer,
|
|
(paddle.optimizer.Adam, paddle.optimizer.AdamW),
|
|
):
|
|
# NOTE(zhiqiu): Since found_inf needs to be on cpu in adam op, we
|
|
# copy it in advance to avoid multiple time copies.
|
|
with self._train_program._optimized_guard([]):
|
|
found_inf = paddle.tensor.creation._memcpy(
|
|
found_inf, paddle.CPUPlace()
|
|
)
|
|
real_optimizer._set_auxiliary_var('found_inf', found_inf)
|
|
elif hasattr(real_optimizer, "_set_auxiliary_var"):
|
|
real_optimizer._set_auxiliary_var('found_inf', found_inf)
|
|
optimize_ops = self._optimizer.apply_gradients(params_grads)
|
|
return optimize_ops
|
|
|
|
def _split_grads(self, params_grads):
|
|
grads = [g for _, g in params_grads]
|
|
fp32_grads = [
|
|
g
|
|
for g in grads
|
|
if g.dtype == paddle.float32 or g.dtype == core.DataType.FLOAT32
|
|
]
|
|
fp16_grads = [g for g in grads if g.dtype == self._amp_vartype]
|
|
assert len(fp32_grads) + len(fp16_grads) == len(grads), (
|
|
"Data types of all grads must be either fp16/bf16 or fp32."
|
|
)
|
|
return grads, fp32_grads, fp16_grads
|
|
|
|
def _check_finite_and_unscale(self, params_grads):
|
|
grads, fp32_grads, fp16_grads = self._split_grads(params_grads)
|
|
found_infs = []
|
|
|
|
if self._is_distributed:
|
|
# if distributed, split check_finite_and_unscale to overlap
|
|
# unscale with communication
|
|
for p, g in params_grads:
|
|
with self._train_program._optimized_guard([p, g]):
|
|
_, found_inf = check_finite_and_unscale(
|
|
[
|
|
g,
|
|
],
|
|
self._loss_scaling,
|
|
name="find_infinite_scale",
|
|
float_status=self._float_status,
|
|
)
|
|
found_infs.append(found_inf)
|
|
elif self._use_pure_fp16:
|
|
if fp32_grads:
|
|
with self._train_program._optimized_guard(fp32_grads):
|
|
_, fp32_found_inf = check_finite_and_unscale(
|
|
fp32_grads,
|
|
self._loss_scaling,
|
|
name="find_infinite_scale_fp32",
|
|
float_status=self._float_status,
|
|
)
|
|
found_infs.append(fp32_found_inf)
|
|
if fp16_grads:
|
|
with self._train_program._optimized_guard(fp16_grads):
|
|
_, fp16_found_inf = check_finite_and_unscale(
|
|
fp16_grads,
|
|
self._loss_scaling,
|
|
name="find_infinite_scale_fp16",
|
|
float_status=self._float_status,
|
|
)
|
|
found_infs.append(fp16_found_inf)
|
|
else:
|
|
with self._train_program._optimized_guard(grads):
|
|
_, found_inf = check_finite_and_unscale(
|
|
grads,
|
|
self._loss_scaling,
|
|
name="find_infinite_scale",
|
|
float_status=self._float_status,
|
|
)
|
|
found_infs.append(found_inf)
|
|
|
|
if len(found_infs) > 1:
|
|
with self._train_program._optimized_guard([]):
|
|
all_infs = paddle.concat(found_infs)
|
|
found_inf = paddle.any(all_infs)
|
|
else:
|
|
found_inf = found_infs[0]
|
|
|
|
return found_inf
|
|
|
|
def _add_dynamic_loss_scaling(self, params_grads, found_inf):
|
|
if self._supports_check_nan_inf():
|
|
with self._train_program._optimized_guard([]):
|
|
update_loss_scaling(
|
|
[],
|
|
found_inf,
|
|
self._loss_scaling,
|
|
self._num_good_steps,
|
|
self._num_bad_steps,
|
|
self._incr_every_n_steps,
|
|
self._decr_every_n_nan_or_inf,
|
|
self._incr_ratio,
|
|
self._decr_ratio,
|
|
stop_update=self._optimizer._get_stop_update_var(),
|
|
name="update_loss_scaling",
|
|
)
|
|
return
|
|
|
|
grads, fp32_grads, fp16_grads = self._split_grads(params_grads)
|
|
if self._use_pure_fp16:
|
|
stop_update = False
|
|
with self._train_program._optimized_guard([]):
|
|
if fp32_grads:
|
|
update_loss_scaling(
|
|
fp32_grads,
|
|
found_inf,
|
|
self._loss_scaling,
|
|
self._num_good_steps,
|
|
self._num_bad_steps,
|
|
self._incr_every_n_steps,
|
|
self._decr_every_n_nan_or_inf,
|
|
self._incr_ratio,
|
|
self._decr_ratio,
|
|
stop_update=stop_update,
|
|
name="update_loss_scaling_fp32",
|
|
)
|
|
stop_update = True
|
|
if fp16_grads:
|
|
update_loss_scaling(
|
|
fp16_grads,
|
|
found_inf,
|
|
self._loss_scaling,
|
|
self._num_good_steps,
|
|
self._num_bad_steps,
|
|
self._incr_every_n_steps,
|
|
self._decr_every_n_nan_or_inf,
|
|
self._incr_ratio,
|
|
self._decr_ratio,
|
|
stop_update=stop_update,
|
|
name="update_loss_scaling_fp16",
|
|
)
|
|
else:
|
|
with self._train_program._optimized_guard([]):
|
|
update_loss_scaling(
|
|
grads,
|
|
found_inf,
|
|
self._loss_scaling,
|
|
self._num_good_steps,
|
|
self._num_bad_steps,
|
|
self._incr_every_n_steps,
|
|
self._decr_every_n_nan_or_inf,
|
|
self._incr_ratio,
|
|
self._decr_ratio,
|
|
name="update_loss_scaling",
|
|
)
|
|
|
|
def apply_optimize(self, loss, startup_program, params_grads):
|
|
program = loss.block.program
|
|
with paddle.static.program_guard(program, startup_program):
|
|
optimize_ops = self.apply_gradients(params_grads)
|
|
return optimize_ops
|
|
|
|
def minimize(
|
|
self, loss, startup_program=None, parameter_list=None, no_grad_set=None
|
|
):
|
|
"""
|
|
Perform optimization by minimizing the given loss.
|
|
|
|
Args:
|
|
loss (Variable): The loss Variable.
|
|
startup_program (Program): startup_program for initializing parameters
|
|
in `parameter_list`.
|
|
parameter_list (list): list of Variables to update.
|
|
no_grad_set (set|None): set of Variables should be ignored.
|
|
|
|
Returns:
|
|
The scaled loss by scaling factor, the list of optimize ops, and a
|
|
list of scaled parameters and gradients.
|
|
"""
|
|
|
|
opt_dict = self._optimizer.__class__.__dict__
|
|
if 'minimize' in opt_dict and isinstance(
|
|
opt_dict['minimize'], types.FunctionType
|
|
):
|
|
warnings.warn(
|
|
"The decorated optimizer has its own `minimize` method, but it will not be executed."
|
|
)
|
|
|
|
with auto_complete_op_role(loss.block.program, op_role=OpRole.Backward):
|
|
scaled_params_grads = self.backward(
|
|
loss,
|
|
startup_program=startup_program,
|
|
parameter_list=parameter_list,
|
|
no_grad_set=no_grad_set,
|
|
)
|
|
|
|
with auto_complete_op_role(loss.block.program, op_role=OpRole.Optimize):
|
|
optimize_ops = self.apply_optimize(
|
|
loss, startup_program, scaled_params_grads
|
|
)
|
|
|
|
return optimize_ops, scaled_params_grads
|
|
|
|
|
|
@overload(key=FunctionType.FP16_ONLY)
|
|
def decorate(
|
|
optimizer,
|
|
amp_lists=None,
|
|
init_loss_scaling=2**15,
|
|
incr_every_n_steps=1000,
|
|
decr_every_n_nan_or_inf=2,
|
|
incr_ratio=2.0,
|
|
decr_ratio=0.8,
|
|
use_dynamic_loss_scaling=True,
|
|
use_pure_fp16=False,
|
|
use_fp16_guard=None,
|
|
use_bf16=False,
|
|
use_promote=False,
|
|
):
|
|
"""
|
|
Decorate the given optimizer to adapt to the mixed-precision training.
|
|
|
|
Args:
|
|
optimizer(Optimizer): A common Optimizer.
|
|
amp_lists (CustomOpLists): An CustomOpLists object.
|
|
init_loss_scaling(float): The initial loss scaling factor.
|
|
incr_every_n_steps(int): Increases loss scaling every n consecutive
|
|
steps with finite gradients.
|
|
decr_every_n_nan_or_inf(int): Decreases loss scaling every n
|
|
accumulated steps with nan or
|
|
inf gradients.
|
|
incr_ratio(float): The multiplier to use when increasing the loss
|
|
scaling.
|
|
decr_ratio(float): The less-than-one-multiplier to use when decreasing
|
|
the loss scaling.
|
|
use_dynamic_loss_scaling(bool): Whether to use dynamic loss scaling.
|
|
use_pure_fp16(bool): Whether to use the pure fp16 training. Default False.
|
|
use_fp16_guard(bool): Whether to use `fp16_guard` when constructing the program.
|
|
Default None, which means that its value equals to `use_pure_fp16`.
|
|
use_bf16(bool): Whether to enable bfloat16 training. Default False.
|
|
|
|
Returns:
|
|
An optimizer acting like a normal one but with mixed-precision training
|
|
enabled.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
:name: example-1
|
|
|
|
# black&white list based strategy example
|
|
>>> import paddle
|
|
>>> import paddle.static as static
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> data = static.data(name='X', shape=[None, 1], dtype='float32')
|
|
>>> hidden = static.nn.fc(x=data, size=10)
|
|
>>> loss = paddle.mean(hidden)
|
|
>>> optimizer = paddle.optimizer.Adam(learning_rate=0.001)
|
|
|
|
>>> mp_optimizer = static.amp.decorate(optimizer=optimizer, init_loss_scaling=8.0)
|
|
|
|
>>> ops, param_grads = mp_optimizer.minimize(loss)
|
|
>>> scaled_loss = mp_optimizer.get_scaled_loss()
|
|
|
|
|
|
.. code-block:: pycon
|
|
:name: example-2
|
|
|
|
# pure fp16 training example
|
|
>>> import numpy as np
|
|
>>> import paddle
|
|
>>> import paddle.nn.functional as F
|
|
>>> paddle.enable_static()
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> def run_example_code():
|
|
... place = paddle.CUDAPlace(0)
|
|
... exe = paddle.static.Executor(place)
|
|
... data = paddle.static.data(name='X', shape=[None, 1, 28, 28], dtype='float32')
|
|
... conv2d = paddle.static.nn.conv2d(input=data, num_filters=6, filter_size=3)
|
|
... # 1) Use fp16_guard to control the range of fp16 kernels used.
|
|
... with paddle.static.amp.fp16_guard():
|
|
... bn = paddle.static.nn.batch_norm(input=conv2d, act="relu")
|
|
... pool = F.max_pool2d(bn, kernel_size=2, stride=2)
|
|
... hidden = paddle.static.nn.fc(pool, size=10)
|
|
... loss = paddle.mean(hidden)
|
|
... # 2) Create the optimizer and set `multi_precision` to True.
|
|
... # Setting `multi_precision` to True can avoid the poor accuracy
|
|
... # or the slow convergence in a way.
|
|
... optimizer = paddle.optimizer.Momentum(learning_rate=0.01, multi_precision=True)
|
|
... # 3) These ops in `custom_black_list` will keep in the float32 computation type.
|
|
... amp_list = paddle.static.amp.CustomOpLists(custom_black_list=['pool2d'])
|
|
... # 4) The entry of Paddle AMP.
|
|
... # Enable pure fp16 training by setting `use_pure_fp16` to True.
|
|
... optimizer = paddle.static.amp.decorate(
|
|
... optimizer,
|
|
... amp_list,
|
|
... init_loss_scaling=128.0,
|
|
... use_dynamic_loss_scaling=True,
|
|
... use_pure_fp16=True,
|
|
... )
|
|
... # If you don't use the default_startup_program(), you should pass
|
|
... # your defined `startup_program` into `minimize`.
|
|
... optimizer.minimize(loss)
|
|
... exe.run(paddle.static.default_startup_program())
|
|
... # 5) Use `amp_init` after FP32 parameters initialization(such as `exe.run(startup_program)`).
|
|
... # If you want to perform the testing process, you should pass `test_program` into `amp_init`.
|
|
... optimizer.amp_init(place, scope=paddle.static.global_scope())
|
|
|
|
>>> if paddle.is_compiled_with_cuda() and len(paddle.static.cuda_places()) > 0:
|
|
... run_example_code()
|
|
"""
|
|
amp_dtype = "bfloat16" if use_bf16 else "float16"
|
|
if amp_lists is None:
|
|
amp_lists = AutoMixedPrecisionLists(dtype=amp_dtype)
|
|
|
|
if use_fp16_guard is None:
|
|
use_fp16_guard = use_pure_fp16
|
|
|
|
amp_level = "O2" if use_pure_fp16 else "O1"
|
|
mp_optimizer = OptimizerWithMixedPrecision(
|
|
optimizer,
|
|
amp_lists,
|
|
level=amp_level,
|
|
dtype=amp_dtype,
|
|
init_loss_scaling=init_loss_scaling,
|
|
use_dynamic_loss_scaling=use_dynamic_loss_scaling,
|
|
incr_every_n_steps=incr_every_n_steps,
|
|
decr_every_n_nan_or_inf=decr_every_n_nan_or_inf,
|
|
incr_ratio=incr_ratio,
|
|
decr_ratio=decr_ratio,
|
|
use_amp_guard=use_fp16_guard,
|
|
use_promote=use_promote,
|
|
)
|
|
|
|
return mp_optimizer
|
|
|
|
|
|
@overload(key=FunctionType.COMMON)
|
|
def decorate( # noqa: F811
|
|
optimizer,
|
|
amp_lists=None,
|
|
level='O1',
|
|
dtype='float16',
|
|
master_weight=None,
|
|
master_grad=False,
|
|
init_loss_scaling=2**16,
|
|
incr_every_n_steps=2000,
|
|
decr_every_n_nan_or_inf=1,
|
|
incr_ratio=2.0,
|
|
decr_ratio=0.5,
|
|
use_dynamic_loss_scaling=None,
|
|
use_amp_guard=False,
|
|
use_promote=False,
|
|
):
|
|
"""
|
|
Decorate the given optimizer to adapt to the mixed-precision training.
|
|
|
|
Args:
|
|
optimizer(Optimizer): A common Optimizer.
|
|
amp_lists(CustomOpLists, optional): An CustomOpLists object. The default
|
|
white_list and black_list will be used for AMP training when it is
|
|
not set. Default is None.
|
|
level(str, optional): Auto mixed precision level. Accepted values are "O1", "O2" and "OD": At the O1 level, operators in the white list
|
|
will use float16/bfloat16 inputs for calculations, and operators in the black list will use float32 inputs for calculations. At the O2
|
|
level, model's parameters will be casted to float16/bfloat16 by using `decorator`, and operators that have all float16/bfloat16 inputs
|
|
will be converted to float16/bfloat16, and that have any float32 input will be converted to float32. For the OD level, operators in
|
|
default white list will compute in float16/bfloat16, and the others will compute in float32. Default is O1.
|
|
dtype(str, optional): Whether to use 'float16' or 'bfloat16'. Default is 'float16'.
|
|
master_weight(bool, optional): For level='O2', whether to use multi-precision
|
|
during weight updating. If master_weight is None, in O2 level optimizer
|
|
will use multi-precision. Default is None.
|
|
master_grad(bool, optional): For level='O2', whether to use master_grad
|
|
during weight updating. If master_grad is False, in O2 level optimizer
|
|
will not use master grad. Default is False.
|
|
init_loss_scaling(float, optional): The initial loss scaling factor.
|
|
Default is 65536.
|
|
incr_every_n_steps(int, optional): Increases loss scaling every n
|
|
consecutive steps with finite gradients. Default is 2000.
|
|
decr_every_n_nan_or_inf(int, optional): Decreases loss scaling every n
|
|
accumulated steps with nan or inf gradients. Default is 1.
|
|
incr_ratio(float, optional): The multiplier to use when increasing the
|
|
loss scaling. Default is 2.
|
|
decr_ratio(float, optional): The less-than-one-multiplier to use when
|
|
decreasing the loss scaling. Default is 0.5.
|
|
use_dynamic_loss_scaling(bool, None): Whether to use dynamic loss
|
|
scaling. Default is None, which means True for float16, and False
|
|
for bfloat16.
|
|
|
|
Returns:
|
|
An optimizer acting like a normal one but with mixed-precision training
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> paddle.enable_static()
|
|
|
|
>>> # doctest: +REQUIRES(env:GPU)
|
|
>>> class SimpleConvNet(paddle.nn.Layer):
|
|
... def __init__(self):
|
|
... super().__init__()
|
|
... self.conv = paddle.nn.Conv2D(in_channels=1, out_channels=6, kernel_size=3)
|
|
... self.linear = paddle.nn.Linear(in_features=26, out_features=10)
|
|
...
|
|
... def forward(self, x):
|
|
... out = self.conv(x)
|
|
... out = paddle.nn.functional.relu(out)
|
|
... out = self.linear(out)
|
|
... out = paddle.nn.functional.softmax(out)
|
|
... return out
|
|
|
|
>>> main_program = paddle.static.Program()
|
|
>>> startup_program = paddle.static.Program()
|
|
>>> with paddle.utils.unique_name.guard():
|
|
... with paddle.static.program_guard(main_program, startup_program):
|
|
... model = SimpleConvNet()
|
|
... x = paddle.static.data(name='input', shape=[None, 1, 28, 28], dtype='float32')
|
|
... out = model(x)
|
|
... loss = paddle.mean(out)
|
|
... optimizer = paddle.optimizer.AdamW()
|
|
... optimizer = paddle.static.amp.decorate(optimizer, level="O2", dtype="float16")
|
|
... optimizer.minimize(loss)
|
|
|
|
>>> if paddle.is_compiled_with_cuda() and len(paddle.static.cuda_places()) > 0:
|
|
... place = paddle.CUDAPlace(0)
|
|
... exe = paddle.static.Executor(place)
|
|
... exe.run(startup_program)
|
|
...
|
|
... # Call `amp_init` after FP32 parameters initialization, such as `exe.run(startup_program)`,
|
|
... # to convert FP32 parameters to low precision FP16 / BF16.
|
|
... optimizer.amp_init(place, scope=paddle.static.global_scope())
|
|
|
|
"""
|
|
# check amp_level: O0-O2
|
|
level = level.upper()
|
|
if level not in ['O0', 'OD', 'O1', 'O2']:
|
|
raise ValueError("level should be O0, OD, O1 or O2.")
|
|
|
|
amp_dtype = check_amp_dtype(dtype)
|
|
if amp_lists is None or level == 'OD':
|
|
amp_lists = AutoMixedPrecisionLists(dtype=amp_dtype)
|
|
|
|
if level == 'OD':
|
|
if amp_lists is not None:
|
|
warnings.warn(
|
|
"If the Amp level is set to OD, the amp list will not be used."
|
|
)
|
|
amp_lists.black_list = amp_lists.all_list - amp_lists.white_list
|
|
|
|
if use_dynamic_loss_scaling is None:
|
|
use_dynamic_loss_scaling = dtype == "float16"
|
|
|
|
if optimizer is not None:
|
|
# support master_weight
|
|
multi_precision = master_weight is not False
|
|
_set_multi_precision(optimizer, multi_precision)
|
|
|
|
mp_optimizer = OptimizerWithMixedPrecision(
|
|
optimizer,
|
|
amp_lists,
|
|
level=level,
|
|
dtype=amp_dtype,
|
|
init_loss_scaling=init_loss_scaling,
|
|
use_dynamic_loss_scaling=use_dynamic_loss_scaling,
|
|
incr_every_n_steps=incr_every_n_steps,
|
|
decr_every_n_nan_or_inf=decr_every_n_nan_or_inf,
|
|
incr_ratio=incr_ratio,
|
|
decr_ratio=decr_ratio,
|
|
use_amp_guard=use_amp_guard,
|
|
use_promote=use_promote,
|
|
use_master_grad=master_grad,
|
|
)
|
|
|
|
return mp_optimizer
|