Scaling TransNormer to 175 Billion Parameters

文章标题：将TransNormer扩展至1750亿参数
作者/机构：Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Fei Yuan, Xiao Luo, Yu Qiao, Yiran Zhong (1上海人工智能实验室, 2OpenNLPLab)

A1 主要贡献

本文旨在解决传统Transformer模型中softmax注意力机制带来的二次方时间复杂度问题，该问题限制了模型在训练和推理阶段的可扩展性和效率。尽管已有研究提出了多种高效的序列建模方法以将复杂度降至线性，但它们在语言建模任务上性能不佳或在实际应用中未展现出速度优势，因此难以应用于大型语言模型（LLM）。

为应对这些挑战，本文提出了TransNormerLLM，这是首个在准确性和效率上均超越传统softmax注意力模型的、基于线性注意力的LLM。该模型基于先前的线性注意力架构TransNormer【32，The devil in linear transformer，2022，EMNLP】，并进行了一系列关键改进：

  1. 位置编码：采用带有指数衰减的线性化相对位置编码（LRPE），以解决注意力稀释问题，同时保留了Token之间的全局交互能力。
  2. 线性注意力加速：引入闪电注意力（Lightning Attention）技术，将线性注意力的运行速度提升两倍以上，并将内存使用量减少四倍。
  3. 门控机制与张量归一化：利用门控机制平滑训练过程，并采用一种新的张量归一化方案（SimpleRMSNorm），使模型整体速度提升超过20%。
  4. 鲁棒的推理算法：开发了一种数值稳定且推理速度与序列长度无关的推理算法，提升了训练和推理效率。

为了验证模型的有效性，研究团队收集并构建了一个超过6TB、包含超过2万亿Token的大规模语料库，并实施了一种新的自清洁（self-cleaning）策略来保证数据质量。基于此，团队将TransNormerLLM从385M扩展至175B参数（如表1所示），并通过大量实验证明其优于基于softmax注意力的模型，并具有更快的训练和推理速度。

最后，作者计划开源预训练模型，以推动社区在高效Transformer结构方面的研究进展。

表 1：TransNormerLLM 模型变体

A3 相关工作

Transformer-based LLMs

近年来大型语言模型的发展。近年来，大型语言模型（LLMs）领域取得了显著进展。遵循缩放定律【20，Scaling laws for neural language models，2020】，各种超过千亿参数的LLM相继问世，如GPT-3【3，Language models are few-shot learners，2020，NeurIPS】、Gopher【38，Scaling language models: Methods, analysis & insights from training gopher，2022】、PaLM【6，Palm: Scaling language modeling with pathways，2022】、GLM【11，Glm: General language model pretraining with autoregressive blank infilling，2022】等。同时，也出现了像Galactica【41，Galactica: A large language model for science，2022】这样专注于科学等特定领域的模型。一个重要的发展是拥有700亿参数的Chinchilla【17，Training compute-optimal large language models，2022】，它重新定义了缩放定律，强调Token数量而非模型权重的重要性。此外，LLaMA【43，Llama: Open and efficient foundation language models，2023，arXiv】因其优异的性能和开源可用性而引起了广泛关注。关于LLM的讨论还包括开源与闭源模型之间的动态关系。以BLOOM【45，Bloom: A 176b-parameter open-access multilingual language model，2023】、OPT【46，Opt: Open pre-trained transformer language models，2022】、LLaMA【43，Llama: Open and efficient foundation language models，2023，arXiv】和Falcon【28，The refinedweb dataset for falcon llm: Outperforming curated corpora with web data, and web data only，2023】为代表的开源模型正在崛起，与GPT-3【3，Language models are few-shot learners，2020，NeurIPS】和Chinchilla【17，Training compute-optimal large language models，2022】等闭源模型展开竞争。为了加速训练，稀疏注意力（Sparse Attention）【4，Generating long sequences with sparse transformers，2019；1，Longformer: The long-document transformer，2020】被引入，但在大型模型中，只有GPT-3采用了它【3，Language models are few-shot learners，2020，NeurIPS；39，What language model to train if you have one million gpu hours?，2022】。

Non-Transformer-based LLMs Candidates

