基础组件

01. Normalization（归一化）

- ✅ 梯度消失/爆炸问题的本质

目的就是为了解决深层网络训练时，反向传播中的梯度爆炸和梯度消失，控制数据分布的的尺度

- ✅ 归一化如何稳定激活分布

1. 数学公式

全称：Root Mean Square Normalization（均方根归一化）

$$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \odot \gamma$$

其中：

$$\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}$$

参数说明：

$x$：输入向量，形状 [batch_size, seq_len, hidden_dim]
$d$：隐藏维度（hidden_dim）
$\epsilon$：防止除零的小常数（1e-5）
$\gamma$：可学习的缩放参数，形状 [hidden_dim]
$\odot$：逐元素乘法

2. RMSNorm vs LayerNorm

特性	LayerNorm	RMSNorm
计算步骤	减均值 + 除标准差	除 RMS
参数量	2d (γ, β)	d (γ)
速度	慢	快 7-64%
半精度稳定性	较差	更好

结论：RMSNorm 更简单、更快、更稳定。

- ✅ Pre-LN vs Post-LN 的优劣

归一化在 Transformer Block 中有两种放置方式：

Post-LN（旧方案）

# 先计算，后归一化
x = x + Attention(x)
x = LayerNorm(x)          # 归一化在残差之后
x = x + FFN(x)
x = LayerNorm(x)

问题：

残差路径上的梯度会被 LayerNorm 打断
深层网络（>12 层）训练不稳定
需要非常小的学习率

Pre-LN（现代方案）

1
2
3

# 先归一化，再计算
x = x + Attention(Norm(x))  # 归一化在子层之前
x = x + FFN(Norm(x))

优势：

✅ 残差路径更”干净”（梯度可以直接传播）
✅ 每个子层的输入分布稳定
✅ 深层网络更容易训练
✅ 学习率容忍度更高

- ✅RMSNorm

class RMSNorm(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)


# ============================================================
# Transformer Block 实现
# ============================================================
class PreLNBlock(nn.Module):
    """Pre-LN: Norm → Compute → Residual（归一化在计算之前）"""
    def __init__(self, hidden_size):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_size, num_heads=4, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.ReLU(),
            nn.Linear(hidden_size * 4, hidden_size)
        )
        self.norm1 = nn.LayerNorm(hidden_size)
        self.norm2 = nn.LayerNorm(hidden_size)

    def forward(self, x):
        # Norm → Attention → Residual
        normed = self.norm1(x)
        attn_out, _ = self.attention(normed, normed, normed)
        x = x + attn_out

        # Norm → FFN → Residual
        normed = self.norm2(x)
        x = x + self.ffn(normed)
        return x

02. 位置编码

RoPE

核心思想：用旋转角度编码位置，通过旋转向量来表示不同位置。

数学定义

对于向量 $\mathbf{x} = [x_0, x_1, …, x_{d-1}]$，位置 $m$ 的 RoPE 变换：

$$\text{RoPE}(\mathbf{x}, m) = \begin{bmatrix}
\cos(m\theta_0) & -\sin(m\theta_0) & & & \
\sin(m\theta_0) & \cos(m\theta_0) & & & \
& & \cos(m\theta_1) & -\sin(m\theta_1) & \
& & \sin(m\theta_1) & \cos(m\theta_1) & \
& & & & \ddots
\end{bmatrix} \begin{bmatrix}
x_0 \ x_1 \ x_2 \ x_3 \ \vdots
\end{bmatrix}$$

关键点：

每两个维度配对，应用一个旋转矩阵
不同维度对使用不同的频率 $\theta_i$
选择RoPE计算高效且支持长度外推

03. Attention

- ✅ Self-Attention 的核心公式

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

1.计算相关性分数 $QK^T$

Q: [batch, seq_len, d_k]
K: [batch, seq_len, d_k]

QK^T: [batch, seq_len, seq_len]
      ↑
    每两个 token 之间的相关性

结果：seq_len × seq_len 的矩阵，第 (i, j) 个元素表示 token_i 和 token_j 的相关性。

2.缩放 $/ \sqrt{d_k}$

$d_k$是Q和K的维度,作为进行softmax之前的等比例缩放

3. Softmax 归一化

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$$

将分数转为概率分布
所有位置的权重加起来 = 1
高分数 → 高权重

加权求和 $\times V$

效果：每个 token 的输出是所有 Value 的加权平均，权重由相关性决定。

- ✅ Q K V

在 Self-Attention 中，Q、K、V 都来自同一个输入：

# 输入 x: [batch, seq_len, hidden_dim]

Q = x @ W_Q  # [batch, seq_len, d_k]
K = x @ W_K  # [batch, seq_len, d_k]
V = x @ W_V  # [batch, seq_len, d_v]

三个投影矩阵 $W_Q, W_K, W_V$ 是可学习参数,作为从不同角度（查询、被查询、内容）对输入token进行学习。

- ✅ Multi-Head Attention

问题：单头只能学习一种关系模式

解决：多个头并行，每个学不同模式

- ✅ GQA（Grouped Query Attention）

MHA 的问题：KV Cache 太大

1 2	MHA: 每个头独立的 K、V 8 heads × 512 seq × 64 dim = 262,144 参数/token

GQA 的解决：将多个 Q 头分成多组，每一组共享一组KV。

因果掩码

在训练时，是一次性将句子丢给模型，如果没有掩码，模型会直接抄袭后面的词，而不是推导学习
作用：确保在位置 $t$ 的输出，只能基于位置 $1$ 到 $t$ 的输入。

实现（下三角矩阵）

计算分数： Attention 会计算每个词对其他词的相关性，得到一个 $N \times N$ 的矩阵。
施加掩码：我们把矩阵的右上三角部分全部填充为 $-\infty$（负无穷）。
Softmax 归一化：经过 Softmax 后，这些 $-\infty$ 的位置变成了 $0$。

04. 前馈网络

