「源码解析」合集：本文与 TRPO 解析：从数学原理到代码实现 同属一个系列，关注「论文/数学原理 → 工程实现」的拆解。

Tokenizer

Step 1：分词
Word-level（问题：未登录词+词表爆炸） -> Character-level（序列过长+寓意稀疏）+ SubWord-level
目前主流subword分词算法：

• BPE（Byte Pair Encoding）：基于频率的合并算法，逐步合并最频繁的字节对，生成新的子词单元。用于GPT系列、Llama、RoBERTa等。
• WordPiece：基于概率的分词算法，使用最大似然估计来选择合适的子词单元。用于BERT、ALBERT等。
• Unigram：基于概率的剪枝策略，从大词表删减，用于SentencePiece（ALBert,T5等）。

Step Into BPE：
Byte-Pair Encoding（BPE）的训练过程本质是一个迭代的数据压缩过程，旨在通过合并最频繁的字节对来构建一个更紧凑的词表。以下是BPE训练过程的详细步骤：

• Initialization 先拆：把文本拆解为基础单元，通常是字节
• Count Frequencies 统计频率：统计相邻单元对在数据中出现的频率
• Merge 最频繁的单元对：找到出现频率最高的单元对，并将其合并成一个新的单元并把所有的旧单元对ID替换为新ID，例如将“h”和“e”合并成“he”
• Iteration 迭代：重复Count和Merge，直到达到预设的词表大小

# core function
## 统计相邻单元对的频率
def get_stats(ids):
    """
    input：ids, type list. 当前的TokenID列表
    return：counts, type dict. 统计的单元对频率, (id1, id2) -> frequency
    """
    counts = {}
    for pair in zip(ids, ids[1:]):
        # (id1, id2) (id2, id3) (id3, id4) ...
        counts[pair] = counts.get(pair, 0) + 1
    return counts

## 执行合并，贪婪算法
def merge(ids, pair, idx):
    """
    input：ids, type list. 当前的TokenID列表
           pair, type tuple. 要合并的单元对 (id1, id2)
           idx, type int. 新单元的ID
    return：newids, type list. 合并后的TokenID列表
    """
    newids = []
    i = 0
    while i < len(ids):
        if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
        # 找到单元对的位置
            newids.append(idx)  
            i += 2
        else:
        # 别的不处理
            newids.append(ids[i])
            i += 1
    return newids

## 主函数 train
def train(text, vocab_size, verbose=False):
    """
    input：text, type str. 训练文本
           vocab_size, type int. 预设的词表大小
           verbose, type bool. 是否打印训练过程
    return：vocab, type dict. 训练得到的词表，id -> token
    """
    # vocab_size是个超参数。
    # 太小导致序列过长，模型推理慢，无法捕捉长距离依赖。
    # 太大导致Embedding矩阵（vocab_size x hidden_dim）参数量激增，增加训练负担；且由于稀有词频次太低，其Embedding可能训练不充分（Undertrained）。
    # 如果有两个pair的频率相同，为了保证Tokenizer的可复现性（Deterministic），在工业级实现中通常会规定：当频率相同时，按照Pair中字符的字典序（Lexicographical order）进行选择。
    assert vocab_size >= 256, "词表大小必须至少为256，因为初始单元是字节"
    num_merges = vocab_size - 256

    # 1. 初始化词表：将文本拆解为字节，并统计初始词表
    text_bytes = text.encode('utf-8')
    ids = list(text_bytes)  

    merges = {}
    print(f"Original length: {len(ids)}")

    for i in range(num_merges) :
        # 2. 统计频率
        stats = get_stats(ids)
        if not stats: 
            break # 没有可合并的就退出

        # 3. 找到频率最高的单元对 分配id并合并
        pair = max(stats, key=stats.get)
        idx = 256+i
        ids = merge(ids, pair, idx)

        merges[pair] = idx
        if verbose:
            print(f"Merge {i+1}/{num_merges}: {pair} -> {idx} (count: {stats[pair]})")  
              
    return merges

# 推理阶段的编码
# 推理的时候按照训练时确定的优先级顺序进行合并
# e.g. name -> n a m e -> (a, m) -> n am e -> (n, am) -> nam e -> (nam, e) -> name
def encode(text, merges):
    """
    input：text, type str. 输入文本
           merges, type dict. 训练得到的词表，(id1, id2) -> new_id
    return：ids, type list. 编码后的TokenID列表
    """
    ids = list(text.encode('utf-8'))

    while len(ids) >= 2:
        stats = get_stats(ids)
        pair_to_merge = None
        min_rank = float('inf')
        # 寻找在merges中存在且ID最小的单元对
        # 因为ID越小的单元对优先级越高

        # 这里可以做加速
        for pair in stats:
            if pair in merges:
                rank = merges[pair]
                if rank < min_rank:
                    min_rank = rank
                    pair_to_merge = pair
        if pair_to_merge is None:
            break  # 没有可合并的单元对了
        ids = merge(ids, pair_to_merge, min_rank)
    return ids

# 解码
# 注意处理无效字节
# e.g. [256, 270] + vocab -> [b'he', b'lo'] -> b'hello' -> 'hello'
def decode(ids, vocab):
    """
    input：ids, type list. TokenID列表
           vocab, type dict. 词表，id -> token
    return：text, type str. 解码后的文本
    """
    tokens = b"".join(vocab[idx] for idx in ids)
    text = tokens.decode('utf-8', errors='replace')  
    return text

BPE, WordPiece, Unigram 对比

WordPiece 和 BPE 都是自底向上合并，主要区别在于如何选择合并那一对

• BPE：选择频次最高的 => 最大化数据压缩比
• wordPiece：选择合并后能让悬链数据似然度/Likelihood增加最多的（基于贝叶斯概率公式）=> 增强词的内在关联性

Unigram 是基于概率图模型的自顶向下剪枝。

• 构建一个巨大的词表
• E-step/期望步 计算出最优切分路径，P(s)=P(s1)···P(sm)的最大值
• M-step/最大化步 重新计算每个子词出现的概率（？为什么会不一样）
• 计算如果从词表里移除某个子词s，总似然度L会下降的值ΔL
• 剪枝，移除ΔL最小的token
• 重复上述过程，直到词表大小到预期

特性	BPE	WordPiece	Unigram
构建方向	自底向上 (Bottom-up)	自底向上 (Bottom-up)	自顶向下 (Top-down)
核心指标	频率 (Frequency)	似然度增益 (PMI)	似然度损失 (Loss)
分词确定性	确定性	确定性	概率性 (可采样多种结果)
计算复杂度	较低	高	较高
优势场景	生成式大语言模型 (GPT, Llama等)	语言理解模型 (BERT等)	需要正则化与抗噪任务(T5, ALBERT)

再发展

预分词与Regex

PE算法本身是“盲目”的。如果不加干预，它可能会跨越标点和单词的边界进行合并。所以在运行BPE合并之前，需要先用**正则表达式（Regex）**将文本切分为一个个基础的“单词块”。BPE只能在这些块内部进行合并，不能跨块。

