模型大了难在哪里？

大家好，我是Tim。

自从GPT模型诞生以来，其参数规模就在不停的扩大。但模型并非简单的直接变大，需要在数据、调度、并行计算、算法和机器资源上做相应的改变。

今天就来总结下，什么是大模型，模型变大的难在哪里以及对于CV/NLP或者搜推广场景上有什么应对策略。

什么是大模型？

大模型，顾名思义主打的就是“大”。主要包括几个方面：

1. 数据规模大，通过大量的数据提高模型的泛化能力和性能。
2. 大规模并行计算能力，随着计算硬件的不断进步，如GPU和TPU的普及，大规模并行计算能力的提升使得训练和推理大模型成为可能。
3. 更“大”模型复杂性：大模型具备更深层次、更复杂的网络结构，可以捕捉更丰富的特征和关系，提高了模型的表达能力。

大模型主要分为两类：一个是稀疏大模型，另一个是稠密大模型。

1. 稀疏大模型：稀疏大模型是指模型中存在大量稀疏参数的情况，一般是搜索、推荐、广告类任务。它的特点是海量样本及大规模稀疏参数（sparse embeddings），适合使用 CPU/GPU 参数服务器模式（PS）进行训练。
2. 稠密大模型：稠密大模型是指模型中的参数大多数都是非零值，没有明显的稀疏性特征，一般是CV、NLP 任务。它的特点是常规样本数据及大规模稠密参数，它适合用纯 GPU 集合通信模式（Collective）进行训练。

对于搜推广类的稀疏大模型来说，一般包含稀疏特征的嵌入（embedding）和稠密模型两个部分。

其中，稀疏特征的嵌入计算是稀疏大模型的关键，而稠密模型部分一般往往较小，可以放到一个GPU内，因此可以进行data并行以及all reduce通讯。

在训练中，需要在特征嵌入表（embedding table）上需要进行复杂的查找、排列等操作，然后生成张量再做稠密模型的计算。

而特征嵌入表往往会占用非常大的存储空间，需要很多台GPU服务器才能完整存放，这就是典型的tensor并行。

在这样的场景下，就会导致典型的alltoall的通讯模式，而alltoall通讯会带来严重的incast通讯（多打一），进而带来网络拥塞，给网络架构、拥塞控制协议、负载均衡算法等都提出了很高的要求。

对于CV、NLP 任务来说，由于模型参数非常多，远远超过了单个GPU显存所能容纳的空间（NVIDIA最新的A100也就是80GB显存）。

所以往往既需要对模型某一层的tensor并行，也需要不同层之间的pipeline并行，才能放得下整个大模型。

在计算过程中，既有单机内通讯，也有不同机器间的通讯，具体的通讯模式取决于模型的切分和放置方法。

为了加速训练过程，往往完整的大模型之间也会采用data并行，每一个完整的大模型会被投喂不同的训练数据，这就会导致大家熟悉的allreduce通信模式。

总结起来，稠密大模型和稀疏大模型在模型特征上有着明显的差异，对计算/存储/通信资源的需求也存在明显的不同。

要达到GPU算力资源的最大化利用和最好的加速效果，需要结合模型特征和实现方式对GPU服务器架构、网络架构、训练数据存储和拉取、分布式训练框架进行全局的考量和设计。

模型大了难在哪里？

大模型带来的挑战主要有两点：海量样本、参数（万亿级别）和较长的收敛时间。

别看只有这区区两点，它会衍生出很多要解决的问题。

1. 大模型训练需要更大的算力。

大模型训练所需的总算力，其实很简单，6 * 模型的参数量 * 训练数据的 token 数就是所有训练数据过一遍所需的算力。这里的 6 就是每个 token 在模型正向传播和反向传播的时候所需的乘法、加法计算次数。

一堆矩阵相乘，简单来想就是左边若干个神经元，右边若干个神经元，组成一个完全二分图。选出其中任意一个左边的神经元 L 和右边的神经元 R。

