第 8 章组装稠密模型¶

我们将第3章的技术组件进行升级

Post-LN       → Pre-LN
LayerNorm     → RMSNorm
GELU FFN      → SwiGLU
绝对位置表     → RoPE
MHA           → GQA
普通 Q/K      → QK-Norm 后再做 RoPE
整段重算生成   → GQA KV Cache

8.1 顺序¶

第 2 章的一层是两段残差：先注意力，再 FFN，外面套 Post-LN；绝对位置在进入第一层前加到词嵌入上。现代稠密模型保留「注意力 → FFN」这副骨架，层内顺序换成下面这条流水线。

步骤	流程	组件
1	\(x\) 留在残差主路；\(\operatorname{RMSNorm}(x)\) 进入注意力支路（Pre-LN）	第 5 章
2	投影 Q、K、V；Q 为 \(h\) 头，K/V 为 \(n_{\mathrm{kv}}\) 头（GQA：只改头数关系）	第 7 章
3	拆头后对 Q、K 做 RMSNorm（QK-Norm）	第 5 章
4	按当前位置对 Q、K 施加 RoPE	第 6 章
5	因果注意力 + 输出投影，加回残差	第 2、4 章
6	再 \(\operatorname{RMSNorm}\) → SwiGLU → 加回第二条残差	第 5 章

\[ H=X+\operatorname{Attn}(\operatorname{RMSNorm}(X)), \qquad Y=H+\operatorname{SwiGLU}(\operatorname{RMSNorm}(H)), \]

其中 \(\operatorname{Attn}\) 表示上表第 2–5 步整段

8.2 训练¶

第 3 章的 CharTokenizer、sample_batch 和 learning_rate 原样保留。这里只替换 GPTConfig、Attention、MLP、Block 与 GPT，再补一份带缓存的生成函数。

8.2.1 RMSNorm、RoPE 与 GQA¶

from dataclasses import dataclass
import math

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


@dataclass
class ModernConfig:
    vocab_size: int
    max_seq_len: int = 128
    n_layer: int = 4
    n_head: int = 4
    n_kv_head: int = 2
    n_embd: int = 128
    n_ff: int = 352          # 约 8/3 × n_embd，再取硬件友好的倍数
    rope_theta: float = 10000.0
    norm_eps: float = 1e-6


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 归一化统计用 fp32，输出再转回原 dtype。
        scale = x.float().pow(2).mean(dim=-1, keepdim=True)
        x_norm = x.float() * torch.rsqrt(scale + self.eps)
        return (x_norm * self.weight.float()).to(x.dtype)


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, max_seq_len: int, theta: float):
        super().__init__()
        if head_dim % 2:
            raise ValueError("RoPE 要求 head_dim 为偶数")
        inv_freq = 1.0 / (
            theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
        )
        pos = torch.arange(max_seq_len, dtype=torch.float32)
        freq = torch.outer(pos, inv_freq)
        self.register_buffer("cos", freq.cos(), persistent=False)
        self.register_buffer("sin", freq.sin(), persistent=False)

    def forward(self, x: torch.Tensor, start_pos: int) -> torch.Tensor:
        # x: [B, h, n, d_h]
        n = x.size(-2)
        if start_pos + n > self.cos.size(0):
            raise ValueError("序列长度超过 max_seq_len")
        cos = self.cos[start_pos : start_pos + n][None, None, :, :]
        sin = self.sin[start_pos : start_pos + n][None, None, :, :]
        x_even, x_odd = x[..., 0::2], x[..., 1::2]
        y_even = x_even * cos - x_odd * sin
        y_odd = x_even * sin + x_odd * cos
        return torch.stack((y_even, y_odd), dim=-1).flatten(-2).to(x.dtype)


KVCache = tuple[torch.Tensor, torch.Tensor]