• GPT-2的regex：r"""'(?:[sdmt]|ll|ve|re)|?p{L}+|?p{N}+|?[^sp{L}p{N}]+|s+(?!S)|s+"""。在处理多空格、大小写时不佳
• GPT-4的regex：r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^rnp{L}p{N}]?p{L}+|p{N}{2,}|[^rnp{L}p{N}]?[^sp{L}p{N}]+[rn]*|s*[rn]+|s+(?!S)|s+"""

字节级BPE的“可视化”实现

bytes_to_unicode 是 BPE 算法能够处理任意字节流（包括 Unicode 乱码、二进制文件、Emoji 等）的函数。但有一些不可见的字节（换行符、制表符）会造成干扰和调试问题。
解决方案：全射映射（Bijective Mapping），构建一个可逆的映射表：可见字符不动，不可见字符映射到Unicode256号之后的可见字符区域。

• e.g. 空格，0x20 -> Ġ (U+0120)

词表大小

GPT-2 50k
GPT-4 100k
Llama 128K

特殊token

Tiktoken (GPT-4) 的处理方式：Tiktoken 要求用户显式传递 allowed_special 参数。相当于自定义token进行保护。
如何在自定义tokenizer中处理？需要在开头先对特殊token进行正则匹配

def encodeWithSpecialToken(text, special_tokens):
    """
    input：special_tokens, type map. 
           text, type str. 输入文本
    return：ids, type list. 编码后的TokenID列表
    """
    special_pattern = create_pattern(special_tokens.keys())

    # 把text分为普通部分和特殊部分
    splits = re.split(special_pattern, text)

    final_ids = 
    for part in splits:
        if part in special_tokens:
            # 直接追加id
            final_ids.append(special_tokens[part])
        else:
            # 其他用BPE
            final_ids.extend(bpe_encode(part))
    return final_ids

生产中使用的

库	作者	语言	算法	特点	模型
Tiktoken	OpenAI	Rust	BPE	极速	GPT-3.5, GPT-4
SentencePiece	Google	C++	BPE, Unigram	无损处理	Llama, ALBERT, T5
Tokenizers	HuggingFace	Rust	All	大一统 功能最全	BERT, RoBERTa, Mistral

未来

1. Token-free Models (Byte-level Transformers)：研究者正在探索直接在字节层面（MegaByte, MambaByte）进行训练。虽然序列长度增加了4倍，但随着Flash Attention等线性注意力机制的发展，这正在成为可能。这将彻底消灭Tokenizer带来的所有偏见和问题。
2. 多模态融合：随着GPT-4o等模型的出现，Tokenizer不仅要处理文本，还要处理图像Patch和音频帧。未来的Tokenizer将是“万物皆Token”的统一体。

minimind实现Tokenizer

train_tokenizer.py

代码中把词表大小设为了 6400
注意给特殊token预留了位置

Minimind的设计目录

整体架构图：

Minimind 参数仅26M，对齐了 Llama 3、DeepSeek-V2 等前沿大模型的设计规范。MiniMind 采用的是标准的 Decoder-Only Transformer 架构。Decoder-Only 架构的核心任务是“Next Token Prediction”（预测下一个词），决定了数据在模型中的流动方式是单向的，即当前的 Token 只能“看到”过去的 Token，而不能“看到”未来的 Token。

架构类型	工作原理与核心区别	核心优点	主要适用任务	代表模型
Encoder-Only	双向可见：同时看到上下文的所有词，将输入文本转化为深层语义表示。	理解能力极强：能够全方位捕捉词语在特定语境下的双向语义关联。	文本分类、情感分析、命名实体识别、阅读理解提取	BERT, RoBERTa, ALBERT
Decoder-Only	单向自回归：从左到右逐词预测下一个词，后面的词看不见前面的词（Mask机制）。	生成能力极强：最适合开放式生成，且在大规模参数下展现出极强的涌现能力和泛化性。	文本生成、自由对话、逻辑推理、代码生成	GPT系列 (GPT-3/4), LLaMA, Qwen
Encoder-Decoder	先理解后生成：Encoder 负责双向理解输入，Decoder 结合 Encoder 的信息单向生成输出。	兼顾理解与生成：完美契合“输入一个完整序列，输出另一个序列”且长度不一致的场景。	机器翻译、文本摘要、复杂指令微调任务	T5, BART, 原始Transformer

MiniMindConfig
总观：从参数规模（hidden_size=768, layers=8）来看，这是一个极小参数量（可能在几千万到一两亿之间）的微型模型。同时，它的架构极为先进，完全对标当前最前沿的工业界大模型标准（GQA + RoPE YaRN + MoE + RMSNorm + SiLU + FlashAttention）。

class MiniMindConfig(PretrainedConfig):
    model_type = "minimind"
    def __init__(
        self, 
        hidden_size=768,        
        num_hidden_layers=8,    
        use_moe=False, 
        **kwargs):
        super().__init__(**kwargs)
        # ----- 基础架构参数 ----- #
        # 维度 d_model，决定了模型的"宽度"
        self.hidden_size = hidden_size  
        # 层数 L，决定了模型的"深度"                    
        self.num_hidden_layers = num_hidden_layers
        # 激活函数          
        self.hidden_act = kwargs.get("hidden_act", 'silu') 
        # 前馈神经网络（FFN）的中间层维度，这里为了确保维度是 64 的整数倍，以实现 GPU 底层张量计算的硬件对齐和加速 
        self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64) 
        # RMSNorm（均方根归一化）的平滑常数，防止分母为零       
        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)    
        self.dropout = kwargs.get("dropout", 0.0)

        # ----- attention 机制 ----- #
        # 默认开启 FlashAttention 加速技术
        self.flash_attn = kwargs.get("flash_attn", True)
        # Q(uery) 注意力头 数量
        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
        # K(ey) V(alue) 头 数量
        # 这里 4 小于 8，说明模型使用了 GQA（Grouped-Query Attention 分组查询注意力
        self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
        # 每个注意力头的维度 768 // 8 = 96
        self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)

        # ----- 长文本与位置编码 ----- #
        # 最大上下文长度为 32K Token
        self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
        # RoPE 旋转位置编码的基础频率参数
        # 1,000,000是为了支持超长上下文
        self.rope_theta = kwargs.get("rope_theta", 1e6)
        self.rope_scaling = {
            "beta_fast": 32,
            "beta_slow": 1,
            "factor": 16,
            "original_max_position_embeddings": 2048,
            "attention_factor": 1.0,
            "type": "yarn"
        } if self.inference_rope_scaling else None
        #用于推理时动态扩展长度的策略配置。
        #如果开启，会使用 YaRN 在不重新微调的情况下支持超出原有训练长度的文本。
        self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)

        # ----- 词表与 Embedding ----- #
        # 词表大小
        self.vocab_size = kwargs.get("vocab_size", 6400)
        # 定义了句子开头 (BOS) 和句子结束 (EOS) 的特殊 Token ID
        self.bos_token_id = kwargs.get("bos_token_id", 1)
        self.eos_token_id = kwargs.get("eos_token_id", 2)
        # 决定输入层的 Embedding 权重和输出层的分类头 (LM Head) 权重是否共享
        # 开启共享可以节省大量模型参数。
        self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True)

        # ----- MoE ----- #
        # 只有当 use_moe=True 时生效)
        self.use_moe = use_moe
        # 专家网络数量 4个
        self.num_experts = kwargs.get("num_experts", 4)
        # 每次路由时，每个 Token 只激活 1 个专家（常见的配置是总计 8 专家激活 2 个）
        self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
        # 对被选中的专家输出概率进行归一化
        self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
        # 路由器的辅助损失系数
        self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)
        # 设置MoE中每个“专家”网络内部的隐藏层维度大小
        self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)

Embedding & Positional Encoding

基于Transformer架构的大语言模型是基于连续向量空间的并行计算系统 => 嵌入 Embedding & 位置编码 Positional Encoding
embedding：文本 => 向量
Positional Encoding：句子 => 坐标

位置编码的演进

绝对位置编码 Absolute Position Encoding, APE

1. 正弦位置编码（Sinusoidal PE）
   开山之作《Attention Is All You Need》中使用的基于三角函数的固定位置编码。其核心思想是利用不同频率的正弦和余弦波来为每个位置生成一个唯一的指纹（Fingerprint）。对于位置 $pos$ 和维度 $i$，编码公式为

   $$P{E}_{(pos,2i)}=\sin (\frac{pos}{10000^{\frac{2i}{d_{model}}}}),P{E}_{(pos,2i+1)}=\cos (\frac{pos}{10000^{\frac{2i}{d_{model}}}})$$

   维度$i=0$对应最高频（High Frequency），波长最短，能够区分相邻位置的微小变化；随着$i$增大，频率降低，波长变长，用于区分更宏观的距离。

2. 可学习的绝对位置编码（Learned APE）

   (GPT-2和BERT)
   初始化了一个大小为$L_{max}\times D$的可学习矩阵，即$L_{max}$个向量，跟随参数一起训练。但外推性差，如果推理最大长度超过训练的$L_{max}$，直接报错或失效。

相对位置编码 Relative Positional Encoding, RPE

1. 偏置相加法（Bias Addition）

   不执着于将position信息加到Embedding上，而是在计算Attension Score $Q{K}^T$时加入一个表示相对距离的可学习bias项$B_{i-j}$

   T5模型进一步使用了对数分桶（Logarithmic Bucketing）策略，即对近距离的位置赋予精确的独立偏置，对远距离的位置共享同一个桶的偏置。

2. ALiBi (Attention with Linear Biases)

   直接在Attention Score矩阵上减去一个与距离成正比的惩罚项：

   $${\text{Score}}_{i,j}=q_i\cdot k_j-m\cdot |i-j|$$

   m是一个特定于Attention Head的斜率系数。
   外推性很强，即使训练时只看512长度，推理时也能在数千长度上保持不崩。但它假设距离越远关注度越低（线性衰减），可能限制了模型捕捉长距离依赖的能力，并且无法优雅融入KV Cache的压缩和优化。

旋转位置编码（RoPE，Rotary Positional Embedding）

RoPE的核心洞察在于：通过绝对位置的旋转，自然诱导（Induce）出相对位置的内积性质。
关键结论：最终的内积结果中，位置信息仅以$m-n$的形式出现。这意味着：尽管我们对每个Token进行了绝对位置的旋转，但它们之间的相互作用（Attention）完全取决于它们的相对距离。 这一性质被称为平移不变性（Translation Invariance），它是RoPE能够统领江湖的根本原因。

在实际的Transformer中，Embedding维度$d$通常很大（如4096）。RoPE采用“分而治之”的策略，将$d$维向量切分为$d/2$个二维子空间。

特性	APE (绝对位置编码)	RoPE (旋转位置编码)	ALiBi (线性偏置注意力)
核心思想	为每个位置生成一个独立的位置向量（可学习或固定正弦波），直接与输入词向量相加。	通过旋转矩阵对Query和Key向量进行变换，在绝对位置的计算中自然体现相对位置的距离。	不在输入层添加位置信息，而是在计算注意力分数（Attention Score）时，根据两个Token的距离直接加上一个线性负偏置。
作用位置	模型输入层（Input Embedding阶段）	注意力计算层（作用于Q和K向量）	注意力计算层（直接作用于Softmax之前的Attention矩阵）
位置类型	纯绝对位置	形式上是绝对位置，本质上体现相对位置	纯相对位置
长度外推性	较差。几乎无法处理超过训练长度的序列。	较好。可通过动态NTK感知等插值技术（如在长上下文模型中）较好地实现长度扩展。	极佳。天生具备优秀的长度外推能力，无需微调即可处理比训练时更长的序列。
计算开销	极小（通常只是加法或查表操作）	中等（需要进行复数乘法或旋转矩阵乘法）	极小（标量加法操作，但需要针对不同注意力头设置不同的斜率）
代表模型	原版Transformer、BERT、GPT-2、GPT-3	LLaMA、Qwen、ChatGLM、PaLM	BLOOM、MPT、Baichuan (部分版本)

虽然RoPE理论上支持无限长度，但在实际应用中，当推理长度超过训练长度（例如训练4k，推理8k）时，模型性能会崩塌（PPL爆炸）。这被称为外推性故障（Extrapolation Failure）。

为了解决这个问题，Meta的研究人员（Chen et al., 2023）提出了线性内插（Linear Interpolation）。思路非常简单粗暴：如果外面的世界太危险（OOD），那我们就把外面的世界缩小塞进已知的世界里。

然而，线性内插并非没有代价。它引发了分辨率（Resolution）缺失的问题。当将位置除以 s 时，相邻两个Token之间的旋转角度差 Δθ 也变成了原来的 1/s。对于低频分量（转得慢的维度），这可能影响不大。但对于高频分量，甚至是中频分量，这种压缩会导致相邻Token在向量空间中“靠得太近”，难以区分。

NTK-Aware Scaling：频域视角的觉醒

NTK理论揭示了深度神经网络在学习过程中的频谱偏差（Spectral Bias）：网络倾向于优先学习低频函数，而难以学习高频剧烈变化的函数 。 在RoPE的语境下：

• 低维度（高频 θ ）：对应文本中的局部高频变化。由于波长短，它们在Embedding空间中变化剧烈。根据NTK理论，如果我们将这些高频信号进一步压缩（如PI所做的），网络将极难捕捉到这些微小的位置差异。
• 高维度（低频 θ ）：对应文本中的长距离低频变化。这些信号波长很长，稍微压缩一下并不会丢失太多信息。

NTK-Aware Scaling的核心思想是：对不同频率的维度应用不同程度的缩放。Base Change 不再直接除以 s，而是通过改变RoPE的基频参数 base（通常是10000）来实现这种非线性缩放。这种变换在数学上等效于：随着维度 d 的增加，缩放因子从1逐渐平滑过渡到 s。

• 高频维度（低维）：不插值（或极少插值）。因为它们负责局部位置，必须保持高分辨率，让模型能分清邻居。而且高频分量旋转得快，本来在训练中就覆盖了整个圆周，外推并不是大问题。
• 低频维度（高维）：强插值。因为它们负责全局位置，旋转得慢，如果不压缩，稍微走远一点就OOD了。（为什么容易 OOD？ 由于其波长极长（旋转极慢），在训练的 Context Window 内，这些维度往往连“半圈”都没转完（例如只覆盖了 0 到 10 度的扇区）。一旦推理长度超过训练长度，旋转角度就会进入模型从未见过的数值区域（例如 20 度），导致严重的分布外（OOD）问题。）

传统的静态NTK会固定一个最大的扩展长度，很大缺点是：即使你只输入一段很短的文本，它依然会应用缩放，这会导致高频分辨率受损，从而降低短文本的表现。而动态NTK会在推理阶段，根据当前输入文本的实际长度（$L_{\text{current}}$）来动态计算缩放倍数 $s$。
它的公式为：

$$s=\max (1,L_{\text{current}}/L_{\text{train}})$$

（其中 $L_{\text{train}}$ 是模型原本在训练时支持的最大长度）

• 当输入较短时（$L_{\text{current}}\le L_{\text{train}}$）：缩放倍数 $s=1$，也就是说完全不缩放。此时模型保留了原汁原味的精度，表现如同未做任何修改一样。
• 当输入变长时（$L_{\text{current}}$ 逐渐增加并超过训练长度）：缩放力度开始平滑介入，以非线性的方式对旋转位置编码（RoPE）的维度进行压缩。

YaRN (Yet another RoPE extensioN)

NTK-Aware Scaling 在处理超长上下文（如100k+）时仍存在理论缺陷。YaRN的提出，旨在修正这些缺陷，成为目前最完善的RoPE扩展框架。它结合了NTK-by-parts（分段NTK）和熵/温度调节（Temperature Scaling）。

NTK-by-parts

YaRN 认为简单的NTK插值不够精准，于是把位置编码的特征维度按频率划分为三个频段，并使用“斜坡函数”（Ramp Function）进行混合处理：

• 高频段（波长很短）：采用**完全不插值（外推）**策略。因为这部分旋转很快，在训练长度内已经转了无数圈，模型对它们的相对关系已经学得很好了，直接外推没问题。
• 低频段（波长很长）：采用线性插值策略。因为这部分在训练时只转了一点点，直接外推极其危险，必须压缩回已知范围。
• 中频段：作为过渡区域，根据公式进行加权混合。

熵理论与温度缩放（Temperature Scaling）

研究发现，当上下文被暴力拉长时，注意力机制（Attention）的分布熵会发生改变，导致 Softmax 后的概率分布变得异常（过尖或过平），使得模型患上“Lost in the Middle”综合征（忘记长文本中间的信息）。

YaRN 引入了一个温度系数 $t$（Temperature）来对点积的模长进行修正：

$$Attention(Q,K)=Softmax\left(\frac{Q{K}^T}{t\sqrt{d_k}}\right)$$

（或者直接将 Embedding 乘以 $1/t$）。

这里的 $t$ 是一个根据扩展倍数s计算出的经验公式（$1/t=0.1\ln (s)+1$）。通过这种温度缩放操作，它升高了 Attention Logits 的数值，使得注意力分布在超长文本下依然保持尖锐，从而恢复模型对关键信息的聚焦能力。

其他

Llama 3 基频参数 base 调高到 500,000。
Mistral 7B 使用Sliding Window Attention (SWA) 降低计算attention的复杂度。
M-RoPE 利用RoPE维度分块独立的特性，将Embedding向量切分为时间、高度、宽度三部分。
DeepSeek-V2 引用了MLA（Multi-Head Latent Attention）以极大地压缩KV Cache。然而，MLA采用了低秩压缩（Low-Rank Compression），这导致RoPE无法直接应用于压缩后的Latent Vector（因为旋转会破坏压缩后的语义空间）。解决方案是解耦RoPE（Decoupled RoPE），将Query和Key向量显式地拆分为语义向量和位置向量。

minimind实现RoPE

model_minimind.py

# 预计算 RoPE (Rotary Position Embedding) 的频率矩阵
# RoPE 通过旋转矩阵将位置信息编码到 Query 和 Key 中，使模型能够理解 token 的相对位置。
# 本函数预计算所有位置的 cos 和 sin 值，避免在每次前向传播时重复计算。
#支持 YaRN 外推方法，可以处理超过训练时最大长度的序列。
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
    # dim: 每一个注意力头(Head)的特征维度大小
    # end: 预计算的最大长度，比如支持最多 32K 个token
    # rope_base: 决定旋转频率变化的基数（通常很大，如1e6，用于控制长周期的波长）
    # rope_scaling: 扩展长文本的配置字典（如果是长文本扩展YaRN才会用到）

    # 1. 基础频率计算
    # 公式：theta_i = 1.0 / (rope_base ^ (2i / dim))
    # 生成一系列从快到慢变化的频率。高维转得慢，低维转得快。
    freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
    
    # 2. 如果开启了长文本扩展（YaRN 技术）
    if rope_scaling is not None: 
        # YaRN: 通过温度插值，解决大模型在超出训练长度时的“分布漂移”问题
        orig_max, factor, beta_fast, beta_slow, attn_factor = (
            rope_scaling.get("original_max_position_embeddings", 2048), # 原本训练的最大长度
            rope_scaling.get("factor", 16),                             # 准备拉伸多少倍
            rope_scaling.get("beta_fast", 32.0),                        # YaRN高频参数，决定了哪些“极高频”的维度应该保持完全不动（不进行插值/压缩）。在这个频率以上的维度，模型直接进行外推（Extrapolation），完全不损失分辨率。
            rope_scaling.get("beta_slow", 1.0),                         # YaRN低频参数，决定了哪些“极低频”的维度必须进行完全的线性插值（完全压缩）。在这个频率以下的维度，必须老老实实地按扩展倍数（factor）进行线性插值，确保所有角度都落在模型见过的范围内。
            rope_scaling.get("attention_factor", 1.0)                   # 注意力缩放温度（热力学修正）
        )
        # 如果要求推理的长度超过了原本的训练长度，就开始魔法操作
        if end / orig_max > 1.0:
            # 这是一个根据公式计算在不同频率区间应该怎么“插值”的函数
            inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
            # 找到低频和高频的边界
            low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
            # 计算一个渐变系数 ramp (斜坡函数)，从 0 平滑过渡到 1
            ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
            # 根据 YaRN 原理混合：高频不插值，低频线性插值，中间混合
            freqs = freqs * (1 - ramp + ramp / factor)
            
    # 3. 得到每个位置对应的旋转角度矩阵
    # t 是位置索引：[0, 1, 2, ..., end-1]
    t = torch.arange(end, device=freqs.device)
    # torch.outer(t, freqs) 计算 t_m * theta_i：得出每个位置、每个维度的实际旋转弧度
    freqs = torch.outer(t, freqs).float()
    
    # 4. 把刚刚算出来的弧度转换成 余弦(cos) 和 正弦(sin)
    # 并在最后一个维度复制拼接一次，因为复数旋转操作要同时作用在实部和虚部（代码里是用特征前后两半实现的）
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
    
    # 返回预先计算好的 cos 和 sin 矩阵
    return freqs_cos, freqs_sin

# 应用旋转位置编码（RoPE）到 Query 和 Key
# RoPE 通过复数旋转将位置信息编码到 Q 和 K 中
# 在实现中，将复数旋转分解为实部和虚部的线性组合，使用 rotate_half 函数实现。
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    # q: Query 向量 (比如形状: [批次, 句子长度, 头数, 维度])
    # k: Key 向量 
    # cos, sin: 刚刚预计算好的角度矩阵
    
    # 一个内部函数：把向量切成前后两半，并把前半部分取反移到后面
    # 比如输入是 [x1, x2, x3, x4]，切半变成 [x1, x2] 和 [x3, x4]
    # 返回 [-x3, -x4, x1, x2]。这就是二维复数乘法在实数域的等价矩阵变化。
    def rotate_half(x): 
        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
    
    # 对 Query 实施旋转 (复数乘法公式：(a+bi)(cosθ+isinθ) = (a·cosθ - b·sinθ) + i(a·sinθ + b·cosθ))
    # 这里通过刚才定义的 rotate_half 非常巧妙地实现了上面括号里的推导
    # unsqueeze_dim=1 是为了把 cos/sin 的维度跟 q 对齐 (比如扩充出批次或多头的维度)
    q_embed = ((q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))).to(q.dtype)
    
    # 对 Key 实施同样的旋转
    k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
    
    # 返回融入了绝对位置信息的 q 和 k
    # 稍后 q 和 k 做点积相乘时，数学特性会自动让他们只关注相对距离！
    return q_embed, k_embed

Normalization 归一化

要解决的核心问题是：内部协变量偏移（Internal Covariate Shift）。在训练过程中，由于前层参数不断更新，每一层网络的输入数据分布在每个 mini-batch 之间持续发生变化，导致后续各层需要不断重新适应新的输入分布。

背景：为什么 NLP 摒弃了 Batch Normalization？

• 变长序列的处理难题：文本数据的长度是不固定的。在 Batch 训练中，短序列通常需要通过 Padding 补齐。如果使用 BN，Padding Token 的存在会严重扭曲均值和方差的统计量，导致归一化参数产生偏差。虽然可以通过掩码（Masking）技术部分缓解，但实现复杂度较高且效率低下 。
• 小 Batch Size 的统计不稳定性：LLM 的训练通常受限于显存容量，导致每个 GPU 上的 Micro-Batch Size 非常小（甚至为 1）。在这种情况下，BN 估算的样本均值和方差具有极大的噪声，无法代表真实的数据分布，导致模型训练震荡甚至发散 。
• 序列维度的独立性需求：在 NLP 任务中，每个 Token 的语义往往依赖于其上下文，但不同样本之间的 Token 并没有直接的对应关系（例如，Batch 中第 i 个样本的第 5 个词与第 j 个样本的第 5 个词可能完全无关）。BN 强制在 Batch 维度上进行统计，破坏了 Token 级特征的独立性 。

Layer Normalization 应运而生。LN 沿着特征维度（Feature Dimension）对单个样本进行归一化，其统计量的计算仅依赖于当前样本自身，完全独立于 Batch Size 和序列长度。这一特性使其成为处理变长序列数据和 RNN/Transformer 架构的天然选择 。

Layer Normalization (LN)

对于一个维度为 $d$ 的输入向量 $\mathbf{x}=(x_1,x_2,\dots ,x_d)$，LayerNorm 分两步计算：标准化（Normalization）和仿射变换（Affine Transformation）。

第一步：计算均值与方差

$$\mu =\frac{1}{d}\sum _{i=1}^d{x}_i$$

$${\sigma }^2=\frac{1}{d}\sum _{i=1}^d(x_i-\mu {)}^2$$

第二步：标准化

$${\hat{x}}_i=\frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}}$$

其中 $ϵ$（如 $1\times 10^{-5}$）用于防止分母为零带来的数值不稳定。

第三步：仿射变换（恢复表达能力）

$$y_i={\gamma }_i{\hat{x}}_i+{\beta }_i$$

引入可学习的缩放参数 $\gamma$（Gain）和偏置参数 $\beta$（Bias），防止归一化操作过度压缩模型的表达能力。初始化时通常令 $\gamma =1,\beta =0$，使初始状态下 LayerNorm 近似为恒等变换。

import torch
import torch.nn as nn

def manual_layer_norm(x, gamma, beta, eps=1e-5):
    # 1. 计算均值，在特征维度（最后一维）上计算
    mu = x.mean(dim=-1, keepdim=True)
    
    # 2. 计算方差（总体方差，对应公式中的 1/d）
    var = x.var(dim=-1, unbiased=False, keepdim=True)
    
    # 3. 标准化
    x_hat = (x - mu) / torch.sqrt(var + eps)
    
    # 4. 仿射变换
    y = gamma * x_hat + beta
    return y

# --- 验证 ---
d = 5
x = torch.randn(2, d)
gamma = torch.ones(d)
beta = torch.zeros(d)

y_manual = manual_layer_norm(x, gamma, beta)
y_official = nn.LayerNorm(d)(x)

print(f"手动实现与官方API误差是否极小: {torch.allclose(y_manual, y_official, atol=1e-6)}")

LayerNorm 为模型带来了两个关键的不变性属性：

• 重中心化不变性（Re-centering Invariance）：对任意常数 $\delta$，有 $LN(\mathbf{x}+\delta )=LN(\mathbf{x})$。模型对输入数据的绝对偏移不敏感，有助于处理不同偏置的输入。

• 重缩放不变性（Re-scaling Invariance）：对任意缩放因子 $\lambda$，有 $LN(\lambda \mathbf{x})=LN(\mathbf{x})$。权重矩阵的模长不会影响输出值的幅度。在反向传播中，梯度的大小与权重的模长成反比，隐式起到了类似**学习率衰减（Learning Rate Decay）**的调节作用，防止权重无限增长。

尽管 LayerNorm 表现优异，但计算均值的减法操作（Mean Subtraction）在硬件层面引入了额外开销。后续的 RMSNorm 正是基于对”均值计算是否真正必要”的反思而提出的——通过消融实验发现，LayerNorm 的成功主要归功于重缩放不变性，而非重中心化，因此 RMSNorm 直接省去均值计算，大幅降低了计算复杂度。

Post-Norm 与 Pre-Norm 的博弈

Post-Norm：经典但脆弱

原始 Transformer（《Attention Is All You Need》）和 BERT 采用 Post-Norm，LN 放在残差连接（Residual Connection）之后：

$$x_{l+1}=LN(x_l+\text{Sublayer}(x_l))$$

核心问题：

• 梯度不稳定：反向传播时梯度需穿越多个 LN 层，浅层易梯度消失/爆炸。
• 必须 Warm-up：训练初期学习率必须极小，否则直接发散。
• 性能上限略高：每层输入都被标准化，非线性表达能力更充分——代价是训练困难。

Pre-Norm：现代大模型的标配

GPT-2/3、LLaMA、PaLM 等转向 Pre-Norm，LN 放在子层输入端：

$$x_{l+1}=x_l+\text{Sublayer}(LN(x_l))$$

核心优势：

• Highway 效应：主干路径 $x_{l+1}=x_l+\cdots$ 构成一条”恒等高速公路”，梯度可直接无损传到底层。
• 无需 Warm-up：梯度流稳定，可使用激进学习率，对 175B 级大模型训练成本极关键。

隐患——“深度诅咒”：主干 $x_l$ 的方差随层数线性累积，而残差分支经 LN 后方差归 1，深层残差对主干贡献被隐式压缩至约 $1/L$，深层网络可能退化为恒等映射。

对比总结

特性	Post-Norm (BERT)	Pre-Norm (LLaMA)
公式	$LN(x+\text{Sub}(x))$	$x+\text{Sub}(LN(x))$
训练稳定性	差，需 Warm-up	优，可无 Warm-up
梯度流	易爆炸/消失	主干直通，稳定
深层贡献	均衡	随深度衰减
最终性能	理论上限略高	略低但更易训练

结论：Pre-Norm 已成为当今 LLM 默认选择；为弥补其深层退化问题，DeepNorm 等方案应运而生（将残差主干乘以 $\alpha >1$ 后再做 Post-Norm，兼顾两者优点）。

Root Mean Square Normalization (RMSNorm)

核心理念：去中心化的归一化。

真正有用的是重缩放（Re-scaling）不变性，而非重中心化（Re-centering）不变性。

深度网络的激活值均值本来就在 0 附近游荡，强制减去均值的收益微乎其微，但这一步却引入了额外的”规约”（Reduction）计算。于是，砍掉均值，只留均方根缩放，RMSNorm 诞生了。

LayerNorm（对比参考）：

$${\bar{x}}_i=\frac{x_i-\mu }{\sigma }\cdot {\gamma }_i+{\beta }_i$$

RMSNorm（砍掉均值和偏置）：

$$\text{RMS}(\mathbf{x})=\sqrt{\frac{1}{d}\sum _{i=1}^d{x}_i^2+ϵ}$$

$${\bar{x}}_i=\frac{x_i}{\text{RMS}(\mathbf{x})}\cdot {\gamma }_i$$

变化一目了然：去掉了 $\mu$、去掉了 $\beta$，只剩下一次平方→均值→开根→缩放。

关键实现细节：

1. 无偏置项 $\beta$：LLaMA、Mistral 的 RMSNorm 均不带 $\beta$，只保留可学习的缩放因子 $\gamma$，参数量更少。
2. $ϵ$ 的位置：必须加在”平方求均值之后、开根之前”，即 $\sqrt{\cdots +ϵ}$，防止分母为零，典型取值 $1\text{e-}6$。
3. Gemma 的微调：将输出改为 $\text{RMSNorm}(x)\times (1+\gamma )$，初始化时 $\gamma \approx 0$，输出趋近恒等映射，保留了 Pre-Norm 的”直通”特性，有利于深层信号传播。

为什么实际快了 10%~40%？表面上 FLOPs 减少不多（矩阵乘法才是大头），但关键在于：LayerNorm 是内存带宽受限（Memory-Bandwidth Bound）操作

• LayerNorm 需要两次全向量扫描：第一次算 $\mu$，第二次才能算 ${\sigma }^2$（因为 ${\sigma }^2$ 依赖 $\mu$），中间还需同步屏障。
• RMSNorm 只需一次扫描算平方和，显存读写次数减半，同步屏障消失。

在推理阶段显存带宽是瓶颈的背景下，这个优化被放大；加上 RMSNorm 结构简单，更容易做算子融合（Kernel Fusion），在 FlashAttention、vLLM 等框架中收益进一步扩大。

	LayerNorm	RMSNorm
公式	$(x-\mu )/\sigma \cdot \gamma +\beta$	$x/\text{RMS}(x)\cdot \gamma$
均值计算	需要	不需要
偏置 $\beta$	有	无
显存扫描次数	2次	1次
代表模型	BERT, T5	LLaMA, Gemma, Mistral
推理速度	基准	快 10%~40%

结论：RMSNorm 以极小的精度代价，换来了显著的工程效率提升，成为现代 LLM 的标配归一化方案。

DeepNorm ：Post-Norm + Pre-Norm

Pre-Norm 解决了基本稳定性问题，但其”深度诅咒”（Curse of Depth）导致深层层对主干贡献随深度衰减，限制了模型扩展到 1000 层以上的潜力。
Microsoft Research 提出的 DeepNorm 目标是：结合 Post-Norm 的高性能 + Pre-Norm 的稳定性，让模型无论多深，每步权重更新幅度的期望值都被约束在一个常数范围内。

DeepNorm 本质是改进的 Post-Norm，在残差主干上乘以缩放系数 $\alpha$：

$$x_{l+1}=LN(\alpha \cdot x_l+G_l(x_l,{\theta }_l))$$

• $G_l$：第 $l$ 层子层函数（Attention 或 FFN）
• $\alpha >1$：主干放大系数

为什么 $\alpha >1$ 能稳定训练？放大主干 $x_l$ 相对于子层输出 $G_l$ 的幅度，使 $G_l$ 的相对贡献被压缩，从而控制了梯度方差，避免标准 Post-Norm 的梯度爆炸。

初始化规则很重要！$\alpha$ 和权重缩放系数 $\beta$ 必须随总层数动态调整：

Encoder-only（如 BERT），共 $N$ 层：

$$\alpha =(2N{)}^{1/4},\beta =(8N{)}^{-1/4}$$

Decoder-only（如 GPT），共 $M$ 层：

$$\alpha =(2M{)}^{1/4},\beta =(8M{)}^{-1/4}$$

Encoder-Decoder（如 T5），Encoder $N$ 层 + Decoder $M$ 层：

Encoder 部分：

$$\alpha =0.81(N^{4}M{)}^{1/16},\beta =0.87(N^{4}M{)}^{-1/16}$$

Decoder 部分：

$$\alpha =(3M{)}^{1/4},\beta =(12M{)}^{-1/4}$$

$\beta$ 用于缩小子层权重的初始化幅度，是一个”一次性”的缩放操作，不是运行时公式的一部分，让$G_l$的初始输出幅度变小。施加对象：

• FFN 的所有权重矩阵
• Attention 的 $W_V$（Value Projection）和 $W_O$（Output Projection）

⚠️ $W_Q$ 和 $W_K$ 不做 $\beta$ 缩放，避免影响注意力分数的计算尺度。

操作	效果
$\alpha \cdot x_l$（放大主干）	梯度有稳定的”高速公路”，不易消失
$\beta$ 缩小权重初始化	子层输出初始幅度小，不会一开始就”乱喷”梯度
最后套 LN	保持 Post-Norm 的表达能力优势

结论：DeepNorm 是当你需要堆到 1000+ 层时的选择，通过数学推导而非经验调参来保证训练稳定性，是超深 Transformer（DeepNet）的核心支撑技术。

QK-Norm 处理注意力熵坍塌

在训练 ViT-22B、PaLM-540B 这类超大模型时，研究人员发现 Attention Logits（即 $Q\cdot K^T$ 的值）会随训练进行膨胀到 $10^4$ 量级。
极大的 Logits 让 Softmax 进入饱和区，输出趋近于 One-hot 分布（某一位置概率≈1，其余≈0），这就是注意力熵坍塌。后果：

• Softmax 饱和区梯度≈0 → 梯度消失
• 模型停止学习，甚至训练发散

QK-Norm的解决方法是在计算点积之前，分别对 Query 和 Key 做归一化：

$$q^{\prime }=LN(q),k^{\prime }=LN(k)$$

$$\text{Attention}(q,k,v)=\text{Softmax}\left(\frac{q^{\prime }(k^{\prime }{)}^T}{\sqrt{d}}\right)v$$

为什么有效？ 根据柯西-施瓦茨不等式，点积上界由两个向量的模长乘积决定：

$$|q^{\prime }\cdot k^{\prime }|\le \Vert q^{\prime }\Vert \cdot \Vert k^{\prime }\Vert$$

LN 强制约束了 $q^{\prime }$ 和 $k^{\prime }$ 的模长，Logits 的数值范围被直接卡死，无法随深度或训练步数无限增长。
部分研究还对 V 向量归一化，或结合 Softmax Capping（将 Logits 截断在 $[-30,30]$ 区间），提供双重保险。Mistral AI 的最新模型也采用了类似 Softmax Capping 技术。

异构归一化方案

1. Sandwich-Norm（三明治归一化）
   背景：清华大学开发 CogView（文生图模型）时发现，即使用了 Pre-Norm，图像 Token 仍会出现数值溢出——图像生成任务对数值的敏感度远高于纯文本。
   解决方法：在 Pre-Norm 基础上，在残差分支末端（与主干相加之前）再加一个 LN，把子层”夹”在两个 Norm 之间：

   $$x_{l+1}=x_l+LN(\text{Sublayer}(LN(x_l)))$$

   形如”面包-子层-面包”，故名三明治。代价是多一次 LN 的计算，换来的是彻底消除数值溢出，让 CogView 能在 FP16 混合精度下稳定训练。
2. NormFormer：通过”过归一化”加速收敛
   策略：在每个 Transformer 层中塞入三个归一化操作：
   1. Pre-LN：标准子层输入归一化
   2. Post-Attention LN：Self-Attention 输出后立即加一个 LN
   3. Head-Scaling：对每个 Attention Head 输出做可学习标量缩放
      效果：1.3B 参数模型实验中，达到目标 Perplexity 的速度加快了 24%。适度”过归一化”让优化器在训练早期更快找到下降方向，参数量的增加通常被收敛加速抵消。

minimind实现RMSNorm

代码在model_minimind.py

"""
RMSNorm(x) = weight * (x / sqrt(mean(x^2) + eps)), 计算均方根
"""
class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        """
        dim: 输入特征的维度
        eps: 防止除零的小常数（默认 1e-5）
        """
        super().__init__()
        self.eps = eps
        # 可学习的缩放参数，初始化为全 1
        self.weight = nn.Parameter(torch.ones(dim))

    def norm(self, x):
        # x.pow(2): 计算 x 的平方
        # .mean(-1, keepdim=True): 在最后一个维度上求均值，保持维度
        # torch.rsqrt: 计算 1/sqrt，比先 sqrt 再除更快
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """
        前向传播
        Args:
            x: 输入张量，可以是任意精度（float16, bfloat16, float32）
        Returns:
            归一化并缩放后的张量，保持原始精度
        """
        # 先转换为 float32 进行归一化计算（提高数值稳定性）
        # 乘以可学习的权重，最后转换回原始精度（type_as(x)）
        return (self.weight * self.norm(x.float())).type_as(x)

大模型优化方法：从KV Cache到Flash Attention

Roofline模型是高性能计算领域用来定义硬件性能理论上限的模型。Roofline模型用公式表示为：

$$P=\min (\pi ,I\times \beta )$$

其中：

• P = 计算性能（能达到的实际性能）
• π = 峰值算力（GPU能做的最大浮点运算数）
• β = 峰值内存带宽（GPU能传输的最大数据量）
• I = 算术强度（Arithmetic Intensity），即每字节内存访问所做的浮点运算次数（FLOPs/Byte）

这个模型揭示了一个重要现象：硬件性能取决于两个因素的最小值——要么受限于计算能力（Compute Bound），要么受限于内存带宽（Memory Bound）。

在LLM推理中的应用：对于一个70B参数的模型（FP16精度约140GB），在H100 GPU（3.35TB/s带宽）上：

$$\text{推理速度}=\frac{3350\text{ GB/s}}{140\text{ GB}}\approx 24\text{ Tokens/s}$$

这说明在解码阶段，模型虽然只进行一次向量-矩阵乘法，但需要加载全部几百GB的权重，导致算术强度极低，硬件性能完全受限于显存带宽——这被称为”内存墙“，是制约LLM推理速度的根本物理障碍。这就是为什么KV Cache、GQA和Flash Attention等技术如此重要——它们都是在Roofline模型的约束下，通过减少内存访问或提高算术强度来优化推理性能的方案。

背景：标准 Attention 的时间与空间复杂度均为 $O(L^2)$，随上下文长度增长显存爆炸。

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

随着大语言模型（LLM）参数规模从数十亿（7B）向数千亿（70B-400B+）乃至万亿级别跃升，以及上下文窗口从标准的4k token扩展至128k甚至1M token，深度学习的计算重心已显著从模型训练阶段的算力（FLOPS）瓶颈，转移至推理阶段的显存带宽（Memory Bandwidth）与显存容量（Memory Capacity）瓶颈。在自回归（Autoregressive）生成范式下，传统的Transformer架构面临着计算效率与存储效率的双重挑战。

KV Cache是实现实时自回归生成的必要前提，但其引入的显存线性增长问题迫使架构设计从MHA向MQA和GQA转型，以在模型精度与推理吞吐量（Throughput）之间寻找最优平衡。与此同时，Flash Attention通过算法层面的IO感知（IO-Aware）优化，打破了传统注意力机制的二次方复杂度存储瓶颈，成为长文本推理的基石。

KV Cache

原理：缓存历史 Token 的 K、V，避免重复计算。

• Prefill 阶段：计算 Prompt 中所有 Token 的 K、V 并存入缓存
• Decode 阶段：只计算当前 Token 的 $q_n,k_n,v_n$，拼接到缓存后做注意力

为什么不缓存 Query (Q)？

在注意力机制中，Query代表“当前关注点”，它随着生成的推进而不断变化。例如，在生成第10个词时，我们关注的是第10个词与前9个词的关系；生成第11个词时，我们关注的是第11个词与前10个词的关系。因此，Query向量在每一步都是全新的，必须重新计算。而历史Token作为被关注的对象（Key）和信息载体（Value），其特征是固定的，因此可以被缓存 。

注意力架构演进：MHA → MQA → GQA

特性	MHA	GQA	MQA
KV 头数	$H$（与 Q 相等）	$G$（典型值 $H/4$）	1
KV Cache	100%	25%	~3%
模型精度	最高	极接近 MHA	有损失
Q:K:V 比例	$H:H:H$	$H:G:G$	$H:1:1$

GQA 是主流选择（$G=8$ 为甜点），当 $G=H$ 退化为 MHA，$G=1$ 退化为 MQA。

Flash Attention

问题：标准实现反复将 $N\times N$ 矩阵写入/读出 HBM，IO 是瓶颈而非计算。

方案：IO 感知算法，通过**分块（Tiling）+ 重计算（Recomputation）**将 Attention 融合为单一 CUDA Kernel，在 SRAM 内完成计算，减少 HBM 读写。

重计算用于反向传播

注意力层的输出 $O$ 定义为：

$$O=PV=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d}}\right)V\in {\mathbb{R}}^{N\times d}$$

其中 $S=Q{K}^{\top }$ 是相似度矩阵，$P=\text{softmax}(S)$ 是注意力权重矩阵，输出的每一行是对所有 Value 向量的加权求和：

$$O_i=\sum _j{P}_{ij}{V}_j$$

标准实现需要在前向传播时保存完整的 $S$ 和 $P$（均为 $N\times N$）用于反向传播，显存复杂度为 $O(N^2)$，序列长时代价极大。

Flash Attention 的实现是前向传播时不保存 $S$ 和 $P$，只保留三样东西：

• 输出 $O\in {\mathbb{R}}^{N\times d}$
• 每行的最大值 $m_i={\max }_j{S}_{ij}$
• 每行的 softmax 分母 ${\ell }_i={\sum }_j{e}^{S_{ij}-m_i}$

$m_i$ 和 ${\ell }_i$ 合称 softmax 归一化统计量，在分块计算时可以增量维护。

需要梯度时，利用仍在显存中的 $Q,K$ 逐块重新计算 $S$ 和 $P$：

$$S_{ij}=q_i\cdot k_j^{\top },P_{ij}=\frac{e^{S_{ij}-m_i}}{{\ell }_i}$$

再用标准链式法则计算梯度。由 $d{P}_{ij}=d{O}_i\cdot V_j^{\top }$，得：

$$d{S}_{ij}=P_{ij}\left(d{P}_{ij}-\sum _k{P}_{ik}d{P}_{ik}\right)$$

进而：

$$dQ=dSK,dK=d{S}^{\top }Q,dV=P^{\top }dO$$

注：$O$ 本身不用于还原 $P$，而是参与 $dO$ 相关梯度项的计算；真正用于重建 $P$ 的是 $Q,K,m_i,{\ell }_i$。

指标	标准 Attention	Flash Attention
前向显存	$O(N^2)$	$O(N)$
计算量	$O(N^{2}d)$	$O(N^{2}d)$（反向多一次前向）

重计算以少量额外计算换取大量显存节省，是 Flash Attention 能支持长序列训练的核心机制。

• Flash Decoding：针对生成阶段（Decode），将 KV Cache 切分为多块（Split-K），并行分发给所有 SM 计算，长上下文（32k+）下提速 8倍以上。

显存占用公式

$$M_{\text{total}}=M_{\text{weights}}+M_{\text{kvcache}}+M_{\text{activation}}$$

$$M_{\text{weights}}\approx P_{\text{model}}\times S_{\text{prec}}$$

$$M_{\text{kvcache}}=2\times N_{\text{layers}}\times N_{kv}\times D_{\text{head}}\times L_{\text{seq}}\times B\times S_{\text{prec}}$$

系数 2 代表同时存储 Key 和 Value。

FP16 下的简化版：

• 静态显存（权重）：$M_{\text{weight}}=\text{参数量}\times 2$
• 动态显存（KV Cache）：$M_{kv}=2\times \text{层数}\times \text{KV头数}\times \text{头维度}\times \text{序列长}\times \text{Batch}\times 2$

示例：Llama 2 7B（MHA，32 KV头）vs Llama 3 8B（GQA，8 KV头），后者 KV Cache 仅为前者的 1/4。

前沿扩展

• MLA（DeepSeek-V2/V3）：通过低秩矩阵分解将 KV 压缩为极小的 Latent Vector，比 GQA 更进一步压缩每头内部的数据表示。
• PagedAttention（vLLM）：类似 OS 虚拟内存分页，允许 KV Block 离散存储，显存利用率从 60% 提升至 95% 以上。

技术	解决什么	作用层面
MQA / GQA	KV 头数太多，存储太大	模型结构
MLA	KV 特征维度冗余，可低秩压缩	模型结构
Flash Attention	Attention 计算时的 HBM IO 瓶颈	算子/计算
PagedAttention	KV Cache 内存碎片，并发利用率低	系统/调度

minimind实现GQA、Flash Attention

实现在model_minimind.py
minimind使用了GQA、Flash Attention。

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

# ============================================================
# 工具函数：repeat_kv
# 作用：实现 GQA（Grouped Query Attention）的 KV 头扩展
#       当 KV 头数 < Q 头数时，把每个 KV 头"复制"多份，
#       使 KV 头数对齐 Q 头数，才能做 Attention 矩阵乘法。
# ============================================================
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    Args:
        x: Key 或 Value 张量 [batch, seq_len, num_kv_heads, head_dim]
        n_rep: 每个 KV head 需要重复的次数（n_rep = num_heads / num_kv_heads）
        
    Returns:
        重复后的张量 [batch, seq_len, num_heads, head_dim]
    """
    bs, slen, num_key_value_heads, head_dim = x.shape
    # n_rep == 1 说明是标准 MHA，KV 头数 == Q 头数，不需要复制，直接返回
    if n_rep == 1: return x
    return (
        x[:, :, :, None, :] # 在第 3 维（头数后面）插入一个新维度,shape: (bs, slen, num_kv_heads, 1, head_dim)
        .expand(bs, slen, num_key_value_heads, n_rep, head_dim) # 沿新维度扩展 n_rep 份（不复制内存，只是广播视图）, shape: (bs, slen, num_kv_heads, n_rep, head_dim)
        .reshape(bs, slen, num_key_value_heads * n_rep, head_dim)) # 把 num_kv_heads × n_rep 合并成最终的头数（== Q 头数）, shape: (bs, slen, num_q_heads, head_dim)

