[博客翻译]INTELLECT–1：启动10B参数模型的分布式训练

由 发布于 去中心化训练框架开源人工智能
 收藏

原文地址：https://www.primeintellect.ai/blog/intellect-1


我们很高兴推出 INTELLECT-1，这是第一次分散训练运行 100 亿参数模型，邀请任何人贡献算力并参与。这使我们更接近开源 AGI 的目标。

最近，我们发布了 OpenDiLoCo，这是一个开源实现和扩展 DeepMind 的分布式低通信（DiLoCo）方法的项目，使全球范围内的人工智能模型分散训练得以实现。我们不仅重现并开源了这项工作，并且成功地将其规模扩大到 10 亿参数大小。

现在我们 进一步将规模提升 10 倍至 100 亿参数规模，比原始研究提升了约 25 倍。 这使我们进入了我们的 总体规划 中的第三步：协同训练前沿的开放基础模型：从语言模型、智能体到科学模型。

我们的目标是逐步解决分散训练问题，确保 AGI 将是开源的、透明的和可访问的，防止少数中心化实体控制 AGI 并加速人类进步。

启动合作伙伴和支持者

我们非常激动并且感谢能够有像 Hugging Face、SemiAnalysis、Arcee、Hyperbolic、Olas、Akash、Schelling AI 等领先的开源人工智能机构加入我们并为这次分散训练提供算力支持。

如何贡献算力

现在任何人都可以通过我们的平台贡献资源来推进开源人工智能的发展，未来还将更容易地使用他们自己的硬件。

分散训练的范式转变

正如 Anthropic 共同创始人 Jack Clark 所强调的，还没有任何模型能够在全球分布式工作者之间以 100 亿参数的规模高效训练。我们最初的 OpenDiLoCo 训练突破了 10 亿参数的大关，并且通过 INTELLECT-1，我们正在达到新的分散训练规模水平。

有没有过分布式训练规模超过 100 亿参数的情况？我读过的所有关于分布式训练的论文（例如 OpenDiLoCo）似乎都只涉及大约 10 亿参数的语言模型。

— Jack Clark (@jackclarkSF) 2024 年 8 月 3 日

DiLoCo 能够在互连性差的设备上训练人工智能模型。这种方法允许在不同的“岛屿”中进行数据并行训练，仅需每几百个步骤同步一次伪梯度。这显著降低了通信频率（最多降低 500 倍），从而减少了分布式训练所需的带宽要求。

Prime：我们的分散训练框架

自从我们最初开源以来，我们在分布式训练框架方面进行了两方面的改进：

1. 算法改进

很多基于我们 OpenDiLoCo 工作的改进表明，在进一步减少通信需求方面具有很大的前景。特别是我们在伪梯度量化实验中，已经将带宽需求降低了最高 2000 倍，结合了伪梯度的 int8 量化以及每 500 步与外层优化器同步。这些结果在较小规模下已经很有效，我们对将其放大到更大模型规模感到兴奋。

2. 可扩展的分散训练框架

分布式训练是一个工程和研究上的挑战。实现跨分布式数据中心的容错 训练 是今天即使是最大的人工智能实验室也正在努力解决的问题。

我们很高兴宣布发布一个新的分散训练框架称为 Prime。Prime 支持容错训练、允许动态增加/减少计算资源，并优化一个全球分布的 GPU 网络中的通信和路由。

该框架构成了我们开源技术栈的基础，旨在支持我们自身和其他超越 OpenDiLoCo 的分散训练算法。通过构建这一基础设施，我们旨在推动全球分布式 AI 训练的可能性边界。

