手写Mini-GPT：从因果语言模型到PyTorch完整实现

上一篇文章从 Encoder-Decoder 结构出发，完整介绍了 Transformer 的 Self-Attention、Multi-Head Attention、位置编码、Mask、Encoder 和 Decoder。

这一次，我们不再调用 nn.Transformer，而是只使用 PyTorch 的基础层，从零手写一个能够训练和生成文本的 Mini-GPT。

它会包含 GPT 最核心的组件：

• 字符级 Tokenizer。
• Token Embedding 和可学习 Position Embedding。
• 手写 Multi-Head Causal Self-Attention。
• Pre-Norm Transformer Block。
• GELU 前馈网络。
• 自回归语言模型损失。
• AdamW、学习率 warmup、余弦衰减和梯度裁剪。
• Temperature、Top-k 与多项式采样。
• 权重绑定和 KV Cache 原理。

最终模型只有约 35 万个参数，可以在 CPU 上训练。它当然不是 ChatGPT，也不会拥有通用知识，但它足以把 GPT 内部最重要的数据流完整跑通。

一句话概括：

GPT 接收一段 token 前缀，通过带因果遮罩的 Transformer Block 预测每个位置的下一个 token，再把新 token 接回输入继续生成。

1. GPT 到底是什么

GPT 是 Generative Pre-trained Transformer 的缩写：

• Generative：通过预测和生成后续 token 建模文本。
• Pre-trained：先在大量文本上进行语言模型预训练。
• Transformer：核心网络使用 Transformer 架构。

本篇文章主要实现其中的“Generative Transformer”。因为我们只使用一小段教学语料，所以更准确地说，它是一个从零训练的 Mini-GPT，而不是经过大规模预训练的通用模型。

1.1 GPT 与上一篇完整 Transformer 的关系

原始 Transformer：

Encoder
  -> 产生 Memory
  -> Decoder 通过 Cross-Attention 读取 Memory

GPT：

只有 Decoder-only Block
  -> Causal Self-Attention
  -> 不需要 Encoder
  -> 不需要 Cross-Attention

这里有一个容易误解的地方：

GPT 的“Decoder-only”并不是把原始 Transformer Decoder 原封不动地拿出来。

原始 Decoder 每层有三个子层：

1. Masked Self-Attention。
2. 与 Encoder 交互的 Cross-Attention。
3. FFN。

GPT 没有 Encoder，因此删除 Cross-Attention，只保留：

1. Causal Self-Attention。
2. FFN。

2. GPT 的学习目标：预测下一个 token

给定 token 序列：

$$
x_1,x_2,\ldots,x_T
$$

自回归语言模型把整段序列的联合概率分解为：

$$
P(x_1,x_2,\ldots,x_T)
=\prod_{t=1}^{T}P(x_t\mid x_1,\ldots,x_{t-1})
$$

也就是：

看到前文                  预测
<开始>                    第 1 个 token
x1                        x2
x1 x2                     x3
x1 x2 x3                  x4
...
x1 x2 ... x(T-1)          xT

GPT 不需要人为提供“主题”“语法规则”或“下一词类别”。训练文本自身就能构造监督信号。

例如文本是：

语言模型

可以得到：

输入 x：[语, 言, 模]
目标 y：[言, 模, 型]

输入和目标只错开一个位置。

3. 为什么一次前向传播能训练多个位置

训练时，模型一次接收：

x = [语, 言, 模]

并同时输出三个位置的词表 logits：

位置 0：根据“语”预测“言”
位置 1：根据“语言”预测“模”
位置 2：根据“语言模”预测“型”

因果遮罩保证位置 0 看不到位置 1 和 2，位置 1 看不到位置 2。因此，虽然整个矩阵并行计算，每个位置仍然只使用合法的历史信息。

如果：

• batch size 为 $B$。
• 上下文长度为 $T$。
• 词表大小为 $V$。

那么：

输入 idx： [B, T]
目标 y：   [B, T]
logits：   [B, T, V]

一次前向传播实际上产生了 $B\times T$ 个 next-token 训练样本。

4. Tokenizer：模型如何看见文字

神经网络不能直接接收字符串。Tokenizer 需要建立：

文字 -> token -> token id

真实 GPT 通常使用子词或字节级 Tokenizer。为了把注意力放在模型结构上，这里使用字符级 Tokenizer。

4.1 构造字符词表

raw_text = """春天来了，花儿在风中开放。
夏天来了，蝉儿在树上歌唱。
秋天来了，果实在枝头成熟。
冬天来了，雪花在天空飞舞。
清晨，太阳从东方升起。
傍晚，月亮在云间出现。
小河穿过安静的村庄。
微风吹过绿色的田野。
我们读书，也观察世界。
我们思考，也记录问题。
机器学习从数据中寻找规律。
神经网络用参数表示知识。
注意力让词语彼此交换信息。
语言模型根据上文预测下一个字符。
"""

text = raw_text * 120
chars = sorted(set(text))

stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for ch, i in stoi.items()}

• stoi 表示 string to integer。
• itos 表示 integer to string。
• 本例实测词表大小为 123。

4.2 编码与解码

def encode(s):
    return [stoi[ch] for ch in s]


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

测试：

ids = encode("春天")
print(ids)
print(decode(ids))

字符级 Tokenizer 的优点：

• 代码简单。
• 词表小。
• 不容易遇到未知字符拆分问题。

缺点：

• 同样一句话需要更多 token。
• 单个 token 的语义较弱。
• 序列变长，Attention 成本上升。

Mini-GPT 使用字符级分词是为了教学，不代表它是实际大模型的最佳方案。

5. 把连续文本切成训练窗口

