版本：Next

从零实现 mini-LLM：完整可训练代码

本篇把前面所有零件拼成一个约 300 行、单文件、可训练可生成的 GPT。默认配置约 2500 万参数，在一张消费级 GPU 上训练几小时即可产出连贯的莎士比亚风格英文（用 tiny-shakespeare 数据集）；改小配置甚至可以在 CPU 上跑通全流程。

训练目标回顾

输入:  x = [t₀, t₁, ..., t_{T-1}]
标签:  y = [t₁, t₂, ..., t_T]        ← 右移一位
损失:  交叉熵(model(x), y)

一次 forward 同时计算 T 个位置的 next-token 损失，这就是训练比推理高效得多的原因（教师强制，teacher forcing）。

完整代码

# mini_llm.py
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

# ---------------- 配置 ----------------
class Config:
    vocab_size = 50257        # 与 GPT-2 BPE 对齐；自制分词器则改你的词表大小
    block_size = 256          # 上下文长度 T
    n_layer = 6
    n_head = 8
    d_model = 512
    dropout = 0.1
    device = "cuda" if torch.cuda.is_available() else "cpu"

# ---------------- 模型 ----------------
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight

def precompute_rope(d_head, max_len, base=10000.0):
    inv = base ** (-torch.arange(0, d_head, 2).float() / d_head)
    t = torch.arange(max_len).float()
    freqs = torch.outer(t, inv)
    return freqs.cos(), freqs.sin()

def apply_rope(x, cos, sin):
    # x: (B, H, T, D); cos/sin: (T, D/2) → 广播到 (1,1,T,D/2)
    cos, sin = cos[None, None], sin[None, None]
    x1, x2 = x[..., ::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)

class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_head = cfg.n_head
        self.d_head = cfg.d_model // cfg.n_head
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
        cos, sin = precompute_rope(self.d_head, cfg.block_size)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).chunk(3, -1)
        q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        q = apply_rope(q, self.cos[:T], self.sin[:T])
        k = apply_rope(k, self.cos[:T], self.sin[:T])
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)

class FFN(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        hidden = int(8 / 3 * cfg.d_model) // 256 * 256   # SwiGLU 惯例
        self.gate = nn.Linear(cfg.d_model, hidden, bias=False)
        self.up = nn.Linear(cfg.d_model, hidden, bias=False)
        self.down = nn.Linear(hidden, cfg.d_model, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg.d_model)
        self.attn = Attention(cfg)
        self.norm2 = RMSNorm(cfg.d_model)
        self.ffn = FFN(cfg)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

class MiniLLM(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer))
        self.norm_f = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight          # 权重绑定
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        x = self.tok_emb(idx)
        for blk in self.blocks:
            x = blk(x)
        logits = self.lm_head(self.norm_f(x))
        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, temperature=1.0, top_k=None):
        for _ in range(max_new):
            ctx = idx[:, -self.cfg.block_size:]            # 截断到上下文窗口
            logits, _ = self(ctx)
            logits = logits[:, -1, :] / temperature        # 只取最后位置
            if top_k is not None:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx

数据管道

用 GPT-2 的 BPE 把 tiny-shakespeare（约 40 万 token）编码后存为内存映射：

# prepare_data.py
import numpy as np, tiktoken, requests

url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
text = requests.get(url).text
enc = tiktoken.get_encoding("gpt2")
ids = np.array(enc.encode_ordinary(text), dtype=np.uint16)

n = int(0.9 * len(ids))
ids[:n].tofile("train.bin")
ids[n:].tofile("val.bin")
print(f"train {n} tokens, val {len(ids)-n} tokens")

# dataloader
def get_batch(split, cfg, batch_size):
    data = np.memmap(f"{split}.bin", dtype=np.uint16, mode="r")
    ix = torch.randint(len(data) - cfg.block_size, (batch_size,))
    x = torch.stack([torch.from_numpy(data[i:i+cfg.block_size].astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy(data[i+1:i+1+cfg.block_size].astype(np.int64)) for i in ix])
    return x.to(cfg.device), y.to(cfg.device)

训练循环

cfg = Config()
model = MiniLLM(cfg).to(cfg.device)
print(f"参数量: {sum(p.numel() for p in model.parameters())/1e6:.1f}M")

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4,
                              betas=(0.9, 0.95), weight_decay=0.1)
max_steps, warmup = 5000, 100

def lr_at(step):
    if step < warmup:
        return 3e-4 * step / warmup
    p = (step - warmup) / (max_steps - warmup)
    return 3e-4 * (0.1 + 0.45 * (1 + math.cos(math.pi * p)))   # cosine 到 10%

for step in range(max_steps):
    for g in optimizer.param_groups:
        g["lr"] = lr_at(step)
    x, y = get_batch("train", cfg, batch_size=32)
    with torch.autocast(device_type=cfg.device, dtype=torch.bfloat16,
                        enabled=(cfg.device == "cuda")):
        _, loss = model(x, y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

    if step % 250 == 0:
        model.eval()
        vx, vy = get_batch("val", cfg, 32)
        with torch.no_grad():
            _, vloss = model(vx, vy)
        print(f"step {step:5d}  train {loss.item():.4f}  val {vloss.item():.4f}")
        model.train()

要点逐一解释：

bf16 autocast：半精度省一半显存、提速 2-3 倍；bf16 比 fp16 动态范围大，不需要 GradScaler。
梯度裁剪 clip_grad_norm_(1.0)：防止偶发的梯度尖峰把参数打飞，LLM 训练标配。
cosine 调度 + warmup：见数学基础。
训练/验证 loss 一起记录：两者同步下降是健康信号；val 回升 = 过拟合。

采样生成

model.eval()
enc = tiktoken.get_encoding("gpt2")
start = torch.tensor([enc.encode("ROMEO:")], device=cfg.device)
out = model.generate(start, max_new=300, temperature=0.8, top_k=50)
print(enc.decode(out[0].tolist()))

训几千步后你会看到语法正确、带戏剧格式、偶尔有莎翁味道的文本——词语基本合法，长程逻辑仍然混乱，这正是 25M 参数 + 40 万 token 数据的合理上限。想要更聪明？这就是 Scaling Law 的地盘：加大 n_layer、d_model、数据量。

常见坑清单

症状	原因与对策
loss = NaN	lr 过大 / 未做梯度裁剪 / fp16 溢出（换 bf16）
loss 不降	targets 没右移 / 学习率过小 / 词表 id 越界
输出复读一句话	温度太低或 top_k=1；提高 temperature
显存爆炸	block_size 减半、batch 减小、加梯度累积
训练快但生成乱码	生成时忘记截断到 block_size，或 encode/decode 不一致

下一步

这个玩具框架离工业级还差三件事：数据（40 万 token → 数万亿 token 的工程）、训练（单机 → 千卡集群）、对齐（续写 → 对话）。依次见：

训练目标回顾​

完整代码​

数据管道​

训练循环​

采样生成​

常见坑清单​

下一步​