Reproduce “SubTrack++ : Gradient Subspace Tracking for Scalable LLM Training” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Sahar Rajabi, Nayeema Nonta, Sirisha Rambhatla
A method developed by researchers at the University of Waterloo, SubTrack++, significantly improves the scalability of Large Language Model (LLM) training by integrating Grassmannian subspace tracking with a projection-aware optimizer and recovery scaling. This approach achieves superior convergence and reduces pre-training wall-time by up to 67% on Llama models while maintaining memory efficiency.
Key insights
- Gradient subspaces in LLM training can be efficiently and robustly tracked using geometric methods on Grassmannian manifolds, providing more stable and controlled updates than SVD-based approaches.
- Maintaining proper alignment between an adaptive optimizer's internal states and an evolving low-rank gradient subspace is crucial for preserving and enhancing performance during training.
- Information lost through low-rank gradient projection can be effectively recovered by leveraging the scaling behavior of adaptive optimizers, thereby boosting model performance without incurring significant additional cost.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.