先把完整文本编码成一维张量：

import torch

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))

train_data = data[:n]
val_data = data[n:]

这里使用前 90% 训练，后 10% 验证。

定义最大上下文长度：

block_size = 64

从随机起点取 65 个连续 token：

t0 t1 t2 ... t63 t64

前 64 个作为输入，后 64 个作为目标：

x = t0 t1 t2 ... t63
y = t1 t2 t3 ... t64

def get_batch(split, batch_size=32):
    source = train_data if split == "train" else val_data
    starts = torch.randint(
        len(source) - config.block_size - 1,
        (batch_size,),
    )

    x = torch.stack([
        source[i:i + config.block_size]
        for i in starts
    ])
    y = torch.stack([
        source[i + 1:i + config.block_size + 1]
        for i in starts
    ])

    return x.to(device), y.to(device)

5.1 block_size 是什么

block_size 是模型一次最多能使用的上下文 token 数，也称 context length 或 context window。

本例中：

block_size = 64

表示模型预测下一个字符时，最多参考最近 64 个字符。

它不是 batch size：

参数	含义
batch_size	一次并行训练多少段文本
block_size	每段文本包含多少个 token
vocab_size	可以预测多少种 token
n_embd	每个 token 的向量维度

6. Mini-GPT 配置

使用 dataclass 集中管理结构参数：

from dataclasses import dataclass


@dataclass
class GPTConfig:
    vocab_size: int
    block_size: int = 64
    n_layer: int = 3
    n_head: int = 4
    n_embd: int = 96
    dropout: float = 0.1

这些参数分别控制：

参数	作用	本例
vocab_size	输出类别数	123
block_size	最大上下文长度	64
n_layer	GPT Block 数量	3
n_head	每层注意力头数	4
n_embd	隐藏维度	96
dropout	Dropout 概率	0.1

每个头的维度为：

$$
d_{head}=\frac{n_{embd}}{n_{head}}
=\frac{96}{4}=24
$$

因此 n_embd 必须能被 n_head 整除。

7. Token Embedding 与 Position Embedding

输入 idx 只是整数编号：

[B, T]

Token Embedding 把每个编号查表转换成向量：

$$
E_{token}=W_{te}[idx]
$$

self.token_embedding = nn.Embedding(
    config.vocab_size,
    config.n_embd,
)

输出形状：

[B, T, C]

其中 $C=n_{embd}$。

7.1 为什么还需要位置向量

Causal Mask 只能限制“能看哪些位置”，并不能单独提供每个位置的完整顺序表示。因此还要加入 Position Embedding：

self.position_embedding = nn.Embedding(
    config.block_size,
    config.n_embd,
)

pos = torch.arange(t, device=idx.device)

token_emb = self.token_embedding(idx)      # [B, T, C]
position_emb = self.position_embedding(pos) # [T, C]
x = token_emb + position_emb               # [B, T, C]

相加时 [T,C] 会在 batch 维广播成 [B,T,C]。

数学上：

$$
x_{b,t}=E_{token}(idx_{b,t})+E_{position}(t)
$$

上一篇文章使用了固定正弦位置编码。本篇采用 GPT 风格的可学习绝对位置向量，它们会与模型其他参数一起通过反向传播更新。

8. GPT 的因果自注意力

上一篇文章已经介绍了普通 Self-Attention：

$$
\operatorname{Attention}(Q,K,V)
=\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
$$

GPT 需要额外加入 Causal Mask：

$$
\operatorname{CausalAttention}(Q,K,V)
=\operatorname{softmax}\left(
\frac{QK^T}{\sqrt{d_k}}+M
\right)V
$$

其中：

$$
M_{ij}=
\begin{cases}
0,&j\le i\
-\infty,&j>i
\end{cases}
$$

含义是第 $i$ 个 Query：

• 可以查看自己和过去位置。
• 不能查看未来位置。

9. 手写多头因果自注意力

下面不使用 nn.MultiheadAttention，完整写出 QKV 投影、拆头、遮罩、Softmax、加权和合并头。

9.1 初始化层和 Mask

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


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()

        if config.n_embd % config.n_head != 0:
            raise ValueError("n_embd must be divisible by n_head")

        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head

        self.qkv = nn.Linear(
            config.n_embd,
            3 * config.n_embd,
        )
        self.proj = nn.Linear(
            config.n_embd,
            config.n_embd,
        )

        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(
            config.block_size,
            config.block_size,
        ))
        self.register_buffer(
            "causal_mask",
            mask.view(
                1, 1,
                config.block_size,
                config.block_size,
            ),
        )

下三角矩阵形如：

$$
\begin{bmatrix}
1&0&0&0\
1&1&0&0\
1&1&1&0\
1&1&1&1
\end{bmatrix}
$$

这里约定：

• 1 表示允许查看。
• 0 表示需要遮住。

causal_mask 不是可训练参数，但需要跟随模型移动到 GPU 并保存到 state_dict，所以使用 register_buffer。

9.2 一次性产生 Q、K、V

q, k, v = self.qkv(x).split(c, dim=-1)

这里使用一个大线性层，而不是三个独立线性层：

x:       [B, T, C]
qkv(x):  [B, T, 3C]
q:       [B, T, C]
k:       [B, T, C]
v:       [B, T, C]

一个 Linear(C, 3C) 在数学上等价于把三组投影矩阵拼在一起：

$$
[Q,K,V]=X[W_Q,W_K,W_V]+b
$$

这样写通常能减少小矩阵算子的调度开销。

9.3 拆成多个头

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)

形状变化：

[B, T, C]
-> [B, T, H, D]
-> [B, H, T, D]

其中：

$$
C=H\times D
$$

