• 周二. 6 月 18th, 2024

GPT-4 炼丹指南MoE参数量训练成本和推理的秘密

这个图表假设无法融合每个操作是低效率的,注意力机制所需的内存带宽以及硬件开销与参数读取相当。在现实中 ,即使使用 NVIDIA FasterTransformer 这类“优化的”库,整体开销会更大。

上图展示了在足够高的吞吐量情况下,服务于单个用户 LLM 所需的内存带宽。从这张图可以看出:

• 即便是 8 倍于 H100 的带宽也无法实现以每秒 33.33 个 token 的速度为 1 万亿参数规模的稠密模型提供服务;

• 此外,8x H100 的 FLOPS 利用率在每秒 20 个 token 时仍低于 5%,这会导致极高的推理成本。

实际上,对于今天的 8-way 张量(tensor)并行的 H100 系统来说,推理约束为约 3000 亿前馈参数。

然而,OpenAI 正在用 A100 以及大于 1 万亿参数的模型实现人类的阅读速度,以每 1000 个 token 0.06 美元的低价广泛提供,之所以能够实现这一点正是因为它的稀疏架构。

接下来,我们会讨论 GPT-4 的模型架构、训练和推理的 infra、参数数量、训练数据集构成、token 数量、层数、并行策略、多模态视觉编码器等一系列不同工程设计背后的考量、实现技术,以及 OpenAI 是如何解决大模型推理过程中的瓶颈的。

02. 模型结构

GPT-4 的规模是 GPT-3 的 10 倍以上,我们估计它有约 1.8 万亿个参数,这些参数分布在 120 个 transformer 层上,作为对比,GPT-3 的参数为大约 1750 亿个。(拾象注:GPT-3 仅有 12 个 transformer 层,层数是 GPT-4 的 1/10。)

为了控制成本,OpenAI 选择使用 MoE 模型。OpenAI 在模型中使用了 16 个 MLP.2 类型的专家,每个专家的参数大约为 1110 亿个。每次前向传递中会调用其中的两个专家模型。

• 专家混合模型(Mixture-of-Experts,MoE):MoE 模型是一种深度学习架构,该架构,通常由多个专家(Experts)组成,每个专家负责处理输入数据的不同方面,并拥有自己的参数集(也有一些参数,例如 embedding,可以被所有专家共享,即共享参数)。在模型的推理过程中,根据输入数据的不同特征,模型会将输入路由到不同的专家,每个专家根据其参数集处理对应分配到的输入后完成输出,最终输出则是各个专家输出的集成。

• MLP:即多层感知机(Multi-Layer Perceptron),MLP 是一种人工神经网络,包含了多个隐藏层,MoE 模型中通常有多个独立的 MLP 专家。

有很多文献中都在讨论如何将每个待处理 token 路由(assign)至某个专家模型的算法,但据说 OpenAI 所采用的这一套算法却相当简单,至少 GPT-4 是这样的。

此外,大约有 550 亿个共享参数被用于注意力机制。

每次前向推理(生成一个 token)仅利用了约 2800 亿个参数和 560 TFLOP,相比之下,如果纯粹使用稠密模型每次前向推理所需的约 18000 亿个参数和 3700 TFLOP 。

03. 数据集

GPT-4 是在约 13 万亿个 tokens 上训练的,考虑到 CommonCrawl RefinedWeb 包含有约 5 万亿个高质量的 tokens,13 万亿这个数字还算合理。作为参考,Deepmind 的 Chinchilla 和 Google 的 PaLM 模型分别是用约 1.4 万亿个 token 和约 0.78 万亿个 token 训练的,PaLM2 据说也是在约 5 万亿个 token 上训练的。

CommonCrawl Refinedweb:

CommonCrawl 是一项非营利项目,旨在构建和维护一个开放、可访问的互联网数据集,可使用网络爬虫技术定期扫描互联网上的网页,并将抓取到的网页和相关的元数据进行整理和存档。CommonCrawl RefinedWeb 是CommonCrawl 通过算法和人类审查后,从原始收集的数据中筛选出高质量文本形成的资料库。

