2. OLMo 2

非营利机构艾伦人工智能研究所（Allen Institute for AI）推出的 OLMo 系列模型，因其训练数据与代码的高度透明，以及相对详尽的技术报告而备受关注。

尽管 OLMo 系列模型或许不会登顶各类基准测试排行榜，但它们的实现规范干净；更重要的是，得益于极高的透明度，它们是大语言模型研发的绝佳参考范本。

OLMo 系列的知名度虽主要来自透明度，但性能表现也并不逊色。实际上，在今年 1 月发布时（早于 Llama 4、Gemma 3 与 Qwen 3），OLMo 2 系列模型就处于算力 – 性能的帕累托最优前沿，如下图 7 所示。

图 7：不同大语言模型的建模基准性能（数值越高越好）与预训练算力成本（FLOPs，数值越低越好）对比。本图标注版摘自 OLMo 2 论文：https://arxiv.org/abs/2501.00656

正如前文所述，为控制文章篇幅，我将只聚焦大语言模型的架构细节（不涉及训练与数据部分）。那么 OLMo 2 有哪些值得关注的架构设计选择？核心要点集中在归一化设计上：RMSNorm 层的放置位置，以及新增的 QK 归一化（QK-norm），我会在下文展开说明。

还有一点值得一提：OLMo 2 仍然采用传统的多头注意力（MHA），而非 MLA 或 GQA。

2.1 归一化层的位置

整体来看，OLMo 2 与同时期的其他大语言模型类似，大体沿用了初代 GPT 的架构，但也存在几处值得关注的差异。我们先从归一化层说起。

与 Llama、Gemma 以及绝大多数大语言模型一样，OLMo 2 将层归一化（LayerNorm）替换为均方根层归一化（RMSNorm）。

不过 RMSNorm 已经是成熟技术（本质是 LayerNorm 的简化版本，可训练参数更少），这里就不展开对比二者的差异了。（感兴趣的读者可以在我的《GPT-2 转 Llama 架构转换指南》中找到 RMSNorm 的代码实现。）

但 RMSNorm 层的放置位置很有讨论价值。《Attention Is All You Need》提出的原生 Transformer 中，Transformer 块内的两个归一化层分别位于注意力模块与前馈网络模块之后。

这种设计也被称为后归一化（Post-LN / Post-Norm）。

而 GPT 以及后续绝大多数大语言模型，都将归一化层放在注意力与前馈网络模块之前，也就是前归一化（Pre-LN / Pre-Norm）。后归一化与前归一化的对比见下图。

图 8：后归一化、前归一化，以及 OLMo 2 采用的后归一化变体对比

2020 年，Xiong 等人的研究表明，前归一化能让模型初始化时的梯度表现更稳定。研究人员还提到，前归一化甚至无需精心设计学习率预热也能正常训练 —— 而学习率预热对后归一化来说是至关重要的训练手段。

之所以提到这一点，是因为 OLMo 2 采用了一种后归一化方案（不过用的是 RMSNorm 而非 LayerNorm，因此统称其为后归一化）。

如上图所示，在 OLMo 2 中，归一化层并非放在注意力层与前馈层之前，而是置于二者之后。但需要注意的是，与原生 Transformer 架构不同，这里的归一化层仍然位于残差层（跳跃连接）内部。

那么团队为什么要调整归一化层的位置？原因是这种改动有助于提升训练稳定性，如下图所示。

图 9：前归一化（GPT-2、Llama 3 等模型采用）与 OLMo 2 后归一化变体的训练稳定性对比。本图标注版摘自 OLMo 2 论文：https://arxiv.org/abs/2501.00656

遗憾的是，图中的结果是归一化层重排与 QK 归一化共同作用的效果，二者是独立的技术点，因此很难单独衡量归一化层位置调整的贡献占比。

2.2 QK 归一化（QK-Norm）

上一节已经提到了 QK 归一化，而我们后续要讨论的其他大语言模型（比如 Gemma 2 与 Gemma 3）也采用了这项技术，这里就简要介绍一下它的原理。

QK 归一化本质上就是额外增加的一层 RMSNorm。它被置于多头注意力（MHA）模块内部，在施加旋转位置编码（RoPE）之前，对查询向量（q）和键向量（k）做归一化处理。为便于理解，下方节选自我从零实现 Qwen3 时编写的分组查询注意力（GQA）层代码（GQA 中的 QK 归一化用法与 OLMo 的 MHA 中一致）：

class GroupedQueryAttention(nn.Module):
    def __init__(
        self, d_in, num_heads, num_kv_groups,
        head_dim=None, qk_norm=False, dtype=None
    ):
        # ...

        if qk_norm:
            self.q_norm = RMSNorm(head_dim, eps=1e-6)
            self.k_norm = RMSNorm(head_dim, eps=1e-6)
        else:
            self.q_norm = self.k_norm = None

    def forward(self, x, mask, cos, sin):
        b, num_tokens, _ = x.shape

        # Apply projections
        queries = self.W_query(x) 
        keys = self.W_key(x)
        values = self.W_value(x) 

        # ...

        # Optional normalization
        if self.q_norm:
            queries = self.q_norm(queries)
        if self.k_norm:
            keys = self.k_norm(keys)

        # Apply RoPE
        queries = apply_rope(queries, cos, sin)
        keys = apply_rope(keys, cos, sin)

        # Expand K and V to match number of heads
        keys = keys.repeat_interleave(self.group_size, dim=1)
        values = values.repeat_interleave(self.group_size, dim=1)

        # Attention
        attn_scores = queries @ keys.transpose(2, 3)
        # ...

如前文所述，QK 归一化与后归一化搭配使用，能够稳定训练过程。需要注意的是，QK 归一化并非 OLMo 2 首创，最早可追溯到 2023 年的《Scaling Vision Transformers》论文。

2.3 OLMo 2 小结

简而言之，OLMo 2 最具代表性的架构设计决策主要集中在 RMSNorm 的排布上：一是将 RMSNorm 置于注意力与前馈模块之后（而非之前），属于后归一化的一种变体；二是在注意力机制内部，为查询向量与键向量额外增加了 RMSNorm（即 QK 归一化）。二者共同作用，起到了稳定训练损失的效果。

下图进一步对比了 OLMo 2 与 Llama 3 的架构；可以看出，除了 OLMo 2 仍采用传统 MHA 而非 GQA 之外，二者的整体架构高度相似。（不过 OLMo 2 团队在 3 个月后推出的 320 亿参数版本，已改用 GQA。）

图 10：Llama 3 与 OLMo 2 的架构对比

【转载】大语言模型架构全面对比02：OLMo 2

2. OLMo 2

2.1 归一化层的位置

2.2 QK 归一化（QK-Norm）

2.3 OLMo 2 小结

Leave a Reply Cancel reply