模型总览: 一个预训练模型由哪些部分组成?

这一节的目标，是先把一个基础的 LLM 模型到底是什么这件事整理清楚。
先不进入训练、loss、优化器这些内容，而是只回答一个更基础的问题：如果只看模型结构，一个预训练语言模型到底由哪些部分组成？这些部分分别在做什么？常见的设计取舍又是什么？

这一节准备围绕下面这些问题来组织：

1. 一个预训练模型大体可以划分为哪些部分？
2. Embedding 层和位置编码层分别在做什么？
3. 激活函数、Norm、Dropout 这类“看起来像细节”的设置，为什么其实很重要？
4. FFN 和 Attention 这两个核心模块分别在做什么？
5. 后来的改进，比如 RoPE、FlashAttention、GQA、KV Cache，又分别是在解决什么问题？

Q1: 一个预训练模型大体可以划分为哪些部分?

如果先把训练流程、数据处理、优化器这些东西都放到一边，只看“模型定义”本身，那么一个典型的 GPT 类预训练模型，大致可以分成下面几部分：

1. 输入嵌入层（Embedding）
2. 位置信息注入（Position Encoding / Positional Representation）
3. 多层 Transformer Block
4. 输出投影层（LM Head）

如果把这条链路写成最简化的形式，可以先记成：

token id -> embedding -> 多层 transformer block -> logits

进一步展开，可以写成：

token id [n,]
-> token embedding [n , dim]
-> 注入位置信息 [n, dim]
-> 重复多层 [attention + norm + ffn] [n,dim]
-> 最终 hidden state [n, dim]
-> 词表投影 [n, V]
-> logits [n, V]

这个流程图本身是对的，但如果只停在这一步，还是会有一个问题：
我们知道模型“经过了哪些部分”，却还不知道每一层到底接收什么、输出什么、以及发生了什么变换。

所以接下来最好把这个流程进一步写成一个更正式的“整体定义”。

MiniMind 里这条主链路的骨架也写得很直接。模型主干 MiniMind_Dense 负责 Embedding、多层 MiniMindBlock 和最后的 Norm；MiniMindForCausalLM 再在外面包一层 lm_head，把 hidden states 投影回词表空间。对应实现来自 model_minimind.py 和 model_minimind.py。

class MiniMind_Dense(torch.nn.Module):
    """
    Dense模型的定义
    """

    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.config = config
        self.vocab_size, self.num_hidden_layers = (
            config.vocab_size,
            config.num_hidden_layers,
        )
        # Embedding
        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size
        )  # [vocab_size , embedding_size]
        # Dropout and norm
        self.dropout = nn.Dropout(config.dropout)
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # Attention Layers
        self.layers = nn.ModuleList(
            [MiniMindBlock(l, config) for l in range(self.num_hidden_layers)]
        )


class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = MiniMindConfig

    def __init__(self, config: MiniMindConfig = None):
        self.config = config or MiniMindConfig()
        super().__init__(self.config)
        self.model = MiniMind_Dense(self.config)
        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
        self.model.embed_tokens.weight = self.lm_head.weight
        self.OUT = CausalLMOutputWithPast()

整体定义

如果只从函数角度看，一个不带 LM head 的 Transformer 模型可以理解成下面这个映射：

$$ f_\theta(t_1,t_2\dots t_n) \in \mathbb{R}^{n\times d} $$

其中：

• \( \theta \) 表示整个模型的参数
• \( d \) 表示 hidden size
• \( n \) 表示输入序列长度
• \( V \) 表示词表大小

这个式子的意思是：

• 输入是一串长度为 \( n \) 的 token id

• 每个 token id 的取值范围在 \( 0 \) 到 \( V-1 \) 之间

• 输出是一个形状为 \( [n, d] \) 的矩阵，用 hidden state 表示整个序列在各个位置上的编码结果

• 最后再经过具体的 LM head，比如使用一个 \( [d, V] \) 的投影矩阵，把它变成形状为 \( [n, V] \) 的 logits 张量，再经过 softmax 变成概率分布。这就是最基本的语言模型

• 不同任务下，最后接的 head 也可能不一样

模型的整体输入是什么?

$$ t_1, t_2, \dots, t_n $$

表示输入 token 序列，其中：

• \( t_i \) 表示第 \( i \) 个 token 的 id
• \( n \) 表示序列长度

在实际代码里，如果考虑 batch 维，通常更常见的 shape 是：

$$ T \in \mathbb{N}^{B \times n} $$

其中：

• \( B \) 表示 batch size
• \( n \) 表示序列长度

也就是说，输入本质上是一批 token id 序列，而不是原始字符串，也不是 embedding 向量。

每一层的输入输出 shape 和发生的变换

1. token id -> token embedding

Embedding 层接收离散 token id：

$$ T \in \mathbb{N}^{B \times n} $$

输出连续表示：

$$ X^{(0)} \in \mathbb{R}^{B \times n \times d} $$

其中：

• \( d \) 表示 hidden size，也就是 embedding 维度

这一层发生的变换，本质上是：

• 对每个 token id 查 embedding 矩阵中的一行
• 把离散符号变成连续向量

到这一步之后，输入已经从 token id 变成了连续向量。后面的大部分计算，基本都不会再改动它的 shape ([n,d])，而是在同一个表示空间里不断做变换，让这些向量逐渐混入上下文信息。从这个角度看，Transformer 这个名字其实很直观，它后面一直在做的就是“变换表示”这件事。

MiniMind 里这一层的实现也很直接，就是一个 nn.Embedding，然后在 forward 里把 input_ids 查表成 [batch_size, seq_len, hidden_size]。对应实现来自 model_minimind.py 和 model_minimind.py。

# Embedding
self.embed_tokens = nn.Embedding(
    config.vocab_size, config.hidden_size
)  # [vocab_size , embedding_size]


# Embedding 要过Dropout!
# [batch_size , seq_len] -> [batch_size , seq_len , hidden_size] -> [batch_size , seq_len , hidden_size]
hidden_states = self.dropout(self.embed_tokens(input_ids))

2. token embedding -> 注入位置信息

如果使用最基本的位置表示思路，那么这一层的输入和输出 shape 都仍然是：

$$ X^{(0)} \in \mathbb{R}^{B \times n \times d} $$

和

$$ \tilde{X}^{(0)} \in \mathbb{R}^{B \times n \times d} $$

这一层发生的事情不是改 shape，而是把位置信息注入表示里。
Position embedding 的重要性在于，模型不只需要知道“有哪些 token”，还需要知道“它们出现在哪些位置”。