Transformer替代架构的探索。尽管基于Transformer的大型模型在研究界大量涌现，但近期部分工作优先解决了其平方时空复杂度问题。这一焦点推动了一系列不同于传统Transformer结构的模型架构的探索和发展。其中，四个重要的竞争者——线性Transformer、状态空间模型、长卷积和线性循环网络——在建模长序列时作为自注意力（SA）模块的替代品显示出有希望的结果。这些替代方案因其更优的渐进时间复杂度和有竞争力的性能而受到青睐。

线性Transformer。线性Transformer将Softmax注意力分解为隐藏表示内积的形式，这使其能够使用“右乘技巧”（Right Product Trick），即先计算键（keys）和值（values）的乘积，以避免生成二次方的n×n矩阵。不同的方法使用不同的隐藏表示。例如，【21，Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】使用1+elu作为激活函数，【33，cosformer: Rethinking softmax in attention，2022，ICLR】使用cos函数来近似softmax的性质，而【22，Rethinking positional encoding in language pre-training，2021，ICLR；48，Linear complexity randomized self-attention mechanism，2022，ICML；49，Efficient attention via control variates，2023，ICLR】则通过理论方法来近似softmax。尽管其理论复杂度为$O(nd^2)$，但由于在因果注意力中需要进行累加求和（cumsum）操作，线性注意力的实际计算效率变得相当低【18，Transformer quality in linear time，2022，arXiv】。另一方面，大多数线性Transformer与传统Transformer相比，仍然存在一定的性能差距【21，Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML；25，Neural architecture search on efficient transformers and beyond，2022，arXiv】。

状态空间模型（State Space Model）。状态空间模型基于状态空间方程进行序列建模【15，Efficiently modeling long sequences with structured state spaces，2022，ICLR】，通过特殊的初始化【13，Hippo: Recurrent memory with optimal polynomial projections，2020；14，On the parameterization and initialization of diagonal state space models，2022，NeurIPS】、对角化假设【16，Diagonal state spaces are as effective as structured state spaces，2022，NeurIPS】和一些技术【9，Hungry hungry hippos: Towards language modeling with state space models，2022，CoRR】实现了与Transformer相当的性能。另一方面，由于状态空间方程的特性，它能以常数复杂度进行推理【15，Efficiently modeling long sequences with structured state spaces，2022，ICLR】。

长卷积（Long Convolution）。长卷积模型【31，Toeplitz neural network for sequence modeling，2023，ICLR；12，Simple hardware-efficient long convolutions for sequence modeling，2023，CoRR】使用的核大小等于输入序列长度，从而比传统卷积能捕获更广的上下文。这些模型的训练利用了高效的$O(n \log n)$快速傅里叶变换（FFT）算法。然而，长卷积也带来了一些挑战，例如因果卷积推理需要缓存所有历史计算，类似于自注意力（SA）的键值（KV）缓存。处理长序列的内存需求，加上比RNN更高的推理复杂度，使其在处理长序列方面不那么理想。

线性RNN（Linear RNN）。相比之下，线性RNN在长序列建模中是替代自注意力（SA）的更合适选择。一个著名的例子是RWKV【29，Rwkv: Reinventing rnns for the transformer era，22023】模型，这是一个基于线性RNN的大型语言模型，与同等规模的GPT模型相比，展现出了有竞争力的性能。

A2 方法细节

架构改进

在本节中，我们深入研究了网络的每个模块，并提出了一些改进，以在效率和性能之间达到最佳平衡。下面，我们概述了每个模块的关键设计以及每次更改背后的灵感。我们在第5节中对每个修改都进行了消融研究。

改进1：位置编码

采用带指数衰减的LRPE-d。在TransNormer中，底层网络使用DiagAttention来避免稀释问题。然而，这导致Token之间缺乏全局交互。在TransNormerLLM中，我们利用带有指数衰减【30，Train short, test long: Attention with linear biases enables input length extrapolation，2022，ICLR；31，Toeplitz neural network for sequence modeling，2023，ICLR】的LRPE【34，Linearized relative positional encoding，2023】来解决这个问题，从而在底层保留了完整的注意力。我们位置编码的整体表达式如下：


我们称之为LRPE-d——带有指数衰减的线性化相对位置编码。与原始的LRPE类似，我们将$\theta$设置为可学习的。

与线性注意力的兼容性及衰减率设计。需要注意的是，这种位置编码与线性注意力完全兼容，因为它可以分别针对$s$和$t$进行分解。第$l$层（假设总共有$L$层）第$h$个注意力头（假设总共有$H$个头）的$\lambda$值由以下公式给出：


