简介

简介

本文基于 Meta LLaMA 3.2 3 B 模型架构，以“你知道王汉三是谁吗”为完整示例，从输入到输出逐步拆解大模型内部的计算流程。覆盖 Chat Template、BPE 分词、Token Embedding、28 层 Transformer Block、RMSNorm 归一化、GQA（分组查询注意力）、因果掩码、ROPE 旋转位置编码、FFN（前馈神经网络）及最终输出层。解析了模型如何将文本转换为数字、在多层网络中流转、最终逐字生成回答的全过程。

大模型完整推理流程

第一步：Chat Template（对话模板包装）

• 作用：为原始输入添加结构化标识，让模型区分“用户问话”与“自己作答”。
• 实现（以 LLaMA 3 为例）：在输入前后添加特殊 token。例如，原句“你知道王汉三是谁吗”被包装为：<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n你知道王汉三是谁吗<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n。
• 为什么需要：模型本身分不清对话角色。Chat Template 是模型在训练时见过的固定格式，模型看到 <|start_header_id|>assistant<|end_header_id|>\n\n 后面为空，就知道该自己回答了。
• 注意：不同模型（如 LLaMA、Qwen、DeepSeek）的 Chat Template 各不相同，不可混用。

第二步：BPE 分词（Tokenization）

• 为什么分词：模型只能处理数字，不能直接处理文字。
• BPE 原理：BBPE（字节对编码）是一种数据驱动的子词切分方法。它会统计海量语料中字符或子词的出现频率，从高频到低频逐步合并相邻的字符对，形成一个子词词表。
• 示例：包装后的长文本被分词器切分为 18 个 token（包括 <|begin_of_text|> 等特殊 token 和“知道”、“王汉三”等中文子词）。每个 token 都有一个唯一的数字 ID。

第三步：Token Embedding（词嵌入）

• 作用：将离散的 token ID 映射为连续的、能表达语义的稠密向量。
• 实现：通过查表实现。模型内部有一个形状为 [vocab_size, embedding_dim] 的权重矩阵（词嵌入表）。例如，对于 vocab_size = 128000，embedding_dim = 3072 的模型，token ID 101538 会直接取出该矩阵的第 101538 行，得到一个 3072 维的向量。
• 输出：输入句子的 18 个 token，各自查表后，堆叠成一个形状为 [18, 3072] 的矩阵。

第四步：RMSNorm（均方根归一化）

• 作用：将向量中的数值缩放到稳定、合理的范围，防止后续计算（尤其是 Softmax）出现数值溢出。
• 计算步骤：
  1. 对一行（即一个 token 的向量）求得均方根（RMS）：RMS = sqrt(1/d * Σ(x_i^2))。
  2. 每个元素除以 RMS，乘以一个可学习的缩放参数γ。
• 示例：假设原始向量数值范围在-18 到 20 之间，经过 RMSNorm 后会被压缩到-1 到 1 的区间。
• 维度不变：输出形状仍然是 [18, 3072]。

第五步：GQA（分组查询注意力机制）

• 核心问题：让每个 token 关注其之前的所有 token，理解上下文关系。
• QKV 概念：
  • Q（Query）：当前 token 想问的问题（想知道什么信息）。
  • K（Key）：每个 token 能提供的“标签”（有什么信息）。
  • V（Value）：每个 token 的实际“内容”（信息本身）。
• QKV 生成：将上一步输出的 [18, 3072] 矩阵，分别与三个不同的权重矩阵 W_Q, W_K, W_V 相乘。本例中，得到 Q 矩阵 [18, 3072]，K 和 V 矩阵 [18, 1024]。K 和 V 的维度仅为 Q 的三分之一，这是 GQA 的关键优化。
• 多头（Multi-Head）：模型从不同角度理解语义。
  • 分头：
    • Q 矩阵被切分成 24 个头（查询头），每个头维度 128。
    • K 和 V 矩阵被切分成 8 个头（KV 头），每个头维度 128。
