您的位置:  首页 > 技术杂谈 > 正文

语言大模型的浮点运算分配

2024-02-28 19:00 https://my.oschina.net/oneflow/blog/11030216 OneFlow深度学习框架 次阅读 条评论

本文通过实证分析展示了实际LLM模型的FLOPS分配情况,并与理论分析进行对比。通过理论和实证相结合的方式,本文为理解和优化语言大模型的性能提供了有益见解。

 

作者Finbarr Timbers是一名机器学习研究员,曾就职于DeepMind。(以下内容由OneFlow编译发布,转载请联系授权。原文:https://www.artfintel.com/p/where-do-llms-spend-their-flops

 

作者 | Finbarr Timbers

OneFlow编译

翻译|宛子琳、杨婷

 

本文对LLM的性能进行了理论分析,然后通过详细分析一个实际的LLM,查看实证结果与理论之间的差异。首先是理论分析部分,我会借助Kipply的优质博文来填充细节。基本结论是:对于标准解码器模型,FLOPS(每秒浮点运算数)的分配如下(按每层计算):

 

  1. 6d^2 用于计算QKV(Query(查询)、Key(键)和Value(值))

  2. 2d^2 用于计算注意力输出矩阵,softmax(Q @ K.T) @ V

  3. 16d^2 用于运行前馈神经网络(FFN) 

 

总计24d^2 FLOPS。从百分比看,25%的时间用于计算QKV,约8%的时间用于计算注意力输出矩阵,约66%的时间用于运行FFN。

 

那么用于注意力机制的时间呢?众所周知,注意力机制方程为:

 

 

假设你正在使用KV缓存,Q(查询)、K(键)和V(值)都是d维向量(等价于(d,1)矩阵)。每个点积大约需要2d个flops(https://www.stat.cmu.edu/~ryantibs/convexopt-F18/scribes/Lecture_19.pdf),加上进行d次除法需要d个flops,总计约为5d个flops,四舍五入为零。

 

 

当d等于4096(在Llama7b中的取值),这仅为0.005%,几乎可以忽略不计。这似乎表明注意力机制不重要,但事实并非如此。我们之所以使用KV缓存(以及flash attention等)正是因为它们非常重要,可以将其类比于米尔顿·弗里德曼的恒温器(https://worthwhile.typepad.com/worthwhile_canadian_initi/2010/12/milton-friedmans-thermostat.html,感谢 @bradchattergoon):

 

假设一个房屋配备了一个运行良好的恒温器,那么我们能看到炉子燃烧的油量(M)与室外温度(V)之间存在强烈的负相关关系,同时炉子燃烧的油量(M)与室内温度(P)之间没有相关性,此外,室外温度(V)与室内温度(P)之间也没有相关性。

 

一位计量经济学家观察数据后得出结论:燃烧的油量对室内温度没有影响,室外温度对室内温度也没有影响。唯一的影响似乎是燃烧油量会降低室外温度。

 

观察相同的数据,第二位计量经济学家得出了完全相反的结论。他认为,室外温度(V)增加唯一的影响是会减少耗油量(M),而不会对室内温度(P)产生任何影响。

 

尽管两位计量经济学家得出了不同的结论,但他们一致认为燃烧油量(M)和室外温度(V)对室内温度(P)没有影响。基于这一共识,他们决定关闭炉子,不再浪费金钱购买燃油。

 

KV缓存需要O(T)的内存(其中T是我们希望生成的词元数),因此内存需求成本较高,这一点可以参考公司股票($NVDA)情况。

 

KV缓存有多大呢?对于每个词元,需要存储以下数量的字节(第一个2是因为我们假设使用bf16精度,因此每个参数占用2个字节;第二个2是因为需要同时存储K和V张量):

 

注意,根据假设,n_heads*d_head=d_model=d,因此字节数为4*层数*d。

 

GPT-3有96层,d_model为12288,每个词元需要4.72MB。因此,生成2048个词元需要5.6GB的内存。

 

尽管如此,要使用给定模型生成给定长度的序列,我们仍需使用与KV缓存相同的内存量,只是在每次前向传播结束时将其丢弃。因此,我们并不需要更多内存。从某种意义上说,KV缓存不占用内存(至少在Jax中是如此,除了一些繁琐的bookkeeping工作)。

 

对于一些新兴架构(例如Mistral 7B)又有何不同呢?Mistral 7b使用了分组查询注意力(Llama2也使用了类似的注意力机制,就好像这两个模型的作者存在某种联系。)和滑动窗口注意力。

 

在分组查询注意力中,我们可以在多头之间共享一个KV投影(MQA),具体而言,可以是所有注意力头之间共享一个KV投影(MQA,https://arxiv.org/abs/1911.02150),或者将其分成多个组(GQA,https://arxiv.org/abs/2305.13245v3)。这两种方法都等同于具有较小d_model的标准多头注意力(MHA)。在之前的KV缓存计算中,我们假设注意力头的数量乘以头的维度等于模型维度,但是在MQA/GQA中,我们放宽了这一假设。KV缓存公式如下:

 

 

可以转换为:

 

 

其中,注意力头的数量乘以头的维度就是模型的有效维度。因此,可以看到,随着KV头数量的减少,KV缓存大小呈线性减小( 这也是GQA/MQA方法背后的关键动机之一)。

 

Llama{1,2} 模型参数如下:

 

 

Llama 2中,每个词元所需的KV缓存如下:

 

在没有分组查询注意力(GQA)的情况下,34B模型需要的KV缓存内存是原来的5倍,而70B模型需要的KV缓存内存是原来的8倍。

 

Llama/Mistral的另一个改进是滑动窗口注意力,它保证我们可以将KV缓存限制在窗口大小,对于Llama7B来说,窗口大小为4096。

 

1

性能驱动的架构变化

 

如前所述,LLM每层使用了24d^2个flops。增加的层数将线性扩展flops和参数数量,增加模型宽度会二次方扩展模型大小。需要注意的是,这是因为参数的数量与d_model的平方成正比,因为我们的大多数层是从一个d_model输入向量转变为一个d_model的输出向量,所以权重矩阵的尺寸为(d_model, d_model)。换句话说,计算量与参数的数量呈正比,增加d_model会使参数数量呈二次方增加。模型深度增加一倍会使参数数量翻倍,但模型宽度增加一倍会使参数数量增加四倍。

 

尽管如此,更宽的模型的优势是能够更好地并行化。要计算第N层,必须首先计算前面的N-1层。这在高效并行化上十分困难,尤其是在训练期间,而通过张量并行化方法,跨多个GPU拆分单个层要容易得多。如果你主要关心时延,那么选择更宽的模型可能更合适。

 

2

实证分析

 

我使用Colab进行了这项分析。(https://colab.research.google.com/drive/1TH6AKsICZqlFoW1ph8h3wsF7q7qVMF8T?usp=sharing

以下是单次前向传播的高级概要(我的网站上有交互式概要:https://finbarr.ca/static/llama-profile.svg):

 

 

可以看到,本次运行的总时间中有4.88%用于单次前向传播。在前向传播中,有1.98%的时间用于注意力机制,有2.58%的时间用于多层感知机(MLP)。在前向传播的总时间中,有40%的时间用于注意力层,53%用于MLP。在注意力层内部,时间分配在4个不同的线性层上,其中有2个线性层花费的时间大致相同(linear_1、linear_2),一个花费的时间多50%(linear_3),另一个则是前两者的两倍(linear_0)。我猜测linear_0正在计算查询嵌入,而linear_1/2正在计算键和值嵌入。请注意,由于KV头的数量较少,计算速度要快得多!GQA(Query-aware Attention)带来了明显的差异,即便所使用的注意力机制(xformers.ops.memory_efficient_attention)要求QKV嵌入被广播到相同的大小。

 

理论分析预测,2/3的时间将用于计算FFN,1/3将用于计算注意力机制。这基本符合我们上面所看到的情况。我们花在计算注意力机制上的时间略多于MLP,但我怀疑这是因为MLP正在为Torch执行一个经良好优化的路径。

 

3

性能变化

 

接下来,我对Llama2进行了一系列实验,涉及模型宽度和深度的调整。以下是实验结果:

 

 

结果非常有趣。可以看到,隐藏维度为1024和1536的两个模型的速度基本没有变化(1.10秒vs1.11秒),隐藏维度为1024和2048的模型只发生了轻微变化(1.15秒vs1.10秒)。然而,当我们比较隐藏维度为2048(1.15秒)、3072(1.41秒)和4096(1.82秒)的模型时,可以看到速度类似于线性扩展!

 

对此,我的看法是,调度kernel和实际执行矩阵乘法中存在较大的开销。这是在T4上运行的,尽管按现在的标准来看有些过时,但仍具有65 TFLOPS的bf16计算能力。如果我们将两个1024x1024的矩阵相乘,这就需要1G FLOP的计算能力,因此,理论上,我们可以每秒乘以 65000个1024x1024的矩阵。实际上,我们只能得到其60-80%的性能,但仍然是每秒40000次矩阵乘法。如今的GPU拥有大量核心,T4有2560个CUDA核心,每个核心的运行频率在585到1590 MHz之间。因此,任何能够并行化的任务都表现良好,但是那些需要顺序计算的任务则不会得到优化。我认为,这就是图中所看到的的情况——没有足够的并行性来充分利用GPU。

 

Transformer的深度使性能与预期完全一致:推理时间与深度呈线性增长。最深的模型可能存在一些噪声,但它的性能表现相当稳定。

 

 

接下来,我计算了生成更多词元所需的成本(对每个词元数量进行了10次运行,以减少噪音):

 

 

正如预期的那样,完全呈线性增长,因为Llama2使用了KV缓存。如果我们查看保留的内存,就会看到KV缓存与预期的一致(在某种程度上):

 

 

可以看到,模型每增加20个词元,内存占用就会增加约2.1MB。由于该模型的 d_model 为1024,有8个隐藏层,因此需要4 * num_layers * d_model字节的内存,即4*8*1024字节=每词元32KB的内存。理论上我们只需要640KB的内存。目前还不清楚额外的3倍开销是从哪里产生的。我怀疑是因为执行还不够高效。

 

 

【语言大模型推理最高加速11倍】SiliconLLM是由硅基流动开发的高效、易用、可扩展的LLM推理加速引擎,旨在为用户提供开箱即用的推理加速能力,显著降低大模型部署成本,加速生成式AI产品落地。(技术合作、交流请添加微信:SiliconFlow01)

 

SiliconLLM的吞吐最高提升近4倍,时延最高降低近4

 

数据中心+PCIe:SiliconLLM的吞吐最高提升近5倍;消费卡场景:SiliconLLM的吞吐最高提升近3

 

System Prompt场景:SiliconLLM的吞吐最高提升11倍;MoE模型:推理 SiliconLLM的吞吐最高提升近10

 

其他人都在看

 

试用OneDiff: github.com/siliconflow/onediff

 

 

本文分享自微信公众号 - OneFlow(OneFlowTechnology)。
如有侵权,请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”,欢迎正在阅读的你也加入,一起分享。

展开阅读全文
  • 0
    感动
  • 0
    路过
  • 0
    高兴
  • 0
    难过
  • 0
    搞笑
  • 0
    无聊
  • 0
    愤怒
  • 0
    同情
热度排行
友情链接