现代大语言模型的架构细节：从 RMSNorm 到 Loss 计算

之前陆续写过 Normalization 的对比、RL 基础，但一直没整理一份 LLM 架构的完整笔记。每次看 LLaMA、Qwen、DeepSeek 的源码，都在脑子里重新过一遍同样的东西，不如一次性写清楚。

这篇覆盖：整体架构 → RMSNorm → 自注意力 + GQA → Causal Mask → Flash Attention → Paged Attention → RoPE → 激活函数 → 残差连接 → MoE → 参数量估算 → Loss 计算（含辅助 Loss 与评估指标） → 权重绑定 → 训练配方 → 计算量估算。每个部分都写清楚为什么这么设计，不只是 what。

整体架构：一个 Transformer Block 里到底发生了什么

先看大局。一个典型的 Decoder-only LLM（LLaMA 系列、Qwen、DeepSeek）的架构长这样：

Token序列 [N, vocab]  →  Embedding [N, dim]
                                │
            ┌───────────────────┴───────────────────┐
            │          Transformer Block × L         │
            │                                       │
            │   input (hidden_states)               │
            │      │                                │
            │      ▼                                │
            │   ┌─────────────────────────┐         │
            │   │       RMSNorm            │  ← 第 1 个 Norm
            │   └───────────┬─────────────┘         │
            │               ▼                       │
            │   ┌─────────────────────────┐         │
            │   │    Self-Attention        │         │
            │   │  (QKV proj + RoPE +      │         │
            │   │   GQA/MHA + output proj) │         │
            │   └───────────┬─────────────┘         │
            │               │                       │
            │               ├──── 残差 ────┤         │
            │               ▼                       │
            │   ┌─────────────────────────┐         │
            │   │       RMSNorm            │  ← 第 2 个 Norm
            │   └───────────┬─────────────┘         │
            │               ▼                       │
            │   ┌─────────────────────────┐         │
            │   │     FFN / MoE            │         │
            │   │  (gate+up → SiLU → down) │         │
            │   └───────────┬─────────────┘         │
            │               │                       │
            │               ├──── 残差 ────┤         │
            │               ▼                       │
            │   output → 进入下一个 Block            │
            │                                       │
            └───────────────────┬───────────────────┘
                                ▼
                          RMSNorm (final)
                                │
                                ▼
                        LM Head [dim, vocab]
                                │
                                ▼
                        logits [N, vocab]

关键设计决策，一个个拆开说。

Pre-Norm vs Post-Norm

Pre-Norm：先归一化，再做子层操作，再加残差。

1
2
3

# Pre-Norm (现代 LLM 都用这个)
x = x + attention(rms_norm(x))
x = x + ffn(rms_norm(x))

Post-Norm（原版 Transformer）：先做子层操作，加残差，再归一化。

1
2
3

# Post-Norm (Vaswani 2017)
x = layer_norm(x + attention(x))
x = layer_norm(x + ffn(x))

为什么 Pre-Norm 赢了：

梯度流更畅通。Pre-Norm 下，残差连接的梯度可以直接经过 identity 路径往回传，不用穿过 LayerNorm。网络越深，这个差别越明显。
训练初期更稳定。子层（Attention、FFN）在训练刚开始时输出很不稳定，Pre-Norm 先把输入压到合理范围再喂给子层，子层不会因为输入尺度炸了而输出得更炸。
降低对 warmup 的依赖。原版 Transformer 需要很长的学习率 warmup 才能训练，Pre-Norm 让深层 Transformer 的梯度更稳定，对 warmup 的依赖显著减弱。但在大规模 LLM 训练中，warmup 仍然是常见标配。

RMSNorm：为什么 LLM 几乎都用它

Pre-Norm 解决了位置问题，但 Norm 本身的选择呢？LayerNorm vs RMSNorm 的详细对比我在之前那篇里写过了，这里只说 LLM 为什么选 RMSNorm。

公式回顾

LayerNorm：

RMSNorm：

RMSNorm 去掉了：

计算均值的求和
减均值的操作
偏移参数

为什么去掉均值

在 Pre-Norm 架构下，残差连接本身就在累积偏移。LayerNorm 减均值只是在这个 block 内部短暂”清零”，下一个 block 的残差又会把偏移加回来。所以减均值这个操作在 Pre-Norm 架构下收益不大。

可学习的可以在一定程度上补偿。RMSNorm 原论文的实验也验证了：加不加效果差不多。

计算开销

对于一个 hidden_dim = 4096 的向量，LayerNorm 比 RMSNorm 多做了：

一次求和（算均值）
D 次减法（减均值）
D 次加法（加）

在 LLaMA-7B 这种模型里，有个 Norm 层。每个 token 都要过这些 Norm。训练时 seq_len 可能是 4096，那每个 step 就是 64 × 4096 × 4096 次多余操作。省一点是一点。

所以现代 LLM 几乎所有 Norm 都换成了 RMSNorm。唯一的例外是少数模型（比如某些版本的 GPT）的最后一层用 LayerNorm，但也在慢慢消失。

自注意力 + GQA：KV cache 驱动的设计

先看标准的多头注意力（MHA）

MHA 就是做 h 组上面这个操作，每组在不同的子空间：

输入，分别投影到
每个 head：
$\text{head}i = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d{\text{head}}}}\right) V_i$
所有 head 拼接后过 output projection：

KV cache 是推理的瓶颈

自回归推理时，每生成一个新 token，所有之前 token 的 K 和 V 都已经算过了。为了避免重复计算，把它们存起来——这就是 KV cache。

KV cache 的大小：

比如 LLaMA-70B，n_layers=80，d_head=128，seq_len=4096，fp16：

MHA（n_kv_heads = n_q_heads = 64）：80 × 2 × 64 × 128 × 4096 × 2 ≈ 10 GB
这只是 KV cache，还不算模型权重

实际 LLaMA-70B 用的是 GQA（64 Q heads / 8 KV heads），按真实 KV heads 算，4K context 单 batch 下 KV cache 约 1.25 GB，远小于 140 GB 的模型权重。但长上下文、大 batch、高并发场景下，KV cache 会线性增长，成为主要显存瓶颈。减少 KV head 数量是直接降低 KV cache 的方法。

MHA → MQA → GQA

三个方案的 head 配置：

MHA (Multi-Head Attention):
  Q: H 个 head     K: H 个 head     V: H 个 head
  每对 Q_i K_i V_i 都独立

MQA (Multi-Query Attention, 2019):
  Q: H 个 head     K: 1 个 head     V: 1 个 head
  所有 Q head 共享同一组 K, V

GQA (Grouped-Query Attention, 2023):
  Q: H 个 head     K: G 个 head     V: G 个 head
  Q head 分组，每组共享一组 K, V

图解：

