用新方法让训练模型少占内存，70亿参数也能跑。
FlashOptim: Optimizers for Memory-Efficient Training
- 通过量化误差约束和压缩函数，优化权重与状态存储方式。
- 每参数内存从16字节降至7字节，检查点大小减半。
- 兼容主流优化器，适配资源有限的研究者使用。
神经网络的混合精度训练需要大量加速器内存，每个参数通常需4字节存储权重、梯度及一个或多个优化器状态变量。即使70亿参数模型，在不足100GB内存环境下也难以训练。我们提出FlashOptim，一套内存优化技术，使每参数内存减少超50%，同时保持模型质量与API兼容性。核心方法包括：改进主权重分片，发现并利用其量化误差的紧界；设计压缩函数，显著降低8位优化器状态量化误差。结合16位梯度，可将AdamW内存从16字节降至7字节（启用梯度释放后为5字节），模型检查点大小减少超过一半。在多个标准视觉与语言基准测试中，包括Llama-3.1-8B微调任务，应用FlashOptim于SGD、AdamW和Lion均未造成可测量的质量下降。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。