dots.llm1 Technical Report

dots.llm1 技术报告

作者/机构: rednote-hilab
链接: https://huggingface.co/rednote-hilab, https://github.com/rednote-hilab/dots.llm1


A1 主要贡献

本文介绍了dots.llm1,这是一款大规模、高性价比的专家混合(Mixture of Experts, MoE)模型。该模型总参数量为1420亿,但每个token仅激活140亿参数,从而在降低训练和推理成本的同时,实现了与当前最先进模型相媲美的性能。

核心问题与研究目标
当前,大型语言模型(LLM)正迅速发展,但其巨大的计算成本成为一个主要挑战。MoE架构通过仅激活模型参数的一个子集来提高计算效率和可扩展性,已成为扩展语言模型的一种有前途的范式。本研究的目标是开发一款开源的MoE模型,它既能在性能上与顶级的密集模型竞争,又能在训练和推理上具有显著的成本效益。

核心创新与贡献
dots.llm1模型通过结合先进的架构、数据处理流程和基础设施优化,实现了其性能和效率目标。主要贡献总结如下:

如下图所示,dots.llm1在MMLUPro基准测试中,其性能与领先的密集模型和MoE模型相当,同时保持了成本优势。


图 1: 开源MoE和密集语言模型的性能与成本比较。圆圈(◦)表示密集模型,菱形(⋄)表示MoE模型。我们使用MMLUPro对模型能力进行基准测试,结果表明dots.llm1取得了与领先模型相当的准确率。


A2 方法细节

架构

模型架构概述。dots.llm1模型采用仅解码器(decoder-only)的Transformer架构【Vaswani et al., Attention is all you need, Advances in neural information processing systems, 2017】,其中每一层由一个注意力层和一个前馈网络(FFN)组成。与Llama【AI@Meta, Llama 3 model card, 2024a】或Qwen【Yang et al., Qwen2. 5 technical report, arXiv, 2024】等密集模型不同,该模型中的FFN被替换为专家混合(MoE)模块【Jiang et al., Mixtral of experts, CoRR, 2024; DeepSeek-AI, Deepseek LLM: scaling open-source language models with longtermism, CoRR, 2024a;b; DeepSeek-AI et al., Deepseek-v3 technical report, 2024】。这一修改使得在保持经济成本的同时,能够训练出能力强大的模型。

注意力层。我们在模型中采用了原始的多头注意力机制(multi-head attention)【Vaswani et al., Attention is all you need, Advances in neural information processing systems, 2017】。遵循Dehghani等人【Dehghani et al., Scaling vision transformers to 22 billion parameters, In International Conference on Machine Learning, 2023】的做法,在计算注意力之前,对查询(query)和键(key)的投影应用了RMSNorm。这种归一化操作降低了出现过大注意力对数(logits)的风险,从而避免了训练过程的不稳定【Wortsman et al., Small-scale proxies for large-scale transformer training instabilities, arXiv, 2023; Tian et al., Nyonic technical report, arXiv, 2024a; OLMo et al., 2 olmo 2 furious, arXiv, 2024】。

专家混合层。我们借鉴了DeepSeek-AI【DeepSeek-AI, Deepseek LLM: scaling open-source language models with longtermism, CoRR, 2024a】和Qwen【Qwen, Qwen1.5-moe: Matching 7b model performance with 1/3 activated parameters, 2024b】的设计,将FFN替换为一个包含共享专家和独立专家的专家混合(MoE)层。我们的实现包含128个路由专家和2个对所有tokens都激活的共享专家,每个专家都实现为一个使用SwiGLU【Shazeer, Glu variants improve transformer, arXiv, 2020】激活函数的细粒度、两层FFN。对于每个token,路由器会选择排名前6的独立专家,再加上2个共享专家,因此每个token总共有8个激活的专家。值得注意的是,我们在门控层(gating layer)的计算中采用了FP32精度而不是BF16,以确保路由过程中的数值稳定性和更准确的专家选择。