本例是：

C = 96
H = 4
D = 24

9.4 计算分数并应用 Mask

scores = q @ k.transpose(-2, -1)
scores = scores / math.sqrt(self.head_dim)

scores = scores.masked_fill(
    self.causal_mask[:, :, :t, :t] == 0,
    float("-inf"),
)

形状：

q:             [B, H, T, D]
k.transpose:   [B, H, D, T]
scores:        [B, H, T, T]
causal_mask:   [1, 1, T, T]

Mask 在 batch 和 head 维自动广播。

为什么填充 $-\infty$？因为：

$$
e^{-\infty}=0
$$

经过 Softmax 后，未来位置的注意力权重会变成 0。

9.5 加权求和并合并头

weights = F.softmax(scores, dim=-1)
weights = self.attn_dropout(weights)

y = weights @ v

y = y.transpose(1, 2).contiguous().view(b, t, c)
y = self.proj(y)
y = self.resid_dropout(y)

return y

形状变化：

weights:       [B, H, T, T]
v:             [B, H, T, D]
weights @ v:   [B, H, T, D]
transpose:     [B, T, H, D]
view:          [B, T, C]
output proj:   [B, T, C]

9.6 完整注意力类

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()

        if config.n_embd % config.n_head != 0:
            raise ValueError(
                "n_embd must be divisible by n_head"
            )

        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head

        self.qkv = nn.Linear(
            config.n_embd,
            3 * config.n_embd,
        )
        self.proj = nn.Linear(
            config.n_embd,
            config.n_embd,
        )
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(
            config.block_size,
            config.block_size,
        ))
        self.register_buffer(
            "causal_mask",
            mask.view(
                1, 1,
                config.block_size,
                config.block_size,
            ),
        )

    def forward(self, x):
        b, t, c = x.shape

        q, k, v = self.qkv(x).split(c, dim=-1)

        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)

        scores = q @ k.transpose(-2, -1)
        scores = scores / math.sqrt(self.head_dim)
        scores = scores.masked_fill(
            self.causal_mask[:, :, :t, :t] == 0,
            float("-inf"),
        )

        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)

        y = weights @ v
        y = y.transpose(1, 2).contiguous().view(b, t, c)

        y = self.proj(y)
        return self.resid_dropout(y)

10. 为什么 transpose 后要调用 contiguous

transpose 通常不会真的复制数据，而是改变张量的 stride，让同一块内存以不同顺序被解释。

例如：

y = y.transpose(1, 2)

此时逻辑形状虽然是 [B,T,H,D]，底层内存排列却可能不是连续的。直接用 view 合并维度可能报错或不符合预期。

y = y.transpose(1, 2).contiguous().view(b, t, c)

contiguous() 会把数据整理成当前维度顺序对应的连续内存，再交给 view。

也可以使用 reshape，它在必要时会复制数据。但在这里显式写出 contiguous().view(...) 更容易看清拆头与合并头的过程。

11. MLP：对每个位置进行非线性变换

Attention 负责 token 之间的信息交换，MLP 负责每个 token 内部的特征变换。

GPT 风格 MLP 通常先把维度扩展到 4 倍，再缩回原维度：

$$
\operatorname{MLP}(x)
=W_2\operatorname{GELU}(W_1x+b_1)+b_2
$$

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(
                config.n_embd,
                4 * config.n_embd,
            ),
            nn.GELU(),
            nn.Linear(
                4 * config.n_embd,
                config.n_embd,
            ),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)

本例形状变化：

[B, T, 96]
-> [B, T, 384]
-> GELU
-> [B, T, 96]

11.1 GELU 的数学原理

GELU 全称 Gaussian Error Linear Unit，常见定义为：

$$
\operatorname{GELU}(x)=x\Phi(x)
$$

$\Phi(x)$ 是标准正态分布的累积分布函数。

常用近似形式是：

$$
\operatorname{GELU}(x)
\approx\frac{x}{2}
\left[
1+\tanh\left(
\sqrt{\frac{2}{\pi}}
\left(x+0.044715x^3\right)
\right)
\right]
$$

与 ReLU 直接丢弃所有负数不同，GELU 会平滑地调节输入，负值附近仍可能保留少量信息。

12. Pre-Norm GPT Block

一个 Mini-GPT Block 包含：

1. LayerNorm。
2. Causal Self-Attention。
3. 残差连接。
4. LayerNorm。
5. MLP。
6. 残差连接。

代码非常短：

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 = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

数学形式：

$$
x’=x+\operatorname{Attention}(\operatorname{LN}(x))
$$

$$
y=x’+\operatorname{MLP}(\operatorname{LN}(x’))
$$

12.1 为什么叫 Pre-Norm

因为 LayerNorm 位于 Attention 或 MLP 之前：

x -> LayerNorm -> 子层 -> 与 x 相加

Post-Norm 则是：

x -> 子层 -> 与 x 相加 -> LayerNorm

上一篇原始 Transformer 主要展示了 Post-Norm。GPT-2 风格模型常使用 Pre-Norm，并在全部 Block 后再增加一次最终 LayerNorm。

12.2 残差连接为什么不能省

残差连接：

$$
y=x+F(x)
$$

让模型可以学习“在已有表示上增加修正”，也为深层网络提供更直接的梯度路径。

如果删除残差连接，每一层都必须完全重建表示，深层模型通常更难优化。

13. 组合完整 Mini-GPT

class MiniGPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.token_embedding = nn.Embedding(
            config.vocab_size,
            config.n_embd,
        )
        self.position_embedding = nn.Embedding(
            config.block_size,
            config.n_embd,
        )
        self.dropout = nn.Dropout(config.dropout)

        self.blocks = nn.ModuleList([
            Block(config)
            for _ in range(config.n_layer)
        ])

        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(
            config.n_embd,
            config.vocab_size,
            bias=False,
        )

        self.apply(self._init_weights)

        self.lm_head.weight = self.token_embedding.weight

nn.ModuleList 能让 PyTorch 正确注册每个 Block 的参数。普通 Python list 虽然能保存对象，但其中模块不会自动出现在 model.parameters() 和 state_dict() 中。

13.1 参数初始化

@staticmethod
def _init_weights(module):
    if isinstance(module, nn.Linear):
        nn.init.normal_(
            module.weight,
            mean=0.0,
            std=0.02,
        )
        if module.bias is not None:
            nn.init.zeros_(module.bias)

    elif isinstance(module, nn.Embedding):
        nn.init.normal_(
            module.weight,
            mean=0.0,
            std=0.02,
        )

这里使用均值为 0、标准差为 0.02 的正态分布初始化线性层和 Embedding。Mini-GPT 教学实现可以使用这个简单策略；更深模型还会针对残差分支进一步缩放初始化。

14. Weight Tying：输入词向量与输出分类器共享参数

输入层把 token id 映射成向量：

token_embedding.weight: [V, C]

输出层把隐藏向量映射回词表 logits：

lm_head.weight: [V, C]

两者形状相同，因此可以共享同一份参数：

self.lm_head.weight = self.token_embedding.weight

输入时：

$$
e_t=W_E[x_t]
$$

输出时：

$$
z_t=h_tW_E^T
$$

可以把它理解为：

• 输入阶段从词表向量表中取出当前 token。
• 输出阶段计算隐藏状态与所有词表向量的匹配程度。

共享参数的好处：

• 减少参数量。
• 让输入表示空间与输出分类空间保持联系。

本例若不共享，会额外增加：

$$
V\times C=123\times96=11808
$$

个参数。

15. 前向传播与语言模型头

def forward(self, idx, targets=None):
    b, t = idx.shape

    if t > self.config.block_size:
        raise ValueError(
            "sequence is longer than block_size"
        )

    pos = torch.arange(t, device=idx.device)

    token_emb = self.token_embedding(idx)
    position_emb = self.position_embedding(pos)
    x = token_emb + position_emb
    x = self.dropout(x)

    for block in self.blocks:
        x = block(x)

    x = self.ln_f(x)
    logits = self.lm_head(x)

    loss = None
    if targets is not None:
        loss = F.cross_entropy(
            logits.reshape(
                -1,
                self.config.vocab_size,
            ),
            targets.reshape(-1),
        )

    return logits, loss

完整形状变化：

idx                         [B, T]
token embedding             [B, T, C]
position embedding          [T, C]
相加并广播                   [B, T, C]
N 个 GPT Block              [B, T, C]
final LayerNorm             [B, T, C]
lm_head                     [B, T, V]

16. 交叉熵为什么就是语言模型损失

对第 $t$ 个位置，模型输出词表 logits：

$$
z_t=(z_{t,1},z_{t,2},\ldots,z_{t,V})
$$

Softmax 转为概率：

$$
p_{t,j}=\frac{e^{z_{t,j}}}{\sum_{k=1}^{V}e^{z_{t,k}}}
$$

如果正确下一个 token 是 $y_t$，该位置的负对数似然为：

$$
\mathcal{L}t=-\log p{t,y_t}
$$

整个 batch 的平均损失：

$$
\mathcal{L}
=-\frac{1}{BT}
\sum_{b=1}^{B}
\sum_{t=1}^{T}
\log P(y_{b,t}\mid x_{b,\le t})
$$

16.1 一个三分类手算例子

假设某位置 logits 为：

$$
z=[2.0,1.0,0.0]
$$

Softmax 分母：

$$
e^2+e^1+e^0
\approx7.389+2.718+1=11.107
$$

概率约为：

$$
p\approx[0.665,0.245,0.090]
$$

如果正确 token 是第 1 类：

$$
\mathcal{L}=-\log(0.665)\approx0.408
$$

如果正确 token 是第 3 类：

$$
\mathcal{L}=-\log(0.090)\approx2.408
$$

模型越确信错误答案，惩罚越大。

16.2 为什么不用手动 Softmax

F.cross_entropy 内部会结合 LogSoftmax 与负对数似然，直接接收 logits：

loss = F.cross_entropy(logits, targets)

错误写法：

loss = F.cross_entropy(
    F.softmax(logits, dim=-1),
    targets,
)

提前 Softmax 不仅重复计算，还会损失数值稳定性。

17. 困惑度 Perplexity 是什么

语言模型还经常报告困惑度：

$$
\operatorname{PPL}=e^{\mathcal{L}}
$$

如果平均交叉熵为：

$$
\mathcal{L}=2.0
$$

那么：

$$
\operatorname{PPL}=e^2\approx7.39
$$

可以粗略理解为，模型在每一步像是在约 7.39 个同等可能候选中犹豫。

如果模型在 $V$ 个 token 上完全均匀随机预测：

$$
p=\frac{1}{V}
$$

则：

$$
\mathcal{L}=-\log\frac{1}{V}=\log V
$$

本例 $V=123$，随机模型的理论损失约为：

$$
\log(123)\approx4.812
$$

困惑度越低通常表示对验证文本预测得越好，但不同 Tokenizer、词表和数据集上的困惑度不宜直接横向比较。

18. Mini-GPT 有多少参数

本例配置：

V = 123
T = 64
C = 96
H = 4
L = 3

18.1 Embedding

Token Embedding：

$$
123\times96=11808
$$

Position Embedding：

