一文揭秘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推理和数据存储场景会在接下来的文章中呈现，敬请关注。

AI训练程序首先将模型参数加载到GPU内存中，之后将经历多个epoch（即使用所有训练集对模型进行一次完整训练），每个epoch的处理过程可以简单描述为4步：

1. 加载训练数据，在每个epoch中，根据batch size将整个数据集分为若干个mini-batch，分批次加载训练数据，直到遍历整个训练数据集。
2. 训练，包括前向传播、计算损失、反向传播和参数/梯度更新，每个mini-batch都进行上述步骤。
3. 评估，使用评估数据集对模型的指标进行评估。这一步是可选的，可以在整个训练完成后单独进行，也可以间隔若干个epoch进行一次。
4. 保存checkpoint，包括模型状态、优化器状态和训练指标等。为了减少存储需求，通常经过多个epoch后保存一次。

在大模型出现之前，整个过程在可在一台AI服务器内部完成，训练程序从服务器本地磁盘读取AI模型和训练集，加载到内存中，完成训练、评估，然后将结果存储回本地磁盘。虽然为了加速训练，也会采用多块GPU同时训练，但所有的I/O均发生在一台AI服务器内部，并不需要网络I/O。

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

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

首先是模型的参数规模超出了单个GPU的内存，采用GPU集群协同计算，则需要相互之间通信以交换信息，这类信息包括参数/梯度、中间激活值等。

庞大的数据集被所有GPU共享，需要集中存放到远端的存储服务器中通过网络调用，分批加载到GPU服务器上。此外，定期保存的参数和优化器状态也需要通过存储服务器共享，在每个训练epoch中，都要通过网络读写数据。

由此，AI大模型训练的网络流量可分为以下两类：

• 第一类是GPU之间同步梯度和中间激活的网络流量，它发生在所有GPU之间，是一种广播式流量，逻辑上需要所有GPU全连接。
• 第二类是GPU和存储服务器之间的流量，它仅仅发生在GPU和存储服务器之间，是一种单播流量，逻辑上仅需要以存储服务器为中心的星型连接。