MHA:                    MQA:                    GQA (G=2):
Q: ■ ■ ■ ■              Q: ■ ■ ■ ■              Q: ■ ■ ■ ■
K: ■ ■ ■ ■              K: ■                    K: ■   ■
V: ■ ■ ■ ■              V: ■                    V: ■   ■

每个 Q 有独立 KV        所有 Q 共享 KV          每 2 个 Q 共享一组 KV

GQA 的 head 映射

设 n_q_heads = H，n_kv_heads = G，H 是 G 的整数倍。映射关系：

# 每个 kv head 对应 H/G 个 q head
kv_head_idx = q_head_idx * G // H

# 在模型中，K,V 投影矩阵的形状是 [dim, G * d_head]
# 而不是 MHA 的 [dim, H * d_head]
# 计算时，K,V 的 head 数少于 Q，需要扩展（repeat）

实际实现中，K,V 只投影出 G 个 head，然后在计算 attention 时通过 repeat_kv 把 K,V 扩展到和 Q 一样的 head 数：

def repeat_kv(kv, n_rep):  # n_rep = H // G
    # kv: [batch, G, seq, d_head]
    # → [batch, G, 1, seq, d_head]
    # → [batch, G, n_rep, seq, d_head]
    # → [batch, G*n_rep, seq, d_head] = [batch, H, seq, d_head]
    return kv.unsqueeze(2).expand(-1, -1, n_rep, -1, -1).reshape(...)

常见的 GQA 比例

模型	Q heads	KV heads	比例
LLaMA 3 8B	32	8	4:1
LLaMA 3 70B	64	8	8:1
Qwen 2.5 7B	28	4	7:1
Qwen 2.5 72B	64	8	8:1
DeepSeek-V2	128	128	1:1 (MLA, 另说)
Mistral 7B	32	8	4:1

规律：模型越大，GQA 比例越激进。小模型用 4:1，大模型可以用到 8:1。因为大模型的 Q 维度（dim）更大，减少 KV head 数能省更多显存，而质量损失相对小。

怎么选 GQA 比例

这是一个显存和质量的权衡：

比例太小（接近 MHA）：KV cache 大，但 attention 质量高
比例太大（接近 MQA）：KV cache 小，但 attention 表达能力受限，可能掉点
经验规律：4:1 到 8:1 之间效果不错，再大可能需要在训练数据或 loss 上看到明显差异

标准 GQA 的做法是：保持 d_head 不变，只减少 KV head 数 G。这样 K,V 总投影维度从变成，Q 和 K 的每个 head 维度相同，点积正常进行。如果在 dim 不变的情况下强行保持 G × d_head 不变（让单个 KV head 的 d_head 变大），会导致 Q 和 K 的 head_dim 不匹配，普通点积就对不上了。

Causal Mask：Decoder-only 的核心约束

Decoder-only 模型在预测第个 token 时，不能看到第及之后的 token。这就是 Causal Mask。

具体做法：在 Attention Score 计算完之后、Softmax 之前，给未来位置加一个，让 softmax 把它们的权重压成 0。

Mask 矩阵（以 seq_len=4 为例）：

    K0  K1  K2  K3
Q0  0  -∞  -∞  -∞
Q1  0   0  -∞  -∞
Q2  0   0   0  -∞
Q3  0   0   0   0

用 PyTorch 实现：

seq_len = 4
# torch.triu: 上三角全是 1，下三角全是 0
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
# tensor([[F, T, T, T],
#         [F, F, T, T],
#         [F, F, F, T],
#         [F, F, F, F]])

# attention 计算
scores = Q @ K.transpose(-2, -1) / math.sqrt(d_head)
scores = scores.masked_fill(mask, float('-inf'))
attn_weights = F.softmax(scores, dim=-1)   # softmax(-∞) → 0
output = attn_weights @ V

观察注意力权重矩阵：第行只有前列有非零值，后面全是 0。这保证了位置的 token 只看及之前的 token。

为什么必须单向：语言生成是天然的单向任务——说下一个词时，你确实不知道后面的词是什么。用 Causal Mask 训练，就是在模拟推理时的真实情况。如果用双向注意力训练但推理时自回归，会出现 train-inference mismatch，效果很差。

Causal Mask vs Padding Mask：两者不是一回事。Causal Mask 屏蔽未来 token（上三角区域置，允许区域是下三角），Padding Mask 屏蔽 batch 内不同长度序列的 padded 位置。实际代码中两个 mask 合并成 mask = causal_mask | padding_mask。

Flash Attention：让 Attention 省下 O(n²) 显存

标准 attention 的显存瓶颈不在输出，在中间结果—— 产生 [batch, heads, n, n] 的矩阵。seq_len=32k 时，单头就是个元素，fp16 占 2GB。

Flash Attention 的核心思路（Dao et al., 2022）：

Tiling（分块计算）：把 Q, K, V 切成小块，每次只算一小块 attention。用 online softmax 保证分块结果和全局计算完全一致。
Recompute（重算不存）：反向传播时不从显存读中间结果（因为没存），而是重新算一遍前向的 attention 矩阵。用额外计算换显存。

Online Softmax 的数学：

标准 softmax 要两趟：一趟遍历找最大值（数值稳定），一趟遍历求和。Online softmax 把两者合并成一趟，允许分块增量更新。

设第块的局部最大值为，局部 sum 为，局部 output 为。

合并两个块时，需要”矫正”之前块的值（因为全局 max 可能变了）：

这个矫正公式保证了分块 softmax 的数值结果和全局计算完全一致，不是近似。

标准 Attention:              Flash Attention:
  HBM ← 读 Q,K,V (O(nd))      SRAM ← 读一个 block 的 Q
  HBM ← 写 QK^T (O(n²))       ├─ HBM ← 读一个 block 的 K,V
  HBM ← 读 QK^T → softmax      ├─ SRAM 内算 partial attention
  HBM ← 写 output (O(nd))      ├─ online softmax 更新
                               └─ HBM ← 写最终 output
  显存: O(n²)                  显存: O(n)
  IO 次: 大                    IO 次: 小（block 在 SRAM 内处理）

实际效果：seq_len=2k 能省 5-10 倍显存，seq_len 越长越明显。现在所有 LLM 训练和推理几乎都依赖 Flash Attention。PyTorch 2.0+ 原生支持：

# PyTorch 2.0+ 自动调用最优实现（FlashAttention / MemEfficient）
attn_output = F.scaled_dot_product_attention(
    query, key, value,
    attn_mask=causal_mask,
    is_causal=True,
)

Paged Attention：KV Cache 的”虚拟内存”

GQA 减少了 KV head 数，但 KV cache 的存储管理本身也有优化空间。

朴素做法的问题：预分配 KV cache 为 max_seq_len 的连续内存。但 batch 里每个请求的长度差异很大——有的短请求早结束了，有的还很长。预分配造成大量空间浪费（碎片化），显存利用率只有 20-40%。

