理解大模型最好的方式是自己训练一个。本文基于nanoGPT,从Transformer原理到数据准备、训练循环,手把手预训练一个字符级语言模型。
为什么要自己训练
读论文和用API是一回事,自己跑通训练循环是另一回事。亲手训练一个小模型能帮你理解:
- Token化到底在做什么
- 注意力机制是怎么学到模式的
- Loss下降的过程中模型在学什么
- 过拟合长什么样
我们用Andrej Karpathy的nanoGPT作为起点,在Shakespeare数据集上训练一个字符级GPT。最终模型大约10M参数,一张消费级GPU几分钟就能训完。
Transformer快速回顾
GPT的核心是Decoder-only Transformer。关键组件:
自注意力(Self-Attention):
- 每个token生成Q、K、V三个向量
- 注意力分数 = softmax(QK^T / sqrt(d_k))
- 输出 = 注意力分数 × V
- Causal mask确保每个token只能看到之前的token
前馈网络(FFN):
- 两层线性变换 + 激活函数
- 负责存储"知识"
Layer Norm + Residual Connection:
- 稳定训练过程
- GPT用Pre-Norm(LayerNorm在attention/FFN之前)
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0
self.n_head = config.n_head
self.n_embd = config.n_embd
self.head_dim = config.n_embd // config.n_head
# Q, K, V 一起算,效率更高
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
def forward(self, x):
B, T, C = x.size()
# 计算Q, K, V
qkv = self.c_attn(x)
q, k, v = qkv.split(self.n_embd, dim=2)
# 拆分为多头
q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
# 使用PyTorch的scaled_dot_product_attention(自动使用FlashAttention)
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
y = y.transpose(1, 2).contiguous().view(B, T, C)
return self.c_proj(y)
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = nn.LayerNorm(config.n_embd)
self.attn = CausalSelfAttention(config)
self.ln_2 = nn.LayerNorm(config.n_embd)
self.mlp = nn.Sequential(
nn.Linear(config.n_embd, 4 * config.n_embd),
nn.GELU(),
nn.Linear(4 * config.n_embd, config.n_embd),
)
def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x
数据准备
用Shakespeare数据集——约100万字符的莎士比亚全集。字符级tokenizer最简单,适合理解原理。
import urllib.request
import os
def download_shakespeare():
url = "https://naw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
path = "shakespeare.txt"
if not os.path.exists(path):
urllib.request.urlretrieve(url, path)
with open(path, 'r') as f:
text = f.read()
return text
text = download_shakespeare()
print(f"数据集大小: {len(text)} 字符")
# 数据集大小: 1115394 字符
# 构建字符级词表
chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"词表大小: {vocab_size}")
# 词表大小: 65
# 编码/解码
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])
# 训练集/验证集分割(90%/10%)
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
print(f"训练集: {len(train_data)} tokens, 验证集: {len(val_data)} tokens")
数据加载器——随机采样固定长度的序列:
def get_batch(split, batch_size, block_size, device):
data = train_data if split == 'train' else val_data
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([data[i:i+block_size] for i in ix])
y = torch.stack([data[i+1:i+block_size+1] for i in ix])
return x.to(device), y.to(device)
block_size就是上下文窗口长度。输入是x[0:T],标签是x[1:T+1]——本质上是预测下一个字符。
模型定义
from dataclasses import dataclass
@dataclass
class GPTConfig:
vocab_size: int = 65
block_size: int = 256 # 上下文窗口
n_layer: int = 6 # Transformer层数
n_head: int = 6 # 注意力头数
n_embd: int = 384 # 嵌入维度
class GPT(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.transformer = nn.ModuleDict(dict(
wte = nn.Embedding(config.vocab_size, config.n_embd),
wpe = nn.Embedding(config.block_size, config.n_embd),
blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
ln_f = nn.LayerNorm(config.n_embd),
))
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
# 权重共享: embedding和output projection用同一套权重
self.transformer.wte.weight = self.lm_head.weight
# 参数数量
n_params = sum(p.numel() for p in self.parameters())
print(f"模型参数量: {n_params/1e6:.2f}M")
def forward(self, idx, targets=None):
B, T = idx.size()
pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
tok_emb = self.transformer.wte(idx) # (B, T, C)
pos_emb = self.transformer.wpe(pos) # (T, C)
x = tok_emb + pos_emb
for block in self.transformer.blocks:
x = block(x)
x = self.transformer.ln_f(x)
logits = self.lm_head(x) # (B, T, vocab_size)
loss = None
if targets is not None:
loss = F.cross_entropy(
logits.view(-1, logits.size(-1)),
targets.view(-1)
)
return logits, loss
@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=0.8, top_k=40):
for _ in range(max_new_tokens):
# 截断到block_size
idx_cond = idx[:, -self.config.block_size:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :] / temperature
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx
参数量约10.6M——比GPT-3的175B小一万多倍,但核心结构完全一样。
训练循环
def train():
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"使用设备: {device}")
config = GPTConfig()
model = GPT(config).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
# 学习率调度: warmup + cosine decay
max_iters = 5000
warmup_iters = 100
def get_lr(it):
if it < warmup_iters:
return 3e-4 * it / warmup_iters
decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters)
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
return 3e-4 * 0.1 + coeff * (3e-4 - 3e-4 * 0.1)
batch_size = 64
block_size = config.block_size
for iter_num in range(max_iters):
# 更新学习率
lr = get_lr(iter_num)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# 前向 + 反向
x, y = get_batch('train', batch_size, block_size, device)
logits, loss = model(x, y)
optimizer.zero_grad(set_to_none=True)
loss.backward()
# 梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
# 日志
if iter_num % 500 == 0:
model.eval()
val_losses = []
for _ in range(20):
xv, yv = get_batch('val', batch_size, block_size, device)
_, val_loss = model(xv, yv)
val_losses.append(val_loss.item())
avg_val_loss = sum(val_losses) / len(val_losses)
print(f"iter {iter_num}: train loss {loss.item():.4f}, val loss {avg_val_loss:.4f}, lr {lr:.6f}")
model.train()
return model
# model = train()
训练过程中你会看到loss从约4.2(随机猜测65个字符 → -ln(1/65)≈4.17)逐渐下降到约1.5左右。
生成效果
训练500步后的输出(模型刚学会基本英语单词拼写):
KING RICHARD:
The swo you beart, my lord, the pood
And not the sears of this grave...
训练5000步后(开始有莎士比亚的味道了):
KING RICHARD III:
What say you, uncle Buckingham? Have we
The power of all the state to make our peace?
BUCKINGHAM:
My lord, the commons is resolved upon.
That we shall have the crown, if once the duke
Be moved to give consent unto your wish.
虽然内容是编的,但格式(角色名+冒号+台词)、用词风格都像模像样。这就是语言模型在做的事——学习文本的统计模式。
常见问题和调试
Loss不下降:
- 检查学习率是否太大或太小
- 检查数据加载是否正确(打印几个batch看看)
- 确认target是否正确偏移了一位
过拟合(train loss远低于val loss):
- 增加dropout
- 减小模型大小
- 增加数据量
生成效果差:
- 训练更多步
- 调整temperature:太低会重复,太高会乱
- top_k采样比纯sampling效果好
从字符级到BPE
真正的大模型不用字符级tokenizer,而是用BPE(Byte Pair Encoding)。区别:
- 字符级:词表小(几十到几百),序列长,模型需要自己学习拼写
- BPE:词表大(32K-100K),序列短,常见词是一个token
实际训练中BPE效率高得多,但字符级更适合学习原理。理解了字符级的训练过程,切换到BPE只是换一个tokenizer,模型结构完全不变。
小结
从这个10M的小模型到GPT-4这样的万亿级模型,核心架构是一样的,差异在于:
- 数据量:从1MB到几TB
- 模型大小:从10M到万亿参数
- 训练技巧:混合精度、数据并行、模型并行、梯度累积
- 后训练:SFT、RLHF/DPO
但训练循环的本质没变:给模型看一段文本,让它预测下一个token,算loss,反向传播,更新参数。理解了这个循环,后面的一切都是工程优化。