[论文解读] Rail-only: A Low-Cost High-Performance Network for Training LLMs with Trillion Parameters

Weiyang Wang, Manya Ghobadi+3 more|arXiv (Cornell University)|Jul 22, 2023

Topic ModelingComputer Science被引用 3

一句话总结

本文提出 Rail-Only，一种通过利用万亿参数大语言模型的独特通信模式，实现低成本、高性能的网络架构，用于训练万亿参数大语言模型。与全连接的任意对任意通信不同，它将 GPU 划分为高带宽（HB）域，并仅通过存在非零流量的跨域连接进行连接，从而在不损失性能的前提下将网络成本降低 37% 至 75%。

ABSTRACT

This paper presents a low-cost network architecture for training large language models (LLMs) at hyperscale. We study the optimal parallelization strategy of LLMs and propose a novel datacenter network design tailored to LLM's unique communication pattern. We show that LLM training generates sparse communication patterns in the network and, therefore, does not require any-to-any full-bisection network to complete efficiently. As a result, our design eliminates the spine layer in traditional GPU clusters. We name this design a Rail-only network and demonstrate that it achieves the same training performance while reducing the network cost by 38% to 77% and network power consumption by 37% to 75% compared to a conventional GPU datacenter. Our architecture also supports Mixture-of-Expert (MoE) models with all-to-all communication through forwarding, with only 8.2% to 11.2% completion time overhead for all-to-all traffic. We study the failure robustness of Rail-only networks and provide insights into the performance impact of different network and training parameters.

研究动机与目标

• 解决训练万亿参数大语言模型时大规模 GPU 集群日益增长的成本与复杂性问题。
• 挑战全连接任意对任意网络互连在大语言模型训练中为必需的假设。
• 设计一种与实际大语言模型通信模式相匹配的成本高效网络，以减少基础设施开销。
• 在最小化网络硬件的前提下保持训练性能，通过移除闲置互连实现。
• 提供一个分析框架，基于大语言模型超参数和基础设施条件，估算训练迭代时间。

提出的方法

• 分析大语言模型训练的通信流量模式，发现高带宽通信仅限于小型、局部的 GPU 组（HB 域）。
• 提出一种新型网络架构 Rail-Only，将集群划分为仅通过必要跨域链路连接的 HB 域。
• 使用分析模型估算训练迭代时间，综合考虑计算、通信和并行化策略。
• 通过消除传统 Clos 网络拓扑中未被利用的链路，实现网络成本降低。
• 为 Rail-Only 设计实现容错机制和故障恢复策略，提升系统容错能力。
• 将分析模型与先前研究的实测结果进行对比验证，浮点利用率估算误差低于 0.15%。

Figure 1: State-of-the-art GPU clusters are based on rail-optimized, any-to-any Clos networks Nvidia ( 2023a ) .

实验结果

研究问题

• RQ1 大语言模型训练是否真的需要集群中所有 GPU 之间都实现任意对任意的网络连接？
• RQ2 能否设计一种网络架构，使其与大语言模型训练的实际通信需求相匹配，从而降低成本？
• RQ3 如何准确建模训练迭代时间，以指导大语言模型网络架构的设计？
• RQ4 在大规模 GPU 集群中，移除非必要互连的性能与成本权衡如何？
• RQ5 所提出的 Rail-Only 架构能否在显著降低网络硬件的同时保持性能？

主要发现

• 大语言模型训练表现出一种通信模式：仅小型、局部的 GPU 组需要高带宽的任意对任意通信，而组间通信量极少且趋于均匀。
• 与最先进 Clos 网络相比，Rail-Only 架构在不损失性能的前提下，将网络成本降低了 37% 至 75%。
• 用于训练迭代时间的分析模型，其硬件浮点利用率估算值与先前研究的实测结果相比误差低于 0.15%。
• 对于 32,768 个 GPU 的集群，Rail-Only 将交换机数量从 SOTA Clos 网络的 2,560 台减少至 1,536 台，实现 37% 的成本降低。
• 对于 128,000 个 GPU 的集群，通过将顶层配置的交换机数量从 1,280 台减少至 256 台，成本降低达到 75%。
• 通过支持集群的逻辑分区为独立域，该设计仍可兼容多任务训练和推理工作负载。

Figure 2: Different parallelization strategy for training LLMs. Each color represents a different set of devices.