一句话总结：model.py 不是一团凭空出现的代码。它只是把 Embedding、位置编码、Causal Self-Attention、FFN、Norm、残差连接和输出投影，按前面推导过的数据流接起来。

📦 2024 年原始演示仓库：github.com/waylandzhang/Transformer-from-scratch。仓库里的历史 demo 把模型、训练和生成放在同一个 model.py；本章是经过本轮审校的独立教学版本，两者不是逐行镜像。

18.1 先划清这个模型的边界

我们要实现的是一个小型 decoder-only Transformer：

Token IDs [B, T]
    ↓
Token Embedding + 固定 Sinusoidal Position Encoding
    ↓
N × Pre-Norm Transformer Block
    ├── LayerNorm → Causal Multi-Head Attention → Residual Add
    └── LayerNorm → FFN → Residual Add
    ↓
Final LayerNorm
    ↓
Tied Output Projection
    ↓
Logits [B, T, vocab_size]

为了把原理露出来，本章手写 Attention，暂时不加入 KV Cache、RoPE、Flash Attention、padding mask 或分布式训练。它支持等长、无 PAD 的 batch；带 padding 的 batch 还需要 Chapter 16 讲过的 key padding mask。

这个模型也不是 GPT-2 的逐项复刻：它采用 GPT 风格的 Pre-Norm decoder block，却保留了原始 Transformer 风格的固定正弦位置编码。这样做是教学选择，后面每个选择都会明说。

18.2 配置先替我们守住形状

from dataclasses import dataclass

@dataclass
class ModelConfig:
    vocab_size: int
    context_length: int = 256
    d_model: int = 512
    n_layers: int = 6
    n_heads: int = 8
    dropout: float = 0.1

    def __post_init__(self):
        if self.vocab_size <= 0 or self.context_length <= 0:
            raise ValueError("vocab_size and context_length must be positive")
        if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0:
            raise ValueError("d_model, n_layers, and n_heads must be positive")
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")

    @property
    def head_dim(self):
        return self.d_model // self.n_heads

head_dim = d_model // n_heads 只有在整除时才成立。旧代码默认它一定成立，却又让输出投影固定接收 d_model，一旦传入不整除的配置，错误会拖到很后面才暴露。

vocab_size 也应该来自 tokenizer 的词表大小，而不是“这份训练文本里出现过的最大 token ID”。后者会让模型是否支持某个合法 token 取决于样本里碰巧有没有出现它。

18.3 Feed Forward Network

每个位置独立通过同一个 FFN：

[B, T, d_model]
       ↓ Linear
[B, T, 4 × d_model]
       ↓ GELU
[B, T, 4 × d_model]
       ↓ Linear + Dropout
[B, T, d_model]

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_dim = 4 * config.d_model
        self.net = nn.Sequential(
            nn.Linear(config.d_model, hidden_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(hidden_dim, config.d_model),
            nn.Dropout(config.dropout),
        )

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

这里用 GELU，是为了与前面以 GPT-2 为参考的完整前向传播保持一致。现代模型也可能改成 SwiGLU，并让中间宽度不再等于简单的 4 × d_model；本章先实现经典两层 FFN。

18.4 先写一个看得清的单头 Attention

带 causal mask 的公式是：

\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_h}}+M\right)V

其中 $M_{ij}=0$ 表示可见， $M_{ij}=-\infty$ 表示第 $i$ 个 Query 不能看第 $j$ 个 Key。

class AttentionHead(nn.Module):
    def __init__(self, d_model, head_dim, context_length, dropout):
        super().__init__()
        self.head_dim = head_dim
        self.query = nn.Linear(d_model, head_dim, bias=False)
        self.key = nn.Linear(d_model, head_dim, bias=False)
        self.value = nn.Linear(d_model, head_dim, bias=False)
        self.attn_dropout = nn.Dropout(dropout)

        mask = torch.tril(
            torch.ones(context_length, context_length, dtype=torch.bool)
        )
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        _, T, _ = x.shape
        if T > self.causal_mask.size(0):
            raise ValueError("sequence is longer than context_length")

        q = self.query(x)                          # [B, T, head_dim]
        k = self.key(x)                            # [B, T, head_dim]
        v = self.value(x)                          # [B, T, head_dim]

        scores = q @ k.transpose(-2, -1)           # [B, T, T]
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:T, :T]
        scores = scores.masked_fill(~allowed, float("-inf"))

        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)
        return weights @ v                         # [B, T, head_dim]

这里把 mask 注册成 bool buffer：它不是参数，不接受梯度，但会跟着 model.to(device) 一起移动。persistent=False 还避免把这个可以重建的三角矩阵塞进 checkpoint。

