【Rust与AI】LLM模型基本架构

本篇是《Rust与AI》系列的第二篇，上一篇我们主要介绍了本系列的概览和方向，定下了一个基调。本篇我们将介绍LLM的基本架构，我们会以迄今为止使用最广泛的开源模型LLaMA为例展开介绍。

LLM背景

Rust 本身是不挑 AI 模型的，但是 LLM 是当下最热的方向，我们就从它开始吧，先了解一些非常基础的背景知识。

Token

LLM 中非常重要的一个概念是 Token，我们输入给 LLM 和它输出的都是 Token。Token 在这里可以看做语言的基本单位，中文一般是词或字（其实字也是词）。比如：”我们喜欢 Rust 语言“，Token 化后会变成类似 ”我们/喜欢/Rust/语言“ 这样的四个词，可以理解为四个 Token。

给定一段任意的自然语言文本，我们可以用一个分词器（Tokenizer）将其 Token 化成一个个连续的 Token。这些 Token 接下来就可以映射成一个个数字，其实是在词表中的索引，索引进而可以找到一个稠密向量，用来表示该位置 Token 的语义输入。

我们以刚刚的”我们喜欢 Rust 语言“为例，假定已有词表如下。

……
1000 Rust
……
2000 我们
2001 喜欢
2002 语言
……

注意，前面的数字是行号，并不是词表内容。刚刚那句话其实就是 [2000, 2001, 1000, 2002]，这就是 LLM 的输入。LLM 拿到这些 ID 后，会在一个非常大的表里查找对应的稠密向量。这个非常大的表就是词表，大小是：词表大小N × 模型维度，如下所示。

……
1000 0.9146, 0.066, 0.4469, 0.3867, 0.3221, 0.6566, 0.2895, ...
……
2000 0.5702, 0.9579, 0.0992, 0.9667, 0.5013, 0.4752, 0.1397, ...
2001 0.2896, 0.7756, 0.6392, 0.4034, 0.3267, 0.9643, 0.4311, ...
2002 0.4344, 0.6662, 0.3205, 0.3929, 0.6418, 0.6707, 0.2414, ...
……

也就是说，输入”我们喜欢Rust语言“这句话，我们实际传递给模型的其实是一个 4×Dim 的矩阵，这里的 4 一般也叫 Sequence Length。

我们可以暂时把模型看作一个函数 f(x)，输入一个 Sequence Length × Dim 的矩阵，经过模型 f(x) 各种运算后会输出 Sequence Length × Vocabulary Size 大小的一个概率分布。有了概率分布就可以采样一个 Token ID（基于上下文最后一个 Token ID 的分布），这个 ID 也就是给定当前上下文（”我们喜欢Rust语言“）时生成的下一个 Token。接下来就是把这个 ID 拼在刚刚的 4 个 ID 后面（输入变成 5 个 ID），继续重复这个过程。

生成

如上所言，生成过程就是从刚刚的概率分布中 “选择” 出一个 Token ID 作为下一个 Token ID。选择的方法可以很简单，比如直接选择概率最大的，此时就是 Greedy Search，或 Greedy Decoding。

不过我们平时用到大模型时一般都用的是采样的方法，也就是基于概率分布进行采样。抛硬币也是一种采样，按概率分布（0.5，0.5）进行采样，但假设正面比较重，概率分布就可能变成了（0.8，0.2）了。基于 Vocabulary Size 个概率值进行采样也是类似的，只不过括号里的值就是词表大小那么多个。

top_p/top_k 采样是概率值太多了，大部分都是概率很小的 Token，为了避免可能采样到那些概率很低的 Token（此时生成的结果可能很不连贯），干脆就只从前面的 Token 里挑。

top_k 就是把 Token 按概率从大到小排序，然后从前 k 个里面选择（采用）下一个 Token；top_p 也是把 Token 按概率从大到小排序，不过是从累积概率大于 p 的 Token 里选。就是这么简单。

这里有个小细节需要说明，因为选择了 top_p/k，所以这些备选的 Token 需要重新计算概率，让它们的概率和为 1（100%）。

开源代表——LLaMA

接下来，我们把重心放在函数 f(x) 上，以最流行的开源 LLM——LLaMA 为例，简单介绍一下模型的结构和参数。

结构