OpenAI 训练 GPT-4 所使用的这个数据集不是 13 万亿个 unique tokens。相反,由于缺乏高质量的 token,该数据集包含了多个 epoch。基于文本的数据有 2 个 epoch,基于代码的数据有 4 个 epoch。(拾象注:这里指部分高质量的文本和代码给模型多次学习过。)这远远没有实现 Chinchilla-optimal(需要在双倍的 token 数上训练模型),这也说明网络上的易获得 token 不足。网络上实际存在的高质量文本 tokens 应该是现如今可利用的 1000 倍,音频和视频类的 token 则要更多,但采集这些 tokens 并不能简单通过网络抓取来实现。比较遗憾的是,我们还没有找到太多关于 OpenAI 的 RLHF 到数据信息。

一个 epoch 指的是将整个训练集(training set)中的所有样本都用于训练模型一次的过程。具体来说,一个 epoch 包括多个训练步骤(training steps),每个训练步骤都是将一个小批量(batch)的样本输入模型中进行训练,并更新模型的参数,以最小化损失函数(loss function)。

如果 epoch 过小,那么模型可能无法充分利用训练集中的信息,导致欠拟合(underfitting),即模型无法很好地拟合训练数据,导致在测试集上的表现不佳。相反的,如果一个 epoch 过大,那么模型可能会过拟合(overfitting),过多地学习训练集中的噪声和局部特征,而忽略了全局特征。

在预训练阶段,上下文长度(seqlen)是 8k。GPT-4 的 32k 上下文版本是在预训练后对 8k 进行微调的基础上实现的。

Batch size 在集群上经过若干天的逐步提升,但到最后,OpenAI 使用了高达 6000 万的 batch size。当然,由于不是每个参数都看到所有参数,这只是每个专家的 batch size 大小为 750 万。

Batch size 指的是每一次迭代 (iteration) 或者前向传递 (forward pass) 的训练样本数量。模型训练中会把数据分批次来训练,Batch size 就表示每一批次中样本的数量。分批训练的好处在于可以规避内存限制、节省用于重复计算中间结果的计算资源。

Batch Size 的大小对模型的训练效果和速度有很大的影响。Batch Size 越大,每次更新参数的计算量就越大,但是训练过程会更加稳定,因为每个 Batch 中的样本可以平均化噪声和不确定性。另一方面,如果 Batch Size 过小,训练过程可能会变得不稳定,并且需要更多的训练步骤才能收敛到最优解。此外,Batch Size 的大小也会受到硬件资源的限制。因此,在实际应用中,选择适当的 Batch Size 是非常重要的。

04. 并行策略

在所有 A100 GPU 上进行并行处理非常重要。

OpenAI 使用了 8 路(8-way)规模的张量并行(Tensor Parallelism),之所以是 8 路(8-way)是因为这是 NVLink 的极限。此外,我们还听说 OpenAI 正在使用 15 路(15-way)的流水线并行策略。从理论上讲,考虑数据通信与计算时间,15路(15-way)显得数量过多,但如果他们受限于内存容量,这一点也很合理。

大模型训练中存在几种经典的分布式并行范式,分别为流水线并行(Pipeline Parallelism),数据并行(Data Parallelism)和张量并行(Tensor Parallesim)。微软开源的分布式训练框架 FastSpeed 就融合了这三种并行范式。

如果只是使用流水线并行和张量并行,每个 GPU 上的参数在 FP16 下大约需要 30GB,而一旦考虑到 KV 缓存和 KV 开销,如果 OpenAI 所采用的大部分 GPU 是 40GB 的 A100 的话,这个架构从理论上也是合理的。OpenAI 可能使用了 ZeRo stage 1、块级 FSDP 或混合共享数据并行。

• KV 开销(KV overhead):是指在 KV 存储系统中由于额外的开销而产生的负担。这些开销可能包括存储和管理键值对的元数据、索引结构、数据复制和同步、网络通信等。KV overhead 的增加可能导致性能下降、存储需求增加以及系统复杂性的增加。

• ZeRo Stage 1:ZeRO(Zero Redundancy Optimizer)是指原来每张卡存储完整的优化器状态。若每张卡只存储一部分优化器状态,所有卡的优化器状态共同构成完整的状态,即 Pos (Partition Optimizer States),被称为 ZeRO-stage1。

