Building an LLM Inference Engine from ScratchPart III / Module 7

Part III · Concrete Candidates

Module 7：模型选择 — Qwen2 / Qwen2.5

把抽象 transformer 落到一个真实 checkpoint：读懂 Qwen2 的 config、张量命名、权重格式与显存规模。

学习目标

读懂 Qwen2 / Qwen2.5 的 config.json，并把字段转成张量 shape。
掌握 Qwen2 相对通用 decoder 的具体差异：rope_theta、GQA head 数、QKV bias、tied embeddings 及其对应张量名。
解释 tied / untied LM head 与词表大小对内存和计算的影响。
理解 safetensors、f32/f16/bf16、memory mapping 在推理加载中的作用。
能估算一个 checkpoint 是否适合教学机器的 VRAM，并知道哪些量随模型大小增长。

7.1Qwen2 decoder 架构与 config.json

前几章使用的是 decoder-only transformer 的通用描述。实现推理引擎时，还必须选择一个具体模型族，因为权重名称、config 字段、attention head 划分、normalization 细节都会进入代码。本项目以 Qwen2 / Qwen2.5 为目标：公开 checkpoint 容易取得，0.5B 规模适合调试，同时保留了现代 LLM 中常见的 GQA、RoPE、RMSNorm、SwiGLU、QKV bias 与 tied embeddings。

xinfer-core 中的 QwenConfig 只保留推理需要的字段；xinfer-loader 从模型目录读取 config.json 并反序列化为该结构。以仓库中的 Qwen2.5-0.5B-Instruct 为例，关键配置如下：

字段	值	含义
`hidden_size`	896	residual stream 宽度 $H$
`intermediate_size`	4864	MLP 的扩展维度 $I$
`num_hidden_layers`	24	decoder layer 数 $L$
`num_attention_heads`	14	query head 数
`num_key_value_heads`	2	K/V head 数（GQA）
`vocab_size`	151936	词表大小，也决定 LM head 输出维度
`rope_theta`	1000000	RoPE 的频率基底
`rms_norm_eps`	1e-6	RMSNorm 稳定项
`tie_word_embeddings`	`true`	LM head 与 embedding table 共享权重

这些字段直接决定运行时 buffer 的大小。单个 attention head 的维度为：

d_{\text{head}} = \frac{H}{n_{\text{heads}}} = \frac{896}{14}=64

因此 Q 的 shape 是 $[S,14,64]$ 。GQA 使 K/V 只保留 2 个 KV head，shape 为 $[S,2,64]$ 。代码中的 head_dim() 与 kv_group_size() 正是从这些字段推导得到，而不是写死在 kernel 或模型逻辑中。

图 7-1：config.json 字段会直接决定推理时的张量 shape 和 buffer 大小。

7.2Qwen2 相对通用 decoder 的差异点

Qwen2 的 decoder layer 与 Module 2 描述的通用结构完全一致——RMSNorm、RoPE、GQA、SwiGLU 的数学定义已在那里给出，这里不再重复。装配一个真实 checkpoint 时，真正需要落实的是每个算子的具体取值与权重张量名称。下表把 Qwen2 的选择与 Module 2 的对应小节并列，作为对照表（张量名均省略 model.layers.{i}. 前缀）：

结构	Qwen2 的选择	config 字段	权重 / 张量名	数学见
归一化	RMSNorm，pre-norm，不减均值、无 bias	`rms_norm_eps=1e-6`	`input_layernorm.weight`、`post_attention_layernorm.weight`、`model.norm.weight`	§2.2
位置编码	RoPE（GPT-NeoX `rotate_half`）	`rope_theta=1000000`	无权重，由位置在线计算	§2.6
注意力	GQA：14 query / 2 KV head，group=7	`num_attention_heads=14`、`num_key_value_heads=2`	`self_attn.{q,k,v,o}_proj.weight`	§2.4–2.5
QKV bias	带 bias（区别于多数 Llama）	由权重是否存在决定	`self_attn.{q,k,v}_proj.bias`	§2.4
MLP	gated SwiGLU（SiLU）	`hidden_act="silu"`、`intermediate_size=4864`	`mlp.{gate,up,down}_proj.weight`	§2.7
LM head	tied embeddings	`tie_word_embeddings=true`	复用 `model.embed_tokens.weight`	§2.8、§7.3

