Prime Intellect 开源 OpenDiLoCo，跨三大洲训练 1.1B 模型

OpenDiLoCo：面向全球分布式低通信训练的开源框架

介绍 OpenDiLoCo，这是 DeepMind 分布式低通信（DiLoCo）方法的开源实现与扩展，支持全球分布式 AI 模型训练。

上周，我们发布了宏伟计划的第一步，推出了 Prime Intellect Compute Exchange，用于聚合和编排全球计算资源。

今天，我们非常激动地宣布在第二部分上迈出了重要一步：开源我们的分布式训练框架，以支持跨全球分布式硬件的协作式模型开发。

我们提供了 Deepmind DiLoCo 实验的可复现实现，并将其置于一个可扩展的去中心化训练框架中。我们通过在两大洲和三个国家训练一个模型来证明其有效性，同时保持 90-95% 的计算利用率。此外，我们将 DiLoCo 扩展到原始工作 3 倍的规模，证明了其对十亿参数模型的有效性。

论文：https://arxiv.org/abs/2407.07852

代码：https://github.com/PrimeIntellect-ai/OpenDiLoCo

OpenDiLoCo 训练运行可视化

大型语言模型已经彻底改变了 AI，但传统上训练它们需要大规模、集中式的计算集群。这种资源集中限制了谁能参与 AI 开发，并减缓了创新步伐。

最近，我们发布了一篇详细的博客文章，探讨了分布式 AI 训练的最新技术水平。在文章中，我们重点介绍了最有前景的方法以及仍需克服的几个关键挑战：

互连带宽缓慢
确保容错训练
非均质硬件环境
以及更多……

OpenDiLoCo 是我们为克服第一个挑战而开展的研究工作之一，它通过促进跨全球多个连接不良设备的高效训练来实现这一目标。

主要贡献

复现与扩展：我们成功复现了原始 DiLoCo 实验，并将其扩展到十亿参数模型规模。
开源实现：我们发布了基于 Hivemind 库构建的可扩展实现，使广大开发者和研究人员都能使用去中心化训练。通过与 PyTorch FSDP 集成，我们的框架使单个 DiLoCo 工作节点能够扩展到数百台机器。
全球去中心化训练：我们通过在两大洲和三个国家训练一个模型，展示了 OpenDiLoCo 的实际潜力，实现了 90-95% 的计算利用率。
效率洞察：我们的消融研究为该算法的可扩展性和计算效率提供了宝贵见解，为未来改进铺平了道路。

DiLoCo

Google DeepMind 最近的工作提出了一种方法，可以在连接不良的设备孤岛上训练语言模型。该方法允许在这些不同孤岛上进行数据并行训练，仅需每 500 步同步一次伪梯度。

DiLoCo 引入了一种内外层优化算法，允许进行局部和全局更新。每个工作节点使用本地 AdamW 优化器独立地多次更新其权重（内层优化）。每约 500 次更新后，算法使用 Nesterov 动量优化器执行一次外层优化，同步所有工作节点的伪梯度（所有局部梯度之和）。

这种方法显著降低了通信频率（最高可达 500 倍），从而降低了对分布式训练的带宽需求。

OpenDiLoCo

为了促进这一有望实现 AI 民主化的研究方向上的协作，我们以开源许可证发布了 OpenDiLoCo 的代码：https://github.com/PrimeIntellect-ai/OpenDiLoCo。

我们的实现构建在 Hivemind 库之上。Hivemind 不使用 torch.distributed 进行工作节点通信，而是利用分布在每个工作节点上的分布式哈希表（DHT）来通信元数据并同步它们。该 DHT 使用开源的 libp2p 项目实现。我们利用 Hivemind 进行 DiLoCo 工作节点之间的节点间通信，并使用 PyTorch FSDP 进行 DiLoCo 工作节点内部的节点内通信。

我们与 Hivemind 的集成使 DiLoCo 能够实现真实世界中的去中心化训练设置，使其许多固有特性得以使用，例如：

