.. _mixed_precision:
Mixed Precision Training
========================
.. epigraph::
Increasing the size of a neural network typically improves accuracy but also
increases the memory and compute requirements for training the model.
We introduce methodology for training deep neural networks using
half-precision floating point numbers, without losing model accuracy or
having to modify hyperparameters. This nearly halves memory requirements
and, on recent GPUs, speeds up arithmetic. ...
DNN operations benchmarked with DeepBench on Volta GPU see 2-6x speedups
compared to FP32 implementations if they are limited by memory or arithmetic
bandwidth. Speedups are lower when operations are latency-limited.
-- "Mixed Precision Training", Micikevicius et al, ICLR, 2018
Prerequisites
-------------
Mixed precision training utilizes Tensor Cores introduced in `NVIDIA Volta GPUs `_
such as `Titan V `_ and `Tesla V100 `_.
NVIDIA Volta GPUs are also available from `AWS on p3.2xlarge, p3.8xlarge, p3.16xlarge instances `_ .
For an optimal mixed precision performance we recommend using NVIDIA's TensorFlow docker containers (version 18.03 and above)
which can be obtained here: `NVIDIA GPU cloud `_ .
Alternatively, you can build TensorFlow yourself with CUDA 9.1 and this `PR `_ included:
How to enable mixed precision
-----------------------------
Enabling mixed precision with existing models in OpenSeq2Seq is simple:
change ``dtype`` parameter of ``model_params`` to "mixed".
You might need to enable loss scaling: either statically, by setting
``loss_scale`` parameter inside ``model_params`` to the desired number, or
you can use dynamic loss scaling by setting ``automatic_loss_scaling`` parameter
to "Backoff" or "LogMax"::
base_params = {
...
"dtype": "mixed",
# enabling static or dynamic loss scaling might improve model convergence
# "loss_scale": 10.0,
# "automatic_loss_scaling": "Backoff",
...
}
.. One can also experiment with more fine precision granularity.
For example set encoder precision in float16 and decoder in float32::
.. "model_params": {
...
"dtype": tf.float16,
...
}
"decoder_params": {
...
"dtype": tf.float32,
...
}
Implementation details
----------------------
For mixed precision training we follow an algorithmic recipe from
Micikevicius et al :cite:`mp-2018`. At a high level it can be summarized
as follows:
1. Maintain and update a float32 master copy of weights (using the float16 copy
for forward and back propagation)
2. Apply loss scaling while computing gradients
It is worth mentioning that (1)-(2) are not always necessary. However, this
method has proven to be robust across a variety of bigger and more complex
models.
Note that while (1) does mean a 50% increase in memory consumption for weights
over a float32 version, in practice, the total memory consumption is often
*decreased*. This is because activations, activation gradients, and other
intermediate tensors can now be kept in float16. This is especially beneficial
for models with a high degree of parameter sharing or reuse, such as recurrent
models with many timesteps.
Optimizer
~~~~~~~~~
Our implementation is different from the one described in
`NVIDIA documentation `_:
instead of a custom variable getter, we introduce a wrapper around standard
TensorFlow optimizers. The model is created with float16 data type, so all
variables and gradients are in float16 by default (except for the layers which
are explicitly redefined as float32; for example data layers or operations on
CPU). The wrapper then converts float16 gradients to float32 and submits them
to TensorFlow's optimizer, which updates the master copy of weights. Updated
weights are converted back to float16, and used by the model in the next
iteration. The :class:`MixedPrecisionOptimizerWrapper `
architecture is graphically illustrated below:
.. figure:: MixedPrecisionOptimizer.png
:scale: 50 %
:align: center
"Mixed precision" optimizer wrapper around any TensorFlow optimizer.
Regularizers
~~~~~~~~~~~~
:class:`MixedPrecisionOptimizerWrapper `
ensures that all float16 variables will have
a master copy in float32 and that their gradients will be cast to the full
precision before computing the weight update. While this is enough in most
situations, in some cases it is important to keep the gradient in float32 from
the beginning of the computation. One common case when this is necessary is
weight decay regularization. This regularization results in the following
addition to the usual gradients with respect to the loss:
:math:`\frac{\partial L}{\partial w} \mathrel{+}= 2\lambda w`,
where :math:`\lambda` is usually on the order of
:math:`\left[10^{-5}, 10^{-3}\right]`.
Given that the weights are commonly initialized with small values, multiplying
them with weight decay coefficient $\lambda$ can result in numerical underflow.
To overcome this problem we implemented the following design principles. First,
all regularizers should be defined on the variable creation level by passing
regularizer function as a regularizer parameter to the ``tf.get_variable``
function or ``tf.layers`` objects (this is a recommended way to do it in
TensorFlow). Second, the regularizer function should be wrapped with
:func:`mp_regularizer_wrapper `
function which will do two things. First, it
will add weights and the user-provided regularization function to the TensorFlow
collection. And second, it will disable the underlying regularization function
by returning None (only if the weights are in float16, otherwise it will not
introduce any additional behavior). The created collection will later be
retrieved by ``MixedPrecisionOptimizerWrapper`` and the corresponding
functions will be applied to the float32 copy of the weights ensuring that their
gradients always stay in the full precision. Since this regularization will not
be a part of the loss computation graph, we explicitly call ``tf.gradients``
and add the result to the gradients passed in the ``compute_gradients``
function of the optimizer.
Automatic Loss Scaling
~~~~~~~~~~~~~~~~~~~~~~
The mixed precision training approach suggests that the user
set a *loss scale* hyperparameter to adjust the dynamic range of backpropagation
to match the dynamic range of float16. OpenSeq2Seq implements an extension to
the mixed precision recipe that we call *automatic loss scaling*. The optimizer
inspects the parameter gradients at each iteration and uses their values to
select the loss scale for the *next* iteration. As a result, the user does not
have to select the loss-scale value manually.
Concretely, OpenSeq2Seq has support for two automatic loss scaling algorithms,
*Backoff* and *LogNormal* scaling.
* *Backoff* scaling begins with a large loss scale and checks for overflow in
the parameter gradients at the end of each iteration. Whenever there is an
overflow, the loss scale decreases by a constant factor (default is 2) and the
optimizer will skip the update. Furthermore, if there has been no overflow for
a period of time, the loss scale increases by a constant factor (defaults are
2000 iterations and 2, respectively). These two rules together ensure both
that the loss scale is as large as possible and also that it can adjust to
shifting dynamic range during training.
* *LogNormal* scaling uses gradient statistics, rather than the presence of
overflow, to set the loss scale. It keeps a running estimate of the mean and
variance of the inter-iteration maximum absolute value of the parameter
gradients. It models the inter-iteration maximum as log-normally distributed
(hence the name), and then chooses the loss scale for the next iteration s.t.
the probability of the maximum overflowing float16 is less than some constant
(default is 0.001). In the rare event of an overflow, the optimizer skips the
update.
.. How to port models from float32 to mixed precision
.. --------------------------------------------------
.. ...
.. bibliography:: refs.bib
:cited:
:style: unsrt