LLaMA 的结构相对而言比较简单，如果我们忽略其中的很多细节，只考虑推理过程，看起来如下图所示。

图中 [] 中的是该位置的张量 shape，B 表示 Batch Size，一般时候都是批量丢给 GPU 计算的，L 就是 Sequence Length，D 就是上面提到的 Dim。这是一个简化了的架构图，但是足以清晰地表达模型了。

两个 Hidden states（以下简称 HS），外面（之上和之下）的部分我们前面已经提到过了（注意上面部分，[B,L,D] 会先变成 [B,L,VS]，然后取最后一个 Token 就得到了 [B,1,VS]），上面的 HS 会传回到 Block 里面，重复 N 次，N 就是模型的层数。接下来我们就把重点放在中间这个 Block 里。

每个 Block 包括两个主要模块，一个 MHA（Multi-Head Attention）模块，一个 FFN（Feedforward Network）模块，每次传给模块之前都需要 Normalization，这个叫 Pre-Normalization，一般用来稳定训练。另外，每个模块结束后会叠加模块之前的输入，这个叫残差连接，一般能加速收敛。

接下来是 MHA 和 FFN，先看 FFN 模块，它的大概流程如下（@ 表示矩阵/张量乘法）。

z1 = ns @ up_weights
z2 = ns @ gate_weights
z3 = z1 * silu(z2)
z4 = z3 @ down_weights

整体来看是先将网络扩大再收缩，扩大时增加了一个激活处理。silu 函数大概长这样：

等价于只激活了一部分参数，这个非线性激活非常重要，可以让模型学习到更丰富的知识和表达。

再就是 MHA 模块了，大概流程如下（为了更直观，去掉了 Batch Size 和 Softmax）。

q = ns @ q_weights # (L, D) @ (D, D) = (L, D)
k = ns @ k_weights # (L, D) @ (D, D) = (L, D)
v = ns @ v_weights # (L, D) @ (D, D) = (L, D)

q = q.reshape(L, NH, HD)
k = k.reshape(L, NH, HD)
v = v.reshpae(L, NH, HD)

attn = q.trans(NH, L, HD) @ k.trans(NH, HD, L)  # (NH, L, HD) @ (NH, HD, L) = (NH, L, L)
v = attn @ v.trans(NH, L, HD) # (NH, L, L) @ (NH, L, HD) = (NH, L, HD)
v = v.reshpe(L, NH*HD) # (L, D)

其中，NH 表示 Attention 的 Head 数，HD 表示 Head 的维度。因为有 NH 个 Head，所以叫 Multi-Head，但其实我们看上面的过程，在实际计算的时候它们是合并一起算的。我们不妨只看一个 Head，如下所示。

q = ns @ hq_weights # (L, D) @ (D, HD) = (L, HD)
k = ns @ hk_weights # (L, D) @ (D, HD) = (L, HD)
v = ns @ hv_weights # (L, D) @ (D, HD) = (L, HD)

attn = q @ k.T # (L, HD) @ (HD, L) = (L, L)
v = attn @ v # (L, L) @ (L, HD) = (L, HD)

上面的多个 Head 的 v 就是下面的每个 Head 的 v 拼接起来的。

Multi-Head 是多个注意力头去执行 Attention，其思想是让每个 Head 去捕获不同角度/层面的 Attention，这些角度/层面是什么？不是特别清楚（但一定是某种特征），但我们可以通过 Attention 的权重看出外在 Token 级别的注意力，知道每个注意力 Head，哪些 Token 之间有比较强的连接。

参数

关于 f(x) 我们已经介绍完了，可以发现这个函数其实还是有点复杂的。接下来，我们看看参数情况。

对一个一元一次方程（比如 f(x) = ax + b）来说，参数就两个：a 和 b，但对于 LLM 来说，参数就非常多了，目前常用的是 7B、13B、20B 的级别，也就是 70亿、130亿和 200亿的参数规模。

在神经网络中，可以把矩阵乘法看作是多元一次方程组的计算过程，输入的 Hidden State 维度是 D，就表示未知变量的维度是 D，也就是 D 元一次方程组。

以前面的但 Head Attention 的 q 为例，q_weights 是一个 DxHD 的参数矩阵，我们把 D 和 HD 设置的小一点（假设为4和2），看一个具体的例子。