• GQA 核心（分组查询）：每 3 个查询头（Q 头）共享一组 KV 头（K 头和 V 头）。即 Q 头 1、2、3 共用 K 头 1 和 V 头 1。
• 注意力计算（以一组 QKV 头为例）：
  1. 计算注意力分数：Score = Q * K^T，结果为一个 [9, 9] 的矩阵（假设有效 token 数为 9）。
  2. 缩放：除以 √d_k（d_k=128），防止 Softmax 进入饱和区，使注意力权重分布更柔和。
  3. 因果掩码：将未来位置的分数设为负无穷，确保模型只能看到当前位置及之前的 token。具体来说，在 [9, 9] 矩阵右上角填入负无穷。
  4. Softmax：对每行执行 Softmax，将分数转换为概率（权重），和为 1。
  5. 加权求和：用得到的权重矩阵与 V 矩阵相乘，得到融合了上下文信息的表示。例如，“知道”这个 token 的新的向量表示，是它自己与“你”的向量的加权和，而“王”及后续 token 的信息被掩码屏蔽。

后续模块

• 残差连接：将注意力层的输出与原始输入相加，帮助梯度直接流通，缓解深层网络退化问题。
• FFN（前馈神经网络）：每个 token 独立进行非线性变换，进一步提取和精炼信息。通常包含一个维度更大的中间层（如 8192 维）。
• 最终输出：经过 28 层 Transformer Block 后，通过一个线性层将 [18, 3072] 的向量映射回词表大小（128000 维），每个维度代表下一个词是相应 token 的概率。模型选择概率最高的 token 作为输出。

难点与补充（补充内容）

以下内容是对原文概念的更深层剖析和补充，帮助你更透彻地理解大模型工作原理。

补充 1：为什么 MHA（标准多头注意力）不如 GQA?[^1]

• 显存瓶颈：在推理时（尤其是长上下文），K 和 V 需要被缓存（KV Cache）。在 MHA 中，num_heads 个 KV 头意味着 num_heads 份 K 和 V 都需要被缓存。随着上下文长度增长，KV Cache 占用显存会线性增长，很快成为瓶颈。
• GQA 的优化：通过将 KV 头数减少为 num_key_value_heads，KV Cache 大小直接缩减为原来的 num_heads / num_key_value_heads 倍。例如，24 个 Q 头、8 个 KV 头，KV Cache 节省 2/3。
• 性能权衡：GQA 在几乎不损失效果的前提下，显著降低了显存占用并提升了推理速度，这对部署大规模模型至关重要。(补充自：GQA 论文)

补充 2：因果掩码(Causal Mask)的细节

• 位置：掩码在 Scale 操作之后、Softmax 之前施加。
• 实现：通常构建一个形状为 [seq_len, seq_len] 的上三角矩阵，上三角部分（j > i）的值为 -inf，下三角（j <= i）为 0。将这个矩阵加到 Q*K^T / √d_k 的结果上。
• 影响：Softmax 对 -inf 的计算结果是 0，因此模型“看不到”未来的 token。这确保了大模型在进行第 t 个 token 的预测时，只依赖于前 t-1 个 token 的信息，这是自回归生成的核心要求。

补充 3：FFN 中的隐藏层维度 (Intermediate Size)

• 原文示例: 原始描述中 FFN 中间维度为 8192，比嵌入维度 3072 高很多。
• 作用：将一个 token 的向量映射到一个更高维度的空间，在该空间中更容易分离和形成复杂的模式。可以理解为在更大的“工作台”上进行深度加工，然后再压缩回原始维度。
• 常用激活函数：原文未提及，但常见的 FFN 结构中会使用激活函数（如 ReLU, GELU, SwiGLU 等）引入非线性。LLaMA 系列常用 SwiGLU。

面试题

