揭秘AI智算中心网络流量：大模型训练的背后技术与挑战

本文链接：https://blog.csdn.net/njbaige/article/details/145240097

一文揭秘AI智算中心网络流量 - 大模型训练篇

前言：自2017年起，AI模型的规模每半年翻一番，从初代Transformer的6500万增长到GPT-4的1.76万亿，预计下一代大语言模型将达到10万亿规模。另一方面，用于模型训练的数据量持续增长，如C4数据集，原始数据量累计超过9.5PB，每月新增200-300TB，目前经过清洗加工后的数据集大小约38.5 TB，训练样本数364.6M。进一步，随着多模态大模型的兴起，训练数据从单一的文本过渡到图像和视频乃至3D点云，数据规模将是文本数据的1万倍以上。
AI模型规模庞大且快速扩张，预示着数据中心流量呈指数级增长。这一独特的流量特征将为网络带来新的挑战与机遇。

深入剖析AI大模型在训练、推理和数据存储方面的优势，助数据中心建设者降低成本、提高速度、优化网络，从而为用户带来更优质的服务。

AI训练的常规工作流

AI训练程序首先将模型参数加载到GPU内存中，之后进行多个epoch的训练。每个epoch的处理过程可以简单描述为4步：加载训练数据，在每个epoch中，根据batch size将整个数据集分为若干个mini-batch,分批次加载训练数据，直到遍历整个训练数据集。

在训练过程中，通过每个epoch按batch size划分数据集为多个mini-batch,逐一加载以遍历整个训练数据集。
在训练过程中，我们采用前向传播、计算损失、反向传播和参数/梯度更新的策略，每个小批量样本都会经历这些步骤。
在训练模型过程中，评估数据集对模型指标至关重要。这一步骤可选择在整个训练结束后进行，或每隔若干epoch进行一次。
"保存训练的关键点，如模型参数、优化器状态和进度指标等。为降低存储负担，每过若干轮迭代后进行一次整体备份。"

在AI大模型出现前，训练过程主要在单台服务器内完成。这包括从服务器本地磁盘读取模型和训练集，加载到内存进行训练和评估，最后将结果存回磁盘。尽管采用多GPU并行训练以加速，但所有I/O都在服务器内部进行，无需网络I/O。

AI大模型训练的网络流量有哪些？

进入大模型时代，AI训练的流量路径和其网络需求发生了巨大变革。

首先是模型的参数规模超出了单个GPU的内存，采用GPU集群协同计算，则需要相互之间通信以交换信息，这类信息包括参数/梯度、中间激活值等。庞大的数据集被所有GPU共享，需要集中存放到远端的存储服务器中通过网络调用，分批加载到GPU服务器上。

此外，定期保存的参数和优化器状态也需要通过存储服务器共享，在每个训练epoch中，都要通过网络读写数据。由此，AI大模型训练的网络流量可分为以下两类：第一类是GPU之间同步梯度和中间激活的网络流量，它发生在所有GPU之间，是一种广播式流量，逻辑上需要所有GPU全连接。第二类是GPU和存储服务器之间的流量，它仅仅发生在GPU和存储服务器之间，是一种单播流量，逻辑上仅需要以存储服务器为中心的星型连接。

并行训练：AI智算中心的主要流量来源

当前广泛应用于AI训练并行计算模式主要有以下三类：

数据并行	将不同的样本数据分配给不同的GPU，以加快训练速度；用在主机之间
张量并行	将模型的参数矩阵划分为子矩阵，并分配到不同的GPU上，以解决内存限制并加速计算。一般用在主机内部。
流水线并行	将模型分为多个阶段，每个阶段分配给不同的GPU，以改善内存利用率和资源效率。一般用在主机之间

常见的集合通信流量模式（如下图）

数据并行（Data Parallelism）
在数据并行时，主要的网络流量来源于梯度同步，它发生在每次mini-batch处理之后，由一个all-reduce操作计算平均值。理想情况下，所有GPU全连接，每个GPU给其它G-1个GPU单独发送数据，共需发送G x（G-1）份数据。
FSDP(完全分片数据并行)是一种改进的数据并行技术，旨在优化内存使用和通信效率。它通过将模型参数和梯度在多个GPU之间分片（shard）存储，实现更高效的内存利用和通信。
在FSDP时，网络流量来自前向传播的参数收集以及反向传播中的梯度同步。前向传播的参数收集由all-gather操作完成，all-gather的通信复杂度与all-reduce相同。后向传播的梯度同步由all-reduce操作完成，由于每个GPU的参数只有原来的1/G，一个epoch中总的网络流量只有普通数据并行的1/G。

张量并行（Tensor Parallelism）

在张量并行时，模型参数分布到G个GPU上，每个GPU只存储1/G参数。网络流量主要来自前向传播过程的中间激活值的传递以及反向传播过程中的梯度同步。前向传播中，各个GPU计算出的中间激活值需要合并，由一次all-reduce操作进行求和。对于每个Token，在模型的每个layer的中均进行2次合并，共2xTxL次通信。反向传播中，梯度需要在GPU之间同步，这种在每一层的处理中发生2次，由all-reduce操作将各个GPU上梯度求和。这种同步发生在每个mini-batch的每个layer的处理过程中。共2×N×L次通信。

03流水线并行（Pipeline Parallelism）

在流水线并行时，网络流量主要来自前向和反向传播过程的中间激活值的传递。与张量并行不同，这些流量的传递发生在模型的前后两个阶段之间，使用Point-to-point通信而非all-reduce操作，并且频率也大大减小了。
综上，在三种并行方式中，张量并行的网络流量最大、频率最高，流水线并行的流量最低，数据并行的通信频率最低。如下表所示，P为模型参数，T为token数，L为模型层数，H为隐藏状态大小，G为GPU数量，N为mini-batch的数量，采用BFLOAT16数据格式，每个参数占2个字节。在每个epoch过程中：

	流量模式	后向传播总网络流量	反向传播同步次数	前向过程总网络流量	前向过程传递次数
