大模型推理的极限在哪？一起聊一聊其瓶颈及极限理论值

AI生成未来 · 公众号 · · 2024-08-01 00:00

正文

点击下方 卡片 ，关注“ AI生成未来 ”

请加小助理 加入AIGC技术交流群

备注公司/学校+昵称+研究方向， 广告勿扰

本篇讲述大模型推理机制及其极限理论值分析。

大家在做大模型推理优化的时候，可能都会考虑一个核心问题： 推理的极限在哪里？

本文基于文章 [LLM inference speed of light] 及ArthurChiao的中文版文章做了详细解读，分析了大模型推理的速度瓶颈及量化评估方式，希望对小伙伴们理解大模型推理内部工作机制与推理优化有帮助。

下面是一个快捷目录。由于内容比较多，本篇主要介绍一到四，剩下的部分会在下一篇结合一些落地应用进行讲解。

一、常见的浮点运算单位介绍

二、推理机制

三、瓶颈分析

四、以Mistral-7B为例，计算极限推理延迟

五、推理理论极限值的作用

六、GQA带来的启发

一、常见的浮点运算单位介绍

这里简单列举一些比较常见的单位：

FLOPs ：floating point of operations的缩写，是浮点运算次数，可以用来衡量算法/模型复杂度。
一个GFLOPS（gigaFLOPS）= 每秒十亿（=10^9）次的浮点运算
一个TFLOPS（teraFLOPS） = 每秒一万亿（=10^12）次的浮点运算

二、推理机制

1. 语言模型Teacher-forcing的特点：逐token生成，无法并行

当LLM生成文本时，是teacher-forcing，也就是逐个token生成的，我们将其看作一个函数：

输入：一个token
输出：一组概率，每个概率对应词汇表中的一个token
推理过程 ：使用概率来指导抽样，产生（从词汇表中选择）下一个 token 作为最终输出。

此处简单介绍一下词汇表。

词汇表（vocabulary） ：通常由单词、单词片段、中文汉字等组成（这些都称为 token），一般会基于训练语料先基于bpe或者sentence piece等算法切分token生成。

文本生成过程就是不断重复以上过程。可以看出，在生成一个文本序列时，没有并行性的可能性。

2. 生成过程建模：矩阵乘法

广义上，当处理一个 token 时，模型执行两种类型的操作：

1）矩阵-向量乘法

一个大矩阵（例如 8192x8192）乘以一个向量，得到另一个向量。

2） attention 计算

在生成过程中，模型不仅可以看到当前 token 的状态，还可以看到序列中所有之前 token 的内部状态 —— 这些状态被存储在 KV-cache 的结构中，它本质上是 文本中每个之前位置的 key 向量和 value 向量的集合 。

attention 为当前 token 生成一个 query 向量，计算它与所有之前位置的 key 向量之间的点积，然后归一化得到的一组标量，并通过对所有之前的 value 向量进行加权求和来计算一个 value 向量，使用点积得到最终得分。

三、瓶颈分析

上述的矩阵-向量乘法和attention 计算两步计算有一个重要的共同特征：从矩阵或 KV-cache 读取的每个元素，只需要进行 非常少量的浮点运算 。

矩阵-向量乘法对每个矩阵元素执行一次 乘加运算 （2 FLOPs）；
attention 对每个 key 执行一次乘加，对每个 value 执行一次乘加。

1. 一些典型显卡的“算力-带宽”比

现代 CPU/GPU 的 ALU 操作（乘法、加法）内存 IO*速度要快得多 。例如：

AMD Ryzen 7950X： 67 GB/s 内存带宽和 2735 GFLOPS， Flop:byte = 40:1
NVIDIA GeForce RTX 4090： 1008 GB/s 显存带宽和 83 TFLOPS， Flop:byte = 82:1
NVIDIA H100 SXM： 3350 GB/s 内存带宽和 67 TFLOPS，对于矩阵乘法，tensor core 提供 ~494 TFLOPS 稠密算力， Flop:byte = 147:1 。

对于 FP16/FP8 等精度较低的浮点数，比率更夸张：

H100 TensorCore 对于 dense FP8 矩阵的理论吞吐量为 1979 TFLOPS， FLOP:byte = 590:1 。

在这些场景中，无论是否使用 TensorCore 或使用什么浮点格式，ALU 都非常充足。

2. 瓶颈：访存带宽

因此，transformer 这种 只需要对每个元素执行两次操作 的场景，必定受到访存带宽的限制。所以，基于下面几个因素，

1）模型参数配置

2）KV-cache 大小

3）访存带宽

四、以Mistral-7B为例，计算极限推理延迟

大模型推理的极限在哪？一起聊一聊其瓶颈及极限理论值

正文

请到「今天看啥」查看全文