从计算角度解读LLM内部结构与推理过程_llm模型依赖的算子-优快云博客

本文链接：https://blog.youkuaiyun.com/2401_85375151/article/details/145280799

！！！纯手推计算纯干货长文，你一定能看懂，绝对受益匪浅。

相信通过前面的文章，大家对LLM的内部结构已经有了一个认知。

本文将继续深入LLM的内部结构，会从计算的角度分析每个算子的输入数据结构、输出数据结构，以及统计对应的计算量，以此对LLM内部结构有一个更加深刻的认识。

1. 先认识下什么是Tensor Shape
2. 词表Embedding计算细节
3. MHA内部计算细节
```
 3.1. Normalization
```
```
 3.2. Q/K/V Projection
```
```
 3.3. Q/K Matmul => Score
```
```
 3.4. Softmax(Score) => S
```
```
 3.5. SV Matmul
```
```
 3.6. O Projection
```
4. MHA总结
5. MLP内部计算细节
```
 5.1. Gate/Up Projection
```
```
 5.2. Down Projection
```
6. MLP总结

先认识下什么是Tensor Shape

张量（Tensor）是深度学习中用于表示数据的核心结构，它是一个多维数组，可以看作是标量、向量和矩阵的泛化，是深度学习框架如TensorFlow和PyTorch中的核心数据结构，用于表示和处理多维数据，是实现复杂算法的基石。如果你了解Tensor，这部分可以直接跳过。

• 数据处理：张量支持多种数学运算，如逐元素加法、矩阵乘法等，这些操作在数据预处理和特征提取中非常有用。
• 高性能计算：张量可以利用GPU和TPU等硬件加速计算，提高模型训练和推理的效率。

可以这么说，在机器学习里面，一切数据皆Tensor，在机器学习模型中，所有的数据都是以Tensor的格式流转的。

然后是Tensor的Shape属性，Tensor Shape是理解计算过程的重要属性。

定义：张量的形状（Shape）是一个整数数组，描述了张量在每个维度上的大小。

例子1：一维张量（向量）

import torch   # 创建一个一维张量（向量），包含5个元素   vector = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])   print(vector)   print("Shape of vector:", vector.shape)         输出：   tensor([1., 2., 3., 4., 5.])   Shape of vector: torch.Size([5])

例子2：三维张量

tensor([[[3, 7, 2, 5],            [1, 4, 6, 0],            [9, 8, 3, 2]],           [[4, 1, 7, 3],            [5, 0, 2, 8],            [6, 9, 4, 1]]])   Shape of tensor_3d: torch.Size([2, 3, 4])

好了，理解Tensor和Shape属性，后面就进入文章主题了，后面的所有内容都会围绕Shape属性来展开，以此来了解LLM推理的完整计算细节过程。

词表Embedding计算细节

首先是Embedding部分，在前文已经详细介绍过了，

第三篇：数据在大模型里是如何存在的？

这里简单回顾下，不记得的话，可以回看前面的文章。词表Embedding是将词汇表中的单词转换为稠密向量表示的过程，也即是将人类理解的分词词语转换成计算机理解的向量数据。

对应到每个模型，每一个大语言模型都有一个词表，这个词表表示了该LLM能够输出的不同词语的数量，比如下图，Llama-2-7b的词表大小是：128256。

在LLM推理过程中，LLM的第一步，就是将我们输入的文字做分词，并通过Embedding转成向量数据。

输入Tensor Shape：[B, S], [V, H]

输出Tensor Shape：[B, S, H]

Tensor计算量：主要是查表操作，计算量约等于0

其中：
B = Batch Size（批次大小）
S = Seq Len（输入长度）
V = Vocab Size（词表大小）
H = Hidden Size（隐藏层大小）

Embedding过程主要是查表操作，就是把每个字拿去和词表[V, H]去做匹配，拿到对应的Hidden Tensor。

其中Shape = [V, H]的数据就是预训练得到的，Shape = [B, S]是输入数据，一般来说B=1，但是在LLM后端处理并发的时候，可能把多个人的Prompt同时作为输入，这个时候B就会大于1。

MHA内部计算细节

接下来，就是重头戏了，没错，还是这个图！！在MHA部分，将会有几个算子需要重点关注的：

1. Normalization
2. Q/K/V Projection
3. Q/K Matmul ⇒ Score
4. Softmax(Score) ⇒ S
5. SV Matmul
6. O Projection

Normalization

Normalization计算的作用是将输入数据规范化到标准分布，以稳定神经网络的计算过程。

输入Tensor Shape：[B, S, H]
输出Tensor Shape：[B, S, H]
Tensor计算量：k*B*S*H，k是常数项，归一化操作不涉及矩阵乘法，所以整体运算量比较小。

Q/K/V Projection

Q/K/V Projection计算的作用是Attention，将输入Q分别和Q、K、V预训练的权重进行运算。其中QKV的Shape都是[B,S,H]，权重Shape是[H,H]。

输入Tensor Shape：[B, S, H], [H, H]，这里的[H, H]就是训练得到的权重
输出Tensor Shape：[B, S, H]
Tensor计算量：6BSHH = 3*2*B*S*H*H，其中3=Q/K/V，2=矩阵乘法的一次乘法+一次加法，B*S*H*H=[B, S, H]*[H, H]。