如果用 RoPE 这类方案，那么位置信息不一定直接加在输入 embedding 上，而可能延后到 attention 的 Q/K 计算阶段再进入。

MiniMind 没有使用固定Position Embedding，而是使用了RoPE,先预计算 RoPE 所需的 freqs_cos 和 freqs_sin，再在 forward 里按当前位置切出来传给各层 attention。RoPE请参考RoPE小节的代码实现, 对应实现来自 model_minimind.py 和 model_minimind.py。

# RoPE vector
freqs_cos, freqs_sin = precompute_freqs_cis(
    dim=config.hidden_size
    // config.num_attention_heads,  # dim for each attention heads
    end=config.max_position_embeddings,
    rope_base=config.rope_theta,
    rope_scaling=config.rope_scaling,
)
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
self.register_buffer("freqs_sin", freqs_sin, persistent=False)


# [seq_len,]
position_embeddings = (
    self.freqs_cos[start_pos : start_pos + seq_length],
    self.freqs_sin[start_pos : start_pos + seq_length],
)

3. 多层 Transformer Block

Transformer 层是模型的核心模块。它由许多结构相同的 Transformer block 重复堆叠起来，对序列表示一层一层地继续做变换。Transformer 的层数本身也是一个很重要的超参数，而每一层的输入输出 shape 通常都保持一致，这样这个变换过程才能持续往下传。

假设模型一共有 \( L \) 层 Transformer block，那么第 \( \ell \) 层可以写成：

$$ X^{(\ell+1)} = \mathrm{Block}^{(\ell)}\left(X^{(\ell)}\right) $$

其中：

• \( \ell = 0,1,\dots,L-1 \)
• 每一层的输入输出 shape 通常都保持不变

也就是说：

$$ X^{(\ell)} \in \mathbb{R}^{B \times n \times d} $$

这一层内部虽然会发生很多计算，比如：

• attention
• residual
• norm
• ffn

但从外部看，最重要的事情是：

• shape 基本保持不变
• 表示的语义不断被更新
• 上下文信息逐层混合进每个位置的表示里

4. 最终 hidden state -> 词表投影

经过最后一层 Transformer block 后，得到最终 hidden state：

$$ H \in \mathbb{R}^{B \times n \times d} $$

接下来就要进入语言模型最后那一层，也就是常说的 LM head。

对于最基本的语言模型来说，LM head 的作用很直接：
把每个位置的 hidden state 映射成“这个位置对整个词表的预测分数”。

如果只写单个位置的形式，可以记成：

$$ z_i = h_i W_{\text{lm}} + b $$

其中：

• \(h_i\) 表示第 \(i\) 个位置的 hidden state,这里的\(h_i \in \mathbb{R}^{1 \times d} \) 实际是一个行向量. 这是为了方便之后把batch维度加到它前面.
• \(W_{\text{lm}}\) 表示输出投影矩阵
• \(b\) 表示偏置项
• \(z_i\) 表示这个位置对应的词表分数向量，也就是 logits，即还没有经过 softmax 的原始分数（raw scores）

如果一次处理整个序列，那么经过 LM head 之后，最后一个维度就会从 hidden size 变成词表大小：

$$ Z \in \mathbb{R}^{B \times n \times V} $$

也就是说，这一层发生的事情可以理解成：

• 对每个位置的 hidden state
• 输出它对整个词表中每个 token 的原始分数

这些原始分数就是 logits表示未经过 softmax 的原始分数（raw scores）.

这里其实可以顺手记住一句话：

• Transformer block 负责得到“上下文化之后的表示”
• LM head 负责把这个表示重新翻译成词表上的分数

为什么这里叫 LM head?

因为这是 Language Model head，也就是语言模型任务对应的输出头。

最基本的语言模型任务是 next-token prediction。
也就是说，模型最后必须回答这样一个问题：

• 给定当前位置之前的上下文
• 下一个 token 最可能是词表中的哪一个？

所以最自然的输出方式，就是对整个词表打分。
LM head 正是在完成这一步。

为什么任务变了，head 也会跟着变?

这里也能顺手看出一个很重要的事实：
模型主体和最后的输出头，其实是两层不同的东西。

模型主体负责提取表示，而最后的 head 负责把表示变成某个具体任务需要的输出。

所以如果任务变了，最后这一层也经常会跟着变：

• 对于基础语言模型：输出是词表上的 logits，所以使用 LM head
• 对于分类任务：最后可能只需要输出几个类别分数，所以会换成分类 head
• 对于回归任务：最后可能只输出一个或几个连续值，所以也会换成别的输出层

也就是说，LM head 不是“所有模型最后一层的唯一形式”，而是“语言模型任务最基本、最自然的输出头”。

模型的整体输出是什么?

如果先不考虑 batch 维，那么模型对一个长度为 \( n \) 的输入序列，输出可以写成：

$$ Z \in \mathbb{R}^{n \times V} $$

其中：

• 第一个维度对应序列中的位置
• 第二个维度对应词表中的 token 数目
• 矩阵的每一项表示“这个位置对应这个 token 的预测分数(logit)”

如果把 batch 维也补上，更完整的写法通常是：

$$ Z \in \mathbb{R}^{B \times n \times V} $$

其中：

• 第一个维度 \( B \) 是 batch size
• 第二个维度 \( n \) 是序列长度
• 第三个维度 \( V \) 是词表大小

这个张量里的每一个 \( Z_{b,i,:} \)，都表示：

• 第 \( b \) 个样本
• 第 \( i \) 个位置
• 对整个词表中每个 token 的预测分数

所以如果用一句话概括最基本的语言模型输入输出，可以写成：

• 输入：形状为 \( [B, n] \) 的 token id 序列
• 输出：形状为 \( [B, n, V] \) 的 logits 张量

而模型中间所有模块，本质上都在做一件事：
把输入 token 序列逐层变成更有上下文信息的隐藏表示，最后再投影回词表空间。

从这个角度看，模型的核心其实是在做两件事：

1. 把离散 token id 变成高维连续表示
2. 在多层变换中不断让上下文信息彼此交互，最后再投影回词表空间

Q2: Embedding 层在做什么?

Embedding 层的作用，是把离散 token id 映射成连续向量。

如果词表大小记作 \( V \)，隐藏维度记作 \( d \)，那么 token embedding 矩阵通常可以写成如下的形式。事实上，Embedding 层一般就是这个矩阵，它本身也是模型参数的一部分，会参与训练。对于比较小的模型，embedding 矩阵甚至会占参数量的大头：

