NVIDIA / NVIDIA/TransformerEngine
[JAX] Upstream checkpointing of quantizations in TE/JAX to MaxText
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.5k
- Forks
- 831
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 65
Description
MaxText upstream already supports TE quantization, but for optimal performance for NVFP4 and MXFP8, quantization checkpointing support should be upstreamed to MaxText to prevent 2x quantization in fwd and rematerialized in backward,
This work was initially started here but needs continuing: https://github.com/AI-Hypercomputer/maxtext/pull/2773
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
Start by reviewing the linked MaxText pull request #2773 and the existing TE quantization support in MaxText. Trace how NVFP4 and MXFP8 quantizations are checkpointed during forward and rematerialized backward passes. Done means the support is upstreamed so these paths avoid the reported 2x quantization.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100