具有 4 位状态的存效优化器
机器学习
2023-10-30 v3 人工智能
摘要
优化器状态是训练神经网络内存消耗的主要来源,限制了给定内存预算内可训练的最大模型。将优化器状态从 32 位浮点压缩至更低比特宽度有望减少训练内存占用,而当前可实现的最低比特宽度为 8 位。在本工作中,我们通过针对一阶与二阶矩的详细实证分析将优化器状态比特宽度降至 4 位。具体而言,我们发现矩具有复杂的离群值模式,当前分块量化无法精确近似。我们采用更小的块尺寸,并提议利用行向与列向信息以实现更好量化。我们进一步识别出量化二阶矩的零点问题,并通过排除零点的线性量化器解决该问题。我们的 4 位优化器在自然语言理解、机器翻译、图像分类与指令微调等多种基准上进行了评估。在所有任务上,我们的优化器均能取得与全精度对应项相当的准确性,同时具备更优的内存效率。
引用
@article{arxiv.2309.01507,
title = {Memory Efficient Optimizers with 4-bit States},
author = {Bingrui Li and Jianfei Chen and Jun Zhu},
journal= {arXiv preprint arXiv:2309.01507},
year = {2023}
}
备注
v3: camera ready revisions for NeurIPS 2023