AdaPM通过部分动量机制，大幅降低大模型训练内存占用。
AdaPM: a Partial Momentum Algorithm for LLM Training
- 采用非均匀动量设计，仅对关键模块保留完整动量。
- 通过偏差校正技术，使部分动量仍保持优化性能。
- 在1.5B模型上节省超90%动量内存，预训练节省30%显存小时。
在大语言模型训练中，动量常用于加速优化过程，但存储动量会带来显著内存开销。本文提出AdaPM，一种基于部分动量的自适应训练策略，实现高效低内存优化器。AdaPM采用非均匀动量设计：多数模块无需完整动量即可维持优化性能。为缓解部分动量带来的偏差与性能损失，引入偏差校正技术。实验证明，该方法在60M至1.5B规模的语言模型预训练、监督微调及强化学习人类反馈（RLHF）中，动量内存减少超过90%，同时保持效率与性能。结合二阶统计量的内存优化技术，动量状态内存可进一步降低95%，在预训练GPT-2 1.5B时节省超30% GPU小时。
原文摘要 · Abstract (English)
In the training of large language models, momentum is widely used and often demonstrated to achieve significant acceleration. However, storing momentum typically presents memory challenges. In this paper, we propose AdaPM, an adaptive training strategy that leverages partial momentum to implement a memory-efficient optimizer. To this end, AdaPM utilizes a non-uniform momentum design: for most blocks, full momentum is not necessary to preserve the performance of the optimization. In the momentum design of AdaPM, to mitigate the bias and performance loss caused by partial momentum, we enhance the partial momentum by a bias correction technique. Empirically, we verify that our approach reduces memory by over $90\%$ in momentum while maintaining both efficiency and performance for pretraining various language models ranging from 60M to 1.5B, as well as for supervised fine-tuning and RLHF. AdaPM can further reduce memory by up to $95\%$ in optimizer states by combining the memory-efficient technique on the second-order statistic, saving over $30\%$ GPU hours for pretraining GPT-2 1.5B.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。