OpenDiLoCo：分布式AI训练的开源解决方案，低通信成本，全球覆盖！

在这个AI大爆炸的时代，大型语言模型（LLMs）已经成为推动机器学习应用的超级引擎。但是，训练这些庞然大物需要巨大的计算资源。想象一下，如果我们能在世界各地分散的设备上，高效地训练这些模型，那将会怎样?这就是OpenDiLoCo带给我们的惊喜。

传统的分布式训练方法，需要频繁的通信和大量的带宽，这限制了训练的规模和效率。而DiLoCo（分布式低通信）训练方法，通过减少通信需求，让LLMs的全球训练成为可能。

OpenDiLoCo是一个开源框架，它实现了DiLoCo训练方法，并用Hivemind库提供了一个可扩展的、去中心化的训练框架。这个框架的厉害之处在于，它在全球范围内，跨越两大洲、三个国家进行了模型训练，同时保持了90-95%的计算利用率。

主要特性:

资源的动态增减:训练过程中，可动态调整计算资源，新设备和集群可以在训练中途加入或退出。

容错性:去中心化训练中，某些设备可能不可靠。通过Hivemind的容错训练，即使设备随时不可用，训练过程也不会停止。

点对点通信:没有主节点，所有通信均以点对点方式进行。

研究人员们不仅复现了DiLoCo的实验，还将其扩展到了十亿参数模型的规模。他们通过消融研究，展示了DiLoCo算法在计算效率和扩展性方面的优势。更令人印象深刻的是，他们证明了DiLoCo的梯度可以在不降低性能的情况下，使用FP16进行全归约。

主要贡献:

复制与扩展:成功复现了DiLoCo的原始实验，并将其扩展到十亿参数模型的规模。

开源实现:基于Hivemind库构建的可扩展实现，使去中心化训练对广泛的开发者和研究者开放。

全球去中心化训练:通过在两大洲和三个国家之间的模型训练，展示了OpenDiLoCo的实际潜力，同时保持了90-95%的计算利用率。

效率洞察:通过消融研究提供了算法可扩展性和计算效率的宝贵见解。

实验结果:

通过复制DiLoCo的主要实验结果，Prime Intellect公司证明了其方法的有效性。在C4数据集上使用150M参数的模型进行语言建模任务的训练，展示了DiLoCo在通信需求降低500倍的情况下，与基线性能相当。

DeepMind的原始DiLoCo论文仅针对最多4亿参数的模型进行了实验。在这项工作中，Prime Intellect公司将该方法扩展到了1.1亿参数的模型，并采用了与TinyLlama相同的超参数。

为了展示OpenDiLoCo在不同大陆的去中心化训练功能，实验使用了位于加拿大、芬兰和美国两个不同州的四个DiLoCo工作节点，每个节点都配备了八个H100GPU。

Prime Intellect公司成功复现了DiLoCo的主要实验结果，并将该方法扩展到原始工作参数大小的三倍，并在现实世界的去中心化训练环境中展示了其应用。

未来，该公司计划将DiLoCo扩展到更多分布式工作节点上的更大模型，并探索可能提高稳定性和收敛速度的模型合并技术，以及通过实施异步权重平均通信方法来减少计算空闲时间。

论文地址：https://arxiv.org/pdf/2407.07852