Closed james77777778 closed 1 year ago
Patch coverage: 56.00%
and project coverage change: +5.03%
:tada:
Comparison is base (
49e5b06
) 70.90% compared to head (609d711
) 75.94%. Report is 3 commits behind head on main.
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
Thanks for the PR! @mattdangerw was working on this, so I'll let him review.
Thanks for the PR! @mattdangerw was working on this, so I'll let him review.
Sure.
Here are some additional notes:
base_optimizer
, similar to gradient clipping and EMA, because it is easier to implement and configure.loss_scaling="auto"
has been added to compile
to configure loss scaling if the user does not specify use_loss_scaling
in the optimizer.Loss scaling is crucial for mixed precision even in a small model. You can observe the difference when loss scaling is applied versus when it is not.
MNIST classification training with mixed_float16
:
with loss scaling
without loss scaling
Updated colab: https://colab.research.google.com/drive/1T8DpRXu7CI52-67Poq5kKFvEM1P3UwT-?usp=sharing
Ah shoot! Yes I started working on this last week, and think I have this working on all backends.
On the discussion thread, we decided to try adding this as a separate optimizer for the compatibility with tf.keras
. If we end up going with that (separate optimizer) approach, we may need to close this. See https://github.com/keras-team/keras-core/pull/851
Also one random note, I believe bfloat16 does not actually need loss scaling, as it has the same amount of exponent bits as float32. Loss scaling is just for mixed_float16
!
https://www.tensorflow.org/guide/mixed_precision#loss_scaling
Thanks @mattdangerw
I was attempting to train a large model that is challenging to fit into a 24GB memory with float32. So I implemented my own solution. It's good to know that we will have an official solution for loss scaling.
This PR should be closed.
Related to #571
I'm not sure if anyone is addressing this issue, or if it should not be handled externally.
I'm simply trying to implement a backend-agnostic solution. Currently, only tf trainer has the ability to automatically enable the loss scaling when using mixed_float16 and mixed_bfloat16. It should be straightforward to extend this to other backends.
In the colab, the training works well. However, it will fail if the loss scaling technique is not applied. Colab: https://colab.research.google.com/drive/1T8DpRXu7CI52-67Poq5kKFvEM1P3UwT-?usp=sharing
cc @mattdangerw
EDITED: Now, this PR should work with tf and torch backend. However, jax does not support yet.