mlcommons / mlcommons/submissions_algorithms
schedule free adamw jax training results doesn't match pytorch
Nobody has claimed this yet.
- Dominant language
- Jupyter Notebook
- Stars
- 11
- Forks
- 11
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 4
Description
The JAX implementation of the schedule-free AdamW algorithm exhibits significant training curve discrepancies compared to the PyTorch reference across several workloads, including Librispeech DeepSpeech, ImageNet, WMT, and Criteo1TB.
Following an initial debugging session with @priyakasimbeg, a key issue was identified: the JAX code incorrectly used a single variable (y) for both training and validation phases. The intended logic requires using x for validation and y for training.
A fix was issued in pull request #16. However, the training results for Librispeech still do not align with PyTorch, and other issues have emerged for WMT and Criteo1TB, specifically deadlocks and out-of-memory errors.
Further in-depth debugging is necessary to bring the JAX training results in line with PyTorch to finalize this MLCommons submission.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Review the JAX implementation and pull request #16 first, then compare the remaining Librispeech, WMT, and Criteo1TB behavior with the PyTorch reference. Done means matching training results across the listed workloads without the reported deadlocks or out-of-memory failures.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100