文本进了模型

上一章你学会了分词——一段文字变成了一串数字 ID。但然后呢？这些 ID 丢进模型之后，经历了什么变换，才能"预测下一个字"？

答案全在一个文件里：model/model_minimind.py，一共 288 行。GPT、Llama 的核心架构，就在这里。

知识地图

在深入每个组件之前，先建立全局视野。这张思维导图展示了 MiniMind 模型包含的所有核心组件——接下来的每一页都会围绕它展开。

Part 1：全景

先看全貌，再看零件。这一部分建立数据从输入到输出的完整直觉。

数据流水线

文字进去，经过层层处理，最后预测下一个字

"你好世"这 3 个字进到模型里，到底经历了什么？看下面这条流水线——每个箭头都是一个关键变换步骤。

类的组成

MiniMind 用了哪些类，谁包含谁

288 行代码分成了 7 个类/函数。理解它们的包含关系，就像拿到了一张地图。

MiniMindForCausalLM          # 顶层模型
  └── MiniMindModel           # Transformer 主体
        ├── embed_tokens      # 词嵌入层
        ├── layers × N        # N 个 Transformer Block
        │     └── MiniMindBlock
        │           ├── self_attn      # GQA 注意力
        │           ├── mlp            # SwiGLU 前馈 / MoE
        │           ├── input_layernorm
        │           └── post_attention_layernorm
        └── norm               # 最终 RMSNorm
  └── lm_head                 # 输出投影层 (768 → 6400)

最外层是 MiniMindForCausalLM，它包含一个 MiniMindModel（主体）和一个 lm_head（输出头）。Model 内部有 Embedding + 8 层 Block + 最终 Norm。

Config 图纸

所有超参数都在这里定义——模型的"基因"

在盖楼之前先看图纸。MiniMindConfig 定义了模型的所有超参数。以下代码对应源文件第 10-45 行。

# model_minimind.py 第 10-45 行
class MiniMindConfig(PretrainedConfig):
    model_type = "minimind"
    def __init__(self, hidden_size=768, num_hidden_layers=8,
                 use_moe=False, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size          # 模型"宽度"
        self.num_hidden_layers = num_hidden_layers  # 模型"深度"
        self.vocab_size = kwargs.get("vocab_size", 6400)
        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
        self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
        self.head_dim = kwargs.get("head_dim",
            self.hidden_size // self.num_attention_heads)  # = 96
        self.intermediate_size = kwargs.get("intermediate_size",
            math.ceil(hidden_size * math.pi / 64) * 64)    # = 2432

Config 参数

每个参数控制模型的哪个方面

Part 2：归一化

Transformer 的两大基础设施：让数值稳定、让模型知道位置。

RMSNorm

比 LayerNorm 更快，效果几乎相同

定义：对每个向量的所有维度求均方根（RMS），然后做逐元素除法，将向量归一化到单位长度附近。目的是防止训练时数值爆炸或消失。

白话：就像调节音响音量——不管输入多大多小，输出都保持在合理区间。

这个比喻只是为了帮你建立直观感受。在计算机内部，RMSNorm 是对每个向量做纯数值运算——求平方均值、取倒数平方根、逐元素除——没有"音量"这种物理概念。

# model_minimind.py 第 50-60 行
class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # 可学习的缩放参数

    def norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        return (self.weight * self.norm(x.float())).type_as(x)

RMSNorm 解读

每一步在算什么，为什么要这么做

# 核心公式：RMSNorm(x) = weight * (x / sqrt(mean(x²) + eps))
#
# self.norm(x) 做了什么？
#   1. x.pow(2)     → 每个元素平方
#   2. .mean(-1)    → 沿最后一维求均值（得到每个向量的 RMS²）
#   3. + self.eps   → 加一个小数防止除以 0
#   4. torch.rsqrt  → 取倒数平方根（= 1 / sqrt(...)）
#   5. x * ...      → 逐元素相乘，归一化到单位长度附近
#
# forward 做了什么？
#   self.weight * self.norm(x.float())  → 归一化后乘以可学习缩放系数
#   .type_as(x)                         → 转回原始精度（bf16/fp16）

# 维度追踪（假设输入 (2, 5, 768)）：
# x.pow(2)                    → (2, 5, 768)    每个元素平方
# .mean(-1, keepdim=True)     → (2, 5, 1)      每个向量一个均值
# torch.rsqrt(...)            → (2, 5, 1)      标量因子
# x * ...                     → (2, 5, 768)    逐元素除法（广播）

# (x - mean) / std * gamma + beta
# 减均值、有偏移项 beta
# 计算两次求和（均值+方差）

# x / rms * weight
# 不减均值、无偏移项
# 只计算一次求和（平方均值）
# 快 10-15%，效果相当

为何要位置编码

注意力机制本身不知道词的顺序

"我爱你"和"你爱我"在注意力看来完全一样——都是三个向量。自注意力是排列不变的：它计算的是向量之间的相关性，和位置无关。

如果模型不知道"我"在第 1 个位置、"你"在第 3 个位置，它就分不清主语和宾语。位置编码就是给每个位置打上"标签"，让模型能区分顺序。

RoPE 原理

通过旋转角度编码位置，点积自动包含相对距离信息

白话：想象一个指南针，不同位置的向量旋转不同角度。位置 0 转 θ 度，位置 1 转 2θ 度。两个位置做点积时，角度差就自然编码了相对距离。

这个比喻只是为了帮你建立直观感受。在计算机内部，RoPE 通过在 cos/sin 频率矩阵上做外积运算，将位置信息转化为旋转变换的三角函数值，然后对 Q 和 K 向量施加这个旋转变换——没有真正的"指南针"在转，只有矩阵运算。

预计算缓冲

cos/sin 值只需算一次，之后反复使用

# model_minimind.py 第 62-78 行（简化版）
def precompute_freqs_cis(dim, end=32768, rope_base=1e6, rope_scaling=None):
    # 第1步：计算频率 θ_i = 1 / (base^(2i/d))
    # dim=96 时，产生 dim/2=48 个不同频率
    freqs = 1.0 / (rope_base ** (torch.arange(0, dim, 2).float() / dim))

    # 第2步：如果需要长度外推（YaRN），调整频率
    if rope_scaling is not None:
        # 低频维度缩放 1/factor，高频保持不变，中间线性过渡
        ramp = torch.clamp(...)
        freqs = freqs * (1 - ramp + ramp / factor)

    # 第3步：位置 t × 频率 θ → 外积矩阵
    t = torch.arange(end)
    freqs = torch.outer(t, freqs).float()  # (32768, 48)

    # 第4步：计算 cos/sin 并各复制一份拼接
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)  # (32768, 96)
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)  # (32768, 96)
    return freqs_cos, freqs_sin

