【论文阅读】Rail-only:A Low-Cost High-Performance Network for Training LLMs with Trillion Parameters

• 链接：https://arxiv.org/pdf/2205.05198.pdf
• 发表会议：HOTI’24(CCF-C)
• 团队：Meta, MIT

背景

训练大模型的大规模集群往往采用全连接的CLOS架构，而将 Clos 网络扩展到数万个 GPU并不容易并且造价高昂。而实际上训练大模型并不需要所有GPU都可以全连接，通过分析大模型训练的网络流量模式可以发现大部分通信都发生在rank相同的GPU对之间。所以本文构建了一个Rail-Only网络架构，其移除了交换机的脊梁层，仅连接具有高网络流量的GPU组。

方法思路

LLM训练流量分析

• 采用与Megatron-LM论文中相同的3D并行，其主要网络流量如下图所示，有：
  • TP的All Gather、Reduce Scatter流量
  • DP的All Reduce流量
  • PP的点对点流量
• 下图(a)展示了各类流量的分布，TP流量占75%以上，但是TP流量一般局限在主机内部，也就是HB域，而主机间也就是NIC域的流量较低。
• 下图（b）展示了各GPU对之间拥有的通信模式的比例分布，可以看到超过99%的GPU对之间没有流量。
• 下图详细分析了各GPU对之间的通信流量关系，下图（a）展示了一个流水线阶段的流量模式，下图（b）展示了4个流水线阶段的流量模式。

分层集体通信算法

• 借助主机域内通信带宽高，主机间通信带宽低的特点，可以先将部分数据在跨主机的GPU间进行传递，然后再在主机内进行进一步的整合。

Rail-Only设计

• 如下图是传统的Rail优化的GPU集群拓扑：
• Rail-Only设计下去除了脊交换机，如下所示：
• 在这种设计下，全对全的流量都需要转发，rank不对等的GPU间的流量也需要二次转发，但是考虑到主机内的网络带宽更大，并且大模型训练中往往都是大量的主机内流量，少量主机间流量，而且基本也是rank对等的主机间流量，所以代价可接受。

实验效果

• 并没有开展真实实验，而是用理论模型推导，首先与【Reducing activation recomputation in large trans-former models】中的展示的实际硬件浮点运算利用率与理论模型推动结果进行对比，得到随着规模增大，理论模型推导结果就越接近。
• 还测试了改变主机GPU数量下迭代时间的理论优化效果
• 还分析了可用带宽、批次大小对于迭代时间的影响等

总结

• 对LLM训练流量分析确实很全面。
• Rail-only完全是面向LLM训练设计的，但是我怀疑在使用过程中可能需要做一些客制化的操作。并且怀疑多个任务场景下可能性能下降会更严重。
• 最后没有做实际的实现有点可惜，不知道现在是否真的有用上了。