只需一行代码 from transformers import AutoModel 就能加载预训练模型，但真正透彻理解 Transformer 底层原理的开发者却寥寥无几。许多人对 Attention 的认知止步于“Q、K、V 三个矩阵”的口诀——面试时能流畅背诵公式，真要动手实现一个完整的 Block 就寸步难行。本文的目标极为纯粹：不借助任何高层封装库，仅依靠 PyTorch 原生算子，从零手写一个完整的 Decoder-only Transformer 语言模型（即 GPT 架构，也是 GPT-4 / LLaMA / Qwen 的基础）。你看到的每一行代码都能在经典论文中找到对应的公式。通过这场实战，你将从“调包侠”蜕变为真正的模型建造者。
一、架构对齐：Decoder-only Transformer
下面是我们将要逐个攻破的核心模块清单：
- Token Embedding → 将离散 token 映射到连续向量空间
- RoPE 位置编码 → 注入相对位置信息，支持外推
- Multi-Head Attention + Causal Mask → 自回归注意力计算
- Feed-Forward Network → 非线性特征变换
- Pre-Norm 残差连接 → 稳定深层训练
- 自回归生成 → 逐个 token 预测，实现文本生成
二、基础配置与超参数
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class Config:
vocab_size = 10000 # 词表大小
d_model = 512 # 模型维度
n_heads = 8 # 注意力头数
n_layers = 6 # Block 层数
d_ff = 2048 # FFN 中间维度
max_seq_len = 512 # 最大序列长度
dropout = 0.1
pad_id = 0
config = Config()
device = 'cuda' if torch.cuda.is_a vailable() else 'cpu'
三、逐层实现详解
3.1 Token Embedding
将每个离散 token 映射为 d_model 维的连续向量，并按照 Transformer 原始论文的要求，乘以 √d_model 完成尺度对齐，使嵌入向量与后续位置编码的尺度匹配。
class TokenEmbedding(nn.Module):
def __init__(self, vocab_size, d_model):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.d_model = d_model
def forward(self, x):
return self.embedding(x) * math.sqrt(self.d_model)
3.2 旋转位置编码（RoPE）
RoPE 是 LLaMA、Qwen 等主流 LLM 广泛采用的位置编码方案。它通过旋转操作将位置信息编码到 Q 和 K 向量中，天然支持相对位置外推，效果远超传统绝对位置编码，也是位置编码领域的热门选择。
class RotaryPositionalEmbedding(nn.Module):
def __init__(self, d_model, max_seq_len):
super().__init__()
d_half = d_model // 2
# 频率：θ_i = 10000^(-2i/d)
freqs = 1.0 / (10000 ** (torch.arange(0, d_half).float() / d_half))
positions = torch.arange(max_seq_len).float()
angles = torch.outer(positions, freqs) # [seq_len, d_half]
self.register_buffer('cos', angles.cos())
self.register_buffer('sin', angles.sin())
def forward(self, x):
seq_len = x.size(1)
cos = self.cos[:seq_len].unsqueeze(0)
sin = self.sin[:seq_len].unsqueeze(0)
d_half = x.size(-1) // 2
x1, x2 = x[..., :d_half], x[..., d_half:]
# 旋转：(x1 + i*x2) * (cos + i*sin)
return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
3.3 Multi-Head Causal Self-Attention
核心公式铭记于心：Attention(Q,K,V) = softmax(QKT / √dk) V。下面代码严格复现了这一计算，并加入了因果掩码以实现自回归。
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.n_heads = n_heads
self.d_head = d_model // n_heads
self.W_q = nn.Linear(d_model, d_model, bias=False)
self.W_k = nn.Linear(d_model, d_model, bias=False)
self.W_v = nn.Linear(d_model, d_model, bias=False)
self.W_o = nn.Linear(d_model, d_model, bias=False)
self.dropout = nn.Dropout(dropout)
def forward(self, x, rope=None, mask=None):
B, S, _ = x.shape
# 线性投影 → 切分多头
Q = self.W_q(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
K = self.W_k(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
V = self.W_v(x).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
# 注入 RoPE（仅作用于 Q 和 K）
if rope is not None:
Q = rope(Q)
K = rope(K)
# 注意力分数
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)
# 因果掩码：上三角置 -inf
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn = self.dropout(F.softmax(scores, dim=-1))
out = torch.matmul(attn, V) # [B, n_heads, S, d_head]
out = out.transpose(1, 2).contiguous().view(B, S, -1)
return self.W_o(out)
3.4 Feed-Forward Network
经典的 MLP 结构，先升维再降维，通过 ReLU 引入非线性，增强模型表达能力。
class FeedForward(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
self.w1 = nn.Linear(d_model, d_ff)
self.w2 = nn.Linear(d_ff, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
return self.w2(self.dropout(F.relu(self.w1(x))))
3.5 Transformer Block（Pre-Norm 残差结构）
现代 LLM 全部采用 Pre-Norm 设计——先做 LayerNorm 再进入 Attention / FFN，训练稳定性远超原始论文中的 Post-Norm，这也是深度网络得以收敛的关键。
class TransformerBlock(nn.Module):
def __init__(self, config):
super().__init__()
self.attn = MultiHeadAttention(config.d_model, config.n_heads, config.dropout)
self.ffn = FeedForward(config.d_model, config.d_ff, config.dropout)
self.norm1 = nn.LayerNorm(config.d_model)
self.norm2 = nn.LayerNorm(config.d_model)
self.dropout = nn.Dropout(config.dropout)
def forward(self, x, rope=None, mask=None):
x = x + self.dropout(self.attn(self.norm1(x), rope=rope, mask=mask))
x = x + self.dropout(self.ffn(self.norm2(x)))
return x
3.6 完整模型组装
将所有组件拼接成完整的 Decoder-only 语言模型，特别地，Embedding 与 LM Head 共享权重（参数减半），这是 GPT 系列常用的参数节约技巧。
class TransformerLLM(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.token_embedding = TokenEmbedding(config.vocab_size, config.d_model)
self.rope = RotaryPositionalEmbedding(config.d_model, config.max_seq_len)
self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
self.final_norm = nn.LayerNorm(config.d_model)
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
# 权重共享
self.lm_head.weight = self.token_embedding.embedding.weight
self.apply(self._init_weights)
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(self, input_ids):
B, S = input_ids.shape
# 因果掩码：下三角矩阵
mask = torch.tril(torch.ones(S, S, device=input_ids.device)).unsqueeze(0).unsqueeze(0)
x = self.token_embedding(input_ids)
for block in self.blocks:
x = block(x, rope=self.rope, mask=mask)
x = self.final_norm(x)
return self.lm_head(x) # [B, S, vocab_size]
3.7 前向传播验证
用随机输入测试维度是否正确，确保模型搭建无误。
model = TransformerLLM(config).to(device)
n_params = sum(p.numel() for p in model.parameters())
print(f"参数量: {n_params/1e6:.1f}M")
dummy = torch.randint(1, config.vocab_size, (2, 64)).to(device)
logits = model(dummy)
print(f"输出: {logits.shape}") # [2, 64, 10000]
四、训练与推理实战
4.1 训练单步实现
采用标准交叉熵损失，注意输入与目标错位（shift），并使用梯度裁剪防止梯度爆炸。
def train_step(model, batch, optimizer, config):
model.train()
input_ids = batch[:, :-1].to(device) # 去掉最后一个 token
targets = batch[:, 1:].to(device) # 左移一位
logits = model(input_ids)
loss = F.cross_entropy(logits.reshape(-1, config.vocab_size),
targets.reshape(-1),
ignore_index=config.pad_id)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
return loss.item()
4.2 自回归生成流程
逐个 token 采样，每次取概率最高的 token 拼接，直到达到目标长度或遇到结束符。这是 GPT 推理的标准方式。
@torch.no_grad()
def generate(model, start_tokens, max_new_tokens, config):
model.eval()
input_ids = torch.tensor([start_tokens]).to(device)
for _ in range(max_new_tokens):
logits = model(input_ids[:, -config.max_seq_len:]) # 裁剪到最大长度
next_token = logits[:, -1, :].argmax(-1).unsqueeze(-1)
input_ids = torch.cat([input_ids, next_token], dim=-1)
return input_ids[0].tolist()
五、模型规模与数据工程
通过调整 Config 即可随意缩放模型：
| 配置 | d_model | n_heads | n_layers | 参数量 | 对标 |
|---|---|---|---|---|---|
| 教学版 | 512 | 8 | 6 | ~38M | — |
| GPT-2 Small | 768 | 12 | 12 | ~85M | GPT-2 |
| GPT-2 Medium | 1024 | 16 | 24 | ~350M | GPT-2 Medium |
工业级训练的真正瓶颈往往不在模型代码本身，而在数据吞吐——高质量语料的采集和清洗需要完整的数据管线支撑，这部分工作量常被低估。
六、踩坑速查手册
| 陷阱 | 根因 | 解法 |
|---|---|---|
| Loss 变 NaN | 梯度爆炸 | 梯度裁剪 max_norm=1.0 + 采用 Pre-Norm |
| Loss 不下降 | 学习率不匹配 | 前 200 步线性预热，之后使用余弦衰减调度 |
| 生成重复 | 采样策略单一 | 引入 Top-K 采样 + 温度缩放参数 |
| 显存溢出 | batch/seq 过大 | 减小 batch_size 或使用梯度累积技术 |
| 维度报错 | d_model % n_heads ≠ 0 | 确保 d_model 能被 n_heads 整除 |
七、总结与核心收获
手写一遍 Transformer 后，你将获得每个组件的深层理解：
- Embedding：离散 token → 连续空间，乘 √d_model 完成尺度对齐
- RoPE：用旋转编码将位置信息注入 Q/K，天然支持相对位置外推
- Attention：Q·KT 计算相似度 → softmax 归一化 → 加权融合 V，复杂度 O(n²·d)
- Causal Mask：上三角填入 -inf，保证自回归单向信息流
- Pre-Norm 残差：信息直通通道，让深层网络得以稳定训练
- 权重共享：Embedding 与 LM Head 共享权重，参数量减半
当面试官问“Attention 的计算复杂度为什么是 O(n²)”时，你能立刻回答：因为 Q·KT 矩阵乘法的维度是 [seq, d] × [d, seq]，每一对位置都要计算内积。这种底层视角，是仅仅调用 from_pretrained() 永远无法获得的。
