用 transformers 推理 Qwen2-0.5B-Instruct

0. 背景

本文主要介绍 hugging face 托管的 LLM 如何下载，并用 transformers repo 推理

1. 相关工具介绍

hugging face 是一个托管 LLM 的网站，类似 github / dockerhub，可以上传下载大模型，也可以在线推理等
transformers 是一个 hugging face 维护的开源大模型代码库，包含约 300 个开源大模型的网络结构，可以配合 hugging face 下载的大模型，实现一键推理

2. 模型文件下载

2.1 下载工具

hugging face 有专用工具 huggingface-cli，作用类似于低配版 git，可以实现远程 hugging face repo 的登录、登出、上传、下载等
对于 Qwen2-0.5B-Instruct 模型，huggingface-cli 可直接免登录下载，而例如 Llama3.2-1B 等模型需要向 meta 提交申请并获得同意后下载

2.2 下载文件内容

hugging face 模型文件格式较为固定，一般为：
1. config.json：
  1. 这个文件包含了模型的配置信息，比如模型的架构参数、词表大小、激活函数类型等
  2. 它通常用于在加载模型时，确保模型的配置与原始模型一致
2. generation_config.json：
  1. 这个文件包含了生成文本时的配置信息，比如最大长度、停止序列、解码相关配置（温度、top-k、top-p、惩罚系数）等，以及依赖的 transformers 版本
  2. 这些参数影响模型生成文本的行为
3. LICENSE：
  1. 这个文件包含了模型的许可证信息，说明了你可以如何使用这个模型，以及使用时需要遵守的法律条款
4. merges.txt：
  1. 这个文件通常用于字节对编码（Byte Pair Encoding, BPE）分词器，它包含了词汇的合并规则
  2. BPE 是一种用于创建词汇表和编码文本的算法
5. model.safetensors：
  1. 这是一个保存模型权重的文件，使用了 SafeTensors 格式
  2. 这是一种高效的文件格式，用于存储和加载深度学习模型的权重
6. README.md：
  1. 这个文件包含了模型的说明文档，通常包括模型的简介、如何使用模型、训练细节、性能指标等信息
7. tokenizer_config.json：
  1. 这个文件包含了分词器的配置信息，比如分词器的类型、是否使用 BPE 等
8. tokenizer.json：
  1. 这个文件包含了分词器的预训练信息，比如词汇表、特殊标记等
  2. 它用于将文本转换为模型可以理解的数值输入
9. vocab.json：
  1. 这个文件包含了模型的词汇表，即模型在训练时使用的所有单词或标记的列表
  2. 分词器使用这个词汇表将文本分割成模型可以理解的输入
Qwen2-0.5B-Instruct 中的 "Instruct" 是指此模型是经过指令遵循微调后的模型

2.3 模型参数解析

下载的文件中 model.safetensors 就是模型参数
safetensors 是由 Hugging Face 开发的一种新型文件格式，专门用于安全地存储和加载机器学习模型的权重，这种格式特别关注模型的安全性、隐私保护和快速加载
1. 安全性：safetensors 格式的设计目标之一是提高模型存储的安全性，它不允许执行代码，从而减少了模型文件被恶意篡改的风险
2. 快速加载：与传统的 PyTorch 序列化格式（如 .pth 或 .bin）相比，safetensors 可以更快地加载模型权重，尤其是在 CPU 和 GPU 上
3. 跨语言和跨框架兼容性：safetensors 支持在不同的编程语言和机器学习框架之间共享模型权重，例如可以在 Python 中序列化并在 C++ / Java / JavaScript 中加载
4. 懒加载：支持懒加载，即可以选择只读取文件中的一部分张量，这对于分布式设置中的多节点或 GPU 非常有用
如何解析一个 model.safetensors 文件？

from safetensors import safe_open
tensors = {}
# framework="pt" means pytorch
with safe_open("model.safetensors", framework="pt", device="cpu") as f:
    for key in f.keys():
        tensors[key] = f.get_tensor(key)

3. 构建模型并加载预训练参数

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    model_path, torch_dtype="auto", device_map="auto"
)

hugging face 的 transformers 库提供了一个 AutoModelForCausalLM 类，可以根据模型路径，读取配置文件并自动加载模型

4. 输入构建

输入构建分成三个部分：

初始化 tokenizer
prompt 指令化
text to token id （tokenize）

4.1 初始化 `tokenizer`

1 2	from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(model_path)

transformers 会自动读取 tokenizer 相关配置文件，创建对应的 tokenizer 对象

4.2 `prompt` 指令化

prompt = "艾萨克牛顿是谁？"
messages = [
    {
        "role": "system",
        "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
    },
    {"role": "user", "content": prompt},
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)

