Repository navigation
Add schedule-free adamw submission in JAX #809
Description
Description
Currently we have been unable to reproduce the schedule free adamw results with JAX.
There seem to be differences between the optax implementation of schedule-free adamw and the pytorch submission.
Activity
adefazio commentedon Nov 1, 2024
I can help debug any issues here. Do you have any code you can share? If there are issues with the optax jax implementation I want to get it fixed asap.
adefazio commentedon Nov 6, 2024
There are many small differences between the behavior of schedule-free jax wrapper and the original algoperf submission. Some differences I'm aware of:
So overall I expect the jax wrapper version to give as good results on all problems (maybe slightly slower on fastmrI), so if there is a difference it would be from some sort of bug.
priyakasimbeg commentedon Nov 19, 2024
Hi Aaron! thanks for weighing in on this. I seemed to have missed your messages on this thread.
We have a slightly modified version based on the optax code here: https://github.com/priyakasimbeg/algorithmic-efficiency/blob/compare_schedule_free/tests/test_algorithms/schedule_free_adamw/jax/schedule_free_optax.py.
This code adds r and we tested it with 0.75 on our google internal codebase.
I'm working on a test to compare the pytorch and jax implementations side by side on the algoperf github code but the test is still in progress. I can perhaps run a full training run on some of the workloads.
But in the meantime feel free to weigh in again if you spot any other differences
adefazio commentedon Nov 19, 2024
Ok, I take a look and see if I spot any differences.
adefazio commentedon Nov 19, 2024
It looks like the z buffer my be initialized with zeros:
https://github.com/priyakasimbeg/algorithmic-efficiency/blob/5556015054e3dda681e2a25e05a2f217d933453d/tests/test_algorithms/schedule_free_adamw/jax/submission.py#L58C51-L59C1
It needs to be initialized the same as the main parameter buffer. I think this line is a copy-paste error from the Jax version of NAdamW and other methods, where all optimizer state is normally initialized at zero.
Suggestion: you might want to set z on the first call to the main optimizer update, that's what we do in the pytorch version.
adefazio commentedon Nov 22, 2024
@priyakasimbeg Let me know if that initialization issue was the problem.
priyakasimbeg commentedon Nov 25, 2024
Hi Aaron thanks for spotting that!
We did the end-to-end training with an internal codebase using the optax implementation: https://github.com/google-deepmind/optax/blob/3ba9822c2a8d5fa7d046180b2574e108094523b4/optax/contrib/_schedule_free.py. It's not immediately obvious to me how this is initialized but will investigate and report back here
init-22 commentedon Apr 10, 2025
Hello @adefazio
Just following up on this thread, me and @priyakasimbeg were actually debugging this issue and found that z buffer initialization to zeros might actually be the reason for this, instead now we are passing model_params as suggested and it seems to solve the problem...
This is how the val loss of Jax vs Pytorch looked when we initialize the z buffer as the above lines (https://github.com/priyakasimbeg/algorithmic-efficiency/blob/5556015054e3dda681e2a25e05a2f217d933453d/tests/test_algorithms/schedule_free_adamw/jax/submission.py#L58C51-L59C1)
and This is how it looks now after passing the model_params in opt_init_fn function:
FYI, We are still going to run some more tests on it and observe the difference
adefazio commentedon Apr 10, 2025
Awesome, let me know if there are any other issues I can help debug.
wyfEmma commentedon Oct 21, 2025
I ran serveral other workloads comparing the diff between the jax implementation with proper initialized z-buffer and the pytorch implementation.
The result we see shows diff for different workloads, here's the plot bwteen global_step and log-scaled validation loss:
will look into it more to find where the diff comes from
wyfEmma commentedon Nov 6, 2025
interestingly for imagenet pytorch is unstable:
wyfEmma commentedon Mar 2, 2026
we have a finding that in the previous implementation, y is used in both training and evaluation phase, after fixing, here's the latest result(I also removed the grad clipping for jax because pytorch doesn't have it):
wyfEmma commentedon Mar 23, 2026
after fixing the replication error and memory issue:
wyfEmma commentedon Mar 23, 2026
wyfEmma commentedon Mar 24, 2026
wyfEmma commentedon Mar 24, 2026
wyfEmma commentedon Mar 25, 2026
wyfEmma commentedon Apr 21, 2026
ok for the previous results, librispeech deepspeech was using ctc_loss, it looks very much like a scaling issue within ctc_loss, here's the plot for wer, it shows that the two curves are aligning
wyfEmma commentedon May 10, 2026
to see why ctc diff between Jax and pytorch in librispeech deepspeech: #924
wyfEmma commentedon May 10, 2026
why jax performed better than pytorch in finewebedulm?
I noticed in pytorch, there's a hyperparam r that controls as steps increase, how much z got averaged into x. if r>0, with step increase, z will be taken into account more when calculating x. pytorch set it to be 0.75 and jax set it as 0(no such logic)
but after i align the r hyperparameter, and the batch size, the diff is still there: