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

数据格式的创新对于机器学习（ML）的扩展至关重要，这进而推动了具有突破性的ML能力。然而，在低精度格式存在的情况下，模型权重在训练过程中通常同时存储在高精度和低精度中。此外，随着新兴的定向数据格式（例如MX9、MX6等），可能需要多个低精度权重副本。为了降低权重的存储需求，我们探索了即时量化（JIT-Q），即只在内存中存储高精度权重，并仅在需要时生成低精度权重。为了有效地执行JIT-Q，在这项工作中，我们评估了新兴的内存处理（PIM）技术来执行量化。通过PIM，我们可以将量化卸载到内存计算单元中，使量化能够在不产生昂贵的数据移动的情况下执行，同时允许量化与加速器计算并发进行。我们提出的PIM卸载量化可以跟上GPU计算，并在较小的吞吐量损失（最多2.4\%）的情况下实现了相当大的容量节省（高达24\%）。所述的内存容量节省可以解锁多种好处，例如在同一系统中放置更大的模型，减少模型并行性要求，并提高整体ML训练效率。

论文旨在通过just-in-time quantization（JIT-Q）来降低模型权重的内存需求，探索使用处理内存（PIM）技术执行量化以提高效率。

使用PIM技术将量化操作转移到内存中进行，避免了昂贵的数据移动，同时允许量化与加速器计算并发进行，从而实现JIT-Q。该方法在GPU计算的同时实现了显著的内存容量节省（高达24%），对吞吐量的损失也很小（高达2.4%）。

实验设计了使用PIM技术实现JIT-Q的方法，并将其与GPU计算进行了比较。结果表明，该方法可以实现显著的内存容量节省，并且对吞吐量的损失也很小。这可以为模型训练带来多个好处，例如在同一系统中安装更大的模型、减少模型并行性要求，提高整体ML训练效率。

相关论文包括： 1. 'Training Deep Neural Networks with 8-bit Floating Point Numbers' by Suyog Gupta, Ankur Agrawal, Kailash Gopalakrishnan, and Pritish Narayanan from NVIDIA Corporation. 2. 'Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference' by Matthieu Courbariaux, Yoshua Bengio, and Jean-Pierre David from Université de Montréal.