• 正向传播的时候： L把它的输出乘上 L和 R 之间的权重 w，发给 R；R不可能只连一个神经元吧，总要把多个 L的加到一起，这就是 reduce，需要一次加法。
• R把它收到的梯度乘上 L和 R 之间的权重 w，发给 L；L也不可能只连一个 R，需要把梯度 reduce 一下，做个加法；别忘了权重 w 需要更新，那就要计算 w 的梯度，把 R 收到的梯度乘上 L正向传播的输出（activation）；一个 batch 一般有多个 sample，权重 w 的更新需要把这些 sample 的梯度加到一起。
• 一共 3 次乘法，3 次加法，不管 Transformer 多复杂，矩阵计算就是这么简单，其他的向量计算、softmax 之类的都不是占算力的主要因素，估算的时候可以忽略。

有了模型训练所需的总算力，除以每个 GPU 的理论算力，再除以 GPU 的有效算力利用比例，就得到了所需的 GPU-hours。

如果训练一个通用大语言模型的基座需要半年或几个月的时间，同时还占用非常多的机器资源，这就使得大模型的训练非常“贵”，使得大模型的训练成为了个别大企业的专用。

例如，如果你有100个实验想试试，而模型训练需要半年，那你只能在其中选择优先级高进行实验。

这也就是为什么现在大家的大语言模型基座都是追踪最新开源的模型。因为自己搞通用大模型基座，一方面搞半天成本上耗不起，另一方面可能还没啥效果。

2. 大模型训练需要更多的显存内存资源。

深度学习训练需要的内存包括模型参数、反向传播的梯度、优化器所用的内存、正向传播的中间状态（activation），显存占用 = 模型参数大小 + Batch Size * 优化器参数与中间变量的大小。

• 优化器所用的内存的计算其实也很简单，如果用最经典的 Adam 优化器，它需要用 32 位浮点来计算。即使我们使用mixed-precision进行计算，每个参数需要也要存 4 字节的 32 位版本（正向传播时用 16 位版本，优化时用 32 位版本），还需要存 4 字节的 momentum 和 4 字节的 variance，一共 12 字节。如果是用类似 SGD 的优化器，可以不存 variance，只需要 8 字节。
• 正向传播的中间状态（activation）是反向传播时计算梯度必需的，而且跟 batch size 成正比。Batch size 越大，每次读取模型参数内存能做的计算就越多，这样对 GPU 内存带宽的压力就越小。划重点：正向传播的中间状态数量是跟 batch size 成正比的。

当然也有节省内存资源的办法，例如算力换内存，时间换内存等。

算力换内存的把戏，就是不要存储那么多梯度和每一层的正向传播的中间状态，而是在计算到某一层的时候再临时从头开始重算正向传播的中间状态。

如果每一层都这么干，那么就只要 2 个字节来存这一层的梯度，但是计算中间状态的算力开销会很大。

因此实际中一般是把整个 Transformer 分成若干组，一组有若干层，只保存每组第一层的中间状态，后面的层就从该组第一层开始重新计算，这样就平衡了算力和内存的开销。

时间换内存的把戏，按顺序执行Mini-Batch数据的前向计算梯度，同时对梯度进行累积，累积的结果在最后一个Mini-Batch计算后求平均更新模型变量。

此外，还可以多级缓存，GPU 内存放不下可以换出到 CPU 内存。

例如，对于 LLaMA-2 70B 模型，模型参数需要 140 GB，反向传播的梯度需要 140 GB，优化器的状态（如果用 Adam）需要 840 GB。

3. 对数据的数量和数据的质量要求极高

对于海量数据样本来说，并不是都喂进去就效果好，哪些数据有价值，哪些没价值。由于数据量的增加，分辨数据的价值也带来很大的困难。

此外，大量的数据存储在哪里，一般可以存储在HDFS或S3。但怎么保证存取能不把机器塞满且能快速调取，对于相同数据的模型多次训练，是否可以通过cache来加速模型训练的时间。

4. Transformer的“不可能三角”特性

图片

Transformer在训练并行性、推理效率和竞争性能之间很难取得平衡被称为“不可能三角”。

Transformers 由于其固有的每步 O(N) 复杂度和内存限制的键值缓存，在推理过程中表现出次优效率。这种低效率使它们的实际部署变得复杂，特别是对于长序列。

5. 深度学习框架上不是很友好