class Attention(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()

        # 头数配置 :如果没有单独指定 KV 头数，就等于 Q 头数（退化为标准 MHA）
        self.num_key_value_heads = (
            config.num_attention_heads 
            if config.num_key_value_heads is None 
            else config.num_key_value_heads)
        self.n_local_heads    = config.num_attention_heads   # Q 的头数
        self.n_local_kv_heads = self.num_key_value_heads     # KV 的头数（GQA 中 < Q 头数）
        self.n_rep            = self.n_local_heads // self.n_local_kv_heads  # 每个 KV 头需要被复制的份数
        self.head_dim         = config.head_dim              # 每个头的维度

        self.is_causal = True # 因果掩码标志 语言模型训练时必须是因果的（只能看到当前及之前的 token）

        # ── 线性投影层 ────────────────────────────────────────
        # Q 投影：hidden_size -> (Q头数 × head_dim)
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        # K 投影：hidden_size -> (KV头数 × head_dim)，GQA 时输出比 Q 小
        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        # V 投影：同 K
        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        # 输出投影：把所有头拼接后的结果映射回 hidden_size
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)

        s# ── QK Norm（对每个头单独做 RMSNorm）────────────────────
        # 稳定训练，防止 QK 点积数值爆炸（在大模型中逐渐成为标配）
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

        self.attn_dropout  = nn.Dropout(config.dropout)   # 作用在 attention 权重上
        self.resid_dropout = nn.Dropout(config.dropout)   # 作用在输出残差上
        self.dropout       = config.dropout

        # ── 是否启用 Flash Attention ──────────────────────────
        # 只有 PyTorch >= 2.0 且 config 允许时才用
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn

    # 执行入口：投影拆分多头-计算QK-softmax+dropout-加权求V-拼接多头结果-输出层投影
    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
        """
        Args:
            x: 输入张量 [batch, seq_len, hidden_size]
            position_embeddings: RoPE 位置编码 (cos, sin) 元组
            past_key_value: 缓存的 KV 值，用于增量解码 [batch, past_len, num_kv_heads, head_dim]
            use_cache: 是否返回 KV 缓存供下次使用
            attention_mask: 注意力掩码 [batch, seq_len]，1 表示有效位置，0 表示掩码位置
            
        Returns:
            output: 注意力输出 [batch, seq_len, hidden_size]
            past_kv: 新的 KV 缓存（如果 use_cache=True），否则为 None
        """
        bsz, seq_len, _ = x.shape
        # ── Step 1：QKV 线性投影 ──────────────────────────────
        xq = self.q_proj(x)   # (bsz, seq_len, n_q_heads * head_dim)
        xk = self.k_proj(x)   # (bsz, seq_len, n_kv_heads * head_dim)
        xv = self.v_proj(x)   # (bsz, seq_len, n_kv_heads * head_dim)

        # ── Step 2：Reshape 成多头格式 ────────────────────────
        xq = xq.view(bsz, seq_len, self.n_local_heads,    self.head_dim)  # (bsz, seq_len, n_q_heads, head_dim)
        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)  # (bsz, seq_len, n_kv_heads, head_dim)
        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)  # (bsz, seq_len, n_kv_heads, head_dim)

        # ── Step 3：QK Norm（逐头归一化）─────────────────────
        xq = self.q_norm(xq)
        xk = self.k_norm(xk)

        # ── Step 4：施加 RoPE 旋转位置编码 ───────────────────
        cos, sin = position_embeddings
        xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
        # 注意：V 不做 RoPE，只有 Q 和 K 需要位置信息

        # ── Step 5：拼接 KV Cache（推理阶段）────────────────
        # past_key_value[0] shape: (bsz, past_len, n_kv_heads, head_dim)
        if past_key_value is not None:
            xk = torch.cat([past_key_value[0], xk], dim=1)  # 把历史 K 拼接到当前 K 前面
            xv = torch.cat([past_key_value[1], xv], dim=1)  # 同理 V
        # 此后 xk/xv 的 seq 维度 = past_len + 当前 seq_len（完整上下文）

        # ── Step 6：保存当前 KV 供下次使用 ───────────────────
        past_kv = (xk, xv) if use_cache else None

        
        # ── Step 7：repeat_kv 扩展 KV 头数 ──────────────────
        # 将 KV 从 n_kv_heads 复制到 n_q_heads，使维度对齐
        # 同时 transpose(1,2) 将 heads 维度提前：(bsz, n_heads, seq_len, head_dim)
        xq = xq.transpose(1, 2)                               # (bsz, n_q_heads,  seq_len,  head_dim)
        xk = repeat_kv(xk, self.n_rep).transpose(1, 2)       # (bsz, n_q_heads,  full_len, head_dim)
        xv = repeat_kv(xv, self.n_rep).transpose(1, 2)       # (bsz, n_q_heads,  full_len, head_dim)


        # ── Step 8：计算 Attention ────────────────────────────
        # 判断是否走 Flash Attention 路径：
        # 条件：① PyTorch 支持；② seq_len > 1（单步解码不值得用）；
        #       ③ 不是 KV Cache 续写的推理步；④ 没有 padding mask
        use_flash = (
            self.flash
            and (seq_len > 1)
            and (not self.is_causal or past_key_value is None)
            and (attention_mask is None or torch.all(attention_mask == 1))
        )

        if use_flash:
            output = F.scaled_dot_product_attention(
                xq, xk, xv,
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=self.is_causal
            )
        else:
            # 手动实现 Scaled Dot-Product Attention
            # scores shape: (bsz, n_heads, seq_len, full_len)
            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
            # ↑ Q × K^T 再除以 sqrt(head_dim) 防止点积数值过大

            if self.is_causal:
                # 构造上三角为 -inf 的因果掩码，只看当前及之前位置
                # triu(1) 把主对角线及以下置 0，只保留右上角为 -inf
                causal_mask = torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
                scores[:, :, :, -seq_len:] += causal_mask
                # 注意：只掩 scores 的最后 seq_len 列（当前步的新 Q 对完整 KV 序列）

            if attention_mask is not None:
                # padding mask：把无效位置（mask=0）的分数设为极小值
                # unsqueeze 扩展维度以广播：(bsz,1,1,full_len)
                scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9

            # Softmax 归一化得到 Attention 权重，然后加权求和 V
            # .float() 防止 fp16 下 softmax 数值溢出，type_as 再转回原 dtype
            output = self.attn_dropout(
                F.softmax(scores.float(), dim=-1).type_as(xq)
            ) @ xv
            # output shape: (bsz, n_heads, seq_len, head_dim)

        # ── Step 9：合并多头 & 输出投影 ──────────────────────
        # transpose 回 (bsz, seq_len, n_heads, head_dim)，
        # reshape 合并最后两维 -> (bsz, seq_len, n_heads * head_dim)
        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)

        # 线性投影回 hidden_size，并加残差 dropout
        output = self.resid_dropout(self.o_proj(output))
        # output shape: (bsz, seq_len, hidden_size)

        return output, past_kv
        # 返回：① 当前层的输出 hidden states；② 更新后的 KV Cache

