Closed qlzh727 closed 11 months ago
Patch coverage: 12.50%
and project coverage change: -0.06%
:warning:
Comparison is base (
1704ecf
) 79.75% compared to head (3fd4ad4
) 79.70%.
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
Let's park this PR for now, since there are some underlying issues (discovered in the unit test) need to be address.
ok. I think i figure out the issue.
Humm, its more tricky than I expected, esp for RNN. In the RNN (LSTM/GRU) with dropout, the seed generator get updated within the jax.lax.scan() in https://github.com/keras-team/keras-core/blob/b4019bc1bfa2c2f15decb533d054cbcf31562124/keras_core/backend/jax/rnn.py#L188 (we even have a comment about the function need to be stateless). Since the step function doesn't return the RNG state as explicit result, when we captured the updated RNG from train function and reuse it for jax constraint, JAX complain about this leaked state from the scan function.
I think the current repo is probably not working correctly, since the StatelessScope for the scan function was throw away, which mean the RNG seed update is lost.
Will take a closer look tomorrow.
I think the current repo is probably not working correctly, since the StatelessScope for the scan function was throw away, which mean the RNG seed update is lost.
We previously observed some weirdness in JAX RNNs (though not connected to dropout it seems). Maybe some relationship there? https://github.com/keras-team/keras-core/issues/322
Since the step function doesn't return the RNG state as explicit result, when we captured the updated RNG from train function and reuse it for jax constraint, JAX complain about this leaked state from the scan function.
I think we might want to add a stateless version of the step function / etc.
I think the current repo is probably not working correctly, since the StatelessScope for the scan function was throw away, which mean the RNG seed update is lost.
We previously observed some weirdness in JAX RNNs (though not connected to dropout it seems). Maybe some relationship there? #322
The particular issue we hit is for the RNG state, which gets updated when using dropout. It runs fine if I disable the dropout for the layer. Need to dig a bit more.
Having said that, I think my PR will actually fix the issue in https://github.com/keras-team/keras-core/issues/322. Due to the stateless scope in the jax.lax.scan function, it will only read the staled variable value, which is probably why it is not trained properly.
So the stateless RNN test failure should be addressed by https://github.com/keras-team/keras-core/pull/924, and the test for this PR should pass now.
Update the JAX trainer to use less memory during fit/eval.
Currently we keep 3 copies of all the model variable state during training:
the 2 and 3 will keep getting update during the training process, but 1 will be a stale copy and unnecessarily occupying the heap size. In the large model case, this will be huge.
This PR will purge the 1, and it will restore the KerasVariable at the end of the epoch. This save the memory size by 33%. From some early test result internally, the per device memory usage is reduce from 9.49G to 6.65G for a OPT2 model.
I will send follow up PR to address additional memory usage for eval and predict functions.