• Block-level FSDP:是指基于块的全精度动态量化(Full Precision Dynamic Quantization)技术。可以在训练和推理过程中保留更高的模型精度,使模型 inference 的时候成本更低。

之所以不使用全模型 FSDP 的原因可能在于过高的通信成本。虽然 OpenAI 在大多数节点之间有高速网络,但可能不是所有的节点都这样,我们认为至少存在一些集群的连接带宽要比其他集群低得多。

我们暂时还不清楚 OpenAI 是如何在如此高的管道并行性下避免出现巨大的 bubble。很可能他们只是承担了这个成本。

Bubble: 在每个批次(batch)中由于高度的管道并行性而导致的延迟或等待时间。它表示在进行高度并行的计算过程中,由于不同部分的计算速度不同,可能会导致某些部分需要等待其他部分完成计算,从而产生延迟或空闲时间。这种情况下,bubble 指的是这些空闲或等待的时间间隔。这句话意味着他们可能只是接受了在计算过程中存在一些空闲时间或延迟。

Google 在 PaLM 推理论文中也提到了对上述三个问题的处理。值得注意的是,这是针对 PaLM 类的稠密模型,而不是 GPT4 这类的稀疏模型。

如果一个应用程序需要尽可能低的延迟,我们需要更多的芯片,并以尽可能多的方式对模型进行分割,这样才能有经济效益。较小的批处理量可以实现较低的延迟,但较小的批处理量也会导致较差的 MFU [利用率],导致每个 token 的总成本(以芯片秒数或美元计算)较高。

如果一个应用需要离线推理,并且延迟不是问题,那么主要目标是最大限度地提高每个芯片的吞吐量(即最小化每个 token 的总成本)。增加批处理量是最有效的,因为较大的批处理量通常会带来更好的 MFU [利用率],但是某些对小批处理量无效的分区策略会随着批次处理量的增长而变得有效。

更多的芯片和更大的批次量反而是更便宜的,因为它们提高了利用率,但这也引入了第三个变量,即联网时间(Networking Time)。把模型部署在多个芯片上的方法可以有效解决延迟,但要对利用率做出牺牲。

存储时间的权重加载部分和非注意计算时间都与模型大小成正比,与芯片数量成反比。对于一个给定的分区布局,芯片与芯片之间的通信所需时间随着所用芯片数量的增加而减少得不那么快(或根本不减少), 所以随着芯片数量的增加,它成为一个越来越重要的瓶颈。

我们注意到,随着批处理量和规模的增长,KV 缓存的内存需求量会激增。

如果一个应用程序需要生成具有长注意力上下文(long attention contexts)的文本,就会大大增加推理时间。对于一个具有多头注意力 500B 以上的模型,注意力的 KV 缓存会变得很大:对于批量大小为 512、上下文长度为 2048 的模型,KV 缓存的总量为 3TB, 是模型参数大小的 3 倍。片上存储器(on-chip memory)需要从片外存储器(off-chip memory)中加载这个 KV 缓存,每产生一个 token 就加载一次,在此期间芯片的计算核心基本上是空闲的。

较长的序列长度对内存带宽和内存容量来说特别麻烦。OpenAI 16k 上下文的 GPT-3.5 turbo 和 32k 上下文的 GPT-4 很贵的原因是由于内存的限制,它们不能采用更大的批次。

较小的批次导致较低的硬件利用率。此外,随着序列长度的增加,KV 缓存会膨胀。KV 缓存不能在用户之间共享,所以需要单独的内存读取,这进一步降低了内存带宽。关于 MQA 的更多信息,请见下文。

08. 推理的 Infra 与成本

Infra

MoE 的架构则让 GPT-4 的推理在延迟、吞吐量和利用率上都面临挑战。因为每个 token 的前向传递可以路由到不同的专家模型,这种情况下要尽可能实现低延迟、高吞吐量和高利用率就十分困难,尤其是在高 batch-size 的情况下。

OpenAI 的 GPT-4 架构中包含了 16 个专家模型,每个前向通道有 2 个路由(router)。这意味着在batch-size 为 8 的情况下,每个专家的参数读取可能只占批量大小的“1”。更严重的是,这还会导致某个专家的 batch-size 为 8,而其他专家的 batch-size 可能只有 4、1 或 0。