class GQAAttention(nn.Module):
    def __init__(self, cfg: ModernConfig):
        super().__init__()
        if cfg.n_embd % cfg.n_head:
            raise ValueError("n_embd 必须被 n_head 整除")
        if cfg.n_head % cfg.n_kv_head:
            raise ValueError("n_head 必须被 n_kv_head 整除")

        self.n_head = cfg.n_head
        self.n_kv_head = cfg.n_kv_head
        self.group_size = cfg.n_head // cfg.n_kv_head
        self.d_h = cfg.n_embd // cfg.n_head

        self.q_proj = nn.Linear(cfg.n_embd, cfg.n_head * self.d_h, bias=False)
        self.k_proj = nn.Linear(cfg.n_embd, cfg.n_kv_head * self.d_h, bias=False)
        self.v_proj = nn.Linear(cfg.n_embd, cfg.n_kv_head * self.d_h, bias=False)
        self.o_proj = nn.Linear(cfg.n_head * self.d_h, cfg.n_embd, bias=False)

        # 每个头各自归一化最后一个维度，增益在头之间共享。
        self.q_norm = RMSNorm(self.d_h, cfg.norm_eps)
        self.k_norm = RMSNorm(self.d_h, cfg.norm_eps)
        self.rope = RotaryEmbedding(
            self.d_h, cfg.max_seq_len, cfg.rope_theta
        )

    def forward(
        self,
        x: torch.Tensor,
        past_kv: KVCache | None = None,
        use_cache: bool = False,
    ) -> tuple[torch.Tensor, KVCache | None]:
        B, n, _ = x.shape
        past_len = 0 if past_kv is None else past_kv[0].size(-2)

        q = self.q_proj(x).view(B, n, self.n_head, self.d_h).transpose(1, 2)
        k = self.k_proj(x).view(B, n, self.n_kv_head, self.d_h).transpose(1, 2)
        v = self.v_proj(x).view(B, n, self.n_kv_head, self.d_h).transpose(1, 2)

        q = self.rope(self.q_norm(q), past_len)
        k = self.rope(self.k_norm(k), past_len)

        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat((past_k, k), dim=-2)
            v = torch.cat((past_v, v), dim=-2)
        new_cache = (k, v) if use_cache else None

        # 参考实现显式复制；正式 kernel 会按头号映射，不真的复制 KV。
        k_attn = k.repeat_interleave(self.group_size, dim=1)
        v_attn = v.repeat_interleave(self.group_size, dim=1)

        score = (q @ k_attn.transpose(-2, -1)) * (self.d_h ** -0.5)
        key_pos = torch.arange(k.size(-2), device=x.device)
        query_pos = past_len + torch.arange(n, device=x.device)
        allowed = key_pos[None, :] <= query_pos[:, None]
        score = score.masked_fill(~allowed[None, None, :, :], float("-inf"))
        prob = score.float().softmax(dim=-1).to(q.dtype)
        y = prob @ v_attn
        y = y.transpose(1, 2).contiguous().view(B, n, -1)
        return self.o_proj(y), new_cache

这段代码同时覆盖三种路径：

训练：past_kv=None, use_cache=False，整段并行并使用因果掩码。
Prefill：输入完整 prompt，use_cache=True，返回每层的 K/V。
Decode：输入一个新 token 和旧缓存，past_len 决定 RoPE 位置，新的 K/V 追加到缓存尾部。

为了把形状写清楚，参考实现用 repeat_interleave 把 K/V 逻辑展开到全部 Q 头。仅做数学参考；正式 GQA kernel 直接让第 \(i\) 个 Q 头读取第 \(\lfloor i/g\rfloor\) 个 KV 头。

8.2.2 SwiGLU、Block 与整机¶