Softmax 后再做 Dropout 时，训练中的 Attention 权重行和不必严格等于 1；inverted Dropout 保持的是期望值。调用 model.eval() 后 Dropout 关闭。

18.5 从“拆开的头”走到融合实现

18.5.1 教学版：每个头都是一个 Module

class TeachingMultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.heads = nn.ModuleList([
            AttentionHead(
                config.d_model,
                config.head_dim,
                config.context_length,
                config.dropout,
            )
            for _ in range(config.n_heads)
        ])
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        y = torch.cat([head(x) for head in self.heads], dim=-1)
        return self.resid_dropout(self.out_proj(y))

数学上，各个头没有先后依赖；但这段 Python 的 list comprehension 会分别调用每个 module。它容易阅读，却不等于 GPU 上最高效的一次融合计算。

18.5.2 实际使用版：一次投影，再 reshape

原论文定义了每个头各自的 $W_i^Q,W_i^K,W_i^V$ ，并没有规定 PyTorch 类必须写成“一头一个对象”。把这些矩阵沿输出维拼在一起，就能用一次大的矩阵乘法得到所有 Q、K、V；这是与论文公式代数等价的常见融合实现。

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.head_dim
        self.context_length = config.context_length

        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(
            config.context_length,
            config.context_length,
            dtype=torch.bool,
        )).view(1, 1, config.context_length, config.context_length)
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        if T > self.context_length:
            raise ValueError("sequence is longer than context_length")

        q, k, v = self.qkv(x).split(C, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        scores = q @ k.transpose(-2, -1)
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:, :, :T, :T]
        scores = scores.masked_fill(~allowed, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)

        y = weights @ v                              # [B, H, T, D]
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.out_proj(y))

形状变化是：

x                 [B, T, C]
q, k, v           [B, T, C] each
split heads       [B, H, T, D]   where C = H × D
scores / weights  [B, H, T, T]
weighted values   [B, H, T, D]
merge heads       [B, T, C]

生产代码还可以换成 PyTorch 的 scaled_dot_product_attention，让后端在条件合适时选择 Flash Attention 等 kernel。本章保留手算版本，是为了让 mask、Softmax 和形状都看得见。

18.6 Transformer Block

Pre-Norm Transformer Block 中 Attention、FFN 与两条残差路径

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

这是 Pre-Norm：Norm 在子层之前，残差支路保留未归一化的 x。GPT-2 使用 LayerNorm 的 Pre-Norm；LLaMA 也采用 Pre-Norm 拓扑，但换成了 RMSNorm。Pre-Norm 往往更容易优化深层网络，不代表它在所有指标上无条件优于 Post-Norm。

18.7 Token 与位置怎样进入模型

18.7.1 位置编码只建一次

旧版每次 forward() 都重新创建整张正弦表，还把设备写在配置字典里。这样既重复计算，又可能在 model.to(...) 后出现设备不一致。

def sinusoidal_positions(max_length, d_model):
    position = torch.arange(max_length, dtype=torch.float32).unsqueeze(1)
    frequency = torch.exp(
        torch.arange(0, d_model, 2, dtype=torch.float32)
        * (-math.log(10000.0) / d_model)
    )
    angles = position * frequency

    table = torch.zeros(max_length, d_model)
    table[:, 0::2] = torch.sin(angles)
    table[:, 1::2] = torch.cos(angles[:, :table[:, 1::2].shape[1]])
    return table.unsqueeze(0)  # [1, context_length, d_model]

最后一行的切片让奇数 d_model 也不会在 cos 赋值时报 shape mismatch。虽然常见配置通常是偶数，这个 helper 没必要偷偷依赖它。

在 Model.__init__ 中注册：

self.register_buffer(
    "position_encoding",
    sinusoidal_positions(config.context_length, config.d_model),
    persistent=False,
)

之后模型移到 CPU、CUDA 或其他 dtype 时，位置表会一起移动，不再需要 self.device。

18.7.2 Weight tying

self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
self.apply(self._init_weights)
self.lm_head.weight = self.token_embedding.weight

最后一行让输入词嵌入与输出投影共享同一份参数。不是复制一份数值，而是两个 module 真正引用同一个 Parameter。

PyTorch 文档说明 nn.Embedding 默认从 $\mathcal{N}(0,1)$ 初始化。这里又有 $\sqrt{d_{model}}$ 缩放和共享输出权重，如果照搬这个默认值，初始 logits 会异常大。因此在绑定权重之前，先统一初始化：

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

0.02 是这个教学模型明确采用的起点，不是所有架构的初始化定律；更深的 GPT 配方还会按残差层数缩放某些输出投影。

18.7.3 Forward 与 loss

