一种新方法让大语言模型训练提速 37%,减少 34% 训练时间
MoARa: Module-Aware Rank Allocation and Structure-Preserving Decomposition for Low-Rank LLM Pre-training
朋友A对朋友B说:刚看到个挺有意思的论文,叫MoARa,是DeepSeek团队搞的,专门针对大语言模型训练优化,能让训练速度提升37%,时间缩短34%,挺实用的。
研究人员提出 MoARa 方法,通过模块感知的投影秩分配和块级分解来优化低秩 LLM 预训练。在 Llama 2 7B 模型上,该方法使 GaLore 达到相同 perplexity 所需步骤减少 37%,训练时间减少 34%,同时仅增加 0.2% 的峰值内存开销。
MoARa: Module-Aware Rank Allocation and Structure-Preserving Decomposition for Low-Rank LLM Pre-training
Low-rank gradient projection reduces the optimizer-state memory cost of large language model (LLM) pretraining, but the steps and wall-clock time needed to reach a target quality remain a meaningful axis for improvement. We attribute this to two design choices in existing methods: the projection-rank budget is allocated uniformly across Transformer modules with heterogeneous projection sensitivity, and projecting a raw gradient attenuates its magnitude and direction jointly. We propose MoARa, which combines a static profiling-based module-aware projection-rank allocation with a block-wise magnitude-direction decomposition; the default block size is set in the neighborhood of the attention head dimension. Across five Transformer architectures spanning Llama, Qwen, and DeepSeek at 300M to 7B scales, GaLore with MoARa reaches standard GaLore's final perplexity in 37% fewer steps and 34% less wall-clock time on Llama 2 7B, with only 0.2% peak reserved memory overhead under standard graph compilation. Across the six low-rank pretraining methods we evaluate, module-aware rank allocation alone delivers directionally consistent step reductions on all six. On compatible hosts, the two-component design reaches up to 41.7% step reduction and 37.1% wall-clock reduction.