负载均衡。MoE层中不均衡的专家负载会导致路由崩溃,从而降低模型在训练和推理过程中的容量和计算效率。为了解决这个问题,我们采用了一种无辅助损失(auxiliary-loss-free)的方法【Wang et al., Auxiliary-loss-free load balancing strategy for mixture-of-experts, CoRR, 2024a】,该方法也应用于DeepSeek-AI等人【DeepSeek-AI et al., Deepseek-v3 technical report, 2024】的模型中。它为每个专家引入一个偏置项(bias term),该偏置项被加到相应的亲和度分数(affinity scores)上,以确定top-k路由。这个偏置项在训练过程中动态调整,以维持专家间的负载均衡。此外,我们还采用了一种序列级(sequence-wise)的平衡损失,以防止任何单个序列内出现极端的不平衡。由于采用了有效的负载均衡策略,dots.llm1在整个训练过程中保持了良好的负载均衡,并且在训练期间没有丢弃任何token。

基础设施

训练框架Cybertron。dots.llm1的训练由内部的Cybertron框架支持,这是一个基于Megatron-Core构建的轻量级训练框架。我们利用Megatron-Core精心构建了一套用于模型预训练和后训练的完整工具包。针对预训练、监督微调(SFT)和强化学习(RL)等不同训练阶段,我们以统一的方式封装了独立的训练器,以保证训练过程的连贯性和高效率。

3.1 基于交错式1F1B的通信与计算重叠

一种创新的重叠方案。我们提出了一种创新的基于交错式1F1B的all-to-all通信与计算重叠解决方案【NVIDIA, MoE A2A Interleaved 1F1B based Computation and Communication Overlap, 2024b】,并与NVIDIA合作将其集成到Megatron-Core中。与原始的交错式1F1B流水线相比,我们在预热(warm-up)阶段增加了一个额外的步骤,这个步骤既不产生额外的气泡(bubble)开销,也不增加激活函数的GPU内存消耗。在这种精细化的流水线调度方案下,在稳定的1F1B阶段,前向和后向步骤对内的all-to-all通信与计算可以被有效重叠,这与DeepSeek提出的DualPipe【DeepSeek-AI et al., Deepseek-v3 technical report, 2024】类似。与DualPipe相比,我们的方法在内存消耗方面表现出显著优势,尽管气泡率略高。


图 2: 基于交错式1F1B的通信与计算重叠

3.2 Grouped GEMM 的高效实现

Grouped GEMM的重要性与现有方法。分组GEMM(Grouped GEMMs)在MoE架构的计算中扮演着重要角色。一个直接的方法是在流式多处理器(streaming multiprocessors)上统一调度子GEMM问题。一些框架在主机端预先计算瓦片(tiles)和线程块(thread blocks)之间的映射,以减少设备端的调度开销【Li et al., A coordinated tiling and batching framework for efficient GEMM on gpus, In PPoPP, 2019】。其他工作探索了针对小型或可变大小批处理GEMM的调度策略【Yang et al., A batched GEMM optimization framework for deep learning, J. Supercomput., 2022; Zhang et al., Accelerating small matrix multiplications by adaptive batching strategy on GPU, In IEEE Int Conf on High Performance Computing & Communications, 2022】。此外,NVIDIA在cuBLAS【Hejazi, Introducing grouped gemm apis in cublas and more performance updates, 2024】和Transformer Engine (TE)【NVIDIA, pytorch - transformer engine 2.1.0 documentation, 2024a】中引入了Grouped GEMM API。

我们的实现方法。为了支持更大的瓦片尺寸并减少调度开销,我们将每个专家的token段$M_i$对齐到一个固定的块大小。这个固定的块大小必须是异步warpgroup级矩阵乘法累加(WGMMA)指令(wgmma.mma async)的瓦片形状修饰符$mMnNkK$中$M$的倍数。因此,单个线程块中的warpgroups将拥有统一的瓦片划分,并且由一个线程块处理的整个token段($M_i$)必然属于同一个专家,这使得调度过程与普通GEMM非常相似。

性能对比。与NVIDIA在Transformer Engine (v2.1)中的Grouped GEMM API相比,我们的实现展现出显著优势。表1展示了在H800上,当tokens被均匀路由到专家时,前向和后向计算的性能比较。我们的方法在前向计算上实现了平均14.00%的性能提升,在后向计算上实现了平均6.68%的性能提升。

表 1: Transformer Engine 2.1 与我们在H800上的实现性能比较。其中,g, m, n, 和 k 分别代表组(专家)的数量以及每个子GEMM问题的维度。


A3 预训练

4.1 预训练数据

