📚 课程目录第1章从零实现的预训练流程

1. 从零实现的预训练流程

任务说明

本案例使用PyTorch从零构建一个简化版模型（2层Transformer解码器、4头注意力、128维嵌入），在小型语料库上完成自回归语言模型预训练。通过训练损失下降和生成文本质量验证模型已学会基本的词间依赖关系。

环境依赖

Python 3.8+，PyTorch 1.12+，NumPy
无需预训练权重，纯随机初始化训练

完整实现代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import random

# ---------- 超参数 ----------
VOCAB_SIZE = 128      # 小型字符级词表
EMB_DIM = 128
N_LAYERS = 2
N_HEADS = 4
SEQ_LEN = 32
BATCH_SIZE = 16
EPOCHS = 200
LR = 3e-4

# ---------- 训练语料（小型英文句子） ----------
corpus = [
    "the cat sat on the mat",
    "the dog ran in the park",
    "birds fly in the sky",
    "fish swim in the sea",
    "the sun is bright today",
    "a star shines at night",
    "the moon glows above",
    "rain falls from clouds",
    "snow covers the ground",
    "wind blows through trees",
]

# 构建字符词表
chars = sorted(list(set("".join(corpus) + " ")))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
VOCAB_SIZE = len(chars)

def encode(text):
    return [stoi[c] for c in text]

def decode(ids):
    return "".join(itos[i] for i in ids)

# ---------- 简化版模型 ----------
class SimplifiedGPT2(nn.Module):
    def __init__(self):
        super().__init__()
        self._emb = nn.Embedding(VOCAB_SIZE, EMB_DIM)
        self.pos_emb = nn.Embedding(SEQ_LEN, EMB_DIM)
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=EMB_DIM, nhead=N_HEADS, dim_feedforward=256,
            dropout=0.1, batch_first=True, activation="gelu"
        )
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=N_LAYERS)
        self.ln_f = nn.LayerNorm(EMB_DIM)
        self.lm_head = nn.Linear(EMB_DIM, VOCAB_SIZE, bias=False)
        self._init_weights()

    def _init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, idx):
        B, T = idx.shape
        tok = self._emb(idx)                    # (B,T,C)
        pos = self.pos_emb(torch.arange(T, device=idx.device))
        x = tok + pos
        # 因果掩码：每个位置只能看到当前位置及之前
        causal_mask = torch.triu(torch.ones(T, T, device=idx.device), diagonal=1).bool()
        # TransformerDecoderLayer的tgt_mask为True表示屏蔽
        x = self.decoder(x, x, tgt_mask=causal_mask)
        x = self.ln_f(x)
        logits = self.lm_head(x)                     # (B,T,VOCAB)
        return logits

# ---------- 数据准备 ----------
def get_batch():
    batch = []
    for _ in range(BATCH_SIZE):
        text = random.choice(corpus)
        ids = encode(text)
        # 填充或截断到SEQ_LEN
        if len(ids) < SEQ_LEN:
            ids = ids + [stoi[" "]] * (SEQ_LEN - len(ids))
        else:
            ids = ids[:SEQ_LEN]
        batch.append(ids)
    x = torch.tensor(batch)
    y = torch.zeros_like(x)
    y[:, :-1] = x[:, 1:]
    y[:, -1] = stoi[" "]  # 最后一个位置预测空格（填充）
    return x, y

# ---------- 训练 ----------
model = SimplifiedGPT2()
optimizer = torch.optim.AdamW(model.parameters(), lr=LR)
loss_fn = nn.CrossEntropyLoss()

model.train()
for epoch in range(EPOCHS):
    x, y = get_batch()
    logits = model(x)                     # (B,T,VOCAB)
    loss = loss_fn(logits.view(-1, VOCAB_SIZE), y.view(-1))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (epoch + 1) % 50 == 0:
        print(f"Epoch {epoch+1}/{EPOCHS}, Loss: {loss.item():.4f}")

# ---------- 生成验证 ----------
def generate(prompt, max_new_tokens=20):
    model.eval()
    ids = encode(prompt)
    ids = ids[:SEQ_LEN]
    with torch.no_grad():
        for _ in range(max_new_tokens):
            if len(ids) >= SEQ_LEN:
                ids = ids[-(SEQ_LEN-1):]
            x = torch.tensor([ids])
            logits = model(x)
            next_logits = logits[0, -1, :]
            probs = F.softmax(next_logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1).item()
            ids.append(next_id)
            if itos[next_id] == " " and len(ids) > 3:
                break
    return decode(ids)

print("\n生成示例：")
print(generate("the cat"))
print(generate("birds"))

运行方式

python _pretrain.py

验证方法与通过标准

训练损失：初始损失约为log(VOCAB_SIZE)≈4.5，训练200轮后应降至1.5以下，说明模型已从随机猜测变为能预测下一个字符。
生成质量：输入"the cat"应生成类似"the cat sat on the mat"的合理续写；输入"birds"应生成包含"fly"或"sky"的句子。生成文本应基本符合英语词序（主谓宾结构），而非随机字符序列。
判定通过：当训练损失稳定低于2.0，且生成文本中大部分单词来自训练语料、词序合理时，即可判定预训练流程实现正确。

常见问题

损失不下降：检查学习率是否过大（建议3e-4），或语料是否过小导致模型欠拟合。可增加EPOCHS至300。
生成全是空格：填充符" "占比过高导致模型倾向预测空格。可减少填充比例或使用更长的语料句子。
因果掩码方向错误：PyTorch的TransformerDecoderLayer中，tgt_mask=True的位置会被屏蔽，因此上三角矩阵（屏蔽未来位置）是正确的。