Paged Attention（vLLM, Kwon et al., 2023）：把 KV cache 切成固定大小的 Page（比如 16 tokens 一页），按需分配。类似操作系统的虚拟内存分页。

传统方案: 每个序列预分配一整块连续内存
  序列1: [████████████████░░░░░░░░] max_len=32, 实际只用 16
  序列2: [████████░░░░░░░░░░░░░░░░] max_len=32, 实际只用 8
  序列3: [██████████████░░░░░░░░░░] max_len=32, 实际只用 14
  浪费了大量未使用的空间

Paged Attention: 以 page 为单位动态分配
  物理 pages:  [Page0][Page1][Page2][Page3][Page4][Page5][Page6]...
  序列1 的页表: Page0 → Page3 → Page5 (3 pages, 48 tokens)
  序列2 的页表: Page1 → Page4 (2 pages, 32 tokens, 已结束可回收)
  序列3 的页表: Page2 → Page6 (2 pages, 32 tokens, 继续分)

三个核心好处：

零碎片浪费：用多少分配多少，不预分配。显存利用率从 20-40% 提升到 90%+。
内存共享：beam search 时多个 beam 共享同一段 prefix 的 KV cache，只需指向同一组物理页即可。并行生成时同理。
动态调整：请求变长时追加新页，请求结束回收所有页。

概念上的数据结构（注意：以下为概念示意，真实高性能实现不会每次 torch.cat 成连续 tensor，而是在 attention kernel 内直接根据 block 索引读取对应物理块，避免拷贝开销）：

# KV cache 不是一整块 tensor，而是 page table（概念示意）
class KVCache:
    def __init__(self, num_blocks, block_size, num_layers, num_heads, head_dim):
        # 物理存储：固定大小的 pages
        self.blocks = torch.zeros(num_blocks, block_size, num_layers,
                                   num_heads, head_dim)
        self.free_blocks = list(range(num_blocks))  # 空闲页列表

    def alloc(self, request_id, num_tokens):
        needed = ceil(num_tokens / self.block_size)
        pages = [self.free_blocks.pop() for _ in range(needed)]
        self.page_table[request_id] = pages
        return pages

    def free(self, request_id):
        self.free_blocks.extend(self.page_table.pop(request_id))

vLLM 等推理引擎在 attention kernel 内部直接读取 page table，按 block 索引访问 KV cache 物理块，无需先拼成连续 tensor。torch.cat 只是概念示意。

Paged Attention 配合 GQA，是目前 LLM 推理引擎（vLLM、SGLang、TensorRT-LLM）的标准方案。

RoPE：为什么每层都加，加在 QK 上

旋转位置编码（Rotary Position Embedding, RoPE）是苏剑林 2021 年提出的，现在几乎所有主流 LLM 都在用。

核心思想

RoPE 的思路：用旋转矩阵对 Q 和 K 进行编码，使得 QK^T 的点积结果天然包含相对位置信息。

对位置 m 的向量 x，RoPE 施加一个旋转：

看起来复杂，实际操作很简单——把向量按维度两两分组，每组用不同的角度旋转：

def apply_rotary_emb(x, cos, sin):
    # x: [batch, heads, seq, dim]
    # cos, sin: [seq, dim]
    # 把 dim 分成两半
    d2 = x.shape[-1] // 2
    x1, x2 = x[..., :d2], x[..., d2:]
    # 旋转
    rotated = torch.cat([
        x1 * cos - x2 * sin,
        x2 * cos + x1 * sin,
    ], dim=-1)
    return rotated

频率设计：大角度捕捉局部，小角度捕捉全局

旋转角度的设计：

不同维度的旋转频率不同（LLaMA 1/2 默认 base=10000，LLaMA 3 系列默认 base=500000，Qwen2.5 7B 默认 base=10000）：

低维（i 小，θ 大）：高频旋转，对位置变化敏感，捕捉局部关系
高维（i 大，θ 小）：低频旋转，对位置变化迟钝，捕捉全局关系

这个性质很重要。后来的方法（如 YaRN、NTK-aware scaling）就是通过调整 base 来扩展上下文窗口。

为什么每层都加

RoPE 是逐层对 Q 和 K 施加的，不是只加一次。每层对应 Transformer 的一个处理阶段，不同层的位置信息需求不同：

浅层更关注局部语法（相邻 token 之间的关系），需要精细的位置信息
深层更关注全局语义（远距离依赖），需要长程位置信息

如果在 Embedding 层加位置编码然后所有层共享，那位置信息经过多层非线性变换后会衰减和变形。每层重新注入 RoPE，保证位置信息在每层都是”新鲜的”。

代码上就是每次计算 attention 之前，在当前层的 Q 和 K 上 apply RoPE：

q, k = self.q_proj(x), self.k_proj(x)  # 投影
q, k = apply_rotary_emb(q, k, cos, sin)  # 每层都对本层的 Q/K 施加 RoPE
v = self.v_proj(x)  # V 不加 RoPE
attn_output = attention(q, k, v)
# 注：cos/sin 频率表通常是全局共享的（所有层用同一组频率），不是每层独立学习

为什么只加在 Q 和 K 上，不加在 V 上

Attention 的公式：

注意这里位置信息只在 QK^T 这一步起作用——attention score 决定了每个 V 的权重。V 本身不需要位置信息，因为 V 的加权求和已经是位置相关的结果（权重就是从位置敏感的 QK^T 来的）。

数学上：设经过 RoPE 后的 Q’ = R_m Q, K’ = R_n K，则

因为旋转矩阵满足（旋转的叠加性），所以 QK^T 的点积结果只依赖于相对位置 n-m。这就是 RoPE 的精妙之处：不做任何显式的相对位置计算，但结果天然编码了相对位置。

对 V 加 RoPE 是多余的，而且会干扰 Attention 的加权求和结果。

激活函数：从 ReLU 到 SwiGLU

三种主流激活函数对比

ReLU：

优点简单，缺点是 x < 0 时梯度为 0，神经元可能”死掉”（永远不激活）。Transformer 里很少单独用 ReLU。

GELU（Gaussian Error Linear Unit）：

其中是标准正态分布的 CDF。GELU 是 BERT、GPT-2 时代的标准激活函数。它比 ReLU 平滑，x < 0 时也有小幅输出（不是严格 0），梯度不会完全消失。

SiLU / Swish：

Google 在 2017 年搜出来的激活函数。和 GELU 形状很像（看图），但计算更简单（只有一个 sigmoid，不需要）。在深层网络里比 GELU 略好。

SwiGLU：给激活函数加个”门”

SwiGLU 是 PaLM 论文里验证效果最好的一种 Gated Linear Unit 变体。

GLU（Gated Linear Unit）的思路：让两个线性投影的输出做逐元素乘法，其中一个过激活函数当”门”：

