【斯坦福 CS336】第一章：模型架构与基础

# 概述

本章是斯坦福 CS336 课程的第一章，也是本系列笔记的开篇章节。我们将介绍模型架构与基础知识，探讨模型设计的核心原则和常见的架构模式，为后续章节的学习奠定基础。

事先声明，本笔记并不会覆盖课程中的所有细节，不属于手把手的教程，而是重点总结和提炼出核心内容，帮助大家更好地理解和应用所学知识。

好了，废话不多说，我们直接进入正题。

# Tokenizer

Tokenizer 是自然语言处理中的一个重要组件，负责将文本转换为模型可以理解的输入格式。这部分内容属于模型架构的基础知识，理解 Tokenizer 的工作原理对于设计和优化模型至关重要，因此也是本章的课后作业之一。我们将这部分内容放在下一篇笔记中进行详细讲解。

# PyTorch 基本知识

这一部分汇总了课程中提及的部分 PyTorch 基础知识，这只是对于 PyTorch 的一个简要 补充，建议读者结合官方文档进行更深入的学习。

# 浮点数类型

PyTorch 张量中最常见的数据类型是浮点数类型，它们用于表示实数。不同的浮点数类型在存储空间、精度和计算性能上有所差异，选择合适的浮点数类型对于模型的训练和推理效果有重要影响。

以下是 PyTorch 浮点类型的对比表：

一般会使用 FP32 (torch.float32) 进行训练，而在推理阶段可以选择 FP16 (torch.float16) 或 BF16 (torch.bfloat16) 来加速计算，具体选择取决于硬件支持和精度需求。

# 张量切分

PyTorch 中有许多操作可以切分张量，但得到的是原始张量的视图，共享内存，而不是新的张量副本。这意味着对切分后的张量进行修改会影响原始张量。

以下面张量 x 为例：

x = torch.tensor([[1, 2, 3],

                  [4, 5, 6]])

• y = x[0] ：获取张量 x 的第一行，即 [1, 2, 3] 。
• y = x[:,1] ：获取张量 x 的第二列，即 [2, 5, 8] 。
• y = x.view(3,2) ：将张量 x 重塑为形状为 (3, 2) 的张量，即 [[1, 2], [3, 4], [5, 6]] 。
• y = x.transpose(0,1) ：交换张量 x 的第 0 维和第 1 维，即得到 [[1, 4], [2, 5], [3, 6]] 。

以上操作都不会创建新的张量副本，而是返回原始张量的视图。因此，对 y 进行修改会影响 x 。

有些操作会使得张量存储空间不再连续，例如 transpose 。在这种情况下，如果需要对张量进行进一步的操作（如传递给某些函数），可能需要调用 contiguous() 方法来创建一个新的连续张量副本。

所谓张量的 “连续性”，是指张量在内存中的存储方式是否是连续的。如果一个张量是连续的，那么它的元素在内存中是按顺序存储的，这对于某些操作（如矩阵乘法）来说是必要的。

# 浮点运算

FLOP（Floating Point Operations）是衡量计算复杂度的一个重要指标，表示浮点运算的次数。在深度学习中，FLOP 通常用于评估模型的计算需求和性能。对应的，FLOPS（Floating Point Operations Per Second）则表示每秒钟可以执行的浮点运算次数，是衡量计算设备性能的一个关键指标。通过估计模型的 FLOP，以及了解硬件的 FLOPS，可以帮助我们评估模型在特定硬件上的运行速度和效率。

下面以一个例子来说明如何计算 FLOP：考虑一个简单的矩阵乘法操作 C = A @ B ，其中矩阵 A 的形状为 (m, n) ，矩阵 B 的形状为 (n, p) 。在这种情况下，计算矩阵 C 需要进行 m * n * p 次乘法和 m * n * p 次加法，总共需要 2 * m * n * p 次浮点运算。因此，这个矩阵乘法操作的 FLOP 为 2 * m * n * p 。

事实上，硬件的 FLOPS 通常远高于实际应用中的计算速度。这是因为实际应用中存在许多其他因素会影响性能，例如内存带宽限制、缓存命中率、并行化效率等。因此，MFU (Model FLOPS Utilization) 被用来衡量模型在实际运行中利用硬件 FLOPS 的效率。MFU 计算公式如下：

