Reproduce “Just-in-time Quantization with Processing-In-Memory for Efficient ML Training” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Mohamed Assem Ibrahim, Shaizeen Aga, Ada Li, Suchita Pati, Mahzabeen Islam
Researchers at AMD developed Just-In-Time Quantization (JIT-Q) that uses Processing-in-Memory (PIM) to dynamically generate low-precision weight copies from a single high-precision master copy. This approach significantly reduces the memory footprint for large models like LLMs, achieving an average 12.5% memory capacity saving for BF16 to MX6 quantization with a 1.6% average training throughput loss.
Key insights
- Processing-in-Memory (PIM) units can effectively alleviate memory capacity bottlenecks in ML training by performing data format conversions directly within memory.
- Careful hardware-software co-design, including tailored data placement strategies and minimal PIM ALU augmentations, is essential to enable efficient on-demand quantization.
- Eliminating redundant low-precision weight copies via Just-In-Time Quantization (JIT-Q) yields substantial memory savings with only marginal impact on overall training throughput.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.