资源的动态增减：可用计算量可以在训练过程中变化，新的设备和集群可以在训练中途加入和离开。
容错性：对于去中心化训练，某些设备可能不如其他设备可靠。通过 Hivemind 的容错训练，设备可以随时变得不可用而不会停止训练过程。
点对点：没有主节点。所有通信都以点对点的方式进行。

主要结果

作为第一步，我们复现了 DiLoCO 的主要实验结果。我们使用 C4 数据集在语言建模任务上训练了一个具有 1.5 亿参数的模型。

我们表明，具有 8 个副本的 DiLoCo 显著优于没有任何副本的基线，并且在相同计算预算下达到了更强基线的性能，尽管其通信需求降低了 500 倍。

DiLoCo 8-worker 150M experiment — 使用 8 个 DiLoCo 工作节点预训练 1.5 亿参数的 Llama 模型，其困惑度显著低于没有 DiLoCo 的基线，甚至与在相同计算预算下使用 8 倍更大批量大小的基线相比也是如此，同时通信量减少了 500 倍。

‍

除了复现结果之外，我们还对 DiLoCo 进行了多项消融研究，重点关注该算法在工作节点数量方面的可扩展性和计算效率。我们还证明了 DiLoCo 伪梯度可以有效地以 FP16 进行 all-reduce 而不会造成任何性能下降。

更多详情，请查看 arXiv 上的论文。

将 DiLoCo 扩展到十亿参数模型

DeepMind 最初的 DiLoCo 论文仅对最多 4 亿参数的模型规模进行了实验。在我们的工作中，我们将该方法扩展到了具有 11 亿参数的模型。我们采用与 TinyLlama 相同的超参数，并使用包含 800 万个 token 的总批量大小（批量大小为 8192，序列长度为 1024）。由于批量大小是之前实验的 4 倍，我们决定在此实验中仅训练 44k 步。

我们将结果与两个基线进行比较：一个没有 DiLoCo 且没有副本的弱基线，以及一个使用 4 倍批量大小和数据并行的更强基线。

当使用 500 的本地步长（每 500 步在工作节点之间同步）时，与 1.5 亿参数模型的实验类似，我们观察到训练早期阶段收敛效果不佳。训练动态在后期阶段有所改善，如果我们进一步训练到 88k 步，可能达到基线水平。

我们还进行了一个本地步长为 125 的实验。在这种模式下，训练动态在早期阶段实际上更好。使用 125 个本地步长的 DiLoCo 运行几乎达到了相同计算预算下更强基线的性能，同时通信量减少了 125 倍。

DiLoCo 1.1B scaling experiment — **1.1B 规模实验：** 比较使用 OpenDiLoCo 的 1.1B 训练，其中 4 个工作节点每 500 个本地步长和每 125 个本地步长同步一次，与两个基线进行对比。

虽然我们证明了 DiLoCo 在十亿参数规模下有效，但我们认为需要进一步的工作来使其在更大的批量大小和增加的本地步长下有效。

全球分布式训练设置

为了展示 OpenDiLoCo 跨不同大洲的去中心化训练功能，我们使用了四个 DiLoCo 工作节点，每个节点配备八个 H100 GPU，分别位于加拿大、芬兰和美国两个不同的州。图中显示了工作节点之间的网络带宽，范围在 127 到 935 Mbit/s 之间。我们使用 500 个本地步长训练了 11 亿参数的模型，梯度以 FP16 进行全归约。由于本地步长较多，四个工作节点独立运行约 67.5 分钟后才进行梯度平均的通信。对于外部优化器步骤，我们的实验显示工作节点之间的平均全归约时间为 300 秒。

Decentralized training connection speeds — 工作节点之间的网络带宽：在我们的去中心化训练设置中，使用 iperf 测量的四个工作节点之间的平均双向网络带宽（单位：Mbit/s）。这些节点位于三个不同的国家，由不同的云提供商托管：加拿大（Hyperstack）、芬兰（DataCrunch）、美国德克萨斯州（Voltage Park）、美国特拉华州（Runpod）。