这里，$8^h$对应第$h$个头的衰减率，而$1-\frac{l}{L}$对应第$l$层的衰减率。$1-\frac{l}{L}$这一项确保了底层网络的理论感受野（Theoretical Receptive Fields, TRF）【35，Exploring transformer extrapolation，2023】比高层网络小，这与TransNormer的设计动机一致。值得注意的是，最后一层的衰减率设置为1，允许每个Token关注全局信息。我们选择$\lambda$为不可学习的，因为根据经验发现，当$\lambda$可学习时，梯度会变得不稳定，导致出现NaN值。

改进2：门控机制

采用门控线性注意力（GLA）。门控（Gate）可以增强模型的性能并平滑训练过程。在TransNormerLLM中，我们采用了Flash【18，Transformer quality in linear time，2022，arXiv】的方法，在Token混合（token mixing）中使用了门控线性注意力（Gated Linear Attention, GLA）的结构：


其中：

我们选择$\phi$为$1+elu$。我们发现$\phi$的具体选择对结果影响很小，如表8所示。

提出简单门控线性单元（SGLU）。为了进一步加速模型，我们提出了简单门控线性单元（Simple GLU, SGLU），它从原始GLU结构中移除了激活函数，因为门控本身就可以引入非线性。因此，我们的通道混合（channel mixing）变为：


我们通过经验发现，不使用激活函数不会导致任何性能损失，如表9所示。

改进3：张量归一化

采用NormAttention并简化为线性形式。我们采用了TransNormer【32，The devil in linear transformer，2022，EMNLP】中引入的NormAttention，具体如下：


这种注意力机制去除了softmax和缩放操作。此外，它可以通过右乘法转换为线性注意力：

这种线性形式允许以$O(nd^2)$的复杂度进行循环预测，使其在推理过程中非常高效。具体来说，我们只以循环方式更新$K^TV$，而无需计算完整的注意力矩阵。在TransNormerLLM中，我们将RMSNorm替换为一个新的简单归一化函数，称为SimpleRMSNorm，简写为SRMSNorm：

我们根据经验发现，使用SRMSNorm不会导致任何性能损失，这在表10的消融研究中得到了证明。

整体结构

模型整体结构流程。整体结构如图1所示。在此结构中，输入$X$通过两个连续的步骤进行更新：首先，它经过应用了SimpleRMSNorm（SRMSNorm）归一化的门控线性注意力（GLA）。然后，它再次经过应用了SRMSNorm归一化的简单门控线性单元（SGLU）。这种整体架构有助于基于PreNorm方法提高模型性能。整体过程的伪代码如下：



图 1：所提出模型的架构概览。每个Transformer块由一个用于通道混合的简单门控线性单元（SGLU）和一个用于Token混合的门控线性注意力组成。我们对两个模块都应用了Pre-Norm。

训练优化

闪电注意力（Lightning Attention）

为解决因果预测中的并行性问题而提出Lightning Attention。线性注意力的结构允许通过右乘法以$O(nd^2)$的复杂度高效计算注意力。然而，对于因果预测，右乘法效率不高，因为它需要进行累加求和（cumsum）计算【18，Transformer quality in linear time，2022，arXiv】，这阻碍了训练的并行性。

Lightning Attention的IO友好设计。因此，在训练期间，我们继续使用传统的左乘法版本。为了加速注意力计算，我们引入了受【7，FlashAttention-2: Faster attention with better parallelism and work partitioning，2023；8，Flashattention: Fast and memory-efficient exact attention with io-awareness，2022，NeurIPS】启发的闪电注意力（Lightning Attention）算法，使我们的线性注意力对IO（输入/输出）更加友好。它计算以下表达式：


这里，$M$是注意力掩码，用于实现下三角因果掩码和位置编码。在闪电注意力中，我们将输入$Q, K, V$分割成块，将它们从慢速的HBM（高带宽内存）加载到快速的SRAM（静态随机存取存储器）中，然后计算这些块的注意力输出，最后累加得到最终结果。通过避免在慢速HBM上的操作，计算速度得到了加速。闪电注意力的实现细节分别在算法1（前向传播）和算法2（后向传播）中展示。请注意，存在一种更快的梯度计算实现方式，将在未来发布。

算法1 闪电注意力前向传播

算法2 闪电注意力后向传播

模型并行