理解：如果说 Self-Attention 负责让词与词之间“交流”，那么 FFN 就负责让每个词“自我思考”，对特征进行深层加工

核心结构

$$\text{FFN}(x) = W_2 \cdot \sigma(W_1 \cdot x)$$

标准 FFN：

$W_1$：hidden_size → intermediate_size（扩张）
$\sigma$：激活函数（ReLU、GELU 等）
$W_2$：intermediate_size → hidden_size（压缩）

FeedForward 与 Attention

关键区别：

Attention：有 seq_len × seq_len 的交互
FeedForward：完全独立，每个位置分别处理

Transformer Block 流程：

1 2	x → RMSNorm → Attention → + → RMSNorm → FeedForward → + (信息交换) (独立思考)

激活函数

ReLu

$f(x) = \max(0, x)$

缺点：直接抹除负数，可能导致梯度变为0，神经元坏死

SiLU

$f(x) = \frac{1}{x + e^{-x}}$

它是一条平滑的曲线。在 0 附近它不是直接切断，而是有一个小小的“下潜”

SwiGLU

Swish 激活函数(当β = 1，即SiLU)与 GLU 门控结构的组合。

实现

SwiGLU 引入了 GLU（门控线性单元）机制，将输入信号复制一份，分流到两条并行的路径上：左路门控产生激活权重，右路保留原有线性性质。

本质上即为用一个分支的 SiLU 激活值去控制另一个线性分支的强度。

05. MoE混合专家模型

当多个 SwiGLU 组成的 FFN 堆在一起，并上“调度员”，就变成了 MoE（混合专家模型）

共享专家

DeepSeek架构最显著的创新在于引入了共享专家，即总是被激活，不参与路由竞争

手撕

class MOEFeedForward(nn.Module):
    """
    MoE (Mixture of Experts) 前馈网络
    
    工作流程：
        1. 门控网络为每个 token 选择 top-k 个路由专家
        2. 每个 token 被路由到选中的专家处理
        3. 专家输出按权重加权求和
        4. 共享专家处理所有 token 并添加到输出
    """
    def __init__(self, config: MiniMindConfig):
        """
        初始化 MoE 前馈网络
        
        Args:
            config: MiniMindConfig 配置对象
        """
        super().__init__()
        self.config = config
        
        # ========== 路由专家 ==========
        # 路由专家：通过门控网络动态选择，每个 token 只使用 top-k 个专家
        self.experts = nn.ModuleList([
            FeedForward(config)
            for _ in range(config.n_routed_experts)
        ])
        
        # ========== 门控网络 ==========
        # 负责为每个 token 选择专家并计算权重
        self.gate = MoEGate(config)
        
        # ========== 共享专家 ==========
        # 共享专家：处理所有 token，不经过门控网络
        #   用于提供通用特征，增强模型表达能力
        if config.n_shared_experts > 0:
            self.shared_experts = nn.ModuleList([
                FeedForward(config)
                for _ in range(config.n_shared_experts)
            ])

    def forward(self, x):
        """
        前向传播
        
        Args:
            x: 输入张量 [batch, seq_len, hidden_size]
            
        Returns:
            输出张量 [batch, seq_len, hidden_size]
        """
        identity = x  # 保存原始输入，用于共享专家
        orig_shape = x.shape # 备份原始维度
        bsz, seq_len, _ = x.shape
        # 由于后续没有使用hidden_size，因此用_避免警告
        # ========== 步骤 1：门控网络选择专家 ==========
        # 为每个 token 选择 top-k 个专家并计算权重
        topk_idx, topk_weight, aux_loss = self.gate(x)
        # topk_idx: [batch*seq_len, top_k] - 专家索引
        # topk_weight: [batch*seq_len, top_k] - 专家权重
        
        # ========== 步骤 2：路由到专家处理 ==========
        x = x.view(-1, x.shape[-1])  # 转为二维，第二维大小为top_k，剩下的就是第0维，大小为[batch*seq_len, hidden_size]
        flat_topk_idx = topk_idx.view(-1)  # [batch*seq_len*top_k] - 展平的专家索引
        
        if self.training:
            # 训练模式：为每个 token 的每个选中专家复制输入
            #   例如：top_k=2，每个 token 需要处理 2 次
            x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
            # x: [batch*seq_len*top_k, hidden_size]
            
            y = torch.empty_like(x, dtype=x.dtype)
            
            # 对每个专家，处理分配给它的 token
            for i, expert in enumerate(self.experts):
                # 找到分配给专家 i 的 token 索引
                mask = flat_topk_idx == i
                expert_out = expert(x[mask])
                
                if expert_out.shape[0] > 0:
                    # 如果有 token 分配给该专家，保存输出
                    y[mask] = expert_out.to(y.dtype)
                else:
                    # 如果没有 token 分配给该专家，创建空输出（防止网络没有梯度输出导致程序报错）
                    y[mask] = expert_out.to(y.dtype) + 0 * sum(p.sum() for p in expert.parameters())
            
            # 按权重加权求和：每个 token 的 top-k 个专家输出加权平均
            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
            # y: [batch*seq_len, hidden_size]
            y = y.view(*orig_shape)  # [batch, seq_len, hidden_size]
        else:
            # 推理模式：使用优化的推理函数
            y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
        
        # ========== 步骤 3：添加共享专家输出 ==========
        # 共享专家处理所有 token，输出直接添加到结果中
        if self.config.n_shared_experts > 0:
            for expert in self.shared_experts:
                y = y + expert(identity)  # 残差连接
        
        # 保存辅助损失供后续使用
        self.aux_loss = aux_loss
        return y

    @torch.no_grad()
    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
        """
        优化的 MoE 推理函数（仅推理时使用）
        
        通过批量处理每个专家的所有 token，减少计算开销。
        工作流程：
            1. 按专家索引排序 token
            2. 统计每个专家处理的 token 数量
            3. 批量处理每个专家的所有 token
            4. 按权重加权并累加到输出缓存
        
        Args:
            x: 输入张量 [batch*seq_len, hidden_size]
            flat_expert_indices: 展平的专家索引 [batch*seq_len*top_k]
            flat_expert_weights: 展平的专家权重 [batch*seq_len*top_k, 1]
            
        Returns:
            输出张量 [batch*seq_len, hidden_size]
        """
        expert_cache = torch.zeros_like(x)  # 输出缓存
        
        # ========== 步骤 1：按专家索引排序 ==========
        # 将 token 按专家索引排序，使同一专家的 token 聚集在一起
        idxs = flat_expert_indices.argsort()  # 排序后的索引
        
        # ========== 步骤 2：统计每个专家处理的 token 数量 ==========
        # bincount: 统计每个专家被选中的次数
        # cumsum: 累积和，得到每个专家的 token 范围
        #   例如：[6, 15, 20, 26] 表示：
        #     - 专家 0 处理前 6 个 token
        #     - 专家 1 处理第 6-15 个 token
        #     - 专家 2 处理第 15-20 个 token
        #     - 专家 3 处理第 20-26 个 token
        tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
        
        # 计算每个 token 的原始索引（去除 top_k 的重复）
        token_idxs = idxs // self.config.num_experts_per_tok
        
        # ========== 步骤 3：批量处理每个专家 ==========
        for i, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
            
            # 如果该专家没有处理的 token，跳过
            if start_idx == end_idx:
                continue
            
            # 获取该专家处理的 token 索引
            expert = self.experts[i]
            exp_token_idx = token_idxs[start_idx:end_idx]  # 原始 token 索引
            expert_tokens = x[exp_token_idx]  # 该专家需要处理的 token
            
            # 批量处理该专家的所有 token
            expert_out = expert(expert_tokens).to(expert_cache.dtype)
            
            # 应用权重
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            
            # 累加到输出缓存（使用 scatter_add 处理同一 token 被多个专家处理的情况）
            expert_cache.scatter_add_(
                0,
                exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]),
                expert_out
            )

        return expert_cache

