# Just-in-time Quantization with Processing-In-Memory for Efficient ML Training

arXiv: https://arxiv.org/abs/2311.05034 · HTML: https://mlsanity.com/paper/2311.05034/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 52 references that are in the map.
- First posted on arXiv 8 Nov 2023.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Multimodal Few-Shot Learning with Frozen Language Models](https://mlsanity.com/paper/2106.13884/) (2021)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [VL-Adapter: Parameter-Efficient Transfer Learning for Vision-and-Language Tasks](https://mlsanity.com/paper/2112.06825/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [A Study of BFLOAT16 for Deep Learning Training](https://mlsanity.com/paper/1905.12322/) (2019)
- [FP8 Formats for Deep Learning](https://mlsanity.com/paper/2209.05433/) (2022)
- [Training Deep Neural Networks with 8-bit Floating Point Numbers](https://mlsanity.com/paper/1812.08011/) (2018)
- [Microscaling Data Formats for Deep Learning](https://mlsanity.com/paper/2310.10537/) (2023)
- [Structured Pruning of Deep Convolutional Neural Networks](https://mlsanity.com/paper/1512.08571/) (2015)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [ZeRO++: Extremely Efficient Collective Communication for Giant Model Training](https://mlsanity.com/paper/2306.10209/) (2023)
- [With Shared Microexponents, A Little Shifting Goes a Long Way](https://mlsanity.com/paper/2302.08007/) (2023)
- [Compressing DMA Engine: Leveraging Activation Sparsity for Training Deep Neural Networks](https://mlsanity.com/paper/1705.01626/) (2017)
- [8-bit Numerical Formats for Deep Neural Networks](https://mlsanity.com/paper/2206.02915/) (2022)
- [Mixed Precision Training With 8-bit Floating Point](https://mlsanity.com/paper/1905.12334/) (2019)
- [MiCS: Near-linear Scaling for Training Gigantic Model on Public Cloud](https://mlsanity.com/paper/2205.00119/) (2022)
- [Mokey: Enabling Narrow Fixed-Point Inference for Out-of-the-Box Floating-Point Transformer Models](https://mlsanity.com/paper/2203.12758/) (2022)
- [Computation vs. Communication Scaling for Future Transformers on Future Hardware](https://mlsanity.com/paper/2302.02825/) (2023)

## Built on by (most cited first, 1 of 1)

- [BitSnap: Checkpoint Sparsification and Quantization in LLM Training](https://mlsanity.com/paper/2511.12376/) (2025)

## Abstract (arXiv)

Data format innovations have been critical for machine learning (ML) scaling, which in turn fuels ground-breaking ML capabilities. However, even in the presence of low-precision formats, model weights are often stored in both high-precision and low-precision during training. Furthermore, with emerging directional data formats (e.g., MX9, MX6, etc.) multiple low-precision weight copies can be required. To lower memory capacity needs of weights, we explore just-in-time quantization (JIT-Q) where we only store high-precision weights in memory and generate low-precision weights only when needed. To perform JIT-Q efficiently, in this work, we evaluate emerging processing-in-memory (PIM) technology to execute quantization. With PIM, we can offload quantization to in-memory compute units enabling quantization to be performed without incurring costly data movement while allowing quantization to be concurrent with accelerator computation. Our proposed PIM-offloaded quantization keeps up with GPU compute and delivers considerable capacity savings (up to 24%) at marginal throughput loss (up to 2.4%). Said memory capacity savings can unlock several benefits such as fitting larger model in the same system, reducing model parallelism requirement, and improving overall ML training efficiency.