上面的 message 是 Qwen2 指令微调后模型的输入格式，包含了 role 和 content 两个字段
具体来说需要两个步骤：
1. prompt to message：将 prompt 转换为 message 格式
2. message to chat：在 message 中插入特殊字符，形成 chat 格式，特殊字符主要包括：
  1. <|im_start|>：表示对话开始
  2. <|im_end|>：表示对话结束
  3. \n：间隔 role 和 content
最终生成的 text 实际内容为：

1	<\|im_start\|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<\|im_end\|>\n<\|im_start\|>user\n艾萨克牛顿是谁？<\|im_end\|>\n<\|im_start\|>assistant

一些额外知识：

prompt 输入模板信息在 tokenizer_config.json 中被定义过，key 为 chat_template：

{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}

模板用到的代码是 Jinja2 写的，Jinja2 是一种模板引擎，tokenizer 会根据这个模板生成 chat 格式的 text
代码中的 messages 和 add_generation_prompt 是模板的输入参数

4.3 `tokenize`(`text to token id`)

1	model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

这个过程就是 tokenize，将 text 转换为 token id
- 首先根据 tokenizer.json 中声明的 token encode 方法（例如 BPE），将 text 分词
- 然后根据 vocab.json 中的词汇表，将分词后的词转换为 token id（查表）
- 最后将 token id 转换为 tensor 格式（tensor of uint64 dtype）

5. 模型前向传播

generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512,
)

generate 方法是 transformers 提供的一个高级封装方法，可以实现 LLM 的推理
- max_new_tokens 参数表示最大生成的 token 数量
- generate 方法会自动迭代调用模型的前向传播方法和解码方法，实现文本生成
本章节的重点关注 generate 的模型推理过程，可以分成：

token id to token embedding
position emebedding 构造
attention mask 构造
decoder layer forward

5.1 `token id to token embedding`

nn.Embedding 本质就是查表，将 token id 转换为 token embedding

5.2 `position emebedding` 构造

Qwen2 模型使用的 position embedding 方法是 RoPE (rotary position embedding)

5.2.1 `RoPE` 介绍

RoPE 的核心思想是将 query embedding 和 key embedding 旋转一定角度，让模型感知到序列中的位置信息。

RoPE 的输入是：
1. position id（例如：[0, 1, 2, ..., seq_len - 1]）
2. query embedding (shape = [seq_len, head_dim])
3. key embedding (shape = [seq_len, head_dim])
RoPE 的输出是：
1. 经过 RoPE 处理后的 query embedding
2. 经过 RoPE 处理后的 key embedding
RoPE 的计算过程是：
1. 构造 $\cos\in\mathbb{R}^{seq\_len\times head\_dim}$ 和 $\sin\in\mathbb{R}^{seq\_len\times head\_dim}$ 两个矩阵
2. 用两个矩阵去旋转 query embedding 和 key embedding
RoPE 的实现代码如下：

# 构造 RoPE 旋转矩阵
inv_freq = 1.0 / (1000000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))    # 1. 计算频率的倒数，shape = [dim // 2]
freq = inv_freq_expanded.float() @ position_ids_expanded.float()    # 2. 计算频率，shape = [seq_len, dim // 2]
emb = torch.cat((freqs, freqs), dim=-1)   # 3. 构造 RoPE 矩阵，shape = [seq_len, dim]
cos = emb.cos()[None, None, ...]    # 4. 计算 RoPE 的 cos，扩展到多 batch 和多头，shape = [1, 1, seq_len, dim]
sin = emb.sin()[None, None, ...]    # 5. 计算 RoPE 的 sin，扩展到多 batch 和多头，shape = [1, 1, seq_len, dim]
# 旋转 query embedding 和 key embedding
def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)
q = (q * cos) + (rotate_half(q) * sin)  # 6. 旋转 query embedding，shape = [batch_size, num_heads, seq_len, head_dim]
k = (k * cos) + (rotate_half(k) * sin)  # 7. 旋转 key embedding，shape = [batch_size, num_heads, seq_len, head_dim]

以上 RoPE 的公式解释：
- $x'_{m,i}=x_{m,i}\cdot \cos (m\theta_i)-x_{m,i+d/2}\cdot \sin (m\theta_i)$
- $x'_{m,i+d/2}=x_{m,i}\cdot \cos (m\theta_i)+x_{m,i+d/2}\cdot \sin (m\theta_i)$
- 其中：
  - $x_{m,i}$ 是 query embedding 或 key embedding 的第 $m$ 个位置的第 $i$ 个维度的值
  - $\theta_i=1000000^{-\frac{2i}{d}}$
  - $m\theta_i=m\cdot 1000000^{-\frac{2i}{d}}$ ， $m$ 是位置 id
  - $d$ 是 head_dim