此外,每一次 token 生成,路由算法都会将前向传递路由到不同的方向,这会导致 token 到 token 之间的延迟和专家批处理量的显著变化。也就是说,处理不同的 token 时,不同的专家可能会被分配到不同的任务,计算负载和批量大小都可能因此发生变化。

推理 infra 是 OpenAI 在 MoE 的设计上选择较少数量专家的主要考量之一。如果他们使用更多的专家,内存带宽将成为推理面临的一个更大瓶颈。OpenAI 在自己的推理集群上常常能达到 4k 以上的 batch-size,这意味着即使在专家之间实现了最佳负载平衡,每个专家的批处理量也只能达到大约 500 个。这需要非常大的使用量才能实现。

我们的理解是,OpenAI 在由 128 个 GPU 组成的集群上运行推理,并且在不同的数据中心和地理区域拥有多个这样的集群。推理是以 8 路张量(tensor)并行和 16 路流水线并行的方式进行的。每个节点使用 8 个 GPU,每个 GPU 只有大约 130B 的参数,或者在 FP16 下每个 GPU 不到 30GB,在 FP8/int8 下不到 15GB。只要所有批次的 KV 缓存大小不会膨胀得太大,这就可以在 40GB 的 A100 上运行推理。

FP16、FP8 和 int8 是不同的数值精度(precision)表示方式,常用于深度学习中的计算过程中,以减少内存和计算资源的使用,从而提高模型训练和推理的效率。

FP16 、FP8、int8 分别指的是 16 位浮点数、8 位浮点数和 8 位整数,它们的精度相对于 32 位的单精度浮点数(FP32)更低,但可以大大减少内存和计算资源的使用,从而加速深度学习中的模型训练和推理。例如,使用 FP16 可以在不损失太多精度的情况下,将计算时间减少一半以上,而使用 int8 可以在不损失太多精度的情况下,将计算时间减少约 4 倍。

需要注意的是,使用低精度计算可能会对模型精度产生一定的影响,因此需要在精度和效率之间进行权衡,并根据具体的任务需求选择最合适的精度表示方式。

为了避免网络通信过于不规则,同时避免在每个 token 生成之间重新计算 KV 缓存的成本过高,包含各种专家的各个层并没有在不同的节点上进行分解,以便共享 KV 缓存。

未来所有的 MoE 模型扩展和条件路由的最大困难。在于如何处理 KV 缓存周围路由层数最大为 120 的限制。

在 MoE 模型中,每个分支的路由层数不能超过 120 层,否则将无法有效地处理 KV 缓存。这是因为在模型的推理过程中,每个分支都需要计算 KV 缓存,这会导致计算成本的增加。

解决这个问题的一个简单方案是,基于 120 的层数限制,在 15 个不同的节点中放置一个横跨路由。这样就可以将计算负载平均分布在不同的节点上,从而提高模型的效率和性能。然而,由于第一个节点需要进行数据加载和嵌入,因此如何在推理集群的头部节点上放置更少的层是很重要的。

另外,在对输入数据进行编码和解码的过程中,可能会存在一些关于推理解码的噪音,这个问题我们将在后面进一步讨论。较为关键的问题在于确定这些噪音是否应该被相信。这也可以解释为什么头部节点上包含较少的层是有意义的。

推理成本

相较于 175B 参数的 Davinchi 模型,GPT-4 有 1.6 倍的前馈参数,但成本却是 Davinchi 的 3 倍。这主要是因为 GPT-4 需要更大的集群,以及实现利用率更低。

我们猜测使用 128 个 A100s 进行 GPT-4 8k 上下文长度(seqlen)的推理,每 1k tokens 的成本约为 0.0049 美元。而使用 128 个 H100s 进行 GPT-4 8k 上下文的推理,每 1k tokens 的成本约为 0.0021 美元。(拾象注:当前 GPT-4-8k 的价格是 0.03/1k input tokens, 0.06/1k output tokens,当前 OpenAI 对推理芯片使用不会如作者推测的一样奢侈,该测算可以作为未来降价的 lower bound。)需要特别注意的是,这些成本的计算是在利用率和 batch-size 都很高的情况下得出的。