主要功能：

  • ElasticDeviceMesh 用于容错训练：
    • 在 Prime 中，我们添加了一种新的分布式抽象叫做 ElasticDeviceMesh，它封装了用于跨互联网进行容错通信的动态全局进程组以及用于节点或数据中心内通信的本地进程组。
    • ElasticDeviceMesh 管理全局进程组在节点加入或离开时的重设大小。不同于标准的 torch 分布式中的 DeviceMesh，在这种情况下，进程组发生故障会崩溃并需要冷重启以重设进程组大小。
    • 为了判断何时重设进程组大小，我们使用心跳机制来发现死节点并将其从进程组中移除。发生故障的节点将尝试尽力发出死亡信号以快速失败其心跳检测，从而节省其他节点等待超时的时间。
  • 异步分布式检查点
    • 由于模型体积庞大，检查点操作可能是一个耗时的操作，在我们的测试节点上花费时间长达 20 分钟。如果这个过程阻塞了主要训练流程，则会影响我们的计算资源利用率。
    • 为最大限度地减少阻塞时间，我们首先将检查点保存至 /dev/shm 中，这是一个基于 RAM 的文件系统。此操作更快，一旦检查点创建完毕，即可解除对主要训练流程的阻塞。
    • 然后使用两个子进程异步地将检查点从 /dev/shm 中复制到磁盘上的检查点目录，并上传至远程位置。
  • 实时检查点恢复
    • 想要在训练中途加入的节点需要能够获取最新的模型和优化器状态，这样才能加入训练。它们必须在此操作在一个外层步骤之间的窗口时间内完成，否则，它们将接收到过期的检查点。
    • 为了快速完成这一过程，我们将加入节点的请求指向那些正在托管一个 sidecar HTTP 服务器的对等节点，这台服务器提供来自 /dev/shm 最新检查点的服务。
    • 一旦加入的节点下载并初始化好模型后，就会跳过内部步骤，并带着零伪梯度加入到外层步骤。这是为了防止加入节点阻塞现有运行中的节点。如果加入节点也执行了内部步骤，那么它将因为下载和加载检查点所需时间而延迟进入外层步骤，从而降低集群的整体计算利用率。
  • 自定义 Int8 全局规约核
    • 在我们的实验中，我们发现可以对伪梯度执行 int8 量化而不影响损失曲线。这意味着如果我们用 int8 而不是 fp32 进行伪梯度通信，可以将每个外层步骤的全局规约负载大小减少 4 倍。
    • 但是我们需要用 fp32 积累规约结果，在全局规约期间对其去量化和重新量化中间结果。这不被任何集体通信库所支持。
    • 因此我们使用 C++ 实现了一个完全管线化的环规约核，并通过torch库 JIT 编译为自定义算子。
    • 但是，鉴于所需执行的量化工作量，使用 torch 操作（quantize_per_tensor、scatter_add、index 等）过于缓慢，导致目标网络带宽利用率不足 4 Gbps。
    • 于是我们使用 C++ 实现了自己的多线程 uint8 操作，提高了量化速度 60 多倍。
  • 最大化带宽利用：
    • 通过对 DiLoCo 伪梯度在节点内部进行分片，可以通过同时打开多个连接来最大化网络带宽利用率，从而在执行全局规约时显著提高传输速度。这一改进使部分节点的传输速度提升了 8 倍。
    • 依赖公共 IP 转发在某些计算提供商中会导致差或者不稳定的 p2p 带宽。为了解决这一问题，我们采用了 VPN 技术以优化节点间的对等连接，通过修改数据包路径以更好地利用节点间可用的互联网带宽。
    • 在类似数据中心环境下的节点间带宽利用率相比我们的 OpenDiLoCo 发布版本提升了最多 40 倍，实现了跨越整个美国的数据中心间的高达 4Gb/s 的连接。
  • PyTorch FSDP2 / DTensor ZeRO-3 实现
    • 为了适应给定内存资源下的 10G 模型训练，我们需要在节点内的 GPU 之间分散模型权重、梯度和优化器状态。
    • 我们通过 PyTorch FSDP2 的 fully_shard API 实现了这一点，此 API 将模型参数包装为 DTensors 并注册钩子来调度所有聚集与规约拆散操作。FSDP2 还通过将参数分桶为 FSDPParamGroups 来优化集体操作。这使得我们可以对更大的张量执行集体操作，改善协议到载荷的比值并提高流水线重叠效果。对于我们的伪梯度，我们也采取相同的方法按层分块。
  • CPU 卸载
    • 我们的 DiLoCo 优化器不增加任何 GPU 开销。所有由 DiLoCo 优化器所需的张量均被移至 CPU 内存。
    • 既然我们每几百个步骤才进行一次全局同步，那么在 CPU 上复制和计算伪梯度的速度降低相对于执行内部步骤和全局规约的时间是微乎其微的。

更多信息，请查看 仓库，并关注我们即将发布的研究论文。

这一切都使得 INTELLECT-1 的训练规模达到了 10B 参数级别，并能实现跨多个分布式工作者 98% 的计算资源利用率。 在我们的生产训练运行中，我们选择每 100 个步骤同步一次，在由 8 台 H100 节点组成的集群上大约需要 40 分钟才能完成。我们将伪梯度量化为 int8，减少了 400 倍的通信需求。利用我们的新框架进行 DiLoCo 外部优化器的全局规约同步耗时不到 1 分钟，将节点间通信时间缩短到了总训练时间的 1%-2%。

INTELLECT-1：首个去中心化训练的 10B 模型

INTELLECT-1 是一个基于 Llama-3 架构的 10B 参数的模型。

这将是首个在最高质量的开源数据集之一 Fineweb-Edu 上广泛训练的大规模模型，该数据集由 Hugging Face 提供。

我们精心挑选的 数据混合 包括如下内容：

由于其中一些数据集中包含相关的数据分段，我们通过从 12 个流数据迭代器中随机抽取并对数据集重新分片进行了预洗牌处理。我们的预洗牌数据集可以在 HuggingFace Hub 找到。通过使用 Llama-3 分词器，我们数据混合中经过处理的 token 总数量超过 6 万亿。

我们使用 WSD 学习率调度器，在初始预热阶段后维持恒定的学习率。这种方法提供了我们根据计算贡献数量训练所需 token 数量的灵活性。在训练末期，我们计划使用高质量数据集实施冷却阶段以进一步提升性能，并执行训练后的优化。

下一步计划：

INTELLECT-1 仅仅是我们迈向第一步。我们将在我们 路线图 上继续前行，以扩展去中心化训练至最大且最强大的科学、推理和编码领域的开放前沿模型。

我们的路线图包括：

  1. 扩展至更强大且规模更大的科学、推理及编码领域的开放前沿模型。
  2. 开发一种允许任何人都能贡献计算资源的系统，并使用证明机制确保其在去中心化训练中的安全性和可验证性。
  3. 创建一个让任何人都可以发起去中心化训练的框架，并开放他人的贡献。