Qwen 模型中的 RoPE 实现代码是 RoPE 的一个简化版本，主要是为了提高计算效率，官方的 RoPE 计算公式如下：
- $x'_{m,2i}=x_{m,2i}\cdot \cos (m\theta_i)-x_{m,2i+1}\cdot \sin (m\theta_i)$
- $x'_{m,2i+1}=x_{m,2i}\cdot \cos (m\theta_i)+x_{m,2i+1}\cdot \sin (m\theta_i)$
- $\mathbf{R}_{d\Theta,m}\mathbf{x} = \begin{bmatrix} \mathbf{x}_1 \\ \mathbf{x}_2 \\ \mathbf{x}_3 \\ \mathbf{x}_4 \\ \vdots \\ \mathbf{x}_{d-1} \\ \mathbf{x}_d \end{bmatrix} \otimes \begin{bmatrix} \cos(m\theta_1) \\ \cos(m\theta_1) \\ \cos(m\theta_2) \\ \cos(m\theta_2) \\ \cdots \\ \cos(m\theta_{d/2}) \\ \cos(m\theta_{d/2}) \end{bmatrix} + \begin{bmatrix} -\mathbf{x}_2 \\ \mathbf{x}_1 \\ -\mathbf{x}_4 \\ \mathbf{x}_3 \\ \vdots \\ -\mathbf{x}_d \\ \mathbf{x}_{d-1} \end{bmatrix} \otimes \begin{bmatrix} \sin(m\theta_1) \\ \sin(m\theta_1) \\ \sin(m\theta_2) \\ \sin(m\theta_2) \\ \cdots \\ \sin(m\theta_{d/2}) \\ \sin(m\theta_{d/2}) \end{bmatrix}$
- 区别在于：
  - 官方的 RoPE 是同一个位置的相邻维度（2i 和 2i+1）之间旋转
  - Qwen 的 RoPE 是同一个位置的跳跃维度（i 和 i+d/2）之间旋转

5.3 `attention mask` 构造

实现中，attention mask 实际上并不需要构造，因为 torch.nn.functional.scaled_dot_product_attention 的 is_causal 参数可以自动实现 attention mask
具体来说：
1. torch.nn.functional.scaled_dot_product_attention 的 is_causal == True 时，会自动构造一个 attention bias，shape = [query_len, key_len]，其中 bias[i, j] = -inf if i > j else 0
2. 在 query @ key.transpose(-2, -1) 得到 attention score 后，会自动加上这个 attention bias
3. 然后再经过 softmax，这样就实现了 causal attention
而且 causal attention 在 GPT-like LLM 中，只在 prefill 阶段使用，generate 阶段不使用（或者说隐式使用）

5.4 `decoder layer forward`

LLM 的计算量几乎都体现在 decoder layer forward 上
从 layer norm 位置角度来看，Qwen2 模型属于 pre-LM，即：
1. $x=x+MHA(LayerNorm(x))$
2. $x=x+FFN(LayerNorm(x))$
3. $x = LayerNorm(x)$
模型包含 24 层 decoder layer，每层包含 2 个子层：
1. multi-head self-attention（MHA）
2. feed forward network（FFN）
为了减小计算量和参数量，Qwen2 模型使用了 GQA (grouped query attention) 方法：
1. 对于 query 来讲，每层 MHA 包含 14 个头，每个头的 head_dim = 64
2. 对于 key 和 value 来讲，每层 MHA 包含 2 个头，每个头的 head_dim = 64
3. 计算时，需要将 key 和 value 的 head 通过 repeat 扩展到 14 个头
同时，Qwen2 使用了 KV cache 算法，即：
1. 在 prefill 阶段，key 和 value 会被缓存下来
2. 在 generate 阶段，之前的 key 和 value 会被直接使用不再计算，只计算最后一个 token 的 query / key / value
在 prefill 阶段：
- 输入为 hidden state，shape = [1, 36, 896]，其中 36 是 seq_len，896 是 hidden_dim
- 输入逐层经过 decoder layer，最后输出 hidden state，和输入 shape 一样（shape = [1, 36, 896]）
- 在这一阶段，每一层的每个头的每个位置（token position）的 key 和 value 都会被缓存下来，shape = [2, 1, 24, 2, 36, 64]，其中：
  - 2 是 key / value
  - 1 是 batch_size
  - 24 是 layers
  - 36 是 seq_len
  - 2 是 key_value_head_num
  - 64 是 head_dim