$$
64\times96=6144
$$

18.2 每个 Block

QKV 投影：

$$
96\times288+288=27936
$$

输出投影：

$$
96\times96+96=9312
$$

两次 LayerNorm：

$$
2\times(96+96)=384
$$

MLP 第一层：

$$
96\times384+384=37248
$$

MLP 第二层：

$$
384\times96+96=36960
$$

每个 Block 合计：

$$
27936+9312+384+37248+36960=111840
$$

三个 Block：

$$
3\times111840=335520
$$

最终 LayerNorm：

$$
96+96=192
$$

由于 lm_head 与 Token Embedding 权重绑定，总参数量是：

$$
11808+6144+335520+192=353664
$$

代码验证：

parameter_count = sum(
    p.numel() for p in model.parameters()
)
print(parameter_count)  # 353664

19. 使用 AdamW 配置参数组

并非所有参数都适合权重衰减。一个实用的教学规则是：

• 二维及以上的权重矩阵使用 weight decay。
• 一维 bias 和 LayerNorm 参数不使用 weight decay。

decay = [
    p for p in model.parameters()
    if p.requires_grad and p.dim() >= 2
]

no_decay = [
    p for p in model.parameters()
    if p.requires_grad and p.dim() < 2
]

optimizer = torch.optim.AdamW(
    [
        {
            "params": decay,
            "weight_decay": 0.1,
        },
        {
            "params": no_decay,
            "weight_decay": 0.0,
        },
    ],
    lr=3e-3,
    betas=(0.9, 0.95),
)

为什么 LayerNorm 参数通常不衰减？它们主要负责缩放与平移归一化结果，对这些一维参数施加同样衰减往往没有必要。

按维数分组适用于这个结构规整的 Mini-GPT。复杂生产模型更适合根据模块类型和参数名显式分组，避免误分类特殊参数。

20. Warmup 与余弦学习率衰减

训练刚开始时，参数和 AdamW 的动量统计都还不稳定。直接使用最大学习率可能让更新过猛，因此先 warmup：

$$
\eta_t=\eta_{max}\frac{t+1}{T_{warmup}}
$$

warmup 结束后使用余弦衰减：

$$
r=\frac{t-T_{warmup}}
{T_{max}-T_{warmup}}
$$

$$
\eta_t=\eta_{min}
+\frac{1}{2}
(1+\cos(\pi r))
(\eta_{max}-\eta_{min})
$$

max_steps = 1000
warmup_steps = 100
max_lr = 3e-3
min_lr = 3e-4


def get_lr(step):
    if step < warmup_steps:
        return max_lr * (step + 1) / warmup_steps

    ratio = (
        (step - warmup_steps)
        / (max_steps - warmup_steps)
    )
    coeff = 0.5 * (
        1.0 + math.cos(math.pi * ratio)
    )

    return min_lr + coeff * (max_lr - min_lr)

训练每一步更新参数组学习率：

lr = get_lr(step)

for group in optimizer.param_groups:
    group["lr"] = lr

21. 完整训练循环

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

config = GPTConfig(vocab_size=len(chars))
model = MiniGPT(config).to(device)

model.train()

for step in range(max_steps):
    lr = get_lr(step)
    for group in optimizer.param_groups:
        group["lr"] = lr

    x, y = get_batch("train")
    _, loss = model(x, y)

    optimizer.zero_grad(set_to_none=True)
    loss.backward()

    grad_norm = torch.nn.utils.clip_grad_norm_(
        model.parameters(),
        max_norm=1.0,
    )

    optimizer.step()

    if (step + 1) % 200 == 0:
        print(
            f"step={step + 1}, "
            f"loss={loss.item():.4f}, "
            f"lr={lr:.6f}, "
            f"grad={grad_norm:.3f}"
        )

21.1 为什么先清梯度再反向传播

PyTorch 默认累加梯度。若不清理：

loss.backward()

会把当前梯度加到上一步的 .grad 中。

标准单步更新顺序是：

forward
-> loss
-> zero_grad
-> backward
-> clip_grad_norm
-> optimizer.step

set_to_none=True 让梯度变回 None，通常比逐元素填 0 更节省一些操作。

21.2 梯度裁剪做了什么

如果所有参数梯度的总范数为：

$$
|g|_2
=\sqrt{\sum_i|g_i|_2^2}
$$

当它超过阈值 $c$ 时，统一缩放：

$$
g_i\leftarrow g_i\frac{c}{|g|_2}
$$

这样保持梯度方向大致不变，只限制单次更新幅度。

22. 正确评估训练集与验证集

评估时要关闭 Dropout 和梯度：

@torch.no_grad()
def estimate_loss(eval_iters=20):
    model.eval()
    result = {}

    for split in ("train", "val"):
        losses = []

        for _ in range(eval_iters):
            x, y = get_batch(split)
            _, loss = model(x, y)
            losses.append(loss.item())

        result[split] = sum(losses) / len(losses)

    model.train()
    return result

为什么平均多个 batch？单个随机窗口的损失波动可能较大，多次采样平均更能反映趋势。

实测 1000 步结果：

step=200,  train=0.0191, val=0.0187
step=400,  train=0.0156, val=0.0157
step=600,  train=0.0124, val=0.0130
step=800,  train=0.0118, val=0.0127
step=1000, train=0.0115, val=0.0115

这个损失看起来很低，主要因为教学语料只有 14 行，并被重复了 120 次。模型很容易记住其中模式。

而且这里是在重复之后再按位置切分，训练集和验证集包含相同句式。因此验证损失主要用于检查训练流程，并不能衡量对陌生文本的真正泛化。正式实验应先准备互不重复的训练语料与验证语料，再分别编码。