MoE 核心知识

传统 Dense Transformer 的 FFN 层是全局共享的巨大矩阵；MoE 将其拆解为 $N$ 个独立子网络（”专家”）$E_1,E_2,\dots ,E_N$。

MoE 层输出公式：

$$y=\sum _{i=1}^{N}G(x{)}_i\cdot E_i(x)$$

其中 $G(x)$ 是门控网络（Router）输出，稀疏激活时只有 $K$ 个元素非零（$K\ll N$）。

两个关键度量：

• 总参数量：决定模型知识库大小
• 激活参数量：处理单个 Token 实际参与计算的参数量，决定推理 FLOPs

模型	总参数量	激活参数量	稀疏比
DeepSeek-V3	671B	37B	~5.5%
DeepSeek-V2	236B	21B	~8.9%
Mixtral 8x7B	47B	13B	~27%

基础组件：SwiGLU

Swish 激活函数：

$$\text{Swish}(x)=x\cdot \sigma (x)=\frac{x}{1+e^{-x}}$$

SwiGLU FFN 公式：

$${\text{FFN}}_{\text{SwiGLU}}(x)=\left(\underset{\text{门控信号}}{\underbrace{\text{Swish}(x{W}_g)}}\odot \underset{\text{内容信息}}{\underbrace{(x{W}_u)}}\right)W_d$$

