# From O(mn) to O(r2): Two-Sided Low-Rank Communication for Adam in Distributed Training with Memory Efficiency

arXiv: https://arxiv.org/abs/2602.08007 · HTML: https://mlsanity.com/paper/2602.08007/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 33 references that are in the map.
- First posted on arXiv 8 Feb 2026.

## Builds on (most cited first, 27 of 27)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [DoRA: Weight-Decomposed Low-Rank Adaptation](https://mlsanity.com/paper/2402.09353/) (2024)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training](https://mlsanity.com/paper/1712.01887/) (2017)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [8-bit Optimizers via Block-wise Quantization](https://mlsanity.com/paper/2110.02861/) (2021)
- [AdaLoRA: Adaptive Budget Allocation for Parameter-Efficient Fine-Tuning](https://mlsanity.com/paper/2303.10512/) (2023)
- [Error Feedback Fixes SignSGD and other Gradient Compression Schemes](https://mlsanity.com/paper/1901.09847/) (2019)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization](https://mlsanity.com/paper/1905.13727/) (2019)
- [1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed](https://mlsanity.com/paper/2102.02888/) (2021)
- [Memory-Efficient LLM Training with Online Subspace Descent](https://mlsanity.com/paper/2408.12857/) (2024)
- [Subspace Optimization for Large Language Models with Convergence Guarantees](https://mlsanity.com/paper/2410.11289/) (2024)
- [From Low Rank Gradient Subspace Stabilization to Low-Rank Weights: Observations, Theories, and Applications](https://mlsanity.com/paper/2407.11239/) (2024)
- [GaLore 2: Large-Scale LLM Pre-Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2504.20437/) (2025)
- [Greedy Low-Rank Gradient Compression for Distributed Learning with Convergence Guarantees](https://mlsanity.com/paper/2507.08784/) (2025)
- [Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining](https://mlsanity.com/paper/2502.05790/) (2025)
- [Natural GaLore: Accelerating GaLore for memory-efficient LLM Training and Fine-tuning](https://mlsanity.com/paper/2410.16029/) (2024)

## Abstract (arXiv)

As foundation models continue to scale, pretraining increasingly relies on data-parallel distributed optimization, making bandwidth-limited gradient synchronization a key bottleneck. Orthogonally, projection-based low-rank optimizers were mainly designed for memory efficiency, but remain suboptimal for communication-limited training: one-sided synchronization still transmits an O(rn) object for an m× n matrix gradient and refresh steps can dominate peak communicated bytes. We propose TSR, which brings two-sided low-rank communication to Adam-family updates (TSR-Adam) by synchronizing a compact core U^top G V∈ℝr× r, reducing the dominant per-step payload from O(mn) to O(r2) while keeping moment states in low-dimensional cores. To further reduce the peak communication from subspace refresh, TSR-Adam adopts a randomized SVD-based refresh that avoids full-gradient synchronization. We additionally extend low-rank communication to embedding gradients with embedding-specific ranks and refresh schedules, yielding additional communication and memory savings over keeping embeddings dense. Across pretraining from 60M to 1B model scales, TSR-Adam reduces average communicated bytes per step by 13×, and on GLUE fine-tuning it reduces communication by 25×, while achieving comparable performance; we further provide a theoretical stationarity analysis for the proposed update. Code is available at https://github.com/DKmiyan/TSR-Adam.