def forward(self, idx, targets=None):
    if idx.ndim != 2:
        raise ValueError("idx must have shape [batch, time]")
    _, T = idx.shape
    if T == 0:
        raise ValueError("sequence must contain at least one token")
    if T > self.config.context_length:
        raise ValueError("sequence is longer than context_length")
    if targets is not None and targets.shape != idx.shape:
        raise ValueError("targets must have the same shape as idx")

    token = self.token_embedding(idx) * math.sqrt(self.config.d_model)
    position = self.position_encoding[:, :T, :]
    x = self.embedding_dropout(token + position)

    for block in self.blocks:
        x = block(x)
    x = self.final_norm(x)
    logits = self.lm_head(x)

    loss = None
    if targets is not None:
        loss = F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            targets.reshape(-1),
        )
    return logits, loss

原始 Transformer 会在与正弦位置编码相加前，把 token embedding 乘以 $\sqrt{d_{model}}$ ；这里明确采用同一个约定。

F.cross_entropy 直接接收未经 Softmax 的 logits和 token ID targets。训练数据要在 Chapter 19 中先做好一位错位：idx=[x_0,...,x_{T-1}]，targets=[x_1,...,x_T]。

18.8 生成函数也要写清边界

@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
    if idx.ndim != 2 or idx.size(1) == 0:
        raise ValueError("idx must have shape [batch, time] with time > 0")
    if max_new_tokens < 0:
        raise ValueError("max_new_tokens must be non-negative")
    if temperature < 0:
        raise ValueError("temperature must be non-negative")
    if top_k is not None and top_k <= 0:
        raise ValueError("top_k must be positive")

    for _ in range(max_new_tokens):
        idx_crop = idx[:, -self.config.context_length:]
        logits, _ = self(idx_crop)
        next_logits = logits[:, -1, :]

        if temperature == 0:
            idx_next = torch.argmax(next_logits, dim=-1, keepdim=True)
        else:
            next_logits = next_logits / temperature
            if top_k is None:
                probs = F.softmax(next_logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1)
            else:
                k = min(top_k, next_logits.size(-1))
                top_logits, top_indices = torch.topk(next_logits, k, dim=-1)
                top_probs = F.softmax(top_logits, dim=-1)
                choice = torch.multinomial(top_probs, num_samples=1)
                idx_next = top_indices.gather(-1, choice)

        idx = torch.cat((idx, idx_next), dim=1)
    return idx

三处边界不能含糊：

temperature == 0 要显式走 greedy，不能真的除以 0；
这里直接在 Top-K 的 K 个 logits 内采样，因此即使边界分数并列，也不会意外保留超过 K 个索引；
@torch.no_grad() 关闭梯度记录，但不会关闭 Dropout。调用者生成前仍要执行 model.eval()。

这个教学函数每一步都会重新计算当前窗口，没有 KV Cache。裁掉最旧 token 后，固定绝对位置也从窗口内的 0 重新开始；这是这里明确选择的滑动窗口策略，不是所有服务系统的默认行为。

18.9 完整 `model.py`

import math
from dataclasses import dataclass

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


@dataclass
class ModelConfig:
    vocab_size: int
    context_length: int = 256
    d_model: int = 512
    n_layers: int = 6
    n_heads: int = 8
    dropout: float = 0.1

    def __post_init__(self):
        if self.vocab_size <= 0 or self.context_length <= 0:
            raise ValueError("vocab_size and context_length must be positive")
        if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0:
            raise ValueError("d_model, n_layers, and n_heads must be positive")
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")

    @property
    def head_dim(self):
        return self.d_model // self.n_heads


def sinusoidal_positions(max_length, d_model):
    position = torch.arange(max_length, dtype=torch.float32).unsqueeze(1)
    frequency = torch.exp(
        torch.arange(0, d_model, 2, dtype=torch.float32)
        * (-math.log(10000.0) / d_model)
    )
    angles = position * frequency
    table = torch.zeros(max_length, d_model)
    table[:, 0::2] = torch.sin(angles)
    table[:, 1::2] = torch.cos(angles[:, :table[:, 1::2].shape[1]])
    return table.unsqueeze(0)


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_dim = 4 * config.d_model
        self.net = nn.Sequential(
            nn.Linear(config.d_model, hidden_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(hidden_dim, config.d_model),
            nn.Dropout(config.dropout),
        )

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


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.head_dim
        self.context_length = config.context_length
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(
            config.context_length,
            config.context_length,
            dtype=torch.bool,
        )).view(1, 1, config.context_length, config.context_length)
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        if T > self.context_length:
            raise ValueError("sequence is longer than context_length")

        q, k, v = self.qkv(x).split(C, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        scores = q @ k.transpose(-2, -1)
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:, :, :T, :T]
        scores = scores.masked_fill(~allowed, 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)
        return self.resid_dropout(self.out_proj(y))


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x