• $W_g$：门控投影矩阵（Gate）
• $W_u$：升维投影矩阵（Up）
• $W_d$：降维投影矩阵（Down）
• $\odot$：逐元素乘法

工程权衡：SwiGLU 比标准 FFN 多用一个矩阵，参数量多 50%。为保持参数预算不变，需将隐藏层宽度乘以 2/3。

在minimind model_minimind.py中，SwiGLU的写法如下：

注：以前版本FFN里的dropout 被移到了 Attention 内部（attn_dropout 和 resid_dropout），FFN 本身不再做 dropout。

import torch
import torch.nn.functional as F
from torch import nn
from transformers.activations import ACT2FN

# ============================================================
#  FeedForward —— 单个 SwiGLU 前馈网络（稠密 FFN / 单个 MoE 专家）
# ============================================================
class FeedForward(nn.Module):
    def __init__(self, config, intermediate_size: int = None):
        super().__init__()

        # intermediate_size 优先用外部传入的值，否则用 config 里的默认值
        # MoE 场景下每个专家的中间维度通常比稠密 FFN 小（补偿多专家带来的参数量增加）
        intermediate_size = intermediate_size or config.intermediate_size

        # ── 三个线性投影，bias=False 是主流 LLM 的标配，减少参数并稳定训练 ──

        # 门控信号 gate_proj: hidden_size → intermediate_size
        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)

        # 内容信号up_proj:   hidden_size → intermediate_size
        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)

        # down_proj: intermediate_size → hidden_size
        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)

        # 激活函数，默认 'silu'（即 Swish），与 GLU 结合就是 SwiGLU
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        # x: (batch_size, seq_len, hidden_size)
        # 输出 shape: (batch_size, seq_len, hidden_size)
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