$$ E \in \mathbb{R}^{V \times d} $$

其中：

• \( V \) 表示词表大小
• \( d \) 表示 hidden size，也可以理解成 embedding 维度

当一个 token id 输入 embedding 层时，本质上就是在这个矩阵里取出对应的一行。
如果输入序列长度是 \( n \)，那么 embedding 层的输出可以写成：

$$ X \in \mathbb{R}^{n \times d} $$

从直觉上看：

• 词表决定“有哪些离散符号”
• embedding 决定“这些离散符号如何在连续空间里表示”

经验问题 1: embedding 维度通常设多少?

这件事通常不会单独决定，而是直接跟整体 hidden size 绑定。

在大多数 Transformer 里，token embedding 维度通常直接等于 hidden size，也就是：

$$ d_{\text{embed}} = d_{\text{model}} $$

原因也很直接：

• embedding 输出之后，后面所有 Transformer block 都在同一个维度空间里工作
• 如果 embedding 维度和 hidden size 不一致，就需要额外投影层

所以在大多数 decoder-only Transformer 里，embedding 维度通常就直接等于模型主干的 hidden size。

这里直有几组真实模型配置：

从这些配置里可以直接看到两件事：

• embedding 维度通常就是 hidden size，本身并不会单独设成另一套数
• hidden size 的取值通常会和头数严格对齐，保证 \( d_{\text{model}} / n_{\text{heads}} \) 是一个整齐的整数

经验问题 2: hidden size 一般怎么定?

这一节里，最值得记住的不是“hidden size 有什么经验最优值”，而是它其实是一个被多种约束一起卡住的量。先看最直接的一条结构约束：

$$ d_{\text{model}} = n_{\text{heads}} \cdot d_{\text{head}} $$

其中：

• \(d_{\text{model}}\) 表示 hidden size
• \(n_{\text{heads}}\) 表示注意力头数
• \(d_{\text{head}}\) 表示每个 head 的维度

也就是说，hidden size 首先要和注意力结构对齐，而不是随便写一个数。除此之外，它还会非常直接地影响参数量。粗略看：

• attention 里的投影矩阵参数大致和 \(d_{\text{model}}^2\) 同阶
• FFN 的参数大致和 \(d_{\text{model}} \cdot d_{\text{ff}}\) 同阶
• 如果 \(d_{\text{ff}} \approx 4 d_{\text{model}}\)，那么 FFN 参数通常也会近似随 \(d_{\text{model}}^2\) 增长

所以 hidden size 不是线性地影响模型大小，而往往会以接近平方的方式影响参数量、显存占用和计算量。

如果真要落到操作上，我觉得比较实用的顺序是：

1. 先定模型总规模 比如是 20M、100M，还是 1B 以上。这个规模先把 hidden size 的上限大致卡住。

2. 再定层数和头数 同样的总参数量，可以做成更宽更浅，也可以做成更窄更深，所以 hidden size 不能脱离层数单独看。

3. 最后让 \( d_{\text{model}} \) 和 \( n_{\text{heads}} \) 对齐 也就是让它能自然整除，这样 \( d_{\text{head}} \) 会比较规整，实现也更自然。

如果只是做教学模型或者小模型，那么最实用的结论其实可以压缩成一句话：先定模型规模，再定层数和头数，最后让 hidden size 去匹配 attention 结构。

经验问题 3: 输出层为什么常常和 embedding 共享权重?

很多 GPT 类模型里，会把输入 embedding 矩阵和输出投影层权重做共享 (直接作为LM head 的投影矩阵)，也就是常说的 weight tying。

如果输出层写成：

$$ z = h W_{\text{out}} + b $$

其中：

• \( h \in \mathbb{R}^{1 \times d} \) 表示某个位置的 hidden state
• \( W_{\text{out}} \in \mathbb{R}^{d \times V} \) 表示输出投影矩阵

其正好就是Embedding矩阵的转置,这个操作在数学上也可以理解成,输出的hidden state直接和embedding矩阵的行做内积,得到每个词的分数就是在衡量和此表里面每个词的相似度.也就是说embedding不但在编码时刻进行训练,同样也在输出的时候受到约束.

$$ W_{\text{out}} = E^\top $$

这样做的好处通常包括：

• 参数更省
• 输入表示空间和输出词表空间更一致
• 在实践中常常有不错的效果

所以 embedding 层虽然看起来很基础，但它其实已经和模型参数规模、输出层设计直接连在一起了。

MiniMind 这里也直接使用了 weight tying，也就是让输入 embedding 的权重和 lm_head 共享。对应实现来自 model_minimind.py。

self.model = MiniMind_Dense(self.config)
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
self.model.embed_tokens.weight = self.lm_head.weight

Q2.1: 位置编码层在做什么?

如果只有 token embedding，而没有位置信息，那么模型虽然能知道“有哪些 token”，却不知道“它们以什么顺序出现”。

这也是为什么位置表示是必须的。

假设 token embedding 序列写成：

$$ X = [x_1, x_2, \dots, x_n] $$

其中：

• \( x_i \in \mathbb{R}^{d} \) 表示第 \( i \) 个位置的 token embedding

最早的一种思路，是直接构造位置向量：

$$ p_1, p_2, \dots, p_n $$

然后做相加：

$$ \tilde{x}_i = x_i + p_i $$

这种思路很直观：
token 内容由 \( x_i \) 提供，位置信息由 \( p_i \) 提供。

经验问题 1: 为什么今天很多模型不再用最早那种绝对位置 embedding?

这个问题其实也可以直接从真实模型配置里看出来。

从这些配置里可以很直接地看到：早期 GPT-2 这一类模型更接近“给每个位置单独学一个位置向量”；而更现代的 Llama、Qwen 这类模型，已经转向了和 attention 计算绑定得更紧的 RoPE。

绝对位置 embedding 当然也能工作，但它有几个比较明显的局限：

• 位置表示通常和最大上下文长度绑定
• 超出训练长度时泛化可能较弱
• 位置关系是“加在输入上”的，不是直接进入 attention 匹配过程

这也是为什么后来越来越多模型转向 RoPE 这类位置表示方式。
不过 RoPE 本身已经不只是“位置编码层的一个小变体”了，它和 attention 的计算过程紧密绑定，所以我把它单独拆成了一节：RoPE: 位置编码如何进入 Attention?

Q3: 激活函数、Norm、Dropout: FNN的重要细节?

