Learning@home：大模型分布式训练范式

文章链接：Towards Crowdsourced Training of Large Neural Networks using Decentralized MoE

1.背景介绍

在过去的十年中，深度学习社区通过在大量数据上训练大型模型取得了显着的进步，但代价是网络模型的不断加深膨胀，模型的体积以及训练的数据量与目标效益呈指数倍增长:

• 在CV领域，从ImageNet挑战赛出现130万个标记图像以来，卷积神经网络的典型大小从几MB增加到数百MB;
• 在NLP领域，目前一大趋势是训练大型的类Transformer语言模型。这项任务的数据几乎是无限的，允许研究人员使用数十甚至数百GB的参数训练模型。

试想一下，如果将当下最火热的NLP模型Chat-GPT4.0(1w亿 parameters)放大100倍，我们将得到一个100w亿的庞然大物;根据使用的是单精度还是半精度，仅存储模型就需要170-340 TB的内存，没有现代大规模生产的GPU能够完成这样的任务。即使是配备16倍V100加速器的高端服务器，也只能在组合GPU内存中存储该模型的0.15%，更不用说训练它了。

目前训练大模型的一个范式思路是在专门的紧密互连的集群中部署更强大的GPU（堆算力）。比如T-NLG和Megatron-LM等模型已经在DGX-SuperPOD上进行了训练，其中DGX-SuperPOD是一台拥有数百个Tesla V100 GPU的超级计算机，分布在数十台服务器上。这种模式的主要挑战是费用昂贵，导致只有那些拥有大量预算的实验室能登台献唱。

另一个训练范式是不需要使用超级计算机，而可以通过普通PC从志愿者进行众包计算，这种模式被称为志愿者计算（volunteer computing，VC) 。这种模式的主要挑战在于如何利用每一台志愿者的性能。VC最适合那些可以分成许多独立块的任务，任务间不需要太多通讯，VC的案例有Folding@home、SETI@home、Einstein@home等。

自BOINC发展以来，有足够公众宣传的研究组织已经能够在志愿者提供的设备上运行大规模的科学计算。Folding@home等成功的项目可以拥有超过105名活跃参与者，可以与千万亿级别超级计算机的浮点性能相媲美。事实上，Folding@home是第一台达到1~10 PFlops里程碑的“超级计算机”。
有关超级计算机的知识可以参考本人这篇博客：Super-ComputerTop10

现有的大模型分布式训练需要节点不断传输大量的中间数据，并且采用的都是高性能计算机（High Performance Computer，HPC），这对于消费级PC来说通信成本过高，所带来的时延是难以接受的，也就不适合采用志愿者计算这一模式。本文跳出现有的分布式训练策略，通过确定每个志愿者计算的优势，加以利用。

---

2.内容解读

本文讨论了使用众包计算能力训练大型神经网络的挑战和潜在解决方案。作者强调了在专用硬件上训练大型模型的高成本，并提出了一种名为Learning@home的新颖训练范式。他们分析了该范式的性能、可靠性和架构约束，并将其与现有的分布式训练技术进行了比较。具体贡献有如下几点:

1. 提出了分散式专家混合（DMoE）:一个设计用于使用大量不可靠的消费级硬件进行训练的层
2. 设计了用于训练由DMoE层组成的大型神经网络的框架
3. 进行仿真实验检验方法效率和可靠性，并开源了复现代码

2.1 场景搭建

1. 大量性能不佳的计算机：基础架构由$10^3-10^6$台异构PC组成
2. 通信：节点以家庭互联网连接的速度和可靠性进行通信
3. 节点故障频繁：计算节点可能由于各种原因而无法处理任务。我们预计，在正常操作条件下，每天至少有5-20%的计算机出现一次故障。

2.2 关键技术

A 常见分布式训练测略 ×

分布式训练测略	处理方式
同步数据并行训练	每个工作器存储模型参数的副本，计算训练批次的一部分的梯度。然后，梯度在工人之间平均并应用于模型，在所有机器上构成相同的更新。
异步训练	单个参数服务器和多个计算节点获取最新参数，处理批次，并将更新提交回服务器。
模型并行训练	每个节点存储一部分模型层，每个训练批次由所有节点按照层分布方案确定的顺序处理。训练批次可以分为几个微批次并以流水线方式处理，从而显著提高硬件利用率
联邦学习	通过将数据保存在设备上，训练模型的本地版本，并仅发送参数更新