DeepSeek MoE 架构创新

细粒度专家（Fine-Grained Experts）

传统 MoE 专家数少（如 Mixtral 只有 8 个），每个专家被迫承担异构知识。DeepSeek 采用更多、更小的专家：

• DeepSeek-V2：160 个路由专家，每次激活 6 个，中间维度 1536
• DeepSeek-V3：256 个路由专家，每次激活 8 个

共享专家（Shared Experts）

部分专家固定为”共享专家”，始终激活，不参与路由竞争，负责捕获公共知识。

共享专家 MoE 输出公式：

$$y=\sum _{i\in {\mathcal{A}}_{\text{shared}}}{E}_i(x)+\sum _{j\in \text{TopK}(G(x))}{g}_j\cdot E_j(x)$$

• ${\mathcal{A}}_{\text{shared}}$：共享专家集合（负责通识）
• 路由专家：专注长尾/领域知识（知识解耦）

Minimind的MoE设计，以前的 moe_infer 函数已经完全移除了，新版代码直接在 MOEFeedForward.forward() 里一体处理训练和推理，不再是单独的函数。现在的 forward() 已经用 self.training 判断来区分训练/推理。
model_minimind.py

import torch
import torch.nn.functional as F
from torch import nn
from transformers.activations import ACT2FN

# ============================================================
#  MOEFeedForward —— 混合专家前馈网络
#  结构：N 个 FeedForward 专家 + 1 个路由门控
#  每个 Token 只激活 Top-K 个专家，输出加权求和
# ============================================================
class MOEFeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        # 路由门控：把每个 Token 的隐向量映射为对所有专家的"偏好分数"
        # 输入 hidden_size，输出 num_experts 个分数
        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)

        # 创建 num_experts 个独立的 FeedForward 专家
        # 每个专家的中间维度是 moe_intermediate_size（通常小于稠密 FFN 的 intermediate_size）
        # 这样 N 个专家的总参数量不会太大
        self.experts = nn.ModuleList([
            FeedForward(config, intermediate_size=config.moe_intermediate_size)
            for _ in range(config.num_experts)
        ])

        # 激活函数（实际上 FeedForward 内部已有，这里备用）
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        # x: (batch_size, seq_len, hidden_size)
        batch_size, seq_len, hidden_dim = x.shape

        # ── Step 1：把 3D 张量展平为 2D，方便对每个 Token 独立做路由 ──
        # x_flat: (batch_size * seq_len, hidden_size)，每行是一个 Token
        x_flat = x.view(-1, hidden_dim)

        # ── Step 2：路由打分 ──
        # gate(x_flat): (num_tokens, num_experts) 原始 logits
        # softmax 后得到每个 Token 对每位专家的"归一化概率"
        scores = F.softmax(self.gate(x_flat), dim=-1)
        # scores: (num_tokens, num_experts)

        # ── Step 3：选出 Top-K 专家 ──
        # topk_weight: (num_tokens, K) 被选中专家的概率权重
        # topk_idx:    (num_tokens, K) 被选中专家的索引编号
        # sorted=False 不排序，节省一点时间
        topk_weight, topk_idx = torch.topk(
            scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False
        )

        # ── Step 4：（可选）对 Top-K 权重重新归一化 ──
        # 因为 softmax 是对所有专家算的，Top-K 截取后概率之和 < 1
        # 重新归一化让 K 个专家的权重之和 = 1，输出幅度更稳定
        if self.config.norm_topk_prob:
            topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)

        # ── Step 5：逐专家计算，把结果加权累加回输出 ──
        y = torch.zeros_like(x_flat)  # 初始化输出为全 0

        for i, expert in enumerate(self.experts):
            # mask: (num_tokens, K)，标记哪些 (token, slot) 选择了专家 i
            mask = (topk_idx == i)

            if mask.any():
                # 找出所有选择了专家 i 的 Token 的行索引
                # mask.any(dim=-1): (num_tokens,) bool，该 token 是否选了专家 i
                # .nonzero().flatten(): 取出这些 token 的行号
                token_idx = mask.any(dim=-1).nonzero().flatten()

                # 取出这些 token 对应的权重，shape: (num_selected, 1)，乘 1 列方便后面广播
                weight = topk_weight[mask].view(-1, 1)

                # 专家 i 只处理被选中的那部分 token，避免全量计算
                # expert(x_flat[token_idx]): (num_selected, hidden_size)
                # * weight: 按权重缩放
                # index_add_: 把结果累加到 y 对应行（一个 token 可能被多个 slot 选，但每个专家只被选一次）
                y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))

            elif self.training:
                # 训练时若某个专家完全没被任何 token 选中，
                # 用 0 * 参数之和 挂一个假梯度，防止该专家参数因无梯度而被优化器跳过
                # （不影响数值结果，仅保持梯度流动）
                y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())

        # ── Step 6：计算辅助负载均衡损失（仅训练阶段）──
        # MoE 的常见问题：路由坍塌（所有 token 都选同一两个专家）
        # 辅助损失惩罚负载不均衡，鼓励所有专家被均匀使用
        if self.training and self.config.router_aux_loss_coef > 0:
            # load: (num_experts,) 各专家被选中的平均次数（one-hot 统计）
            load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)

            # scores.mean(0): (num_experts,) 各专家的平均路由概率
            # 两者点积：负载高 × 概率高 → 惩罚集中路由
            # * num_experts 做归一化使损失量级与专家数无关
            self.aux_loss = (
                (load * scores.mean(0)).sum()
                * self.config.num_experts
                * self.config.router_aux_loss_coef
            )
        else:
            # 推理阶段不需要辅助损失，置 0 即可
            self.aux_loss = scores.new_zeros(1).squeeze()

        # ── Step 7：还原形状 ──
        # (num_tokens, hidden_size) → (batch_size, seq_len, hidden_size)
        return y.view(batch_size, seq_len, hidden_dim)

为什么MoE的训练和推理代码实现不同？训练时需要做推理不需要的事：

1. 反向传播保活：未被选中的专家参数没有梯度流过，会”休眠”导致训练不充分，所以要加 +0 * expert.parameters() 这个假梯度让所有专家都参与梯度计算。而推理不需要算梯度，只需要以前向传播最快的方式得到结果。
2. 负载均衡损失：需要计算 aux_loss 惩罚”某几个专家被过度选中”的现象，推理时不需要优化任何参数，算了也没用。
3. 策略不同：训练时的策略是布尔掩码，Batch Size 通常很大；而推理时的策略是排序与装箱，Batch Size 可能非常小

推理只需要”选 top-k 专家、算结果”，其他都是训练专属的优化手段。

路由机制

Top-K Gating（经典方法）

对输入 $x$ 和路由权重矩阵 $W_r\in {\mathbb{R}}^{d_{\text{model}}\times N}$：

Step 1 — 计算亲和度分数：

$$h(x)=x\cdot W_r$$

Step 2 — Top-K 截断（保持稀疏性）：

$$\text{KeepTopK}(h(x),K{)}_i=\{\begin{array}{ll}h(x{)}_i& \text{if }h(x{)}_i\in \text{Top-}K(h(x))\\ -\infty & \text{otherwise}\end{array}$$

Step 3 — Softmax 归一化：

$$G(x)=\text{Softmax}(\text{KeepTopK}(h(x),K))$$

Step 4 — 加权求和：

$$y=\sum _{i\in \text{Top-}K}G(x{)}_i\cdot E_i(x)$$

负载均衡辅助损失（传统方法）

定义 $f_i$：路由给专家 $i$ 的 Token 比例；$P_i$：路由器预测给专家 $i$ 的平均概率。

$${\mathcal{L}}_{\text{aux}}=\alpha \cdot N\sum _{i=1}^N{f}_i\cdot P_i$$

其中，$\alpha$是超参，训练前手动设置，$\alpha$越大，负载均衡越强，但可能牺牲主任务性能。$N$ 是专家数量。

或方差形式，让每个专家被选中的频率尽量接近均匀分布，即每个专家的理想负载是 $\frac{1}{N}$（$N$ 为专家数）：

$$L_{aux}=\sum _{j=1}^N{\left(\underset{\text{理想负载}}{\underbrace{\frac{1}{N}}}-\underset{\text{实际负载}{f}_j}{\underbrace{\frac{1}{T}\sum _{i=1}^T{g}_{ij}}}\right)}^2$$

其中，

• $g_{ij}$：第 $i$ 个 Token 分配给专家 $j$ 的 gate 值（被选中为 1/权重，未选中为 0）
• $\frac{1}{T}{\sum }_{i=1}^T{g}_{ij}$：专家 $j$ 在这批 $T$ 个 Token 上的实际平均负载 $f_j$
• $\frac{1}{N}$：均匀分布时每个专家应得的理想负载

DeepSeek 创新：无辅助损失负载均衡

在路由 Logits 上添加动态偏置项 $b_i$（不参与梯度下降）：

$${\text{Score}}_i=x\cdot W_{r,i}+b_i$$

动态更新规则（类 PID 控制）：

• 若 ${\text{Load}}_i>\text{Target}$（过载）：$b_i\leftarrow b_i-\gamma$
• 若 ${\text{Load}}_i<\text{Target}$（空闲）：$b_i\leftarrow b_i+\gamma$

优势：$W_r$ 仅由主任务 Loss 优化，$b_i$ 仅负责流量调控，完全解耦，模型可自由探索最优路由策略。

minimind实现aux_loss

现在的minimind代码里去除了MoeGate。现在的逻辑是 RMSNorm → Attention → FeedForward → MOEFeedForward → MiniMindBlock → MiniMindModel → MiniMindForCausalLM

路由功能就是 MOEFeedForward 里的一行 self.gate = nn.Linear(…) 加上 forward() 里的 softmax + topk，没有单独封装成类。

现在aux_loss 在新版中的完整流程是

step 1: MOEFeedForward.forward() 里计算并存储：

if self.training and self.config.router_aux_loss_coef > 0:
    # 公式来源：Switch Transformer，L_aux = α · N · Σ(f_i · P_i)
    # load[i]：专家 i 在当前 batch 中被"硬选中"的频率（0/1 离散）
    # F.one_hot(topk_idx) 把每个 token 选中的专家编号转为 one-hot 向量
    # .mean(0) 对所有 token 取平均 → 每个专家被选中的比例
    # 形状：[num_experts]，理想均匀时每个值 = top_k / num_experts
    load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)

    # scores.mean(0)[i]：专家 i 的平均路由概率（softmax 输出，连续可微）
    # load 是离散不可微的，scores.mean(0) 是可微的软版本
    # .sum()：对所有专家求和 → 得到标量
    # * num_experts：归一化补偿，使损失量级不随专家数变化
    #   （均匀时 Σ f_i·P_i = top_k/N，乘以 N 后变为常数 top_k）
    # * router_aux_loss_coef：超参数 α，控制负载均衡损失的强度
    self.aux_loss = (load * scores.mean(0)).sum() * num_experts * router_aux_loss_coef