SwiGLU 把 sigmoid 换成 SiLU：

然后再过一个 down projection：

为什么用 SwiGLU

和标准 FFN（ReLU(xW1)W2）相比：

门控机制：SiLU(xW_gate) 作为一个门，控制 xW_up 的信息流。这在直觉上类似 LSTM/GRU 的门控，让网络自己决定哪些信息通过。
非线性加倍：SiLU 是非单调的（x 很负时输出接近 0，x 负但不很负时有负输出），提供了更丰富的梯度信号。
PaLM 的消融实验：在同等参数量下，SwiGLU >> GELU ≈ ReLU > 无激活。

SwiGLU 的参数变化

标准 FFN 是 2 个矩阵（dim → 4×dim → dim），SwiGLU 是 3 个矩阵：

# 标准 FFN: 两个矩阵，中间 dim 通常是 4 * hidden_dim
self.w1 = Linear(dim, 4*dim)   # 升维
self.w2 = Linear(4*dim, dim)   # 降维

# SwiGLU: 三个矩阵，中间 dim 通常缩小到 8/3 * hidden_dim
# 保持总参数量不变
self.gate_proj = Linear(dim, intermediate_dim)  # 门
self.up_proj = Linear(dim, intermediate_dim)    # 值
self.down_proj = Linear(intermediate_dim, dim)  # 输出

为什么 intermediate_dim 要缩小？因为 SwiGLU 多了 gate_proj，如果保持原来的 4×dim，总参数量就多了 50%。一般把 intermediate_dim 调到 8/3 * dim 保持总参数不变：

$标准$

这个”保持总参数不变”的换算是 LLM 设计中的一个基本原则：替换一个组件时，要让总参数量等于原来的量，再对比效果。否则分不清效果差异是来自参数多了还是结构好了。

实际实现中，LLaMA 的做法：

1
2
3

# LLaMA 的 intermediate_dim 计算
intermediate_dim = int(8/3 * hidden_dim)
intermediate_dim = intermediate_dim + (256 - intermediate_dim % 256) % 256  # 取整到 256 的倍数

比如 LLaMA-7B dim=4096：8/3 × 4096 = 10922.67 → 取整到 256 的倍数 → 11008。

残差连接：信息流的”高速公路”

残差连接本身很简单：

但在 Transformer 里，残差连接承担了一个原论文没有明说但很重要的功能：残差流（residual stream）。

把 Transformer 的计算看成两条路径：

残差路径：直接复制，不加任何变换。信息靠残差连接在层间原样传递。
子层路径：经过 Attention 或 FFN 的计算。子层读取残差流里的信息，做变换，再写入残差流（通过加法）。

这种设计让梯度可以直接顺着残差路径回传，不受子层内部非线性变换的影响。这也是为什么即使 100+ 层的 Transformer 也能训练——残差连接保证了”梯度高速公路”不断。

Pre-Norm 在这个框架下的作用就是：在子层”读取”残差流之前，先把残差流的”幅度”控制住，防止子层看到太大或太小的输入。

MoE：用”专家”换参数

MoE（Mixture of Experts）是可选的，但近两年越来越主流（Mixtral、DeepSeek-V2/V3、Qwen-MoE）。

核心思想

把 FFN 从”一个大网络”变成”多个小专家 + 一个路由器”：

传统 FFN:                    MoE:
x → FFN → output            x → Router → 专家选择
                                  │
                            ┌─────┼─────┐
                            ▼     ▼     ▼
                          E0    E1    E2 ... (每个都是一个 FFN)
                            │     │     │
                            └─────┼─────┘
                                  ▼
                           加权求和 → output

路由（Routing）

每个 token 经过一个路由层（就是个 Linear 层），输出每个专家的 logit：

然后选 top-k 个专家（通常 k=2）：

router_logits = self.router(x)            # [N, num_experts]
scores = F.softmax(router_logits, dim=-1)
topk_weights, topk_indices = torch.topk(scores, k=2, dim=-1)
# 只用选中的专家做计算
for expert_idx in selected_experts:
    expert_output += expert(x) * weight

实际计算

不写循环，用矩阵运算实现：

# 1. 路由
router_logits = self.gate(x)                 # [batch*seq, n_experts]
routing_weights = F.softmax(router_logits, dim=-1)
weights, indices = torch.topk(routing_weights, k, dim=-1)
weights = weights / weights.sum(dim=-1, keepdim=True)  # 归一化

# 2. 把 token 分到对应的 expert
# 每个 expert 得到分配给它的 token，做完 FFN，再按原位置拼回去
# 这一步通常用 einsum 或者 scatter 实现

为什么要用 MoE

具体倍数取决于 expert 大小和 Attention 等非 FFN 参数占比。以 FFN 占总参数约 2/3 的典型模型为例，8 个 expert（每个与原始 FFN 同大小）、top-2 路由：

设 Dense 参数 = A（Attention 等）+ F（FFN）
MoE 总参数 ≈ A + 8F + Router，激活参数 ≈ A + 2F + Router
若 F ≈ 2A，则总参数约 5.7 倍 Dense，激活参数约 1.7 倍 Dense

若每个 expert 缩小（如 Mixtral 的做法），倍率会降低。这里统一用近似值示意——关键是总参数量涨了，但每个 token 的激活参数和计算量涨得少，用少量额外计算换大量额外容量。

负载均衡

MoE 有个头疼的问题：如果 Router 总是选某几个专家，其他专家就”学废了”。需要加一个辅助 Loss：

其中是分配给专家 i 的 token 比例，是 Router 给专家 i 的平均概率。

这个 loss 鼓励 Router 均匀分配 token，但不强制。因为太强制（如 top-1 + 严格均衡）会影响模型质量。

和 GQA 的关系

MoE 和 GQA 解决不同的问题：

GQA：减少 KV cache 显存（针对 Attention）
MoE：增加模型参数但不增加太多计算（针对 FFN）

两者不冲突，经常一起用。比如 DeepSeek-V3 就是 MoE + MLA（MLA 是比 GQA 更进一步的东西）。

参数量估算：从 config 算出总参数量

给定模型 config：{vocab_size V, dim d, n_layers L, n_q_heads H, n_kv_heads G, d_head, intermediate_dim}。

Embedding 层

每个 Transformer Block

RMSNorm（通常不带 bias，只有可学习）：

Self-Attention（无 bias）：

当 G=H（MHA）时：Attention = 4d²；GQA 比例 8:1 时：Attention ≈ 2.25d²。

SwiGLU FFN：

标准配置下 intermediate_dim = 8d/3，则 FFN = 8d²，与标准 FFN 参数量等同。

MoE（有 E 个 expert，每个 expert 大小同上）：

完整公式

Dense 模型（不绑定权重）：

Dense 模型（绑定权重）：