class SwiGLU(nn.Module):
    def __init__(self, cfg: ModernConfig):
        super().__init__()
        self.gate_proj = nn.Linear(cfg.n_embd, cfg.n_ff, bias=False)
        self.up_proj = nn.Linear(cfg.n_embd, cfg.n_ff, bias=False)
        self.down_proj = nn.Linear(cfg.n_ff, cfg.n_embd, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class ModernBlock(nn.Module):
    def __init__(self, cfg: ModernConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.n_embd, cfg.norm_eps)
        self.attn = GQAAttention(cfg)
        self.ffn_norm = RMSNorm(cfg.n_embd, cfg.norm_eps)
        self.ffn = SwiGLU(cfg)

    def forward(
        self,
        x: torch.Tensor,
        past_kv: KVCache | None = None,
        use_cache: bool = False,
    ) -> tuple[torch.Tensor, KVCache | None]:
        attn_out, new_kv = self.attn(
            self.attn_norm(x), past_kv=past_kv, use_cache=use_cache
        )
        x = x + attn_out
        x = x + self.ffn(self.ffn_norm(x))
        return x, new_kv


class ModernGPT(nn.Module):
    def __init__(self, cfg: ModernConfig):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.blocks = nn.ModuleList(ModernBlock(cfg) for _ in range(cfg.n_layer))
        self.norm = RMSNorm(cfg.n_embd, cfg.norm_eps)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        self.apply(self._init_weights)
        self._scale_residual_outputs()

    @staticmethod
    def _init_weights(m: nn.Module) -> None:
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def _scale_residual_outputs(self) -> None:
        scale = (2 * self.cfg.n_layer) ** -0.5
        for block in self.blocks:
            block.attn.o_proj.weight.data.mul_(scale)
            block.ffn.down_proj.weight.data.mul_(scale)

    def forward(
        self,
        idx: torch.Tensor,
        labels: torch.Tensor | None = None,
        past_key_values: list[KVCache] | None = None,
        use_cache: bool = False,
    ):
        B, n = idx.shape
        past_len = 0 if past_key_values is None else past_key_values[0][0].size(-2)
        if past_len + n > self.cfg.max_seq_len:
            raise ValueError("本次输入与缓存之和超过 max_seq_len")
        if labels is not None and past_key_values is not None:
            raise ValueError("训练 loss 不与历史 KV Cache 混用")

        x = self.tok_emb(idx)  # 不再加绝对位置表
        new_key_values = [] if use_cache else None
        for i, block in enumerate(self.blocks):
            past_kv = None if past_key_values is None else past_key_values[i]
            x, new_kv = block(x, past_kv=past_kv, use_cache=use_cache)
            if use_cache:
                new_key_values.append(new_kv)

        logits = self.lm_head(self.norm(x))
        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.size(-1)),
                labels[:, 1:].reshape(-1),
                ignore_index=-100,
            )
        return logits, loss, new_key_values

ModernGPT 没有 pos_emb，位置只进入每一层的 Q/K。o_proj 与 down_proj 在普通初始化后乘 \(1/\sqrt{2L}\)，对应第 5 章的残差分支缩放。输入嵌入和输出头这里没有绑权重，以便贴近 Qwen3、Llama 3.1 这档模型的常见配置；若小模型希望节省参数，仍可令 lm_head.weight = tok_emb.weight。

8.2.3 训练和推理¶

训练数据、batch 切片和优化器都沿用第 3 章，只有模型返回值多了一份缓存：

def train_modern(
    model: ModernGPT,
    token_ids: torch.Tensor,
    steps: int,
    batch_size: int,
    max_lr: float = 3e-4,
    warmup: int = 100,
    clip: float = 1.0,
) -> None:
    device = next(model.parameters()).device
    decay, no_decay = [], []
    for p in model.parameters():
        if p.requires_grad:
            (decay if p.ndim >= 2 else no_decay).append(p)
    opt = torch.optim.AdamW(
        [
            {"params": decay, "weight_decay": 0.1},
            {"params": no_decay, "weight_decay": 0.0},
        ],
        lr=max_lr,
        betas=(0.9, 0.95),
    )

    model.train()
    for step in range(steps):
        lr = learning_rate(step, warmup, max_lr)  # 复用第 3 章
        for group in opt.param_groups:
            group["lr"] = lr
        idx, labels = sample_batch(              # 复用第 3 章
            token_ids, batch_size, model.cfg.max_seq_len, device
        )
        _, loss, _ = model(idx, labels=labels)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        opt.step()
        if step % 50 == 0:
            print(
                f"step {step:4d}  lr {lr:.2e}  "
                f"loss {loss.item():.4f}  grad {grad_norm:.2f}"
            )