这里特别说明下，在本文中，如果输入的Tensor有两个，那表达的意思就是两个Tensor进行矩阵乘法，这种运算，要求第一个Tensor的最后一维和第二个Tensor的倒数第二维一致，得到的结果就是消掉这两维度的结果。

Q/K Matmul ⇒ Score

Q/K matmul是计算Attention分数的关键步骤，这里引入一个新的变量 a ，它表示多头注意力的头数，由H进行分组a得到的；同时K进行了一个转置，所以Q、K的Shape有以下变化：

Q: [B, S, H] ⇒ [B, a, S, H/a]   K: [B, S, H] ⇒(转置) [B, H, S] ⇒(多头) [B, a, H/a, S]

输入Tensor Shape：[B,a, S, H/a], [B, a, H/a, S]
输出Tensor Shape：[B, a, S, S]
Tensor计算量：2BSSH = 2*B*S*S*H = 2*B*a*S*H/a*S

同理，[B,a, S, H/a]的Q Shape和[B, a, H/a, S]的K Shape矩阵相乘，消掉Q的S维和K的H/a维，最终得到[B, a, S, S] Shape的输出。

Softmax(Score) ⇒ S

输入Tensor Shape：[B, a, S, S]
输出Tensor Shape：[B, a, S, S]
Tensor计算量：k*B*a*S*S，k为常数项，Softmax操作不涉及矩阵乘法，所以整体运算量比较小。

SV Matmul

V Shape同样做了多头注意力的 a 分组，所以V的Shape会变成：

[B, S, H] ⇒ [B, a, S, H/a]

输入Tensor Shape：[B, a, S, S], [B, a, S, H/a]
输出Tensor Shape：[B, a, S, H/a]
Tensor计算量：2BSSH = 2*B*a*S*S*H/a

O Projection：

O Projection和Q/K/V Projection差不多，不同的是O Projection的输入是SV Matmul出来的，Shape不一样。所以O的输入会有一个Reshape操作：

Reshape([B, a, S, H/a]) => [B, S, H]

输入Tensor Shape：[B, S, H]， [H, H]
输出Tensor Shape：[B, S, H]
Tensor计算量：2BSHH = 2*B*S*H*H=2*[B, S, H]*[H, H]

MHA总结

从MHA算子来看，输入和输出的Tensor Shape都是[B, S, H]，包括后面的MLP也是，在设计上是挺规整的，在MHA内部，主要实现的就是这个公式：

我们通过拆分计算过程，更清晰地理解了MHA的内部细节。

MLP内部计算细节

有了前面MHA的基础之后，MLP就比较好理解了。

1. Gate/Up Projection
2. Normalization：和MHA介绍的一样，就不重复介绍了
3. SiLU：和Softmax类似，这里也就不展开了
4. Down Projection

所以，重点还是Gate/Up/Down Projection。

Gate/Up Projection

在MLP层，有一个新的参数ffn_dim，它是MLP内部的隐藏层大小，和前面介绍的Hidden Size不一样。

输入Tensor Shape：[B, S, H], [H, ffn_dim]，这[H, ffn_dim]是预训练的权重
输出Tensor Shape：[B, S, ffn_dim]
Tensor计算量：22BSH*ffn_dim

输入[B, S, H]和[H, ffn_dim]，消掉两个H维之后，得到结果[B, S, ffn_dim]。

Gate/Up Projection过程相当于一个放大的计算逻辑，把Hidden Size的隐藏层大小放大到ffn_dim的隐藏层大小，扩大数据特征。

Down Projection

Down Projection做的事相反的计算逻辑，缩小的计算操作，把ffn_dim缩放到Hidden Size大小，返回原形。

输入Tensor Shape：[B, S, ffn_dim], [ffn_dim, H]
输出Tensor Shape：[B, S, H]
Tensor计算量：2BSH*ffn_dim

MLP总结

从计算过程可以看到MLP做的主要工作就是一个放大和缩小，把Hidden Size放大到ffn_dim，再缩小到Hidden Size，这个线性变换的过程有利于提取语义特征，捕捉计算过程的语义信息，增强表达能力。

不过，在这里也需要说明一下，不同模型的MLP层是有些不一样的，Llama模型是引入了Gate机制，对于更加一般的模型来说，MLP所做的事情大概就是下面这张图所呈现的结构了。

最后总结

可以看到，不管是进入Transformer Block还是出来，数据的Shape都是[B, S, H]，而这样的数据计算过程还需要经过Layers层，每层的计算过程都是一样的，不同的是预训练的权重不一样，也即是[H, H]矩阵和[ffn_dim, H]矩阵。

从Transformer Block出来之后，还需要将[B, S, H]反映射成词语，这里和Embedding做的事情是类似的但是是相反的，就不展开了。

本文从计算角度，详细解析了大型语言模型（LLM）的内部结构，并且分析了各个算子的输入输出数据结构及其计算量，结合前面的一系列文章，再加上这一轮从计算角度的介绍，相信大家对LLM的内部结构已经非常清晰明了了。

在这里插入图片描述

如何学习AI大模型？

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

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

在这里插入图片描述