LLM推理速度快速预估

最新推荐文章于 2024-09-13 19:41:31 发布

江湖人称麻花滕

最新推荐文章于 2024-09-13 19:41:31 发布

阅读量548

点赞数 9

文章标签：迁移学习人工智能机器学习语言模型自然语言处理

本文链接：https://blog.csdn.net/m0_59235699/article/details/142186210

版权

原文：www.jinghong-chen.net/estimate-vram-usage-in-llm-inference

通过几行计算，您就能估算出大语言模型（LLM）推理所需的首token时间（TTFT）、每输出token时间（TPOT）以及显存的需求。我将以Llama-7B为实例，向您展示具体操作方法。

LLM推理基础

LLM推理主要分为两个阶段：预填充和解码。在预填充阶段，模型处理输入上下文，并计算其embedding表示。而在解码阶段，模型则基于输入上下文和之前已生成的所有令牌，逐个生成新的令牌。

预填充的速度直接影响首token时间（TTFT），因为在处理完输入上下文之前，无法生成任何内容。解码速度则关系到每输出token时间（TPOT），即预填充后单位时间内生成的token数。TTFT和TPOT对于用户体验至关重要，也是衡量LLM推理性能的关键指标。

VRAM的使用主要由模型参数和键值缓存（KV-cache）决定。KV-cache存储了输入上下文和之前所有token的键值embedding，这些embedding只计算一次并存储在内存中，以便在解码阶段重复使用。VRAM的使用量直接影响到最大批次大小和序列长度，从而对LLM推理的吞吐量产生显著影响。

总而言之，VRAM使用量、TTFT和TPOT是描述LLM推理性能的三大要素。为了估算这些指标，我们需要确定加载的数据量和浮点运算次数（FLOPs）。接下来，让我们看看具体如何操作。

预估方法

GPU主要执行两项任务：加载数据和进行浮点运算（FLOPs）。要估算时间，我们需要估算LLM推理每个阶段所需的总内存和计算工作量，并将其除以GPU的处理速率。为此，您需要以下常量：

s：序列长度
b：批次大小
h：隐藏维度
L：Transformer层数
N：模型参数数量
GPU的FLOPs速率。对于A100 GPU，这是XX FLOPs/秒
GPU高带宽内存（HBM）速率。对于A100 GPU，这是XX TB/秒

我们假设使用16位精度（每参数2字节），则计算公式如下：

在这里插入图片描述

解释：2 * N 是处理每个token所需的大致计算量。考虑乘积Wx = w11*x1 + w12*x2+… ，其中W是权重矩阵。每个参数都参与一次乘法和一次加法。b*s表示需要处理的token总数。在解码阶段，我们逐个生成令牌，因此s=1 。

在这里插入图片描述

解释：为确保计算过程顺利进行，所有模型参数均需加载至GPU中。该转换基于16位精度的假设，将模型参数的数量换算为对应的字节量。

让我们一步步解析KV缓存表达式，从左到右进行。每个注意力头中，每个标记都配有一个缓存的键向量和一个缓存的值向量。将所有注意力头的参数聚合后，得到的维度就是模型的维度(*h）。一个标记中包含了（b * s）个这样的标记。最后，基于16位精度的假设，我们计算出所需的字节量来表示这些参数。

Llama-7B

让我们来估算一下Llama-7B推理模型的两个关键指标：首token时间（TTFT）和显存（VRAM）需求，然后看看我们的估算是否接近实际实验值。不过要注意，TPOT的解释很灵活，得根据具体应用场景来定，所以这次实验我们主要就盯着TTFT这个指标来估算。

在这里插入图片描述

预估与实际的VRAM使用量对比

我们的计算预测，如果将峰值VRAM使用量与批量大小进行绘图，斜率应接近每批量大小400MB。实验数据的线性拟合给出的斜率为[350MB]。我们的预估相当准确。正如所预测的那样，当批量大小设置为200时，报告了内存不足（OOM）错误。

预估与实际的首token（TTFT）对比

计算结果显示，每增加一个批量大小，TTFT将增加20.8毫秒。我们通过在Hugging Face的generate()调用中设置max_new_token=1来估算TTFT。在图2中，我们对实验得到的TTFT与批量大小进行了线性模型拟合。截距为109.84毫秒，这反映了推理过程中的固定开销。拟合得到的斜率为每批量大小增加19.68毫秒。再次验证，这与我们的预估值20.8非常接近！

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述