GaLore：在消费级硬件上推进大模型训练

GaLore 通过显著降低优化器状态的内存占用，使得在消费级硬件（如 NVIDIA RTX 4090）上训练参数量高达 70 亿的大型语言模型（LLM）成为可能。这种 AI 研究的民主化让从业者无需高端工业计算资源即可训练大规模模型。

通过低秩梯度投影实现内存效率

GaLore 通过在梯度被优化器处理之前将其投影到低维子空间，从而降低内存消耗。这一方法利用深度神经网络中梯度固有的低秩结构，最小化训练过程中需要存储和操作的数据量。

对于 Adam 等自适应优化算法，优化器状态通常占据内存的大部分。通过应用此投影，GaLore 在训练期间实现了对优化器状态存储内存超过 82.5% 的降低。

动态子空间切换

为了保持全参数学习的能力，避免模型被限制在参数空间的有限部分，GaLore 使用了动态子空间切换机制。该机制使模型在整个训练过程中能够在不同的低秩子空间之间切换。

这些切换的频率经过平衡，以确保在适应梯度不断演化的低秩结构的同时保持一致的优化轨迹。这使得在内存效率与优化性能之间的权衡能够得到细致的控制。

与 8 位优化器的集成

将 GaLore 与 8 位精度优化器结合，可通过量化优化器状态进一步最大化内存效率。这种协同作用使得在相同硬件限制下能够训练更大的模型或使用更大的批量大小，而不影响模型精度或收敛速度。

GaLore 与 8 位优化的算法流程

  1. 梯度投影：使用投影矩阵将全精度梯度投影到低秩子空间，然后量化为 8 位格式。
  2. 量化：将投影后的梯度、模型权重以及优化器状态（如 Adam 的动量均值）从 32 位浮点量化为 8 位整数表示。
  3. 优化器更新：对 8 位量化的梯度进行更新；这包括将梯度反量化为浮点数，应用更新规则，然后将更新后的优化器状态重新量化为 8 位。
  4. 反量化与权重更新：权重被反量化为浮点数进行处理。随后 GaLore 使用最终投影，将反量化后的低秩更新映射回原始参数空间，随后再应用权重更新。

在 Hugging Face Transformers 中的实现

GaLore 已集成到 Hugging Face transformers 库（版本 4.39.0 或更高）以及 galore-torch 库中。用户可以在 TrainingArguments 中指定优化器选项，如 galore_adamw、galore_adamw_8bit 或 galore_adafactor，并通过 optim_target_modules 定义目标模块，以实现 GaLore。

分层更新

为进一步降低内存占用，GaLore 支持分层权重更新。优化器不再在反向传播后一次性更新所有层，而是使用 PyTorch 的后累加钩子逐层更新权重。通过在优化器名称后添加 _layerwise（例如 galore_adamw_layerwise）即可启用此功能。

Sources