Fused AdamW

[AakashKumarNain]

Does optax provide fused implementation of adamw? If not, is there any reason why?

[vroulet]

Hello @AakashKumarNain,

No there is no ``fused'' implementation and none foreseen in the future.
Optax focuses on simple gradient transformation, no additional low-level optimization.
It could be interesting to first do benchmarks between pytorch and jitted compilations offered in JAX before trying to add a fuse implementation. You may also take a look at libraries specialized in e.g. transfomers for languages (maybe https://github.com/google/maxtext for example) to see best current practices.

[AakashKumarNain]

Thanks for the prompt response.

It could be interesting to first do benchmarks between pytorch and jitted compilations offered in JAX

Well, I asked because on my side the time difference between a torch run and jax run is drastic which is both surprising and concerning to me. And on jax side, the big overhead seems to lie in the backward pass.

Re Maxtext: Though it's a great library, but certainly not a readable one. But I can use as a ref for checking things on my end.

[vroulet]

Hello @AakashKumarNain,

Following on your question: are you using gpus or tpus?
A priori HLO/XLA should take care of fusing for loops etc... but it may only be optimized for tpus.

[AakashKumarNain]

Hey @vroulet

I am using A100 atm. I noticed that my backward pass is very slow. I have double checked everything on the model side, and I see no bottleneck anywhere except for the backward pass. If you are interested in benchmarking the code, I can send it to you.

[vroulet]

Thanks for the answer! A MWE for the record could be good to have but no emergency here, I may not have time to dig into that in the close future.