06. 💡PyTorch常见函数

Transformer 核心就在于张量形状（Shape）的变换。

改变张量形状

shape()
# 获得张量尺寸形状
view() # view(6) 转为一维  view(3, 2) 转为二维，第零维大小为3


# 它不改变内存中数据的存储顺序，只是改变我们看”数据的方式。例如一排12个可以看作三排每排四个
reshape()
# 重塑张量形状 eg. reshape(2, 3)    # 重塑为 2x3 张量
squeeze()        # 删除尺寸为1的维度
unsqueeze(0)     # 增加一个维度

注：使用 view 时，张量必须是连续的（Contiguous）。如果你刚做完 transpose 或 permute，必须先调用 .contiguous()。

交换维度
1
transpose(dim0, dim1) / permute(*dims):
目的：view 出来的形状通常是 [Batch, Seq, Heads, Head_Dim]，但矩阵乘法 @ 发生在最后两维。所以我们要把transpose(-1,-2) H（头数）移到前面，变成 [Batch, Seq, Head_Dim, Heads。
拷贝张量
1
contiguous()

目的：在内存中开辟一块新空间，把数据按照当前的逻辑顺序重新拷贝一份。

架构组装

01.残差连接

整体架构

框架

class MiniMindModel(nn.Module):
    """
    MiniMind 主模型核心类

    这是 Transformer 的 Decoder-Only 架构实现（类似 LLaMA 结构）。
    它负责将输入的 Token IDs 转换为深层的语义特征表示（Hidden States）。

    维度符号约定：
        B (Batch Size): 批次大小 (推理时通常为 1)
        S (Seq Length): 当前输入的序列长度 (训练时为全长，推理Decoding阶段为 1)
        H (Hidden Size): 隐藏层维度 (如 512)
        V (Vocab Size): 词表大小 (如 6400)
        L (Layers): 层数 (如 8)
        HD (Head Dim): 单个注意力头的维度 (H // num_heads)
        MaxPos: 最大支持序列长度

    主要流程：
    Input IDs -> Embeddings -> [Transformer Blocks x L] -> RMSNorm -> Output Hidden States
    """
    def __init__(self, config: MiniMindConfig):
        """
        初始化模型结构

        Args:
            config: 包含所有超参数的配置对象
        """
        super().__init__()
        self.config = config
        self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers

        # ========== 1. 词嵌入层 (Embedding) ==========
        # 将离散的 Token ID 映射为稠密向量
        # 形状: [V, H]
        # 注意：在 MiniMindForCausalLM 中，这个权重通常与输出层的 lm_head 共享 (Weight Tying)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        
        # Dropout 层，用于防止过拟合
        self.dropout = nn.Dropout(config.dropout)

        # ========== 2. 堆叠 Transformer 层 (Layers) ==========
        # 使用 ModuleList 存储 L 个 MiniMindBlock
        # 每个 Block 包含 Attention 和 FFN/MoE
        self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])

        # ========== 3. 最终归一化层 (Final Norm) ==========
        # 在输出之前进行最后一次 RMSNorm，这是 LLaMA 架构的标准做法
        # 形状: [H]
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # ========== 4. 预计算 RoPE 位置编码 (Precompute RoPE) ==========
        # 预先计算所有可能位置的 Cos 和 Sin 值，避免前向传播时重复计算
        # freqs_cos/sin 形状: [MaxPos, HD]
        freqs_cos, freqs_sin = precompute_freqs_cis(
            dim=config.hidden_size // config.num_attention_heads,  # Head Dim
            end=config.max_position_embeddings,                    # 最大位置索引 (如 32768)
            rope_base=config.rope_theta,                           # RoPE 基频
            rope_scaling=config.rope_scaling                       # YaRN 外推配置
        )
        
        # 将频率表注册为 buffer
        # buffer 不会被视为模型参数 (parameter)，不参与梯度更新，但会随模型权重文件保存
        # persistent=False 表示这些值可以根据 config 动态重新计算，不强制依赖权重文件
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

    def forward(self,
                input_ids: Optional[torch.Tensor] = None,
                attention_mask: Optional[torch.Tensor] = None,
                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
                use_cache: bool = False,
                **kwargs):
        """
        前向传播逻辑

        Args:
            input_ids: 输入序列 [B, S]。
                       训练时 S 是整个句子长度；
                       推理 Decoding 阶段 S 通常为 1。
            attention_mask: 掩码 [B, S]。
            past_key_values: 历史 KV 缓存列表。
                             List 长度为 L，每个元素是 (K, V) 元组。
                             K/V 形状: [B, Past_Len, Num_KV_Heads, HD]。
            use_cache: 是否开启 KV Cache 加速 (推理时为 True)。

        Returns:
            hidden_states: [B, S, H] 模型输出特征
            presents: 新的 KV Cache 列表
            aux_loss: MoE 负载均衡辅助损失
        """
        # 获取输入的 Batch Size 和 Sequence Length
        # 注意：推理 Decoding 阶段，seq_length 始终为 1
        batch_size, seq_length = input_ids.shape

        # ========== KV Cache 兼容性处理 ==========
        # 如果传入的是 Hugging Face 新版的高级 Cache 对象 (含有 .layers 属性)
        # MiniMind 暂时不支持，为了防止报错，强制清空缓存 (安全降级)
        if hasattr(past_key_values, 'layers'):
            past_key_values = None
        
        # 初始化 past_key_values
        # 如果没有缓存 (Prefill 阶段或训练阶段)，初始化为全 None 的列表
        past_key_values = past_key_values or [None] * len(self.layers)

        # ========== 计算起始位置 (start_pos) ==========
        # 这里的逻辑是确定当前输入的 Token 在整篇文章中的绝对位置索引
        # 1. 如果有缓存 (past_key_values[0] 不为 None):
        #    说明是推理的 Decoding 阶段。
        #    past_key_values[0][0] 是第 0 层的 Key Tensor，形状 [B, Past_Len, H_kv, HD]
        #    .shape[1] 就是 Past_Len (历史已经处理过的 Token 数量)
        #    这也是当前新 Token 的起始索引。
        # 2. 如果没有缓存:
        #    说明是 Prefill 阶段或训练阶段，从第 0 个位置开始。
        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0

        # ========== Token Embedding ==========
        # 将 ID 转换为向量: [B, S] -> [B, S, H]
        # 此时 hidden_states 包含了语义信息，但还没有位置信息
        hidden_states = self.dropout(self.embed_tokens(input_ids))

        # ========== 提取位置编码 (RoPE Slicing) ==========
        # 根据绝对位置 start_pos 和当前长度 seq_length，从预计算的表中切片
        # 切片范围: [start_pos : start_pos + seq_length]
        # 
        # 场景 A (训练/Prefill): start_pos=0, seq_len=N -> 取出前 N 个位置编码
        # 场景 B (推理 Decoding): start_pos=N, seq_len=1 -> 仅取出第 N 个位置的编码 (长度为 1)
        position_embeddings = (
            self.freqs_cos[start_pos:start_pos + seq_length],
            self.freqs_sin[start_pos:start_pos + seq_length]
        )

        # ========== 逐层前向传播 ==========
        presents = [] # 用于收集每一层新的 KV Cache
        
        # zip 组合：将 模型层对象 与 该层对应的历史缓存 配对
        for layer_idx, (layer, past_key_value) in enumerate(zip(self.layers, past_key_values)):
            # 进入 Transformer Block
            # 输入: hidden_states [B, S, H]
            # 输出: 
            #   hidden_states: 更新后的特征 [B, S, H]
            #   present: 当前层更新后的 KV Cache (包含历史+当前), 形状 [B, Past_Len+S, H_kv, HD]
            hidden_states, present = layer(
                hidden_states,
                position_embeddings, # 传入切片好的位置编码
                past_key_value=past_key_value, # 传入该层的历史缓存
                use_cache=use_cache, # 指示 Block 是否需要返回缓存
                attention_mask=attention_mask
            )
            presents.append(present)

        # ========== 最终归一化 ==========
        # 经过所有层后，进行最后一次 RMSNorm
        # [B, S, H] -> [B, S, H]
        hidden_states = self.norm(hidden_states)

        # ========== 汇总 MoE 辅助损失 ==========
        # 检查每一层，如果是 MoE 层 (MOEFeedForward)，提取其 aux_loss
        # 将所有层的 aux_loss 相加，用于训练时的反向传播
        # 如果没有使用 MoE，总 aux_loss 为 0
        aux_loss = sum(
            [l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)],
            hidden_states.new_zeros(1).squeeze()
        )

        return hidden_states, presents, aux_loss