MoE 模型（不绑定，总参数）：

MoE 的激活参数（实际参与计算的参数）只有 top-k 个 expert 被激活：

实算例子：LLaMA-7B

config：V=32000, d=4096, L=32, H=32, G=32 (MHA), d_head=128, intermediate=11008

组件	计算	参数量
Embedding	32000 × 4096	131.1M
每层 Norm ×2	2 × 4096	8.2K
每层 Attention	4 × 4096²	67.1M
每层 FFN	3 × 4096 × 11008	135.3M
每层合计	8192 + 67.1M + 135.3M	202.4M
32 层总计	32 × 202.4M	6.48B
Final Norm	4096	4K
LM Head	4096 × 32000	131.1M
总计		≈ 6.74B

通常说”7B”，一部分是因为还有少量的其他参数（如 RoPE 的 cos/sin 频率表很小忽略不计），另一部分是习惯性向上取整。

参数量如何随超参变化

理解参数量公式后，调整模型大小就很直观了：

超参变化	参数量变化	备注
dim × 2	~×4 (因为 FFN ∝ d², Attention ∝ d²)	主要的 scaling 维度
L × 2	约 ×2	线性增长
V × 2	Embedding + LM Head 翻倍	占总参数比例通常不大
GQA: H/G 增大	轻微减少 Attention 参数	K,V 投影缩小
SwiGLU: intermediate_dim × 2	FFN 参数翻倍	需考虑整体匹配

常见模型的参数量对应

模型	dim	L	总参数
LLaMA-7B	4096	32	6.7B
LLaMA-13B	5120	40	13.0B
LLaMA-70B	8192	80	~70B
Qwen2.5-7B	3584	28	7.6B
Qwen2.5-72B	8192	80	72.7B
Mistral-7B	4096	32	7.2B
DeepSeek-V3	7168	61	671B (总) / 37B (激活)

Loss 计算：语言模型的”考试题”

预测目标

LLM 的训练目标是下一个 token 预测。给定序列，最大化：

实际操作：输入 [BOS] 今天天气真好，模型要预测 今天、天气、真好、[EOS]。

代码实现

# 前向
logits = model(input_ids)        # [batch, seq_len, vocab_size]

# 计算 loss
# 方案 A：labels 直接用 input_ids，由代码做 shift
# input_ids:  [今天, 天气, 真好, EOS]
# labels 即 input_ids
shift_logits = logits[..., :-1, :].contiguous()   # [batch, seq_len-1, vocab]
shift_labels = input_ids[..., 1:].contiguous()     # [batch, seq_len-1]
# shift 后：logits 位置 0 预测 token 1，位置 1 预测 token 2，...
#          即 logits 的每个位置预测下一个位置的 token

loss = F.cross_entropy(
    shift_logits.view(-1, vocab_size),
    shift_labels.view(-1),
    ignore_index=-100,   # 忽略 padding token 的 loss
)
# HuggingFace 风格：labels 传原始 input_ids，模型内部做 shift

Cross-Entropy 展开

Cross-Entropy Loss = Log-Softmax + NLL：

其中是正确 token 的 logit，V 是词表大小（通常 30k~150k）。

Perplexity 的关系

Perplexity 是 LLM 领域常用的评估指标，就是 loss 的指数：

直观理解：PPL 是模型在预测每个 token 时”平均候选数”。如果模型对每个 token 都完全确定（概率=1），PPL=1；如果每次都在 100 个等概率选项里猜，PPL=100。

训 LLaMA-7B 的时候，loss 大概从 10+ 降到 2.x，对应 PPL 从 20000+ 降到 7~8。

Loss 计算的稳定性

大 vocab_size 下，Log-Softmax 的指数可能溢出。标准做法是减去最大值：

PyTorch 的 F.cross_entropy 已经内置了这个优化，不用自己写。

训练模式

有两个常见模式：

Teacher Forcing（标准）：输入是真实的前缀，预测下一个 token。loss 在所有位置计算。
自回归（推理时）：用自己的输出作为下一步的输入，loss 不适用。

训练时不逐 token 自回归生成，因为所有位置的真实 token 已知，直接用 teacher forcing 并行计算所有位置的 loss。但训练目标本身就是自回归 next-token prediction——预测第个 token 时只能看到（通过 causal mask 保证），不能看到及之后。

Gradient Accumulation

LLM 训练中常听到 gradient_accumulation_steps。因为 batch size 需要很大（如 4M tokens），但显存放不下：

# 等效 batch_size = per_device_batch_size * grad_accum_steps * num_gpus
for i, batch in enumerate(dataloader):
    loss = model(batch) / grad_accum_steps
    loss.backward()  # 梯度累积在 param.grad 里
    if (i + 1) % grad_accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

注意 loss 要除以 grad_accum_steps，不然梯度会被放大 N 倍。

辅助 Loss：不只 Cross-Entropy

现代 LLM 训练不只是优化 Cross-Entropy，总 Loss 通常包含多个组件：

1. 语言模型 Loss（主 Loss）

这是核心，占主导地位。

2. MoE 负载均衡 Loss

之前 MoE 章节提过，这里展开完整公式：

其中：

：实际分配给 expert i 的 token 比例（每个 token 只看 top-1 路由结果）
：Router 给 expert i 的 softmax 概率在所有 token 上的平均

$𝟙$

如果 token 均匀分配到各 expert，，，则，达到理论最小值。权重系数通常取 0.01，大小要仔细调——太大会强制均衡损害质量，太小则负载不均。

3. MoE z-Loss（Router Logit 尺度惩罚）

DeepSeek-V2/V3 引入了 z-Loss，惩罚 Router logit 的 scale：

其中是 token t 到 expert j 的 router logit（softmax 之前的 raw logit）。

为什么需要：Router logit 的 scale 如果越来越大，softmax 会趋向 one-hot，Router 失去”软路由”能力，同时数值也不稳定。z-Loss 惩罚 log-sum-exp 的平方，本质上是在惩罚 logit 的整体幅度。

典型权重，很小但足以防止 logit 发散。

4. 其他可选辅助 Loss

Embedding 正则化（少见）：$\mathcal{L}{\text{emb}} = \lambda \cdot |W{\text{emb}}|_2^2$，防止 embedding 矩阵过拟合
Auxiliary z-Loss for shared expert（DeepSeek-V3）：对共享 expert 也加 z-Loss，防止其输出 scale 和稀疏 expert 不一致
Token-level auxiliary z-Loss（某些 Qwen 版本）：每层都加，但权重极小

评估指标：不只是看 Loss

Loss/PPL 反应的是语言建模能力，但最终用户关心的是能不能回答问题、写代码、做推理。因此 LLM 的评估体系分三层：

第一层：训练监控指标

