RFC: Keras Mixed Precision API

[reedwm]

This RFC will be open for comment until October 14, 2020.

Status	Proposed
RFC #	293
Author(s)	Reed Wanderman-Milne (reedwm@google.com
Sponsor	Francois Chollet (fchollet@google.com)
Updated	2020-09-29

Objective

Make mixed precision easy to use in Keras.

[lgeiger]

@reedwm Thank you very much for the in-depth writeup. I made a few minor comments below, I am currious about what you think.

[lgeiger]

How do you feel about calling it dtype instead of name to avoid confusion? "mixed_float16" and "mixed_bfloat16" could be interpreted as a string representation that are convertible to a custom DType.
policy_name would be an alternative option, although a bit too verbose for my taste.

[reedwm]

The issue is the argument is also exposed as as a property, so if it were called dtype instead of name, Policy would have a dtype field, a compute_dtype field, and a variable_dtype field. It would be nonobvious that the dtype field was mostly unused and only used to determine the compute_dtype and variable_dtype fields.

[lgeiger]

Thanks for the explanation, that makes sense 👍

[lgeiger]

The deprecated mixed precision graph rewrite opted out of using mixed precision on non-optimal hardware.
I expect users will now need to handle environment detection manually or write custom code to detect the current hardware? I think this might lead to confusion when running in heterogeneous computing environments which requires users to be aware which policy is optimal on which hardware platform.

How do you feel about adding an auto policy (alternatively called optimal or mixed) that would autoselect between either float32, mixed_float16 or mixed_bfloat16 depending on the hardware support? This would mean that there are no code changes required when switching between CPU only, Volta GPUs or Ampere GPUs systems.

[reedwm]

I strongly considered this earlier. My main concern was that I believed users should be aware of whether they were using float16 or bfloat16. Users need to explicitly use a LossScaleOptimizer if using a custom training loop with float16 (when using Model.fit, a LossScaleOptimizer will automatically be used but otherwise the user must explicitly wrap their optimizer). Additionally, the SavedModel and checkpoint format will save the loss scale if mixed_float16 is used, and I don't think the SavedModel/checkpoint format should depend solely on the device that is used.

I added a paragraph summarizing why we don't do this.

Also, the batch size you should run at depends on the device as well, as different GPUs have different amounts of memory, so already users must be somewhat aware of their devices.

In practice, I expect most models supporting multiple devices to have a flag to choose between float32, mixed_float16, and mixed_bfloat16, and that flag is directly passed to set_global_policy. Within Google, we do this but also have a per-device per-model config file specifying what flags to pass, so running the GPU config file causes mixed_float16 to be passed to the flag and running the TPU config file causes mxied_bfloat16 to be passed.

[lgeiger]

I see your point that users might need to customize other training parameters depending on the device anyway which is indeed very common.

I think the main reason why I was missing this autoselection is that I often end up having model configs which enable mixed precision by default. However, sometimes it can be useful to run a few steps locally on a CPU only machine for debugging which now would require users to remember to disable mixed precision as it would lead to greatly decreased performance on CPU (at least that was the case a few months ago when I last tested this API).
I don't think this is a huge issue since one can easily check for the existance of a GPU in user code and disable mixed precision if needed, but if this becomes a common source of confusion it might be interesting to rethink this behaviour.

[lgeiger]

What is the benefit of skipping update steps instead of recomputing the gradients with a scaled down loss until the gradients become finite (or until fixed number of retries is reached)? Intuitively this would enable slightly higher performance as the forward pass could be reused, or am I missing something?

[reedwm]

This is a good idea, but the issue is that we cannot run the loop if the user calls apply_gradients instead of minimize because apply_gradients does not compute the gradients. For example:

vars = ...
opt = tf.keras.mixed_precision.LossScaleOptimizer(...)
with tf.GradientTape() as tape:
  loss = get_model_loss()
  grads = tape.gradient(loss, vars)
  opt.apply_gradients(zip(grads, vars))

Since the gradients are computed outside the LossScaleOptimizer, LossScaleOptimizer has no way of repeatedly recomputing the gradients.

However, we can do this in minimize. I added a paragraph suggesting we should add an option to do this in the future (starting with "Instead of skipping steps when there are NaNs"). The reason this won't be done initially is that minimize is rarely used, but I anticipate it will be used more frequently due to #234 allowing users to pass a tensor to minimize once TF 2.4 is released.

I think the performance gain will be negligible, due to steps only being skipped 1/2000 of the time on average, but the behavior will be more intuitive. The skipping behavior is confusing because for the first few steps, the variables are not actually updated since the loss scale starts out so high.

[lgeiger]

Yes, sorry I should've been clearer. I was mainly suggesting this for the use in keras.Model.fit() which indeed calls .minimize(). But since for custom training loops the skipping behaviour still needs to be supported anyway, I agree it makes sense to keep the skipping behaviour for now.

[nluehr]

Op-based autocasting with down-graph type inheritance is important for providing robust accuracy in mixed precision training. This is particularly true in experimental and research applications. For example, early implementations of batch- and layer-norm used a sequence of low-level operations such as the following.

bn(x) = y * (x-mean(x)) / sqrt(var(x)+epsilon) + beta

Without computing the mean and variance ops in fp32, training will either diverge or achieve reduced accuracy for many important networks such as InceptionV3, ResNet50, and Xception.

Of course, once the utility of an experimental layer is proven, it may be replaced with a “fused” custom kernel that provides optimized performance and handles intermediate precisions internally. However, relatively few layers will merit this effort. Novel model architectures are an important area of innovation, and mixed precision needs to be robust so that researches and early adopters can benefit from the full performance of modern architectures.

[fchollet]

We did the internal design review for this API today. Since there were no outstanding issues, it was very quick.

Notes:

• Reed: Let's present the RFC and go over the changes that happened since we sent it out.
• Before the changes: there were 3 dynamic loss scale classes. They were glorified named tuples
• Now: everything related to loss scaling moved to LossScaleOptimizer
  • There's a dynamic boolean argument to use dynamic scaling or not
    • if true, you can use initial_scale and dynamic_growth_steps
    • It's a simple API
• Q from Tomer: do you always have to use loss scale optimizer when using mixed precision?
  • Reed: it's automated with compile/fit, but with a custom training loop you need to use it
  • This is well documented

Status:

Design approved after the minor changes described during the meeting.