Low-rank rank allocation speeds up large language model training
MoARa: Module-Aware Rank Allocation and Structure-Preserving Decomposition for Low-Rank LLM Pre-training
Machine LearningComputation and Language
Summary
Training large language models takes a lot of computer memory and time. The authors found that spreading computer effort evenly across parts of the model doesn't work well because different parts affect the training differently. They made a new method called MoARa that gives more effort to important parts and better handles how changes are made during training. This method lets models like Llama 2 reach good quality faster while using almost the same amount of memory.
What this means in practice
- •For machine learning engineers: Speed up large language model pretraining by allocating computation to more sensitive model parts, reducing training steps and wall-clock time.
- •For ai infrastructure teams: Lower memory demands during training of large transformers by integrating module-aware low-rank gradient projections with minimal overhead.
Authors
Keunyoung Kim, Nojun Kwak
Abstract
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.