FNN 在 LLM 里不只是一个“顺手接在 attention 后面的两层线性层”，它往往占据了大量参数，也承担了非线性表达能力。
所以在真正进入 FNN 结构本身之前，先把和它密切相关的几个细节讲清楚是很有必要的：激活函数、Norm 和 Dropout。

这些东西看起来像细节，但它们其实直接影响：

• 模型是否容易训练
• 梯度是否稳定
• 参数能不能有效发挥作用
• 最终收敛出来的效果如何

3.1 激活函数: 为什么今天更常见是 GELU、SiLU、SwiGLU?

经验上，今天的语言模型更常见的选择往往不是最早期的 ReLU，而是：

• GELU
• SiLU / Swish
• 或者进一步用于 gated MLP 的 SwiGLU

它们更常见，通常不是因为“名字更新”，而是因为在大模型训练里，它们往往表现得更平滑、更稳定，也更适合和现代 FFN 结构配合。

GELU

GELU（Gaussian Error Linear Unit）常见写法是：

$$ \mathrm{GELU}(x)=x ,\Phi(x) $$

其中：

• \( \Phi(x) \) 表示标准高斯分布的累积分布函数

实践中常见的近似写法是：

$$ \mathrm{GELU}(x)\approx 0.5x\left(1+\tanh\left(\sqrt{\frac{2}{\pi}}\left(x+0.044715x^3\right)\right)\right) $$

这也是很多代码里真正实现时更常看到的形式。

SiLU / Swish

SiLU 的形式写成：

$$ \mathrm{SiLU}(x)=x \cdot \sigma(x) $$

其中：

• \( \sigma(x) \) 表示 sigmoid 函数

也就是说：

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

所以 SiLU 也可以展开写成：

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

SwiGLU

SwiGLU 常见于更现代的 gated FFN 结构里。它的一个常见写法可以记成：

$$ \mathrm{SwiGLU}(x,W,V)=\mathrm{SiLU}(xW)\odot(xV) $$

其中：

• \( W, V \) 表示两组不同的投影矩阵
• \( \odot \) 表示逐元素乘法

它的直觉是：
一条分支负责产生门控，另一条分支负责提供内容，最后逐元素相乘，让 FFN 的表达更细致。

一个经验性结论

如果只说一个非常实用的结论，我会更倾向于这么记：

• 小模型时代，ReLU 很常见
• 现代 LLM 里，更主流的是 GELU、SiLU，以及进一步用于 gated MLP 的 SwiGLU

3.2 Norm: 为什么它几乎是必须的?

Norm 的核心作用，是稳定训练过程。
如果没有这一步，随着层数加深，表示的尺度和梯度都更容易失控。

在真实的 LLM 代码里，Norm 并不是作用在一个孤立的“抽象向量”上，而通常是作用在形状为：

$$ X \in \mathbb{R}^{B \times n \times d} $$

的张量上，其中：

• \( B \) 表示 batch size
• \( n \) 表示序列长度
• \( d \) 表示 hidden size

对于 Transformer 里的 LayerNorm 或 RMSNorm，最常见的做法都是：
对最后一个维度 \( d \) 做归一化，也就是对每个 batch、每个位置对应的那一个 hidden vector 单独归一化。

也就是说，如果把某个位置的表示记成：

$$ x_{b,t,:} \in \mathbb{R}^{d} $$

那么 Norm 实际上是分别作用在每一个 \( x_{b,t,:} \) 上，也就是只沿着 feature 维度做归一化，而不会跨 batch 维或跨序列位置去做。这一点和 CNN 里的 BatchNorm、InstanceNorm 很不一样。一个直观理解是：在语言模型里，真正对应“同一个位置内部特征”的是最后这个 feature 维度；而 batch 维和序列位置维本身并不是应该拿来混在一起做归一化的量。

LayerNorm

如果先只看某一个位置的 hidden vector，可以写成：

$$ x=(x_1,x_2,\dots,x_d) $$

其中：

• \( d \) 表示特征维度

那么 LayerNorm 里先计算均值和方差：

$$ \mu=\frac{1}{d}\sum_{i=1}^{d}x_i,\qquad \sigma^2=\frac{1}{d}\sum_{i=1}^{d}(x_i-\mu)^2 $$

再做归一化：

$$ \mathrm{LayerNorm}(x)_i=\gamma_i\frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta_i $$

其中：

• \( \gamma_i,\beta_i \) 是可学习参数
• \( \epsilon \) 是数值稳定项

如果放回真实张量 \( X \in \mathbb{R}^{B \times n \times d} \) 来看，那么 LayerNorm 做的事情就是：

• 固定 \( b ) 和 \( t \)
• 对 \( X_{b,t,:} \) 这个长度为 \( d \) 的向量求均值和方差
• 再对这个向量本身做归一化

所以它是“对每个 token 的特征维度做归一化”，而不是“在整段序列上做归一化”。

RMSNorm

在很多较新的 LLM 里，RMSNorm 更常见。它通常不显式减去均值，而是直接根据均方根来做缩放：

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

$$ \mathrm{RMSNorm}(x)_i=\gamma_i\frac{x_i}{\mathrm{RMS}(x)+\epsilon} $$

其中：

• \( \gamma_i \) 是可学习缩放参数
• \( \epsilon \) 是数值稳定项

如果放回真实张量形状去理解，它和 LayerNorm 一样，最常见的也是对：

$$ X_{b,t,:} \in \mathbb{R}^{d} $$

这个最后一维向量逐个做处理。

为什么很多现代 LLM 喜欢 RMSNorm?

经验上，RMSNorm 常被认为：

• 形式更简单
• 参数和计算都更省一点
• 在大模型里表现很稳定

所以如果你在 MiniMind 或更现代的开源模型里看到 RMSNorm，这通常不是偶然，而是今天常见的设计取向。

Norm 在训练和推理时有区别吗?

对于 LayerNorm 和 RMSNorm 这类 Transformer 里常见的 Norm 来说，训练和推理阶段的行为通常没有本质区别。
也就是说：

• 训练时怎么按当前输入算
• 推理时也还是怎么按当前输入算

它不像 BatchNorm 那样会显式依赖“训练时统计量”和“推理时统计量”的切换。
这也是为什么在 Transformer / LLM 里，LayerNorm 和 RMSNorm 会比 BatchNorm 自然得多。