大规模预训练面临的系统挑战。为了有效进行TransNormerLLM的大规模预训练，我们从多个角度集中精力进行了系统优化。需要强调的是，“大规模”一词指两个方面：首先，模型参数量巨大；其次，计算基础设施拥有大量的GPU节点。这些因素导致了各种实际挑战，包括GPU内存限制、计算效率下降、通信速度慢等。解决这些问题对于成功实施大规模预训练过程至关重要。

采用的系统优化策略。在我们的研究中，我们使用了完全分片数据并行（Fully Sharded Data Parallel, FSDP）【47，Pytorch fsdp: experiences on scaling fully sharded data parallel，2023，arXiv】方法，将所有模型参数、梯度和优化器状态张量分布到整个集群中。这种策略性分区减少了每个独立GPU上的内存占用，优化了内存利用率。为了进一步提高效率，我们使用了激活检查点（Activation Checkpointing）【40，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】，它减少了在反向传播过程中缓存在内存中的激活数量。相反，在计算梯度时，这些激活被移除并重新计算。这项技术有助于更高效的计算和资源节约。此外，我们使用自动混合精度（Automatic Mixed Precision, AMP）【26，Mixed precision training，2017，arXiv】来减少GPU内存消耗，同时加速计算速度。值得注意的是，在我们的整个实验过程中，我们使用了BFloat16【19，A study of bfloat16 for deep learning training，2019，arXiv】，因为它被观察到对大型TransNormerLLM模型的训练稳定性有益。

针对SGLU和GLA的模型并行化。除了上述努力，我们通过在线性Transformer上执行模型并行化，进一步深入系统工程优化，这在很大程度上受到了Nvidia的MegatronLM模型并行化【40，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】的启发。在传统Transformer模型中，每个Transformer层由一个自注意力块和一个两层多层感知机（MLP）块组成。当使用Megatron-LM模型并行化时，它会独立地应用于这两个块。类似地，在TransNormerLLM结构中，它也由两个主要块SGLU和GLA组成，我们分别对它们执行模型并行化。详细的模型并行化策略如下所示。

SGLU上的模型并行。回顾公式5中的SGLU结构：


其模型并行版本如下：

这将权重矩阵$W_v$和$W_u$沿其列进行拆分，并得到一个同样沿列拆分的输出矩阵。然后，分离的输出$[O_1, O_2]$与另一个沿行拆分的矩阵相乘，如下所示：

整个过程将SGLU块中的三个GEMM（通用矩阵乘法）操作分布到多个GPU上，并且在前向和后向传播中分别只引入了一次all-reduce操作。

GLA上的模型并行。回顾公式3和4中的GLA块，GLA上的模型并行如下：


其中：



注意，在实现中，我们使用合并的QKVU投影来提高计算效率。得到的拆分输出矩阵$[O_1, O_2]$再次与一个沿其列轴拆分的权重矩阵相乘，这与公式14类似。

鲁棒的推理

原始RNN式推理的数值精度问题。在本节中，我们讨论TransNormerLLM中的推理问题。需要注意的是，公式1可以分解为以下形式：


这使得TransNormerLLM能够以RNN的形式进行推理【21，Transformers are rnns: Fast autoregressive transformers with linear attention，2020，ICML】。该过程的详细信息如算法3所示。然而，值得注意的是，由于$\lambda < 1$，会导致：

这会引发数值精度问题。

算法3 原始推理算法
输入：$q_t, k_t, v_t, t = 1, \dots, n$；
输出：$o_t, t = 1, \dots, n$；
初始化：$[kv]_0 = 0$；
for $t = 1, \dots, n$ do
$[kv]_t = [kv]_{t-1} + k_t \lambda^{-t} v_t^T$,
$o_t = q_t \lambda^t [kv]_t$.
end for

鲁棒推理算法及其等价性证明。为了避免这些问题，我们在算法4中提出了一个鲁棒的推理算法。由于$\|q_s \exp^{i\theta_s}\| = \|q_s\|$，$\|k_t \exp^{i\theta_t}\| = \|k_t\|$，为简化起见，我们在后续讨论中将省略LRPE【34，Linearized relative positional encoding，2023】，仅考虑$a_{st} = q_s^T k_t \lambda^{s-t}$。
算法4 鲁棒推理算法
输入：$q_t, k_t, v_t, t = 1, \dots, n$；
输出：$o_t, t = 1, \dots, n$；
初始化：$[\widetilde{kv}]_0 = 0$；
for $t = 1, \dots, n$ do
$[\widetilde{kv}]_t = \lambda [\widetilde{kv}]_{t-1} + k_t v_t^T$,
$o_t = q_t [\widetilde{kv}]_t$.
end for