推理不能照第 3 章那样每步重算整个前缀。先做一次 Prefill，此后每步只送入刚生成的 token：

def sample_next(logits: torch.Tensor, temperature: float) -> torch.Tensor:
    if temperature <= 0:
        return logits.argmax(dim=-1, keepdim=True)
    prob = (logits / temperature).softmax(dim=-1)
    return torch.multinomial(prob, num_samples=1)


@torch.no_grad()
def generate_cached(
    model: ModernGPT,
    prompt: torch.Tensor,
    max_new: int,
    temperature: float = 0.8,
) -> torch.Tensor:
    model.eval()
    if prompt.size(1) + max_new > model.cfg.max_seq_len:
        raise ValueError("prompt + max_new 超过 max_seq_len")
    if max_new == 0:
        return prompt

    # Prefill：整段 prompt 只算一次，并建立每层 KV Cache。
    logits, _, cache = model(prompt, use_cache=True)
    nxt = sample_next(logits[:, -1, :], temperature)
    out = torch.cat((prompt, nxt), dim=1)

    # Decode：每轮只送上一步生成的一个 token。
    for _ in range(max_new - 1):
        logits, _, cache = model(
            nxt, past_key_values=cache, use_cache=True
        )
        nxt = sample_next(logits[:, -1, :], temperature)
        out = torch.cat((out, nxt), dim=1)
    return out

入口只需把第 3 章的 GPTConfig/GPT/train/generate 换掉：

cfg = ModernConfig(
    vocab_size=tok.vocab_size,
    max_seq_len=64,
    n_layer=4,
    n_head=4,
    n_kv_head=2,
    n_embd=128,
    n_ff=352,
)
model = ModernGPT(cfg).to(device)
train_modern(model, token_ids, steps=300, batch_size=16)
out = generate_cached(model, prompt, max_new=32, temperature=0.8)
print(tok.decode(out[0].tolist()))

这份实现首先用来核对语义与形状。训练时显式生成注意力分数矩阵，推理时用 torch.cat 扩展缓存，都会产生额外内存读写；工业实现会换成 FlashAttention、预分配或分页 KV Cache，以及融合后的 RMSNorm/RoPE/GQA kernel。

8.3 继续优化¶

代码已经有现代稠密模型的模样，但是和「现代大模型已经完成」距离甚远。Qwen3-8B、Llama 3.1 8B 的基本层结构确实与此同源。更强大的能力还依赖几十亿参数、万亿级训练 token、完整训练系统与后训练等。

优化FFN¶

稠密模型的每个 token 都经过每层全部注意力与 FFN 参数。扩大 \(L\)、\(d\) 和 \(d_{\mathrm{ff}}\) 可以装入更多知识，Decode 时也要在每一步读取这些权重，计算量仍约为 \(2P\)。GQA 减少的是 KV Cache，不会让 FFN 少算；

在第三部分我们从 FFN 入手，第 9 章把一块 FFN 换成许多专家，每个 token 只选少数几个，让总参数量继续增长而激活参数量不按同样比例增长。第 10 章用更高视角看待扩容这个问题，当然如果只了解MoE也已经足够。

减少计算复杂度¶

GQA 把写入 KV 缓存的头数从 \(h\) 降到 \(n_{\mathrm{kv}}\)，显存

\[ M_{\mathrm{KV}}=2BLn_{\mathrm{kv}}nd_hb_{\mathrm{kv}} \]