class Model(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.embedding_dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.n_layers)
        ])
        self.final_norm = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.apply(self._init_weights)
        self.lm_head.weight = self.token_embedding.weight

        self.register_buffer(
            "position_encoding",
            sinusoidal_positions(config.context_length, config.d_model),
            persistent=False,
        )

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

    def forward(self, idx, targets=None):
        if idx.ndim != 2:
            raise ValueError("idx must have shape [batch, time]")
        _, T = idx.shape
        if T == 0:
            raise ValueError("sequence must contain at least one token")
        if T > self.config.context_length:
            raise ValueError("sequence is longer than context_length")
        if targets is not None and targets.shape != idx.shape:
            raise ValueError("targets must have the same shape as idx")

        token = self.token_embedding(idx) * math.sqrt(self.config.d_model)
        position = self.position_encoding[:, :T, :]
        x = self.embedding_dropout(token + position)

        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                targets.reshape(-1),
            )
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        if idx.ndim != 2 or idx.size(1) == 0:
            raise ValueError("idx must have shape [batch, time] with time > 0")
        if max_new_tokens < 0:
            raise ValueError("max_new_tokens must be non-negative")
        if temperature < 0:
            raise ValueError("temperature must be non-negative")
        if top_k is not None and top_k <= 0:
            raise ValueError("top_k must be positive")

        for _ in range(max_new_tokens):
            idx_crop = idx[:, -self.config.context_length:]
            logits, _ = self(idx_crop)
            next_logits = logits[:, -1, :]

            if temperature == 0:
                idx_next = torch.argmax(next_logits, dim=-1, keepdim=True)
            else:
                next_logits = next_logits / temperature
                if top_k is None:
                    probs = F.softmax(next_logits, dim=-1)
                    idx_next = torch.multinomial(probs, num_samples=1)
                else:
                    k = min(top_k, next_logits.size(-1))
                    top_logits, top_indices = torch.topk(next_logits, k, dim=-1)
                    top_probs = F.softmax(top_logits, dim=-1)
                    choice = torch.multinomial(top_probs, num_samples=1)
                    idx_next = top_indices.gather(-1, choice)

            idx = torch.cat((idx, idx_next), dim=1)
        return idx

18.10 不要只看“能运行”，还要验证语义

torch.manual_seed(7)
config = ModelConfig(
    vocab_size=32,
    context_length=8,
    d_model=12,
    n_layers=2,
    n_heads=3,
    dropout=0.0,
)
model = Model(config).eval()

idx = torch.tensor([[1, 2, 3, 4]])
targets = torch.tensor([[2, 3, 4, 5]])

with torch.inference_mode():
    logits, loss = model(idx, targets)
assert logits.shape == (1, 4, 32)
assert loss.ndim == 0

# 改动未来 token，不应改变前两个位置的 logits。
future_changed = torch.tensor([[1, 2, 9, 10]])
with torch.inference_mode():
    a, _ = model(idx)
    b, _ = model(future_changed)
torch.testing.assert_close(a[:, :2], b[:, :2])

# Greedy 生成 3 个 token 后，长度应增加 3。
generated = model.generate(idx[:, :2], 3, temperature=0)
assert generated.shape == (1, 5)

这三个断言分别验证输出契约、causal 语义和生成长度。只检查“没有报错”还不够。

18.11 参数量：绑定与不绑定差很多

假设 d_model=512、n_heads=8、n_layers=6、vocab_size=50,000：

组件	参数量
Tied token embedding / output	25,600,000
Attention，6 层	6,291,456
FFN，6 层（含 bias）	12,598,272
13 个 LayerNorm	13,312
Sinusoidal position encoding	0
总计	44,503,040

如果不绑定输出权重，还要再加 512 × 50,000 = 25,600,000，总量会变成 70,103,040。旧版一边把输出层算成独立参数，一边又用“约 7000 万”掩盖 bias 与 Norm；现在假设和精确结果都写清楚了。

本章交付物

学完这一章，你应该能够：

从形状约束开始配置模型，而不是等矩阵乘法报错
手写带 bool causal mask 的单头 Attention
解释“一头一个 module”和融合 QKV 为什么代数等价、执行方式不同
写出 Pre-Norm Block、固定位置 buffer 与真正的 weight tying
正确处理 temperature=0、Top-K 并列值和 model.eval()
用 causal invariance 断言验证模型没有偷看未来

下一章预告

模型现在可以前向传播，也可以用随机参数生成一串 token，但它还没有从数据中学到东西。

下一章我们写 train.py：构造错位后的输入与 targets，切分训练/验证数据，反向传播并让 optimizer 真正更新参数。