因此这个结果说明：

• 数据构造和目标错位是正确的。
• 因果注意力能够学习局部与长距离关系。
• 完整训练链路已经跑通。

它不说明：

• 模型理解了开放世界知识。
• 模型能对陌生主题稳定泛化。
• 重复小语料是训练真实语言模型的合理方案。

23. 自回归生成的基本过程

训练完成后，给模型一个提示词：

春天

生成流程是：

春天       -> 预测“来”
春天来     -> 预测“了”
春天来了   -> 预测“，”
春天来了， -> 预测“花”
...

每次只使用最后一个位置的 logits：

logits = logits[:, -1, :]

因为前面位置对应的是“过去各位置的下一个 token 分布”，我们现在只关心当前完整前缀之后是什么。

24. Temperature、贪心与随机采样

模型输出 logits $z$ 后，可以通过 Temperature 调整分布：

$$
p_i=\frac{\exp(z_i/\tau)}
{\sum_j\exp(z_j/\tau)}
$$

其中 $\tau$ 是 temperature。

24.1 Temperature 小于 1

temperature = 0.7

logits 被放大，概率分布更尖锐，模型更偏向高概率 token，输出通常更稳定但更保守。

24.2 Temperature 大于 1

temperature = 1.3

logits 被缩小，概率分布更平坦，低概率 token 更容易被采到，输出更多样但更容易混乱。

24.3 Temperature 不能为 0

因为代码要计算：

logits = logits / temperature

如果希望完全确定性地选最大概率 token，应直接使用：

next_token = logits.argmax(dim=-1, keepdim=True)

而不是把 temperature 设为 0。

25. Top-k 采样

Top-k 只保留概率最高的 $k$ 个候选，其余 logits 设为 $-\infty$：

k = min(top_k, logits.size(-1))
threshold = torch.topk(
    logits,
    k,
).values[:, [-1]]

logits = logits.masked_fill(
    logits < threshold,
    float("-inf"),
)

再计算概率并采样：

probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(
    probs,
    num_samples=1,
)

Top-k 的作用是避免从词表长尾中采到极低概率 token。

策略	做法	特点
Greedy	每次选最大概率	稳定、确定，但可能重复
Temperature	调整分布尖锐程度	控制保守与随机程度
Top-k	只保留前 $k$ 个	去除低概率长尾
Top-p	保留累计概率达到 $p$ 的集合	候选数量动态变化

26. 完整生成函数

@torch.no_grad()
def generate(
    self,
    idx,
    max_new_tokens,
    temperature=1.0,
    top_k=None,
):
    self.eval()

    if temperature <= 0:
        raise ValueError("temperature must be greater than 0")

    for _ in range(max_new_tokens):
        idx_cond = idx[:, -self.config.block_size:]

        logits, _ = self(idx_cond)
        logits = logits[:, -1, :] / temperature

        if top_k is not None:
            k = min(top_k, logits.size(-1))
            threshold = torch.topk(
                logits,
                k,
            ).values[:, [-1]]

            logits = logits.masked_fill(
                logits < threshold,
                float("-inf"),
            )

        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(
            probs,
            num_samples=1,
        )

        idx = torch.cat((idx, next_token), dim=1)

    return idx

调用：

prompt = "春天"
context = torch.tensor(
    [encode(prompt)],
    dtype=torch.long,
    device=device,
)

result = model.generate(
    context,
    max_new_tokens=100,
    temperature=0.8,
    top_k=8,
)

print(decode(result[0].tolist()))

一次实测输出为：

春天来了，花儿在风中开放。
夏天来了，蝉儿在树上歌唱。
秋天来了，果实在枝头成熟。
冬天来了，雪花在天空飞舞。
清晨，太阳从东方升起。
傍晚，月亮在云间出现。
小河穿过安静的村庄。
微风吹过绿色的田野。

它基本复现了训练语料。这再次说明模型学会并记忆了语料模式，不应把它误解为开放式创作能力。

本例词表没有单独定义 <EOS>，因此使用 max_new_tokens 强制结束生成。需要让模型自主停止时，应把 <EOS> 加入词表和训练文本，并在采到它时退出循环。

26.1 为什么每次截取最后 block_size 个 token

位置 Embedding 只定义了：

0, 1, ..., block_size - 1

如果输入无限增长，既会超过位置表范围，也会让注意力矩阵持续增大。

idx_cond = idx[:, -self.config.block_size:]

表示模型始终只看最近的上下文窗口。更早的 token 会被遗忘。

27. 训练与生成为什么一个并行、一个串行

训练时完整目标已知：

输入：春 天 来 了
目标：天 来 了 ，

利用 Causal Mask，可以一次并行计算所有位置。

生成时未来 token 不存在：

先生成 token t
才能把它接回输入
再生成 token t+1

所以自回归生成在 token 维度上天然串行。

这不意味着生成内部没有并行：

• batch 中多个请求可以并行。
• 每层的大矩阵乘法可以并行。
• 多个注意力头可以并行。

但同一条序列的第 $t+1$ 个 token 必须等待第 $t$ 个 token 产生。

28. KV Cache：为什么生成不必重复计算全部历史

当前教学版 generate 每一步都把最近的完整上下文重新传入模型。

假设已经有：

x1 x2 x3 x4

生成 x5 时计算了前四个位置的 K、V。接下来生成 x6 时，x1 到 x4 的 K、V 没有变化，却又被计算一遍。

KV Cache 会保存每层历史 token 的 Key 和 Value：

第 1 层：K_cache¹、V_cache¹
第 2 层：K_cache²、V_cache²
...
第 L 层：K_cacheᴸ、V_cacheᴸ

