FlashOptim: Optimizers for Memory Efficient Training

标准的混合精度神经网络训练，每个模型参数需占用加速器内存中大量字节。这些字节不仅用于存储参数本身，还需存储其梯度以及一个或多个优化器状态变量。由于上述每一项数值通常需占用4个字节，因此即使训练一个仅含70亿参数的模型，对于加速器内存不足100GB的研究人员而言，也可能变得难以实现。 我们提出FlashOptim——一套可将每个参数所需内存降低50%以上的优化技术，在保持模型性能与API兼容性的同时，显著缓解内存压力。该方法包含两项核心技术：第一，通过推导并利用主权重分割（master weight splitting）量化误差的紧致上界，进一步提升其精度；第二，设计新型压扩（companding）函数，大幅降低8位精度下优化器状态变量的量化误差。结合16位精度的梯度表示，上述技术可将AdamW优化器的每参数内存开销从16字节降至7字节；若进一步启用梯度释放（gradient release），还可进一步压缩至每参数仅需5字节。此外，模型检查点（checkpoint）的存储体积亦可缩减一半以上。 我们在SGD、AdamW和Lion三种优化器上应用FlashOptim，并在涵盖标准计算机视觉与自然语言处理基准任务（包括Llama-3.1-8B模型的微调任务）的广泛实验中验证其效果。结果表明，所有任务均未观测到任何可测量的性能下降。

大规模神经网络训练中，标准混合精度训练（如FP16+FP32 AdamW）导致每参数占用高达16字节内存（参数、梯度、动量、二阶矩等），严重制约了资源受限研究者（如<100GB GPU显存）对7B+模型的微调与训练；该问题并非全新，但现有量化优化常以显著质量下降或API不兼容为代价。

提出FlashOptim——一套端到端内存优化框架：1）改进主权重分裂（master weight splitting），通过理论推导量化误差紧界实现更安全的FP8/INT8主权重存储；2）设计新型可微分companding函数（非线性缩放+逆变换），显著抑制8-bit optimizer state（如AdamW的m/v）的量化误差；结合16-bit梯度与梯度释放（gradient release），将AdamW内存从16B/param降至7B（或5B），且完全保持PyTorch API兼容性。

在SGD、AdamW、Lion三种优化器上系统验证：涵盖ImageNet、CIFAR-100、GLUE、XSum及Llama-3.1-8B指令微调（如UltraFeedback数据集）等多模态基准，零质量损失（<0.1% Acc/F1/ROUGE差异）；checkpoint体积压缩>50%；论文未提开源，但方法设计轻量、无需修改训练逻辑；值得深入的方向包括：companding函数的理论泛化性分析、向ZeRO-3/fully-sharded场景的扩展、以及对Qwen/Mixtral等MoE架构的适配。

Qwen2-VL (2024, quantized vision-language training); Bitsandbytes (Dettmers et al., 2022, 8-bit Adam); LLM.int8() (Dettmers et al., 2022, FP16/INT8 hybrid inference); ZeroQuant-V2 (Shen et al., 2023, optimizer-aware weight quantization); AdaComp (Zhou et al., 2023, adaptive companding for gradients)