谁告诉过你分布式训练是不可能的？民主化AI：Psyche网络架构

# 推动AI民主化：Psyche网络架构 - NOUS RESEARCH 来源：https://nousresearch.com/nous-psyche ## 引言 先进AI模型的开发已变得**高度集中**，需要大量计算资源，只有大型企业才能负担得起。这种权力集中威胁着创新，并限制了能为AI进步做出贡献的人群，几乎确保了所构建的系统遵循封闭实体对“对齐”的愿景。与此同时，全球范围内大量计算能力处于闲置或未充分利用状态。 Psyche通过创建**去中心化基础设施**改变了AI开发方式，允许任何人参与训练大型语言模型。Psyche无需在单一地点部署数千个加速器的大型基础设施，而是协调分布式、异构硬件在全球范围内进行训练。 日益增长的未充分利用硬件，为无需高昂投资成本即可训练和使用LLM提供了机会。 开发路线图包含两个主要阶段→ 1. **合作训练**：从许可测试网开始，逐步过渡到完全去中心化环境 2. **可访问的推理与高级能力**：包括强化学习和创建推理模型的能力 ## 回顾DeMo与DisTrO 我们首次在DeMo：解耦动量优化（https://arxiv.org/pdf/2411.19870）中介绍了高效去中心化训练的工作。DisTrO在这些核心概念基础上进一步提升了性能。您可以在以下博客文章（https://blog.lambdaclass.com/introducing-demo-decoupled-momentum-optimization-for-efficient-distributed-llm-training/）中阅读更完整的解释。 在LLM训练过程中，优化器决定如何计算梯度以更新模型参数（AdamW（https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html）是著名的优化器）。在大型设置中，数千个加速器计算并共享梯度（外加一些额外信息）到整个网络，以便每个节点都能了解其他节点在该训练步骤中学习到的内容。 **DisTrO是一系列优化器，利用机器学习训练中的某些意外特性，大幅压缩加速器之间传递的信息。** 这一思路模仿了JPEG压缩的底层原理。压缩图像时，我们将一个矩阵（MxN像素网格）转换到频域（https://www.mathworks.com/help/signal/ug/practical-introduction-to-frequency-domain-analysis.html）。一旦转换到频域，许多自然信号（例如声音和图像）的绝大部分信息都集中在低频部分。这是因为自然图像在像素间变化平缓——图像或录音中通常不存在急剧的硬边缘。大系数（最强波）通常出现在低频（缓慢、平滑的变化），而只有小系数出现在高频（尖锐、快速的变化）。因此，我们可以丢弃高频部分而不会损失太多信息，因为低频部分反映了原始矩阵的全局变化。 DCT压缩示意图可以看出，即使图像质量几乎不变，高频部分（图谱边缘部分）已经消失。**DeMo将类似原理应用于AI训练。**大多数优化器追踪一个“动量张量”——本质上是对过去梯度的运行平均值，随时间平滑更新。 **DeMo**从带动量的随机梯度下降（https://github.com/PsycheFoundation/psyche/blob/08cf9258e1d592ea409d42aea219d825a5f88c93/shared/modeling/src/distro.rs#L154-L181）开始，解耦各节点的动量，让每个节点积累独特的动量。我们使用离散余弦变换（DCT）（https://github.com/PsycheFoundation/psyche/blob/08cf9258e1d592ea409d42aea219d825a5f88c93/shared/modeling/src/distro.rs#L154-L181）来近似最优的频域转换，因为DCT计算成本低且精度足够（JPEG压缩也使用同样的方法）。在频域中，我们提取能量最高的前k个动量分量。 然后将这些动量分量以数值及其对应索引列表的形式传递给其他节点。 为了避免训练随时间产生系统性偏差，我们需要改变信息压缩方式，与典型的JPEG压缩方案不同。在典型JPEG压缩方案中，高频分量通常比低频分量压缩得更厉害（因为预期低频分量包含更多信息）。当表示为矩阵时，这些元素总是出现在相同的坐标位置。而DeMo则从**幅度最大**的分量中提取信息，无论它们在矩阵中的位置。这使得方案能够从频谱的任何位置提取更新，而不是持续偏向固定坐标的信息。 当准备同步各节点时，我们使用逆离散余弦变换（https://github.com/PsycheFoundation/psyche/blob/08cf9258e1d592ea409d42aea219d825a5f88c93/shared/modeling/src/distro.rs#L190-L239）恢复近似的梯度分量，并用这些近似分量更新模型参数。 为了提高效率，DCT基矩阵可以分块预生成，从而减少压缩和解压的墙钟时间开销。实际上，这种开销不到总训练时间的1%。 **DisTrO的核心由以下权重更新方程表示：** \\\[w_{t+1} = w_t - \\eta Q_t\\\] 关于算法的完整描述，请参见算法1（https://arxiv.org/pdf/2411.19870）以及论文中的解释。 ## Psyche对DeMo的改进 **在Psyche代码库中，我们对原始DeMo算法实施了几项改进。** 首先，我们加入了重叠训练，使得节点无需等待前一步的所有更新下载完毕即可开始下一步。在共享前一步的DisTrO结果时，它们已经在训练下一步了。由于DisTrO结果随模型参数规模呈次线性增长，当我们将DisTrO应用于更大模型时，训练下一步所花费的墙钟时间将超过通信前一步DisTrO结果所花费的时间。这意味着理论上，随着规模扩大，通信延迟不会成为瓶颈。这显著提高了GPU利用率，理论上可使DisTrO设置的效率与集中式设置相当。 通过重叠DisTrO进行3步训练 应用更新生成更新见证活动第0步跳过生成0见证（虚拟）第1步跳过生成1见证0第2步应用0生成2见证1第3步应用1生成3见证2第4步应用2跳过见证3第5步应用3跳过见证（虚拟） 其次，我们通过对动量张量的离散余弦变换进行“量化”，**进一步降低了之前DisTrO实现的带宽**。事实证明，在DisTrO结果中仅传达DCT的**符号**就足以重建正确的动量矩阵，从而获得正确的梯度更新。换句话说，DCT的幅度几乎没有带来额外有用信息。通过只发送符号（正或负，1或-1）而不带幅度（以及这些符号在矩阵中对应位置的索引），我们将结果压缩了3倍以上。当 `quantize_1bit` 标志开启时，此压缩功能即生效。 Psyche在基于Rust的系统中实现了DisTrO，具备P2P网络能力，能够协调在任意时刻进行多次运行。 Psyche网络概览 图1：Psyche网络系统架构**Psyche架构包含三个主要角色：** 1 ### 协调器 协调器作为训练运行状态的权威实体，存在于链上智能合约中。协调器负责： - 存储运行的元数据和参与者列表 - 处理训练运行各状态之间的转换 - 提供随机性以确定分配和见证者 - 为所有参与者提供同步点 2 ### 客户端 客户端是负责部分训练的GPU节点。此外，它还可以在运行过程中充当见证者或上传检查点。客户端在运行期间执行三种操作： - **训练：**计算梯度，与其他客户端共享梯度，并更新模型参数 - **见证：**验证其他客户端的活跃性和正确性 - **验证：**重新计算结果并与其他参与者比较，以检测恶意行为者 3 ### 数据提供者 数据提供者（https://github.com/PsycheFoundation/psyche/tree/main/shared/data-provider）负责为训练提供必要的数据。它可以是本地的（每个客户端负责自己的训练数据副本），也可以是HTTP或TCP提供者。 Psyche网络架构图 图2：Psyche网络架构，展示不同组件之间的交互 ## 训练运行生命周期 在Psyche中，一个训练运行由多个“epoch”组成。此处，epoch是一组训练步骤（例如500个训练步骤），我们将整个训练过程分解为这样的片段（这个定义与传统机器学习术语有所不同）。在加入过程中，运行会短暂暂停，以便新进入者下载模型的检查点。 **我们设计了带有epoch的Psyche，以减少贡献算力的机会成本**，让客户端能够安全地加入和退出。算力提供者可以在相对短的时间内停止运行Psyche，转而投向更有利可图的机会，这有助于通过鼓励原本闲置的GPU参与来降低训练成本。 在epoch开始时（包括新运行的开始时），新客户端有两个模型下载选项： 1. Huggingface标志 **HuggingFace检查点**：一组客户端负责在每个epoch后将信息上传到HuggingFace，并将该检查点的URL发送给协调器。 2. **P2P检查点**：客户端可以从可用的对等节点获取最新参数，向不同客户端请求每层的参数集。 训练阶段图 图3：展示客户端通过P2P获取模型权重的示意图**一旦加入，一个epoch的行为如下：** 1. **等待成员阶段**：协调器等待达到最低连接客户端阈值。 2. **预热阶段**：客户端下载并加载模型。如果客户端断开连接且数量低于最小值，则过程回退至等待。 3. **训练阶段**：协调器提供随机种子以确定使用哪些训练数据。客户端获取数据，计算自己的训练部分，并与其他客户端共享结果。 4. **见证阶段**：指定的见证者验证客户端活动，生成“见证证明”（通过布隆过滤器压缩）并转发给协调器。 5. **继续或冷却**：协议要么继续下一个训练步骤，要么在以下情况下进入冷却：- 参与率低于最低阈值 - 当前步骤完成一个epoch - 未达到见证法定人数 6. **冷却阶段**：某些客户端对模型进行检查点保存，然后返回等待阶段。 Psyche训练工作流图 图4：详细工作流，展示客户端在训练过程中的交互方式 ## P2P与网络 为了让客户端的网络连接更简单，**我们利用UDP打洞技术**，使得NAT或防火墙后的节点能够直接相互连接，无需手动协调端口。打洞技术通过协调同时连接尝试，在防火墙中创建临时开口，允许直接通信，无需手动端口转发或中继服务器，从而让防火墙后的对等节点之间建立直接连接。 P2P网络图 图5：P2P网络图，展示UDP打洞和节点通信Psyche使用Iroh（https://www.iroh.computer/）进行P2P网络通信。使用Iroh的应用程序通过`NodeId`（一个32字节的Ed25519（https://ed25519.cr.yp.to/index.html）公钥）连接对等节点，而非其IP地址。这种设计将路由与物理网络地址解耦：`NodeId`全局唯一、自我生成，并在网络变化时保持稳定。由于它们是加密密钥，所有连接默认都是端到端加密和认证的（每个对等节点在握手期间验证对方的`NodeId`）。 默认情况下，Iroh会尽可能通过UDP建立直接点对点连接（使用QUIC（https://en.wikipedia.org/wiki/QUIC）），必要时回退到中继。这确保了大约90%的情况下能建立直接连接（其余由中继处理），高于libp2p（70%）或BitTorrent的通用UDP打洞（60-70%）。 此外，Iroh可以无缝地在实时中继模式和直接模式之间切换。如果网络条件发生变化（例如节点的某个网络接口断开，但另一个仍然工作），Iroh将动态迁移QUIC连接——利用QUIC的连接迁移能力——或根据需要回退到中继，而不会中断会话。这使得连接高度可靠，能够透明地处理IP更改和网络转换。 参与训练运行的Psyche节点之间使用iroh-gossip（https://github.com/n0-computer/iroh-gossip）通信关于训练结果的元数据，该协议基于HyParView（https://asc.di.fct.unl.pt/~jleitao/pdf/dsn07-leitao.pdf）和PlumTree（https://asc.di.fct.unl.pt/~jleitao/pdf/srds07-leitao.pdf），并包含一个标识包含结果的二进制blob哈希的`iroh-blob`票据。节点然后可以通过Iroh blob协议下载结果并继续训练。 ## 容错与鲁棒性 我们预计在任何运行期间，节点可能不会在epoch结束时“优雅地”退出。在训练过程中，如果某个客户端没有收到来自其他客户端的数据，它就会向协调器发送健康检查。如果协调器收到足够多的来自客户端的负面健康检查消息，它会移除崩溃的节点。无响应的节点将不再从运行协调器接收分配给它的数据，其数据将被重新训练。 我们还通过暂停和重启那些很可能无法在合理时间内完成的epoch来避免过多节点离线的情况。 ## 布隆过滤器实现高效验证 Psyche使用布隆过滤器（https://en.wikipedia.org/wiki/Bloom_filter）来高效检查客户端的DisTrO结果是否已共享。布隆过滤器是用于检查集合中成员关系的廉价数据结构。它们偶尔会产生误报，即指示元素在集合中但实际不在，但我们为了效率而接受这种可能性（与更确定性的数据结构如默克尔树相比）。 布隆过滤器基本上创建一个初始化为0的位数组。然后，当我们看到一个要添加到集合中的新元素时，我们在该元素上运行多个不同的哈希函数，并将结果哈希函数对应的索引位置设置为1。因此，在检查某个元素是否在集合中时，我们只需要运行这些哈希函数。如果其中任何一个函数在该索引处返回0，我们就可以停止检查并确定该元素不在集合中。