由于 DiLoCo 显著减少了通信时间，全归约瓶颈仅占训练时间的 6.9%，对整体训练速度影响极小。

在我们的场景中，最快的工作节点会闲置额外的时间。在未来的工作中，我们将通过探索异步设置中的 DiLoCo 来解决这个问题。

运行 OpenDiLoCo

运行代码很简单。唯一的要求是至少有两块 GPU，它们不必位于同一地点。设置环境后，使用以下命令创建初始 DHT 节点：

python ../hivemind_source/hivemind/hivemind_cli/run_dht.py
--identity_path fixed_private_key.pem
--host_maddrs /ip4/0.0.0.0/tcp/30001

在另一个终端中，您可以使用以下命令启动 DiLoCo 工作节点，确保正确设置 PEER、NUM_DILOCO_WORKERS 和 WORLD_RANK：

export PEER=/ip4/192.168.100.20/tcp/30001/p2p/Qmbh7opLJxFCtY22XqwETuo6bnWqijs76YXz7D69MBWEuZ
# change the IP above to your public IP if using across nodes connected via internet

export NUM_DILOCO_WORKERS=4
export WORLD_RANK=0

torchrun --nproc_per_node=8 \
     train_fsdp.py \
     --per-device-train-batch-size 16 \
     --total-batch-size 2048 \
     --total-steps 88_000 \
     --project OpenDiLoCo \
     --lr 4e-4 \
     --model-name-or-path PrimeIntellect/llama-1b-fresh \
     --warmup-steps 1000 \
     --hv.averaging_timeout 1800 \
     --hv.skip_load_from_peers \
     --hv.local_steps 500 \
     --hv.initial-peers $PEER \
     --hv.galaxy-size $NUM_DILOCO_WORKERS \
     --hv.world-rank $WORLD_RANK \
     --checkpoint_interval 500 \
     --checkpoint-path 1b_diloco_ckpt

你可以在 GitHub 仓库的 README 中找到更多关于运行 OpenDiLoCo 的信息。

在 PI 计算平台上运行 OpenDiLoCo

设置全局编排层来运行 DiLoCo 训练仍然相当具有挑战性。我们的 PI 计算平台通过预构建的 OpenDiLoCo docker 镜像让这一切变得容易得多。该镜像预装了所有依赖项，让你可以轻松启动 DiLoCo 工作节点。

在未来的工作中，我们很期待在计算平台中构建一个集成的开源技术栈，为跨多个集群的编排、效率优化、节点故障处理、基础设施监控等提供流畅的解决方案。

结论与未来方向

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

在未来的工作中，我们计划将 DiLoCo 扩展到更大规模的模型和更多的分布式工作节点上。一些有趣的方向包括可能提升稳定性和收敛速度的模型合并技术。此外，可以通过实现异步执行权重平均通信的方法来减少计算空闲时间，将其与下一个外部优化步骤的计算交错进行。

我们对这项技术立即可用的实际应用感到兴奋，并期待很快在我们的总体规划第三部分中继续推进：在语言、智能体、代码和科学等高影响力领域协作训练并贡献开放 AI 模型，实现 AI 模型的集体所有权。

加入我们，共建 AI 的开放未来

塑造 AI 未来的力量不应集中在少数人手中，而应向任何有能力做出贡献的人开放。我们邀请你加入我们，共同为 AI 构建一个更加分布式、更具影响力的未来：

如果你有不懈的雄心并希望实现这一目标，请申请我们的开放职位。
在我们的 AI 模型计划和开源框架上展开合作。
贡献算力，并获得最先进 AI 模型的所有权。

我们的 OpenDiLoCo 工作也已被 ICML 的 ES-FoMo 研讨会接收。如果你要来维也纳，请联系我们！