表中唯一需要展开的差异是 QKV bias：多数 Llama 系模型的 Q/K/V projection 不带 bias，而 Qwen2 带 bias。因此模型装配必须额外加载 q_proj.bias、k_proj.bias、v_proj.bias（o_proj 仍无 bias）。xinfer-model 的 CPU reference 与 GPU forward 都按这些名称读取 bias；遗漏会使 logits parity 失效。

图 7-2：GQA 让 Q head 数保持较多，但 K/V cache 只需要较少 head。

7.3Tied vs untied LM head 与词表规模

tie_word_embeddings=true 表示输出 LM head 与输入 embedding table 共享同一份权重：

W_{\text{lm}} = E

对 Qwen2.5-0.5B，这一张表的元素数为：

|V|\cdot H = 151936 \times 896 \approx 1.36\times 10^8

若以 f16 存储，大小为：

151936 \cdot 896 \cdot 2 \approx 272\text{ MB}

按 MiB 计约为 260 MiB。tied embeddings 使这份权重同时用于 token embedding 与 logits projection；xinfer-model 在 tie_word_embeddings 为 true 或 checkpoint 中没有 lm_head.weight 时，使用 model.embed_tokens.weight 作为 LM head。

性能提醒

大词表不仅占内存，也影响每个 token 的 LM head 计算和 logits 读回。xinfer 后来的 GPU argmax 优化正是为了避免每步读回 $[151936]$ 个 logits。

7.4权重格式：safetensors、dtype 与 memory mapping

HuggingFace checkpoint 通常由 config.json、tokenizer.json 和一个或多个 *.safetensors 文件组成。safetensors 的文件头记录 tensor 名称、dtype、shape 与 byte offset，数据区连续存放原始 bytes。xinfer-loader 遍历模型目录中的全部 *.safetensors 文件，用 memmap2 建立 memory mapping，再交给 safetensors 解析。

图 7-3：safetensors 适合推理加载：结构明确，可 memory-map，避免 pickle 类格式的安全问题。

加载阶段支持 f32、f16、bf16 三种权重 dtype，并统一转换为 f32 host tensor。GPU 侧再根据用途选择存储格式：decoder layer 中 q/k/v/o 与 gate/up/down 七类线性层权重以 f16 上传，bias 与 norm weight 保持 f32；shared LM head 也以 f16 上传。这样可以保留 CPU reference 的简单性，同时降低 GPU 权重显存和带宽。

dtype	字节/元素	特点	本项目中如何用
f32	4	精度高，内存/带宽大	CPU reference、activation
f16	2	半精度，显存减半	GPU matmul weights / LM head
bf16	2	指数范围接近 f32，尾数更少	加载时可转成 f32/f16

7.5选择一个合适大小：为什么从 0.5B 开始？

教学模型需要同时满足两个约束：结构必须接近实际 LLM，规模又要允许频繁运行 parity test、GPU kernel test 与端到端生成。Qwen2.5-0.5B 有 24 层、真实 tokenizer、真实 Qwen2 结构，适合作为从 CPU reference 迁移到 D3D12 backend 的基线。README 中的运行示例显示，f16 权重常驻 GPU 后，0.5B checkpoint 的 GPU resident weights 约为 942 MiB。

随模型规模增长，主要增长项包括：

层数 $L$ ：所有 per-layer 权重、KV cache 都随 $L$ 线性增长；
hidden size $H$ ：线性层权重通常随 $H^2$ 或 $H\cdot I$ 增长；
intermediate size $I$ ：MLP gate/up/down 权重随 $H\cdot I$ 增长；
词表大小 $|V|$ ：embedding / LM head 随 $|V|\cdot H$ 增长；
上下文长度 $T$ ：KV cache 随 $T$ 线性增长。

这些量的增长方式决定了课程中的实验顺序。先在 0.5B 上建立正确性与性能测量，再讨论 layer streaming、copy queue 与更大模型，比直接从大模型开始更容易定位错误。

图 7-4：Qwen2.5-0.5B 的权重显存中，embedding/LM head 和 MLP 权重都很重要（示意）。

Lab 7解析真实 config 并检查 safetensors

本实验要求你写一个小工具，读取模型目录中的 config.json 和 model.safetensors，输出关键字段和若干 tensor 的 shape。若 checkpoint 被拆成多个 *.safetensors 文件，也应逐个读取并合并名称列表。

# 示例输出（格式可自定义）
model_type: qwen2
hidden_size: 896
layers: 24
heads: 14
kv_heads: 2
head_dim: 64
vocab_size: 151936

model.embed_tokens.weight: [151936, 896]
model.layers.0.self_attn.q_proj.weight: [896, 896]
model.layers.0.self_attn.k_proj.weight: [128, 896]
model.layers.0.mlp.gate_proj.weight: [4864, 896]

你也应该验证这些 shape 能从 config 推导出来。例如：

\operatorname{shape}(W_Q) = [H,H],\qquad \operatorname{shape}(W_K) = [n_{\text{kv}}d_{\text{head}}, H]

小结

本章把通用 decoder-only transformer 约束到 Qwen2 / Qwen2.5 的具体实现。config.json 决定 hidden size、层数、head 数、GQA 比例、词表大小、RoPE 基底和 tied embeddings；safetensors 提供可 memory-map 的权重布局，loader 将 f16/bf16/f32 权重转换为 f32 host tensor；GPU 路径再将主要线性层权重和 LM head 压到 f16。选择 0.5B 模型的目的在于：保留真实结构，同时控制显存、加载时间和调试成本。

思考与练习

基础从 Qwen2.5-0.5B 的 config 推导 head dimension。

$d_{\text{head}}=\text{hidden\_size}/\text{num\_attention\_heads}=896/14=64$ 。K/V head 的单头维度相同，但 K/V head 数只有 2，因此 K/V projection 的输出维度是 $2\times64=128$ 。

基础解释 tie_word_embeddings=true 的含义。

它表示输入 embedding 矩阵与输出 LM head 共享同一份参数。对 Qwen2.5-0.5B，这避免了额外保存一张 $151936\times896$ 的大矩阵；f16 下约节省 260 MiB。xinfer 在该字段为 true 时使用 model.embed_tokens.weight 作为 LM head。

进阶为什么 Qwen2 的 K/V projection 输出维度是 128，而 Q projection 是 896？

因为该模型采用 GQA：query 有 14 个 head，K/V 只有 2 个 head。每个 KV head 服务 7 个 query heads，从而降低 K/V projection 计算量和 KV cache 大小。attention 计算时按 kv_group_size=14/2 将 query head 映射到对应的 KV head。

进阶估算 embedding table 的 f16 显存大小。

元素数为 $151936\times896$ ，f16 每个元素 2 bytes，因此大小为 $151936\times896\times2=272{,}269{,}312$ bytes，约 260 MiB。由于 tied embeddings，这份权重也用于 LM head，不需要再为输出投影保存第二份同形状矩阵。

挑战写一个脚本列出 safetensors 中所有 tensor，并按总字节数排序前 10。

脚本可以用 safetensors 库打开文件，遍历 f.keys()，读取每个 tensor 的 shape 与 dtype，并按 $\prod\text{shape}\times\text{bytes(dtype)}$ 计算字节数。若模型目录中有多个 *.safetensors 文件，应合并全部文件的结果后排序。预期前列包括 embedding、各层 MLP 的 gate/up/down 权重，以及 attention 的 q/o projection 权重。