指标	含义	作用
Loss / PPL	核心训练信号	判断训练是否收敛
Gradient Norm	梯度 L2 范数	检测梯度爆炸/消失
Learning Rate	当前 lr	配合 scheduler 观察
Token Throughput	tokens/sec/GPU	训练效率，影响成本
MFU	Model FLOPs Utilization	GPU 发挥了多少算力

第二层：验证集 Loss / PPL

在 held-out 数据上测 PPL，判断是否过拟合。如果训练 PPL 持续下降但验证 PPL 上升，说明过拟合——但 LLM 通常训练一个 epoch（甚至不到），很少出现过拟合。

第三层：下游 Benchmark

每个 benchmark 测不同的能力维度。单一指标无法代表模型质量，通常看 benchmark suite 的平均分或最高分。

Benchmark	测什么	形式	典型模型得分区间
MMLU	57 个学科知识广度	4 选 1	60-90
HellaSwag	常识推理	4 选 1 补全句子	70-95
ARC-Challenge	科学推理	4 选 1	55-96
GSM8K	小学数学应用题	自由生成数字答案	50-95
MATH	竞赛级数学	自由生成	20-80
HumanEval	代码生成	补全函数，过单元测试	30-95 (pass@1)
MBPP	代码生成（基础）	同上	40-85
BBH	综合推理（23 项子任务）	多种格式	50-90
TruthfulQA	事实真实性	多选题	40-65 (MC1)
IFEval	指令遵循	判断格式/约束是否满足	30-85

这些 benchmark 的评分逻辑各不相同。HumanEval 用 pass@k（k 次采样，有一次过测即算对），MMLU 是 accuracy，GSM8K 需要提取最终答案数值后比较。

Prompt 格式方面，分两类：

固定 prompt 的 benchmark（如 MMLU、HellaSwag）：评测框架内置了标准化的 few-shot 模板，每个学科/question 用同一个格式，只替换题目内容。例如 MMLU 的 prompt：
1
2
3
4
The following are multiple choice questions about {subject}.
Q: {question}
A. ... B. ... C. ... D. ...
Answer:
自由生成的 benchmark（如 GSM8K、HumanEval）：prompt 相对固定但模型需要自由输出，评测代码按规则（正则提取、单元测试判定）来给分。GSM8K 要求模型在末尾输出 #### 42 格式，评测脚本提取 #### 后数字与答案对比。

开源评测框架（lm-evaluation-harness、OpenCompass）已经封装好了这些 prompt 模板和评分逻辑，调用即可，不需要手写。

如何综合看这些指标

Loss 低 ≠ 下游好：模型可能在语言建模上很精确但缺乏推理能力。Loss 反映的是”预测下一个词的能力”，不是”理解后再输出的能力”。
Benchmark 之间不独立：MMLU 高通常意味着其他知识类 benchmark 也高，但和代码/数学不一定相关。
评测有方差：同一模型的同一 benchmark，不同 seed、不同 prompt、不同题库拆分，结果可能差 1-3 个点。所以 benchmark 变动 <2 点通常不给人看到。
“涌现”现象：某些能力（如复杂推理、多步逻辑）在模型较小（<10B）时几乎为零，但到一定规模突然出现。所以评测小模型和大模型时要关注能力是否存在质变。

权重绑定：Embedding 和 LM Head 共享参数

输入 Embedding 矩阵和输出 LM Head 矩阵形状天然相同，语义上也对称（都是 token ↔ 向量的映射）。一个自然的想法：能不能共用？

三种方案

1. 不绑定（LLaMA 风格）

1
2
3

self.embed_tokens = nn.Embedding(vocab_size, dim)     # [V, d]
self.lm_head = nn.Linear(dim, vocab_size, bias=False) # [d, V]
# 两个独立矩阵，总参数: 2 × V × d

LLaMA、Qwen、DeepSeek 都选这个。理由简单：多出来的参数量带来的灵活性比省的那点参数值。

2. 输入-输出绑定（GPT-2 风格）

self.embed_tokens = nn.Embedding(vocab_size, dim)
self.lm_head = nn.Linear(dim, vocab_size, bias=False)
self.lm_head.weight = self.embed_tokens.weight  # 指向同一个 tensor
# 总参数: V × d，省了一半

GPT-2、OPT、BLOOM、XLNet 用这个。直觉上合理：一个 token 的 embedding 向量和它的输出 logit 向量应该有某种关联——如果”猫”的 embedding 靠近”狗”，那模型预测”猫”时对”狗”也应该给高分。权重绑定强制了这种约束。

3. 全绑定（极端，LLM 不用）

ALBERT 试过不仅 embedding 和 LM Head 共享，所有 Transformer Block 的参数也共享——整个模型只有一层参数但执行 L 遍。效果有限，现代 LLM 不会这么搞。

不绑定的理由

输入输出分布差异太大。Embedding 的输入是离散 token_id，LM Head 的输入是经过 L 层变换的 continuous hidden state。两者需要的映射关系不同，强制共享权重限制了表达能力。
参数占比太小。LLaMA-7B 中 embedding + LM Head 约

占 7B 总参数的不到 4%。省这 130M 参数意义不大，不如留着增加灵活性。

梯度冲突。绑定后 embedding 同时接收两个梯度：一个来自 LM Head 的 loss 反向，一个来自第一层的隐式梯度。两者混合在一起优化，方向可能冲突。如果一定要绑定，建议给 LM Head 的梯度乘以缩放因子。

训练配方：AdamW + Cosine Schedule + Warmup

LLM 的训练超参出奇地一致，几乎所有模型的配方都差不多。记住了就可以直接套。

AdamW：把 weight decay 从 gradient 里拆出来

Adam 原版把 L2 正则化混在梯度更新里做：

其中是修正过的一阶动量，是修正过的二阶动量。L2 正则化的效果等价于在每个梯度步减去，但这个被分母调制了——参数的不同维度，wd 的实际力度不一样。

AdamW（Loshchilov & Hutter, 2019）的做法：把 weight decay 独立出来，直接对参数衰减：

区别很小但很重要：weight decay 不再被自适应学习率”缩放”，效果更一致。

LLM 的 AdamW 超参：

超参	值	为什么
	0.9	一阶动量系数，常规值
	0.95	二阶动量系数（比常规 0.999 小），让模型更快适应新梯度分布
	或	数值稳定，bf16 下可适当调大
weight decay	0.1	比 CNN 时代的大得多（），LLM 需要更强的正则化
grad clip	1.0	裁剪全局 L2 norm，防止单步更新过大

代码：

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=lr,
    betas=(0.9, 0.95),
    eps=1e-5,
    weight_decay=0.1,
)

# 每个 step 之后 clip gradient
for batch in dataloader:
    loss = model(batch) / grad_accum_steps
    loss.backward()
    if (step + 1) % grad_accum_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()

Learning Rate Schedule：Cosine Decay