我们将使用数学归纳法来证明：


因此，原始推理算法和鲁棒推理算法产生相同的结果。

基础情况 (n = 1):


因此，基础情况成立。让我们假设该命题对$n=m-1$成立，即$[kv]_{m-1} = \lambda^{-(m-1)}[\widetilde{kv}]_{m-1}$。那么，当$n=m$时：

该命题成立。因此，通过数学归纳法，该命题对所有$n \ge 1$都成立。

A3 语料库

语料库的构建与构成。我们从互联网上收集了大量可公开访问的文本，总大小超过700TB。收集到的数据经过我们的数据预处理流程（如图2所示）处理后，留下一个6TB的清洁语料库，大约包含2万亿个Token。我们对数据源进行了分类，以提供更好的透明度和理解。这些类别的具体细节在表2中列出。
表 2：我们的语料库统计数据。对于每个类别，我们列出了在2万亿个Token上进行训练时，在该子集上执行的epoch数，以及Token数量和磁盘大小。右侧的表格则根据语言分布进行排列。

数据预处理

数据预处理流程。我们的数据预处理流程包括三个步骤：1) 基于规则的过滤，2) 去重，以及 3) 自清洁方案。在被添加到训练语料库之前，清洁后的语料库需要经过人工评估。


图 2：数据预处理流程。收集的数据经过基于规则的过滤和去重，然后是我们自创的数据自清洁策略：基于模型的过滤、人工评估和评估模型。经过上述循环的多次迭代，我们获得了大约2T Token的高质量训练数据。

基于规则的过滤。我们用于过滤收集数据的规则如下：
* 移除HTML标签和URL：我们流程的第一步是消除文本中的HTML标签和网页URL。这通过正则表达式技术实现，识别并移除这些模式，确保语言模型专注于有意义的文本内容。
* 消除无用或异常字符串：随后，清洁后的数据集会进行第二层精炼，识别并删除没有价值的字符串，如异常字符串或乱码文本。这个过程依赖于预定义的规则，将某些字符串模式归类为非贡献元素。
* 标点符号去重：我们处理数据中冗余标点符号的问题。多个连续的标点符号在训练模型时会扭曲句子的自然流畅度和结构。我们采用一个基于规则的系统，将这些重复项削减为每个标点符号的单个实例。
* 处理特殊字符：识别出不常属于该语言文本语料库的异常或特殊字符，并将其移除或替换为标准化表示。
* 数字标准化：数字在不同文本中可能以各种格式呈现。这些数字被标准化为一种通用格式以保持一致性。
* 保留Markdown/LaTeX格式：在移除非文本元素时，对Markdown和LaTeX格式的文本做了例外处理。鉴于它们的结构化性质以及在学术界和文档中的普遍使用，保留这些格式可以增强模型理解和生成类似格式文本的能力。

去重。为了确保我们数据的独特性并避免过拟合的风险，我们使用MinHash和局部敏感哈希（LSH）算法，在文档或行级别采用高效的去重策略。MinHash和LSH的结合确保了去重过程中计算效率和准确性之间的平衡，为数据去重和文本水印移除提供了强大的机制。

自清洁方案。我们的数据自清洁过程涉及以下三个步骤的迭代循环，以持续提炼和增强我们数据集的质量。使用基于模型的数据过滤器的一个问题是，过滤后的数据将与评估模型的分布相似，这可能会对训练数据的多样性产生重大影响。假设大部分预处理数据是高质量的，我们可以在整个预处理数据集上训练一个评估模型，该模型将自动平滑数据流形分布，并剔除低质量数据，同时保留大部分的多样性。

自清洁方案展开如下：
* 评估模型：我们在预处理的语料库上训练一个385M的模型，作为数据质量过滤器。
* 基于模型的数据过滤：我们使用评估模型通过困惑度（perplexity）评估每条数据。只有得分高于某个阈值的数据才被保留到下一步。低质量数据在此阶段被剔除。
* 人工评估：我们对过滤后的数据进行小部分抽样，并手动评估其质量。

这些步骤以循环方式重复，每次迭代都提高了数据的整体质量，确保最终模型在相关、高质量的文本上进行训练。这个自清洁过程为维护数据完整性提供了强大的机制，从而增强了最终语言模型的性能。