MiniMind 里 Norm 的实现和这里的讨论是直接对应的。它既保留了 RMSNorm，也额外实现了一个 LayerNorm 版本，两者都清楚标了 [batch_size, seq_len, dim] 这类 shape。对应实现来自 model_minimind.py。

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        """
        dim: embedding dim
        weights: [dim,]
        """
        super().__init__()
        self.eps = eps
        self.weights = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        # [batch_size, seq_len, dim] * [batchsize, seq_len , 1]
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """
        x [batch_size, seq_len, dim]
        return [barch_size,seq_len,dim]
        """
        return self.weights.type_as(x) * self._norm(x.float()).type_as(
            x
        )  # 处理不同的类型转化很重要


class LayerNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.beta = nn.Parameter(torch.zeros(dim))
        self.gamma = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        mean = torch.mean(x, -1, keepdim=True)  # [batchsize, seq_len , 1]
        var = (x - mean).pow(2).mean(-1, keepdim=True)  # [batchsize, seq_len , dim]
        inv_std = torch.rsqrt(var + self.eps)  # [batchsize, seq_len , dim]
        return (x - mean) * inv_std

    def forward(self, x):
        """
        x [batch_size, seq_len, dim]
        return [barch_size,seq_len,dim]
        """
        norm_x = self._norm(x.float()).type_as(x)
        return self.beta.type_as(x) + self.gamma.type_as(x) * norm_x

3.3 Dropout: 它到底在控制什么?

Dropout 的作用更偏向训练阶段的正则化。
它最粗略的形式可以写成：

$$ \tilde{x}=m\odot x $$

其中：

• \( x \) 表示输入
• \( m \) 表示随机 mask
• \( \odot \) 表示逐元素乘法

如果把保留概率写成 \( 1-p \)，那么训练时更常见的写法可以记成：

$$ m_i \sim \mathrm{Bernoulli}(1-p) $$

并且常常还会做缩放：

$$ \tilde{x}_i=\frac{m_i}{1-p}x_i $$

这样做的目的是让训练和推理阶段的期望尺度更一致。

Dropout 的经验性设置

Dropout 最值得问的问题，通常是：

• 放在哪些位置？
• 比例设多少？
• 在今天的小模型和大模型里，这个设置是不是一致？

经验上：

• 早期 Transformer 往往更依赖 dropout
• 今天很多大模型，尤其数据量非常大时，dropout 可能会设得很小，甚至某些位置干脆不用

所以 dropout 并不是“越多越安全”，而是要和模型规模、数据规模、训练稳定性一起考虑。

需要注意的是,Dropout可能会在很多地方被使用,它虽然简单,但是出现的位置却很可能对训练稳定性有很大的影响,不得不进行注意,下面是Dropout可能会出现的几个位置

• Embedding Dropout（输入扰动）
• Attention Dropout（注意力权重扰动）
• Residual Dropout（残差连接前的扰动）
• FFN Dropout（前馈网络内部）
• Output Dropout（输出层前）

MiniMind 里实际上用到了下面这几种：

• Embedding Dropout
• Attention Dropout
• Residual Dropout
• FFN Dropout

如果严格对照前面那张列表，它没有再单独定义一个额外的 Output Dropout；输出侧主要是直接接 lm_head 做词表投影。

对应实现都在 model_minimind.py 里：

• MiniMind_Dense.dropout：embedding 后的 dropout，见 model_minimind.py
• Attention.attn_dropout：attention 权重上的 dropout，见 model_minimind.py
• Attention.resid_dropout：o_proj 之后、回到残差支路之前的 dropout，见 model_minimind.py
• FeedForward.dropout：FFN 内部、down_proj 输出后的 dropout，见 model_minimind.py

class FeedForward(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        ...
        self.dropout = nn.Dropout(config.dropout)
        ...

    def forward(self, x: torch.Tensor):
        ...
        return self.dropout(
            self.down_proj(middle)
        )  # [...,hidden_dim] -> [..., intermediate_dim]


class Attention(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        ...
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        ...

    def forward(...):
        ...
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        scores = self.attn_dropout(scores)
        ...
        output = self.resid_dropout(
            self.o_proj(output)
        )  # -> [batch_size, seq_len_q ,hidden_dim]


class MiniMind_Dense(torch.nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        ...
        self.dropout = nn.Dropout(config.dropout)
        ...

    def forward(...):
        ...
        hidden_states = self.dropout(self.embed_tokens(input_ids))

Dropout 在训练和推理时分别怎么表现?

这是一个非常值得单独说清楚的问题。

在训练阶段：

• Dropout 会随机把一部分神经元输出置零
• 同时对保留下来的部分做缩放
• 这样做的目的，是减少特征之间过度共适应，增强泛化

在推理阶段：

• Dropout 会被关闭
• 不再随机丢弃任何神经元
• 模型会直接使用完整的前向路径

也就是说，Dropout 在推理时通常不会继续“随机掉点”。
这也是为什么我们会在代码里看到：

• model.train() 时 Dropout 生效
• model.eval() 时 Dropout 关闭

Dropout 对推理有什么影响?

从直接机制上说：

• 推理阶段 Dropout 本身通常是关闭的
• 所以它不会也不应该在推理时继续引入随机性

这一小节最值得记住什么?

如果只保留最关键的结论，我会更倾向于这样总结：

1. 激活函数不是一个无关紧要的小选择，它直接影响 FFN 的非线性表达方式。
2. Norm 不是“锦上添花”，而几乎是深层 Transformer 稳定训练的基础部件；而且Transformer的Norm一般都是作用在 [B, n, d] 张量的最后一个feature维度上面.
3. Dropout 也不是固定模板，它和模型规模、数据规模、训练目标都有关系；并且它通常只在训练时生效，在推理时会被关闭。

所以这一节虽然在讲“细节”，但这些细节本身就是现代 LLM 能不能稳定训练起来的重要组成部分。

Q4: FFN / MLP 层在做什么?

Attention 负责不同 token 之间的信息交互，但如果只有 attention，模型的非线性表达能力还是不够。

Transformer block 里通常还会有一个前馈网络，也就是 FFN 或 MLP。最基础的形式可以写成：

$$ \mathrm{FFN}(x) = W_2\phi(W_1 x + b_1) + b_2 $$

其中：

• \( x \in \mathbb{R}^{d_{\text{model}}} \)
• \( W_1 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}} \)
• \( W_2 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}} \)
• \( d_{\text{ff}} \) 表示 FFN 中间层维度

可以把它粗略理解成：

• Attention 负责“跨 token 聚合信息”
• FFN 负责“在每个 token 自己的特征维度上做更强的非线性变换”

经验问题 1: FFN 中间层为什么通常更大?

因为 FFN 的核心作用，就是先把表示映射到一个更高维空间里做变换，再投回原维度。
所以很多实现里都会取：