数据并行	all-reduce	2 × N × P × G × (G-1)	1	0	0
FSDP	all-gather + all-reduce	2 × N × P × (G-1)	L	2 × N × P × (G-1)	L
张量并行	all-reduce	4 × N × P × L × (G-1)	2 × L	4 × L × T × H × (G-1) × G	2 × L × T
流水线并行	Point-to-point	2 × T × H × (G-1)	G-1	2 × T × H × (G-1)	G-1

以Llama3 70B(P)模型和C4数据集为例，采用BFLOAT16数据格式，每个参数占2个字节，隐藏层维度设为8192(H),使用8个GPU(G)进行数据并行。C4数据集的token(T)总数约为156B,样本数为364.6 million;batch size为2048,每个epoch包含约178,000个mini-batch(N)。

计算可得每个epoch过程中：

	反向传播总网络流量（PB）	反向传播同步次数	前向过程总网络流量（PB）	前向过程传递次数
数据并行	1396 PB	1	0	0
FSDP	175	80	175	80
张量并行	26622	160	21840	16015610^9
流水线并行	17.9	7	17.9	7

3D并行技术下的网络流量

"通过结合数据并行、张量并行和流水线并行，我们能更有效地训练大型模型，同时提高其效率和扩展性。这便是GPU集群的构成方式，它在这三个维度上发挥着重要作用。"

在假设有G(tp)×G(pp)×G(dp)个GPU组成的3D并行阵列中，P个参数将被分割为G(tp)×G(pp)份，每份大小为P/G(tp)/G(pp)。模型并行、流水线并行和数据并行在三个维度上都会带来网络流量。接下来，我们将深入探讨每个epoch的训练过程，详细计算各阶段的网络流量组成和规模。

反向传播中的网络流量

在每个mini-batch中，反向传播时的梯度同步分为：

在模型的每一层和数据维度的每一组中进行张量维度上的梯度同步，总共 LxG(dp) 次，每次包含2个all-reduce操作。
在数据维度上，梯度同步在流水线阶段和张量组中进行，总计G(tp)xG(pp)次。每次操作包括一个all-reduce操作。

如下图所示：

这样，在一个epoch中，梯度同步的总网络流量为：

3、流水线并行维度的中间激活梯度传播，流量为：

因此，在一个epoch中，整个反向传播的总流量为：

02前向传播中的网络流量

在前向传播过程中，中间激活层通过张量并行和流水线并行维度进行交替计算。具体来说，张量并行的激活传递每次涉及两个all-reduce操作。以下图为例，展示了一个Token的前向传播过程：

因此，在一个epoch中，前向传播总网络流量为：

即：

以Llama3-70B模型为例，我们采用了8路张量并行、8路流水线并行和16路数据并行的模式进行训练。在总共1024个GPU的支持下，一个epoch产生的总流量约为$85\times10^9$次方字节(EB)。如此巨大的流量规模，若使用1台交换容量为$51.2\times10^12$字节(TB)的交换机，在24小时满负荷运行的情况下，需要大约20天才能传输完毕。
考虑到一次预训练通常包含100个左右epoch,如果需要在100天完成训练，至少需要20台51.2T交换机来传输训练过程产生的数据。

AI训练对智算中心网络的要求

超低时延：一个训练样本处理，瞬间产生100GB以上数据，传输速度超越1000个800G接口。
"探索集合通信：All-reduce和All-gather操作在GPU服务器间引发的广播式流量，跨越万个GPU,即上亿个GPU-GPU对的同步挑战。"
根据木桶原理，零容忍丢包会导致整个集体通信的延迟，进而造成大量GPU进入空闲等待时间。为了解决这个问题，可以使用零拷贝技术，它可以将数据从内核空间直接传输到用户空间，避免了数据在内核和用户空间之间的拷贝，从而提高了数据传输效率。此外，还可以使用多线程技术，将多个任务分配给不同的GPU进行处理，从而提高GPU的利用率。
时间同步关键：木桶原理下，若GPU时钟不同步，计算量大者将耗时更多，从而使计算快的GPU等待计算慢的GPU。

星融元CX-N系列交换机，专为智算中心AI训练场景打造，具备超低时延以太网交换功能。在保持卓越性能的同时，实现可编程、可升级特性，与计算设备协同作业，共同构建10万级别计算节点互联，将数据中心重塑为与超级计算机相媲美的AI超级工厂。

最大支持64个800G以太网接口，共51.2T交换容量。
超低时延，在800G端口上实现业界最强的560ns cut-through时延。
这句话的意思是：这个设备支持RoCEv2,Rail-only,全连接Clos以及200G/400G混合组网，可以灵活适应不同的算力中心建设方案。
以下是

200+ MB大容量高速片上包缓存，显著减小集体通信时RoCE流量的存储转发时延。
"搭载Intel至强处理器，内存容量巨大且可扩展，我们的企业级SONiC网络操作系统——AsterNOS,借助DMA直接访问包缓存，对网络流量进行即时处理，确保持续高效的运行。"
"FLASHLIGHT精准流量分析引擎，实时监控每个数据包的延迟与往返时间，借助CPU智能解析，实现路由自适应与拥塞控制。"
10纳秒级别的PTP/SyncE时间同步，保证所有GPU同步计算。
开放API,助力AI数据中心管理。通过REST API无缝集成计算设备，实现GPU集群自动化部署，提升运营效率。