应用 RoPE

把预计算的 cos/sin 应用到 Q 和 K 上

# model_minimind.py 第 80-84 行
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    def rotate_half(x):
        # 把向量切成两半，交叉互换：[a,b] → [-b,a]
        # 这等价于二维旋转矩阵 [cos -sin; sin cos] 的一半
        return torch.cat((-x[..., x.shape[-1]//2:],
                           x[..., :x.shape[-1]//2]), dim=-1)

    # q_embed = q * cos + rotate_half(q) * sin
    q_embed = (q * cos.unsqueeze(unsqueeze_dim)
             + rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
    k_embed = (k * cos.unsqueeze(unsqueeze_dim)
             + rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
    return q_embed.to(q.dtype), k_embed.to(k.dtype)

Part 3：注意力

Transformer 最核心的组件——决定"我该关注谁"。

注意力全景

Query、Key、Value 三个角色如何协作

想象你去图书馆找书。Query 是"你想找什么"，Key 是每本书的"标签"，Value 是"书的内容"。注意力机制就是：用你的 Query 和所有 Key 比对（点积），找出最相关的几本，然后把对应的 Value 加权求和。

比喻到此为止。在数学上，注意力就是 Q·K^T 得到权重矩阵，softmax 归一化后乘以 V。纯粹的矩阵运算。

注意力初始化

四个投影矩阵和 QK-Norm

# model_minimind.py 第 91-109 行
class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_local_heads = config.num_attention_heads     # 8 个 Q 头
        self.n_local_kv_heads = config.num_key_value_heads  # 4 个 KV 头
        self.n_rep = self.n_local_heads // self.n_local_kv_heads  # = 2
        self.head_dim = config.head_dim                     # 96
        # 四个投影矩阵（无偏置）
        self.q_proj = nn.Linear(768, 8 * 96, bias=False)    # 768 → 768
        self.k_proj = nn.Linear(768, 4 * 96, bias=False)    # 768 → 384
        self.v_proj = nn.Linear(768, 4 * 96, bias=False)    # 768 → 384
        self.o_proj = nn.Linear(8 * 96, 768, bias=False)    # 768 → 768
        # QK-Norm：对每个头的表示做归一化（提升训练稳定性）
        self.q_norm = RMSNorm(96)
        self.k_norm = RMSNorm(96)

投影与 QK-Norm

Attention forward 的前半段（第 111-123 行）

def forward(self, x, position_embeddings, past_key_value=None,
            use_cache=False, attention_mask=None):
    bsz, seq_len, _ = x.shape
    # ① 线性投影：x → Q, K, V
    xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
    xq = xq.view(bsz, seq_len, 8, 96)   # 拆成 8 个 Q 头
    xk = xk.view(bsz, seq_len, 4, 96)   # 拆成 4 个 KV 头
    xv = xv.view(bsz, seq_len, 4, 96)

    # ② QK-Norm：稳定训练（防止 Q/K 幅度过大导致注意力退化）
    xq, xk = self.q_norm(xq), self.k_norm(xk)
    # ③ 应用 RoPE 旋转位置编码
    cos, sin = position_embeddings
    xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)

    # ④ KV-Cache：拼接历史 KV
    if past_key_value is not None:
        xk = torch.cat([past_key_value[0], xk], dim=1)  # 沿序列维拼接
        xv = torch.cat([past_key_value[1], xv], dim=1)
    past_kv = (xk, xv) if use_cache else None

注意力与 GQA

4 个 KV 头复制成 8 个，和 Q 头匹配后计算注意力

    # model_minimind.py 第 124-134 行
    # ④ repeat_kv：4 个 KV 头复制成 8 个，和 Q 头匹配
    xq = xq.transpose(1, 2)                           # (B, 8, S, 96)
    xk = repeat_kv(xk, self.n_rep).transpose(1, 2)    # (B, 8, S, 96)
    xv = repeat_kv(xv, self.n_rep).transpose(1, 2)    # (B, 8, S, 96)

    # ⑤ 计算注意力（Flash Attention 或手动实现）
    if self.flash and seq_len > 1:
        output = F.scaled_dot_product_attention(
            xq, xk, xv, dropout_p=self.dropout, is_causal=True)
    else:
        scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(96)
        # 因果掩码：上三角填 -∞
        scores[:, :, :, -seq_len:] += torch.full(
            (seq_len, seq_len), float("-inf")).triu(1)
        output = F.softmax(scores.float(), dim=-1).type_as(xq) @ xv

    # ⑥ 合并多头，投影输出
    output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
    output = self.resid_dropout(self.o_proj(output))
    return output, past_kv

GQA 分组查询

8 个 Q 头共享 4 组 KV——用更少显存获得几乎相同的效果

# model_minimind.py 第 86-89 行
# repeat_kv：把 4 个 KV 头复制成 8 个
def repeat_kv(x, n_rep):
    # x: (batch, seq, 4, 96) → (batch, seq, 8, 96)
    bs, slen, num_kv_heads, head_dim = x.shape
    if n_rep == 1: return x
    return (x[:, :, :, None, :]
            .expand(bs, slen, num_kv_heads, n_rep, head_dim)
            .reshape(bs, slen, num_kv_heads * n_rep, head_dim))

repeat_kv 不是真正复制数据（expand 是视图操作，不占额外内存）。只是在计算时让每个 Q 头"看到"对应的 KV 头。4 组 KV 被 8 个 Q 头共享，显存减半。

因果掩码

当前 token 只能看到它之前的 token，不能"偷看"未来

语言模型是逐字预测的——预测第 3 个字时，不能偷看第 4 个字。因果掩码就是把"未来"的注意力权重强制归零。

      你  好  世  界
你 [  0  -∞  -∞  -∞ ]    softmax(-∞) = 0，完全忽略
好 [  0   0  -∞  -∞ ]    "好"能看到"你"和自己
世 [  0   0   0  -∞ ]    "世"能看到"你""好"和自己
界 [  0   0   0   0 ]    "界"能看到所有之前的字

# 源码第 129 行：
scores[:, :, :, -seq_len:] += torch.full(
    (seq_len, seq_len), float("-inf")).triu(1)
# triu(1) = 上三角矩阵（对角线以上为 1）
# 加到 scores 上，上三角全变 -∞
# 注意 triu(1) 中的 1：对角线保留为 0（自己看自己），
# 如果写成 triu(0) 则对角线也变 -∞，模型连自己都看不到

KV-Cache

缓存历史 KV，每步只算新 token

Flash Attention

分块计算，显存从 O(S²) 降到 O(S)

# 源码第 109 行：自动判断是否可用
self.flash = hasattr(torch.nn.functional,
    'scaled_dot_product_attention') and config.flash_attn

# 源码第 125-126 行：满足条件就用 Flash Attention
# 条件：可用 + 序列长度>1 + 首次前向（无KV-Cache）+ 无attention_mask
if self.flash and seq_len > 1 and past_key_value is None:
    output = F.scaled_dot_product_attention(
        xq, xk, xv, dropout_p=self.dropout, is_causal=True)

Part 4：前馈

注意力决定"关注谁"，前馈网络决定"怎么处理"。MoE 让不同 token 由不同专家处理。

FeedForward

SwiGLU 结构——Llama 系列使用的 FFN

定义：对每个位置独立做非线性变换，增强模型的表达能力。输入先升维（768→2432），经过激活函数，再降维回来（2432→768）。

白话：注意力决定"看谁"，FFN 决定"怎么想"。注意力是"看"的动作，FFN 是"思考"的动作。

比喻卸甲：在数学上，FFN 就是两个线性变换中间夹一个非线性激活函数。"看"和"思考"只是拟人化的辅助理解。

# model_minimind.py 第 136-146 行——整个 FFN 只有 6 行！
class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate_proj = nn.Linear(768, 2432, bias=False)  # 门控投影
        self.down_proj = nn.Linear(2432, 768, bias=False)  # 降回原维度
        self.up_proj   = nn.Linear(768, 2432, bias=False)  # 升维投影
        self.act_fn    = ACT2FN[config.hidden_act]         # SiLU 激活

    def forward(self, x):
        # SwiGLU = SiLU(gate_proj(x)) * up_proj(x)，然后 down 投影回去
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

SwiGLU 计算图

gate 和 up 两条路，门控机制让模型"选择性地通过信息"

SiLU：SiLU(x) = x × sigmoid(x)。比 ReLU 更平滑——ReLU 在 x<0 时梯度直接为 0，SiLU 在 x<0 时仍有微小梯度，训练更稳定。

为什么用 gate + up 两条路？ gate 路经过 SiLU 激活后输出 0~1 范围的"门控信号"，控制"哪些信息通过"；up 路提供原始内容。两者相乘实现"选择性过滤"。这比简单的 ReLU(Linear(x)) 表达能力更强。

MoE 混合专家

不同 token 由不同"专家"处理

定义：MoE（Mixture of Experts）把一个大的 FFN 拆成多个小的"专家"FFN，每次只激活少数几个。参数量增大了，但计算量几乎不变。

白话：像一个医院，全科医生什么都能看（Dense FFN），但专科医生更专业（MoE）。每个 token 来了先挂号（Router），然后被分到对应科室（Expert）。

比喻卸甲：在数学上，MoE 就是多个并行的 FFN，一个线性层（Router）输出每个 token 对各专家的偏好分数，只取 top-k 个专家计算并加权求和。

# model_minimind.py 第 148-154 行
class MOEFeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate = nn.Linear(768, config.num_experts, bias=False)  # 路由器
        self.experts = nn.ModuleList([
            FeedForward(config, config.moe_intermediate_size)
            for _ in range(config.num_experts)  # 默认 4 个专家
        ])

MoE 路由过程

Router 打分 → Top-1 选择 → 专家计算 → 加权合并

# model_minimind.py 第 156-176 行（简化版）
def forward(self, x):
    batch_size, seq_len, hidden_dim = x.shape
    x_flat = x.view(-1, hidden_dim)          # 展平: (N, 768)

    # ① 路由器给每个 token 对 4 个专家打分
    scores = F.softmax(self.gate(x_flat), dim=-1)  # (N, 4)

    # ② 选得分最高的 1 个专家（Top-1）
    topk_weight, topk_idx = torch.topk(scores, k=1, dim=-1)

    # ③ 只让选中的专家处理对应的 token
    y = torch.zeros_like(x_flat)
    for i, expert in enumerate(self.experts):
        mask = (topk_idx == i)               # 哪些 token 选了专家 i
        if mask.any():
            token_idx = mask.any(dim=-1).nonzero().flatten()
            weight = topk_weight[mask].view(-1, 1)
            y.index_add_(0, token_idx, expert(x_flat[token_idx]) * weight)

    # ④ 计算负载均衡损失
    load = F.one_hot(topk_idx, 4).float().mean(0)
    self.aux_loss = (load * scores.mean(0)).sum() * 4 * 5e-4
    return y.view(batch_size, seq_len, hidden_dim)

负载均衡损失

防止所有 token 涌向同一个专家

解决方案：aux_loss 惩罚负载不均。如果某个专家被选中太多（load 高）或得分太高（score 高），损失就变大。训练会自动纠正这种不均衡。

# aux_loss = N * sum(load_i * score_i) * coef
# load_i = 专家 i 被选中的比例
# score_i = 专家 i 的平均得分
# coef = 0.0005（很小，不干扰主损失）
#
# 为什么乘 load * score 而不是只看 load？
# 只看 load：专家被均分了但每个都不太确定 → 没解决问题
# load * score：惩罚"既被频繁选中又自信"的专家
# 这迫使 Router 的决策更分散、更均匀

Dense vs MoE

# 所有 token 经过同一个 FFN
# 参数量：~5.6M
# 每个 token 激活全部参数
# 优点：简单稳定
# 缺点：参数利用率低

# token 被 Router 分配给不同专家
# 参数量：~5.6M x 4 = 22.4M
# 每个 token 只激活 1/4 参数
# 优点：参数量大、计算量小
# 缺点：需要 aux_loss 均衡

MiniMind 默认 use_moe=False，使用 Dense FFN。在配置中把 use_moe 设为 True 即可启用 MoE。两种模式在 MiniMindBlock 中通过一个 if-else 切换，其他代码完全不变。

Part 5：组装

把所有零件组装成完整的 Transformer Block，堆叠 8 层，加上输出头，实现自回归生成。

Transformer块

Pre-Norm 残差结构——信息层层叠加

# model_minimind.py 第 178-194 行
class MiniMindBlock(nn.Module):
    def __init__(self, layer_id, config):
        super().__init__()
        self.self_attn = Attention(config)
        self.input_layernorm = RMSNorm(config.hidden_size)
        self.post_attention_layernorm = RMSNorm(config.hidden_size)
        self.mlp = FeedForward(config) if not config.use_moe \
                   else MOEFeedForward(config)

    def forward(self, hidden_states, position_embeddings, ...):
        residual = hidden_states                              # 保存副本
        hidden_states, present_kv = self.self_attn(
            self.input_layernorm(hidden_states),              # 先归一化再注意力
            position_embeddings, ...)
        hidden_states += residual                             # 残差连接

        hidden_states = hidden_states + self.mlp(
            self.post_attention_layernorm(hidden_states))     # 归一化+FFN+残差
        return hidden_states, present_kv

残差连接的意义

输出 = 输入 + 变化量——让信息有"捷径"传递

即使某一层没学到有用的东西（变化量 = 0），信息也不会丢失——直接从残差路径传过去。这就是恒等映射。

模型组装

Embedding + 8 层 Block + 最终 Norm（第 196-232 行）

# model_minimind.py 第 196-207 行
class MiniMindModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embed_tokens = nn.Embedding(6400, 768)     # 词嵌入
        self.dropout = nn.Dropout(config.dropout)
        self.layers = nn.ModuleList([
            MiniMindBlock(l, config) for l in range(8)]) # 8 层 Block
        self.norm = RMSNorm(768)                         # 最终归一化
        # 预计算 RoPE 的 cos/sin 缓冲区（不参与梯度更新）
        freqs_cos, freqs_sin = precompute_freqs_cis(
            dim=96, end=32768, rope_base=1e6, rope_scaling=config.rope_scaling)
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

forward 流程

数据从 Embedding 到最终 hidden_states 的完整旅程

# model_minimind.py 第 209-232 行（简化版）
def forward(self, input_ids, attention_mask=None,
            past_key_values=None, use_cache=False, **kwargs):
    batch_size, seq_length = input_ids.shape

    # ① 词嵌入：离散 ID → 连续向量
    hidden_states = self.dropout(self.embed_tokens(input_ids))  # (B, S, 768)

    # ② 取出当前位置的 RoPE cos/sin（支持 KV-Cache 偏移）
    start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
    position_embeddings = (
        self.freqs_cos[start_pos:start_pos + seq_length],
        self.freqs_sin[start_pos:start_pos + seq_length])

    # ③ 逐层通过 Transformer Block
    presents = []
    for layer, past_kv in zip(self.layers, past_key_values):
        hidden_states, present = layer(
            hidden_states, position_embeddings,
            past_key_value=past_kv, use_cache=use_cache)
        presents.append(present)

    # ④ 最终归一化 + 收集 MoE 的 aux_loss
    hidden_states = self.norm(hidden_states)
    aux_loss = sum([l.mlp.aux_loss for l in self.layers
                    if isinstance(l.mlp, MOEFeedForward)])
    return hidden_states, presents, aux_loss

CausalLM 头

在 MiniMindModel 之上加 lm_head 做词汇预测

# model_minimind.py 第 234-253 行
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
    def __init__(self, config):
        self.model = MiniMindModel(config)
        self.lm_head = nn.Linear(768, 6400, bias=False)  # 768 → 词汇表大小
        # 关键：embed 和 lm_head 共享权重！
        if config.tie_word_embeddings:
            self.model.embed_tokens.weight = self.lm_head.weight

    def forward(self, input_ids, labels=None, ...):
        hidden_states, past_kv, aux_loss = self.model(input_ids, ...)
        logits = self.lm_head(hidden_states)      # (B, S, 6400)

        if labels is not None:
            # 预测左移一位，和真实标签右移一位对齐
            loss = F.cross_entropy(
                logits[..., :-1, :].contiguous().view(-1, 6400),
                labels[..., 1:].contiguous().view(-1),
                ignore_index=-100)
        return MoeCausalLMOutputWithPast(
            loss=loss, aux_loss=aux_loss, logits=logits, ...)

损失计算的秘密

为什么 logits 取 [:-1] 而 labels 取 [1:]？

输入:  [你,  好,  世,  界]    (input_ids)
       ↓ embed → ×8层 → lm_head
logits:[好', 世', 界', ?' ]   模型对每个位置的预测
       取[:-1] ↑ ↑ ↑
labels:[好,  世,  界,  ! ]    真实的下一个字
              取[1:] ↑ ↑ ↑

对齐后：
logits[:-1] = [好', 世', 界']  vs  labels[1:] = [好, 世, 界]
→ 交叉熵衡量每个位置的预测与真实的差距

本质就是：模型在每个位置预测下一个字，损失衡量预测和真实下一个字的差距。ignore_index=-100 让 padding 和不需要计算损失的位置被忽略。

自回归生成

逐步生成 token 直到遇到 EOS（第 256-288 行）

自回归生成是模型推理的核心过程：每次前向传播只预测一个 token，把它拼到输入后面，再预测下一个。如此循环直到遇到结束符或达到最大长度。

# model_minimind.py 第 256-288 行（简化版）
@torch.inference_mode()
def generate(self, input_ids, max_new_tokens=8192,
             temperature=0.85, top_p=0.85, top_k=50,
             eos_token_id=2, repetition_penalty=1.0, ...):
    for _ in range(max_new_tokens):
        # ① 前向传播，只取最后一个 token 的 logits
        past_len = past_key_values[0][0].shape[1] if past_key_values else 0
        outputs = self.forward(input_ids[:, past_len:], ...)  # 只算新 token
        logits = outputs.logits[:, -1, :] / temperature       # 温度缩放

        # ② Repetition Penalty：惩罚已出现的 token
        if repetition_penalty != 1.0:
            for i in range(input_ids.shape[0]):
                seen = torch.unique(input_ids[i])
                logits[i, seen] = torch.where(
                    logits[i, seen] > 0,
                    logits[i, seen] / repetition_penalty,
                    logits[i, seen] * repetition_penalty)

        # ③ Top-K：只保留概率最高的 K 个
        # ④ Top-P：累积概率超过 P 后面的都去掉
        # ⑤ 从过滤后的分布中采样
        next_token = torch.multinomial(torch.softmax(logits, dim=-1), 1)
        input_ids = torch.cat([input_ids, next_token], dim=-1)
        if (next_token == eos_token_id).all(): break  # 遇到结束符就停
    return input_ids

采样策略详解

Temperature / Top-K / Top-P 三重过滤

# Top-K 实现（第 271-272 行）
logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')

# Top-P 实现（第 273-277 行）
sorted_logits, sorted_idx = torch.sort(logits, descending=True)
mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0  # 保留第一个
logits[mask.scatter(1, sorted_idx, mask)] = -float('inf')

重复惩罚

惩罚已生成的 token，防止重复循环

# model_minimind.py 第 268-270 行
if repetition_penalty != 1.0:
    for i in range(input_ids.shape[0]):
        seen = torch.unique(input_ids[i])
        score = logits[i, seen]
        # 正分数除以 penalty（变小），负分数乘以 penalty（变更负）
        logits[i, seen] = torch.where(
            score > 0, score / repetition_penalty,
            score * repetition_penalty)

为什么正负分数要分别处理？如果统一除以 penalty：正分数变小（好），但负分数会变大（变接近 0，反而增加被选中的概率）。所以正的除、负的乘，确保已出现的 token 在两个方向都被"推开"。

Part 6：总结

回顾全景，清点参数，动手验证。

参数量追踪

每一层有多少参数，它们花在哪了

维度变化总览

数据在模型中每一步的形状

# 假设输入 input_ids: (2, 5)，即 batch=2, seq_len=5
# 假设已缓存 past_len=3 个 token 的 KV-Cache

# ===== MiniMindModel.forward =====
input_ids                              # (2, 5)
embed_tokens(input_ids)                # (2, 5, 768)

# ===== Attention.forward（单层）=====
q_proj(x)                              # (2, 5, 768)
view → xq                              # (2, 5, 8, 96)    8个Q头
k_proj(x) → view                       # (2, 5, 4, 96)    4个KV头
QK-Norm                                # (2, 5, 8, 96) / (2, 5, 4, 96)
RoPE                                   # 维度不变
KV-Cache 拼接后 xk                     # (2, 8, 4, 96)    seq=3+5=8
repeat_kv(xk, 2)                       # (2, 8, 8, 96)    复制成8个头
xq @ xk^T                              # (2, 8, 5, 8)     注意力分数
softmax @ xv                           # (2, 8, 5, 96)
reshape → o_proj                       # (2, 5, 768)

# ===== FeedForward.forward =====
gate_proj(x)                           # (2, 5, 2432)
up_proj(x)                             # (2, 5, 2432)
SiLU(gate) * up                        # (2, 5, 2432)
down_proj                              # (2, 5, 768)

# ===== MiniMindForCausalLM.forward =====
lm_head(hidden_states)                 # (2, 5, 6400)     logits
logits[:, -1, :]                       # (2, 6400)         最后一个token的预测

常见错误排查

排查方法：在 forward 的每一步加 print(x.shape)，和上面的维度追踪表对照。最常见的错误是把 view(bsz, seq_len, n_heads, head_dim) 中的 n_heads 和 head_dim 写反。

练习：维度

x = input                               # (2, 5, 768)
xq = self.q_proj(x)                     # (2, 5, 768)
xq = xq.view(2, 5, 8, 96)              # (2, 5, 8, 96)
xk = self.k_proj(x).view(2, 5, 4, 96)  # (2, 5, 4, 96)
xk = cat([past_k, xk], dim=1)          # (2, 8, 4, 96)    3+5=8
xq = xq.transpose(1, 2)                # (2, 8, 5, 96)
xk = repeat_kv(xk, 2).transpose(1, 2)  # (2, 8, 8, 96)
scores = xq @ xk.transpose(-2, -1)     # (2, 8, 5, 8)
output = softmax(scores) @ xv           # (2, 8, 5, 96)
output = output.transpose(1,2).reshape()# (2, 5, 768)

练习：理解

练习：修改

# head_dim = 512 // 8 = 64
# intermediate_size = ceil(512 * pi / 64) * 64 = ceil(25.13) * 64 = 26 * 64 = 1664
# q_proj: 512 * 512 = 0.26M
# k_proj: 512 * 256 = 0.13M
# FFN: 3 * 512 * 1664 = 2.56M
# 单层: ~3.1M, 4 层: ~12.4M + embed 3.3M ≈ 15.7M

# MoE 参数 = Router(768*4) + 4 * FFN(5.6M) ≈ 22.4M
# 每层从 ~7.37M 变为 ~24.2M（约 3 倍）
# 但每次只激活 1 个专家，计算量 ≈ 1 * FFN + Router ≈ 不变
# 这就是 MoE 的核心优势：参数多、计算少

延伸阅读