大模型训练入门:从零开始预训练一个小模型

理解大模型最好的方式是自己训练一个。本文基于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,反向传播,更新参数。理解了这个循环,后面的一切都是工程优化。