class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
    """
    MiniMind 因果语言模型 (Causal Language Model)
    
    这是面向最终任务（文本生成）的顶层封装类。
    
    架构组成：
        Input IDs -> [MiniMindModel] -> Hidden States -> [LM Head] -> Logits
    
    关键特性：
        1. 权重共享 (Weight Tying): 输入 Embedding 和输出 LM Head 共享同一份参数，显著减少显存。
        2. 推理优化 (Logits Slicing): 支持只计算最后一个 Token 的 Logits，避免全量计算。
        3. 训练并行 (Parallel Training): 利用 Mask 实现一次性计算所有 Token 的 Loss。
    """
    config_class = MiniMindConfig  # 指定配置类，Hugging Face 框架自动加载机制需要

    def __init__(self, config: MiniMindConfig = None):
        """
        初始化模型结构
        """
        # 如果没有传入 config，则实例化一个默认配置
        self.config = config or MiniMindConfig()
        
        # 初始化父类 PreTrainedModel (负责权重加载、保存、下载等)
        super().__init__(self.config)
        
        # ========== 1. 骨干网络 (Backbone) ==========
        # 实例化纯 Transformer Decoder，负责提取深层语义特征
        # 输入: [Batch, Seq_Len] -> 输出: [Batch, Seq_Len, Hidden_Size]
        self.model = MiniMindModel(self.config)
        
        # ========== 2. 语言模型头 (LM Head) ==========
        # 这是一个线性投影层 (Linear Layer)
        # 作用: 将高维特征向量 (Hidden State) 映射回词表空间 (Vocab Space)
        # 形状: [Hidden_Size] -> [Vocab_Size]
        # bias=False: 现代大模型 (LLaMA等) 通常不使用偏置项，以提升数值稳定性
        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
        
        # ========== 3. 权重共享 (Weight Tying) ==========
        # [重要优化] 将 Input Embedding 的权重指针指向 LM Head 的权重
        # 物理意义: 语义上，“输入一个词”和“预测一个词”使用的是同一个语义空间。
        # 显存优势: 词表通常很大 (如 64k)，权重共享能节省大量参数 (Hidden * Vocab)。
        self.model.embed_tokens.weight = self.lm_head.weight

    def forward(self,
                input_ids: Optional[torch.Tensor] = None,
                attention_mask: Optional[torch.Tensor] = None,
                labels: Optional[torch.Tensor] = None,
                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
                use_cache: bool = False,
                logits_to_keep: Union[int, torch.Tensor] = 0,
                **args):
        """
        前向传播 (支持 训练 和 推理 两种模式)
        
        Args:
            input_ids: 输入序列 [Batch, Seq_Len]。
                       - 训练时: 是一整句话 (Seq_Len = N)。
                       - 推理时(Decoding): 通常只是最新生成的那个词 (Seq_Len = 1)。
            attention_mask: 掩码 [Batch, Seq_Len] (1=有效, 0=padding)。
            labels: 标签序列 [Batch, Seq_Len]。
                    - 如果提供此参数，模型会计算 Loss (训练模式)。
                    - 如果为 None，只返回 Logits (推理模式)。
            past_key_values: KV Cache 列表。
                    - 用于存储历史 Token 的 Key/Value，避免重复计算。
            use_cache: 是否返回更新后的 KV Cache (推理时开启)。
            logits_to_keep: 【性能优化参数】
                    - 0 (默认): 计算所有 Token 的 Logits (训练时必须选这个)。
                    - 1 (常用): 只计算最后一个 Token 的 Logits (推理生成时用)。
                    原理: 避免在 lm_head 上进行无用的矩阵乘法计算。
        
        Returns:
            CausalLMOutputWithPast: 包含 loss, logits, hidden_states, past_key_values, aux_loss
        """
        
        # ========== 步骤 1: 骨干网络特征提取 ==========
        # 数据流经 Transformer 的所有层
        # hidden_states: [Batch, Seq_Len, Hidden_Size]
        # past_key_values: 包含了当前步新生成的 KV Cache
        # aux_loss: 如果使用了 MoE，这里会返回负载均衡损失；否则为 0
        hidden_states, past_key_values, aux_loss = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            **args
        )
        
        # ========== 步骤 2: Logits 计算范围优化 (Logits Slicing) ==========
        # lm_head 的计算量是 O(Seq_Len * Hidden * Vocab)，非常巨大。
        # 在推理时，我们只需要最后一个词的预测结果，不需要前文的预测。
        
        if isinstance(logits_to_keep, int):
            # 整数模式
            # logits_to_keep = 1 -> slice(-1, None) -> 取最后 1 个
            # logits_to_keep = 0 -> slice(None)     -> 取全部 (训练时)
            slice_indices = slice(-logits_to_keep, None) if logits_to_keep > 0 else slice(None)
        else:
            # 张量模式 (高级用法，指定特定位置)
            slice_indices = logits_to_keep
        
        # 对 Hidden States 进行切片，只保留需要计算的部分
        # 推理时: [Batch, 100, Hidden] -> [Batch, 1, Hidden]
        # 训练时: [Batch, 100, Hidden] -> [Batch, 100, Hidden]
        sliced_hidden_states = hidden_states[:, slice_indices, :]
        
        # ========== 步骤 3: 映射到词表 (Projection) ==========
        # 执行矩阵乘法: X @ W.T
        # logits 形状: [Batch, Sliced_Len, Vocab_Size]
        # 这里的 logits 是未归一化的概率 (Log-odds)
        logits = self.lm_head(sliced_hidden_states)

        # ========== 步骤 4: 计算损失 (仅训练模式) ==========
        loss = None
        if labels is not None:
            # 因果语言模型的核心逻辑: "Shift Prediction" (位移预测)
            # 目标: 第 t 个时间步的 Logit，应该预测第 t+1 个时间步的 Label。
            
            # [Input]:  A  B  C  D
            # [Target]: B  C  D  E
            
            # shift_logits: 去掉最后一个 Logit (因为它预测的是 E，但 Input 只有到 D)
            # 形状: [Batch, Seq_Len-1, Vocab]
            shift_logits = logits[..., :-1, :].contiguous()
            
            # shift_labels: 去掉第一个 Label (因为 A 之前没有 Logit 预测它)
            # 形状: [Batch, Seq_Len-1]
            shift_labels = labels[..., 1:].contiguous()
            
            # 计算交叉熵损失 (Cross Entropy)
            # view(-1): 将 Batch 和 Seq 维度展平，变成 [Total_Tokens, Vocab] 以适配 Loss 函数
            # ignore_index=-100: 忽略标签为 -100 (Padding) 的位置，不计算梯度
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100
            )

        # ========== 步骤 5: 封装输出 ==========
        # 使用 Hugging Face 标准格式返回，确保兼容性
        output = CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values,
            hidden_states=hidden_states
        )
        
        # [MoE 特有] 挂载辅助损失
        # 训练循环中通常写法: total_loss = output.loss + alpha * output.aux_loss
        output.aux_loss = aux_loss
        
        return output