Tokenization

BPE分词与中文支持。我们使用字节对编码（Byte-Pair Encoding, BPE）算法对数据进行分词。值得注意的是，为了增强与中文内容的兼容性，我们将大量常见和不常见的汉字纳入了我们的词汇表。在词汇表项不存在于字典中的情况下，单词会被分解为其组成的UTF-8字符。这一策略确保了在模型训练期间对多样化语言输入的全面覆盖和灵活性。

A4 实验环境

  • 软件配置：
    • 实现框架：PyTorch【27，Pytorch: An imperative style, high-performance deep learning library，2019，NeurIPS】 和 Trition【42，Triton: an intermediate language and compiler for tiled neural network computations，2019】
    • 基础框架：Metaseq【46，Opt: Open pre-trained transformer language models，2022】
    • 优化器：Adam【23，Adam: A method for stochastic optimization，2017】
    • 并行策略：FSDP，模型并行
  • 硬件配置：
    • GPU：NVIDIA A100 80G 集群
  • 数据集：
    • 对于消融研究，使用了一个从自建语料库中抽样的300B Token的子集。
  • 评估指标：
    • 为了减少损失（Losses）和困惑度（PPLs）的波动，采用最后1000次迭代的平均值作为最终指标。

A5 实验结果

架构消融实验

  • Transformer vs TransNormerLLM (表3)

    • 实验内容：在不同模型规模下比较TransNormerLLM和标准Transformer的性能。
    • 实验结果：在相同配置下，TransNormerLLM表现更优。在385M规模下，性能高出5%；在1B规模下，优势扩大到9%。
    • 结论：TransNormerLLM架构优于标准Transformer。

    表 3：Transformer 与 TransNormerLLM 的对比。在385M和1B规模下，TransNormerLLM在相同配置下分别比Transformer性能高出5%和9%。

  • TransNormer vs TransNormerLLM (表4)

    • 实验内容：比较原始TransNormer与改进后的TransNormerLLM。
    • 实验结果：TransNormerLLM性能分别提升了2%和1%，且速度显著加快。
    • 结论：架构改进有效提升了性能和效率。

    表 4：TransNormer 与 TransNormerLLM 的对比。TransNormerLLM 取得了最好的结果。

  • 位置编码 (表5, 表6)

    • 实验内容：比较LRPE-d、APE（绝对位置编码）、LRPE和仅指数衰减（Exp-Decay）的效果，并对衰减温度项进行消融。
    • 实验结果：LRPE-d（LRPE + Exp-Decay）的组合效果最佳，比LRPE方法提升了2%。添加衰减温度项 $1 - \frac{l}{L}$ 能够降低困惑度。
    • 结论：LRPE-d位置编码方案及其衰减温度设计是有效的。

    表 5：位置编码。LRPE+LRPE-d的组合带来了最优的结果。


    表 6：衰减温度的消融实验。带有衰减温度的结果被证明更优。

  • 门控机制 (表7)

    • 实验内容：评估门控机制对模型性能的影响。
    • 实验结果：启用门控机制后，损失值从2.263降低到2.247。
    • 结论：门控机制能有效提升模型性能。

    表 7：门控机制的消融实验。带有门控的性能被证明更优。

  • GLA与GLU激活函数 (表8, 表9)

    • 实验内容：分别测试GLA中不同激活函数的影响，以及在GLU中移除激活函数的影响。
    • 实验结果：GLA中的激活函数选择对结果影响不大，最终选择$1+elu$。移除GLU中的激活函数（即SGLU）对最终结果影响可以忽略。
    • 结论：可以简化GLU结构为SGLU以提升效率而不损失性能。

    表 8：GLA激活函数的消融实验。不同激活函数得到的结果几乎相同。


    表 9：GLU激活函数的消融实验。排除激活函数对结果没有负面影响。

  • 归一化函数 (表10, 图3)

    • 实验内容：比较SRMSNorm, RMSNorm, LayerNorm的性能，并评估SRMSNorm的Triton实现速度。
    • 实验结果：三种归一化方法在TransNormerLLM上性能差异极小。然而，自研的SRMSNorm Triton实现在较大维度下比PyTorch实现有显著的速度提升。
    • 结论：SRMSNorm在性能上无损，且通过优化实现可获得更高计算速度。

    表 10：归一化函数。以下归一化函数之间的结果偏差极小。



    图 3：SRMSNorm实现的性能评估。上图展示了在固定特征维度为3072时，不同序列长度下前向传播（左）和后向传播（右）的运行时间对比。下图展示了在固定序列长度为4096时，不同特征维度下的运行时间对比。

  • 闪电注意力（Lightning Attention） (图4)

    • 实验内容：比较Lightning Attention与基线（PyTorch实现的NormAttention）的速度和内存占用。
    • 实验结果：在训练（前向+后向）中，Lightning Attention的速度至少是基线的2倍。内存占用随序列长度线性增长，在序列长度为8192时比基线节省高达4倍内存。
    • 结论：Lightning Attention实现了卓越的效率。


    图 4：线性注意力与闪电注意力的内存和速度对比。左图：在固定特征维度为2048时，不同序列长度下前向+后向传播的运行时间（毫秒）。右图：在固定特征维度为2048时，不同序列长度下前向+后向传播的内存占用。

  • 推理效率 (图5)

    • 实验内容：评估TransNormerLLM在不同序列长度下的推理时间和内存占用。
    • 实验结果：随着序列长度增加，TransNormerLLM的推理时间和内存占用保持恒定。
    • 结论：得益于RNN式的推理模式，TransNormerLLM在长序列推理上具有高效和低资源消耗的优势。


    图 5：推理时间与内存占用。左图：在不同序列长度下，前向和后向传播的推理运行时间（毫秒）。右图：不同序列长度下的推理内存消耗。值得注意的是，随着序列长度的增加，TransNormerLLM表现出恒定的推理时间和内存占用。