$$\text{MFU} = \frac{\text{模型实际使用的 FLOP}}{\text{硬件的理论 FLOPS} \times \text{运行时间}}$$

除了前向传播阶段的 FLOP 计算外，训练阶段的 FLOP 计算还需要考虑反向传播阶段的计算量。通常情况下，反向传播的 FLOP 大约是前向传播的两倍（分别计算输入的梯度和参数的梯度），因此在估计训练阶段的总 FLOP 时，可以将前向传播的 FLOP 乘以三来得到一个近似值。

综上，对于矩阵乘法而言，模型在训练阶段的 FLOP 可以表示为 6 倍的数据点数量乘以参数数量。

# 模型显存

模型在训练和推理过程中，显存占用是一个重要的考虑因素。显存主要用于存储模型参数、激活值（中间计算结果）以及优化器状态等。

优化器状态的显存占用取决于所使用的优化器类型。例如，Adam 优化器需要存储每个参数的动量和二阶矩估计，因此其显存占用通常是参数数量的两倍。而 SGD 优化器则只需要存储参数本身，因此显存占用为参数数量本身。

激活值的显存占用取决于模型的架构和输入数据的大小。对于深度神经网络，激活值的显存占用通常会随着网络深度的增加而增加。需要根据每一层的输出形状来计算总的激活值显存占用。

# 大模型架构设计

众所周知，当下的大模型，基本采样的都是 Transformer 架构，从中衍生出了各种变体，例如 GPT 系列、BERT 系列等。这些模型在自然语言处理任务中表现出色，成为了当前的主流选择。

# Transformer

首先我们来回顾以下最初的 Transformer 架构。Transformer 由 Vaswani 等人在 2017 年提出，彻底改变了自然语言处理领域。

如图所示，我们可以用几个关键词来总结初版 Transformer：编码器 - 解码器架构、多头自注意力机制、sine/cosine 位置编码、ReLU 前馈网络、LayerNorm、Post-LN 等。至今 Transformer 已经出现了许多变体和改进版本，区别主要集中在这些核心组件的选择上。

下图展示了本课程所用的 Transformer 变体，其中很多选择也是目前主流大模型的方案，包括：

• 使用 Pre-LN 结构以提升训练稳定性。
• 使用 RMSNorm 替代 LayerNorm。
• 使用 RoPE（旋转位置编码）替代传统的绝对位置编码。
• 采用 SwiGLU 激活函数而非 ReLU。
• 移除模型参数中的偏置项（Bias-free）。

本课程 Assignment 1 的其中一个任务就是实现上述 Transformer 架构，这部分内容会在后续的笔记中进行详细讲解。本节将从理论的角度介绍 Transformer 的核心组件和设计原则，包括但不限于上述提到的内容。

# Pre-LN 与 Post-LN

下图展示了 Post-LN（左）与 Pre-LN（右）的模型结构和计算流程：

Post-LN 将 LayerNorm 置于残差相加之后，而 Pre-LN 则将其置于子层计算之前。Pre-LN 的核心设计理念是让 LayerNorm 不干扰主残差信号路径，使得输入信号和梯度可以在主干路径上更顺畅地传递。相比之下，Post-LN 在每一层都会重新归一化主路径上的输出，这在深层网络中容易导致梯度流动的复杂化。

在优劣势方面，Post-LN（如 BERT）虽然在小模型上可能具有更强的表达力，但其训练极不稳定，通常需要极其细致的学习率预热（Warm-up）。Pre-LN 则显著提升了训练稳定性，支持更深层的架构（如 GPT-3、Llama 等现代大模型）稳定收敛，因此已成为当前工业界的事实标准。

# RMSNorm

RMSNorm（Root Mean Square Normalization）是一种归一化技术，与 LayerNorm 计算均值和方差不同，RMSNorm 仅计算输入的均方根，从而简化了归一化过程。

LayerNorm 对神经元输出进行 “中心化” 与 “缩放”。其公式为：

$$\text{LayerNorm}(x) = \frac{x - E[x]}{\sqrt{Var[x] + \epsilon}} \cdot \gamma + \beta$$

RMSNorm 认为 “中心化” 非必须，仅保留 “缩放” 以简化计算。其公式为：