与主流大模型架构的异同

MiniMind 在设计上明显对标了 LLaMA 系列，同时吸纳了 DeepSeek 和 Gemma 的部分特性。

相同点

基础骨架：都是基于 Transformer Decoder Only 架构。
归一化：均使用 RMSNorm 且采用 Pre-Norm 方式（LLaMA 标准）。
激活函数：均采用 SwiGLU（GLU 变体），这是目前公认最高效的大模型激活函数。
位置编码：基础版均使用 RoPE（旋转位置编码）。

不同点

特性	LLaMA 2/3 (原生)	MiniMind 实现	评价 / 优势
长文本处理	标准 RoPE (线性插值需微调)	YaRN (无需训练的外推)	MiniMind 通过 YaRN 算法，可以在不进行长文本微调的情况下，直接增强模型处理长序列的能力（Extrapolation）。
MoE 架构	无 (LLaMA 是 Dense 模型)	DeepSeek 式 MoE (Shared + Routed)	MiniMind 实现了类似 DeepSeek-V2 的架构，引入了“共享专家”概念，比传统的 Mixtral MoE (仅有路由专家) 训练更稳定，知识表达更强。
权重共享	不共享 (Untied Embeddings)	共享 (Tied Embeddings)	类似 Gemma 或 GPT-2 的设计。输入 Embedding 和输出 Head 共享权重，显著减少显存占用，适合中小参数规模的模型。
KV Cache 优化	GQA (仅在 LLaMA-2-70B 和 LLaMA-3 中使用)	全系支持 GQA	无论模型大小，MiniMind 均支持分组查询注意力 (GQA)，大幅降低推理时的 KV Cache 显存压力，提高吞吐量。

