[控场AI]
· 5 分钟阅读· 2,907 字

AMD RCCL GPU发起网络集成TorchTitan:MoE通信性能提升15%

AMD RCCL GPU发起网络集成TorchTitan:MoE通信性能提升15%

AMD将GPU发起网络(GIN)集成至TorchTitan,使MoE模型端到端训练性能提升12%-15%。

AMD工程师将RCCL团队研发的GPU发起网络(GIN)技术集成到PyTorch官方分布式训练框架TorchTitan中,重点针对MoE(混合专家)模型高频、全量的All-to-All通信进行优化。GIN的核心思路是让GPU内核直接发起网络传输,绕过传统流程中CPU调度的中间环节,从而压缩通信延迟。对于节点间通信,GIN采用网络put操作;对于节点内通信,则利用本地对称访问实现高速数据搬运。实测结果显示,all-to-all内核层面性能最高提升30%,折算至端到端的整体训练流程后仍有12%至15%的收益。这一成果验证了"降低通信开销可大幅提升训练效率"的判断,也展示了AMD在构建完整AI训练软件栈方面的进展,相关细节将在PyTorch Conference上进一步披露。

分布式训练的通信瓶颈难题

在大规模分布式训练中,通信往往是制约端到端性能的核心瓶颈。随着模型参数规模不断膨胀,GPU之间频繁的数据交换会占用大量时间,这部分开销如果不能有效压缩,就会直接拖累整体训练效率。来自AMD的工程师Rishi在最新分享中指出,缩短通信耗时能够显著提升端到端性能,而这正是其团队工作的出发点。

限制通信耗时可大幅提升性能

AMD团队的做法是将RCCL团队(视频中称RICL团队)在GPU发起网络方面的成果,集成到TorchTitan中。TorchTitan是PyTorch官方推出的分布式训练框架,专为大模型训练设计。将底层的通信优化能力与上层训练框架打通,意味着这项技术能够直接惠及实际的训练工作负载,而不只是停留在实验室基准测试阶段。

什么是GPU发起网络(GIN)

GPU发起网络(GPU-Initiated Networking,简称GIN)是这项工作的技术核心。传统的网络通信通常需要CPU参与调度和发起数据传输,而GIN的关键区别在于:GPU内核本身直接向远程GPU发出网络put操作。这种机制省去了CPU的中间环节,让数据传输的发起点更贴近计算本身。

来自RCCL团队的GPU发起网络技术

这种设计带来两方面优势。对于节点间(internode)通信,GPU直接发起网络传输能够优化跨节点的数据交换效率;对于节点内(intranode)传输,则通过本地对称访问(local symmetric access)实现快速的数据搬运。两种场景的协同优化,构成了GIN性能提升的基础。

GPU发起网络(GIN)的定义

节点内与节点间的差异化优化

理解GIN的价值,需要区分两类通信场景。节点间通信跨越物理服务器,受网络带宽和延迟影响较大;节点内通信则发生在同一服务器的多块GPU之间,通常走更高带宽的互联通道。GIN针对这两种场景分别采用网络put和本地对称访问,针对性地压缩各自的通信开销。

本地对称访问实现快速节点内传输

传统通信流程中,CPU负责调用网络驱动、准备描述符并触发DMA传输,GPU必须等待CPU完成这些准备工作后才能继续推进计算图。这种"CPU中介"模式在集群规模扩大时会产生显著的调度延迟,尤其是当通信操作频率极高时,CPU成为整条流水线的节拍器,反而拖慢了本可并行的GPU计算。GIN通过在GPU内核中直接调用RDMA(远程直接内存访问)原语,将通信的发起权完全交还给GPU线程,消除了CPU往返调度的固定开销,也让通信与计算的重叠(overlap)更易实现。这与NVIDIA在NVLink Direct和SHARP等技术上的探索方向一致,代表了高性能互联领域的共同趋势:让数据尽可能在计算侧自主流动,减少控制面的干预。

MoE模型中的All-to-All通信

这项工作的落脚点是MoE(Mixture of Experts,混合专家)模型中的通信模式。MoE架构通过将不同的token路由到不同的专家网络,实现了参数规模的大幅扩展,但同时也引入了密集的all-to-all通信需求——每个GPU需要向其他所有GPU发送和接收数据。这种全交换模式对通信性能极为敏感,也因此成为优化的理想切入点。

AMD团队将GIN应用于TorchTitan的all-to-all操作后,取得了明确的性能收益:内核(kernel)层面的性能提升最高可达30%,而端到端的整体训练性能提升了12%到15%。从内核加速到端到端收益之间的差距,反映出通信在整个训练流程中的占比——即便单个通信内核快了30%,由于计算等其他环节的存在,整体提升会被摊薄到12%至15%这一区间。

All-to-All是分布式集合通信中开销最重的原语之一。在MoE架构里,路由机制(router)决定每个token由哪位专家处理,而专家通常分散在不同GPU上,因此每轮前向传播都需要一次全量的token重分配——每块GPU既要向其余所有GPU发送属于对方专家的token,又要接收属于自己专家的token。与AllReduce(主要用于梯度同步)不同,All-to-All的流量模式高度不规则且难以预测,负载均衡稍有偏差就会出现热点,导致部分GPU等待时间远长于其他GPU。这也是为什么MoE场景是GIN优化的"天然靶场":通信量大、频率高、对延迟敏感,任何降低单次通信开销的手段都能被快速放大为可测量的端到端收益。

性能数据背后的意义

12%到15%的端到端提升在大规模训练场景中并非小数目。对于动辄需要数周、消耗大量GPU算力的大模型训练任务而言,这一比例意味着可观的时间与成本节省。更重要的是,这验证了"降低通信开销能够大幅提升端到端性能"这一假设,为后续在AMD硬件生态上的通信优化指明了方向。

从技术栈的角度看,这项工作打通了RCCL通信库、GIN网络机制与TorchTitan训练框架三个层次。RCCL作为AMD对标NVIDIA NCCL的集合通信库,其GPU发起网络能力的成熟,对于AMD在AI训练市场的竞争力有着实际意义。AMD表示将在即将举行的PTC(PyTorch Conference)上进一步分享这项工作的细节。

RCCL(ROCm Collective Communications Library)是AMD为ROCm平台开发的集合通信库,功能上对标NVIDIA的NCCL(NCCL, NVIDIA Collective Communications Library)。两者都封装了AllReduce、AllGather、All-to-All等分布式训练常用原语,并针对各自的硬件互联(AMD的xGMI/Infinity Fabric vs. NVIDIA的NVLink)做了深度优化。对于希望在AMD GPU集群上运行PyTorch训练任务的用户,RCCL通常通过环境变量或插件机制透明替换NCCL,上层框架代码无需修改。此次GIN能力被纳入RCCL并进一步集成到TorchTitan,意味着AMD正在将底层硬件优势系统性地传导至主流训练框架,缩短从硬件创新到开发者可用之间的距离。

小结

AMD通过将RCCL的GPU发起网络能力集成到PyTorch的TorchTitan框架,针对MoE模型的all-to-all通信实现了内核最高30%、端到端12%至15%的性能提升。这一成果体现了GPU直接发起网络传输在分布式训练中的潜力,也展示了AMD在构建完整AI训练软件栈方面的持续投入。对于关注大模型训练效率的开发者和研究者来说,GIN这类绕过CPU、让GPU直接掌控网络通信的技术路线,值得持续关注。

分享:

相关推荐