- 每个 decoder layer 层的 MHA 都需要 causal attention，即 is_causal = True
- 输出的 hidden state 的 最后一个 token position 的 feature (shape = [1, 896]) 会被输入到 fc 层，输出 shape = [1, 151936]，151936 是 vocab_size
- 然后通过 decode method 得到最终的 token id
在 generate 阶段：
- 输入为 预填充阶段最终输出的 token id 对应的 token embedding，shape = [1, 1, 896]
- 输入逐层经过 decoder layer，最后输出 hidden state，和输入 shape 一样（shape = [1, 1, 896]）
- 在这一阶段，每一层的每个头的之前的位置（token position）的 key 和 value 都会被直接使用，不再计算；只计算最后一个 token 的 query / key / value，并将计算得到的 key / value 缓存下来，缓存的 key / value 会拼接到之前的 key / value 上，形成新的 key / value，shape = [2, 1, 24, 14, 36 + 1, 64]
- 每个 decoder layer 层的 MHA 的输入 shape 如下：
  - query: shape = [1, 14, 1, 64]
  - key / value: shape = [1, 14, 37, 64]
  - is_causal = False，因为 query 长度为 1，不需要 causal attention
- 最后一层 decoder layer 的输出经过 fc 和 decode method 得到最终的 token id
- 重复直到生成的 token id 数量达到 max_new_tokens 或生成的 token id 为 eos 为止

6. `decode method`

decode method 实际上属于模型 forward 过程中的一部分，但是为了方便理解，这里单独拎出来讲
decode method 的作用是将 vocab logits 映射到 token id
最简答的 decode method 是 argmax 方法，即取 vocab logits 中最大的值对应的 token id，但是这种方法会导致生成的文本过于单一
实际上 Qwen2 使用了五个串联的 decode method，分别是：
- RepetitionPenaltyLogitsProcessor
- TemperatureLogitsWarper
- TopKLogitsWarper
- TopPLogitsWarper
- SoftmaxSampler

6.1 `RepetitionPenaltyLogitsProcessor`

RepetitionPenaltyLogitsProcessor 的作用是惩罚重复的 token id，即：

在预测的 vocab logits 中，找到之前所有 token id 对应的 logits
如果 logits > 0，则将 logits 除以一个 penalty，penalty 的值取 1.1
如果 logits < 0，则将 logits 乘以一个 penalty，penalty 的值取 1.1
目的是让模型更倾向于生成不重复的 token id

6.2 `TemperatureLogitsWarper`

TemperatureLogitsWarper 的作用是调整 vocab logits 的温度，即：

将 vocab logits 除以一个 temperature，temperature 的值取 0.7
目的是提高模型的生成多样性

6.3 `TopKLogitsWarper`

TopKLogitsWarper 的作用是截断 vocab logits，即：

保留 vocab logits 中最大的 k 个值，其他值置为 -inf，k 的值取 20
目的是降低模型的生成多样性，提高生成的准确性

6.4 `TopPLogitsWarper`

TopPLogitsWarper 的作用是截断 vocab probs，即：

将 vocab logit 通过 softmax 变成 vocab probs
将 vocab probs 从大到小排序，累加到 p 大于 top-p 为止，保留这些 vocab logits，其他值置为 -inf
top-p 的值取 0.8，最终至少需要保留一个 token id
目的是降低模型的生成多样性，提高生成的准确性

6.5 `SoftmaxSampler`

将 vocab logits (此时只有少量元素不是 -inf) 通过 softmax 变成 vocab probs
根据 vocab probs 采样一个 token id，作为最终的输出

7. 总结

本文主要介绍了 hugging face 和 transformers 的使用方法，以及 Qwen2-0.5B-Instruct 模型的推理过程
GPT-like 模型的推理过程主要包括 tokenize、RoPE、attention mask、decoder layer forward、decode method 等步骤
Qwen2 模型使用了 RoPE、GQA、KV cache 等技术，提高了模型的计算效率和参数量
decode method 使用了 RepetitionPenaltyLogitsProcessor、TemperatureLogitsWarper、TopKLogitsWarper、TopPLogitsWarper、SoftmaxSampler 等方法，提高了模型的生成多样性和准确性

0. 背景

1. 相关工具介绍

2. 模型文件下载

2.1 下载工具

2.2 下载文件内容

2.3 模型参数解析

3. 构建模型并加载预训练参数

4. 输入构建

4.1 初始化 tokenizer

4.2 prompt 指令化

4.3 tokenize(text to token id)

5. 模型前向传播

5.1 token id to token embedding

5.2 position emebedding 构造

5.2.1 RoPE 介绍

5.3 attention mask 构造

5.4 decoder layer forward

6. decode method

6.1 RepetitionPenaltyLogitsProcessor

6.2 TemperatureLogitsWarper

6.3 TopKLogitsWarper

6.4 TopPLogitsWarper

6.5 SoftmaxSampler