arXiv:2602.08007cs.LGcs.AI2026-02

将Adam优化器通信量从O(mn)降至O(r²),大幅节省分布式训练开销

From $O(mn)$ to $O(r^2)$: Two-Sided Low-Rank Communication for Adam in Distributed Training with Memory Efficiency

  • 双向低秩通信:只同步r×r核心矩阵,取代原始梯度
  • 通信量降低13倍(预训练)至25倍(微调),性能相当
  • 支持嵌入层低秩压缩,适合大模型高效训练场景

随着基础模型持续扩大,预训练日益依赖数据并行分布式优化,导致带宽受限的梯度同步成为主要瓶颈。投影式低秩优化器虽以内存效率为目标,但在通信受限训练中仍不理想:单边同步仍需传输O(rn)大小的梯度,且刷新步骤占主导通信量。本文提出TSR,首次将双向低秩通信引入Adam族更新(TSR-Adam),通过同步紧凑的核心矩阵UᵀGV∈ℝʳˣʳ,将每步主通信负载从O(mn)降至O(r²),同时保持动量状态在低维核心中。为减少子空间刷新带来的峰值通信,TSR-Adam采用基于随机SVD的刷新机制,避免完整梯度同步。此外,我们还将低秩通信扩展至嵌入梯度,采用嵌入特异性秩和刷新调度,相较保留嵌入稠密表示进一步降低通信与内存开销。在60M至1B规模的预训练中,TSR-Adam平均每步通信量减少13倍;在GLUE微调任务中,通信量减少25倍,且性能相当。我们还提供了所提更新的理论平稳性分析。代码已开源:https://github.com/DKmiyan/TSR-Adam。

原文摘要 · Abstract (English)

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\times 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\in\mathbb{R}^{r\times r}$, reducing the dominant per-step payload from $O(mn)$ to $O(r^2)$ 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\times$, and on GLUE fine-tuning it reduces communication by $25\times$, 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.

分布式训练低秩通信Adam优化器内存效率

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。