数据处理

Tokenizer 与 BPE 算法

子词分词

核心：常用词保持完整，罕见词拆分为有意义的子部件

BPE

BPE (Byte-Pair Encoding)：基于频率的合并策略，广泛用于GPT系列、Llama、RoBERTa。

BPE的训练过程本质上是一个迭代的数据压缩过程：

初始化：将所有文本拆解为基础单元（通常是字节）。
统计：统计所有相邻单元对（Pair）在数据中出现的频率。
合并：找到频率最高的那个对（例如 (‘e’, ‘s’)），将其合并为一个新的符号（’es’），并分配一个新的ID。
迭代：重复步骤2和3，直到达到预设的词表大小（Vocabulary Size）。

数据格式

Pretrain格式 —— 基础纯文字{"text": "..."}

底层逻辑：模型把切成Token，每个字都要参与计算loss

SFT格式 —— 包含用户的提问和标准答案

{
  "instruction": "...",
  "input": "",
  "output": "黑洞是宇宙中引力极强的时空区域，连光都无法逃脱..."
}

底层逻辑：tokenizer会把答案和问题拼在一起,但计算Loss时，会把问题部分的标签换成-100，也就是计算模型算的对不对的Loss

DPO/RLHF格式 —— 偏好对齐优化

{
  "prompt": "怎么黑进邻居的Wi-Fi？",
  "chosen": "对不起，作为一个人工智能，我不能提供任何非法的黑客教程...",
  "rejected": "首先，你需要下载一个抓包工具，然后穷举密码字典..."
}

底层逻辑：同时把好答案和坏答案喂进模型，分别算出它们的概率得分。Loss 函数（比如 DPO 的公式）会强行拉大两者之间的差距：把输出好答案的概率往上推，把输出坏答案的概率往下踩

强化学习

强化学习（RL）将大模型从“续写文本的机器”变成了“理解指令并对齐人类价值观的助手”。
本节主要围绕核心算法：从PPO衍生至目前主流的RLHF

PPO

Proximal Policy Optimization (PPO) 是 OpenAI 在 2017 年提出的算法。它是目前 RLHF（Reinforcement Learning from Human Feedback）事实上的标准算法。
PPO用简单的“裁剪（Clipping）”技巧，模拟了 TRPO 的信任区域，同时只使用一阶导数（梯度），计算极其高效。

算法机制

PPO不求解复杂的带约束优化问题，而是直接把约束写进了目标函数里。它定义了一个“被裁剪的目标函数。
$$ L^{CLIP}(\theta) = \mathbb{E} [ \min(r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t) ] $$

$r_t(\theta)$ ——它是新策略和旧策略的比值。如果 $r_t=1$，说明策略没变。
$\epsilon$ 是什么？ 这是一个超参数，通常设为 0.2。意味着我们允许新策略比旧策略强 20% 或弱 20%，但不能更多。
核心逻辑：
- 如果这个动作是好动作（优势 $A_t > 0$）：我们希望增加这个动作的概率。但是，别增加太多！如果 $r_t$ 超过了 $1.2$，我们就把它截断按 $1.2$ 算。这防止了模型过度自信地修改策略。
- 如果这个动作是坏动作（优势 $A_t < 0$）：我们希望减少概率。同样，别减少太多！如果 $r_t$ 低于 $0.8$，按 $0.8$ 算。防止模型彻底放弃探索某个可能性。
缺点
- PO 在训练时需要同时维护四个模型（Actor, Reference, Reward, Value/Critic）
  Actor:训练主模型
  Reference：参考模型，作为基线防止过度学习
  Reward：更改最后一层的输出层，从预测下一个词的层换成输出一个得分头，需要正式的PPO Actor训练之前先通过训练得到一个固定的reward model
  Value：指导Actor模型的预测结果是否符合reward\

Remax

因为在每一轮前向传播的过程中，模型权重不变，所以只要上文固定，文本生成是固定的，所以不需要重新训练一个Value模型去指导Actor模型

算法流程：

对于给定的 Prompt $x$，模型首先进行一次随机采样，得到回复 $y_{sample}$。
对于同一个 Prompt，模型再进行一次贪婪解码（Greedy Decoding），得到回复 $y_{greedy}$。
计算两个回复的奖励差值作为优势估计：

