搭建NanoGPT

这一页我们把模型从零写出来。分两部分：数据预处理 prepare.py 和 模型定义 model.py。代码完整可运行，跟着敲一遍，前面每一块理论都会「落地」成你写的某一行。

Part 1：数据预处理（prepare.py）

字符级玩法：把整篇文本的每个不同字符当作一个 token，建立 字符 ↔ 整数 双射，再把全文编码成一串整数存成二进制。

# prepare.py
import os, pickle
import numpy as np

# ① 读入整篇纯文本(把任意语料放成 data/input.txt 即可)
with open("data/input.txt", "r", encoding="utf-8") as f:
    text = f.read()
print("语料长度:", len(text), "字符")

# ② 构建字符词表
chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}   # 字符 -> 整数
itos = {i: c for i, c in enumerate(chars)}   # 整数 -> 字符
encode = lambda s: [stoi[c] for c in s]       # 文本 -> id 列表
decode = lambda l: "".join(itos[i] for i in l)  # id 列表 -> 文本
print(f"词表大小(vocab_size): {vocab_size}")

# ③ 编码全文并划分训练/验证集(9:1)
data = np.array(encode(text), dtype=np.uint16)
n = int(0.9 * len(data))
train, val = data[:n], data[n:]

# ④ 存成二进制 + 保存词表
os.makedirs("data", exist_ok=True)
train.tofile("data/train.bin")
val.tofile("data/val.bin")
with open("data/meta.pkl", "wb") as f:
    pickle.dump(dict(vocab_size=vocab_size, itos=itos), f)
print("已生成 data/train.bin, data/val.bin, data/meta.pkl")

中文语料天然按字切，字符级完全适用；英文会退化到「字母+标点」级，也没问题。运行一次即可。

Part 2：模型定义（model.py）

2.1 配置

# model.py
import math, torch, torch.nn as nn
import torch.nn.functional as F

@torch.no_grad()
def set_seed(seed):
    torch.manual_seed(seed)

class Config:
    def __init__(self, **kw):
        # 一个适合 CPU 快跑的小配置, 可按需放大
        self.n_layer   = 4      # Transformer 层数
        self.n_head    = 4      # 注意力头数 (n_embd 需能被整除)
        self.n_embd    = 128    # = d_model, 主干宽度
        self.block_size = 256   # 最大上下文长度
        self.dropout   = 0.0
        self.bias      = False  # 线性层是否用偏置(现代做法关掉)
        self.__dict__.update(kw)

2.2 因果自注意力（对应「注意力机制」+「多头注意力」+「因果掩码」）

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        assert cfg.n_embd % cfg.n_head == 0
        self.n_head = cfg.n_head
        self.d_head = cfg.n_embd // cfg.n_head
        # 把 Q,K,V 打包成一次线性变换(等价于三个矩阵拼接)
        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=cfg.bias)
        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias)
        self.drop = nn.Dropout(cfg.dropout)
        # 预注册一个下三角因果掩码(buffer 不参与训练)
        mask = torch.full((cfg.block_size, cfg.block_size), float("-inf"))
        mask = torch.triu(mask, diagonal=1)     # 上三角(未来)置 -inf
        self.register_buffer("mask", mask)

    def forward(self, x):                        # x: [B, T, n_embd]
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)    # 各 [B,T,C]
        # 拆成多头: [B,T,C] -> [B,n_head,T,d_head]
        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)
        # 缩放点积注意力(PyTorch 内置 flash 实现), 传入因果掩码
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=self.mask[:T, :T],
                                           is_causal=False)
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # 合并多头
        return self.proj(self.drop(y))

注意 attn_mask 里我们手动把「未来」设为 -inf，这就是 注意力机制 里那张下三角掩码的真实落地。

2.3 前馈网络（对应「前馈网络与残差连接」）

class MLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        inner = 4 * cfg.n_embd            # 先升维 4 倍
        self.c_fc  = nn.Linear(cfg.n_embd, inner, bias=cfg.bias)
        self.c_proj = nn.Linear(inner, cfg.n_embd, bias=cfg.bias)  # 再降回
        self.act = nn.GELU()
        self.drop = nn.Dropout(cfg.dropout)
    def forward(self, x):
        return self.drop(self.c_proj(self.act(self.c_fc(x))))

2.4 一个 Transformer Block（Pre-Norm + 双残差）

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.ln_1 = nn.LayerNorm(cfg.n_embd)
        self.attn = CausalSelfAttention(cfg)
        self.ln_2 = nn.LayerNorm(cfg.n_embd)
        self.mlp  = MLP(cfg)
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # 残差①: 先 norm 再注意力
        x = x + self.mlp(self.ln_2(x))    # 残差②: 先 norm 再 FFN
        return x

正是 前馈网络与残差连接 里那张 Pre-Norm Block 图的代码化。

2.5 把一切装配成 GPT

class GPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)   # 词嵌入
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)   # 可学习位置嵌入
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.Sequential(*[Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.n_embd)                       # 最后的 norm
        self.head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)  # 输出到词表
        self.head.weight = self.tok_emb.weight     # 权重共享(见"词嵌入"页)
        self.apply(self._init)                      # 初始化

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None: nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok = self.tok_emb(idx)                                  # [B,T,C] 语义
        pos = self.pos_emb(torch.arange(T, device=idx.device))   # [T,C] 位置
        x = self.drop(tok + pos)                                 # 相加=合体
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.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=1.0, top_k=None):
        for _ in range(max_new_tokens):
            idx = idx[:, -self.cfg.block_size:]          # 截断到上下文内
            logits, _ = self(idx)                        # 取最后一步(缓存思想简化)
            logits = logits[:, -1, :] / temperature      # 只看下一个词的分布
            if top_k is not None:                        # 只保留概率最高的 k 个
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits = logits.masked_fill(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

tok_emb 与 head 共享权重（self.head.weight = self.tok_emb.weight），正是 词嵌入 提到的 tied embeddings。generate 里的「自回归拼接」就是 GPT架构 的串行生成过程。

2.6 实例化与自检

if __name__ == "__main__":
    import pickle
    meta = pickle.load(open("data/meta.pkl", "rb"))   # 先跑 prepare.py
    cfg = Config(vocab_size=meta["vocab_size"])
    torch.manual_seed(1337)
    model = GPT(cfg)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"参数量: {n_params/1e6:.2f}M, vocab_size={cfg.vocab_size}")

    x = torch.randint(0, cfg.vocab_size, (2, 32))     # 随便造 [B,T]
    logits, loss = model(x, targets=x)                 # 训练模式带损失
    print("logits:", logits.shape, "loss:", float(loss))  # loss≈ln(vocab)属正常

跑通这段、看到参数量和合理 loss，模型骨架就立起来了。

本页把哪些理论落了地

理论页                     →  本实战里的代码
词嵌入                    →  self.tok_emb
位置编码                  →  self.pos_emb (可学习)
注意力机制 + 因果掩码       →  CausalSelfAttention + mask
多头注意力                 →  view/transpose 成 n_head 份
前馈网络与残差/Pre-Norm     →  MLP + Block
GPT架构(自回归生成)         →  GPT.generate

下一步

模型有了，还不会「学习」。下一页写训练循环，让它从数据里进化 → 训练与调参。