论文

MoARa:低秩 LLM 预训练的模块感知等级分配和结构保留分解

MoARa: Module-Aware Rank Allocation and Structure-Preserving Decomposition for Low-Rank LLM Pre-training

模型训练参数高效训练

摘要

低秩梯度投影降低了大型语言模型 (LLM) 预训练的优化器状态内存成本,但达到目标质量所需的步骤和挂钟时间仍然是有意义的改进轴。我们将其归因于现有方法中的两个设计选择:投影等级预算在具有异构投影敏感性的 Transformer 模块之间均匀分配,并且投影原始梯度共同衰减其大小和方向。我们提出了 MoARa,它将基于静态分析的模块感知投影秩分配与逐块幅度方向分解相结合;默认块大小设置在注意力头维度的附近。在涵盖 Llama、Qwen 和 DeepSeek 300M 至 7B 规模的五种 Transformer 架构中,GaLore 与 MoARa 在 Llama 2 7B 上以减少 37% 的步骤和减少 34% 的挂钟时间达到标准 GaLore 的最终困惑度,在标准图编译下峰值保留内存开销仅为 0.2%。在我们评估的六种低秩预训练方法中,仅模块感知等级分配就可以在所有六种方法上提供方向一致的步骤减少。在兼容主机上,两组件设计可减少高达 41.7% 的步长和 37.1% 的时钟减少。