Low Precision Training Methods
The release of new kinds of hardware led to the emergence of new training paradigms that better utilize them. Currently, this is in the form of training in 8-bit precision using packages such as TransformersEngine (TE) or MS-AMP.
For an introduction to the topics discussed today, we recommend reviewing the low-precision usage guide as this documentation will reference it regularly.
A Quick Chart
Below is a quick chart from the MS-AMP documentation showing the different bit-precisions for each solution during training:
Optimization Level | Computation(GEMM) | Comm | Weight | Master Weight | Weight Gradient | Optimizer States |
---|---|---|---|---|---|---|
FP16 AMP | FP16 | FP32 | FP32 | N/A | FP32 | FP32+FP32 |
Nvidia TE | FP8 | FP32 | FP32 | N/A | FP32 | FP32+FP32 |
MS-AMP O1 | FP8 | FP8 | FP16 | N/A | FP8 | FP32+FP32 |
MS-AMP O2 | FP8 | FP8 | FP16 | N/A | FP8 | FP8+FP16 |
MS-AMP O3 | FP8 | FP8 | FP8 | FP16 | FP8 | FP8+FP16 |
TransformersEngine
TransformersEngine
is the first solution to trying to train in 8-bit floating point. It works by using drop-in replacement layers for certain ones in a model that utilizes their FP8-engine to reduce the number of bits (such as 32 to 8) without degrading the final accuracy of the model.
Specifically, 🤗 Accelerate will find and replace the following layers with TransformersEngine
versions:
nn.LayerNorm
forte.LayerNorm
nn.Linear
forte.Linear
As a result we wind up with a model that has most of its layers in BF16, while some layers are in FP8 reducing some of the memory.
Anecdotally, we have noticed that performance gains don’t really start showing when using TransformerEngine
until a large majority of the layers
in the model are made up of those two layers to replace. As a result, only larger models have shown performance improvements when the number of parameters is around and upwards of a few billion.
The TransformerEngine
can receive many different arguments that customize how it performs FP8 calculations and what they do. A full list of the arguments is available below:
margin
: The margin to use for the gradient scaling.interval
: The interval to use for how often the scaling factor is recomputed.fp8_format``: The format to use for the FP8 recipe. Must be one of
E4M3or
HYBRID`.amax_history_len
: The length of the history to use for the scaling factor computationamax_compute_algo
: The algorithm to use for the scaling factor computation. Must be one ofmax
ormost_recent
.override_linear_precision
: Whether or not to executefprop
,dgrad
, andwgrad
GEMMS in higher precision.
You can customize each of these as part of utils.FP8RecipeKwargs to help optimize performance of your models.
If we notice in the chart mentioned earlier, TE simply casts the computation layers into FP8, while everything else is in FP32. As a result this winds up utilizing the most memory but does so with the benefit of guaranteeing the least amount of loss in end accuracy during training.
## MS-AMP
MS-AMP takes a different approach to TransformersEngine
by providing three different optimization levels to convert more operations in FP8 or FP16.
The base optimization level (
O1
), passes communications of the weights (such as in DDP) in FP8, stores the weights of the model in FP16, and leaves the optimizer states in FP32. The main benefit of this optimization level is that we can reduce the communication bandwidth by essentially half. Additionally, more GPU memory is saved due to 1/2 of everything being cast in FP8, and the weights being cast to FP16. Notably, both the optimizer states remain in FP32.The second optimization level (
O2
) improves upon this by also reducing the precision of the optimizer states. One is in FP8 while the other is in FP16. Generally it’s been shown that this will only provide a net-gain of no degraded end accuracy, increased training speed, and reduced memory as now every state is either in FP16 or FP8.Finally, MS-AMP has a third optimization level (
O3
) which helps during DDP scenarios such as DeepSpeed. The weights of the model in memory are fully cast to FP8, and the master weights are now stored in FP16. This fully reduces memory by the highest factor as now not only is almost everything in FP8, only two states are left in FP16. Currently, only DeepSpeed versions up through 0.9.2 are supported, so this capability is not included in the 🤗 Accelerate integration
Combining the two
More experiments need to be performed but it’s been noted that combining both MS-AMP and TransformersEngine can lead to the highest throughput by relying on NVIDIA’s optimized FP8 operators and utilizing how MS-AMP reduces the memory overhead.
< > Update on GitHub