Warmup + Cosine Decay 是标配：

lr
^
|   ████████
|  █        █
| █          █
|█             ███
|                   ███████████ η_min
+----------------------------------> steps
    |warmup|
            |---- cosine decay ----|

：峰值学习率，LLaMA-7B 约，更大模型用更小的 lr
：最终学习率，一般是的 1/10
：warmup 步数，通常占总步数的 1%-5%

def cosine_schedule(step, total_steps, lr_max, lr_min, warmup_steps):
    if step < warmup_steps:
        # 线性 warmup: lr_min → lr_max
        return lr_min + (lr_max - lr_min) * step / warmup_steps
    # Cosine 衰减
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * progress))

为什么需要 warmup：训练初期 Adam 的二阶动量还没稳定，梯度方差大。从大 lr 直接开始容易让模型走偏。Warmup 让 lr 从很小线性爬到峰值，给优化器一个”探路”的时间。

Mixed Precision：fp16 vs bf16

不用 fp32 训练 LLM，显存放不下。

	fp32	fp16	bf16
总位数	32	16	16
指数位	8	5	8
尾数位	23	10	7
范围	±3.4×10³⁸	±65504	±3.4×10³⁸
精度	高	中等	低但范围大
需要 loss scaling	不需要	需要	不需要

bf16 的范围和 fp32 一样（指数位相同），意味着小的梯度值不会下溢成 0，大的值也不会溢出。不需要 loss scaling，代码更简单。

大部分 LLM（LLaMA、Qwen、Mistral）用 bf16。只有 V100 等老 GPU 不支持 bf16，才只能用 fp16 + loss scaling。

# bf16 训练（推荐，简单）
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    output = model(input_ids)

# fp16 训练（老 GPU，需要 GradScaler 防下溢）
scaler = torch.cuda.amp.GradScaler()
with torch.autocast(device_type='cuda', dtype=torch.float16):
    output = model(input_ids)
loss = loss_fn(output, labels)
scaler.scale(loss).backward()       # loss 放大
scaler.step(optimizer)              # optimizer step 前缩小
scaler.update()                     # 动态调整 scale 因子

完整训练伪代码

model = LLM(config).to(device).bfloat16()

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4, betas=(0.9, 0.95), eps=1e-5, weight_decay=0.1,
)

scheduler = CosineSchedule(
    warmup_steps=2000, total_steps=100000,
    lr_max=3e-4, lr_min=3e-5,
)

for step, batch in enumerate(dataloader):
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        logits = model(batch.input_ids)
        loss = F.cross_entropy(
            logits[:, :-1].reshape(-1, vocab_size),
            batch.labels[:, 1:].reshape(-1),
            ignore_index=pad_token_id,
        )
        loss = loss / grad_accum_steps

    loss.backward()

    if (step + 1) % grad_accum_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

计算量估算：训练一个 LLM 要多少算力

知道模型多大后，自然要问：训它要多少算力？推理一次要多少？

单个 Token 前向 FLOPs

Transformer 的计算量主要来自矩阵乘法。矩阵乘法 A_{m×k} @ B_{k×n} 的 FLOPs = （一次乘一次加）。

近似公式：每个 token 经过模型一次，前向 FLOPs ≈ （P 为总参数），因为每个参数参与一次乘法和一次加法。

更精细的逐组件分解：

组件	FLOPs (per token)	占比(典型 7B 模型)
QKV 投影 (MHA)		~24%
QK^T Attention Score		短序列可忽略，长序列显著
Output 投影		~8%
FFN gate+up	$'_' allowed only in math mode4d \cdot \text{im_dim}$	~44%
FFN down	$'_' allowed only in math mode2d \cdot \text{im_dim}$	~22%
RMSNorm		<0.1%
Embedding / LM Head		<0.1%

主要计算量在 FFN 和 QKV 投影的线性层。当 seq_len > 4k 时，Attention Score 的 O(n²d_head) 项开始不可忽略。

一次训练步骤的 FLOPs

为什么是 6P？前向 ≈ 2P，反向 ≈ 4P（梯度计算约等于 2× 前向），总共 ≈ 6P。

总训练计算量

如果模型总参数为 P，训练 token 数为 D：

这就是著名的 “6ND” 公式（Kaplan et al., 2020; Hoffmann et al., 2022）。

实算：LLaMA-7B

P = 6.74B, D = 1T tokens

用 2048 张 A100 (312 TFLOPS fp16)，假设 50% MFU（真实训练很难超过 60%）：

$天数天$

实算：LLaMA-70B（~1.4T tokens）：

FLOPs，约 GPU-hours (A100)。2048 张 A100 约 108 天，10000 张则约 22 天。

缩放定律：训练 Token 数与参数量的关系

给了固定预算，模型应该做多大、训多少数据？这个问题决定了一切。

Kaplan 缩放定律 (2020)

OpenAI 的 Kaplan et al. 首次系统研究了 Transformer LM 的 scaling 行为。核心发现是 Loss 与 P（参数量）和 D（token 数）都呈幂律关系：

他们测得 , 。直觉上：数据量翻倍对 loss 的提升比参数量翻倍更大。由此他们建议：

即模型变大的速度应该快于数据增长的速度——10× 算力预算下，模型应放大 ~5×，数据只需增加 ~2×。

Chinchilla 缩放定律 (2022)

DeepMind 的 Hoffmann et al. 用更大的实验范围（扫描了 70M 到 16B 的模型，各用 4 段不同的 token 量训练），推翻了 Kaplan 的结论。他们发现：在充分调参的情况下，P 和 D 对 loss 的贡献几乎是平等的。

Chinchilla 提出的 loss 预测公式（Approach 3, parametric fit）：

其中：

：模型参数量（non-embedding）
：训练 token 数
：不可约 Loss（irreducible loss，语言的固有不确定性。即使无限大的模型+数据也无法降到 0）
：模型/数据各自的幂律指数

Chinchilla 实际估出的值：, , , , 。

给定固定计算预算，最优分配是让两项的边际收益相等。数学上：

当时，，两者等比增长。

核心结论：

即param:data = 1:20。

两个定律的分歧及原因

	Kaplan (2020)	Chinchilla (2022)
结论	D ∝ P^0.74（模型长得更快）	D ∝ 20P（等比增长）
10× 算力	模型 ×5，数据 ×2	模型 ×3.16，数据 ×3.16
实验方法	固定模型大小，只变 token 数	同时扫描模型大小和 token 数
学习率调度	每个模型独立调 lr	所有模型统一 Cosine schedule
Loss 计算	只看每个序列最后一个 token	计算序列内所有 token 的 loss

关键分歧来源：Kaplan 的实验设计导致小模型训了大量数据但没有真正收敛就被截断了（lr decay 提前结束），因此低估了数据的作用。Chinchilla 用了统一的学习率调度（所有训练到 loss 稳定），结论更可靠。业界普遍接受 Chinchilla。

