Closed mattdangerw closed 12 months ago
Patch coverage: 86.30%
and project coverage change: +0.09%
:tada:
Comparison is base (
ab45558
) 75.99% compared to head (fd58cfb
) 76.09%. Report is 6 commits behind head on main.
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
A few points of awkwardness/discussion:
ops.cond
needs to be stateless for jax. Autoscaling has two cond branches with variables updates in all branches. To do this I overrode stateless_apply
separately and had to do a lot of StatelessScope
s. This feels very verbose and awkward, but I wasn't able to think of a great way around it. Suggestions welcome!learning_rate
to proxy the inner optimizer learning rate. I overrode the learning_rate
property to do so. But the base optimizer still has to be created with a learning rate, so I ended up just passing a zero valued variable which never gets used. This feels awkward and a bit confusing.scale_loss
and stateless_scale_loss
. This is probably fine, just feels like a lot of code.Looks like there is some sort of device placement issue for tensorflow
GPU that isn't picked up in our CPU only tests? Will poke around. Confirmed torch/jax are working.
[Edit: now working on all backends]
Addressed the initial round, though I may play with a test that deliberately triggers the underflow in trainer and asserts that variable updates appear. I don't think that should be too hard? But we will see.
Added a test end to end test only looking at variables updates across fit()
, which will hopefully keep us from accidentally breaking this.
This is the big missing piece we need for feature parity when running mixed precision training compared to tf.keras.
Fixes #571