因为这个领域最近几年才开始热门，而之前的框架pytorch、tensorflow等是早就出现的，当时并没有针对大模型的分布式训练的需求场景做深入的抽象设计和优化。

所以这个领域需要通过大数据框架到深度学习框架的端到端打通，形成这样的一套新的编程范式和对应的计算框架来解决掉。

下面我们以CV、NLP场景和搜推广场景进行分别详细说明。

CV和NLP场景：

对CV和NLP场景来说，其特点主要有：

• 模型一般复杂，单机性能要求高。业界主要使用高性能的GPU进行计算，并采用All-reduce的通信拓扑进行参数的同步更新。
• 模型大（Dens.NET部分），比如NLP领域，GPT-3这样的模型高达1750亿参数，显存占用高达2.8 TB，单机内存无法容纳。

当面对GPT-3这种Dense部分大的模型，Allreduce 单卡内存无法容纳，我们需要采用模型并行(model parallelism)的方式将计算图划分到不同的设备上构建有向无环图(DAG)进行分布式训练，其中Gpipe, Megatron, Oneflow和Whale都提出模型并行的相关解决方案。

相比于数据并行每个worker只有一部分数据，模型并行下每个node使用所有数据。

下面我们简单说明几种模型并行的方法：

• Tensor Parallelism，主要是将一层Layer中的矩阵计算分别拆分到不同的机器上进行运算，比如1D Megatron。
• Pipeline Parallelism，会将模型的layers拆分到不同的机器上，则一次forward就需要跨过不同的机器串行地进行运算，而流行并行通过将batch size切分为更小的mirco batch，减少数据依赖，从而将整个计算过程异步起来，最大化资源利用率。

CTR推广搜场景：

对于CTR大模型场景来说，其具有模型小，词表大的特点。

1. 模型中的Dense部分，一般很小，往往一台机器的内存就可以容纳。但是其特征量级可能高达成百上千亿，造成Sparse部分或者Embedding table高达TB级别，使得单机无法容纳。
2. 一个Batch的embedding lookup量级大，造成查询耗时大。由于特征数量多，一个Batch可能包含几十万个ID类特征，TF原生的embedding lookup查询耗时大，造成训练和inference性能低。
3. 数据具有大规模稀疏的特点。不同于CV和NLP场景，数据是稠密的图像和文本，搜广推的数据非常稀疏的，第一这来源于很多数据无法对所有用户和场景有效采集到，第二是因为建模使用的特征量级大造成的高维稀疏性，这会影响了数据的存储格式和计算效率。

因此，解决CTR大模型的这种稠密参数较大的模型，关键是将Sparse参数由单机存储改造为分布式存储，并主要通过数据并行提高吞吐。

下面我们说明下对训练框架优化点。核心的两点，一个在于分布式通信拓扑的设计，还有一个在于Embedding Lookup的性能优化。

1. 稀疏参数，借助参数服务器（Param Server），将 embedding 存储和更新负担转嫁到PS。稀疏参数 Partition 存放，每个 Worker 只有部分分片，梯度更新时进行 AlltoAll，想办法解决 稀疏tensor 的存储、通信成本。
2. 稠密参数，借助于allreduce，将稠密参数 Replication 存放，每个 Worker 都有副本，梯度更新时进行 allreduce。allreduce 和 alltoall 都会使用 nccl 进行同步通信，效率较高。hb 进行 alltoall 时，通信的是稀疏梯度，而不是完整的参数，通信量上和 ps 是一致的，但是通信效率会更高。

在实现上，可以通过替换TF原生算子进行Sparse参数的读取过程（核心算子是GatherV2算子）。

该算子的作用是从Embedding表中读取ID列表索引对应的Embedding数据并返回，本质上是一个Hash查询的过程；

通过替换该算子，并同时修改上下游节点的Input/Output，将从本地Embedding表中查询，改造为从分布式KV中查询。

总结

可以说，训练大模型，不仅需要充足的计算资源和数据，还需要深厚经验和技能，还需要一定的耐心和定力，就像“炼丹”一样。

每次炼丹师的出手都有着巨大时间和经济成本，如何在最小成本下找到最优解，就是我们一直在探索的。