Standard mixed-precision training of neural networks requires many bytes of accelerator memory for each model parameter. These bytes reflect not just the parameter itself, but also its gradient and one or more optimizer state variables. With each of these values typically requiring 4 bytes, training even a 7 billion parameter model can be impractical for researchers with less than 100GB of accelerator memory. We introduce FlashOptim, a suite of optimizations that reduces per-parameter memory by over 50% while preserving model quality and API compatibility. Our approach introduces two key techniques. First, we improve master weight splitting by finding and exploiting a tight bound on its quantization error. Second, we design companding functions that greatly reduce the error in 8-bit optimizer state quantization. Together with 16-bit gradients, these techniques reduce AdamW memory from 16 bytes to 7 bytes per parameter, or 5 bytes with gradient release. They also cut model checkpoint sizes by more than half. Experiments with FlashOptim applied to SGD, AdamW, and Lion show no measurable quality degradation on any task from a collection of standard vision and language benchmarks, including Llama-3.1-8B finetuning.

翻译：标准的神经网络混合精度训练要求每个模型参数占用加速器内存的多个字节。这些字节不仅反映参数本身，还包括其梯度以及一个或多个优化器状态变量。由于每个值通常需要4字节，训练一个70亿参数的模型对于加速器内存少于100GB的研究者而言往往难以实现。本文介绍FlashOptim，一套优化方案，可在保持模型质量和API兼容性的同时，将每个参数的内存占用降低50%以上。我们的方法引入了两项关键技术。首先，我们通过发现并利用主权重分量的量化误差紧界，改进了主权重分割技术。其次，我们设计了压扩函数，大幅降低了8位优化器状态量化的误差。结合16位梯度，这些技术将AdamW的内存占用从每个参数16字节降至7字节（若释放梯度则为5字节），同时将模型检查点大小缩减一半以上。将FlashOptim应用于SGD、AdamW和Lion优化器的实验表明，在包括Llama-3.1-8B微调在内的标准视觉与语言基准测试任务中，均未出现可观测的质量下降。