$$\text{RMSNorm}(x) = \frac{x}{\sqrt{E[x^2] + \epsilon}} \cdot \gamma$$

现代主流大模型基本首选 RMSNorm 而不是 LayerNorm，原因在于 RMSNorm 舍弃了计算均值和减去均值的步骤，在大规模预训练中，这种微小的简化累积起来能显著提升显存吞吐量和训练速度。此外，实验表明，中心化对模型表达能力的提升极小，归一化的核心价值在于稳定梯度，RMSNorm 并不会损失模型性能。

# 激活函数

激活函数一般用于前馈神经网络中，下面列出了常用的各种激活函数对应的表达式：

ReLU：

$$FFN_{ReLU}(x) = max(0, xW_1)W_2$$

GeLU：

$$FFN_{GeLU}(x) = GeLU(xW_1)W_2\\ GeLU(x) = x \cdot \Phi(x) = x \cdot P(X \leq x), \quad X \sim N(0, 1)$$

GeGLU：

$$FFN_{GeGLU}(x) = (\text{GeLU}(xW_1) \otimes xV)W_2$$

SwiGLU：

$$FFN_{SwiGLU}(x) = (\text{Swish}(xW_1) \otimes xV)W_2\\ Swish(x) = x \cdot \text{sigmoid}(x) = \frac{x}{1 + e^{-x}}$$

GeGLU 和 SwiGLU 均采用了门控机制，通过引入额外的线性变换 $xV$ 来调节激活值，从而提升模型的表达能力和性能。实验表明，GeGLU 和 SwiGLU 在多个自然语言处理任务中均优于传统的 ReLU 和 GeLU，成为现代大模型的首选激活函数。

# 位置编码

Transformer 需要引入位置编码来捕捉输入序列中元素的位置信息。常见的绝对位置编码方法是使用正弦和余弦函数生成固定的编码向量，公式如下：

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)\\ PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

其中 $pos$ 是位置索引，$i$ 是维度索引，$d_{model}$ 是模型的隐藏层维度。绝对位置编码为每个位置生成一个固定的向量，使得模型能够区分不同位置的输入。

后来，**RoPE（Rotary Position Embedding）** 被提出作为一种相对位置编码方法，其优势在于它能够更自然地捕捉输入序列中元素之间的相对位置关系，并且在长序列处理时表现更好，因此被现代大模型广泛采用。

RoPE 通过旋转输入向量来引入位置信息。在计算注意力权重时，RoPE 会将输入向量（通常是查询向量和键向量）与一个基于位置的旋转矩阵相乘，公式如下：

$$RoPE(x, pos) = x \cdot R(pos)$$

其中 $R(pos)$ 是一个旋转矩阵，以二维为例，旋转矩阵的形式如下：

$$R(pos) = \begin{bmatrix}\cos(\theta_{pos}) & -\sin(\theta_{pos}) \\ \sin(\theta_{pos}) & \cos(\theta_{pos})\end{bmatrix}$$

对于高维输入，$R(pos)$ 就是个块对角矩阵，每个块都是一个上述的二维旋转矩阵。其中的旋转角度 $\theta_{pos,i}$ 通常与位置索引和维度索引相关，例如 $\theta_{pos,i} = pos / 10000^{2i/d_{model}}$，与绝对位置编码中的频率设置类似。通过这种方式，RoPE 能够捕捉输入序列中元素之间的相对位置关系，而不仅仅是绝对位置。

# 超参数

首先介绍前馈神经网络（FFN）的隐藏层维度 $d_{ff}$。在经典 Transformer 中，$d_{ff}$ 通常设置为模型隐藏层维度 $d_{model}$ 的 4 倍，即 $d_{ff} = 4 \times d_{model}$。

但在 GLU 变体中，$d_{ff}$ 通常设置为 $\frac{8}{3} \times d_{model}$。这是因为 GLU 结构引入了门控机制，增加了额外的模型参数，为了保持整体的参数规模不变，$d_{ff}$ 需要相应地调整。

从数学上来说就是，在标准的 FFN 中，包含两个线性变换层，总参数量为 $\approx 2 \times d_{model} \times d_{ff}$，当 $d_{ff} = 4d_{model}$ 时，总参数量 $\approx 8d_{model}^2$。而在 GLU 中，包含三个线性变换层，总参数量为 $\approx 3 \times d_{model} \times d_{ff}$，当 $d_{ff} = \frac{8}{3} d_{model}$ 时，总参数量同样 $\approx 8d_{model}^2$。

