![](data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" width="32" height="32"><text y="28" font-size="28">🤑</text></svg>)
02 模型结构代码解读
学习理念：通过手写 GPT2 和 Qwen3 的完整代码，理解 Transformer 组件在真实模型中的组装方式。GPT2 是经典 Decoder-only（MHA + 可学习位置编码），Qwen3 是现代 LLM 代表（GQA + RoPE + RMSNorm + SwiGLU）。两相对比就能看清架构演进。
海外对标：GPT2（OpenAI 2019）、Qwen3（阿里 2025）
本节 AI 替代率：~75% | 人工干预率：~25%
| 角色 | 能力范围 |
|---|---|
| 🤖 AI 擅长 | 生成模型代码骨架、解释前向传播维度变化 |
| 👤 人类需理解 | KV Cache 的 prefill/decode 两阶段逻辑、GQA 中 K/V 的 repeat_interleave 机制 |
📌 来源说明：以下代码来自
3.代码/fine_tune_proj/00_gpt2_model_structure.py和01_qwen3_model_structure.py，提取核心部分，省略测试代码。
📖 阅读优先级
| 等级 | 章节 | 说明 |
|---|---|---|
| 🟢 直接跳过 | 全部 | HuggingFace 一行 from_pretrained 就能用，除非你要改模型结构 |
一、GPT2：经典 Decoder-only
1.1 整体结构
python
class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
# 输入层：token embedding + 可学习位置 embedding
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.drop_emb = nn.Dropout(cfg["drop_rate"])
# Transformer Block 堆叠
self.trf_blocks = nn.ModuleList(
[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
)
self.final_norm = LayerNorm(cfg["emb_dim"])
# 输出层：投影到词表
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
self.current_pos = 0 # KV Cache 位置追踪
def forward(self, in_idx, use_cache=False):
tok_embeds = self.tok_emb(in_idx)
# 位置编码：use_cache 时用累加位置，否则从 0 开始
if use_cache:
pos_ids = torch.arange(self.current_pos, self.current_pos + seq_len)
self.current_pos += seq_len
else:
pos_ids = torch.arange(0, seq_len)
pos_embeds = self.pos_emb(pos_ids).unsqueeze(0)
x = tok_embeds + pos_embeds
for blk in self.trf_blocks:
x = blk(x, use_cache=use_cache)
x = self.final_norm(x)
logits = self.out_head(x)
return logits🔥 【P0 必须理解】
use_cache控制是 Prefill（全量计算）还是 Decode（增量计算）。self.current_pos追踪当前已缓存的 token 数量。
1.2 MultiHeadAttention + KV Cache
🔥 【P0 必须理解】 这是整个代码中最核心的部分。KV Cache 的三段逻辑：
cache_k is None→ Prefill 阶段，直接存入cache_k exists→ Decode 阶段，torch.cat拼接新 K/V 到缓存- mask 切片：从
ptr_current_pos开始，只对当前要计算的 token 做 mask
python
class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, num_heads):
self.W_query = nn.Linear(d_in, d_out)
self.W_key = nn.Linear(d_in, d_out)
self.W_value = nn.Linear(d_in, d_out)
self.out_proj = nn.Linear(d_out, d_out)
# 因果 mask（上三角），持久化在显存中
self.register_buffer("mask",
torch.triu(torch.ones(context_length, context_length), diagonal=1),
persistent=False)
# KV Cache 缓冲区
self.register_buffer("cache_k", None, persistent=False)
self.register_buffer("cache_v", None, persistent=False)
self.ptr_current_pos = 0
def forward(self, x, use_cache=False):
batch_size, num_tokens, d_in = x.shape
# 1. 生成 Q, K, V
keys_new = self.W_key(x).view(batch_size, num_tokens, self.num_heads, self.head_dim)
values_new = self.W_value(x).view(batch_size, num_tokens, self.num_heads, self.head_dim)
queries = self.W_query(x).view(batch_size, num_tokens, self.num_heads, self.head_dim)
# 2. KV Cache 管理
if use_cache:
if self.cache_k is None: # Prefill：全量缓存
self.cache_k, self.cache_v = keys_new, values_new
else: # Decode：拼接新 token
self.cache_k = torch.cat([self.cache_k, keys_new], dim=1)
self.cache_v = torch.cat([self.cache_v, values_new], dim=1)
keys, values = self.cache_k, self.cache_v
else:
keys, values = keys_new, values_new
# 3. 注意力计算（含因果 mask）
keys = keys.transpose(1, 2) # (b, n_heads, n_kv, head_dim)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)
attn_scores = queries @ keys.transpose(2, 3)
# mask 切片：use_cache 时从 ptr_current_pos 开始
if use_cache:
mask_bool = self.mask.bool()[
self.ptr_current_pos:self.ptr_current_pos + num_tokens_Q, :num_tokens_K]
self.ptr_current_pos += num_tokens_Q
else:
mask_bool = self.mask.bool()[:num_tokens_Q, :num_tokens_K]
attn_scores.masked_fill_(mask_bool, -torch.inf)
attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
context_vec = (attn_weights @ values).transpose(1, 2).contiguous()
context_vec = context_vec.view(batch_size, num_tokens, self.d_out)
return self.out_proj(context_vec)1.3 生成函数（Prefill + Decode 两阶段）
python
def generate_text_simple_cached(model, idx, max_new_tokens, use_cache=True):
model.eval()
with torch.no_grad():
if use_cache:
model.reset_kv_cache() # 清空旧缓存
logits = model(idx, use_cache=True) # Prefill：全量计算 prompt
for _ in range(max_new_tokens):
next_idx = logits[:, -1].argmax(dim=-1, keepdim=True)
idx = torch.cat([idx, next_idx], dim=1)
logits = model(next_idx, use_cache=True) # Decode：只输入一个 token
else:
for _ in range(max_new_tokens):
logits = model(idx, use_cache=False) # 每次全量计算
next_idx = logits[:, -1].argmax(dim=-1, keepdim=True)
idx = torch.cat([idx, next_idx], dim=1)
return idx二、Qwen3：现代 LLM 代表
2.1 Qwen3 配置
python
QWEN_CONFIG_06_B = {
"vocab_size": 151936, # 词表大小
"context_length": 40960, # 训练上下文长度
"emb_dim": 1024, # 嵌入维度
"n_heads": 16, # 注意力头数
"n_layers": 28, # 层数
"hidden_dim": 3072, # FFN 中间维度
"head_dim": 128, # GQA 头维度
"qk_norm": True, # 是否对 Q/K 额外归一化
"n_kv_groups": 8, # GQA 的 KV 组数（每组 2 个头）
"rope_base": 1_000_000.0, # RoPE 的 theta 基数
"dtype": torch.bfloat16,
}GPT2 vs Qwen3 关键差异：
- 位置编码：可学习 → RoPE
- 注意力：MHA → GQA（16 头 / 8 KV 组）
- FFN：GELU → SwiGLU（门控双分支）
- 归一化：LayerNorm → RMSNorm
- 额外：Qwen3 新增 QK Norm
2.2 Qwen3 整体结构
python
class Qwen3Model(nn.Module):
def __init__(self, cfg):
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"], dtype=cfg["dtype"])
self.trf_blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
self.final_norm = RMSNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False, dtype=cfg["dtype"])
# 预计算 RoPE 参数
head_dim = cfg["head_dim"] or (cfg["emb_dim"] // cfg["n_heads"])
cos, sin = compute_rope_params(head_dim, theta_base=cfg["rope_base"],
context_length=cfg["context_length"])
self.register_buffer("cos", cos, persistent=False)
self.register_buffer("sin", sin, persistent=False)
self.current_pos = 02.3 TransformerBlock（Qwen3）
python
class TransformerBlock(nn.Module):
def __init__(self, cfg):
self.att = GroupedQueryAttention(d_in=cfg["emb_dim"], num_heads=cfg["n_heads"],
head_dim=cfg["head_dim"], num_kv_groups=cfg["n_kv_groups"],
qk_norm=cfg["qk_norm"], dtype=cfg["dtype"])
self.ff = FeedForward(cfg) # SwiGLU
self.norm1 = RMSNorm(cfg["emb_dim"])
self.norm2 = RMSNorm(cfg["emb_dim"])
def forward(self, x, mask, cos, sin, start_pos=0, cache=None):
shortcut = x
x = self.norm1(x) # PreNorm
x, next_cache = self.att(x, mask, cos, sin, start_pos, cache)
x = x + shortcut
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = x + shortcut
return x, next_cache2.4 GQA 实现
🔥 【P0 必须理解】 GQA 的核心就是 K/V 的输出维度为
num_kv_groups * head_dim（而不是n_heads * head_dim），然后在计算前用repeat_interleave复制到每组内所有 Query 头。
python
class GroupedQueryAttention(nn.Module):
def __init__(self, d_in, num_heads, num_kv_groups, head_dim, qk_norm=False, dtype=None):
self.num_heads = num_heads
self.num_kv_groups = num_kv_groups
self.group_size = num_heads // num_kv_groups # 每组几个 Query 头
self.head_dim = head_dim
self.d_out = num_heads * head_dim
# Query：全量（每个头独立）
self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)
# Key/Value：缩量（每组共享）
self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)
self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)
def forward(self, x, mask, cos, sin, start_pos=0, cache=None):
queries = self.W_query(x).view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
keys_new = self.W_key(x).view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
values_new = self.W_value(x).view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
# QK Norm
if self.q_norm:
queries = self.q_norm(queries)
keys_new = self.k_norm(keys_new)
# RoPE：在 Q 和 K 上应用旋转位置编码
queries = apply_rope(queries, cos, sin, offset=start_pos)
keys_new = apply_rope(keys_new, cos, sin, offset=start_pos)
# KV Cache 管理
if cache is not None:
keys = torch.cat([cache[0], keys_new], dim=2)
values = torch.cat([cache[1], values_new], dim=2)
else:
keys, values = keys_new, values_new
next_cache = (keys, values)
# 核心：将 K/V 复制到每组内的所有 Query 头
keys = keys.repeat_interleave(self.group_size, dim=1)
values = values.repeat_interleave(self.group_size, dim=1)
# 注意力计算
attn_scores = queries @ keys.transpose(2, 3)
attn_scores = attn_scores.masked_fill(mask, -torch.inf)
attn_weights = torch.softmax(attn_scores / self.head_dim**0.5, dim=-1)
context = (attn_weights @ values).transpose(1, 2).reshape(b, num_tokens, self.d_out)
return self.out_proj(context), next_cache2.5 RoPE 实现
python
def compute_rope_params(head_dim, theta_base=10000, context_length=4096):
inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2).float() / head_dim))
positions = torch.arange(context_length)
angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)
angles = torch.cat([angles, angles], dim=1) # (context_length, head_dim)
return torch.cos(angles), torch.sin(angles)
def apply_rope(x, cos, sin, offset=0):
# x: (batch, n_heads, seq_len, head_dim)
x1 = x[..., :head_dim // 2] # 前半
x2 = x[..., head_dim // 2:] # 后半
cos = cos[offset:offset+seq_len].unsqueeze(0).unsqueeze(0)
sin = sin[offset:offset+seq_len].unsqueeze(0).unsqueeze(0)
rotated = torch.cat((-x2, x1), dim=-1) # 旋转
return (x * cos) + (rotated * sin)2.6 SwiGLU FFN
python
class FeedForward(nn.Module):
def __init__(self, cfg):
self.fc1 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
self.fc2 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False) # 门控
self.fc3 = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False)
def forward(self, x):
return self.fc3(nn.functional.silu(self.fc1(x)) * self.fc2(x))三、代码文件索引
| 优先级 | 文件 | 路径 | 核心内容 |
|---|---|---|---|
| 🔥 P0 | 00_gpt2_model_structure.py | 387 行 | GPT2 完整实现：MHA + KV Cache + 生成 |
| 🔥 P0 | 01_qwen3_model_structure.py | 443 行 | Qwen3 完整实现：GQA + RoPE + SwiGLU |
| 🟡 P1 | 04_single_gpu_optimizations_pseudocode.py | 568 行 | 训练优化综合：梯度累积/ZeRO/混合精度 |