新 token 到来时，只计算它自己的：

$$
q_t, k_t, v_t
$$

然后追加缓存：

$$
K_{1:t}=\operatorname{Concat}(K_{1:t-1},k_t)
$$

$$
V_{1:t}=\operatorname{Concat}(V_{1:t-1},v_t)
$$

当前 Query 与全部缓存 Key 计算：

$$
\operatorname{Attention}
(q_t,K_{1:t},V_{1:t})
$$

28.1 KV Cache 改变了什么

没有缓存：

每一步重新计算所有历史 token 的 Q、K、V 和 Block 输出

有缓存：

每一步主要计算新 token，并复用历史 K、V

KV Cache 显著减少重复计算，但会占用额外显存。缓存规模与以下量近似成正比：

$$
2\times L\times B\times H\times T\times D
$$

其中前面的 2 代表 Key 和 Value。

如果缓存元素使用 $s$ 字节，显存近似为：

$$
M_{KV}
\approx2LBHTDs
$$

因为 $H\times D=C$，也可以写成：

$$
M_{KV}\approx2LBTCs
$$

上下文越长、层数越多、batch 越大，KV Cache 越占显存。

28.2 为什么通常不缓存 Query

历史 Query 只在对应历史位置产生输出时使用。生成当前 token 时，只需要当前 Query 去匹配所有历史 Key，因此长期保存历史 Query 通常没有必要。

29. 用 PyTorch 高效 Attention 替换手写矩阵

为了教学，我们显式创建了 [B,H,T,T] 分数矩阵。理解原理后，可以使用：

from torch.nn import functional as F

y = F.scaled_dot_product_attention(
    q,
    k,
    v,
    attn_mask=None,
    dropout_p=(
        self.attn_dropout.p
        if self.training
        else 0.0
    ),
    is_causal=True,
)

它仍然计算同一个核心公式，但 PyTorch 会根据设备与输入选择更合适的实现。

需要注意两个细节：

1. 使用 is_causal=True 时，不要再同时传普通 attn_mask。
2. 函数式接口不会自动根据 model.eval() 关闭 Dropout，所以推理时要显式传 dropout_p=0.0。

Mask 的布尔语义也要看具体 API：本篇手写代码中 1/True 表示允许查看；上一篇 nn.Transformer 的布尔屏蔽 Mask 中 True 常表示禁止查看。不要凭记忆混用，应检查所调用函数的文档。

30. 如何验证因果 Mask 真的有效

只看训练损失下降还不够。可以构造一个针对未来信息泄漏的测试。

先取一段前缀：

prefix = x[:, :20]

只修改第 15 个位置之后的 token：

changed = prefix.clone()
changed[:, 15:] = torch.randint(
    0,
    config.vocab_size,
    changed[:, 15:].shape,
)

比较前 15 个位置的 logits：

model.eval()

base_logits, _ = model(prefix)
changed_logits, _ = model(changed)

max_diff = (
    base_logits[:, :15]
    - changed_logits[:, :15]
).abs().max()

print(max_diff)

实测结果：

tensor(0.)

这说明修改未来 token 不会影响过去位置的输出。

注意测试时必须调用 model.eval()。如果 Dropout 仍然开启，两次前向传播即使输入相同也可能有随机差异。

31. 常见错误

31.1 输入和目标没有错开一位

错误：

x = t0 t1 t2 t3
y = t0 t1 t2 t3

模型只需复制当前 token，不是在预测下一个 token。

正确：

x = t0 t1 t2 t3
y = t1 t2 t3 t4

31.2 没有 Causal Mask

模型训练时能直接看到未来答案，损失会异常漂亮，但生成时完全失效。

31.3 Softmax 维度写错

weights = F.softmax(scores, dim=-1)

最后一维是 Key 位置。每个 Query 对所有 Key 的权重之和应为 1。

31.4 缩放用了 n_embd

错误：

scores /= math.sqrt(config.n_embd)

正确：

scores /= math.sqrt(self.head_dim)

点积发生在每个头的 $D=head_dim$ 维空间中。

31.5 合并头时忘记转置

[B,H,T,D] 不能直接当成 [B,T,C]。必须先把 T 移回第二维：

y = y.transpose(1, 2).contiguous()
y = y.view(b, t, c)

31.6 使用普通 list 保存 Block

错误：

self.blocks = [Block(config) for _ in range(n_layer)]

正确：

self.blocks = nn.ModuleList([
    Block(config) for _ in range(n_layer)
])

31.7 block_size 与输入长度不匹配

输入长度超过位置 Embedding 范围会报错。训练数据和生成函数都必须遵守最大上下文长度。

31.8 训练时忘记 model.train()

Dropout 会保持关闭，训练行为与预期不一致。

31.9 生成时忘记 model.eval()

Dropout 会使同一前缀的 logits 额外波动，也浪费不必要的随机计算。

31.10 Temperature 设为 0

会发生除零。需要确定性生成时使用 argmax。

31.11 Top-k 大于词表大小

应先限制：

k = min(top_k, logits.size(-1))

31.12 把小语料低损失当成智能

重复数据极易记忆。要观察泛化，应使用独立且不重复的验证文本，并检查模型对未见组合的预测。

32. 完整模型代码如何组织

完整可运行脚本：/code/nlp-blog/mini_gpt.py

推荐按下面的依赖顺序放在一个脚本中：

1. import 与随机种子
2. 准备语料
3. encode / decode
4. train / val 切分
5. GPTConfig
6. get_batch
7. CausalSelfAttention
8. MLP
9. Block
10. MiniGPT
11. optimizer 与 get_lr
12. estimate_loss
13. 训练循环
14. generate
15. 保存模型

