[Paper] 异构低带宽预训练 LLMs

Source: arXiv - 2601.02360v1

Overview

训练当今的大型语言模型（LLM）需要海量的分布式计算，但模型并行通信所需的网络带宽很快就会成为瓶颈——尤其是在设备不完善的数据中心之外。本文研究了如何结合两种低通信技巧——SparseLoCo（不频繁、稀疏的梯度同步）和流水线并行中的激活‑梯度压缩，使得异构硬件（高速节点加上带宽受限的参与者）能够共同预训练 LLM，而不会过多牺牲模型质量。

关键贡献

• 异构训练框架：在高带宽节点上混合使用完整模型副本，在低带宽节点上通过多个参与者构建流水线并行副本。
• 将 SparseLoCo（稀疏、低频梯度交换）与流水线并行中使用的子空间投影激活/梯度压缩相结合。
• 选择性压缩策略：仅对带宽受限的流水线副本进行通信压缩，而完整副本保持未压缩通信。
• 实证验证：在参数规模从 1.78 亿到 10 亿的语言模型预训练任务上进行实验，展示了适度的开销和改进的损失‑通信权衡。
• 部署指南：为在真实异构计算集群中实际部署低带宽模型并行提供实用建议。

方法论

1. SparseLoCo recap – 与在每个小批次后同步完整稠密梯度不同，每个工作节点只发送一个 伪梯度，该伪梯度 (a) 经过稀疏化（仅保留 top‑k 条目）且 (b) 仅每 N 步交换一次。这样大幅削减了跨网络传输的数据量。

2. Pipeline parallelism with compression – 将模型划分为多个阶段；每个阶段在不同的设备上运行。前向激活和反向梯度通常以全精度在阶段之间传递。作者采用 子空间投影：将激活投影到低维基上（例如通过随机高斯矩阵），传输后在接收端重建。梯度使用相同的投影方式。

3. Heterogeneous composition –

   • High‑bandwidth nodes 保持模型的完整副本，并使用标准（未压缩）的数据并行更新。
   • Low‑bandwidth nodes 被聚合在一起；它们共同形成一个 虚拟 副本，采用流水线并行。它们的阶段间消息通过子空间投影进行压缩。
   • 两组通过 SparseLoCo 的稀疏同步共享相同的优化器状态，从而保证整体训练的一致性。

4. Adaptations for compatibility – 作者对投影矩阵以及 SparseLoCo 同步的时机进行微调，以避免陈旧更新并保持压缩流水线误差的有界性。

结果与发现

• 激活压缩 与 SparseLoCo 紧密配合：子空间投影引入的额外误差并未显著降低模型质量。
• 选择性压缩（仅对 pipeline 副本）始终优于“全部压缩”方案，尤其在高压缩比（≥16×）时表现更佳。
• 每个 epoch 的训练时间与通信节省成比例提升，投影操作带来的计算开销仅略有增加（约 2‑3 % 的额外开销）。

实际影响

• 成本效益的扩展 – 组织可以利用更便宜、带宽受限的硬件（例如边缘服务器、旧的 GPU 集群）来参与 LLM 预训练，减少对昂贵高速互连的依赖。
• 混合云/本地部署 – 数据中心拥有少量高速节点可作为“锚”，而大量低成本实例运行流水线阶段，实现更灵活的资源分配。
• 节能 – 网络中传输的数据减少，网络设备的功耗降低，符合可持续发展目标。
• 易于集成 – 该框架基于现有的 PyTorch 风格的数据并行和流水线并行 API；开发者只需指定哪些工作节点属于压缩流水线组。

限制与未来工作

• 投影开销 随模型深度增加；对于极深的模型，额外计算可能抵消通信收益。
• 本研究聚焦于 预训练 语言模型；在异构压缩下的微调动态尚未探讨。
• 共享压缩激活跨不可信节点的安全性和隐私影响未被考虑。
• 未来研究可探索自适应压缩比（随层或训练阶段变化）以及对稀疏梯度与子空间压缩更新所引入误差的更紧理论界限。

作者

• Yazan Obeidi
• Amir Sarfi
• Joel Lidin
• Paul Janson
• Eugene Belilovsky

论文信息

• arXiv ID: 2601.02360v1
• 分类: cs.LG
• 出版日期: January 5, 2026
• PDF: 下载 PDF