【Datawhale组队学习202601】Base-NLP task03 深入大模型架构

系列文章目录

• task01 注意力机制与Transformer
• task02 预训练语言模型
• task03 深入大模型架构

---

文章目录

• 系列文章目录
• 前言
• [一、手搓一个 Llama2](#一、手搓一个 Llama2)
• • [1.1 Llama2 架构总揽](#1.1 Llama2 架构总揽)
  • [1.2 代码实践](#1.2 代码实践)
  • [1.3 归一化模块 norm](#1.3 归一化模块 norm)
  • • [1.3.1 预归一化](#1.3.1 预归一化)
    • [1.3.2 RMSNorm](#1.3.2 RMSNorm)
    • [1.3.3 参考代码](#1.3.3 参考代码)
  • [1.4 旋转位置编码模块 rope](#1.4 旋转位置编码模块 rope)
  • • [1.4.1 RoPE介绍](#1.4.1 RoPE介绍)
    • [1.4.2 通俗理解](#1.4.2 通俗理解)
    • [1.4.3 参考代码](#1.4.3 参考代码)
    • • [1.4.3.1 预计算和复数乘法](#1.4.3.1 预计算和复数乘法)
      • [1.4.3.2 广播机制和GQA优化](#1.4.3.2 广播机制和GQA优化)
  • [1.5 分组查询注意力 gqa](#1.5 分组查询注意力 gqa)
• [二、MoE 架构解析](#二、MoE 架构解析)
• • [2.1 MoE 的源来 - 自适应局部专家混合](#2.1 MoE 的源来 - 自适应局部专家混合)
  • • [2.1.1 干扰效应------多任务冲突](#2.1.1 干扰效应——多任务冲突)
    • [2.1.2 怎么解决------分治](#2.1.2 怎么解决——分治)
    • [2.1.3 损失函数设计](#2.1.3 损失函数设计)
    • [2.1.4 MoE vs 集成学习](#2.1.4 MoE vs 集成学习)
  • [2.2 深度神经网络中的 MoE](#2.2 深度神经网络中的 MoE)
  • [2.3 稀疏门控 MoE](#2.3 稀疏门控 MoE)
  • • [2.3.1 从浅层到深层的变革](#2.3.1 从浅层到深层的变革)
  • [**层级化的门控（Hierarchical Gating）**：输入 x x x 首先经过第一层的门控 g 1 g^1 g1，被路由到第一层的专家 f i 1 f^1_i fi1。第一层的输出 z 1 z^1 z1 接着作为第二层门控 g 2 g^2 g2 的输入，再次被路由到第二层的专家 f j 2 f^2_j fj2。](#层级化的门控（Hierarchical Gating）：输入 x x x 首先经过第一层的门控 g 1 g^1 g1，被路由到第一层的专家 f i 1 f^1_i fi1。第一层的输出 z 1 z^1 z1 接着作为第二层门控 g 2 g^2 g2 的输入，再次被路由到第二层的专家 f j 2 f^2_j fj2。)
  • • [2.3.2 学习分解的特征表示](#2.3.2 学习分解的特征表示)
    • [2.3.3 通俗理解](#2.3.3 通俗理解)
  • [2.4 大模型时代的 MoE](#2.4 大模型时代的 MoE)
• 三、手撕大模型生成策略
• 总结

---

前言

• Datawhale组队学习教程地址

---

一、手搓一个 Llama2

1.1 Llama2 架构总揽

Llama2 遵循了 GPT 系列开创的 Decoder-Only 架构。这意味着它完全由 Transformer 解码器 层堆叠而成，天然适用于自回归的文本生成任务。

1.2 代码实践

• 目录结构：

  llama2-scratch/
  │
  ├── src/
  │ ├── init.py #初始化
  │ └── attention.y # 注意力模块
  │ ├── ffn.py # 前馈神经网络模块
  │ │── norm.py # 归一化模块
  │ │── rope.py # 旋转位置编码模块
  │ └── transformer.py # llama2transfomer
  │
  └── main.py # 主函数

• 下面开始逐个模块的学习。

1.3 归一化模块 norm

1.3.1 预归一化

Llama2 中采用了预归一化 Pre-LN 的策略，即相对于原始 Transformer 归一化位置在残差后的 后归一化 策略，Llama2 中的归一化位置在残差前。 这被认为是提升大模型训练稳定性的关键。
Llama 2 预归一化
输入
RMSNorm
Attention
Residual
输出
RMSNorm
FFN
Residual
经典Transformer 后归一化
输入
Attention
Residual
LayerNorm
输出
FFN
Residual
LayerNorm

1.3.2 RMSNorm

• 标准的 Layer Normalization 在 Transformer 中用于稳定训练，但它的计算（减去均值、除以标准差）相对复杂。

• 为了在保证性能的同时提升计算效率，Llama2 采用了它的变体 RMSNorm（Root Mean Square Layer Normalization）

• 公式如下，其中 x x x 是输入向量， γ \gamma γ 是可学习的缩放参数：

y = x 1 n ∑ i = 1 n x i 2 + ϵ ⋅ γ y = \frac{x}{\sqrt{\frac{1}{n}\sum_{i=1}^{n}x_i^2 + \epsilon}} \cdot \gamma y=n1∑i=1nxi2+ϵ x⋅γ

• 通俗对比 LayerNorm 和 RMSNorm：

• 想象你有一群同学参加考试，他们的分数差距特别大：

  • 小明：100分
  • 小红：50分
  • 小刚：10分

• 问题来了 ：如果直接用这些分数训练AI模型，100分会"压垮"10分，就像大象踩在蚂蚁身上------模型学不会关注小分数！

• 传统方法 ：LayerNorm（较复杂），就像老师既要调整平均分，又要调整分数差距：

  • 先算平均分：(100+50+10)/3 = 53分
  • 每人减去平均分 → 小明:47, 小红:-3, 小刚:-43
  • 再除以标准差（衡量分数分散程度）→ 最终得到标准化分数

• 好处：数字变得很整齐

• 坏处：计算太麻烦！大模型有几百层，每层都要算，超级卡！

• Llama2 的聪明解法 ：RMSNorm : "既然调整平均分那么麻烦，我们干脆只调整分数差距！"

• 怎么做？超简单三步：

  • 把所有分数平方：100²=10000, 50²=2500, 10²=100
  • 算平方的平均值：(10000+2500+100)/3 = 4200
  • 开平方根：√4200 ≈ 65
  • 每人分数除以65 → 小明:1.54, 小红:0.77, 小刚:0.15

• 神奇效果：

  • 数字大小变得差不多（1.54, 0.77, 0.15）
  • 省掉了"减平均分"的步骤 → 计算速度快30%！
  • 实验证明：对AI学习效果几乎没有影响

1.3.3 参考代码

• 首先，由之前学习可知，原始的文本数据首先会被分词器（Tokenizer）转换成一个由整数ID组成的序列。为了进行批处理，我们会将多个这样的序列打包在一起，形成一个形状为 [batch_size, seq_len] 的二维张量。随后，这个张量会经过一个词嵌入层（Embedding Layer） ，将每个整数ID映射成一个高维向量。这个向量的维度就是 dim。这样，我们就得到了一个 [batch_size, seq_len, dim] 形状的三维张量，这就是 Transformer Block 的标准输入。

• 接口定义

  • 输入: 一个形状为 [batch_size, seq_len, dim] 的张量 x x x。
  • 输出: 一个与输入形状相同的张量，其中每个词元 (dim 维度) 都被独立归一化。

# code/C6/llama2/src/norm.py
class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim)) # 对应公式中的 gamma

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        # 核心计算：x * (x^2的均值 + eps)的平方根的倒数
        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self._norm(x.float()).type_as(x)
        return out * self.weight

• 解读：

1. 初始化：准备"调节器"（初始设为1），就像音响的音量键，可以整体调大调小声音

   • 为什么需要：有些特征需要放大，有些要缩小（AI自己学怎么调）

2. 核心计算 _norm 函数（重点！），拆解成4个超简单步骤：

   • 步骤1: x.pow(2) → 把所有分数平方
     • 为什么平方？：放大差距！让大数字变得更大，小数字变得超小 → 这样我们能看清"整体规模"
   • 步骤2: .mean(...) → 算平方的平均值，keepdim=True：保持维度（结果还是[3]个数字，不是单个数字）
   • 步骤3: + self.eps → 加个保险丝。为什么？防止万一平均值=0（除以0会爆炸！），eps = 1e-6 就是0.000001，超级小的保险丝
   • 步骤4: torch.rsqrt(...) → 开平方根再取倒数，结果：得到一个缩放因子
   • 步骤5: x * ... → 用原分数乘缩放因子

3. 最后一步：out * self.weight。用"调节器"微调每个分数。

1.4 旋转位置编码模块 rope

1.4.1 RoPE介绍

• 在 Transformer 章节中已经知道，++模型需要位置信息来理解词元的顺序 ++。++++
• Llama2 则采用了更先进的 旋转位置编码（Rotary Positional Embedding, RoPE） ，它是一种相对位置编码。
• 与 传统位置编码 通过加法直接注入词嵌入的方式不同，RoPE 的策略是：位置信息不再是"加"到词嵌入上，而是在计算注意力时，通过复数"乘法"的方式"旋转" Query 和 Key 向量。

1.4.2 通俗理解

• 通俗解释就是：RoPE 给每个词装一个指南针，按位置逆时针旋转，AI通过指针夹角的正负号，天然知道词语的先后顺序。

• 假设班级排队拍照：第1位: 小明 第2位: 小红 第3位: 小刚 第4位: 小美。

• 问题：AI怎么知道"小明"在"小红"前面，而不是后面？

• 传统方法：给小明贴标签："位置1"，给小红贴标签："位置2"...但这样万一有个人叫"位置3"呢，而且这样还会增大向量。

• RoPE 革命：让词向量"旋转"起来！ Llama2 说：不要贴标签，要旋转！

• 第一步：设定旋转规则（像拧螺丝）

  • 固定方向：所有词按逆时针旋转（像拧开瓶盖的方向）

  • 旋转量：位置越靠后，旋转角度越大

    第1位: 0°（不转）→ 小明
    第2位: 30° → 小红
    第3位: 60° → 小刚
    第4位: 90° → 小美

• 第二步：每个同学变成一个带指针的小人

  • 指针方向 = 他们的旋转角度

    小明 (0°): 指针 →
    小红 (30°): 指针 ↗
    小刚 (60°): 指针 ↗↗
    小美 (90°): 指针 ↑

• 第三步：看指针夹角，知道前后关系！（核心！）

  • 当小红想知道和小明的关系：

  • 小红指针（30°）看小明指针（0°）

  • 角度差 = 30° - 0° = +30（正数！）

  • 正数 = 小明在小红前面*

  • 当小红看小刚：

  • 小红指针（30°）看小刚指针（60°）

  • 角度差 = 30° - 60° = -30°（负数！）

  • 负数 = 小刚在小红后面

• 正负号就是时间箭头

• +30° = 前面的人**，**-30° = 后面的人

1.4.3 参考代码

• RoPE代码设计的两个模块

1.4.3.1 预计算和复数乘法

• precompute_freqs_cis:
• 功能: 预计算一个包含旋转角度信息的复数张量 freqs_cis。这个张量在模型初始化时计算一次即可。
• 输入:
  • dim: head 的维度。
  • end: 序列最大长度。
  • theta: 一个用于控制频率范围的超参数。
• 输出: 形状为 [end, dim / 2] 的复数张量。

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor:
		# 1. 生成"旋转速度表"
		# 两两分组，dim//2
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 2. 生成0到end-1的位置编号
    t = torch.arange(end, device=freqs.device)
    # 3. 计算每个位置x每个速度的旋转角度
    freqs = torch.outer(t, freqs).float()
    # 4. 转换成"旋转指令"（复数形式）
    # 复数乘法比较快
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

1.4.3.2 广播机制和GQA优化

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    ndim = x.ndim
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)


def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    # 允许 GQA：Q/K 头数可不同，但最后一维 head_dim 应一致
    head_dim = xq.shape[-1]

    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    # 分别对 Q/K 广播以兼容不同头数
    freqs_q = reshape_for_broadcast(freqs_cis, xq_)
    freqs_k = reshape_for_broadcast(freqs_cis, xk_)
    xq_out = torch.view_as_real(xq_ * freqs_q).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_k).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xq)


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    bsz, seqlen, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bsz, seqlen, n_kv_heads, n_rep, head_dim)
        .reshape(bsz, seqlen, n_kv_heads * n_rep, head_dim)
    )

• apply_rotary_emb:
• 功能: 将预计算的 freqs_cis 应用于输入的 Query 和 Key 向量。
  输入:
  • xq: Query 向量，形状 [batch_size, seq_len, n_heads, head_dim]。
  • xk: Key 向量，形状 [batch_size, seq_len, n_kv_heads, head_dim]。
  • freqs_cis: 预计算的旋转矩阵切片。
• 输出: 旋转后的 xq 和 xk，形状不变。

1.5 分组查询注意力 gqa

---

二、MoE 架构解析

• 稠密模型 Dense Model ， 像 Llama 2、GPT-3这样的模型。对于每一个输入的 Token，模型中 所有的 参数(从第一层到最后一层)都会参与计算。
• 但是，随着模型参数规模实在是越来愈巨大，全量参数计算带来了高昂的算力成本。
• 这就引出了本节的主角------混合专家模型 Mixture of Experts MoE ，它通过一种 稀疏激活的机制，兼具了大规模参数的知识容量与较低的推理成本。

2.1 MoE 的源来 - 自适应局部专家混合

最早的 MoE 思想可以追溯到 1991 年 Michael Jordan 和 Geoffrey Hinton 发表的经典论文 《Adaptive Mixture of Local Experts》 自适应局部专家混合 。

2.1.1 干扰效应------多任务冲突

• 在传统的单体神经网络中，如果我们尝试让一个网络同时学习多个截然不同的子任务（例如既学做菜又学修车），往往会出现 "强干扰效应（Strong Interference Effects）" 。

• 这是因为 ++网络的所有权重都参与了所有任务的计算++。

• 现象就是：当网络调整参数以适应任务 A 时，可能会破坏它在任务 B 上已经学到的特征表示。从而导致学习速度变慢，泛化能力变差。

• 通俗来讲就是，举个栗子：假设一个AI要识别"猫"和"汽车"。猫毛茸茸的，汽车硬邦邦的，但万一遇到一只趴在车顶的胖橘猫？单一模型就会懵圈："这到底是生物还是机械？"它试图用同一套"脑回路"处理所有数据，结果学得四不像------猫识别率暴跌，汽车也认成玩具车。

• 根本难题：这叫 "多任务冲突"。就像你让语文老师教物理，他可能会把牛顿定律写成一首诗。神经网络也一样，不同任务需要不同"思维方式"，硬塞进一个大脑，只会互相拖后腿。

2.1.2 怎么解决------分治

• 文章中提出了基于分治 Divide-and-Conquer 策略的系统架构(计算机的同学肯定很熟悉 )：
  • 专家网络 Expert，
  • 门控网络 Gating ，

2.1.3 损失函数设计

2.1.4 MoE vs 集成学习

• 虽然结构看似相似，但两者有本质区别。
• 集成学习（如随机森林） 通常假设基模型是独立或互补的，预测时所有模型都参与，通过投票或加权平均得出结果。
• 而 MoE 强调动态的条件计算，它根据输入数据本身（Data-driven）动态地划分任务空间，不同的输入激活不同的子网络路径。

2.2 深度神经网络中的 MoE

2013 年，Ilya Sutskever 等人发表了论文 《Learning Factored Representations in a Deep Mixture of Experts》深度混合专家模型中的因子表示学习 ，将 MoE 与深度学习进行了开创性的结合。

2.3 稀疏门控 MoE

Google Brain 团队（包括 Geoffrey Hinton 和 Jeff Dean 等）于 2017 年发表了论文 《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》惊人的大型神经网络：稀疏门控的专家混合层 ，正式将 MoE 带入了超大规模模型（百亿参数级）的时代。

2.3.1 从浅层到深层的变革

在此之前，MoE 通常作为一种独立的浅层模型存在。

• Ilya 等人的工作打破了这一局限，他们提出 Deep Mixture of Experts（DMoE），将 MoE 结构"模块化"并嵌入到深度神经网络的多个层级中。

• 意味着 MoE 不再是一个孤立的架构，而成为了一种可插拔的层 。我们可以在一个深层网络的不同位置（例如第 1 1 1 层和第 2 2 2 层）分别插入 MoE 模块，每一层都有自己独立的门控网络和专家集合。

• 通俗的讲，IIya 就是将 之前一个万能的浅层的MoE 搞成了个深层的专家流水线。

层级化的门控（Hierarchical Gating） ：输入 x x x 首先经过第一层的门控 g 1 g^1 g1，被路由到第一层的专家 f i 1 f^1_i fi1。第一层的输出 z 1 z^1 z1 接着作为第二层门控 g 2 g^2 g2 的输入，再次被路由到第二层的专家 f j 2 f^2_j fj2。

• 指数级增长的组合路径 ：通过这种堆叠，网络能够表达的有效"专家组合"数量呈指数级增长。如果第一层有 N N N 个专家，第二层有 M M M 个专家，那么网络潜在的组合路径就高达 N × M N \times M N×M 种。每个输入样本都会根据其特性，动态地选择一条最适合的处理路径。

2.3.2 学习分解的特征表示

• 论文的标题强调了"Factored Representations"。通过在不同层级引入混合专家，模型能够自发地在不同层级学习到数据的不同维度的特征。

• 论文在"Jittered MNIST"（带随机平移的手写数字）数据集上观察到了有趣的现象：

• 第一层专家倾向于根据数字的**位置（Location）**进行分工，成为了"Where Experts"。

• 第二层专家倾向于根据数字的**类别（Class）**进行分工，成为了"What Experts"。

• 这种自动的特征解耦证明了深度 MoE 能够有效地利用其深层结构，将复杂任务分解为多个正交的子问题进行处理，为后来 MoE 在 Transformer 中的广泛应用奠定了重要的理论基础。

2.3.3 通俗理解

• 想象你开了一家披萨店（当做是深度神经网络）。以前，店里只有一个"万能厨师"（传统浅层MoE），不管客人要海鲜披萨还是水果披萨，都他一个人做 ，您猜怎么着，累不死他。

• 但2013年，Ilya团队说："NO！咱们搞个'专家流水线'！"

• 升级后：店里分成两层车间（Layer 1 和 Layer 2）。

  • 第一层车间：全是"饼底专家"（比如老王专门揉面，老李专攻薄脆底）。
  • 第二层车间：全是"创意 topping 专家"（比如小美只放水果，小明只撒辣条）。

• 每个车间门口还有个"智能分拣机器人"（门控网络），它看一眼订单（输入数据），就喊："这个客人要薄底，老王上！下一个要水果味，小美准备！"

  • 这就是Deep MoE的核心：MoE不再是个孤胆英雄，而是变成可插拔的"智能模块"，你想插在第几层车间都行！（论文里管这叫"模块化"，简单说就是乐高积木）

• 神奇之处1：分拣机器人玩"接力赛"（层级化门控）,AI竟然还真的做成了

  • 输入一个数据（比如一张订单），它先冲进第一层车间：
  • 第一层分拣机器人（ g 1 g^1 g1） 扫一眼："哎哟，这个要薄脆！交给'专家'老王 f i 1 f^1_i fi1！" → 老王 f i 1 f^1_i fi1处理完输出中间结果（ z 1 z^1 z1）。
  • z 1 z^1 z1马上冲进第二层车间，第二层分拣机器人（ g 2 g^2 g2） 再扫："哈！这明明是要水果，交给'专家'小美 f 2 1 f^1_2 f21！" → 小美输出最终答案。

• 神奇之处2：组合技能爆炸！像乐高宇宙

• 假设第一层有3个专家（老王、老李、老张），第二层有2个专家（小美、小明）。

  • 传统AI：只有3+2=5种技能。
  • Deep MoE：技能直接 3×2=6种组合路径！比如：
    • 老王 + 小美 → 薄底水果披萨
    • 老李 + 小明 → 脆底辣条披萨
    • 甚至老张 + 小美 → 全麦底草莓披萨（？）

• 指数级增长的意思是：层数越多，组合越疯！10个专家×10个专家=100条路，100×100=10,000条路......像你玩《我的世界》，用基础方块能搭出城堡、飞船甚至你家小区！AI再也不用"死脑筋"，每个输入自己选最优路径！

• 神奇之处3：AI自动学会"分工哲学"（分解特征表示）

• 论文用了一个超接地气的实验：Jittered MNIST（就是手写数字图，但故意把数字乱挪位置）。

• 但是实验让科学家们惊呆了：

  • 第一层专家：集体变身"地图控"！它们不关心数字是几，只盯位置------"这个数字在左上角？归我管！""在右下角？归他管！" 科学家起名为"Where Experts"（位置专家）。
  • 第二层专家：秒变"分类狂魔"！它们说："我才不管位置，我只认数字是3还是8！" 于是起名成了"What Experts"（内容专家）。

• 没人指挥，但AI自己悟了：复杂任务=拆解小任务！ 这就是论文标题"Learning Factored Representations"（学习分解的特征表示）------说白了，AI学会了"先分地盘，再干专业事"。

2.4 大模型时代的 MoE

进入 Transformer 时代后，MoE 技术成为了突破模型规模瓶颈的关键。Google 在这一领域进行了密集的探索，通过 GShard、Switch Transformer 和 GLaM 等一系列工作，确立了现代大规模 MoE 的技术范式。

---

三、手撕大模型生成策略

• 前两节通过 Llama2 和 MoE，深入理解了大模型的网络架构（即"大脑"是如何构造的）。
• 但仅有架构还不够，模型前向传播输出的仅仅是概率分布（Logits） ，如何将这些概率一步步转化为流畅的文本，就是本节要探讨的核心------解码策略。

---

总结