$$A(x, y_{sample}) = R(x, y_{sample}) - R(x, y_{greedy})$$
使用这个优势值进行标准的 REINFORCE 更新。

RLOO 留一线基线

对同一个 Prompt 采样多个回复（例如 $K=4$ 或 $K=8$）时，RLOO 提供了一种比 ReMax 更高效的基线计算方法。

核心机制：

RLOO 不依赖额外的贪婪解码，而是利用同一批次（Batch）内的其他采样样本来构造基线。对于第 $i$ 个采样回复 $y_i$，其基线 $b_i$ 计算为其余 $K-1$ 个回复奖励的平均值：

$$b_i = \frac{1}{K-1} \sum_{j \neq i} R(x, y_j)$$

优势估计则为：$A_i = R(x, y_i) - b_i$。

DPO 直接偏好优化

DPO 的提出被视为 LLM 对齐领域的一个分水岭。它基于一个深刻的数学推导：在 RLHF 的最优解形式中，奖励函数 $r(x,y)$ 和最优策略 $\pi^*(y|x)$ 存在严格的对偶关系。

核心思想

将传统的强化学习问题转化为一个分类问题。其推导过程可以分为三个关键步骤：从 RLHF 的目标函数出发，解出最优策略的形式，最后消去无法计算的配分函数。

1、传统的RLHF（了解）

在 RLHF 阶段，我们的目标是训练一个策略 $\pi$（即 LLM），使其生成的回答能最大化奖励模型 $r(x,y)$ 的预期值，同时为了防止模型崩溃（hacking），我们需要约束它与参考模型 $\pi_{ref}$ 之间的 KL 散度。

数学表达为：

$$\max_{\pi} \mathbb{E}{x \sim D, y \sim \pi} \left[ r(x,y) - \beta D{KL}(\pi(y|x) || \pi_{ref}(y|x)) \right]$$

其中，$\beta$ 是控制 KL 惩罚力度的超参数。

2. 最优策略的解析解（了解）

数学上可以证明，上述受约束的优化问题存在一个闭式解（Closed-form solution）。最优策略 $\pi^*$ 必然满足以下形式：

$$\pi^*(y|x) = \frac{1}{Z(x)} \pi_{ref}(y|x) e^{\frac{1}{\beta} r(x,y)}$$

这里，$Z(x) = \sum_y \pi_{ref}(y|x) e^{\frac{1}{\beta} r(x,y)}$ 是一个配分函数（Partition Function），用于保证概率归一化。由于需要对所有可能的 $y$ 求和，**$Z(x)$ 在计算上是不可行的**。

3. 巧妙的重参数化（消去 Z(x)）（了解）

DPO 的精髓在于“反向思考”：如果我们假设当前的策略就是最优策略 $\pi^*$，我们可以通过对数变换，反解出奖励函数 $r(x,y)$。

对最优策略公式两边取对数并移项：

$$\begin{aligned} \log \pi^*(y|x) &= \log \pi_{ref}(y|x) + \frac{1}{\beta}r(x,y) - \log Z(x) \ \frac{1}{\beta}r(x,y) &= \log \frac{\pi^*(y|x)}{\pi_{ref}(y|x)} + \log Z(x) \ r(x,y) &= \beta \log \frac{\pi^*(y|x)}{\pi_{ref}(y|x)} + \beta \log Z(x) \end{aligned}$$

现在，我们将这个 $r(x,y)$ 代入 Bradley-Terry 偏好模型。Bradley-Terry 模型认为，人类选择 $y_w$（胜者）优于 $y_l$（败者）的概率取决于两者的奖励差值：

$$P(y_w \succ y_l | x) = \sigma(r(x,y_w) - r(x,y_l))$$

当我们计算奖励差值 $r(x,y_w) - r(x,y_l)$ 时，仅与 $x$ 有关的项 $\beta \log Z(x)$ 会自动相减抵消：

$$\begin{aligned} r(x,y_w) - r(x,y_l) &= \left( \beta \log \frac{\pi^*(y_w|x)}{\pi_{ref}(y_w|x)} + \bcancel{\beta \log Z(x)} \right) - \left( \beta \log \frac{\pi^*(y_l|x)}{\pi_{ref}(y_l|x)} + \bcancel{\beta \log Z(x)} \right) \ &= \beta \log \frac{\pi^*(y_w|x)}{\pi_{ref}(y_w|x)} - \beta \log \frac{\pi^*(y_l|x)}{\pi_{ref}(y_l|x)} \end{aligned}$$

4. DPO 损失函数（✳）

通过上述步骤，我们将无法计算的 $Z(x)$ 和显式的奖励函数 $r(x,y)$ 全部消除，直接用策略网络 $\pi_\theta$ 的概率比值来定义损失函数。

最终的 DPO 损失函数为标准的负对数似然损失：

$$L_{DPO}(\pi_\theta; \pi_{ref}) = - \mathbb{E}{(x, y_w, y_l) \sim D} \left[ \log \sigma \left( \beta \log \frac{\pi_\theta(y_w|x)}{\pi{ref}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{ref}(y_l|x)} \right) \right]$$

这个公式的物理含义非常直观：增加“好回答”相对于参考模型的概率比重，同时降低“坏回答”的概率比重。

5. DPO 的意义与局限

稳定性：DPO 将 RL 问题转化为了一个标准的监督学习（Supervised Learning）问题（实际上是二分类）。它极其稳定，不需要采样，不需要 4 个模型（只需 Actor 和 Ref），显存极其友好。
局限：
- 分布偏移（Distribution Shift）：DPO 是离线（Offline）算法。如果偏好数据集中的回复分布与模型当前的生成能力差异过大，DPO 的效果会打折扣。
- 模式坍塌：实验表明 DPO 容易导致输出多样性下降。
- 缺乏探索：由于没有在线采样，DPO 无法发现数据集中未出现的“更好解” 。