训练数据的数量、质量和多样性对语言模型的性能起着决定性作用。为实现这些目标,我们将方法分为三个关键阶段:文档准备、基于规则的处理和基于模型的处理。这种设计使我们能够用有限的计算资源高效地管理大量数据。文档准备阶段专注于对原始数据进行预处理和组织。基于规则的处理旨在通过自动过滤和清洗数据,最大限度地减少对大量人工筛选的需求。基于模型的处理则进一步确保最终的数据集既高质量又多样化。为确保安全性,我们应用了过滤器来排除可能包含不安全内容或大量个人可识别信息(PII)的网站数据,以及被安全分类器标记为有害的域名。在预训练期间,我们保持中英文数据1:1的平衡比例,并且没有使用任何合成数据。

数据处理流程的详细信息记录在附录C中。我们的数据处理流程有两个关键创新点:

4.2 超参数

dots.llm1模型使用AdamW优化器【Loshchilov & Hutter, Decoupled weight decay regularization, CoRR, 2017】进行训练,其中$\beta_1 = 0.9$,$\beta_2 = 0.95$。我们实施了0.1的权重衰减,并将梯度裁剪设置为1.0。遵循Dai等人【Dai et al., Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, CoRR, 2024】、DeepSeek-AI【DeepSeek-AI, Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, CoRR, 2024b】和DeepSeek-AI等人【DeepSeek-AI et al., Deepseek-v3 technical report, 2024】的做法,权重从标准差为0.006的正态分布中初始化。dots.llm1包含62个层,第一层使用普通的密集FFN,其余层使用MoE。

我们在预训练期间将最大序列长度设置为8K,并在11.2万亿个tokens上训练dots.llm1。遵循Hu等人【Hu et al., MiniCPM: Unveiling the potential of small language models with scalable training strategies, CoRR, 2024】的方法,我们使用了“预热-稳定-衰减”(warmup-stable-decay)的学习率调度策略。学习率在4000步内预热,然后在稳定训练阶段(包含10万亿个tokens的数据)稳定在$3 \times 10^{-4}$。我们逐步增加批量大小,从最初的6400万tokens,到6万亿tokens时增加到9600万tokens,最后在8.3万亿tokens时增加到1.28亿tokens。

在主训练阶段之后,该过程包括两个退火(annealing)阶段,总共包含1.2万亿个tokens的数据。
* 阶段1:我们训练1万亿个tokens,同时将学习率从$3 \times 10^{-4}$逐渐降低到$3 \times 10^{-5}$。在此退火阶段,我们将与推理和知识相关的数据比例显著提高到90%。
* 阶段2:我们继续训练2000亿个tokens,同时将学习率从$3 \times 10^{-5}$降低到$1 \times 10^{-5}$,并增加代码、数学和推理数据的比例。

4.3 长上下文扩展

我们在退火阶段后实施了上下文长度扩展。在此阶段,我们使用UtK策略【Tian et al., Untie the knots: An efficient data augmentation strategy for long-context pre-training in language models, 2024b】在1280亿个tokens上训练,同时保持学习率不变,将序列长度扩展到32K。UtK策略不是修改数据集,而是尝试将训练文档分块成更小的片段,然后训练模型从打乱的块中重建相关片段。通过学习解开这些打结的块,模型可以有效处理更长的输入序列,同时保持其在短上下文任务上的性能。


A4 实验环境


A4 实验结果

4.4 预训练评估

4.4.1 基准测试集

为了全面评估在中英文上预训练的dots.llm1模型,我们在多种领域的基准测试集上对其性能进行了评估。所有评估均使用vLLM框架进行。基准测试分类如下:

4.4.2 结果

与基线模型的比较:如表2所示,我们将dots.llm1与其他领先的开源基础模型在相同条件下进行了比较。结果表明,仅激活140亿参数的dots.llm1性能优于DeepSeek-V2,并与Qwen2.5 72B相当,使其成为迄今为止最经济、最强大的MoE模型之一。

表 2: dots.llm1与其他代表性开源基础模型的比较。所有模型都在相同条件下进行评估。结果表明,dots.llm1的性能优于DeepSeek-V2,同时与Qwen2.5 72B相当。

各领域性能细分(以Qwen2.5 72B为基准):
1. 语言理解:dots.llm1在中文语言理解基准测试上表现更优,这得益于其数据处理流程。
2. 知识任务:虽然在英语知识基准上的得分略低,但在中文知识任务上表现稳健。
3. 代码与数学:在HumanEval和CMath上得分更高。有趣的是,在数学方面,dots.llm1在零样本(zero-shot)设置下的表现优于少样本(few-shot)设置,提升超过4个点。

4.5 分析

5. 后训练

