论文

OASIS:用于内存高效训练的在线激活子空间学习

OASIS: Online Activation Subspace Learning for Memory-Efficient Training

模型训练参数高效训练

摘要

训练 大语言模型 (LLM) 受到内存要求的限制,激活占总占用空间的很大一部分。现有方法使用低秩权重参数化或优化器状态的低秩梯度子空间来减少内存,而激活内存则通过基于定期更新的投影的架构修改或压缩方案来解决。我们提出了 OASIS,一种用于内存高效训练的在线激活子空间学习算法,可在训练期间跟踪并持续更新低维激活子空间。中间激活被投影到这个不断发展的子空间上,在不修改前向计算的情况下减少内存。不断发展的激活子空间会产生低秩梯度表示,使梯度和优化器状态都可以直接在该子空间中维护,而投影感知优化器在子空间更新中一致地传输优化器状态以实现稳定的训练。在各种微调和预训练任务中,OASIS 的峰值内存比完整的 微调 低 2 倍,同时匹配其性能并优于先前的低秩方法。