考虑到 OpenAI 集群的利用率有时可能会非常低,我们的这一假设也有可能是错的。

我们假设,OpenAI 在低谷时期会关闭集群,并重新利用将这些节点重新分配给其他任务,例如恢复小型测试模型的检查点训练,或尝试各种新技术。这样的做法有助于保持低推理成本,否则 OpenAI 的利用率可能会更低,这意味着成本估算会是现在的 2 倍以上。

恢复小型测试模型的检查点训练,通常是指在训练深度学习模型时,使用之前保存的模型参数检查点,重新开始训练一个较小的模型(例如,一个仅使用模型的一部分参数的子集),以便在较短的时间内快速测试新的模型结构或算法。这种方法可以帮助研究人员快速迭代模型设计,并寻找最佳的模型结构和参数。

09. 多查询注意力机制

多查询注意力机制(Multi-Query Attention)的使用已经相当普遍,但我们想强调的是 OpenAI 也是这么做的。总的来说就是只需要 1 个注意力头(head),内存容量就能够显著减少以用于 KV 缓存。即便如此,32k 上下文的 GPT-4 肯定无法在 40GB 的 A100 上运行, 而 8k 的最大批量大小已经达到上限。如果没有 MQA,8k 的最大批处理量将被大大限制,经济效益大打折扣。

• 多查询注意力机制 (Multi-Query Attention,MQA):Fast Transformer Decoding: One Write-Head is All You Need这篇论文于 2019 年提出了 MQA 这一概念,后来成为了自然语言处理中经常使用的注意力机制。

传统的注意力机制中,将一个查询(query)和一组键值对进行匹配,以获得对每个键的加权表示。而在多查询注意力中,有多个查询,并且每个查询都会与键值对进行匹配,以获得对每个键的不同加权表示。这个过程可以看作是在多个不同的“视角”下对输入进行编码,从而获得更全面和准确的表示。

• 注意力头(Head):在深度学习模型中,通常包含多个层(layers)和一个头(head),head 用于将模型的输出映射到期望的输出空间。头部层通常是为了满足特定的任务而添加到模型中的,例如在自然语言处理任务中,head 通常用于将模型的输出转换为文本,进行文本分类等任务。在深度学习模型中,head 通常是紧跟在最后一个层的后面,用于将最后一层的输出转换成期望的输出形式。

10. 连续批处理

为了允许一定程度的最大延迟和优化推理成本,OpenAI 同时使用了可变批次大小和连续批次的技术。这种方法可以在不牺牲模型性能的情况下提高计算资源的利用率,并在模型的推理过程中实现更低的延迟和更高的吞吐量。如果不了解连续批处理这一概念的话,AnyScale 官方发布的 How continuous batching enables 23x throughput in LLM inference while reducing p50 latency 这篇文章非常值得一读。(拾象注:Anyscale 开发的分布式计算框架 Ray 是 OpenAI 在模型的 infra pipeline 中使用的,拾象之前发布过对这家公司的研究。)

连续批处理(Continuous batching):一种在深度学习训练过程中使用的技术,旨在提高训练效率和通过硬件的资源利用率。传统的批处理方法是将一定数量的训练数据一次性加载到内存中,然后对这些数据进行训练,这种方式可以提高训练效率,但也可能会浪费内存空间。

而连续批处理则是将训练数据分成若干个小批次,每次只加载一个小批次进行训练,训练完成后再加载下一个小批次,以此类推,直到完成整个训练数据集的训练过程。使用连续批处理技术可以在减少内存使用的同时提高训练效率,还可以提高模型的稳定性和泛化能力。

AI Agents大爆发:软件2.0雏形初现,OpenAI的下一步

拾象硅谷见闻系列:打破围绕开源LLM的6大迷思

ChatGPT Plugin:被高估的“App Store时刻”,软件和SaaS生态的重组开端

Stability AI:AI开源商业化试验田,Killer Model能成长为Killer App吗?

Inflection创始人:从DeepMind到Pi,AI智能体如何迎来寒武纪大爆发