本文通过实证分析展示了实际LLM模型的FLOPS分配情况，并与理论分析进行对比。通过理论和实证相结合的方式，本文为理解和优化语言大模型的性能提供了有益见解。 作者Finbarr Timbers是一名机器学习研究员，曾就职于DeepMind。以下内容由OneFlow编译发布，转载请联系授权。原文https://www.artfintel.com/p/where-do-llms-spend-their-flops 作者 |Finbarr Timbers OneFlow编译 翻译宛子琳、杨婷 本文对LLM的性能进行了理论分析，然后通过详细分析一个实际的LLM，查看实证结果与理论之间的差异。首先是理论分析部分，我会借助Kipply的优质博文来填充细节。基本结论是对于标准解码器模型，FLOPS每秒浮点运算数的分配如下按每层计算 - 6d^2 用于计算QKVQuery查询、Key键和Value值 - 2d^2 用于计算注意力输出矩阵，softmax(Q @ K.T) @ V - 16d^2 用于运行前馈神经网络FFN 总计24d^2 FLOPS。从百分比看，25%的时间用于计算QKV，约8%的时间用于计算注意力输出矩阵，约66%的时间用于运行FFN。 那么用于注意力机制的时间呢？众所周知，注意力机制方程为 假设你正在使用KV缓存，Q查询、K键和V值都是d维向量等价于d，1矩阵。每个点积大约需要2d个flopshttps://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头的数量较少，计算速度要快得多！GQAQuery-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秒。然而，当我们比较隐藏维度为20481.15秒、30721.41秒和40961.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倍 数据中心+PCIeSiliconLLM的吞吐最高提升近5倍消费卡场景SiliconLLM的吞吐最高提升近3倍 System Prompt场景SiliconLLM的吞吐最高提升11倍MoE模型推理 SiliconLLM的吞吐最高提升近10倍 其他人都在看 - 800+页免费“大模型”电子书 - 语言大模型的推理技巧 - 语言大模型的推理演算 - 语言大模型推理加速指南 - 语言大模型推理性能工程最佳实践 - 迈向100倍加速全栈Transformer推理优化 - Mistral AI:LLM推理的吞吐、时延及成本空间 试用OneDiff:github.com/siliconflow/onediff 本文分享自微信公众号 - OneFlowOneFlowTechnology。 如有侵权，请联系 support@oschina.cn 删除。 本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。