从 $O(mn)$ 到 $O(r^2)$：Adam 分布式训练中基于双侧低秩通信的内存高效方法

随着基础模型的持续规模化，预训练越来越依赖于数据并行分布式优化，使得带宽受限的梯度同步成为关键瓶颈。从另一个角度看，投影基的低秩优化器主要是为内存效率而设计，但在面对通信受限的训练时仍不够次佳：单向同步仍会传输 $O(rn)$ 大小的对象用于 $m\times n$ 矩阵梯度，刷新步骤也可能主导峰值通信字节。我们提出 TSR 方法，通过在 Adam 族更新中引入双侧低秩通信（TSR-Adam），通过同步紧凑核心 $U^\top G V\in\mathbb{R}^{r\times r}$ 来实现，从而将主要的每步负载从 $O(mn)$ 降低到 $O(r^2)$，同时保持矩状态在低维核心中。为进一步减少子空间刷新导致的峰值通信，TSR-Adam 采用基于随机奇异值分解的刷新，避免完整梯度的同步。我们还扩展了低秩通信到嵌入梯度，采用嵌入特定的秩和刷新方案，实现了相比保持嵌入稠密的额外通信和内存节省。对于从 60M 到 1B 模型规模的预训练，TSR-Adam 每步平均通信字节降低 13 倍；在 GLUE 微调任务中通信降低 25 倍，同时保持可比的性能；我们进一步提供了该方法更新的理论稳定性分析。代码已在 https://github.com/DKmiyan/TSR-Adam 上提供。