1. 基础理解题：请从头到尾简述一个问句（如，“什么是注意力机制？”）在大模型 Decoder-only 架构中的完整数据流。它经历了哪些主要模块？每一步输出的形状是什么？
2. 对比分析题：请对比 Multi-Head Attention (MHA)、Multi-Query Attention (MQA) 以及 Grouped Query Attention (GQA) 的异同。重点说明它们在计算量、显存占用（特别是 KV Cache）和模型效果上的区别。
3. 归一化原理题：RMSNorm 和 LayerNorm 都是归一化层，它们的主要区别是什么？为什么现代大模型更倾向于使用 RMSNorm？
4. 因果掩码理解题：为什么大模型在预测下一个 token 时需要使用因果掩码？如果不使用，会产生什么后果？请结合训练与推理两个阶段解释。
5. 数值稳定性题：为什么注意力机制中的 Q*K^T 结果要除以 √d_k？如果不做这个操作，直接进行 Softmax 会导致什么问题？请用数学或举例说明。
6. Embedding 深究题：Token Embedding 和 Positional Embedding（如 ROPE）的作用分别是什么？它们是如何结合的？ROPE 旋转位置编码为什么能有效处理相对位置信息？
7. 模型结构设计题：假设你需要在资源受限的设备上部署一个 LLM，你会如何设计其 Transformer Block 的架构？请从注意力头数、KV 头数（GQA）、FFN 中间层维度、层数等角度阐述你的设计思路和理由。

[^1]: ## 对比 Multi-Head Attention (MHA)、Multi-Query Attention (MQA) 和 Grouped Query Attention (GQA)

### 相同点

- 都属于**自注意力机制**的变体，用于Transformer模型。
- 都包含**Query（Q）、Key（K）、Value（V）** 的线性投影，然后计算注意力得分。
- 都可以在训练和推理中使用。

### 不同点

#### 1. 多头结构

- **MHA**：Q、K、V都拆分成多个头（例如8个头），每个头独立计算注意力，最后拼接。
- **MQA**：K和V**只有一个头**（即所有Query头共享同一个K、V），而Q仍然有多个头。
- **GQA**：将K和V分成**若干组**（例如4组），每组对应多个Query头。也就是介于MHA和MQA之间。

#### 2. 计算量（FLOPs）

- 在**生成阶段（自回归解码）**：
    - MHA每个头都要计算自己的K、V，计算量最大。
    - MQA因为K、V只有一个共享头，计算量最小（约减少为MHA的 1/num_heads）。
    - GQA的K、V分组数（g）通常小于头数（h），计算量比MHA小，但比MQA大：约是MHA的 g/h。
- 在**训练阶段**：三者计算量接近（因为训练时并行处理所有序列位置），但MQA和GQA因为K、V分头少，也有一些微小节省。

#### 3. 显存占用（特别是KV Cache）

- **KV Cache** 是自回归推理时为每个已生成token缓存的K、V张量，用于加速后续计算。
    - MHA：每个头都要缓存K和V，所以缓存大小为 `(batch_size, num_heads, seq_len, head_dim)` × 2。
    - MQA：所有头共享同一份K和V，缓存大小仅为 `(batch_size, 1, seq_len, head_dim)` × 2，**显存降低约 num_heads 倍**。
    - GQA：缓存大小为 `(batch_size, g, seq_len, head_dim)` × 2（g为分组数），显存占用介于MHA和MQA之间，降低为MHA的 g/num_heads。

#### 4. 模型效果

- **MHA**：表达能力最强，每个头可以关注不同的模式，通常效果最好。
- **MQA**：由于K、V共享，会降低表达能力，导致模型效果略差，尤其是在需要K、V区分度大的任务上（如长文本、细粒度语义）。
- **GQA**：通过分组折中，在保持部分多样性的同时大幅降低显存。实践中（如LLaMA 2 70B使用GQA），效果接近于MHA，而显存效率远高于MHA。分组数越大（越接近MHA）效果越好，但显存节省越少。

### 总结表格

|特性|MHA|MQA|GQA|
|---|---|---|---|
|K、V头数|等于Q头数（h）|1|分组数g（通常1 < g < h）|
|训练计算量|高|略低（可忽略）|介于中间|
|推理KV Cache显存|高（h倍）|极低（1/h）|低（g/h）|
|模型效果|最好|略差|接近MHA|
|适用场景|训练、大显存|极端显存受限|主流高效推理（如LLaMA 2/GPT-4）|

实际选择：**MHA适合训练**，**MQA适合极小显存设备**，**GQA是目前平衡效果和效率的最佳方案**。