$$ d_{\text{ff}} \approx 4 d_{\text{model}} $$

标准 Transformer 里，\( d_{\text{ff}} \approx 4d_{\text{model}} \) 是最经典的一条经验规则。到了更现代的 gated MLP 设计里，这个具体比例会变化，但核心思路没有变：中间层通常都会明显大于 hidden size。

直接看几组真实配置会更具体：

这里最值得注意的不是“所有模型都刚好是 4 倍”，而是：

• 标准 Transformer 教科书式的 FFN 比例确实常写成 4 倍
• 到了 Llama、Qwen 这类更现代的结构里，具体数字会变
• 但中间层显著大于 hidden size 这件事基本没变

经验问题 2: 为什么第二个线性变换后通常不再接激活函数?

最基本的 FFN 写成：

$$ \mathrm{FFN}(x) = W_2\phi(W_1x+b_1)+b_2 $$

从这个形式可以看到：

• 第一层线性变换 \( W_1 \) 先把表示投影到更高维空间
• 中间的激活函数 \( \phi \) 提供非线性
• 第二层线性变换 \( W_2 \) 再把表示投影回 hidden size

也就是说，真正负责“打破纯线性结构”的关键步骤，已经发生在中间那一个激活函数上了。
第二个线性层更像是在做“重新压回原维度”和“重新组织特征”的工作。

这是 LLM 里一个很有意思的设计。它的思路更像是：先把表示投影到一个更适合做非线性变换的空间里，做完之后再投影回来，而不是单纯从“多堆几层 FNN 就能拟合任意函数”这种角度去理解。这个结构之所以一直被保留下来，核心还是因为它在长期实践里被证明是有效而且稳定的。

如果在第二个线性层后面再立刻接一个激活函数，通常会带来几个问题：

1. 会改变残差连接希望接收的表示分布
   Transformer block 里 FFN 的输出一般要和残差分支相加，所以最后输出如果再经过一次非线性，往往会让这一层的接口变得没那么干净。更准确地说，到了这一步，我们通常更希望拿到的是一个已经组织好的 feature 表示，而不是再额外套一层激活去继续改它的分布。

2. 非线性已经够用了
   两层线性层中间插一个激活函数，本身就已经让 FFN 具备了足够的非线性表达能力，通常不需要在最后再补一次。

3. 工程上最稳定的主流结构就是这种形式
   现代 Transformer / LLM 的主流 FFN 结构，基本都是“线性 -> 激活 / 门控 -> 线性”，而不是“线性 -> 激活 -> 线性 -> 激活”。

所以这里更准确的说法不是“绝对不能再接激活”，而是主流 LLM 基本都不这么做，因为没有明显收益，而且会让和残差分支的对接变得更别扭。

经验问题 3: FNN 一般做几层?

如果说最经典的 Transformer FFN，它本质上就是一个 两层 MLP：

$$ x \rightarrow W_1x \rightarrow \phi(\cdot) \rightarrow W_2(\cdot) $$

也就是说，最常见的答案其实非常简单：

• 标准 FFN 一般就是两层线性变换

这也是为什么大家平时说 Transformer 里的 FFN，默认通常就是在说这种“两层结构”。更复杂的变体当然也有，比如：

• gated MLP（例如 GLU、SwiGLU）
• 更复杂的 MoE 前馈层

但即使这些结构看起来更复杂，它们大多数也还是围绕“先升维做变换，再回到 hidden size”这个核心思路在发展。

如果只从实践经验出发，可以把结论压缩成：

• 对于标准 Transformer / GPT 类模型，FFN 默认就是两层
• 真正更值得花时间调的，通常不是“FFN 要不要做三层还是四层”，而是
  • 中间层宽度多大
  • 用什么激活函数
  • 是否使用 gated FFN

也就是说，在现代 LLM 里，FFN 的“深度”通常不是首要可调项，FFN 的“宽度”和“具体结构形式”反而更关键。

MiniMind 这里的 FFN 不是最朴素的两层 MLP，而是更接近 LLaMA 系列的 gated FFN：gate_proj 和 up_proj 做逐元素乘法，再经过 down_proj 投回 hidden size。激活函数和 Dropout 也都在这里落地了。对应实现来自 model_minimind.py。

class FeedForward(nn.Module):
    """
    GLU Gate Linear Unit的变体
    From LLaMA 系列
    LLaMA2 首次引入这种结构作为默认 FFN
    Meta 的论文中称之为 Gated Linear Units with SiLU activation

    更强的非线性建模能力:门控乘法能动态调节信息流
    更好的训练稳定性:SiLU 激活 + 无 bias + 64 对齐
    更高的参数利用率:相比单路径 FFN，双路径乘法更充分利用中间维度
    """

    def __init__(self, config: MiniMindConfig):
        super().__init__()

        if config.intermediate_size is None:
            intermediate_size = int(config.hidden_size * 8 / 3)
            # 64 padding!
            config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)

        self.gate_proj = nn.Linear(
            config.hidden_size, config.intermediate_size, bias=False
        )
        self.up_proj = nn.Linear(
            config.hidden_size, config.intermediate_size, bias=False
        )
        self.down_proj = nn.Linear(
            config.intermediate_size, config.hidden_size, bias=False
        )
        self.dropout = nn.Dropout(config.dropout)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x: torch.Tensor):
        """
        x: [batch_size,seq_length,hidden_dim]
            hidden_states already applied  post_layernorm
        Retuen: [batch_size,seq_length,hidden_dim]
        """
        middle = self.up_proj(x) * self.act_fn(
            self.gate_proj(x)
        )  # [...,hidden_dim] -> [..., intermediate_dim]

        return self.dropout(
            self.down_proj(middle)
        )  # [...,hidden_dim] -> [..., intermediate_dim]

Q5: Attention + FNN = Transformer层模型的核心!

到这一步，其实整个模型的架构已经基本清楚了。此时再回头看 Transformer 的结构，就会发现它的主干并不复杂：Attention + FFN 就构成了一个基本的 Transformer block，然后这些 block 一层一层堆起来，像是千层饼一样形成整个模型的 Transformer 主干。

关于 Attention 的更完整直觉和历史背景，可以参考 理解 Attention 机制。
在这一节里，我们只抓住模型定义里最关键的部分：每个 Transformer block 的输入和输出 shape 都保持一致，模型始终在同一个 hidden space 里不断做变换。

也就是说，随着层数增加，真正变化的不是 shape，而是表示本身：

• token 的编码会不断更新
• 上下文信息会不断混合进去
• 每个位置的 hidden state 会越来越“知道”整段序列里发生了什么