else:
    self.aux_loss = scores.new_zeros(1).squeeze()  # 推理时为 0

step 2: MiniMindModel.forward() 里汇总所有层的 aux_loss：

aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], ...)
return hidden_states, presents, aux_loss

step 3: MiniMindForCausalLM.forward() 里返回给训练循环：

return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, ...)

注意这里 aux_loss 只是返回出去，没有在模型内部和 loss 相加。实际加权相加 (loss + α * aux_loss) 是在训练脚本里完成的。

专家容量与 Token 丢弃

$$C=\frac{\text{Tokens per Batch}}{N}\times \text{Capacity Factor}$$

Capacity Factor 通常取 1.0~1.2。超出容量的 Token 被丢弃，直接通过残差连接传递。DeepSeek-V3 因负载均衡极好，实际丢弃率极低。

训练与工程优化

并行策略：

• 数据并行（DP）：复制模型、划分数据，不适用于 MoE（总参数太大）
• 张量并行（TP）：切分矩阵计算，适用于 Attention 层
• 专家并行（EP）：不同 GPU 持有不同专家，引入 All-to-All 双向通信开销（Dispatch + Combine 两阶段）

FP8 混合精度训练（DeepSeek-V3 首创）：

相比 BF16/FP16：显存减少 50%，带宽减少 50%，计算速度提升 2x（理论值）。

通过细粒度量化（Block-wise Scaling） 解决 FP8 动态范围窄的问题，训练成本仅为同级别模型的 1/10（278 万 H800 机时）。

核心结论

1. 稀疏性是后摩尔定律时代突破算力墙的关键路径（37B 激活参数撬动 671B 知识）
2. 共享+路由专家实现通识与专才解耦，是 DeepSeek 参数效率极高的根本原因
3. 现代 MoE 竞争已进入系统工程阶段（FP8、All-to-All 通信优化与算法协同）

主流大模型架构参数对比

模型	类型	总参数	激活参数	稀疏比	层数	专家数	TopK	注意力	精度	负载均衡	备注
DeepSeek-V3	MoE	671B	37B	~5.5%	61	256R+1S	Top-8	MLA	FP8	Bias调整（无Aux）	原文
DeepSeek-V3-0324	MoE	671B	37B	~5.5%	61	256R+1S	Top-8	MLA	FP8	Bias调整（无Aux）	2025.03更新
Llama-3.1-405B	Dense	405B	405B	100%	126	N/A	N/A	GQA	BF16	N/A	原文
Llama 4 Scout	MoE	109B	17B	~15.6%	48	16R	Top-1	GQA+iRoPE	BF16	Aux Loss	2025.04
Llama 4 Maverick	MoE	400B	52B	~13%	48	128R	Top-1	GQA+iRoPE	BF16	Aux Loss	2025.04
Mixtral 8x22B	MoE	141B	39B	~27.7%	56	8R	Top-2	GQA	BF16	Aux Loss	原文
Qwen3-235B-A22B	MoE	235B	22B	~9.4%	94	128R	Top-8	MHA+QKNorm	BF16	Aux Loss	2025.04 旗舰MoE
Qwen3-32B	Dense	32B	32B	100%	64	N/A	N/A	GQA+QKNorm	BF16	N/A	2025.04 旗舰Dense

R = 路由专家（Routed）；S = 共享专家（Shared）
Llama 4 Scout/Maverick 为 Meta 2025.04 发布，交错 MoE 架构（非每层都有专家）
Dense 就是普通的全连接前馈网络，对应传统 Transformer 架构，每次前向传播所有参数都参与计算

minimind 训练

大模型训练的三个阶段：

阶段	名称	目的	数据
第一阶段	预训练 Pre-training	学会语言本身	海量无标注文本（网页、书籍）
第二阶段	监督微调 SFT	学会按指令回答	人工标注的问答对
第三阶段	强化学习 RLHF	让回答更符合人类偏好	人类打分/偏好数据

pretrain

代码在lm_dataset.py，train_pretrain.py，trainer_utils.py

dataset链接是https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files,

预训练的文本格式：

~/minimind$ head -n 2 /home/root/minimind/dataset/pretrain_hq.jsonl | jq

{
  "text": "鉴别一组中文文章的风格和特点，例如官方、口语、文言等。需要提供样例文章才能准确鉴别不同的风格和特点。<|im_end|> 好的，现在帮我查一下今天的天气怎么样?今天的天气依据地区而异。请问你需要我帮你查询哪个地区的天气呢？<|im_end|> 打开闹钟功能，定一个明天早上七点的闹钟。好的，我已经帮您打开闹钟功能，闹钟将在明天早上七点准时响起。<|im_end|> 为以下场景写一句话描述：一个孤独的老人坐在公园长椅上看着远处。一位孤独的老人坐在公园长椅上凝视远方。<|im_end|> 非常感谢你的回答。请告诉我，这些数据是关于什么主题的？这些数据是关于不同年龄段的男女人口比例分布的。<|im_end|> 帮我想一个有趣的标题。这个挺有趣的：\"如何成为一名成功的魔术师\" 调皮的标题往往会吸引读者的注意力。<|im_end|> 回答一个问题，地球的半径是多少？地球的平均半径约为6371公里，这是地球自赤道到两极的距离的平均值。<|im_end|> 识别文本中的语气，并将其分类为喜悦、悲伤、惊异等。\n文本：“今天是我的生日！”这个文本的语气是喜悦。<|im_end|>"
}
{
  "text": "根据输入的内容，编写一个类别标签。\n这是一篇介绍如何阅读心电图的文章类别标签: 医学/心电图阅读指南<|im_end|> 帮我搜索一下最近的天气情况。当然，我可以帮您搜索最新的天气情况。请问您需要查询哪个城市的天气情况呢？<|im_end|> 帮我讲一个令人开心的笑话。好的，我帮您讲一个关于细菌的笑话。为什么细菌不会上网？因为连接总是断开了！<|im_end|> 现在给我生成一首关于大海的五言诗。碧波万顷月满天，海天相接处天地间。波涛滚滚江山美，海鸟翱翔日月闲。<|im_end|> 谢谢你，这篇文章很有用。不客气，我很高兴能够为您提供帮助。如果您还有其他问题或需求，随时可以对我说。<|im_end|> 你好，我想下载一个视频编辑软件，你有什么推荐吗？您好！当然，有很多选择。您想要免费软件还是愿意付费？<|im_end|> 为什么我的程序不输出正确结果？可能是代码逻辑有误，或者输入数据有误，需要仔细调试代码逻辑和输入数据。<|im_end|> 谢谢你的回答。现在我想知道这场比赛的具体时间和地点。这场比赛的时间是北京时间10月4日，地点是上海。<|im_end|>"
}

有监督的微调（SFT）数据集，用于指令微调（Instruction Tuning）。
SFT文本格式：

{
    "conversations": [
        {"role": "user", "content": "你好"},
        {"role": "assistant", "content": "你好！"},
        {"role": "user", "content": "再见"},
        {"role": "assistant", "content": "再见！"}
    ]
}
{
    "conversations": [
        {"role": "system", "content": "# Tools ...", "tools": "[...]"},
        {"role": "user", "content": "把'你好世界'翻译成english"},
        {"role": "assistant", "content": "", "tool_calls": "[{\"name\":\"translate_text\",\"arguments\":{\"text\":\"你好世界\",\"target_language\":\"english\"}}]"},
        {"role": "tool", "content": "{\"translated_text\":\"Hello World\"}"},
        {"role": "assistant", "content": "Hello World"}
    ]
}

train_pretrain.py 负责组装模型、数据和优化器，并执行训练循环。核心训练循环 train_epoch 包括

1. 动态学习率：每个 Step 都调用 get_lr （余弦学习率衰减 调度函数）更新学习率。
2. 前向传播 (Forward)：
   • 使用 autocast 开启混合精度（bfloat16 / float16），节省显存。
   • Loss 计算：loss = res.loss + res.aux_loss
   • • 显式加上 MoE 的辅助 Loss（负载均衡 Loss），防止专家坍缩。
3. 反向传播 (Backward)：
   • 使用 scaler.scale(loss) 放大 Loss，防止 FP16 下梯度下溢。
   • 梯度累积：loss / accumulation_steps
   • 允许在小显存显卡上模拟大 Batch Size 训练。
4. 优化器步进：
   • 每 accumulation_steps 步执行一次 optimizer.step()。
   • 执行梯度裁剪 (clip_grad_norm_)，防止梯度爆炸。

前向传播 (Forward Pass):把输入数据从头到尾跑一遍模型，得到输出和 Loss。
执行流程是

输入 Token
↓
Embedding 层          # Token ID → 向量
↓
× N 层 Transformer Block
├── RMSNorm
├── Attention（Q/K/V 计算、softmax）
├── RMSNorm
└── FFN / MoE FFN   # 含 aux_loss 计算
↓
输出层 Linear (lm_head, Language Model Head)
↓
CrossEntropyLoss → loss 值

反向传播 (Backward Pass)：拿到 Loss 之后，用链式法则从后往前计算每个参数的梯度，告诉模型”每个参数该往哪个方向调整”。

loss.backward()
   ↓
lm_head ← 计算梯度
   ↓
× N 层（反向顺序）
   ├── FFN / MoE 梯度
   ├── Attention 梯度
   └── Embedding 梯度
   ↓
optimizer.step()      # 用梯度更新参数