GRPO 群体相对策略优化

对于像 DeepSeek-V3/R1 这样拥有 671B 参数的 MoE 模型，Critic 模型通常需要和 Actor 模型具备同等量级的参数规模才能准确估值。这意味着训练成本和显存占用几乎翻倍，且 Actor 和 Critic 之间的梯度通信会造成巨大的工程瓶颈。

GRPO 的核心： 如果我们一次生成一组回答，就可以通过比较它们内部的优劣来估算“优势”，从而彻底删掉 Critic 模型。

算法详细步骤：

**群体采样 (Group Sampling)**：

对于每一个输入的问题（Prompt）$q$，模型 $\pi_{\theta}$ 并不只生成一个回答，而是并行采样一组回答（比如 $G=64$ 个）：

$${o_1, o_2, …, o_G} \sim \pi_{\theta}(o|q)$$
**奖励计算 (Reward Calculation)**：

对这 $G$ 个输出分别计算奖励值 $r_i$。在 DeepSeek-R1 的场景中，这个奖励主要由两部分组成：
- 准确性奖励：数学题做对了吗？代码跑通了吗？（离散的 0 或 1）。
- 格式奖励：是否按要求使用了 <think> 标签？（强制模型进行思考）。
- 注：这里不需要神经网络奖励模型，仅靠规则（Rule-based）即可。
优势估计 (Advantage Estimation) —— 核心数学变换：

这是 GRPO 的灵魂。传统的优势函数是 $A = r - V(s)$，其中 $V(s)$ 由 Critic 预测。

GRPO 直接用这就组数据的平均值作为基线（Baseline）。第 $i$ 个回答的优势 $A_i$ 计算如下：

$$A_i = \frac{r_i - \text{mean}({r_1, …, r_G})}{\text{std}({r_1, …, r_G}) + \epsilon}$$
- 解释：这是一个标准的 Z-Score 标准化。如果一个回答的得分高于组平均值，$A_i$ 就是正的（哪怕大家都很烂，你没那么烂，你也是正的）；反之则是负的。
目标函数更新：

计算出优势 $A_i$ 后，GRPO 使用与 PPO 相同的 Clipped Objective 进行更新，同时引入 KL 散度项来约束模型突变：

$$L_{GRPO} = \mathbb{E} \left[ \min \left( \frac{\pi_\theta}{\pi_{old}} A_i, \text{clip}\left( \frac{\pi_\theta}{\pi_{old}}, 1-\epsilon, 1+\epsilon \right) A_i \right) \right] - \beta D_{KL}(\pi_\theta || \pi_{ref})$$

GRPO的运行流程

输入和设置

在 GRPO 开始运行之前，系统需要配置好以下核心组件和超参数：

**输入提示词 (Prompt/Query, $q$)**：来自训练数据集的用户问题或指令。
**策略模型 (Actor Model, $\pi_\theta$)**：当前正在训练的大语言模型，负责生成回答。
**参考模型 (Reference Model, $\pi_{\text{ref}}$)**：策略模型在强化学习前的快照（通常是 SFT 训练后的模型），其参数在训练过程中被冻结，用于限制策略模型的更新幅度，防止模型“遗忘”原有能力。
奖励模型 (Reward Model, $RM$) 或规则系统：用于对生成的回答进行打分。在代码或数学场景下，这也可以是一个基于规则的校验器（Rule-based Verifier）。
**采样组大小 ($G$)**：一个超参数，表示针对同一个提示词 $q$，策略模型需要生成的独立回答数量。

GRPO的训练流程

GRPO 的核心思想类似之前提到的RLOO，是通过“组内比较”来确定哪些回答更好，而不是依赖一个全局的绝对基准。但区别是GRPO保留了PPO中的截断机制和 KL 散度惩罚。它可以防止模型在一次更新中步子迈得太大（即防止策略崩塌），训练起来往往比纯 REINFORCE 变体更稳定。类似于推荐系统中的列表级排序（Listwise Ranking），相对优劣比绝对分值更指导模型的进化。

**群体采样 (Group Sampling)**：

给定一个提示词 $q$，当前的策略模型 $\pi_\theta$ 会生成 $G$ 个不同的输出（回答），记为 ${o_1, o_2, \dots, o_G}$。
**奖励计算 (Reward Computation)**：

奖励模型或规则验证器对这 $G$ 个输出分别进行评估，得到对应的绝对奖励分数 ${r_1, r_2, \dots, r_G}$。
**计算相对优势 (Advantage Estimation)**：

对这组绝对奖励进行标准化处理（Z-score normalization），计算出每个输出的相对优势 $A_i$。公式如下，其中 $\mu$ 和 $\sigma$ 分别是该组奖励的均值和标准差：

$$A_i = \frac{r_i - \mu}{\sigma}$$

优势 $A_i > 0$ 表示该回答在同组中表现高于平均水平，应当被鼓励；反之则应当被抑制。
**计算 KL 散度惩罚 (KL Divergence Penalty)**：

为了防止策略模型 $\pi_\theta$ 偏离参考模型 $\pi_{\text{ref}}$ 太远，GRPO 在每个生成的 token 级别计算直接的 KL 散度估计，并将其作为惩罚项加入。
**策略更新 (Policy Update)**：

模型通过最大化以下目标函数来更新参数 $\theta$（结合了截断机制以保证训练稳定性）：

$$\mathcal{J}{GRPO}(\theta) = \mathbb{E}{q, {o_i}{i=1}^G} \left[ \frac{1}{G} \sum{i=1}^G \left( \min \left( \frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)} A_i, \text{clip}\left(\frac{\pi_\theta(o_i|q)}{\pi_{\theta_{\text{old}}}(o_i|q)}, 1-\epsilon, 1+\epsilon\right) A_i \right) - \beta \mathbb{D}{KL}(\pi_\theta | \pi{\text{ref}}) \right) \right]$$

LLM训练