在多头注意力机制中，头数 $h$ 的选择通常是 $d_{model}$ 的一个因子，以确保每个头的维度 $d_k = \frac{d_{model}}{h}$ 是一个整数。常见的选择包括 8 头、16 头等。例如，对于 $d_{model} = 512$，可以选择 $h = 8$，使得每个头的维度 为 $d_k = 64$。实际上，所有头的维度之和并不一定非得等于 $d_{model}$，这只是一个经验做法。

# MoE

MoE（Mixture of Experts）是一种稀疏模型架构，通过将传统稠密模型的参数拆分为多个独立的专家子模块，并由门控网络对输入进行动态路由，实现按需激活和计算。如下图所示，在 Transformer 中，MoE 通常替换 FFN 部分，每个专家本质上是一个独立的 FFN，输入由门控网络选择少量专家参与计算，从而在扩大模型容量的同时控制计算成本。

# 路由函数

在 MoE 架构中，路由机制决定了输入的 Token（令牌）如何分配给不同的 Expert（专家）。根据分配维度的不同，主要分为以下三类：

• Token Chooses Expert：每个 Token 独立地从所有专家中选择评分最高的 Top-K 个专家进行计算。
• Expert Chooses Token：设定每个专家固定的容量，由每个专家从所有 Token 中选择评分最高的 Top-K 个进行处理。
• Global Routing via Optimization：不再是局部的 Token 或专家单向选择，而是将路由看作一个全局分配问题（通常建模为线性规划或最优传输问题）。

至于路由函数的具体实现，常见的方法包括：

• 基于点积的注意力机制：计算输入 Token 与专家权重的点积，选择得分最高的专家。
• 基于 MLP 的评分网络：使用一个小型的前馈神经网络对输入进行评分，选择得分最高的专家。
• 基于哈希的方法：通过哈希函数将输入映射到特定的专家，适用于大规模专家池。

无论采用哪种路由实现方法 ，目标都是在保持模型容量的同时，最大限度地减少计算开销，从而提升模型的效率和性能。

# 训练

MoE 训练的主要挑战在于稀疏决策的不可微性。虽然常用的 Top-K 路由在选中专家后，梯度可以顺着权重回传给路由器 ，但 “选择专家” 这一离散动作本身无法提供梯度，导致以下问题：

1. 缺乏探索 (Lack of Exploration)： 如果没有额外机制，路由器会迅速锁定初始状态下表现稍好的专家，导致其他专家永远得不到训练。
2. 优化困难： 离散的专家切换逻辑使得标准梯度下降难以优化路由策略 。

为了解决不可微性，业界探索了以下几种主要路径：

• Reinforce 算法：将路由视为策略选择，利用强化学习算法优化。虽然这种方法在理论上是 “正确解”，但由于梯度方差大且实现复杂度高，目前在主流大规模模型中并不常用。
• Stocastic Approximation：通过在门控的 Logits 计算中引入噪声或随机扰动使模型在训练早期能够探索不同的专家。
• Heuristic Balancing Losses：为了防止某些专家过载而另一些闲置，在总损失中加入负载均衡损失，鼓励路由器在专家之间分配负载，从而促进所有专家的训练。这里的负载均衡损失可以是基于专家被选择的频率，也可以从系统的角度，如基于设备利用率进行设计。

# 一个有趣的问题

如果我们尝试调用 GPT-4 的 API 来回答问题，并且设置了 temperature=0 ，理论上模型应该每次都给出相同的答案，因为温度为零意味着模型在采样时会选择概率最高的输出。但是事实上，即使在这种设置下，GPT-4 仍然可能给出不同的答案。我们猜测这很可能与 MoE 架构中的路由机制有关。

要知道，大模型厂商通过会将若干用户的请求打包成一个批次（batch）来进行并行处理，以提高计算效率。在这种情况下，即使你的输入是相同的，由于批次中的其他请求可能不同，路由器在选择专家时可能会受到影响，从而导致不同的专家被激活，最终产生不同的输出结果。