def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
    """
    训练一个 Epoch 的核心函数。
    
    参数：
        epoch       : 当前是第几个 Epoch（从 0 开始）
        loader      : DataLoader，每次迭代返回 (input_ids, labels)
        iters       : 本 Epoch 的总 Step 数
        start_step  : 断点续训时的起始 Step（默认 0）
        wandb       : wandb/swanlab 日志对象（可选）
    """
    start_time = time.time()
    last_step = start_step

    # ── 遍历每个 Batch ──────────────────────────────────────────────────────────
    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):

        # 将数据移动到目标设备（GPU/CPU）
        input_ids = input_ids.to(args.device)
        labels = labels.to(args.device)
        last_step = step

        # ── 1. 动态学习率 ───────────────────────────────────────────────────────
        # 根据当前全局 Step 计算余弦衰减学习率
        # 全局 Step = epoch * 每Epoch步数 + 当前步
        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr  # 实时更新优化器的学习率

        # ── 2. 前向传播（Forward Pass）───────────────────────────────────────────
        # autocast_ctx：开启混合精度（bfloat16/float16），节省显存、加速计算
        # 如果是 CPU 训练，autocast_ctx = nullcontext()（不启用混合精度）
        with autocast_ctx:
            res = model(input_ids, labels=labels)
            # res.loss     : 主 Loss（CrossEntropyLoss，预测下一个 token 的准确度）
            # res.aux_loss : MoE 辅助 Loss（负载均衡 Loss，防止专家坍缩）
            loss = res.loss + res.aux_loss

            # 梯度累积：将 loss 除以累积步数
            # 等效于把 accumulation_steps 个小 Batch 合并成一个大 Batch
            # 作用：在小显存显卡上模拟大 Batch Size 训练
            loss = loss / args.accumulation_steps

        # ── 3. 反向传播（Backward Pass）─────────────────────────────────────────
        # scaler.scale(loss)：在 FP16 模式下放大 loss，防止梯度下溢（变成 0）
        # .backward()：计算所有参数的梯度，梯度会累积在 param.grad 上
        scaler.scale(loss).backward()

        # ── 4. 优化器步进（每 accumulation_steps 步执行一次）───────────────────
        if step % args.accumulation_steps == 0:
            # unscale_：还原被放大的梯度，回到真实数值
            scaler.unscale_(optimizer)

            # 梯度裁剪：将所有参数梯度的全局范数限制在 grad_clip（默认 1.0）以内
            # 防止梯度爆炸导致训练崩溃
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

            # optimizer.step()：用当前梯度更新模型参数（内部由 scaler 代理执行）
            # 如果本步梯度异常（inf/nan），scaler 会跳过本次更新
            scaler.step(optimizer)

            # scaler.update()：根据本次是否跳过，动态调整下一步的放大系数
            scaler.update()

            # 清空梯度，set_to_none=True 比置零更省显存
            optimizer.zero_grad(set_to_none=True)

        # ── 5. 日志记录 ─────────────────────────────────────────────────────────
        if step % args.log_interval == 0 or step == iters:
            spend_time = time.time() - start_time

            # 还原被除以 accumulation_steps 的 loss，用于展示真实 loss 值
            current_loss = loss.item() * args.accumulation_steps
            current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
            current_logits_loss = current_loss - current_aux_loss  # 主 loss（去掉辅助loss）
            current_lr = optimizer.param_groups[-1]['lr']

            # 预估剩余时间（分钟）
            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60

            Logger(
                f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
                f'loss: {current_loss:.4f}, '
                f'logits_loss: {current_logits_loss:.4f}, '
                f'aux_loss: {current_aux_loss:.4f}, '
                f'lr: {current_lr:.8f}, '
                f'epoch_time: {eta_min:.1f}min'
            )
            # 同步到 wandb/swanlab 可视化平台
            if wandb:
                wandb.log({
                    "loss": current_loss,
                    "logits_loss": current_logits_loss,
                    "aux_loss": current_aux_loss,
                    "learning_rate": current_lr,
                    "epoch_time": eta_min
                })

        # ── 6. 定期保存模型权重 ──────────────────────────────────────────────────
        if (step % args.save_interval == 0 or step == iters) and is_main_process():
            model.eval()  # 切换到推理模式（关闭 Dropout 等）

            moe_suffix = '_moe' if lm_config.use_moe else ''
            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'

            # 如果是 DDP 分布式训练，需要取出原始模型（去掉 DDP 包装）
            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
            # 如果使用了 torch.compile，需要取出原始模型
            raw_model = getattr(raw_model, '_orig_mod', raw_model)

            state_dict = raw_model.state_dict()
            # 保存时转为 half()（FP16）并移到 CPU，节省磁盘空间
            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)

            # 同时保存完整 checkpoint（含优化器状态，用于断点续训）
            lm_checkpoint(lm_config, weight=args.save_weight, model=model,
                          optimizer=optimizer, scaler=scaler,
                          epoch=epoch, step=step, wandb=wandb,
                          save_dir='../checkpoints')

            model.train()  # 切回训练模式
            del state_dict

        # 释放本 Step 的中间变量，及时回收显存
        del input_ids, labels, res, loss

    # ── 7. 处理末尾未执行的梯度 ──────────────────────────────────────────────────
    # 如果总 Step 数不能被 accumulation_steps 整除，
    # 最后几步的梯度还没有被 optimizer.step()，这里补充执行一次
    if last_step > start_step and last_step % args.accumulation_steps != 0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

SFT

SFT 的核心逻辑是如何做 Masking，如何构建 Prompt

sft的 train_epoch 和 pretrain 的 train_epoch 函数几乎完全一样，核心逻辑没有差别。真正的区别在于数据集如何构造 labels：

• PretrainDataset：labels = 全部 token（模型学预测每一个词）
• SFTDataset：labels 中问题部分被 mask 掉（设为 -100），模型只学回答部分。loss 计算时 -100 的位置会被 CrossEntropyLoss 自动忽略，不参与梯度计算。

举例：SFT对应到数据，强迫模型只针对 Assistant 的回复产生梯度更新。这让模型从一个原始的语言概率模型，变为一个能和人沟通的助手。

• User 说的话 (“写一个关于秋天的诗吧。”) -> Label设为 -100 (不计算梯度)。
• Assistant 说的话 (“红黄绿交织…”) -> Label设为 Token ID (计算梯度)。
• User 说的话 (“给我推荐一些…”) -> Label设为 -100。
• Assistant 说的话 (“摇滚歌曲推荐…”) -> Label设为 Token ID。

RLHF

minimind RL的数据格式如下：

{
  "chosen": [
    {"content": "Q", "role": "user"}, 
    {"content": "good answer", "role": "assistant"}
  ], 
  "rejected": [
    {"content": "Q", "role": "user"}, 
    {"content": "bad answer", "role": "assistant"}
  ]
}

信任区域算法

TRPO（信任区域策略优化）Trust Region Policy Optimization

核心思想：不限制参数数值的变化，而是限制策略行为分布的变化（用 KL 散度衡量）。

带约束的优化问题：

$$\underset{\theta }{\max }{\mathbb{E}}_t\left[\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{old}}(a_t|s_t)}{A}_t\right]\text{subject to}{\mathbb{E}}_t[\text{KL}({\pi }_{{\theta }_{old}}(\cdot |s_t)\Vert {\pi }_{\theta }(\cdot |s_t))]\le \delta$$

瓶颈：求解需对 KL 散度做二阶泰勒展开，引入 Hessian 矩阵（7B 模型大小为 70亿×70亿），计算量不可接受。

PPO（近端策略优化）—— 工业级标准

核心思想：用简单的”裁剪”技巧模拟信任区域，只用一阶梯度。

裁剪目标函数：

$${\mathcal{L}}^{CLIP}(\theta )=\mathbb{E}\left[\min \left(r_t(\theta ){\hat{A}}_t,\text{clip}(r_t(\theta ),1-ϵ,1+ϵ){\hat{A}}_t\right)\right]$$

• $r_t(\theta )=\frac{{\pi }_{\theta }}{{\pi }_{old}}$ 为新旧策略比值；$ϵ$ 通常取 0.2
• 好动作：增加概率但不超过 $1+ϵ$；坏动作：降低概率但不低于 $1-ϵ$

工业级架构需同时维护 4 个模型：

模型	角色	是否可训
Actor ${\pi }_{\theta }$	正在训练的 LLM	✅
Critic $V_{\varphi }$	预测状态价值，降低方差	✅
Reference ${\pi }_{ref}$	SFT 后的冻结副本，计算 KL 防止跑偏	❌
Reward Model	给生成结果打分	❌

PPO 瓶颈：4 个大模型占用极高显存；On-Policy 需实时生成文本，90% 时间花在推理上。

去 Critic 的 REINFORCE 变体

ReMax

利用 LLM 环境的确定性，用贪婪解码结果作为基线：

$$A(x,y_{sample})=R(x,y_{sample})-R(x,y_{greedy})$$

效果与 PPO 持平，但显存减少 1/4~1/2。

RLOO

对同一 Prompt 采 $K$ 个回复，用其余 $K-1$ 个的平均奖励作为基线：

$$b_i=\frac{1}{K-1}\sum _{j\ne i}R(x,y_j),A_i=R(x,y_i)-b_i$$

无偏估计，随 $K$ 增大方差显著降低。

Reinforce++

将 PPO 的工程技巧（Token 级 KL 惩罚、梯度裁剪、优势归一化）移植回 REINFORCE，无需 Critic 即可达到 PPO 稳定性。

RL-Free：隐式偏好优化

核心动机：跳过显式 Reward Model，直接用人类偏好数据优化策略。

DPO（直接偏好优化）

推导路径：RLHF 目标 → 最优策略解析解 → 消去配分函数 $Z(x)$

RLHF 目标：

$$\underset{\pi }{\max }{\mathbb{E}}_{x\sim \mathcal{D},y\sim \pi }\left[r(x,y)-\beta D_{KL}(\pi (y|x)\Vert {\pi }_{ref}(y|x))\right]$$

最优策略解析解：

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

最终损失（$Z(x)$ 在差值中自动抵消）：

$${\mathcal{L}}_{DPO}({\pi }_{\theta };{\pi }_{ref})=-{\mathbb{E}}_{(x,y_w,y_l)\sim \mathcal{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]$$

局限：离线算法，存在分布偏移、模式坍塌、无法探索新解等问题。

IPO

修正 DPO 的过拟合倾向，用 MSE 控制更新幅度，将偏好差距回归到固定目标值 $\gamma$：

$${\mathcal{L}}_{IPO}={\mathbb{E}}_{(x,y_w,y_l)}\left[{\left(\log \frac{\pi (y_w|x)}{{\pi }_{ref}(y_w|x)}-\log \frac{\pi (y_l|x)}{{\pi }_{ref}(y_l|x)}-\frac{\gamma }{2}\right)}^2\right]$$

KTO

灵感来源于前景理论（损失厌恶），不需要配对数据，只需单条点赞/踩：

$${\mathcal{L}}_{KTO}\approx \text{WeightedLoss}(y_{good})+\lambda \cdot \text{WeightedLoss}(y_{bad})$$

对”避免坏结果”的敏感度高于”生成好结果”。

ORPO

将 SFT 和对齐合二为一，在交叉熵 SFT Loss 上增加基于**赔率比（Odds Ratio）**的惩罚项。无需参考模型 ${\pi }_{ref}$，显存最友好。

GRPO 及推理增强算法

GRPO（群组相对策略优化）—— DeepSeek 核心算法

核心洞见：一次生成一组回答，通过组内比较估算优势，彻底删除 Critic 模型。

算法流程：

1. 群体采样：对 Prompt $q$ 并行采样 $G$ 个回答：$o_1,o_2,\dots ,o_G\sim {\pi }_{\theta }(o|q)$
2. 奖励计算：规则打分（准确性 + 格式），无需神经网络 RM
3. 优势估计（Z-Score 标准化）：

   $$A_i=\frac{r_i-\text{mean}(r_1,\dots ,r_G)}{\text{std}(r_1,\dots ,r_G)+ϵ}$$

4. 目标函数：

$${\mathcal{L}}_{GRPO}=\mathbb{E}\left[\min \left(\frac{{\pi }_{\theta }}{{\pi }_{old}}{A}_i,\text{clip}\left(\frac{{\pi }_{\theta }}{{\pi }_{old}},1-ϵ,1+ϵ\right)A_i\right)\right]-\beta D_{KL}({\pi }_{\theta }\Vert {\pi }_{ref})$$

GRPO 变体

Dr. GRPO：发现并修正 GRPO 的三类偏差：

• 基线偏差：使用留一均值替代含自身的均值
• 长度偏差：归一化不同长度序列的梯度贡献
• 难度偏差：引入历史感知锚点，处于”学习区”的问题给予更强信号

DAPO：解耦裁剪参数、动态采样，非对称 Clip 鼓励探索，适合大规模推理训练。

GSPO：序列级一致性约束，防止长链条崩塌，适合 MoE / 长文本。

SAPO：软门控提供连续信任区域，平滑优化，适合多模态 / 复杂推理。

过程监督与自举学习

PRM（过程奖励模型）

对推理的每一步打分（vs. ORM 只看最终结果），提供密集奖励信号，极大加速复杂推理收敛。

STaR（自我教学推理）

自进化循环：生成 → 过滤正确样本 → 微调 → 对错误题目用正确答案反推推理过程（Rationalization）→ 扩充训练集。

算法横向对比

算法	需要 Critic	显存	核心优势
PPO	✅	极高	稳定，在线探索，理论成熟
ReMax	❌	中	极简，效果持平 PPO
RLOO	❌	中	无偏估计，高效利用采样
DPO	❌	低	稳定，无采样开销
KTO	❌	低	只需单条反馈，无需配对数据
ORPO	❌	极低	SFT+对齐一步到位
GRPO	❌	中低	大规模推理，DeepSeek 首选
DAPO	❌	中低	动态采样，非对称 Clip
GSPO	❌	中低	序列级一致性

未来可能是多算法融合——STaR 生成数据 → ORPO 冷启动 → GRPO/DAPO 大规模推理增强 → PRM 推理时搜索引导。RL 是弥合预训练模型与 AGI 之间”最后一公里”的关键桥梁。