按 \(n_{\mathrm{kv}}/h\) 缩小，但是计算图仍是全序列注意力：Prefill 对长度 \(n\) 做 \(n\times n\) 配对，Decode 每步仍与长度为 \(n\) 的历史算注意力。GQA 不减少平方复杂度；在第四部分我们继续构建稀疏注意力，线性注意力来解决这个问题。

解决token串行生成¶

可以不让下一个 token 依赖上一个 token 么，或者生成更多的token

在第 15 章中，我们让多 Token 预测让一次前向同时猜更远的位置，投机解码让小模型先写，主模型成批验收，让一次昂贵的权重读取确认多个 token。

进入混合架构¶

短序列里，全注意力能精确访问任意历史位置；序列很长时，这份精确访问变得昂贵。现代大模型会把两者混排：便宜层负责连续状态更新，少量全局或稀疏层负责精确检索。

第 14 章开始我们讨论这种现代混合架构，之后在第 16 章我们会看到结合 MoE、线性/稀疏注意力、部分维 RoPE、多 Token 预测等技术的终极模型。

大规模训练¶

字符词表、随机窗口、单卡 AdamW 足以检查 loss 能否下降。正式预训练还需要 BPE、去重与质量过滤、文档边界和数据混合、长短序列调度、分布式并行、混合精度、重计算、容错存盘与稳定性监控。模型从 4 层放到几十层、参数从几十万放到几十亿，工程问题会从代码外面反过来限制形状和超参数。

在第二册我们继续了解工业级模型的实际训练系统架构。

核心问答¶

Q1：这份代码和第 3 章是什么关系？ 数据、标签错位、交叉熵、AdamW、warmup 与梯度裁剪都相同。替换的是模型内部；推理额外加入 KV Cache。最稳妥的检查方式是先用同一份小语料分别过拟合，确认两者 loss 都能下降，再比较参数量和缓存形状。

Q2：QK-Norm 和 RoPE 能不能对调？ 加载既有权重时必须遵守原实现顺序。代码按 Qwen3 的「投影 → 拆头 → QK-Norm → RoPE → 点积」来写。正交旋转保持整头 L2 范数，但逐维 RMSNorm 的可学习增益未必与二维旋转可交换，只看形状发现不了差异。

Q3：为什么训练函数不用 KV Cache？ 训练时整段目标已知，所有位置一次并行前向；缓存不会减少这次计算，还会妨碍正常的反向传播图。KV Cache 服务于右侧 token 尚未生成、下一步只能追加一个位置的推理阶段。

Q4：repeat_interleave 之后还省显存吗？ 参考代码只在缓存中保存 \(n_{\mathrm{kv}}\) 个头，但计算前临时做了逻辑复制，峰值显存并不理想。正式 GQA kernel 不物化副本，才能同时兑现缓存和带宽收益。

Q5：为什么不用 MLA，明明第 7 章说它压得更多？ GQA 是现代稠密模型里更通用、也更容易写对的基线；MLA 需要低秩 KV、解耦 RoPE 与矩阵吸收，代码量会遮住稠密组件的装配主线。第 16 章读旗舰配置时再处理这条分支。

Q6：做到这里能加载 Qwen3-8B 权重了吗？ 只能说主要组件已经对齐。真正加载还要逐项匹配层数、宽度、FFN 中间维、词表、RoPE 参数、RMSNorm epsilon、QK-Norm 参数形状、是否绑嵌入，以及权重键名。实现正确的判据是同一输入下逐层张量和最终 logits 对齐，不能靠生成文本「看起来像」来判断。

Q7：为什么 FFN 是下一章改造的重点？ SwiGLU 有 gate/up/down 三块大矩阵，通常占一层参数的大部分。稠密模型每个 token 都读完它们。MoE 把这块复制成专家并只激活少数几个，最直接地把容量与单 token 计算拆开。

第 8 章 组装稠密模型¶