LLM训练实战手册
1. 前言
表面上看,已发表的研究论文似乎把一切都说得轻描淡写:无非是战略性的架构选择、精心策划的数据集,以及充足的算力。结果光鲜亮丽,消融实验(ablations)结构清晰,事后看来每一步决策都显得理所当然。
但是,这些报告往往只展示了成功的“果”,并进行了一番“玫瑰色的回顾”(Rosy Retrospection)。它们并没有记录下凌晨两点的 dataloader 调试、突然爆发的 Loss 尖刺,或者那个默默破坏你训练的细微张量并行(Tensor Parallelism)Bug(后面我们还会专门提到!)。现实更加混乱、迭代性更强,充满了那些最终没有写进论文的“纠结瞬间”。
这一次,我们将带领大家深入幕后,直击 SmolLM3 的训练过程——这是一个在 11 万亿(11T)Token 上训练出来的 30 亿(3B)参数多语言推理模型。
这绝非一篇普通的博客文章,而是对一系列决策、发现和死胡同的“蜘蛛网”式梳理,它将为你揭示构建世界级语言模型所需要的核心洞察。
同时,本文也是我们模型训练长文系列的“收官之作”(或“压轴大戏”):我们之前已经深入探讨了大规模数据集构建(FineWeb)、编排数千块 GPU 协同工作(Ultra Scale Playbook),以及在过程的每一步选择最佳评估方法(Evaluation Guidebook)。现在,我们将所有这些元素整合起来,共同打造一个强大的 AI 模型。
我们将全程陪伴大家,不仅仅是分享最终成功的“秘方”,更包括那些塑造了每一个决策的失败、基础设施故障以及调试过程。
这个故事读起来就像一出扣人心弦的“戏剧”:
- 你会看到那些有前景的小规模消融实验为什么在大规模上却可能 “水土不服”;
- 为什么我们在训练了 1 万亿(1T)Token 之后不得不重新启动;
- 如何有效地在保持强大的英语性能的同时,平衡多语言、数学和代码这几个相互竞争的目标;
- 以及最终,我们是如何后训练(post-train)出一个混合推理模型的。
我们努力将这趟冒险之旅组织成一个连贯的故事,而非仅仅是冰冷的步骤列表。请将它视为一本实战指南,献给所有想从“我们有优秀的数据集和 GPU”迈向“我们构建了一个真正强大的模型”的实践者。
我们希望这种毫无保留的开放分享,能够弥合研究与生产之间的鸿沟,让你的下一次训练跑起来少一点混乱,多一点笃定。
1.1 如何阅读这篇文章?(你可能不需要“一次性看完”)
老实说,这篇文章篇幅非常、非常长,想要一口气从头读到尾,在现实中可能不太可行。
好消息是,我们已经将文章结构化为几个独立且清晰的板块,你可以根据自己的兴趣和需求选择性地跳过或单独阅读。
以下是文章的结构指南:
-
🧭 训练指南针(Training compass):
- 核心内容: 这是一个高层次的讨论,关于你 是否应该自行预训练(Pretrain)一个模型。
- 适用人群: 在你“烧光”所有风投资金之前,我们为你列出了几个必须问自己的基本问题,并系统地指导你完成决策过程。
- 温馨提示: 这是一个比较宏观的部分。如果你是 “技术党”,想要直接看硬核技术内容,请快速略过此部分。
-
🚀 预训练(Pretraining):
- 核心内容: 在这之后的部分,涵盖了构建你自己 预训练运行的可靠“配方”所需了解的一切:如何运行消融实验(Ablations)、选择评估指标、混合数据源、进行架构决策、调整超参数,以及最终 “熬过” 这场训练马拉松。
- 适用人群: 无论是计划 从头开始预训练,还是对 持续预训练(Continued Pretraining,或称 Mid-training 感兴趣的朋友,这个板块都适用。
-
🍒 后训练(Post-training):
- 核心内容: 在这部分,你将学到如何最大限度地利用你的预训练模型所需的所有“技巧”。从 SFT (监督式微调)、DPO (直接偏好优化)、GRPO 等一整套后训练“字母表”开始,一直到模型合并(Model Merging)这种充满“黑暗艺术与炼金术”的操作。
- 价值所在: 关于如何让这些算法真正起效的知识,往往是通过痛苦的教训换来的。我们将在这里分享我们的经验,希望能让你少走一些弯路。
-
🏭 基础设施(Infrastructure):
- 核心比喻: 如果说预训练是蛋糕,后训练是上面的糖霜和樱桃,那么基础设施就是那台工业级的烤箱。没有它,什么都做不了;如果它坏了,你愉快的周日烘焙时间就会变成一场火灾隐患。
- 核心内容: 理解、分析和调试 GPU 集群的知识分散在各种库、文档和论坛中。本节将详细讲解 GPU 布局、CPU/GPU/节点/存储之间的通信模式,以及如何识别和克服瓶颈。
那么,我们从何处开始呢? 很简单,选择你最感兴趣的那个板块,让我们出发吧!
2. 训练指南针:灵魂三问(Why → What → How)
机器学习领域对于 “优化” 似乎有着一种近乎痴迷的关系。我们的注意力总是集中在 Loss 曲线、模型架构和训练吞吐量上;毕竟,机器学习从根本上说就是关于 优化模型的损失函数。然而,在深入这些技术细节之前,有一个更根本的问题却常常被忽略:我们真的有必要训练这个模型吗?
开源 AI 生态系统几乎每天都在发布世界级的模型:Qwen、Gemma、DeepSeek、Kimi、Llama(当然,还有它的继任者们)、Olmo……这份名单每个月都在持续增长。它们不仅仅是研究原型或“玩具”示例,而是 可以直接用于生产的模型,涵盖了令人震惊的广泛用例——从多语言理解到代码生成,再到复杂的推理能力。更重要的是,它们大多带有 宽松的许可协议(Permissive Licenses),并拥有活跃的社区随时准备为你提供帮助。
这也提出了一个令人不安的真相:也许,你根本不需要训练自己的模型。
这听起来像是一个“大模型训练指南”最奇怪的开场白。但事实是,许多失败的训练项目,并非因为超参数设置不当或代码有 Bug,而是因为 有人决定训练一个他们根本不需要的模型。所以,在你投入资金和时间开始训练,并钻研“如何执行”之前,你必须先回答两个核心问题:你为什么要训练这个模型? 你应该训练一个什么样的模型? 如果对这两个问题没有清晰的答案,你很可能会浪费数月的算力投入和工程师时间,最终却造出了一个“全世界早已拥有”的东西,甚至更糟—— 一个根本没人需要的东西。
让我们从 “Why”(为什么) 开始。因为如果不懂得你的目标,你后续的所有决策都将是 毫无章法。
2.1 Why:那个没人愿意回答的问题
让我们直言不讳地谈谈实践中经常发生的事情。
某个人(如果他们足够幸运)获得了 GPU 集群的访问权限——可能是通过研究经费,也可能是利用公司闲置的算力,而他们的思路大致是这样的:“我们有 100 块 H100s,可以用三个月。来训练一个模型吧!”模型大小被随意选择,数据集则从各种可用的资源中拼凑起来。训练开始了。六个月后,在烧光了算力预算和团队士气之后,这个模型却无人问津,因为从没有人问过“为什么”。
这里有一些你绝对不应该训练模型的常见理由:
- “我们恰好有空闲的算力。”(这只是资源,不是目标。)
- “其他人都在做。”(这是同侪压力,不是战略。)
- “AI 是未来。”(这是陈词滥调,不是计划。)
- “我们想要最强大的模型。”(这个目标不够具体,无法指导任何实际决策。)
“我们训练了自己的模型”这种诱惑力是强大的。但在投入大量时间和资源之前,有必要问一句:你到底为什么需要训练这个模型?
下面的流程图提供了一个思路指导,这是你在开始一个大型预训练项目前应该经历的思考过程。从技术角度看,你首先应该搞清楚的是:是否已经有一个现成的模型,通过简单的提示(Prompt)或微调(Fine-tune)就能完成你的工作?
本质上,只有在以下三个常见领域中,定制化的预训练才可能真正有意义:你想进行开创性的研究,你的生产用例有非常特殊的需求,或者你想填补开源模型生态系统中的空白。让我们快速看看每一种情况:
2.1.1 科研(Research):你到底想理解什么?
在大模型(LLM)领域,你可以做的研究课题非常多。
这些 LLM 研究项目的共同点是:你通常都会从一个清晰定义的问题开始。例如:
- 我们能否将基于这种新型优化器的训练扩展到 100 亿(10B)参数以上的模型?(参考自:Muon 在 LLM 训练中的可扩展性)
- 仅使用强化学习,不依赖 SFT,能否产出推理能力?(参考自:DeepSeek-R1:通过强化学习激励 LLM 的推理能力)
- 我们能否仅靠纯合成的“教科书式”数据,训练出优秀的小模型?(参考自:Textbooks Are All You Need)
- 我们能否通过仅训练公开许可(Openly Licensed)的数据来达到具有竞争力的性能?(参考自:The Common Pile v0.1:一个 8TB 的公共领域和公开许可文本数据集)
把假设设定得尽可能具体,并提前思考所需的实验规模,能够大大增加你成功的几率。
2.1.2 产品化(Production):为什么你不能直接用现成的模型?
企业不能直接使用现成的通用模型(off-the-shelf models)来解决其特定用例,主要有三个原因。其中两个是技术性的,另一个则关乎治理。
第一个要自己训练模型的原因是:领域特殊性(Domain Specificity)。 当你的数据或任务涉及 高度专业化的词汇或结构,而现有通用模型无法很好地处理时。例如:
- 一个关于 DNA 的模型,它需要独特的词汇表和处理长距离依赖关系的能力。
- 一个法律或金融模型,要求对领域特定的术语和逻辑有深入的理解。
第二个相关原因是:部署约束(Deployment Constraints)。 当你需要一个模型来适配你的硬件、延迟或隐私要求时。例如,一个需要在 无人机上运行,或在 配备定制硬件(如 FPGA)的本地(on-prem)系统 上运行的大模型。
这里有一个简单的测试方法:花几天时间,基于 Qwen3、Gemma3 或其他当前的 SOTA 模型进行构建。你能通过 Prompting(提示)、工具调用(Tool-use)或后训练(Post-training) 达到你的性能目标吗?如果不能,那么,可能就是时候自己训练一个了。
- 一个小小的提醒: 即使为了满足你的要求,所需的 后训练预算非常庞大,它仍可能比从头开始训练更经济。 毕竟,为你的模型微调 1 万亿(1T)Token,仍然比从头开始训练 10 万亿(10T)Token 更划算。
- (也是从这个时候开始,大模型训练者开始神奇地称之为 “Mid-training”(中途训练),而不是 Post-training 了。)
第三个建立内部语言模型的原因是:安全与治理(Safety and Governance)。 由于你身处一个受严格监管的行业或处理高风险的应用,你需要对训练数据、模型行为和更新周期拥有完全的控制权。你需要确切地知道模型中包含了什么,并能够向监管机构证明这一点。在某些情况下,你可能别无选择,只能自建模型。
以上是公司训练内部模型的主要原因。那么,那些发布开源模型的公司或组织又是怎么考虑的呢?
2.1.3 战略性开源(Strategic Open-Source):你看到了可以填补的空白吗?
经验丰富的 AI 实验室发布新的开源模型,最常见的原因之一是:他们识别出了开源生态系统中的一个特定空白或一个新的 AI 用例。
这种模式通常是这样的:你注意到一个 未被充分探索的领域。也许现在没有强大且具备超长上下文能力的 设备端(on-device)模型;或者现有的多语言模型在 低资源语言上表现很弱;又或者,领域正在转向像 Genie3 那样的 交互式世界模型,但还没有好的开源模型出现。
你有理由相信自己可以做得更好。也许你整理出了 更优质的训练数据,开发出了 更棒的训练“配方”,或者拥有其他机构无法达到的 算力优势来支持“过度训练”(Overtrain)。你的目标是具体的:不是“有史以来最好的模型”,而是“最适合设备端使用的 3B 模型”,或“第一个具备 1M 上下文的小模型”。
这是一个真实且有价值的目标。成功会创造价值:开发者会采用你的模型,它会成为其他人的基础设施,或者为你建立技术信誉。但成功需要经验。在一个竞争激烈的领域中,你需要知道什么才是真正可行的,以及如何可靠地执行。
为了让这个思路更具体,让我们来看看 Hugging Face 是如何思考这个问题的。
2.1.4 Hugging Face 的思考:我们为什么要训练开源模型?
那么,Hugging Face 为什么要训练和发布开源模型呢?答案很简单:我们致力于构建对开源生态系统有用的东西,并填补那些极少人涉足的空白。
这包括数据集、工具,当然也包括训练模型。我们启动的每一个 LLM 训练项目,都是始于 发现一个空白,并坚信我们能做出有意义的贡献。
我们最初的 LLM 项目是在 GPT-3 (Brown et al., 2020) 发布之后启动的。当时,感觉没有人愿意构建一个开放的替代品,我们担心相关的知识最终会被锁定在少数几个工业实验室中。因此,我们发起了 BigScience 研讨会,旨在训练一个开源版本的 GPT-3。最终诞生的模型就是 Bloom,它汇集了数十位贡献者一年的努力,从构建训练堆栈、分词器(tokenizer)到预训练语料库,最终预训练了一个 1750 亿(175B)参数的模型。
Bloom 的继任者是 2022 年的 StarCoder (Li et al., 2023)。当时 OpenAI 为 GitHub Copilot 开发了 Codex (Chen et al., 2021),但它是闭源的。很明显,构建一个开源替代品将为整个生态系统带来巨大的价值。因此,我们与 ServiceNow 合作,在 BigCode 的框架下,构建了 The Stack 数据集,并训练了 StarCoder 15B 来复现 Codex 的能力。StarCoder2 (Lozhkov et al., 2024) 的诞生,则是源于我们认识到可以进行更长时间的训练,并意识到 更小但训练更久的模型可能比一个巨型模型更有价值。我们训练了一个模型家族(3B/7B/15B),使用了数万亿(Trillions)的 Token,远远超过当时任何开放代码模型的训练量。
SmolLM 系列遵循了类似的模式。我们注意到当时 强大的小型模型非常稀缺,而我们刚刚构建了 FineWeb-Edu (Penedo et al., 2024) 这个强大的预训练数据集。SmolLM (135M/360M/1.7B) 是我们的第一个版本。SmolLM2 (Allal et al., 2025) 则专注于更好的数据和更长的训练,在多个方面达到了 SOTA 性能。而 SmolLM3 则扩展到了 30 亿参数,同时加入了 混合推理(hybrid reasoning)、多语言能力和长上下文 等社区在 2025 年高度重视的功能。
这种模式甚至延伸到了预训练之外:我们训练了 Zephyr (Tunstall et al., 2023) 来证明 DPO 可以进行大规模应用;启动了 Open-R1 来复现 DeepSeek R1 的蒸馏管线;并发布了用于 竞技编程(competitive programming) 的 OlympicCoder,在国际信息学奥林匹克竞赛中达到了 SOTA 性能。我们还探索了其他模态,例如用于视觉的 SmolVLM (Marafioti et al., 2025) 和用于机器人技术的 SmolVLA (Shukor et al., 2025)。
希望这一部分已经成功说服你:深入思考你为什么要训练一个模型是极具价值的。
在本文接下来的部分,我们将假设你已经完成了这种“灵魂拷问”,并且有了一个合法的、充分的理由去启动你的训练项目。
好的,这是下一段的翻译和润色:
2.2 What:将目标转化为实际决策
既然你已经明确了 为什么 要训练,那么接下来就该确定 应该训练什么 了。
这里的“做什么”(What)指的是:模型类型(密集型 Dense、MoE 专家混合、混合型 Hybrid,还是全新的类型)、模型规模、架构细节和数据混合配比。
一旦你确定了“为什么”(Why),你就可以由此推导出“做什么”(What)。举例来说:
| 训练目的 (Why) | $\rightarrow$ | 模型规格 (What) |
|---|---|---|
| 追求设备端运行的快速模型 | $\rightarrow$ | 小且高效的模型 |
| 构建强大的多语言模型 | $\rightarrow$ | 大型分词器词汇表(Large Tokenizer Vocab) |
| 需要超长上下文能力 | $\rightarrow$ | 混合型架构(Hybrid Architecture) |
除了受用例驱动的决策外,还有一些选择是旨在优化训练本身的,例如让训练更稳定、更具样本效率,或速度更快。这些决策并非总是非黑即白,但你可以大致将决策过程分为两个阶段:
规划阶段(Planning): 在进行实验之前,你需要将你的用例映射到需要确定的组件上:你的部署环境决定了模型规模的限制。你的时间表决定了你可以承担哪些架构风险。你的目标能力决定了数据集的要求。这个阶段的核心就是:将“Why”中的每一个约束条件,都与“What”中的具体规格紧密连接起来。
验证阶段(Validation):一旦你有了一个起始点和一份潜在修改清单,就要进行系统性的测试。由于测试成本高昂,你应将精力集中在那些能 有意义地提升你的用例性能,或能 显著优化你的训练过程的改动上。这就是 消融实验(Ablations) 发挥作用的地方,我们将在后续的“消融实验”章节详细介绍。
在接下来的章节中,你将了解到定义模型的所有选项,以及如何通过系统性实验来缩小选择范围。但在那之前,我们想分享一些关于如何组织团队和项目的经验——这些经验来自于我们自己训练模型的实践,以及观察那些成功构建优秀 LLM 团队的优秀做法。
好的,这是下一段的翻译和润色:
2.3 How:迭代速度与数据质量
通往罗马的道路当然不止一条,但我们发现,成功的大模型(LLM)训练团队之所以能脱颖而出,其关键因素在于 迭代速度。
训练 LLM 本质上是一种 “边训边学” 的学科,你训练的次数越多,你的团队就会变得越优秀。因此,那些一年只训练一个模型的团队,和那些一个季度就能训练一个模型的团队相比,后者的进步速度会快得多。你可以参考 Qwen 和 DeepSeek 等团队,他们现在已是家喻户晓的名字,凭借的就是长期以来持续快速地发布新模型。
除了迭代速度之外,数据策划(Data Curation)是迄今为止对 LLM 训练最具影响力的方面。人们天然倾向于钻研架构选择来改进模型,但那些在大模型训练中表现优异的团队,却是对高质量数据痴迷程度高于一切的团队。
另一个与迭代速度紧密相关的因素是团队规模:对于主要的预训练任务,你只需要少数几个人,并为他们配备足够的算力来执行即可。例如,如今要预训练一个像 Llama 3 这样的模型,你可能只需要 2 到 3 个人。只有当你开始涉足更多元化的训练和下游任务(如多模态、多语言、后训练等)时,才需要慢慢增加更多的人手,以在每个领域做到精通。
因此,秘诀就是:从一个小型、装备精良的团队开始,每隔两到三个月就构建一个新模型,你将在短时间内攀升至行业的顶端。
好了,接下来的文章将专注于这个团队的日常技术细节!
3. 每个大模型的诞生,都始于一场小小的“消融实验”
在我们开始训练一个大型语言模型(LLM)之前,我们需要做出无数将影响模型性能和训练效率的决策。我们该选择什么样的架构最适合我们的用例?使用哪种优化器和学习率调度?如何混合不同的数据源?
“这些决策是如何做出的?”这是一个被频繁问及的问题。人们有时期望它们是通过深入的思考得出的。虽然战略性思维至关重要——正如我们在前一节中讨论的,它能帮你识别出哪些架构更改值得测试——但仅仅依靠推理是不够的。在 LLM 领域,事物并不总是直觉化的,那些“应该有效”的假设在实践中往往会落空。
举个例子,使用看起来 “质量最高的数据” 并不一定能产出更强的模型。以 arXiv 为例,它是人类科学知识的巨大宝库。直觉上,用如此丰富的 STEM 数据进行训练应该能产出更优秀的模型,对吗?但实际上,它并非总是如此,特别是对于小模型,甚至可能损害性能 (Shao et al., 2024)。为什么会这样? 原因是虽然 arXiv 论文知识丰富,但它们高度专业化,并以一种狭隘的学术风格撰写,这与模型最擅长学习的多样化、通用性文本截然不同。
那么,如果苦思冥想没有帮助,我们如何知道什么才是有效的呢?答案是:像优秀的经验主义者一样,运行大量的实验! 机器学习不是纯粹的数学,它更像是一门实验科学。由于这些实验将指导我们许多关键决策,因此将它们设置好至关重要。我们本质上希望从这些实验中获得两个主要属性:
- 速度(Speed): 它们应该尽可能快地运行,以便我们能够频繁迭代。我们能运行的消融实验越多,能验证的假设就越多。
- 可靠性(Reliability): 它们应该提供强大的判别力(discriminative power)。如果我们关注的指标无法在早期有意义地分辨不同设置之间的优劣,那么我们的消融实验可能提供的洞察就很少(如果结果充满噪音,我们就有追逐噪音的风险!)。
但在设置消融实验之前,我们需要对架构类型和模型规模做出一些基础性的选择。这些由我们的“指南针”指导的决策,会影响我们使用哪种训练框架、如何分配算力预算,以及从哪个基线开始。
对于 SmolLM3,我们选择了 30 亿参数的密集型(dense)Llama 风格架构,因为我们的目标是小型设备端模型。但正如你将在 “设计模型架构” 一章中看到的,MoE 或混合模型可能更适合你的用例,不同的模型规模也伴随着不同的权衡取舍。稍后我们将深入探讨这些选择,并向你展示如何做出这些决定。现在,让我们从最实际的第一步开始:选择你的基线(Baseline)。
3.1 选择你的基线(Baseline)
每一个成功的模型都是建立在一个经过验证的基础之上,然后根据自身需求进行修改的。
- 当 Qwen 训练他们的第一个模型家族时 (Bai et al., 2023),他们以 Llama 的架构为起点。
- 当 Meta 训练 Llama 3 时,他们从 Llama 2 开始。
- Kimi K2 则始于 DeepSeek-V3 的 MoE 架构。
这种“继承”不仅适用于架构,也适用于训练超参数和优化器。
为什么呢? 优秀的架构和训练设置需要多年的迭代,并汇集众多组织的智慧。标准的 Transformer 结构和 Adam 等优化器,都是经过数千次实验才被完善的。人们已经发现了它们的失败模式,调试了不稳定性,并优化了实现。
从一个经过验证的基础开始,意味着你继承了所有这些积累的知识。 而从零开始,则意味着你需要自己重新发现每一个问题。
以下是一个好的架构起点应该具备的条件:
- 符合你的约束条件: 契合你的部署目标和实际用例。
- 经过大规模验证: 在相似或更大的规模上,跑过数万亿(multi-trillion)Token 的训练。
- 文档完善: 有明确的、被证明在开源模型中有效的超参数配置。
- 框架支持良好: 理想情况下,它应该被你考虑使用的训练框架和计划用于推理的推理框架所支持。
下面列出了一份非详尽的清单,展示了 2025 年针对不同架构类型和模型规模的一些强大基线选项:
| 架构类型 | 模型家族 | 常见规模 (Sizes) |
|---|---|---|
| 密集型(Dense) | Llama 3.1 | 8B, 70B |
| 密集型(Dense) | Llama 3.2 | 1B, 3B |
| 密集型(Dense) | Qwen3 | 0.6B, 1.7B, 4B, 14B, 32B |
| 密集型(Dense) | Gemma3 | 12B, 27B |
| 密集型(Dense) | SmolLM2, SmolLM3 | 135M, 360M, 1.7B, 3B |
| MoE (专家混合) | Qwen3 MoE | 30B-A3B, 235B-A122B |
| MoE (专家混合) | GPT-OSS | 21B-A3B, 117B-A5B |
| MoE (专家混合) | Kimi Moonlight | 16B-A3B |
| MoE (专家混合) | Kimi-k2 | 1T-A32B |
| MoE (专家混合) | DeepSeek V3 | 671B-A37B |
| 混合型(Hybrid) | Zamba2 | 1.2B, 2.7B, 7B |
| 混合型(Hybrid) | Falcon-H1 | 0.5B, 1.5B, 3B, 7B, 34B |
| MoE + 混合型 | Qwen3-Next | 80B-A3B |
| MoE + 混合型 | MiniMax-01 | 456B-A46B |
因此,请找到你的架构类型,并选择一个参数数量接近你目标模型的基线。不要过度纠结于此,因为你最初选择的架构并非一成不变。在下一节中,我们将看到如何从一个基线出发,一步步构建出最适合你的最终架构。
3.1.1 修改你的基线:“去风险化”的原则
现在你拥有了一个有效且符合你用例的基线模型。你大可以停在这里,用你的数据混合集(假设数据质量不错)进行训练,很可能会得到一个体面的模型。许多成功的项目正是这样做的。然而,基线模型并非为你的特定约束而优化,它们是为其构建者的用例和部署目标而设计的。因此,你很可能需要进行一些修改,使其更好地契合你的目标。但请注意,每一次架构上的更改都伴随着风险:它可能提升性能、彻底搞砸,或者什么也没做,只是浪费了你的消融实验算力。
让你保持在正轨上的纪律是 “去风险化”(Derisking):除非你已经测试并确认它有帮助,否则绝不更改任何东西。
棘手之处在于,你的基线和训练设置有太多可修改的组件:注意力机制、位置编码、激活函数、优化器、训练超参数、归一化方案、模型布局等等。每一个都代表着一个潜在的实验,而这些组件往往以非线性的方式相互作用。你既没有时间,也没有算力来测试所有组合或探索所有的交互。
正确的做法是:从测试有前景的改动开始,并以当前基线为参照。 当某个改动有效时,就将其集成进来,创建一个新的基线,然后针对这个新基线测试下一个改动。如果你的算力预算允许,你可以独立测试多项改动,并运行 “留一法分析”(leave-one-out analysis)。
千万不要掉入陷阱: 避免对每一个超参数进行详尽的网格搜索(Grid Searches),也避免测试每一个新出现的架构变体。
现在你知道了如何通过战略规划来确定哪些改动是有前景的,接下来就该进入经验验证了。在接下来的部分中,我们将向你展示如何在实践中真正测试这些更改。我们将涵盖如何设置可靠的实验、如何解读结果以及避免常见的陷阱。随后,在后续章节中,我们将通过具体案例,讲解如何测试流行的架构、数据、基础设施和训练决策。
那么,让我们先搭建一个可用于实验的简单消融设置。第一步,我们需要决定选择哪个训练框架。
3.2 挑选训练框架
我们需要做的第一个决策是:选择哪个框架来训练模型,进而也决定了用来运行所有消融实验的框架。这个选择需要平衡以下三个关键考量点:
- 架构兼容性: 框架必须支持我们的目标架构,或能让我们轻松地进行扩展。
- 稳定性和生产就绪: 框架需要稳定、成熟,不会在训练中途神秘地崩溃。
- 高吞吐量: 它应该能提供强大的吞吐量,以便我们快速迭代,最大限度地利用算力预算。
在实践中,这些要求可能相互掣肘,形成权衡取舍。让我们来看看可用的选项。
| 框架 | 特性覆盖 | 实战检验 | 优化程度 | 核心/总代码行数 | 扩展性与调试难度 |
|---|---|---|---|---|---|
| Megatron-LM | ✅ 功能全面 | ✅ Kimi-K2, Nemotron | ✅ 3D 并行先驱 | 93k / 269k | ⚠️ 难度大,不适合新手 |
| DeepSpeed | ✅ 功能全面 | ✅ BLOOM, GLM | ✅ ZeRO & 3D 并行先驱 | 94k / 194k | ⚠️ 难度大,不适合新手 |
| TorchTitan | ⚡ 功能集不断增长 | ⚠️ 较新,但经 PyTorch 团队检验 | ⚡ 针对密集模型优化,MoE 正在改进中 | 7k / 9k | ⚡ 适中:需了解并行知识 |
| Nanotron | 🎯 极简,为 HF 预训练定制 | ✅ 是 (StarCoder, SmolLM) | ✅ 高度优化 (UltraScale Playbook) | 15k / 66k | ⚡ 适中:需了解并行知识 |
上表总结了流行框架之间的关键权衡。(前三个框架的代码行数数据来自 TorchTitan 技术报告 (Liang et al., 2025))。让我们详细讨论每一个框架:
Megatron-LM Nvidia 的 Megatron-LM 已经存在多年,久经沙场。它是 Kimi 的 K2 (Team et al., 2025) 等模型的幕后功臣,提供了可靠的吞吐量和我们想要的大多数生产功能。但这种成熟也带来了复杂性:当你刚接触它时,代码库可能会让人难以理解和修改。
DeepSpeed DeepSpeed 属于类似的类别。它是 ZeRO 优化的先驱,为 BLOOM 和 GLM 等模型提供了动力。和 Megatron-LM 一样,它经过了广泛的实战检验和优化,但面临着同样的复杂性挑战。庞大的代码库(总计 194k 行)在入门时可能令人生畏,尤其是当你需要实现自定义功能或调试意外行为时。
TorchTitan 另一方面,PyTorch 最近推出的 TorchTitan 库则轻量且更易于导航,这得益于其紧凑和模块化的代码库。它具备预训练所需的核心功能,非常适合快速实验。然而,由于它相对较新,实战检验不如前两者充分,并且由于仍在积极开发中,稳定性可能略逊一筹。
Nanotron 我们 Hugging Face 采取了不同的路径,从头构建了自己的框架 Nanotron。这为我们带来了完全的灵活性和对大规模预训练的深入理解——这些洞察后来演变成了《Ultra Scale Playbook》。虽然我们开源了该库并获得了社区的宝贵反馈,但在大多数情况下,我们不得不先自己对功能进行实战检验。该框架现在支持我们训练所需的所有生产功能,但仍在构建如 MoE 支持等区域。
对于我们而言,从头构建当时是合理的,但这需要对团队专业知识和时间进行重大投入,用于调试问题和添加缺失功能。一个强大的替代方案是分叉(fork)一个现有框架,并根据你的需求进行增强。例如,Thinking Machines Lab 就是将他们的内部预训练库作为 TorchTitan 的一个分叉来构建的(来源)。
最终,你的选择取决于团队的专业知识、目标功能,以及你愿意投入多少时间进行开发,而不是直接使用最成熟的生产选项。
如果多个框架都能满足你的需求,那么请在你的特定硬件上比较它们的吞吐量。对于快速实验和竞速运行,更简洁的代码库通常会获胜。
3.3 消融实验设置
既然训练框架已定,我们现在就需要设计我们的消融实验设置。我们需要实验足够快以便快速迭代,但又需要足够大,以确保结果能提供有价值的信号,并能 外推(extrapolate) 到最终的模型。让我们来看看如何设置它。
3.3.1 搭建消融实验框架
消融实验的目标是在小规模上运行实验,并获得可以自信地外推到最终生产运行的结论。
主要有两种方法:
- 方法一(减 Token): 保持目标模型大小不变,但在更少的 Token 上进行训练。例如,在 SmolLM3 的消融实验中,我们用完整的 30 亿参数模型,但在 1000 亿(100B)Token 上进行训练,而不是最终的 11 万亿 Token。
- 方法二(减模型): 如果目标模型太大,我们可以训练一个更小的代理模型来进行消融。例如,Kimi 在开发他们拥有 32 亿激活参数的 1 万亿(1T)参数 Kimi K2 模型时,对所有消融实验都使用完整大小显然过于昂贵,因此他们用一个 30 亿参数 MoE 模型(0.5B 激活参数)运行了部分消融实验 (Team et al., 2025)。
一个关键问题是:这些小规模的发现真的能迁移吗?根据我们的经验,如果某个改动在小规模上损害了性能,你可以放心地将其排除在更大规模的训练之外。但如果某个改动在小规模上有效,你仍然需要确保在合理的 Token 数量上进行了训练,才能高概率地得出这些发现可以外推到更大规模的结论。训练时间越长,消融模型与最终模型越接近,结果就越可靠。
在本文中,我们将使用一个基线版(Vanilla)Transformer 来进行所有消融实验。我们的主要设置为:一个遵循 Llama 3.2 1B 架构的 10 亿参数 Transformer,在 450 亿 Token 上训练。这在一个配备 8 块 H100 的节点上大约需要 1.5 天(使用 nanotron 配置,速度约为每 GPU 每秒 42k Token)。在 SmolLM3 的训练过程中,我们是在一个 3B 模型上用 100B Token 运行这些消融的(配置)。我们会在每个章节的末尾分享这些结果(你会看到结论是吻合的)。
我们的基线 1B 配置以结构化的 YAML 格式捕获了所有重要的训练细节。以下是关键部分:
## 数据集及其混合权重
data_stages:
- data:
dataset:
dataset_folder:
- fineweb-edu # FineWeb 教育数据集
- stack-edu-python # Python 代码数据集
- finemath-3plus # 针对 3 级以上数学数据集
dataset_weights:
- 0.7 # 权重 70%
- 0.2 # 权重 20%
- 0.1 # 权重 10%
## 模型架构,Llama3.2 1B 配置
model:
model_config:
hidden_size: 2048 # 隐藏层大小
num_hidden_layers: 16 # 隐藏层数量
num_attention_heads: 32 # 注意力头数量
num_key_value_heads: 8 # K/V 头数量 (GQA)
intermediate_size: 8192 # 中间层大小
max_position_embeddings: 4096 # 最大位置嵌入
rope_theta: 50000.0 # RoPE 旋转角度
tie_word_embeddings: true # 绑定词嵌入
## 训练超参数,带余弦调度的 AdamW 优化器
optimizer:
clip_grad: 1.0 # 梯度裁剪
learning_rate_scheduler:
learning_rate: 0.0005 # 最大学习率
lr_decay_starting_step: 2000
lr_decay_steps: 18000
lr_decay_style: cosine
lr_warmup_steps: 2000
lr_warmup_style: linear
min_decay_lr: 5.0e-05 # 最小衰减学习率
optimizer_factory:
adam_beta1: 0.9
adam_beta2: 0.95
adam_eps: 1.0e-08
name: adamW
## 并行化,单节点配置
parallelism:
dp: 8 # 数据并行,跨 8 块 GPU
tp: 1 # 1B 规模不需要张量或流水线并行
pp: 1
## 分词器
tokenizer:
tokenizer_max_length: 4096
tokenizer_name_or_path: HuggingFaceTB/SmolLM3-3B
## Batch size, 序列长度和总共 30B Token 的训练量
tokens:
batch_accumulation_per_replica: 16 # 每个副本的批次累积
micro_batch_size: 3 # GBS (全局批次大小)=dp * batch_acc* MBS * sequence=1.5M tokens
sequence_length: 4096
train_steps: 20000 # GBS * 20000 = 30B
...
在我们的消融实验中,我们会根据测试的内容修改不同的部分,同时保持其他一切不变:测试架构选择时修改 model 部分,测试优化器和训练超参数时修改 optimizer 部分,测试数据策略时修改 data_stages 部分。
运行消融实验时,某些架构更改会显著改变参数数量。例如,从 绑定(tied) 词嵌入切换到 不绑定(untied) 会使我们的嵌入参数翻倍,而从 MHA 切换到 GQA 或 MQA 则会大幅减少我们的注意力参数。为了确保公平比较,我们需要跟踪参数数量,并偶尔调整其他超参数(如隐藏层大小或层数)以保持模型大小大致相同。
3.3.2 理解效果:评估至关重要
一旦我们启动了消融实验,我们如何判断哪些改动有效,哪些无效呢?
任何训练模型的人第一直觉可能是查看 Loss 曲线,这确实很重要。你希望看到它平滑下降,没有剧烈的尖刺或不稳定。对于许多架构选择,Loss 与下游性能有很好的相关性,可能就足够了 (Y. Chen et al., 2025)。
然而,仅看 Loss 并不总是可靠的。以数据消融实验为例,你会发现用维基百科训练比用网页训练得到更低的 Loss(预测下一个 Token 更容易),但这并不意味着你会得到一个更有能力的模型。同样,如果我们在不同训练时更改了分词器,Loss 就没有直接可比性,因为文本被分割的方式不同了。某些更改可能还会专门影响推理和数学等特定能力,而这些影响在平均 Loss 中会被稀释。最后但同样重要的一点是,模型即使在预训练 Loss 收敛后,仍可能在下游任务上持续改进 (Liu et al., 2022)。
我们需要更细致的评估,才能看清全貌并理解这些细微的影响。一个自然的方法是使用 下游评估(downstream evaluations) 来测试知识、理解、推理以及对我们重要的任何其他领域。
对于这些消融实验,最好关注那些 能提供良好早期信号 并 避免嘈杂基准(noisy benchmarks) 的任务。在 FineTasks 和 FineWeb2 中,可靠的评估任务由四个关键原则定义:
- 单调性(Monotonicity): 基准分数应随着模型训练时间的延长而持续改善。
- 低噪音(Low noise): 当我们使用相同设置但不同随机种子训练模型时,基准分数不应有剧烈波动。
- 高于随机水平的性能(Above-random performance): 许多能力只有在训练后期才会浮现,因此长期保持随机水平性能的任务对消融实验没有用处。例如,我们稍后将解释的多项选择格式的 MMLU 就是这种情况。
- 排名一致性(Ranking consistency): 如果某种方法在早期阶段优于另一种方法,那么随着训练的继续,这种排序应该保持稳定。
任务的质量还取决于任务表述(task formulation)(我们如何向模型提问)和指标选择(metric choice)(我们如何计算答案得分)。
三种常见的任务表述是:多项选择格式 (MCF)、完形填空表述 (CF) 和自由形式生成 (FG)。
- MCF 要求模型从 Prompt 中明确给出并标有 A/B/C/D 的选项中进行选择(例如 MMLU 的做法)。
- CF 中,我们比较不同选项的似然性(Likelihood),以查看哪个更可能,而无需在 Prompt 中提供它们。
- FG 中,我们查看模型对给定 Prompt 进行贪婪生成(greedy generation)的准确性。FG 需要模型具备大量的潜在知识,对于进行完整的预训练之前的短时消融实验来说,通常难度太大而没有太大用处。
因此,在运行小型消融实验时,我们主要关注 MCF 或 CF。
研究表明,模型在训练早期阶段难以应对 MCF,只有经过大量训练后才能掌握这项技能,这使得 CF 更适合用于获取早期信号 (Du et al., 2025; Gu et al., 2025; J. Li et al., 2025)。因此,我们使用 CF 进行小型消融实验,并在主运行中集成 MCF,因为它在模型达到一定阈值、信噪比足够高时,能提供更好的中期训练信号。快速说明: 为了在 CF 这样的序列似然评估中对模型的答案进行评分,我们将准确率计算为正确答案具有最高对数概率(Log Probability)(并按字符数标准化)的问题百分比。这种标准化可以防止偏向较短的答案。
我们的消融评估套件包括 FineWeb 消融实验中的基准,但排除了 SIQA(我们发现它噪音太大)。我们增加了 GSM8K 和 HumanEval 等数学和代码基准,以及用于长上下文消融的 RULER 长上下文基准。如下表所示,这些任务聚合在一起,以多种格式测试世界知识、推理和常识。为了加快评估速度,以牺牲一些额外噪音为代价,我们只评估每个基准的 1,000 个问题(GSM8K、HumanEval 和 RULER 除外,我们在 3B SmolLM3 消融中使用了完整集合,但在下面的 1B 实验中省略)。我们还对所有多项选择基准使用了完形填空表述(CF)的评估方式,如上所述。请注意,对于多语言消融和实际训练,我们增加了更多基准来测试多语言能力,这将在后面详述。下表总结了每个基准的关键特征:
| 基准 | 领域 | 任务类型 | 问题数量 | 测试能力 |
|---|---|---|---|---|
| MMLU | 知识 | 多项选择 | 14k | 跨 57 个学科的广泛学术知识 |
| ARC | 科学与推理 | 多项选择 | 7k | 小学级别的科学推理 |
| HellaSwag | 常识推理 | 多项选择 | 10k | 关于日常情景的常识推理(叙事补全) |
| WinoGrande | 常识推理 | 二元选择 | 1.7k | 需要世界知识的代词消解 |
| CommonSenseQA | 常识推理 | 多项选择 | 1.1k | 关于日常概念的常识推理 |
| OpenBookQA | 科学 | 多项选择 | 500 | 包含推理的小学科学事实 |
| PIQA | 物理常识 | 二元选择 | 1.8k | 关于日常物体的物理常识 |
| GSM8K | 数学 | 自由形式生成 | 1.3k | 小学数学应用题 |
| HumanEval | 代码 | 自由形式生成 | 164 | 根据文档字符串合成 Python 函数 |
注意 MMLU 和 ARC 如何用多项选择来测试事实知识,GSM8K 如何需要计算数学问题的数值答案,以及 HumanEval 如何需要生成完整的 Python 代码。这种多样性确保我们在整个消融实验中测试了模型能力的各个方面。
3.3.2.1 消融实验使用什么数据混合配比?
对于架构消融实验,我们使用固定混合的高质量数据集进行训练,这些数据集能在广泛的任务上提供早期信号。我们使用了英语 (FineWeb-Edu)、数学 (FineMath) 和代码 (Stack-Edu-Python)。架构上的发现应该可以很好地外推到其他数据集和领域,包括多语言数据,因此我们可以保持我们的数据混合相对简单。
对于数据消融实验,我们采取相反的方法:我们固定架构,并系统地改变数据混合配比,以了解不同的数据源如何影响模型性能。
有时评估结果的差异可能很小。如果你的算力充足,可能值得用 不同的随机种子(Seeds) 重新运行相同的消融实验,以观察结果的变化程度。
一个可靠的消融设置的真正价值,并不仅仅在于构建一个好的模型。当我们在主训练运行中不可避免地出现问题时(无论我们准备得多么充分,问题总会发生),我们希望对我们所做的每一个决定都充满信心,并能快速识别哪些组件没有经过充分测试,可能是问题的根源。这种准备工作可以节省调试时间,并 “武装” 我们未来的心理健康。
3.3.3 估算消融实验成本
消融实验非常棒,但它们需要 GPU 时间,因此有必要理解这些实验的成本。下表显示了 SmolLM3 预训练的完整算力分解:主运行(包括偶尔的停机时间)、训练前后的消融实验,以及用于解决意外的扩展问题(迫使我们重启训练)和一些调试所花费的算力(我们将在后面详细介绍)。
| 阶段 | GPU 数量 | 天数 | GPU-小时 |
|---|---|---|---|
| 主预训练运行 | 384 | 30 | 276,480 |
| 消融实验(预训练前) | 192 | 15 | 69,120 |
| 消融实验(训练中) | 192 | 10 | 46,080 |
| 训练重启与调试 | 384/192 | 3/4 | 46,080 |
| 总成本 | - | - | 437,760 |
这些数字揭示了一个重要的事实:消融实验和调试总共消耗了 161,280 个 GPU 小时,超过了我们主训练运行(276,480 个 GPU 小时)的一半成本。 在 SmolLM3 的开发过程中,我们总共进行了超过 100 次消融实验:我们在预训练消融上花费了 20 天,在中途训练消融上花费了 10 天,并花费了 7 天从意外的训练问题中恢复,导致重启和一些调试。
这突出表明了为什么消融实验的成本必须计入你的算力预算:你需要为训练成本 + 消融成本 + 应对意外的缓冲进行规划。如果你以 SOTA 性能为目标,正在实施新的架构更改,或者还没有一个经过验证的配方,那么消融实验将成为一个主要的成本中心,而不仅仅是小实验。
在我们进入下一节之前,让我们先确立一些每个进行实验的人都应该遵守的基本规则。
3.4 实验守则
验证你的评估套件。 在训练任何模型之前,请确保你的评估套件能够复现你将要比较的模型的已发布结果。如果任何基准测试本质上是生成式的(例如 GSM8K),请额外保持警惕,手动检查几个样本,以确保 Prompt 格式正确,并且任何后处理步骤都能提取到正确的信息。由于评估结果将指导你做出每一个决策,因此正确完成这一步对于项目的成功至关重要!
测试每一个更改,无论多小。 不要低估那个看似无害的库升级,或者那个声称“只修改了两行代码”的 Commit 所带来的影响。这些微小的变化可能会引入微妙的 Bug 或性能偏移,从而污染你的结果。你需要一个在对你重要的用例上具有强大测试套件的库,以避免回归。
一次只改变一件事。 在实验之间,保持所有其他设置完全相同。有些变化可能会以意想不到的方式相互作用,所以我们首先要评估每个变化的单独贡献,然后才能尝试将它们结合起来,看看它们的总体影响。
训练足够的 Token,并使用充分的评估。 正如我们前面提到的,我们需要确保评估套件有良好的覆盖,并训练足够长的时间以获得可靠的信号。在这里“偷工减料”将导致嘈杂的结果和错误的决策。
遵守这些规则可能会让你感觉过于谨慎,但另一种选择是花费数天时间调试神秘的性能下降,而结果发现这只是几天前一个不相关的依赖更新造成的。黄金原则:一旦你拥有一个良好的设置,任何更改都应该经过测试!
4. 设计模型架构
既然我们已经有了实验框架,是时候做出定义我们模型的重大决策了。我们做出的每一个选择,从模型大小到注意力机制再到分词器(Tokenizer)的选择,都会创造出各种约束和机遇,直接影响模型的训练和最终的使用。
请记住 “训练指南针”:在做出任何技术选择之前,我们都需要对 “为什么”(Why)和“做什么”(What) 有清晰的认识。我们为什么要训练这个模型?它应该是什么样子?
这听起来很显而易见,但正如我们在“训练指南针”中所解释的,在这里保持深思熟虑将塑造我们的决策,并避免我们在无尽的实验空间中迷失方向。我们的目标是构建一个英语 SOTA 模型吗?长上下文是我们的首要任务吗?还是我们试图验证一种新的架构?虽然所有这些情况下的训练循环可能看起来相似,但我们运行的实验以及我们接受的权衡取舍将截然不同。尽早回答这个问题有助于我们决定如何平衡我们在数据和架构工作之间的时间分配,以及在开始运行之前,要在每个方面创新多少。
因此,让我们以身作则,回顾一下指导 SmolLM3 设计的目标。我们想要一个强大的设备端应用模型,同时具备有竞争力的多语言性能、可靠的数学和编码能力,以及强大的长上下文处理能力。正如我们前面提到的,这引导我们选择了一个 30 亿参数的密集型模型:它足够大以提供强大的能力,但又足够小以舒适地安装在手机上。考虑到边缘设备的内存限制和我们的项目时间线(大约 3 个月),我们选择了密集型 Transformer,而不是 MoE 或混合模型。
我们从 SmolLM2 获得了一个小规模(17 亿参数)的英语训练工作“配方”,但扩大规模意味着重新验证一切,并解决多语言和扩展上下文长度等新挑战。这就是明确的目标如何塑造我们的方法的一个清晰例子。例如,在 SmolLM2 中,我们在预训练结束时努力扩展上下文长度,因此对于 SmolLM3,我们从一开始就做出了架构选择——比如使用 NoPE 和文档内掩码 (intra-document masking)(后续会详细介绍)——以最大化我们成功的几率,而事实证明这种方法奏效了。
一旦我们的目标明确,我们就可以开始做出将目标变为现实的技术决策。在本章中,我们将系统地探讨这些核心决策:架构、数据和超参数。请将此视为我们的战略规划阶段,把这些基础工作做好,将使我们在真正的训练马拉松中避免代价高昂的错误。
4.1 架构选择
如果你观察最近的模型,比如 Qwen3、Gemma3 或 DeepSeek v3,你会发现尽管它们存在差异,但它们都共享着同一个基础——2017 年引入的 Transformer 架构。多年来发生变化的不是基本结构,而是对其核心组件的精炼和改进。无论你是构建密集型模型(Dense Model)、专家混合模型(Mixture of Experts, MoE)还是混合架构(Hybrid Architecture),你都在使用这些相同的构建块。
这些改进源于各大团队对更优性能的追求,以及对特定挑战的攻克:推理时的内存限制、大规模训练时的不稳定性,或处理更长上下文的需求。一些修改,比如从多头注意力(MHA)转向计算效率更高的注意力变体,如分组查询注意力(GQA),现已被广泛采纳。而其他修改,比如不同的位置编码方案,仍在争论之中。最终,今天的实验将凝结成明天的基线架构。
那么,现代 LLM 今天实际在使用什么呢?让我们看看领先模型所趋同的关键点。不幸的是,并非所有模型都公开了其训练细节,但我们从 DeepSeek、OLMo、Kimi 和 SmolLM 等模型家族获得的透明度,足以让我们一窥当前的技术格局:
| 模型 | 架构 | 参数量 | 训练 Token | 注意力机制 | 上下文长度(最终) | 位置编码 | 精度 | 初始化 (Std) | 优化器 | 最大学习率 | 学习率调度 | Warmup 步数 | Batch Size |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| DeepSeek LLM 7B | Dense | 7B | 2T | GQA | 4K | RoPE | BF16 | 0.006 | AdamW | $4.2 \times 10^{-4}$ | Multi-Step | 2K | 9.4M |
| DeepSeek LLM 67B | Dense | 67B | 2T | GQA | 4K | RoPE | BF16 | 0.006 | AdamW | $3.2 \times 10^{-4}$ | Multi-Step | 2K | 18.9M |
| DeepSeek V2 | MoE | 236B (21B 激活) | 8.1T | MLA | 128K | Partial RoPE | - | 0.006 | AdamW | $2.4 \times 10^{-4}$ | Multi-Step | 2K | 9.4M -> 37.7M (Warmup 225B) |
| DeepSeek V3 | MoE | 671B (37B 激活) | 14.8T | MLA | 129K | Partial RoPE | FP8 | 0.006 | AdamW | $2.2 \times 10^{-4}$ | Multi-Step + Cosine | 2K | 12.6M -> 62.9M (Warmup 469B) |
| MiniMax-01 | MoE + Hybrid | 456B (45.9B 激活) | 11.4T | Linear + GQA | 4M | Partial RoPE | - | Xavier Init + Deepnorm | AdamW | $2 \times 10^{-4}$ | Multi-Step | 500 | 16M -> 32M -> 64M -> 128M |
| Kimi K2 | MoE | 1T (32B 激活) | 15.5T | MLA | 128K | Partial RoPE | BF16 | 可能是 0.006 | MuonClip | $2 \times 10^{-4}$ | WSD | 500 | 67M |
| OLMo 2 7B | Dense | 7B | 5T | MHA | 4K | RoPE | BF16 | 0.02 | AdamW | $3 \times 10^{-4}$ | Cosine | 2K | 4.2M |
| SmolLM3 | Dense | 3B | 11T | GQA | 128K | NoPE | BF16 | 0.02 | AdamW | $2 \times 10^{-4}$ | WSD | 2K | 2.3M |
如果你还不理解其中的一些术语,比如 MLA、NoPE 或 WSD,请不用担心。我们将在本节中解释每一个。现在,你只需要注意其中的多样性:不同的注意力机制(MHA, GQA, MLA)、位置编码(RoPE, NoPE, partial RoPE)以及学习率调度(Cosine, Multi-Step, WSD)。
看到如此长的架构选择清单,你可能会感到不知所措,不知道从何处开始。与大多数类似情况一样,我们将循序渐进,逐步建立所有必要的专业知识。我们将首先关注最简单的基础架构(密集型模型),并详细研究每个架构方面。稍后,我们将深入探讨 MoE 和混合模型,并讨论何时使用它们是一个好的选择。最后,我们将探索分词器(Tokenizer)——一个经常被忽视和低估的组件。我们应该使用现有的分词器还是自己训练?我们又如何评估我们的分词器是否优秀?
但现在,让我们从每个 LLM 的核心开始:注意力机制(Attention Mechanism)。
4.1.1 注意力机制(Attention)
Transformer 架构周围最活跃的研究领域之一就是注意力机制。虽然在预训练期间前馈层(feedforward layers)占据了大部分计算,但注意力机制在推理时成为了主要的瓶颈(尤其是在长上下文场景中)。在这种情况下,它会推高计算成本,并且 KV 缓存(KV cache)会迅速消耗 GPU 内存,从而降低吞吐量。让我们快速回顾一下主要的注意力机制,以及它们如何在容量和速度之间进行权衡。
4.1.1.1 我的注意力需要多少个注意力头?
多头注意力(Multi-Head Attention, MHA) 是最初的 Transformer 引入的标准注意力机制。其核心思想是,你有 N 个注意力头,每个头独立地执行相同的检索任务:将隐藏状态转换为查询(Queries)、键(Keys)和值(Values),然后使用当前的查询来匹配键,检索出最相关的 Token,最后转发与匹配 Token 相关联的值。在推理时,我们不需要重新计算过去 Token 的 KV 值,可以直接重用它们。存储过去 KV 值的内存被称为 KV-Cache。随着上下文窗口的增长,这个缓存会迅速成为推理的瓶颈,并消耗 GPU 内存的很大一部分。以下是一个简单的计算,用于估算 Llama 3 架构在 MHA 和 8192 序列长度下的 KV-Cache 内存:
\(s_{KV} = 2 \times n_{bytes} \times seq \times n_{layers} \times n_{heads} \times dim_{heads}\) \(= 2 \times 2 \times 8192 \times 32 \times 32 \times 128 \approx 4 \text{ GB (Llama3 8B)}\) \(= 2 \times 2 \times 8192 \times 80 \times 64 \times 128 \approx 20 \text{ GB (Llama3 70B)}\)
(注意:最前面的因子 2 来自于同时存储键和值缓存。)
正如你所见,缓存大小与序列长度线性增加,但上下文窗口却以指数级增长,现在已经达到了数百万 Token。因此,提高缓存的效率将使推理时的上下文扩展变得更容易。
一个自然而然的问题是:我们真的需要为每个头都计算新的 KV 值吗?可能不需要。多查询注意力(Multi-Query Attention, MQA) (Shazeer, 2019) 和分组查询注意力(Grouped Query Attention, GQA) (Ainslie et al., 2023) 都解决了这个问题。
最简单的情况是在所有头之间共享 KV 值,从而将 KV 缓存的大小除以 $\text{n}_{heads}$。例如,对于 Llama 3 70B,这可以带来 64 倍的减少!这就是 MQA 的思想,并被 StarCoder 等一些模型用作 MHA 的替代方案。
然而,我们可能会因此牺牲一些我们不愿意放弃的注意力容量。因此,我们可以考虑中间立场:在分组的头之间共享 KV 值,例如 4 个头共享相同的 KV 值。这就是 GQA 的方法,它在 MQA 和 MHA 之间取得了平衡。
最近,DeepSeek-v2(并在 v3 中使用)引入了 多潜变量注意力(Multi-Latent Attention, MLA) (DeepSeek-AI et al., 2024),它使用了一种不同的策略来压缩缓存:它不减少 KV 值的数量,而是减少它们的尺寸,简单地存储一个潜变量(latent variable),该变量可以在运行时解压缩为 KV 值。通过这种方法,他们成功地将缓存压缩到相当于 GQA 具有 2.25 个组的等效值,同时提供了比 MHA 更强大的性能!为了让它与 RoPE 配合工作,需要一个小小的调整,即增加一个额外的微小潜向量。在 DeepSeek-v2 中,他们选择了 $4 \times \text{dim}{head}$ 作为主潜变量,和 $1/2 \times \text{dim}{head}$ 用于 RoPE 部分,因此总共是 $4.5 \times \text{dim}_{head}$,它同时用于 K 和 V,从而消除了最前面乘数的 2。
下表比较了我们刚才讨论的注意力机制。为了简化,我们比较了每个 Token 使用的参数量。如果你想计算总内存,只需乘以每个参数的字节数(通常为 2)和序列长度即可:
| 注意力机制 (Attention Mechanism) | 每个 Token 的 KV-Cache 参数量 |
|---|---|
| MHA (多头注意力) | $= 2 \times n_{heads} \times n_{layers} \times dim_{head}$ |
| MQA (多查询注意力) | $= 2 \times 1 \times n_{layers} \times dim_{head}$ |
| GQA (分组查询注意力) | $= 2 \times g \times n_{layers} \times dim_{head}$ (通常 $g=2, 4, 8$) |
| MLA (多潜变量注意力) | $= 4.5 \times n_{layers} \times dim_{head}$ |
现在,让我们看看这些注意力机制在实际实验中的表现如何!
4.1.1.2 消融实验 - GQA 胜过 MHA
在这里,我们比较了不同的注意力机制。我们的基线模型使用了 32 个查询头(Query Heads) 和 8 个 KV 头(KV Heads),这对应于 GQA(分组查询注意力),其比率为 32/8=4。如果我们使用 MHA,或者如果我们使用更少的 KV 头和更高的 GQA 比率,性能将如何变化?
请注意,更改 KV 头的数量会影响参数量,特别是对于 MHA 的情况。为了保持一致性,我们调整了 MHA 运行的层数,因为它否则会有超过 1 亿参数的差异;对于其他设置,我们保持了默认的 16 层。
| 注意力类型 | 查询头(Query Heads) | KV 头(KV Heads) | 层数 | 参数量 | 备注 |
|---|---|---|---|---|---|
| MQA | 32 | 1 | 16 | 1.21B | |
| GQA (比率 16) | 32 | 2 | 16 | 1.21B | |
| GQA (比率 8) | 32 | 4 | 16 | 1.22B | 我们的基线 |
| GQA (比率 4) | 32 | 8 | 16 | 1.24B | |
| GQA (比率 2) | 32 | 16 | 15 | 1.22B | 减少了层数 |
| MHA | 32 | 32 | 14 | 1.20B | 减少了层数 |
| GQA (比率 2) | 32 | 16 | 16 | 1.27B | 参数量过大 - 未进行消融 |
| MHA | 32 | 32 | 16 | 1.34B | 参数量过大 - 未进行消融 |
因此,我们比较了 MHA、MQA 和 4 种 GQA 设置(比率分别为 2、4、8、16)。
观察消融结果,我们发现 MQA 和 GQA(16 个分组)(分别只留下 1 个和 2 个 KV 头)的性能显著低于 MHA。另一方面,GQA 配置(2、4、8 个分组)的性能大致与 MHA 持平。
这个结果在 Loss 曲线和下游评估中都保持一致。我们在 HellaSwag、MMLU 和 ARC 等基准测试中清楚地观察到了这一点,而 OpenBookQA 和 WinoGrande 等基准则显示出少许噪音。
基于这些消融实验,GQA 是 MHA 的一个可靠替代品。它在保持性能的同时,在推理时更加高效。一些最新的模型采用了 MLA 来实现更进一步的 KV 缓存压缩,尽管它尚未被广泛采用。由于在进行消融实验时 MLA 尚未在 nanotron 中实现,我们没有对其进行消融。
对于 SmolLM3,我们最终使用了 GQA,分组数量为 4。
除了注意力架构本身,我们在训练中使用的注意力模式(Attention Pattern)也很重要。接下来,让我们看看注意力掩码(Attention Masking)。
4.1.1.3 文档掩码(Document Masking)
我们如何在训练序列中应用注意力,直接影响着计算效率和模型性能。这就引出了 文档掩码(Document Masking) 以及更广泛的问题:我们在数据加载器(dataloader)中如何构建训练样本?
在预训练期间,我们使用固定长度的序列进行训练,但我们的文档长度是可变的。一篇研究论文可能有 10k 个 Token,而一个简短的代码片段可能只有几百个 Token。我们如何将可变长度的文档放入固定长度的训练序列中呢?
将较短的文档填充(Padding)到目标长度会浪费算力在无意义的填充 Token 上。相反,我们使用打包(Packing):将文档与序列结束(End-of-Sequence, EOS)Token 一起打乱并连接起来,然后将结果分割成与序列大小匹配的固定长度块。
实际操作看起来是这样的:
File 1: "Recipe for granola bars..." (400 tokens) <EOS>
File 2: "def hello_world()..." (300 tokens) <EOS>
File 3: "Climate change impacts..." (1000 tokens) <EOS>
File 4: "import numpy as np..." (3000 tokens) <EOS>
...
After concatenation and chunking into 4k sequences:
Sequence 1: [File 1] + [File 2] + [File 3] + [partial File 4]
Sequence 2: [rest of File 4] + [File 5] + [File 6] + ...
如果一个文件足够长能填满我们的 4k 上下文,那么一个训练序列可能只包含一个完整文件。但在大多数情况下,文件都很短,所以序列包含了多个随机连接在一起的文件。
在标准的 因果掩码(Causal Masking) 下,一个 Token 可以关注打包序列中的所有先前 Token。在上面的例子中,文件 4 中那个 Python 函数里的 Token 可以关注燕麦棒食谱、气候变化文章以及其他碰巧打包在一起的内容。
让我们快速看看一个典型的 4k 预训练上下文会包含什么。一项快速分析显示,在 CommonCrawl 和 GitHub 中,绝大多数(约 80-90%)的文件都短于 2k Token。
下方的图表检查了本文中使用的较新数据集的 Token 分布:
这意味着,在一个 2k 或 4k 的训练序列和标准因果掩码下,绝大多数 Token 将浪费算力去关注那些被打包在一起的、不相关的文档内容。
除了计算效率低下之外,Zhao et al. (2024) 发现这种方法引入了来自不相关内容的噪音,可能降低性能。他们建议使用文档内掩码(intra-document masking):我们修改注意力掩码,使得 Token 只能关注同一文档内的先前 Token。下方可视化图(注:原文指代的图表)展示了这种差异。
Zhu et al. (2025) 在 SkyLadder 中也发现了文档内掩码的类似益处,但提供了另一种解释。他们发现较短的上下文长度对训练更有利,而文档内掩码有效地降低了平均上下文长度。
Llama 3 (Grattafiori et al., 2024) 也使用文档内掩码进行训练,他们发现在短上下文预训练期间影响有限,但对于长上下文扩展来说益处显著,因为在那种情况下注意力开销变得更加重要。此外,ProLong 论文 (Gao et al., 2025) 表明,在持续预训练中利用文档掩码来扩展 Llama 3 8B 的上下文,对长上下文和短上下文的基准都有益处。
我们决定在我们的 1B 基线模型上进行一项消融实验,测试文档掩码是否会影响短上下文性能。你可以在这里找到配置。
结果显示,与标准因果掩码相比,Loss 曲线和下游评估得分完全相同(如下方图表所示)。我们唯一观察到的一个微小改进是在 PIQA 上。
要在 nanotron 中启用文档掩码,只需在模型配置中将以下标志设置为 true:
model_config:
_attn_implementation: flash_attention_2
_fused_rms_norm: true
_fused_rotary_emb: true
- _use_doc_masking: false
+ _use_doc_masking: true # 启用文档内掩码
与 Llama 3 类似,我们没有在短上下文任务上观察到明显的性能影响,除了 PIQA 的微小改进。然而,文档掩码在扩展到长序列时变得至关重要,因为它可以加快训练速度。这对于我们的长上下文扩展尤其重要,我们将序列从 4k 扩展到 64k Token(详情请见“训练马拉松”章节)。因此,我们在 SmolLM3 的整个训练过程中都采用了它。
在本节中,我们讨论了注意力如何处理序列。现在,让我们看看 Transformer 中的另一个主要参数块:嵌入层(Embeddings)。
好的,这是关于“嵌入层共享”和“消融实验”部分的翻译和知乎风格的润色:
4.1.2 嵌入层共享(Embedding Sharing)
如果你观察我们基线消融模型的配置,你会发现与标准 Transformer 不同的地方之一是,它通过 tie_word_embeddings 标志启用了嵌入层共享。
LLM 有两个嵌入组件:
- 输入嵌入(Input Embeddings): 作为 Token 到向量的查找表(大小为 $\text{vocab_size} \times \text{hidden_dim}$)。
- 输出嵌入(Output Embeddings): 最后一个线性层,将隐藏状态映射到词汇表 Logits(大小为 $\text{hidden_dim} \times \text{vocab_size}$)。
在经典情况下,如果这两个矩阵是分开的,总嵌入参数量为 $2 \times \text{vocab_size} \times \text{hidden_dim}$。因此,在小型语言模型中,嵌入层可以占据总参数量的很大一部分,尤其是在词汇表大小很大的情况下。这使得嵌入层共享(在输出层重用输入嵌入)成为小型模型的一种自然优化。
大型模型通常不使用这种技术,因为嵌入层只占其参数预算的一小部分。例如,Llama 3.2 8B 中不共享的总嵌入参数仅占 13%,而在 Llama 3.1 70B 中仅占 3%。
4.1.2.1 消融实验 - 绑定嵌入的模型可媲美参数量更大的非绑定变体
现在我们将评估嵌入层共享对我们消融模型的影响。我们借鉴了 MobileLLM 在 125M 规模上对该技术的全面消融实验所获得的洞察,该实验表明:共享嵌入在参数量减少 11.8% 的同时,准确率退化极小。
由于非绑定嵌入会将我们的参数量从 1.2B 增加到 1.46B,我们将训练另一个具有非绑定参数但层数更少的模型,使其参数量与基线的 1.2B 相匹配。我们将比较三个模型:
- 基线模型 (1.2B): 绑定嵌入(16 层)。
- 非绑定-减少层数模型 (1.2B): 非绑定嵌入,但层数更少(12 层)以保持相同的参数预算。
- 非绑定-相同层数模型 (1.46B): 非绑定嵌入,层数与基线相同(16 层),作为额外的参考点。
Loss 和评估结果表明,我们的 1.2B 绑定嵌入基线模型,在所有基准测试(WinoGrande 除外)上,都实现了与 1.46B 非绑定等效模型相当的性能,尽管其参数少了 18%。
而非绑定嵌入且层数减少的 1.2B 模型(12 层对比 16 层)性能不如前两者,表现出更高的 Loss 和更低的下游评估分数。这表明,在参数预算相等的情况下,增加模型深度比解绑嵌入层带来了更大的益处。
基于这些结果,我们为 SmolLM3 3B 模型保留了绑定嵌入(Tied Embeddings)。
至此,我们探索了嵌入层共享策略及其权衡。但仅靠嵌入层本身并不能捕获序列中 Token 的顺序;提供这些信息是位置编码(Positional Encodings)的作用。在下一节中,我们将探讨位置编码策略是如何演变的,从标准的 RoPE 到像 NoPE (No Position Embedding) 这样更新颖的方法,后者能更有效地进行长上下文建模。
4.1.3 位置编码与长上下文(Positional Encodings & Long Context)
当 Transformer 处理文本时,它们面临一个根本性的挑战:它们天生对词序没有感知,因为它们通过并行注意力操作同时处理整个序列。这使得训练高效,但产生了一个问题。在没有明确的位置信息的情况下,从模型的角度来看,“亚当击败了穆恩”和“穆恩击败了亚当”看起来是相似的。
解决方案是位置嵌入(Positional Embeddings):一种数学编码,赋予序列中的每个 Token 一个独特的“地址”。但是,随着我们不断将上下文推向更长——从早期 BERT 的 512 个 Token 到今天的百万级 Token 模型——位置编码的选择对于性能和计算效率变得越来越关键。
4.1.3.1 位置编码的演变
早期的 Transformer 使用简单的绝对位置嵌入(Absolute Position Embeddings, APE) (Vaswani et al., 2023),它本质上是一个学习到的查找表,将每个位置 $(1, 2, 3…)$ 映射到一个向量,然后将其添加到 Token 嵌入中。这对于短序列运行良好,但有一个主要限制:模型的最大输入序列长度被限制在它所训练的最大长度内。它们不具备开箱即用的泛化到更长序列的能力。
该领域转向了相对位置编码(Relative Position Encodings),它捕获的是 Token 之间的距离,而不是它们的绝对位置。这在直觉上是合理的:两个词相隔 3 个位置,比它们是处于位置 (5, 8) 还是 (105, 108) 更重要。
ALiBi (Attention with Linear Biases) (Press et al., 2022) 通过 Token 距离修改注意力得分。两个 Token 距离越远,它们的注意力得分就会通过应用于注意力权重的简单线性偏差受到越大的惩罚。
但主导近期大型语言模型的技术是旋转位置嵌入(Rotary Position Embedding, RoPE) (Su et al., 2023)。
4.1.3.2 RoPE:将位置编码为旋转
RoPE 的核心洞察是:将位置信息编码为高维空间中的旋转角度。RoPE 不是将位置向量添加到 Token 嵌入中,而是通过依赖于其绝对位置的角度来旋转查询(Query)和键(Key)向量。
其直觉是,我们将嵌入向量中的每对维度视为圆上的坐标,并根据以下因素确定的角度旋转它们:
- Token 在序列中的位置 $p$。
- 我们正在处理的维度对 $k$(不同的维度对以不同的频率旋转,这些频率是基础/参考频率的指数)。
(此处省略了原文提供的 RoPE 简化 Python 代码,以保持文章的可读性,读者可以直接查阅原文代码。)
这个代码可能看起来复杂,所以让我们用一个具体的例子来分解它。考虑句子 “The quick brown fox” 中的“fox”这个词。在我们 1B 基线模型中,每个注意力头都使用一个 64 维的查询/键向量。RoPE 将这个向量分成 32 对:$(x_1, x_2), (x_3, x_4), (x_5, x_6)$,依此类推。我们对对进行操作,因为我们在二维空间中围绕圆旋转。为了简单起见,我们关注第一对 $(x_1, x_2)$。词“fox”出现在句子中的位置 3,因此 RoPE 将旋转第一对维度:
\[\text{rotation\_angle} = \text{position} \times \theta_0 = 3 \times \left(1 / 10000^{(0/32)}\right) = 3 \times 1.0 = 3.0 \text{ 弧度} = 172^{\circ}\]我们的基础频率是 10000,但对于第一对维度 ($k=0$),指数为零,因此基础频率不影响计算(我们将其提高到 0 次方)。
现在,当两个 Token 通过注意力相互作用时,奇迹发生了。它们的旋转表示之间的点积(Dot Product) 直接通过它们旋转角度之间的相位差(Phase Difference) 来编码它们的相对距离(其中 $m$ 和 $n$ 是 Token 的位置):
\[\text{dot\_product}(\text{RoPE}(x, m), \text{RoPE}(y, n)) = \sum_{k} \left[x_k \times y_k \times \cos((m-n) \times \theta_k)\right]\]注意力模式仅取决于 $(m-n)$,因此相隔 5 个位置的 Token 将始终具有相同的角度关系,无论它们在序列中的绝对位置如何。因此,模型学习了基于距离的模式,这些模式适用于序列中的任何绝对位置,并可以外推到更长的序列。
4.1.3.3 如何设置 RoPE 频率?
在实践中,大多数 LLM 预训练都是从相对较短的上下文长度(2K-4K Token)开始的,使用的 RoPE 基础频率是几万,例如 $10\text{K}$ 或 $50\text{K}$。从一开始就用非常长的序列进行训练,由于注意力机制的二次方缩放以及长上下文数据(上下文长度超过 4K 的样本)的有限可用性,计算成本会非常高昂。研究还表明,这可能会损害短上下文性能 (Zhu et al., 2025)。模型通常从学习词语之间的短程相关性开始,因此长序列帮助不大。
典型的方法是用较短的序列完成大部分预训练,然后进行持续预训练(Continual Pretraining),或在最后几千亿 Token 上使用更长的序列。然而,随着序列长度的增长,与 Token 位置成正比的旋转角度也会增长,这可能导致远处 Token 的注意力得分衰减过快 (Rozière et al., 2024; Xiong et al., 2023):
\[\theta = \text{position} \times 1 / (\text{base}^{(k/(\text{dim}/2))})\]解决方案是,随着序列长度的增加而增加基础频率,以防止这种衰减,使用诸如 ABF 和 YaRN 之类的方法。
-
RoPE ABF (RoPE with Adjusted Base Frequency) (Xiong et al., 2023b):通过增加 RoPE 公式中的基础频率来解决长上下文中的注意力衰减问题。这种调整减慢了 Token 位置之间的旋转角度,防止了远处 Token 的注意力得分衰减过快。ABF 可以单阶段应用(直接提升频率)或多阶段应用(随着上下文增长而逐渐增加)。该方法易于实现,并以增加的粒度分布嵌入向量,使模型更容易区分远距离位置。虽然简单有效,但 ABF 对所有维度的统一缩放可能不适用于极长上下文。
-
YaRN (Yet another RoPE extensioN) (Peng et al., 2023):采取了一种更复杂的方法,通过使用斜坡或缩放函数在 RoPE 维度上不均匀地插值频率。与 ABF 的统一调整不同,YaRN 对不同的频率分量应用不同的缩放因子,从而优化了扩展的上下文窗口。它包括动态注意力缩放和注意力 Logits 中的温度调整等额外技术,有助于在极大的上下文尺寸下保持性能。YaRN 支持高效的 “短训练,长测试”策略,只需要更少的 Token 和更少的微调即可实现稳健的外推。尽管比 ABF 更复杂,但 YaRN 通过提供更平滑的缩放和减轻灾难性的注意力损失,通常为极长上下文带来更好的经验性能。它也可以单独在推理中使用,无需任何微调。
这些频率调整方法减缓了注意力得分衰减效应,并保持了远程 Token 的贡献。例如,Qwen3 的训练就涉及在将序列长度从 4k 扩展到 32k 上下文时,使用 ABF 将频率从 10k 增加到 1M(该团队随后应用 YaRN 来达到 131k,即 4 倍外推)。
请注意,对于最佳值目前没有强烈的共识,通常最好在上下文扩展阶段尝试不同的 RoPE 值,以找到最适合你特定设置和评估基准的值。
今天大多数主要的模型都使用 RoPE:Llama、Qwen、Gemma 等等。该技术已被证明在不同模型大小和架构(密集型、MoE、混合型)中都稳健可靠。
4.1.3.4 混合位置编码方法(Hybrid Positional Encoding Approaches)
然而,随着模型推向越来越大的上下文 (Meta AI, 2025; Yang et al., 2025),即使是 RoPE 也开始遇到性能挑战。当在比 Needle in the Haystack (NIAH) (Kamradt, 2023) 更具挑战性的长上下文基准(如 Ruler 和 HELMET (Hsieh et al., 2024; Yen et al., 2025))上进行评估时,在长上下文扩展期间增加 RoPE 频率的标准方法存在局限性。一些更新的技术被引入来提供帮助。
我们以 Transformer 需要位置信息来理解 Token 顺序开始本节,但最近的研究挑战了这一假设。如果明确的位置编码根本不是必需的呢?
NoPE (No Position Embedding) (Kazemnejad et al., 2023) 在没有任何明确位置编码的情况下训练 Transformer,允许模型通过因果掩码和注意力模式隐式学习位置信息。作者表明,与 ALiBi 和 RoPE 相比,这种方法表现出更好的长度泛化能力。由于没有明确的位置编码来外推训练长度之外,NoPE 自然可以处理更长的上下文。然而在实践中,与 RoPE 相比,NoPE 模型在短上下文推理和知识任务上表现较弱 (Yang et al.)。这表明,虽然明确的位置编码可能会限制外推,但它们为训练上下文长度内的任务提供了有用的归纳偏置(inductive biases)。
RNoPE 混合方法: 考虑到这些权衡,B. Yang et al. (2025) 提出结合不同的位置编码策略可能很有趣。他们引入了 RNoPE,它在整个模型中交替使用 RoPE 和 NoPE 层。RoPE 层提供明确的位置信息,并以近因偏置(recency bias)处理局部上下文,而 NoPE 层则改善了跨长距离的信息检索。这项技术最近被用于 Llama 4、Command A 和 SmolLM3。
4.1.3.5 消融实验 - NoPE 在短上下文上与 RoPE 匹配
让我们测试一下混合的 NoPE 方法。我们将比较一个纯 RoPE 1B 消融基线、一个每隔 4 层移除位置编码的 NoPE 变体,以及结合 NoPE 和文档掩码的第三种配置来测试这些技术之间的相互作用。我们的基本问题是:我们能否在保持强大的短上下文性能的同时,获得长上下文能力?
Loss 和评估结果显示,所有三种配置的性能相似,这表明 NoPE 保持了强大的短上下文能力,同时为更好的长上下文处理提供了基础。基于这些结果,我们为 SmolLM3 采用了 NoPE + 文档掩码的组合。
部分/分数 RoPE(Partial/Fractional RoPE): 另一个互补的想法是只在模型维度的一个子集上应用 RoPE。与 RNoPE 在整个层面上交替使用 RoPE 和 NoPE 不同,Partial RoPE 在同一层内混合它们。最近的模型,如 GLM-4.5 (5 Team et al., 2025) 或 Minimax-01 (MiniMax et al., 2025),采用了这种策略,但这在 gpt-j (Wang & Komatsuzaki, 2021) 等较旧的模型中也存在。你也会在每个使用 MLA 的模型中看到这一点,因为它是拥有合理推理成本的必备条件。
4.1.3.6 限制注意力范围以实现长上下文
到目前为止,我们探索了如何处理长上下文的位置信息:启用 RoPE、禁用它 (NoPE)、在某些层上部分应用 (RNoPE) 或在某些隐藏维度上应用 (Partial RoPE),或调整其频率 (ABF, YaRN)。这些方法修改了模型编码位置的方式,以处理比训练期间看到的序列更长的序列。
但还有一种互补的策略:我们可以限制哪些 Token 相互关注,而不是调整位置编码。
为了理解为什么这很重要,考虑一个用 8 个 Token 序列预训练的模型。在推理时,我们想处理 16 个 Token(超过训练长度)。位置 8-15 超出了模型位置编码的分布范围。虽然像 RoPE ABF 这样的技术通过调整位置频率来解决这个问题,但注意力范围方法采取了不同的方法:它们策略性地限制哪些 Token 可以相互关注,将注意力模式保持在熟悉的范围内,同时仍然处理整个序列。这降低了计算成本和内存需求。
下方的图表比较了处理我们的 16 个 Token 序列(预训练窗口为 8)的五种策略:
-
分块注意力(Chunked Attention) 将序列分成固定大小的块,其中 Token 只能在自己的块内关注。在我们的示例中,16 个 Token 被分成两个 8 个 Token 的块(0 到 7 和 8 到 15),每个 Token 只能看到其块内的其他 Token。注意 Token 8 到 15 根本不能关注到前面的块。这创建了在块边界重置的隔离注意力窗口。Llama 4 (Meta AI, 2025) 在 RoPE 层(四分之三的解码器层)中使用了 8192 个 Token 块的分块注意力,而 NoPE 层则保持对完整上下文的访问。这通过限制每层的 KV 缓存大小来减少内存需求,尽管它意味着 Token 不能关注到前面的块,这可能会影响某些长上下文任务。
-
滑动窗口注意力(Sliding Window Attention, SWA),由 Mistral 7B (Child et al., 2019; Jiang et al., 2023) 推广,基于最近的 Token 最相关的直觉,采用了不同的方法。每个 Token 只关注最近的 N 个 Token,而不是硬性的块边界。在图表中,每个 Token 最多可以看到前面 8 个位置,创建了一个在序列中连续移动的滑动窗口。注意 Token 15 可以关注位置 8 到 15,而 Token 10 关注位置 3 到 10。窗口向前滑动,在整个序列中保持局部上下文,没有分块的人为障碍。Gemma 3 将 SWA 与完整注意力在交替层中结合使用,类似于混合位置编码方法混合不同策略的方式。
-
双块注意力(Dual Chunk Attention, DCA) (An et al., 2024) 是一种免训练的方法,它扩展了分块注意力,同时保持了跨块的信息流。在我们的示例中,我们使用块大小 $s=4$,将 16 个 Token 分为 4 个块(沿对角线可视化 $4 \times 4$ 方格)。DCA 结合了三种机制:(1) 块内注意力,Token 在其块内正常关注(对角线模式)。(2) 块间注意力,查询使用位置索引 $c-1=7$ 来关注前面的块,创建相对位置上限为 7。(3) 连续块注意力,具有局部窗口 $w=3$,保留相邻块之间的局部性。这使得所有相对位置都保持在训练分布内(0 到 7),同时保持跨块边界的平滑过渡。DCA 使 Qwen 2.5 等模型能够在推理时支持高达 100 万 Token 的超长上下文窗口,而无需在百万 Token 序列上进行持续训练。
4.1.3.7 注意力汇聚(Attention Sinks)
在具有长上下文的 Transformer 模型中,出现了一种有趣的现象:模型会为序列中的起始 Token 分配异常高的注意力分数,即使这些 Token 在语义上并不重要。这种行为被称为注意力汇聚(Attention Sinks) (Xiao et al.)。这些起始 Token 充当了注意力分布的稳定机制,起到了注意力可以积累的“汇聚点”作用。
实际的洞察是:当上下文长度超过缓存大小时,仅保留起始几个 Token 的 KV 缓存以及最近 Token 的滑动窗口,可以在很大程度上恢复性能。这种简单的修改使模型能够处理更长的序列,而无需微调或性能下降。
现代的实现以不同的方式利用注意力汇聚。最初的研究建议在预训练期间添加一个专用的占位符 Token 作为明确的注意力汇聚点。最近,像 gpt-oss 这样的模型将注意力汇聚实现为学习到的“每头偏差 Logits”(learned per-head bias logits),将其附加到注意力分数上,而不是作为输入序列中的实际 Token。这种方法在不修改分词输入的情况下达到了相同的稳定效果。
有趣的是,gpt-oss 还在注意力层本身使用了偏差单元(bias units),这是自 GPT-2 以来就很少见的设计选择。虽然这些偏差单元通常被认为对于标准注意力操作是多余的(Dehghani et al. 的实证结果显示对测试 Loss 的影响很小),但它们可以服务于实现注意力汇聚的专门功能。
核心洞察是:无论是作为特殊 Token、学习到的偏差还是每头 Logits 实现,注意力汇聚都为长上下文场景中的注意力分布提供了一个稳定的“锚点”,允许模型存储关于整个序列的通用有用信息,即使上下文任意增长。
至此,我们已经涵盖了注意力的核心组件:平衡内存和计算的不同头部配置(MHA、GQA、MLA),帮助模型理解 Token 顺序的位置编码策略(RoPE、NoPE 及其变体),以及使长上下文变得可处理的注意力范围技术(滑动窗口、分块和注意力汇聚)。我们还研究了嵌入层应如何配置和初始化。这些架构选择定义了你的模型如何处理和表示序列。
但拥有正确的架构只是成功的一半。即使是精心设计的模型,也可能遭受训练不稳定的困扰,尤其是在大规模训练时。让我们来看看有助于保持训练稳定的技术。
4.1.4 提高稳定性(Improving Stability)
现在,让我们转向 LLM 预训练中最大的挑战之一:不稳定性(Instabilities)。这些问题通常表现为 Loss 尖刺或训练 Loss 的突然跳跃,在大规模训练时尤为常见。
虽然我们将在“训练马拉松”部分深入探讨不同类型的尖刺以及如何处理它们(深入研究浮点精度、优化器和学习率),但某些架构和训练技术也可以帮助我们减少不稳定性。因此,让我们花点时间在这里研究它们。
我们将介绍最近大规模训练运行中(例如,Olmo2 (OLMo et al., 2025) 和 Qwen3 (A. Yang, Li, et al., 2025))用于提高稳定性的几种简单技术:Z-Loss、从嵌入层中移除权重衰减和 QK-norm。
4.1.4.1 Z-loss
Z-loss (Chowdhery et al., 2022) 是一种正则化技术,它通过在损失函数中添加一个惩罚项来防止最终输出的 Logits 变得过大。这种正则化促使 Softmax 的分母(即 $Z$ 项)保持在一个合理的范围内,这有助于在训练过程中保持数值稳定性。
下方在我们的 1B 模型上进行的消融结果(注:指代原文图表)显示,添加 Z-loss 不会影响训练 Loss 或下游性能。对于 SmolLM3,我们最终没有使用它,因为在开始训练时,我们的 Z-loss 实现引入了一些我们没有优化的训练开销。
4.1.4.2 从嵌入层中移除权重衰减
权重衰减(Weight Decay)通常作为一种正则化技术应用于所有模型参数,但 OLMo et al. (2025) 发现,将嵌入层排除在权重衰减之外可以提高训练稳定性。
其推理是:权重衰减会导致嵌入层范数在训练过程中逐渐减小,这可能导致早期层中的梯度变大,因为层归一化(Layer Normalization)的雅可比矩阵与输入范数成反比 (Takase et al., 2025)。
我们通过训练三种配置来测试这种方法:
- 基线模型: 使用标准权重衰减。
- 变体模型: 嵌入层无权重衰减。
- 组合模型: 结合我们所有已采纳的更改(嵌入层无权重衰减 + NoPE + 文档掩码),以确保技术之间没有负面相互作用。
Loss 曲线和评估结果在所有三种配置中几乎相同。因此,我们在 SmolLM3 训练中采用了所有这 3 项更改。
4.1.4.3 QK-norm
QK-norm (Dehghani et al., 2023) 在计算注意力之前,对查询(Query)和键(Key)向量同时应用层归一化(Layer Normalization)。这项技术有助于防止注意力 Logits 变得过大,并被许多最近的模型用于提高稳定性。
然而,B. Yang et al. (2025) 发现 QK-norm 会损害长上下文任务。他们的分析显示,QK-norm 导致相关 Token(即“针”)上的注意力质量降低,而不相关上下文上的注意力质量增高。他们认为发生这种情况是因为归一化操作消除了query-key点积中的幅度信息(magnitude information),这使得注意力 Logits 的幅度更接近。
由于这个原因,我们没有在 SmolLM3 中使用 QK-norm。此外,作为一个小的 30 亿参数模型,与那些 QK-norm 已被证明最有益处的更大模型相比,它面临的训练不稳定性风险也较小。
4.1.5 其他核心组件(Other Core Components)
除了我们已经涵盖的组件之外,还有一些其他值得一提的架构决策,以求完整性。
-
参数初始化(Initialization): 现代模型通常使用截断正态分布初始化(均值=0,标准差 $\text{std}=0.02$ 或 $\text{std}=0.006$),或者像 $\mu\text{P}$ (G. Yang & Hu, 2022) 这样的初始化方案,例如 Cohere 的 Command A (Cohere et al., 2025)。这也可以是另一个消融实验的主题。
-
激活函数(Activation Functions): SwiGLU 已成为现代 LLM 中的事实标准(除了使用 GeGLU 的 Gemma2 和使用 $\text{relu}^2$ 的 Nvidia (Nvidia et al., 2024; NVIDIA et al., 2025)),取代了像 ReLU 或 GELU 这样较旧的选择。
-
架构布局(Architectural Layout): 从更宏观的层面看,架构布局的选择也对模型的行为起着作用。尽管总参数数量在很大程度上决定了语言模型的容量,但这些参数如何分布在深度(depth)和宽度(width)上也至关重要。Petty et al. (2024) 发现,在语言建模和组合性任务上,更深的模型比参数量相等的更宽的模型表现更优,直到这种益处达到饱和。“深而薄”的策略在 MobileLLM 的子十亿参数 LLM 消融实验中表现良好 (Z. Liu et al., 2024),而更宽的模型由于具有更大的并行性,往往能提供更快的推理速度。现代架构以不同的方式体现了这种权衡取舍,正如本文所指出的。
至此,我们已经涵盖了值得为你的训练运行进行优化的密集型 Transformer 架构的最重要方面。
然而,最近也出现了涉及模型整体的其他架构干预,即 MoE(专家混合)模型和混合(Hybrid)模型。让我们来看看它们能提供什么,从 MoE 开始。
4.1.6 走向稀疏:专家混合模型(MoE)
专家混合模型(MoE)的直觉是:我们不需要为了每一个 Token 的预测都用到整个模型,这类似于我们的大脑会根据手头的任务激活不同的区域(例如视觉或运动皮层)。对于一个 LLM 来说,这意味着模型在执行翻译任务时,那些学习了代码语法的组件就不需要被使用。如果能做好这一点,就意味着我们可以节省大量的计算资源,因为在推理时我们只需要运行部分模型。
在技术层面上,MoE 的目标很简单:增加总参数量,同时不增加每个 Token 的“活跃”参数量。简化来说,总参数量影响模型的总体学习容量,而活跃参数量决定了训练成本和推理速度。这就是为什么现在许多前沿系统(例如 DeepSeek V3、K2,以及 Gemini、Grok 等闭源模型)都在使用 MoE 架构。如果你是第一次接触 MoE,不用担心,其机制并不复杂。让我们从标准的密集型架构开始,看看 MoE 所需的必要改变。在 MoE 中,我们将单个 MLP 替换为多个 MLP(“专家 Experts”),并在这些 MLP 之前添加一个可学习的路由器(Router)。对于每个 Token,路由器会选择一小部分专家来执行计算。总参数量和活跃参数量的区别就源于此:模型拥有很多专家,但任何给定的 Token 只使用其中少数几个。
设计 MoE 层会引出几个核心问题:
- 专家形态与稀疏度: 你应该使用许多小型专家还是少数大型专家?每个 Token 应该有多少专家是活跃的?你总共需要多少专家(即稀疏度或“top-k”)?是否应该让某些专家成为 通用型(universal) 专家而始终保持活跃?
- 利用率与专业化: 如何选择路由的专家,并确保它们被充分利用(避免闲置容量),同时鼓励它们实现专业化?在实践中,这是一个 负载均衡(load-balancing) 问题,对训练和推理效率有着重要影响。
在这里,我们专注于一个目标:给定固定的计算预算,我们如何选择一个能将 Loss 降到最低的 MoE 配置? 这是一个不同于纯系统效率(吞吐量/延迟)的问题,我们稍后会再讨论后者。本节的大部分内容遵循蚂蚁集团 MoE 扩展定律论文 (Tian et al., 2025) 中的分析。我们将使用他们提出的效率杠杆(Efficiency Leverage, EL)的概念。简单来说,EL 衡量了你需要多少密集型计算才能匹配 MoE 设计所达到的 Loss,衡量单位是 FLOPs。更高的 EL 意味着与密集型训练相比,该 MoE 配置每单位计算能带来更多的 Loss 改进。让我们仔细看看如何设置 MoE 的稀疏度来提高效率杠杆。
4.1.6.1 稀疏度/激活比率(Sparsity / Activation Ratio)
在本节中,我们想找出哪种 MoE 设置是最佳的。渐进地来看,很容易看出两个极端情况都不是理想的设置:
- 一方面,始终激活所有专家会使我们回到密集型设置,即所有参数始终被使用。
- 另一方面,如果活跃参数非常少(极端情况下只激活 1 个参数),显然不足以解决任务,即使是在一个狭窄的领域。
因此,我们显然需要找到一个中间点。在我们深入寻找最佳设置之前,定义两个量是有用的:激活比率(activation ratio)及其倒数稀疏度(sparsity):
\[\text{激活比率} = \frac{\text{激活的专家数量}}{\text{总专家数量}}\] \[\text{稀疏度} = \frac{\text{总专家数量}}{\text{激活的专家数量}} = \frac{1}{\text{激活比率}}\]从计算的角度来看,成本仅由活跃参数驱动。如果你保持活跃专家的数量(和大小)固定,并增加专家的总数,你的推理/训练 FLOPs 预算大致保持不变,但你增加了模型的容量,因此只要训练时间足够长,模型通常会变得更好。
如果你对最近的 MoE 论文进行调查,会发现一些有趣的经验性结论:在固定活跃专家的数量和大小的情况下,增加专家的总数(即降低激活比率/增加稀疏度)可以改善 Loss,但当稀疏度变得非常高时,回报会递减。
两个例子:
- Kimi K2 (K. Team et al., 2025):显示了两种效应:更高的稀疏度提高了性能,但随着稀疏度的增长,收益会逐渐减少。
- 蚂蚁集团 (Tian et al., 2025):与 K2 得出相同结论,并额外指出稀疏度越高的 MoE 从增加计算中获得的益处越多。
下面是一些 MoE 模型的稀疏度表格:
| 模型 | 总专家数量 | 每个 Token 激活数量(含共享) | 稀疏度 |
|---|---|---|---|
| Mixtral-8x7B | 8 | 2 | 4.0 |
| Grok-1 | 8 | 2 | 4.0 |
| Grok-2 | 8 | 2 | 4.0 |
| OLMoE-1B-7B-0924 | 64 | 8 | 8.0 |
| gpt-oss 20b | 32 | 4 | 8.0 |
| Step-3 | 48 路由 + 1 共享 = 49 | 3 路由 + 1 共享 = 4 | 12.25 |
| GLM-4.5-Air | 128 路由 + 1 共享 = 129 | 8 路由 + 1 共享 = 9 | 14.3 |
| Qwen3-30B-A3B | 128 | 8 | 16.0 |
| Qwen3-235B-A22B | 128 | 8 | 16.0 |
| GLM-4.5 | 160 路由 + 1 共享 = 161 | 8 路由 + 1 共享 = 9 | 17.8 |
| DeepSeek-V2 | 160 路由 + 2 共享 = 162 | 6 路由 + 2 共享 = 8 | 20.25 |
| DeepSeek-V3 | 256 路由 + 1 共享 = 257 | 8 路由 + 1 共享 = 9 | 28.6 |
| gpt-oss 120b | 128 | 4 | 32.0 |
| Kimi K2 | 384 路由 + 1 共享 = 385 | 8 路由 + 1 共享 = 9 | 42.8 |
| Qwen3-Next-80B-A3B-Instruct | 512 路由 + 1 共享 = 513 | 10 总激活 + 1 共享 = 11 | 46.6 |
最近的趋势很明显:MoE 模型正变得越来越稀疏。尽管如此,最佳稀疏度仍然取决于硬件和端到端效率。例如,Step-3 旨在达到峰值效率,并有意不将稀疏度最大化,以适应其特定的硬件和带宽约束,而 gpt-oss-20b 的稀疏度较低则是由于设备上的内存限制(被动专家仍然占用一些内存)。
4.1.6.2 粒度(Granularity)
除了稀疏度之外,我们还需要决定每个专家应该有多大。这由蚂蚁集团引入的 粒度(Granularity) 来捕获。让我们明确这个术语的含义。术语在不同论文中有所不同,有些使用略有不同的公式。在这里,我们将使用与我们引用的图表相匹配的定义:
\[G = \alpha \times \frac{d_{model}}{d_{expert}}\]更高的粒度值对应于拥有 更多具有更小维度($d_{expert}$)的专家(给定固定的参数数量)。该指标是 专家维度($d_{expert}$) 与 模型维度($d_{model}$) 之间比值的倒数(乘上归一化系数 $\alpha$)。
在密集型模型中,一个常用的经验法则是将 MLP 的维度设置为 $d_{intermediate} = 4 \times d_{model}$。如果 $\alpha=4$(像 Krajewski et al. (2024)),你可以大致将粒度视为匹配密集型 MLP 宽度所需的专家数量($4 \times d_{model} = d_{intermediate} = G \times d_{expert}$)。
这种解释只是一个粗略的启发式:现代 MoE 设计通常分配比单个密集型 MLP 大得多的总容量,因此一对一匹配在实践中会失效。蚂蚁团队的设置选择了 $\alpha=2$,这只是一种不同的归一化选择。为了保持一致性,我们将采用这种约定并坚持下去。
以下是一些 MoE 版本的不同粒度值表格:
| 模型 | $d_{model}$ | $d_{expert}$ | $G = 2 \times d_{model} / d_{expert}$ | 年份 |
|---|---|---|---|---|
| Mixtral-8x7B | 4,096 | 14,336 | 0.571 | 2023 |
| gpt-oss-120b | 2,880 | 2,880 | 0.5 | 2025 |
| gpt-oss-20b | 2,880 | 2,880 | 0.5 | 2025 |
| Grok 2 | 8,192 | 16,384 | 1.0 | 2024 |
| StepFun Step-3 | 7,168 | 5,120 | 2.8 | 2025 |
| OLMoE-1B-7B | 2,048 | 1,024 | 4.0 | 2025 |
| Qwen3-30B-A3B | 2,048 | 768 | 5.3 | 2025 |
| Qwen3-235B-A22B | 4,096 | 1,536 | 5.3 | 2025 |
| GLM-4.5-Air | 4,096 | 1,408 | 5.8 | 2025 |
| DeepSeek V2 | 5,120 | 1,536 | 6.6 | 2024 |
| GLM-4.5 | 5,120 | 1,536 | 6.6 | 2025 |
| Kimi K2 | 7,168 | 2,048 | 7.0 | 2025 |
| DeepSeek V3 | 7,168 | 2,048 | 7.0 | 2024 |
| Qwen3-Next-80B-A3B | 2,048 | 512 | 8.0 | 2025 |
让我们谈谈粒度如何影响行为(来自蚂蚁集团的论文):
粒度看起来不是 EL 的主要驱动因素——它有帮助,尤其是当粒度超过 2 时,但它不是决定 Loss 的主导因素。不过,存在一个最佳点:提高粒度到一定程度会有帮助,然后收益就会趋于平稳。因此,粒度是一个有用的调整旋钮,最近发布的趋势明显倾向于更高的值,但不应该孤立地进行优化。
另一种广泛用于改进 MoE 的方法是 共享专家(Shared Experts) 的概念。让我们来看看!
4.1.6.3 共享专家(Shared Experts)
共享专家设置将每个 Token 路由到一组始终处于激活状态的专家。这些共享专家吸收数据中基本、重复的模式,以便其余的专家可以更积极地进行专业化。
在实践中,你通常不需要很多共享专家;模型设计者通常选择一个,最多两个。随着粒度的增加(例如,从 Qwen3 风格的设置转向更接近 Qwen3-Next 的设置),共享专家往往变得更有用。从下方的图表来看,总体影响是适度的,它不会戏剧性地改变 EL。一个简单的经验法则是只使用一个共享专家,这与 DeepSeek V3、K2 和 Qwen3-Next 等模型的选择相符,并且倾向于在不增加不必要复杂性的情况下最大化效率。
那么,共享专家就是所有 Token 都会始终路由通过的专家。那其他专家呢?我们如何学习何时路由到每个专家,并确保我们不会只使用少数几个专家?接下来我们将讨论负载均衡,它正是解决这个问题的。
4.1.6.4 负载均衡(Load Balancing)
负载均衡是 MoE 中的关键环节。 如果设置不当,它可能会破坏所有其他设计选择。通过以下示例,我们可以看到为什么糟糕的负载均衡会给我们带来很多痛苦。
考虑一个非常简单的分布式训练设置,我们有 4 块 GPU,并将模型的 4 个专家均匀地分布在这些 GPU 上。如果路由崩溃,所有 Token 都被路由到专家 1,这意味着只有 1/4 的 GPU 被利用,这对训练和推理效率来说非常糟糕。除此之外,这也意味着模型的有效学习容量也降低了,因为并非所有专家都被激活。
为了解决这个问题,我们可以向路由器添加一个额外的损失项。你可以在下方看到标准的 基于辅助损失(auxiliary loss–based) 的负载均衡 (LBL):
\[\text{Loss}_{LBL} = \alpha \times \sum_{i} f_i \times P_i\]这个简单的公式只使用了三个因子:
- 系数 $\alpha$ 决定了损失的强度。
- $f_i$ 是流量分数(traffic fraction),即流经专家 $i$ 的 Token 分数。
- $P_i$ 是概率质量(probability mass),简单地是流经该专家的 Token 概率总和。
两者都是必需的:$f_i$ 对应于实际的均衡,而 $P_i$ 是平滑且可微分的,允许梯度流动。如果我们实现了完美的负载均衡,我们会得到 $f_i = P_i = 1/N_r$。然而,我们需要小心如何调整 $\alpha$,因为 $\alpha$ 值太小,我们对路由的引导不够;而如果 $\alpha$ 太大,路由的均匀性就变得比主要的语言模型损失更重要。
一个关键的细节是计算路由统计信息的范围:$f_i$ 和 $P_i$ 是按局部批次(每个 Worker 的 Mini-Batch)计算,还是按全局(跨 Worker/设备聚合)计算?Qwen 团队的分析 (Qiu et al., 2025) 表明,当每个局部批次中没有足够的 Token 多样性时,局部计算可能会损害专家专业化(路由健康状况的良好代理)和整体模型性能。专家专业化是一种现象,即一个或多个专家在特定领域被激活得比其他专家更频繁。换句话说,如果一个局部批次很窄,它的路由统计数据就会嘈杂/有偏差,不会带来好的均衡。这表明,只要可行,我们就应该使用全局统计信息(或至少是跨设备的聚合)。值得注意的是,在那篇论文发表时,许多框架——包括 Megatron——默认都是局部计算这些统计信息的。
下方的图表(注:指代原文图表)来自 Qwen 的论文,说明了 Mini-Batch 与全局批次聚合之间的差异及其对性能和专业化的影响。
通常,围绕 MoE 进行架构选择的消融实验是棘手的,因为它涉及许多方面的相互作用。例如,共享专家的有效性可能取决于模型的粒度。因此,值得花一些时间来确保你有一套好的实验,才能真正获得你正在寻找的洞察!
我们现在已经涵盖了 MoE 的基础知识,但仍有更多内容有待发现。以下是一些非详尽的、可供进一步研究的项目清单:
- 零计算专家、MoE 层重新缩放和训练监控(来自 LongCat-Flash 论文)。
- 正交损失负载均衡(如 ERNIE 4.5 中所示)。
- 在训练过程中调度负载均衡系数。
- 架构/优化与 MoE 的相互作用,例如:
- 优化器排名是否会因 MoE 而改变。
- 如何将 $\mu\text{P}$ 应用于 MoE。
- 如何为 MoE 调整学习率(因为它们在每个批次中看到的 Token 数量不同)。
- 起始处的密集层数量。
- …更多。
我们把进一步深入探索的重任留给你们这些求知若渴的读者,现在我们将转向最后一个主要的架构选择:混合模型(Hybrid Models)!
参考
Enjoy Reading This Article?
Here are some more articles you might like to read next: