推理模型从零构建 附录C：Qwen3 源码解读
本文是 Build a Reasoning Model (From Scratch)（Sebastian Raschka 著）附录 C 的学习笔记。本附录展示了全书使用的 Qwen3 0.6B 模型的完整 PyTorch 源码。虽然 “from scratch” 指的是推理技术而非 LLM 本身，但理解模型架构对于理解第2章的推理生成、第5章的 log-probability 计算、第6章的 GRPO 训练都有帮助。本文逐层解读 Qwen3 的 6 个核心组件：RMSNorm → SwiGLU → RoPE → GQA → TransformerBlock → Qwen3Model，外加 KVCache 和 Tokenizer。
📖 Table of Contents
1. 架构总览
Qwen3 0.6B 的架构和 GPT-2 非常相似，都是 decoder-only transformer。但有几个关键的现代化升级：
Qwen3 0.6B 架构：28 层 TransformerBlock，每层包含 GQA 注意力 + SwiGLU FFN (图源: Reasoning from Scratch)
| 参数 | 值 | 说明 |
|---|---|---|
| vocab_size | 151,936 | 词表大小（支持中英日韩等多语言） |
| context_length | 40,960 | 最大上下文长度（约半本哈利波特） |
| emb_dim | 1,024 | 嵌入维度 |
| n_heads | 16 | 注意力头数 |
| n_kv_groups | 8 | KV 分组数（GQA） |
| head_dim | 128 | 每个 head 的维度 |
| hidden_dim | 3,072 | FFN 中间层维度 |
| n_layers | 28 | Transformer 层数 |
| dtype | bfloat16 | 参数精度 |
与 GPT-2 的主要区别：LayerNorm → RMSNorm，GELU → SwiGLU，Absolute PE → RoPE，MHA → GQA，QK Norm（额外的 query/key 归一化）。这些改进在 2024–2025 年的大多数 LLM 中已成为标配。
2. RMSNorm：归一化层
RMSNorm 只用均方根归一化，省去了 LayerNorm 中的均值计算 (图源: Reasoning from Scratch)
Python — RMSNorm
class RMSNorm(nn.Module):
def __init__(self, emb_dim, eps=1e-6, bias=False):
super().__init__()
self.eps = eps
self.scale = nn.Parameter(torch.ones(emb_dim)) # 可学习的缩放参数
def forward(self, x):
input_dtype = x.dtype
x = x.to(torch.float32) # 归一化在 float32 下计算（精度）
# RMS = sqrt(mean(x^2))，然后 x / RMS
variance = x.pow(2).mean(dim=-1, keepdim=True)
norm_x = x * torch.rsqrt(variance + self.eps)
return (norm_x * self.scale).to(input_dtype) # 转回原精度RMSNorm 比 LayerNorm 简单：省去了均值的减法操作，只用 RMS（均方根）做归一化。实践中效果相当，但计算量减少约 10–15%。torch.rsqrt 是 1/sqrt(x) 的高效实现。
3. FeedForward：SwiGLU 激活
SwiGLU：两路并行的线性变换，一路做 SiLU 激活，然后逐元素相乘 (图源: Reasoning from Scratch)
Python — FeedForward (SwiGLU)
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
# 两路并行投影：emb_dim → hidden_dim
self.fc1 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
self.fc2 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
# 降维投影：hidden_dim → emb_dim
self.fc3 = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False)
def forward(self, x):
x_fc1 = self.fc1(x) # "gate" 分支
x_fc2 = self.fc2(x) # "value" 分支
x = nn.functional.silu(x_fc1) * x_fc2 # SwiGLU = SiLU(gate) * value
return self.fc3(x) # 降维回 emb_dimSwiGLU 是 GLU (Gated Linear Unit) 家族的一员。相比 GPT-2 的单路 GELU，SwiGLU 有两路投影 + 门控：一路做 SiLU 激活当 “门控信号”，另一路当 “值”，两者逐元素相乘。代价是参数量增加（3 个矩阵 vs 2 个），但效果更好。
SiLU(x) = x × sigmoid(x)，是 ReLU 的平滑替代品 (图源: Reasoning from Scratch)
4. RoPE：旋转位置编码
GPT-2 用绝对位置嵌入（一个 position → 一个固定向量），Qwen3 用旋转位置编码 (RoPE)：通过在 attention 的 Q/K 向量上应用旋转变换来编码位置信息。
Python — RoPE 核心
def compute_rope_params(head_dim, theta_base=10_000, context_length=4096):
"""预计算每个位置的 cos/sin 旋转角度"""
# 每两个维度一组，频率按指数衰减
inv_freq = 1.0 / (theta_base ** (
torch.arange(0, head_dim, 2).float() / head_dim
))
positions = torch.arange(context_length)
angles = positions[:, None] * inv_freq[None, :] # 外积
angles = torch.cat([angles, angles], dim=1) # 复制到 full head_dim
return torch.cos(angles), torch.sin(angles)
def apply_rope(x, cos, sin, offset=0):
"""将旋转变换应用到 Q 或 K 向量上"""
# x: (batch, num_heads, seq_len, head_dim)
x1 = x[..., :head_dim // 2] # 前半
x2 = x[..., head_dim // 2:] # 后半
rotated = torch.cat((-x2, x1), dim=-1) # 旋转 90°
return (x * cos) + (rotated * sin) # 应用旋转ℹ️ RoPE 的优势：(1) 天然支持相对位置 — 两个 token 的注意力分数只依赖它们的距离，不依赖绝对位置。(2) 支持 KV Cache — 新 token 只需计算自己位置的旋转，不用重新计算之前所有位置。(3) 可以通过调整 theta_base（Qwen3 用 1,000,000）来扩展上下文长度。
5. GQA：分组查询注意力
GQA：16 个 Q heads 共享 8 个 KV groups，每组 2 个 Q head 共享 1 对 K/V (图源: Reasoning from Scratch)
Python — GroupedQueryAttention（核心逻辑）
class GroupedQueryAttention(nn.Module):
def __init__(self, d_in, num_heads, num_kv_groups, head_dim, qk_norm=False):
super().__init__()
self.num_heads = num_heads # 16 个 Q heads
self.num_kv_groups = num_kv_groups # 8 个 KV groups
self.group_size = num_heads // num_kv_groups # 每组 2 个 Q head
self.W_query = nn.Linear(d_in, num_heads * head_dim, bias=False)
self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False)
self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False)
self.out_proj = nn.Linear(num_heads * head_dim, d_in, bias=False)
# QK Norm：对 Q 和 K 做 RMSNorm（Qwen3 特有）
if qk_norm:
self.q_norm = RMSNorm(head_dim)
self.k_norm = RMSNorm(head_dim)
def forward(self, x, mask, cos, sin, start_pos=0, cache=None):
queries = self.W_query(x) # → (b, seq, 16 * 128 = 2048)
keys = self.W_key(x) # → (b, seq, 8 * 128 = 1024)
values = self.W_value(x) # → (b, seq, 8 * 128 = 1024)
# Reshape + QK Norm + RoPE
queries = queries.view(b, seq, 16, 128).transpose(1, 2)
keys = keys.view(b, seq, 8, 128).transpose(1, 2)
queries = self.q_norm(queries) # Qwen3 特有
keys = self.k_norm(keys)
queries = apply_rope(queries, cos, sin, offset=start_pos)
keys = apply_rope(keys, cos, sin, offset=start_pos)
# KV Cache：拼接历史 K/V
if cache is not None:
keys = torch.cat([prev_k, keys], dim=2)
values = torch.cat([prev_v, values], dim=2)
# 扩展 K/V 以匹配 Q head 数量：8 groups → 16 heads
keys = keys.repeat_interleave(self.group_size, dim=1) # 8 → 16
values = values.repeat_interleave(self.group_size, dim=1)
# 标准缩放点积注意力
attn = (queries @ keys.T) / sqrt(head_dim)
attn = attn.masked_fill(mask, -inf)
attn = softmax(attn)
output = attn @ values
return self.out_proj(output)GQA 是 MHA（Multi-Head Attention）和 MQA（Multi-Query Attention）的折中方案。Qwen3 0.6B 有 16 个 Q heads 但只有 8 个 KV groups，每两个 Q head 共享一组 K/V。这使得 KV Cache 大小减半（省内存），推理速度更快，而效果和 full MHA 几乎无损。
ℹ️ QK Norm 是 Qwen3 的特色：在 RoPE 之前对 Q 和 K 做 RMSNorm。这可以防止 attention score 在训练后期因数值增长过大而导致 softmax 饱和。很多早期模型（包括 GPT-2/3、Llama 1/2）没有这个设计。
6. TransformerBlock 与 Qwen3Model
Pre-Norm Transformer Block：先 Norm 后 Attention/FFN，外加残差连接 (图源: Reasoning from Scratch)
Python — TransformerBlock
class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = GroupedQueryAttention(...)
self.ff = FeedForward(cfg)
self.norm1 = RMSNorm(cfg["emb_dim"]) # Attention 前的 norm
self.norm2 = RMSNorm(cfg["emb_dim"]) # FFN 前的 norm
def forward(self, x, mask, cos, sin, start_pos=0, cache=None):
# Pre-Norm + Attention + 残差
shortcut = x
x = self.norm1(x)
x, next_cache = self.att(x, mask, cos, sin, start_pos, cache)
x = x + shortcut
# Pre-Norm + FFN + 残差
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = x + shortcut
return x, next_cacheQwen3Model 由 28 个 TransformerBlock 堆叠而成，加上输入的 token embedding、最后的 RMSNorm 和输出的 linear head：
Python — Qwen3Model（简化）
class Qwen3Model(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(151936, 1024) # token → 向量
self.trf_blocks = nn.ModuleList(
[TransformerBlock(cfg) for _ in range(28)]
)
self.final_norm = RMSNorm(1024)
self.out_head = nn.Linear(1024, 151936, bias=False) # 向量 → logits
# 预计算 RoPE 的 cos/sin（注册为 buffer，不参与训练）
cos, sin = compute_rope_params(head_dim=128, theta_base=1_000_000,
context_length=40960)
self.register_buffer("cos", cos)
self.register_buffer("sin", sin)
def forward(self, in_idx, cache=None):
x = self.tok_emb(in_idx)
# 构造 causal mask（上三角矩阵）
# 支持 KV Cache：mask 形状动态调整
for i, block in enumerate(self.trf_blocks):
x, new_cache = block(x, mask, self.cos, self.sin, ...)
if cache: cache.update(i, new_cache)
x = self.final_norm(x)
return self.out_head(x) # → (batch, seq, vocab_size) logits
Qwen3 0.6B 完整模型结构：Embedding → 28 × TransformerBlock → RMSNorm → Linear Head (图源: Reasoning from Scratch)
7. KVCache 与 Tokenizer
7.1 KVCache
KVCache 的实现非常简洁 — 就是一个按层索引的列表，每层存储一对 (K, V) tensor：
Python — KVCache
class KVCache:
def __init__(self, n_layers):
self.cache = [None] * n_layers # 28 层，每层一对 (K, V)
def get(self, layer_idx):
return self.cache[layer_idx]
def update(self, layer_idx, value):
self.cache[layer_idx] = value # value = (keys, values) tupleKV Cache 的作用在第2章已经详细讨论：生成第 n 个 token 时，前 n-1 个 token 的 K/V 不需要重新计算，直接从 cache 中取出拼接。这将 attention 的计算复杂度从 O(n²) 降到 O(n)。
7.2 Tokenizer
Qwen3 使用 BPE tokenizer，有 151,936 个 token。关键设计：
- Base model：
eos_token = <|endoftext|> - Reasoning model：
eos_token = <|im_end|>（支持 chat template） - 支持
<think>/</think>等 special tokens（token ID 151667/151668） encode()方法自动处理 special token 和普通文本的混合编码
Python — 使用模型生成文本
from reasoning_from_scratch.qwen3 import Qwen3Model, Qwen3Tokenizer, KVCache
model = Qwen3Model(QWEN_CONFIG_06_B)
model.load_state_dict(torch.load("qwen3-0.6B-base.pth"))
tokenizer = Qwen3Tokenizer("tokenizer-base.json")
prompt = "Explain large language models in a single sentence."
input_ids = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)
# 自回归生成（使用 KV Cache）
cache = KVCache(n_layers=28)
for token in generate_text_basic_stream_cache(model, input_ids, max_new_tokens=200):
print(tokenizer.decode(token.squeeze(0).tolist()), end="", flush=True)
# → "Large language models are artificial intelligence systems
# that can understand, generate, and process human language..."
# 速度：28 tokens/sec (CPU)8. 关键收获
- 现代 LLM 架构已高度标准化：RMSNorm + SwiGLU + RoPE + GQA 几乎是 2024–2025 年所有主流 LLM 的标配（Llama 3、Qwen3、Gemma 2、Mistral 等）。理解了 Qwen3 的架构，就理解了这一代 LLM 的共同基础。
- GQA 是效率与效果的最佳平衡：16 个 Q heads 共享 8 个 KV groups，KV Cache 大小减半，推理速度显著提升，效果几乎无损。这是从 MHA → MQA → GQA 的演进结果。
- RoPE 使 KV Cache 成为可能：绝对位置编码需要在每个位置上都有固定的 embedding，KV Cache 会导致位置信息错乱。RoPE 通过相对旋转自然支持增量 decode。
- QK Norm 是训练稳定性的保障：防止 attention score 在训练后期 “爆炸”，是 Qwen3 相比更早模型的一个改进。
- 0.6B 的模型已经能做很多事：28 层、1024 维、151K 词表 — 虽然参数量不大，但经过 GRPO 训练后 MATH-500 准确率可达 47.4%，证明了推理能力主要依赖后训练而非模型大小。
本文是 Build a Reasoning Model (From Scratch) (Sebastian Raschka) 的学习笔记。代码基于原书附录 C 的实现，有简化和中文注释。完整源码请参考 reasoning_from_scratch.qwen3 模块。所有配图版权归原作者所有。