[2311.08105] DiLoCo: Distributed Low-Communication Training of Language Models

文章提出了一种分布式优化算法DiLoCo，能够在连接不佳的设备集群上训练语言模型，通过减少通信频率，在C4数据集上取得了优异的效果，展现出良好的鲁棒性。

  1. 引言：大规模语言模型在机器学习应用中至关重要，但标准训练方法存在诸多挑战，如需要大量紧密连接的加速器，设备间同步困难且易受设备故障影响，难以利用异构设备。受联邦学习启发，文章提出DiLoCo算法，通过调整联邦平均算法，使用AdamW作为内优化器、Nesterov动量作为外优化器，减少通信次数，解决上述问题。实验表明，DiLoCo在C4数据集上通信量减少500倍的情况下，性能与全同步优化相当 ，且对数据分布和资源变化具有鲁棒性。

  2. DiLoCo算法：基于具有参数的基础模型架构，数据集被划分为个分片 ，由个工作节点分别处理。训练过程包含外优化和内优化。外优化进行步，每步收集各工作节点的外梯度，平均后用外优化器更新共享参数副本，再将其分发给各工作节点。内优化中，每个工作节点在各自数据分片上用内优化器AdamW独立并行地进行步优化，通常较大（如几百步），以此减少工作节点间的通信。整体上，DiLoCo可视为一种低通信的数据并行方法，适用于连接不佳的工作节点。

  3. 实验：在C4数据集上进行语言建模实验，评估模型在验证集上的困惑度。对比了DiLoCo与四个基线模型，结果显示DiLoCo在通信量减少500倍的同时，困惑度更低，泛化性能更好。

    • 消融实验：研究了预训练步数、通信频率、数据分布、副本数量、模型大小、外优化器等因素对模型性能的影响。发现DiLoCo对预训练步数不敏感，通信频率时能平衡性能和通信成本，对不同数据分布具有鲁棒性，增加副本数量可提升性能但存在收益递减，模型越大性能越好，Nesterov动量作为外优化器效果最佳。
    • 其他实验：探索了DiLoCo在计算资源变化、异步通信和单工作节点加速场景下的性能，结果表明模型质量主要受总计算量影响，对通信失败具有鲁棒性，应用于单工作节点时可提升收敛速度和泛化性能。

  4. 相关工作：与局部随机梯度下降（SGD）和联邦学习相关工作相比，DiLoCo在更大规模的模型和更多副本上进行实验，且通信频率更低；与线性模式连接性相关工作相比，DiLoCo在训练过程中利用线性模式连接性，而多数相关工作在模型微调后进行平均 。

  5. 局限性：仅研究了语言建模任务和Transformer架构，在更大规模模型上的效果需进一步验证，未考虑工作节点的异构性，算法在超过8个工作节点时收益递减，计算效率和数据效率有待提高。

  6. 结论：DiLoCo是一种在设备连接不佳时，分布式训练Transformer语言模型的有效方法，具有鲁棒性，但在更大规模模型、其他领域和架构上的适用性仍需研究。