Nou Research宣布使用与架构和网络无关的分布式优化器成功将训练LLM时GPU间的通信量降低了1000到10000倍，并介绍了相关实验和结果。该优化器被称为DisTrO，可用于训练大型神经网络，通过降低通信开销来提高训练效率，同时还具有通用、可扩展和时钟同步等特点。除了训练LLM，DisTrO还可应用于联邦学习和创建虚拟异构GPU集群等领域。
DisTrO通过不依赖摊销分析的方式，将GPU间通信要求降低了四到五个数量级，解决了训练大规模神经网络的通信开销问题。
DisTrO是通用、可扩展和时钟同步的，能够支持分布式数据并行训练（DDP），并且对于神经网络架构和电信网络的拓扑不敏感。此外，它还可以创建完全去中心化且无需许可的网络来协作和共享资源。
实验证明，DisTrO可以在不影响训练效果的情况下，将通信量大幅度降低。例如，在训练LLM的实验中，通信量从74.4GB减少到86.8MB，相当于带宽压力减少了857倍。此外，DisTrO还具有弹性和适应性，可以适应新节点的加入和少量降级或丢弃的节点。
DisTrO可用于联邦学习、创建虚拟异构GPU集群等领域。其进一步大规模应用可能会缓解大型数据中心所带来的能源消耗、基础设施成本和土地使用等相关问题。
【新智元导读】
近日，Nous Research宣布了一项重大突破，通过使用与架构和网络无关的分布式优化器，研究人员成功将训练LLM时GPU间的通信量降低了1000到10000倍！
如果可以使用世界上所有的算力来训练AI模型，会怎么样？
近日，凭借发布了开源的Hermes 3（基于Llama 3.1）而引起广泛关注的Nous Research，再次宣布了一项重大突破——DisTrO（分布式互联网训练）。
通过使用与架构和网络无关的分布式优化器，研究人员成功将训练LLM时GPU间的通信量降低了1000到10000倍！
初步技术报告：https://github.com/NousResearch/DisTrO/
在如此夸张的改进之下，大模型训练的重要成本和瓶颈——带宽，也就不再是问题。
使用DisTrO的方法，你可以将训练负载分布到互联网上，而整个网络世界也就成为了一个巨大的异构的AI服务器集群。
——任何有相关算力的设备都可以参与到训练过程之中。
实验证明，本文的方法基本不会导致模型性能下降，同时DisTrO-AdamW在收敛速度方面，也与标准的AdamW+All-Reduce相当。
一般来说，训练大规模神经网络涉及到大量的通信开销。
比如做数据并行的时候，不同的训练数据在不同的硬件（显卡等）上进行前向和反向计算，之后，同一批数据计算出的梯度需要在显卡之间先完成同步，才能进入下一个epoch。
如果是模型并行，那么中间数据就需要通过All-Reduce进行拼接或者累加。
这些数据通信开销如果不能overlap掉，就会成为模型训练的瓶颈。
而恰好，老黄的显存和带宽又很贵，甚至组多卡时候需要的硬件也很贵。
为了解决这个问题，研究人员开发了DisTrO，在不依赖摊销分析的情况下，将GPU间通信要求降低了四到五个数量级，从而能够在慢速网络上对大型神经网络进行低延迟训练。
DisTrO是通用、可扩展，并且时钟同步的（与SGD、Adam等类似，每个训练步骤使用相同的算术运算并花费相同的时间）。
另外，与之前的ad-hoc低通信优化器相比，DisTrO对电信网络的拓扑和神经网络架构不敏感，能够以最小的开销原生支持分布式数据并行训练（DDP）。
LLM预训练
研究人员使用Nanotron作为预训练框架，且仅在DDP策略下运行（每个GPU都将整个模型加载到VRAM中）。
LLM选择1.2B大小的Llama 2，模型和训练所用的超参数如下：
训练数据使用Dolma v1.7数据集，随机选出的10%代表性样本（前 105B个token）。
优化器采用AdamW，β1=0.9、β2=0.95，峰值学习率为4×10e-4，使用余弦衰减方案，权重衰减设置为0.1。
作为对比的另一组实验，将AdamW替换为DisTrO-AdamW，但不更改超参数，并禁用Nanotron中的All-Reduce操作。
与以前的分布式训练方法不同，DisTrO不同步优化器状态（甚至可以无状态）。
下图是两组实验的训练损失曲线，使用105B数据训练25000步。可以看出，DisTrO的收敛能力与All-Reduce持平。
重要的是，在不影响训练效果的情况下，DisTrO将通信量从74.4GB直接减到了86.8MB！相当于带宽压力减少了857倍。
作者还表示，这857倍只是初期测试，后面调调超参数，减少个1000倍到3000倍也不是问题。
如果是后训练和微调，甚至可以实现高达10000倍的通信优化，且基本不影响训练效果。
最后，为了验证训练效果，作者在训练后的模型上执行了GPT4All零样本基准测试，并与在相同数量的token上训练的TinyLlama（checkpoint）进行了比较。
结果如上表所示，TinyLlama的架构和训练过程与本文的实验非常相似，可以作为对结果进行健全性检查的衡量标准。
未来应用
数据流
在本实验的场景中，32个节点使用最简单的All-Reduce（全连接），每个节点平均传输86.8MB（2.8MB×31），并接收相同数量的数据。
如果使用专用服务器进行数据聚合，则每个节点只需上传2.8MB数据（接收数据不变），通信量进一步减少。
另外，不对称性是有优点的，因为大多数消费互联网的带宽严重偏向于更高的下载速度。
假设稳定的网速为100Mbps下载和10Mbps上传，则最坏情况下的延迟仅为下载6.94秒，上传2.24秒，overlap一下则每步延迟为6.94秒。
ps：以上的数据传输都是原始的向量，如果用上压缩技术还能更快。
带宽
作者表示，目前的实验和研究还比较有限，无法断定随着模型变大，带宽减少的比率是会增加、减少还是保持不变。
不过目前的1.2B似乎是DisTrO能够良好工作的最小尺寸（再小就不收敛了），所以可以假设随着模型大小的增长，需要的通信会相对越来越少。
不过也可能通信量与模型大小没有关系，这时可以在不增加通信带宽的情况下增加模型大小，观察更大的模型是否会改善训练和学习的效果。
如果后一种情况属实，那么未来GPU设计和制造的范式将会被改变（更大VRAM和更窄带宽）。
恰好我们也更喜欢计算密集型负载（而不是I/O密集型），毕竟现在的带宽要比计算贵得多。
