MuLoCo: Muon is a practical inner optimizer for DiLoCo

DiLoCo 是一个强大的框架，用于在网络约束下训练大型语言模型（LLMs），并在数据中心环境中提升并行性和加速器利用率。然而，尽管 DiLoCo 显著减少了通信频率，其通信步骤仍然需要对模型参数的完整副本进行全归约（all-reduce）。虽然已有研究探讨了减少 DiLoCo 通信开销的方法，但误差反馈累加器的作用以及内部优化器对可压缩性的影响仍需进一步探索。在本工作中，我们研究了标准压缩方法（包括 Top-k 稀疏化和量化）在与两种本地优化器（AdamW 和 Muon）结合时，减少 DiLoCo 通信开销的有效性。我们在预训练仅解码器变压器语言模型（LMs）的实验中发现，使用 Muon 作为 DiLoCo 的内部优化器，并结合误差反馈累加器，可以将通信的参数变化量（delta）激进地压缩到 2 位，且几乎不会导致性能下降。至关重要的是，MuLoCo（以 Muon 为内部优化器的 DiLoCo）在通信量减少 8 倍、内存复杂度相同的情况下，显著优于原始的 DiLoCo。

该论文试图解决在分布式训练大语言模型（LLMs）时的通信瓶颈问题，特别是在DiLoCo框架下如何进一步减少通信开销而不影响模型性能。这是一个长期存在的问题，但针对DiLoCo框架下的压缩优化仍是一个相对较新的研究方向。

论文提出通过结合Muon优化器和误差反馈累积器来增强DiLoCo框架的效果。与传统方法不同，它利用了2-bit量化和Top-k稀疏化等压缩技术，在大幅减少通信量的同时保持模型性能。此外，MuLoCo（基于Muon的DiLoCo）实现了比原版DiLoCo更高的效率，同时通信量减少了8倍且内存复杂度保持不变。

1. Muon优化器在减少通信需求方面表现出色，尤其是在结合误差反馈机制后；2. 实验使用了预训练解码器-only Transformer语言模型进行验证；3. 提出了MuLoCo，显著优于原始DiLoCo；4. 论文未明确提及是否开源代码，但其设计为未来研究提供了清晰的方向，例如探索其他优化器或更高效的压缩策略。

相关研究包括：1. Deep Gradient Compression (DGC)，探索梯度稀疏化和量化以减少通信；2. SignSGD，仅传输梯度符号以降低带宽需求；3. PowerSGD，通过低秩分解减少通信量；4. DiLoCo本身的研究及其变体，如引入不同的局部优化器或调整同步频率。这些方法共同构成了分布式训练中通信优化领域的前沿研究。