Chinchilla 最优表

模型大小	Chinchilla 最优 tokens	所需训练 FLOPs
1B	20B
7B	140B
70B	1.4T
175B	3.5T

为什么现在几乎所有模型都”过度训练”

LLaMA-7B 训了 1T tokens（D/P ≈ 143），LLaMA 3 8B 训了 15T tokens（D/P ≈ 1875），都远超 Chinchilla 的 20:1。这不是不知道怎么选，而是Chinchilla 最优的假设已经不成立了：

Chinchilla 假设”一个 FLOP 花在训练上和花在推理上成本相等”。

现实中：

训练是一次性的。训一个 7B 模型花 15 天 × 2048 A100，总预算可控。
推理是持续的。模型部署后被调用亿次级，推理成本直接和用户规模挂钩，线性增长无上限。
小模型推理硬件要求低、速度快、用户体验好。在同等推理成本下，一个小而”饱读诗书”的模型，能服务的用户远多于一个同等能力的大模型。

因此工业界实际采取的是 “推理成本优先法”：

先确定能接受的推理硬件和延迟 → 给定模型大小 P
用尽量多的高质量数据训这个 P，直到 loss 不再下降 → 决定 D
D/P 可以远超 20

理论对此有定量解释（Sardana & Frankle, 2023）：如果推理成本占比足够大，最优 D/P 比可以是几百甚至上千。

比 Chinchilla 实际的表中加入”工业界实际”一列：

Chinchilla P	Chinchilla D	工业界实际 P	工业界实际 D	D/P
5B	100B	7B (LLaMA)	1T+	143
12B	240B	70B (LLaMA)	1.4T	20
25B	500B	8B (LLaMA3)	15T	1875

可以看到：P 小、D 极大的趋势越来越明显。高质量数据的获取能力，已经成了 LLM 竞争的核心壁垒——不是你有没有算力放大模型，而是你有没有 15T 高质量 tokens。

数据重复 vs 新鲜数据

Chinchilla 的另一个发现：重复数据不如新鲜数据。同样 total tokens，用 unique data 训练的 loss 低于重复数据。但在以下场景下，重复数据仍有效：

数据来源有限（如代码、数学等专业语料）
多 epoch 训练（4-8 epoch），loss 仍能继续下降
高质量数据比例小，需要反复利用

经验速查：给定预算怎么选 P 和 D

纯研究/追求最小 Loss：按 Chinchilla 最优， → D ≈ 20P
产品化/有推理成本压力：先定 P（按目标推理硬件和延迟），再尽量多训数据
用已有模型做数据配比实验：如果 Tok/S 已经很贵，使用小模型 + 多数据

GQA 对计算量的影响

GQA 减少了 K 和 V 投影的参数和计算。在 MHA→GQA(8:1) 时：

K, V 投影 FLOPs：从降到
总 Attention FLOPs 减少：约 15%
但因为 Attention 只占总 FLOPs 的 ~24%，整体节省约 3-4%

GQA 的主要收益不在训练计算量，在推理时 KV cache 显存。

MoE 对计算量的影响

MoE 前向时只激活 k 个 expert，所以：

比如 E=8, k=2 的 MoE，总参数约 5-7 倍 Dense，但训练和推理的计算量只增加不到 2 倍。这是 MoE 的核心优势：用少量额外计算换大量额外容量。

推理计算量

推理时每个 token 只需一次前向：

对于 Dense 模型就是 2P，对于 MoE 模型 P_active 远小于 P_total。

首 token（prefill）和后续 token（decode）的计算模式不同：

Prefill：并行处理所有 prompt token，compute-bound
Decode：逐 token 生成，受 KV cache 读写带宽限制，memory-bound

因此推理优化关注的不只是 FLOPs，更关注显存带宽利用率和 KV cache 管理。

总结：一条数据在模型里的”旅程”

作为收尾，把上面的东西串起来，走一遍一个 token 序列通过模型的完整流程：

输入: [今天, 天气, 真好]  →  token_ids: [123, 456, 789]

1. Embedding
   token_ids → Embedding 查表 → [3, dim]

2. Block × L (以 LLaMA 风格为例)
   for each block:
       residual = x
       x = RMSNorm(x)                    # Pre-Norm
       Q, K, V = qkv_proj(x)
       Q, K = apply_RoPE(Q, K, positions) # 每层独立加，只加 Q,K
       # GQA: 如果 n_kv_heads < n_q_heads, expand K,V
       x = attention(Q, K, V, causal_mask)
       x = residual + x                  # 残差
       residual = x
       x = RMSNorm(x)                    # Pre-Norm
       # SwiGLU FFN (或 MoE)
       x = gate_proj(x) → SiLU() → × up_proj(x) → down_proj()
       x = residual + x                  # 残差

3. Final RMSNorm(x)

4. LM Head: logits = x @ lm_head_weight^T  → [3, vocab_size]

5. Loss: $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{LM}} + \alpha \cdot \mathcal{L}_{\text{balance}} + \beta \cdot \mathcal{L}_z$ (cross-entropy + 辅助 Loss)

每个设计决策都对应一个实际的问题：

设计	解决的问题
Pre-Norm + 残差	深层网络训练稳定性
RMSNorm	计算效率（去均值、去 β）
GQA	推理时 KV cache 显存
Causal Mask	自回归生成，防止泄露未来 token
Flash Attention	attention 显存从 O(n²) 降到 O(n)，支持长序列
Paged Attention	KV cache 按页分配，消除显存碎片，利用率 90%+
RoPE (每层 Q,K)	每层独立的位置感知，相对位置编码
SwiGLU	门控 + 更好的非线性，参数量匹配保持总量不变
MoE	增加总参数量但控制激活参数量与计算量
参数量估算	从 config 精确计算总参数，指导模型设计
Cross-Entropy	下一个 token 预测的标准目标
辅助 Loss	MoE 的负载均衡、z-Loss 等，保证训练稳定
评估指标	Loss/PPL + 下游 Benchmark 综合评价模型能力
权重绑定（可选）	节省 Embedding/LM Head 参数
AdamW + Cosine + bf16	稳定的训练配方
计算量估算（6ND）	Flops ≈ 6PD，Chinchilla 最优 D ≈ 20P

写完这篇，后面再看任何 LLM 的源码应该都能一眼看出这些部件在干嘛了。

参考：LLaMA (1/2/3)、Qwen 2.5、PaLM、Mistral、DeepSeek-V2/V3、GQA 论文 (Ainslie 2023)、RoPE 论文 (Su 2021)、RMSNorm 论文 (Zhang & Sennrich 2019)、Chinchilla (Hoffmann 2022)、Scaling Laws (Kaplan 2020)、Mixtral (Jiang 2024)