Attention 的核心公式

先只看单头 attention。

给定某一层的输入表示：

$$ X \in \mathbb{R}^{n \times d_{\text{model}}} $$

通常先通过三个线性变换得到：

$$ Q = XW_Q,\qquad K = XW_K,\qquad V = XW_V $$

其中：

• \( W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k} \)
• \( Q, K \in \mathbb{R}^{n \times d_k} \)
• \( V \in \mathbb{R}^{n \times d_v} \)

最经典的 scaled dot-product attention 写成：

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

这个公式的含义可以拆成两步：

1. 用 \( QK^\top \) 计算每个位置和其它位置的相关性(Attention Score)，再通过 softmax 转成权重(概率)
2. 再按这个权重对 \( V \) 做加权求和(期望)

多头注意力

单头 attention 当然可以工作，但通常不够。
现代 Transformer 更常见的是多头注意力，也就是把 hidden size 切成多个 head，并行地在不同子空间里做 attention。

如果头数记作 \( h \)，每个 head 的维度记作 \( d_{\text{head}} \)，那么通常有：

$$ d_{\text{model}} = h \cdot d_{\text{head}} $$

这时我们不再只得到一组 \( Q,K,V \)，而是得到每个 head 各自的一组表示。
如果把它们写成张量形式，再加上 batch 维，那么在真实代码里更常见的形式是：

$$ Q_i, K_i, V_i \in \mathbb{R}^{B \times h \times n \times d_{\text{head}}} $$

其中：

• \( d_{\text{head}} \) 表示每个 head 的维度

对每个 head，attention 的计算形式其实和单头时完全一样：

$$ \mathrm{head}_j=\mathrm{Attention}(Q_j,K_j,V_j) $$

其中：

• \( j = 1,2,\dots,h \)
• \( \mathrm{head}j \in \mathbb{R}^{n \times d{\text{head}}} \)

把所有 head 的结果拼接起来：

$$ \mathrm{Concat}(\mathrm{head}_1,\dots,\mathrm{head}h)\in \mathbb{R}^{n \times (h\cdot d{\text{head}})} $$

由于：

$$ h\cdot d_{\text{head}} = d_{\text{model}} $$

所以拼接之后的结果又会回到 hidden size 对应的维度。
最后再经过一个输出投影矩阵：

$$ W_O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}} $$

得到最终输出：

$$ \mathrm{MHA}(X)=\mathrm{Concat}(\mathrm{head}_1,\dots,\mathrm{head}_h)W_O $$

多头注意力虽然内部切成了很多 head，但最后输出还是会被重新投影回原来的 hidden size，所以输入输出 shape 仍然保持一致。

经验问题 1: head_dim 一般怎么选?

head_dim 往往不会单独拍脑袋定，而是直接由 hidden size 和头数一起确定：

$$ d_{\text{head}} = \frac{d_{\text{model}}}{n_{\text{heads}}} $$

直接看配置会更直观：

从这些真实模型里可以直接看到一个很明显的现象：

• GPT-2 这一系的 head dim 基本固定在 64
• Llama、Qwen 这一类更现代的模型里，head dim 很常见的是 128

所以实践里通常不是“先定一个 head_dim，再凑别的参数”，而是先定 hidden size 和头数，再让它们自然相除。不能整除虽然理论上可以用投影矩阵去弥补,但是实际中鲜有采用.

head_num 一般怎么选?

这个问题和 head_dim 是绑在一起的。更实际的顺序通常是：

1. 先定 hidden size
2. 再选一个能整除 hidden size 的 head 数
3. 看最后得到的 head dim 是否落在主流实现里常见的范围内

从上面的例子看，真实模型里最常见的结果其实很集中，通常就是 64 或 128 这两档。

经验问题 2: 为什么多头结果通常是拼接，而不是相加?

直觉上可以这样理解：

• 每个 head 都在不同子空间里学习不同关系
• 拼接能保留各个 head 的独立信息
• 然后再通过一个输出投影统一融合

如果一开始就直接相加，会发生什么？

1. 各个 head 的信息会过早混在一起
   这样模型还没来得及在输出投影层里统一组织这些信息，就已经把它们揉平了。

2. 很难保留“每个 head 各学各的”这件事
   多头注意力的一个核心好处，就是不同 head 可以在不同子空间里建模不同类型的关系。拼接能把这些差异先保留下来，相加则更容易把它们平均掉。

3. 输出层的表达空间也会变小
   拼接之后再接 \( W_O \)，相当于给模型一个更大的融合空间；如果先相加，后面的输出投影能操作的信息就少很多。

所以从结构上看：

• 拼接更像是“先分别学，再统一融合”
• 相加更像是“过早合并”

而现代 Transformer 明显更偏向前者。

如果只给一个很直观的总结，我会更倾向于这么记：

• 多头注意力的意义，不只是“多算几遍 attention”
• 更关键的是让不同 head 的结果先独立存在
• 然后再通过输出投影把这些结果重新组织起来

这也是为什么主流实现几乎都会选择 concat，而不是直接 sum。

经验问题 3: Attention 复杂度为什么总被强调?

因为它最核心的计算通常来自：

$$ QK^\top $$

假设：

• \( n \) 表示序列长度
• \( d \) 表示 head 维度

那么复杂度通常可以粗略写成：

$$ O(n^2 d) $$

这也是为什么：

• 序列长度一上去，attention 代价会迅速上升
• 后面才会有 FlashAttention、长上下文优化等一系列工作

这也是为什么 attention 的复杂度总会被单独强调。
它不是一个只存在于公式里的理论问题，而是会立刻体现在：

• 训练速度
• 显存占用
• 长上下文推理代价

上。

如果再结合 attention score 的 shape 去看，这件事会更直观。因为：

$$ QK^\top \in \mathbb{R}^{n \times n} $$

所以只要序列一长，这个中间结果本身就会迅速变大。

从工程角度看，大家后来不断去做各种 attention 相关优化，基本都是在围绕这个瓶颈展开。
大方向上看，常见优化通常是在解决下面几类问题：

1. 怎么减少真实计算和访存开销
   例如 FlashAttention，核心不是改数学定义，而是改高效实现方式。

2. 怎么降低长上下文带来的代价
   也就是为什么长上下文优化会变成一个持续存在的主题。

3. 怎么让推理阶段不要重复算
   例如 KV cache，本质上也是在避免把已经算过的 K/V 重复计算。

所以如果只保留一个最重要的结论，我会更倾向于这么写：