经过监督微调(SFT)后,模型被命名为dots.llm1.inst。

SFT数据与配置
* 数据混合:使用了约40万条指令微调实例,重点关注多语言(中英)多轮对话、知识问答、复杂指令遵循及数理编程推理。
* 微调过程:分为两个阶段。第一阶段,对40万实例进行上采样和多轮对话拼接,微调2个epoch。第二阶段,通过拒绝采样微调(RFT)并结合验证器系统,进一步增强数学和代码等特定领域的能力。

后训练评估
如表5所示,dots.llm1.inst与其它SOTA指令微调模型进行了比较,结果分析如下:
* 英语性能:在MMLU、DROP、GPQA等通用英语基准上表现稳定且全面,展现了扎实的知识基础和推理能力。
* 代码性能:具备有竞争力的代码能力,与Qwen2.5系列相当,但与Qwen3和DeepSeek-V3等前沿模型相比仍有提升空间。
* 数学性能:展现出卓越的数学推理能力。在AIME24上得分33.1,MATH500上得分84.8,优于Qwen2.5系列并接近SOTA水平。
* 中文性能:具有显著优势。在CLUEWSC上得分92.6,C-Eval上得分92.2,全面掌握中文知识,超越了包括DeepSeek-V3在内的所有对比模型。
* 对齐性能:在IFEval、AlpacaEval2和ArenaHard等基准上表现出竞争力,能准确理解并执行复杂指令,同时符合人类意图。

总体结论:dots.llm1.inst在仅激活140亿参数的情况下,在跨中英文通用任务、数学推理、代码生成和对齐基准方面表现出色,与Qwen2.5-72B-Instruct等更大模型相比具有强大的竞争优势,并在双语任务、数学和对齐能力上达到或超过Qwen3-32B的水平。

表 5: dots.llm1.inst与其他代表性指令微调模型的比较。所有模型都在我们的内部框架下于相同条件下进行评估。对于Qwen3系列,我们的评估未启用思考模式。


6. MoE 分析

专家负载的定义。我们检查了dots.llm1模型在Pile测试集上的专家负载。对于一个特定领域D,由MoE处理来自该领域的$N_D$个tokens,专家$E_i$的专家负载定义为:

其中$N_{E_i, D}$是来自领域D并被路由到专家$E_i$的token数量。专家负载表明专家$E_i$对领域D的专业化程度。值为1意味着该领域的所有token都被路由到该专家,而值为0则意味着该专家从未用于该领域。

专家专业化分析。dots.llm1在所有层级都显示出更强的专家专业化。在图5中,我们观察到对于特定领域(如DM Mathematics),某些专家的激活频率显著高于随机选择。相比之下,Wikipedia中的专家负载看起来更均衡,这可能是因为Wikipedia包含广泛的知识,像一个涵盖各种实体的知识图谱。对于包含各类数据的Books领域,专家负载也显得更为均衡。这突显了负载均衡功能按预期工作,使模型能够有效利用所有专家来处理通用数据。


图 5: Pile测试数据集上各层的专家负载热力图。


A5 结论

我们介绍了dots.llm1,一个高性价比的专家混合模型,在其规模下实现了最先进的性能。通过每个token仅激活部分参数,dots.llm1显著降低了训练成本,同时取得了与更大模型相当的结果。我们先进的数据处理流程生产了高质量的训练数据,而中间检查点的开源为模型学习动态提供了宝贵的见解。dots.llm1证明了高效的设计和高质量的数据可以持续扩展大型语言模型的能力边界。

作为未来工作的一部分,我们的目标是训练一个更强大的模型。为了在训练和推理效率之间达到最佳平衡,我们计划集成更高效的架构设计,如分组查询注意力(GQA)、多头潜在注意力(MLA)和线性注意力。此外,我们打算探索使用更稀疏的专家混合(MoE)层来提高计算效率。更重要的是,由于数据是预训练的基础,我们将深化对何为最佳训练数据的理解,并探索实现更像人类学习效率的方法,以最大化从每个训练样本中获取的知识。


A6 附录

B 超参数

预训练超参数配置。我们在表6中展示了预训练的超参数配置,并与其他相关模型进行了比较。

表 6: dots.llm1与近期模型的架构和预训练超参数对比。WSD = 权重稳定衰减 (Hu et al., 2024)。

C 网络数据整理

C.1 文档准备阶段

C.2 基于规则的处理阶段

C.3 基于模型的处理阶段