保存 checkpoint：

from dataclasses import asdict

torch.save(
    {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "config": asdict(config),
        "stoi": stoi,
        "itos": itos,
    },
    "mini_gpt.pt",
)

恢复：

checkpoint = torch.load(
    "mini_gpt.pt",
    map_location=device,
    weights_only=True,
)

config = GPTConfig(**checkpoint["config"])
model = MiniGPT(config).to(device)
model.load_state_dict(checkpoint["model"])

decay = [
    p for p in model.parameters()
    if p.requires_grad and p.dim() >= 2
]
no_decay = [
    p for p in model.parameters()
    if p.requires_grad and p.dim() < 2
]
optimizer = torch.optim.AdamW([
    {"params": decay, "weight_decay": 0.1},
    {"params": no_decay, "weight_decay": 0.0},
], lr=3e-3, betas=(0.9, 0.95))
optimizer.load_state_dict(checkpoint["optimizer"])
stoi = checkpoint["stoi"]
itos = checkpoint["itos"]

Tokenizer 词表必须与训练时一致，否则相同 token id 会被解释成不同字符。

33. Mini-GPT 与真实大模型还差什么

Mini-GPT 已经包含 GPT 的核心计算骨架，但与真实系统仍有巨大差距。

方面	本文 Mini-GPT	真实语言模型
Tokenizer	字符级	BPE、字节或子词 Tokenizer
参数量	35 万	数亿到更大规模
数据	少量重复句子	大规模清洗语料
上下文	64 token	通常远长于 64
位置编码	可学习绝对位置	也可能使用相对位置、RoPE 等
Attention	标准 MHA	还可能使用 MQA、GQA 等
训练	单机 CPU 教学	大规模分布式训练
精度	FP32	BF16、FP16 或混合精度
推理	重算上下文	KV Cache、批处理与高效内核
能力塑造	next-token 训练	还可能经过指令微调与偏好对齐

规模扩大不会改变最核心的目标：

$$
P(x_t\mid x_{<t})
$$

但数据、计算、训练稳定性和推理工程会变得复杂得多。

34. 接下来可以怎样继续学习

完成 Mini-GPT 后，可以按以下路线继续：

字符级 Mini-GPT
  -> 子词 Tokenizer
  -> Top-p 采样
  -> KV Cache 实现
  -> PyTorch SDPA / Flash Attention
  -> RoPE 旋转位置编码
  -> RMSNorm
  -> SwiGLU
  -> GQA / MQA
  -> 混合精度训练
  -> 梯度累积与分布式训练
  -> 指令微调

建议每次只替换一个组件，并比较：

• 参数量。
• 训练速度。
• 显存占用。
• 训练与验证损失。
• 生成质量。

这样才能知道改动真正影响了什么。

35. 总结

Mini-GPT 的完整数据流是：

原始文本
  -> Tokenizer
  -> token id
  -> 随机上下文窗口 x 与右移目标 y
  -> Token Embedding + Position Embedding
  -> 多层 Pre-Norm GPT Block
       -> Causal Multi-Head Self-Attention
       -> Residual Connection
       -> MLP + GELU
       -> Residual Connection
  -> Final LayerNorm
  -> LM Head
  -> 每个位置的词表 logits
  -> Cross Entropy
  -> Backpropagation + AdamW

生成流程是：

提示 token
  -> 取最后 block_size 个 token
  -> 模型预测最后位置的 logits
  -> Temperature / Top-k
  -> 采样下一个 token
  -> 拼接回输入
  -> 重复

需要真正掌握的核心包括：

• GPT 是不含 Cross-Attention 的 Decoder-only 自回归 Transformer。
• 语言模型通过右移一位的输入和目标学习 next-token prediction。
• Causal Mask 让并行训练不会泄漏未来信息。
• QKV 投影后要正确拆头、计算注意力并合并头。
• GPT Block 由 Causal Attention、MLP、LayerNorm 和残差连接组成。
• LM Head 把每个隐藏向量投影成词表 logits。
• 交叉熵等价于最小化正确 token 的负对数似然。
• Temperature 和 Top-k 只改变生成策略，不改变模型参数。
• KV Cache 复用历史 Key 和 Value，以显存换取生成速度。
• 小数据上的低损失很可能来自记忆，不能直接等同于泛化能力。

当你能解释下面三条公式之间的关系时，就已经抓住了 GPT 的数学主线：

$$
P(x_1,\ldots,x_T)
=\prod_{t=1}^{T}P(x_t\mid x_{<t})
$$

$$
\operatorname{CausalAttention}(Q,K,V)
=\operatorname{softmax}\left(
\frac{QK^T}{\sqrt{d_k}}+M
\right)V
$$

$$
\mathcal{L}
=-\frac{1}{BT}
\sum_{b,t}
\log P(y_{b,t}\mid x_{b,\le t})
$$

它们分别回答：

1. GPT 要建模什么概率。
2. GPT 怎样从历史 token 中提取上下文。
3. GPT 怎样通过训练让正确下一个 token 的概率变大。

参考资料

1. Alec Radford et al., Improving Language Understanding by Generative Pre-Training: https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf
2. Alec Radford et al., Language Models are Unsupervised Multitask Learners: https://cdn.openai.com/better-language-models/language-models.pdf
3. Ashish Vaswani et al., Attention Is All You Need: https://arxiv.org/abs/1706.03762
4. PyTorch Documentation, scaled_dot_product_attention: https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
5. PyTorch Documentation, CrossEntropyLoss: https://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
6. PyTorch Documentation, torch.multinomial: https://docs.pytorch.org/docs/stable/generated/torch.multinomial.html