联邦学习为了隐私而牺牲了性能。安全聚合过程需要多个工作者进行通信，并根据其数量进行二次扩展。这使得联邦学习不适合在上述场景下训练大型模型。

B 去中心化MoE（DMoE） √

MoE可以用作神经网络中的一个层，并通过反向传播进行联合训练。根据任务的不同，专家可以使用卷积层、递归层或其他专门的层。这类层可以由多个专家组成，专家数量的拓展可以提高模型的性能表现。但当专家从从数千名扩展到数百万名时，门控系统就会面临巨大的任务分配压力。
为了解决这个问题，一个方案是优化搜索，还有一个是去中心化。后者采用了一个叫分布式哈希表（DHT） 的数据结构，它们在网络中的多台计算机上存储键值对。DHT有多个变体，但这些结构都有以下特征：

1. 权力下放：节点形成和维护DHT而无需任何中央协调;
2. 可扩展性：DHT可以扩展到数百万个不断加入和离开的活动节点;
3. 容错：一个或几个节点的故障不影响DHT的完整性和可用性;

假设有K个专家被选择进行训练，该过程最多需要O(KlogN)DHT查询来定位所选择的专家，并与这些专家进行K次直接交互来进行实际处理。此外这种结构设计可以解决场景中的很多难题：

1. 容错能力。如果所选的k个专家中的一些专家由于硬件或网络错误而未能响应，DMoE可以将这些专家排除在求平均值之外。
2. 志愿者硬件。计算节点可以基于其硬件能力来服务不同数量的专家。如果一个节点离开网络，另一个节点可以通过从DHT检索最新的专家检查点来取代它。
3. 负载平衡。专家混合层可以进行正则化，以平衡他们选择混合中每个专家的速度。
4. 异步训练。专家处理完一轮任务后可以立即开始处理下一个批次，而不是等待一个训练批次的结果。

C 结构化门控

结构化门控功能是Learning@home基础结构的关键组件，旨在选择适当的专家以分散的方式处理输入。DMoE中的专家被组织成一个d维网格，其中每个专家 f 与一个唯一的整数元组uid(f)=[$u_0,u_1...u_{d-1}$],$u_i\in [0,M)$相关联。d和M的选择是为了容纳所有具有一定冗余程度的专家。如果有更多的志愿者在培训中途加入，DMoE就可以分配更多的专家。
门控函数由d个线性层组成，记为g0, g1，…, gd-1。$g(x,f)={\sum }_{i=0}^{d-1}{g}_i(x)[u_i]$。门控函数可以在对数时间内选择得分最高的前k名专家。
一旦所有专家完成处理，工作者通过加权平均来汇总专家输出：
$$DMoE(x)=\sum _{f\in TopK(x)}f(x)\frac{exp(g(x,f))}{{\sum }_{f^{\prime }\in TopK(x)}exp(g(x,f^{\prime }))}$$
如果一些专家已经崩溃或花费太长时间来执行计算，我们可以将他们排除在平均值之外，并重新归一化权重，以便他们仍然加起来为1。使用此排除策略进行培训后，DMoE将学习具有重叠专业化的专家，这些专家对单个节点故障的抵抗力更强。

D 训练框架

• Trainer: 生成批任务并传播。在形成批处理并将其转换为输入向量之后，训练器在DMoE层的序列上迭代并组织向前和向后传递。
• Runtime：负责专家推理和培训。一旦所有的experts都被初始化，runtime就会监听来自trainers的传入连接，并处理两种类型的请求：
  1. Forward：给定输入，计算并返回这些输入的专家输出。
  2. Backward：给定损失函数的输入和梯度。通过梯度下降输入和更新专家参数。
• DHT Node: 每个运行时周期性地向DHT通告其专家，将其标识符与该运行时的地址和当前时间戳相关联。然后，Trainer可以使用这些条目来查找负责所选专家的工作人员。

---

2.3 实验结果

模型吞吐量：

收敛情况：

NLP表现

---

3.文章总结

该文章提出了训练大型神经网络的挑战，并提出了一种利用众包计算能力来克服这些挑战的新范式。提出的方法Learning@home利用分散式混合专家(DMoE)，解决了现有分布式训练方法在志愿计算环境中的局限性。