• attention 很强，是因为它能让任意位置直接和其它位置交互
• attention 也很贵，是因为这种“全连接式交互”天然会带来 \( n^2 \) 级别的代价

这也是为什么后面的 RoPE、FlashAttention、GQA、KV cache、长上下文优化这些主题，虽然看起来各不相同，但很多都可以追溯到 attention 的这一个核心瓶颈。

如果只看这节正文真正对应的“标准 attention 主线”，MiniMind 里最值得看的有两段代码。

第一段是 MiniMindBlock：它把 input_layernorm -> self_attn -> residual -> post_attention_layernorm -> mlp 这条主路径写得很清楚。对应实现来自 model_minimind.py。

class MiniMindBlock(nn.Module):
    def __init__(self, layer_id: int, config: MiniMindConfig):
        super().__init__()
        self.num_attention_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.self_attn = Attention(config)

        self.layer_id = layer_id
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        assert not config.use_moe, "Moe not implemented "
        # self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
        self.mlp = FeedForward(config)

    def forward(
        self,
        hidden_states,
        position_embeddings,
        past_key_value=None,
        use_cache=False,
        attention_mask=None,
    ):
        """
        hidden_states: torch.Tensor [batch_size, seq_length, hidden_dim]
            input embeddings.
        position_embeddings: (freqs_cos,freqs_sin) [seq_length,hidden_dim]
        """
        residual = hidden_states
        hidden_states, present_key_value = self.self_attn(
            self.input_layernorm(hidden_states),
            position_embeddings,
            past_key_value,
            use_cache,
            attention_mask,
        )
        hidden_states += residual
        hidden_states = hidden_states + self.mlp(
            self.post_attention_layernorm(hidden_states)
        )
        return hidden_states, present_key_value

第二段是 Attention.forward 里和本节主线最直接对应的部分。这里只保留了 Q/K/V 投影、RoPE 注入、softmax 权重和输出重排这些和正文公式一一对应的部分；MiniMind 里实际还包含 KV cache、repeat_kv 这些更偏后续专题的实现，这里先不展开。对应实现来自 model_minimind.py。

bsz, seq_len, _ = x.shape
# [...,# q_heads * head_dim] [...,# kv_heads * head_dim] [...,# kv_heads * head_dim]
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
# [seq_length,head_dim],[seq_length,head_dim]
cos, sin = position_embeddings
# No RoPE on V
xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])

# xq xk xv
# [bsc, seq_length, #q_heads , head_dim]

xq, xk, xv = (
    xq.transpose(1, 2),
    xk.transpose(1, 2),
    xv.transpose(1, 2),
)
# xq xk xv 标准的Attention输入
# [bsc, #q_heads, seq_length , head_dim]

# 手写Attention计算实现:
# Q @ K^T / sqrt(d)
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(
    self.head_dim
)  # [bsc, #q_heads, seq_length_q , seq_length_kv]

scores = F.softmax(scores.float(), dim=-1).type_as(
    xq
)  # [batch_size, num_heads, seq_len_q, seq_len_k] 数值变成softmax 权重
scores = self.attn_dropout(scores)
#  scores: [bsc, #num_heads, seq_len_q , seq_len_k]
#  xv    : [bsc, #num_heads, seq_len_k, head_dim]
output = scores @ xv  # -> [bsc, #num_heads, seq_len_q, head_dim]

# Reshape for output
output = output.transpose(
    1, 2
)  # [batch_size, num_heads, seq_len_q, head_dim] -> [batch_size, seq_len_q ,num_heads,head_dim]
output = output.reshape(
    bsz, seq_len, -1
)  # ->  [batch_size, seq_len_q ,num_heads * head_dim]
output = self.resid_dropout(
    self.o_proj(output)
)  # -> [batch_size, seq_len_q ,hidden_dim]

Q6: 更多后期改进和优化?

这一节到这里，基础的模型结构就已经基本讲完了。在主干结构之外，后面还有很多继续优化的方向，它们的目标通常也比较明确：

• 更省显存
• 更快训练
• 更快推理
• 支持更长上下文
• 在不明显损伤表达能力的前提下，把这些优化做进去

这些后续的改进和优化，大多不是把主干结构整个改写掉，而是在已有结构上，围绕某一个明确目标继续往前做。也正因为如此，我会把它们先从主线里拿出来，放到后面的附属小节里单独讲。

这样处理，主要是为了让正文的结构更清楚一点。主线部分先回答“一个基础的预训练模型到底由哪些部分组成”，把最核心的结构梳理清楚；而这些后续优化，更像是一组并列展开的专题。它们之间当然也有关联，但整体上不是那种必须严格按前后顺序读下来的关系。

所以我觉得，这部分内容更适合按问题去读，而不是按顺序去读。比较自然的方式是先想清楚：自己现在更关心的到底是哪一类问题，比如是显存、推理速度，还是长上下文能力；然后再去挑对应的小节来看。

后面我也准备把更多相关的方法逐步补充到这一组专题里。比如：

• 有的在改位置信息进入 attention 的方式，比如 RoPE: 位置编码如何进入 Attention?
• 有的在改 attention 的底层实现方式，比如 FlashAttention: attention 为什么还能更快?
• 有的在改多头结构本身，比如 GQA: 为什么 Query head 和 KV head 可以不一样?
• 有的在改推理阶段的缓存方式，比如 KV Cache: 自回归推理为什么能避免重复计算?
• 有的在改长上下文扩展能力，比如 Long Context: 长上下文能力通常在改什么?

本小节总结.

这一节主要回答的是一个很基础的问题：如果只从模型结构出发看，一个预训练语言模型到底由哪些部分组成。

正文里，我是沿着一条比较顺的主线把这件事梳理下来的：先看模型整体的输入输出，再看输入侧的 Embedding 和位置表示是怎么进入模型的；然后再往里走，到 Transformer block 里最核心的 Attention 和 FFN；同时也顺手把激活函数、Norm、Dropout 这些看起来像细节、但实际上很影响训练和实现的部分一起交代清楚。

如果把这一节压缩成一句话，那么我觉得可以这样记：一个基础的 LLM，本质上就是把输入的 token id 先变成连续表示，再在同一个 hidden space 里经过多层 Transformer block 不断做上下文化变换，最后通过 LM head 投影回词表空间。

后面的小节，则是在这个主干结构的基础上，继续展开各个模块的细节，以及围绕显存、速度、长上下文等目标继续做的优化。也就是说，后续内容大多不是在改“LLM 的主干到底是什么”，而是在这个已经相对稳定的主干上继续做结构和工程上的改进。