Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo

随着我们扩展到更大规模的机器学习模型，数据并行方法中固有的频繁同步需求会引发显著的性能下降，这对进一步扩展构成了关键挑战。近期的研究提出了一种方法（DiLoCo），该方法可以在不损害模型质量的前提下减少同步需求。然而，这些研究并未仔细分析 DiLoCo 的行为如何随模型规模变化。在本研究中，我们探讨了在固定计算预算下训练大语言模型时 DiLoCo 的扩展规律行为。我们重点关注算法因素（包括模型副本数量、超参数和 token 预算）如何以可以通过扩展规律准确预测的方式影响训练过程。我们发现，DiLoCo 在模型规模变化时表现出可预测且稳健的扩展能力。当调优得当时，DiLoCo 的扩展性能优于数据并行训练，并且即使在较小模型规模下也能超越数据并行训练的表现。我们的结果揭示了 DiLoCo 相较于之前记录的更广泛的优势，包括更大的最佳批量大小、随着规模增加下游任务泛化性能的提升，以及在固定 token 预算下的改进评估损失。

该论文试图解决在大规模机器学习模型训练中，数据并行方法因频繁同步需求导致的性能瓶颈问题。这是一个长期存在的挑战，尤其是在超大规模语言模型（LLMs）的训练中。

论文提出通过研究DiLoCo方法的行为来优化大规模模型训练，特别是在固定计算预算下。关键思路是分析算法因素（如模型副本数量、超参数和令牌预算）如何影响训练，并利用扩展定律预测这些影响。与传统数据并行方法相比，DiLoCo通过减少同步需求，能够在更大规模模型上表现更优，甚至在小规模模型中也具有竞争力。

论文展示了DiLoCo方法在多个方面的优势：1) 可以显著增加最优批量大小；2) 随着模型规模扩大，下游任务泛化能力更强；3) 在固定令牌预算下，评估损失更低。实验设计涉及不同模型大小和计算预算条件下的对比测试，并且验证了DiLoCo方法的可预测性和鲁棒性。虽然未明确提及，但通常此类研究会基于公开的大规模语言模型数据集进行，开源代码可能会进一步促进后续研究。未来可以探索DiLoCo与其他优化技术结合的可能性。

相关研究包括但不限于以下几项工作：1) PipeDream（PipeDream: Decoupled Model Parallelism for Training Giant Models, 2020），研究了流水线并行策略；2) ZeRO（ZeRO: Memory Optimization Towards Training Trillion Parameter Models, 2019），专注于内存优化以支持更大模型；3) GShard（GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, 2020），提出了自动分片技术以提高分布式训练效率。此外，还有一些研究探讨了异步更新机制对大规模训练的影响，例如Async-SGD等。