系统优化实验

  • 模型并行 (表11)

    • 实验内容：在单个DGX节点（8x A100）上，对7B参数的TransNormerLLM和Transformer（使用Flash Attention）进行模型并行性能测试。
    • 实验结果：模型并行显著节省内存，当并行度为8时，TransNormerLLM-7B的单GPU内存消耗仅为24.1GB，比并行度为1时减少62.3%。虽然内存节省显著，但训练速度仅轻微下降。TransNormerLLM的速度始终大幅领先于Transformer。
    • 结论：模型并行策略能有效降低内存压力，使得在有限资源下训练大模型成为可能，且TransNormerLLM保持了速度优势。

    表 11：模型并行性能。我们在单个带NVLink的A100节点上比较了带Flash Attention的Transformer-7B和带Lightning Attention的TransNormerLLM-7B的模型并行性能。所有实验均使用批大小为2，上下文长度为2048。

  • 模型规模与上下文长度压力测试

    • 模型规模 (表12)：

      • 实验内容：在大型A100集群上，训练不同规模的Transformer和TransNormerLLM模型，固定上下文长度为2048，并使用最大化吞吐量的批大小。
      • 实验结果：在各种模型规模下，TransNormerLLM的计算速度始终优于对应的Transformer模型。
      • 结论：TransNormerLLM的线性计算复杂度优势在不同模型规模上都得以验证。

      表 12：不同规模模型的训练效率。为进行比较，我们固定上下文长度为2048，并为Transformer和TransNormerLLM增加批大小以在不遇到内存溢出问题的情况下达到其最大速度。


      - 上下文长度 (表13)：
      - 实验内容：在64个A100 GPU上，固定批大小为1，测试不同参数规模的Transformer和TransNormerLLM模型能支持的最大上下文长度。
      - 实验结果：在同等计算资源下，TransNormerLLM能够训练更长的上下文长度，并且计算速度更快。
      - 结论：TransNormerLLM在长上下文训练方面具有显著优势。

      表 13：训练Transformer和TransNormerLLM的最大上下文长度。我们在64个A100 80G GPU上比较了Transformer和TransNormerLLM在不同模型大小下的最大上下文长度。所有实验均使用批大小为1。

A6 结论

本文介绍了TransNormerLLM，一个为大型语言模型量身定制的改进版TransNormer。我们的TransNormerLLM在准确性和效率上都持续优于Transformer，并且能有效扩展至1750亿参数。大量的消融实验证明了我们在位置编码、门控机制、激活函数、归一化函数和闪电注意力方面的修改和创新的有效性。为了支持TransNormerLLM的训练，我们收集了一个超过6TB、包含超过2万亿Token的大规模语料库，并利用一种新颖的自清洁策略来确保数据质量和相关性。我们将发布我们的预训练模型，以促进社区在高效LLM领域的进步。