调用 from transformers import AutoModel 只要一行代码，但真正理解 Transformer 的人少之又少。多数开发者对 Attention 的认知停留在"Q、K、V 三个矩阵"的口号层面——面试时背得出公式，动手时写不出一个完整的 Block。
本文的目标很纯粹：不用任何高层封装库，只用 PyTorch 原生算子，从零实现一个完整的 Decoder-only Transformer 语言模型（GPT 架构，也是 GPT-4 / LLaMA / Qwen 的基础架构）。每一行代码对应论文中的一个公式。
一、架构对齐：Decoder-only Transformer
核心实现清单：Token Embedding → RoPE 位置编码 → Multi-Head Attention + Causal Mask → Feed-Forward → Pre-Norm 残差 → 自回归生成。
二、基础配置
```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_available() else 'cpu'
三、逐层实现
```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):
# 论文要求：embedding 乘以 sqrt(d_model) 做尺度对齐
return self.embedding(x) * math.sqrt(self.d_model)
3.1 Token Embedding
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(QKTdk)VAttention(Q,K,V)=softmax(dkQKT)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
# 1. 线性投影 → 切分多头
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)
# Q/K/V: [B, n_heads, S, d_head]
# 2. RoPE 注入位置信息（仅作用于 Q 和 K）
if rope is not None:
Q = rope(Q)
K = rope(K)
# 3. 注意力分数 = Q @ K^T / sqrt(d_head)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)
# 4. 因果掩码：上三角置 -inf，softmax 后概率为 0
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# 5. Softmax → 加权 V → 合并多头 → 输出投影
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
```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（先 Norm 再 Attention），训练稳定性远优于原始论文的 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):
# Sub-layer 1: Attention + 残差
x = x + self.dropout(self.attn(self.norm1(x), rope=rope, mask=mask))
# Sub-layer 2: FFN + 残差
x = x + self.dropout(self.ffn(self.norm2(x)))
return x
3.6 完整模型
```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_seqlen)
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)
# 权重共享：Embedding 和 LM Head 用同一套权重
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 训练单步
```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 自回归生成
```@torch.no_grad()
def generate(model, prompt_ids, max_new_tokens=100, temperature=0.8, top_k=50):
model.eval()
for _ in range(max_new_tokens):
context = prompt_ids[:, -config.max_seq_len:]
logits = model(context)[:, -1, :] / temperature
# Top-K 采样
v, _ = torch.topk(logits, top_k)
logits[logits < v[:, [-1]]] = float('-inf')
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
prompt_ids = torch.cat([prompt_ids, next_token], dim=1)
return prompt_ids
五、模型规模与训练数据工程
调整 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
工业级训练的真正瓶颈往往不在模型代码，而在数据吞吐——高质量语料的采集和清洗需要完整的数据管线支撑。以亿牛云的数据采集场景为例，从外部站点批量采集训练语料时，IP 限流和并发频率是首要约束：
```import requests
亿牛云 API 代理提取出口 IP
api_url = "http://ip.16yun.cn:817/myip/pl//?s=&u=&format=json"
proxy_info = requests.get(api_url, timeout=10).json()[0]
proxy = f"http://{proxy_info['ip']}:{proxy_info['port']}"
批量采集语料
resp = requests.get("https://corpus.example.com/batch",
proxies={"http": proxy, "https": proxy}, timeout=30)
```
采集返回 403 需检查白名单配置，返回 429 需降低请求频率。稳定的数据采集是 LLM 训练的第一道基础设施保障。
六、踩坑速查
陷阱
根因
解法
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·Kᵀ 算相似度 → softmax 归一化 → 加权 V，计算复杂度 O(n²·d)
Causal Mask：上三角填 -inf，实现自回归生成
Pre-Norm 残差：信息直通通道，让深层网络可训练
权重共享：Embedding 和 LM Head 共享权重，参数减半
当面试官问"Attention 的计算复杂度为什么是 O(n²)"时，你能立刻回答：因为 Q·Kᵀ 那一步矩阵乘法的维度是 [seq, d] × [d, seq]。这种理解层次，是调用 from_pretrained() 永远无法获得的。
