首个线性注意力Transformer大模型：1750亿参数，速度、精度更优( 二 ) _大模型

改进二：门控机制
门控可以增强模型的性能并使训练过程平滑。研究者为 TransNormerLLM 使用了来自论文《Transformer quality in linear time》的 Flash 方法并在 token 混合中使用了门控式线性注意力（GLA）的结构。
为了进一步提升模型速度，他们还提出了 Simple GLU（SGLU），其去除了原始 GLU 结构的激活函数，因为门本身就能引入非线性。
改进三：张量归一化
研究者使用了 TransNormer 中引入的 NormAttention 。在 TransNormerLLM 中，他们使用一种新的简单归一化函数 SimpleRMSNorm（简写为 SRMSNorm）替换了 RMSNorm 。
整体结构
图 1 展示了 TransNormerLLM 的整体结构。

文章插图
图 1：新提出模型的整体架构
在该结构中，输入 X 的更新通过两个连续步骤完成：首先，其通过使用了 SRMSNorm 归一化的门控式线性注意力（GLA）模块。然后，再次通过使用了 SRMSNorm 归一化的简单门控式线性单元（SGLU）模块。这种整体架构有助于提升模型的性能表现。下方给出了这个整体流程的伪代码：

文章插图
训练优化
闪电注意力
为了加快注意力计算速度，研究者引入了闪电注意力（Lightning Attention）算法，这能让新提出的线性注意力更适合 IO（输入和输出）处理。
算法 1 展示了闪电注意力的前向通过的实现细节，算法 2 则是后向通过的。研究者表示，他们还有一个可以更快计算梯度的实现版本，这会在未来发布。

文章插图

文章插图
模型并行化
为了在计算机集群上分散所有模型参数、梯度和优化器状态张量，研究者使用了全分片数据并行（FSDP/Fully Sharded Data Parallel）。这种策略性分区方法可减少对每个 GPU 的内存占用，从而优化了内存利用率。为了进一步提高效率，他们使用了激活检查点（Activation Checkpointing），这可减少后向通过过程中缓存在内存中的激活数量。相反，当计算这些梯度时，这些梯度会被移除并重新计算。该技术有助于提升计算效率和节省资源。此外，为了在减少 GPU 内存消耗的同时加快计算速度，研究者还使用了自动混合精度（AMP）。
除了上述成果外，研究者还更进一步通过对线性 transformer 执行模型并行化而进行了系统工程优化，其灵感很大程度上来自于英伟达的 Megatron-LM 模型并行化，在传统的 Transformer 模型中，每个 transformer 层都有一个自注意力模块，其后跟着一个两层多层感知器（MLP）模块。当使用 Megatron-LM 模型并行性时，是在这两个模块上独立使用。类似地， TransNormerLLM 结构也是由两个主要模块构成：SGLU 和 GLA；这两者的模型并行化分开执行。
稳健推理
这让 TransNormerLLM 能以 RNN 的形式执行推理。算法 3 给出了这个过程的细节。但其中存在数值精度问题。

文章插图
为了避免这些问题，研究者提出了稳健推理算法，见算法 4 。

文章插图
原推理算法和稳健推理算法得到的结果是一样的。
语料库
研究者从互联网收集了大量可公开使用的文本，总大小超过 700TB 。收集到的数据经由他们的数据预处理程序进行处理，如图 2 所示，留下 6TB 的干净语料库，其中包含大约 2 万亿 token 。为了提供更好的透明度，帮助用户更好理解，他们对数据源进行了分门别类。表 2 给出了具体的类别情况。

文章插图
图 2：数据预处理流程

文章插图
表 2：语料库统计数据
实验
研究者在 Metaseq 框架中使用 PyTorch 和 Trition 实现了 TransNormerLLM 。模型的训练使用了 Adam 优化器， FSDP 也被用于高效地将模型扩展到 NVIDIA A100 80G 集群。他们也适当地使用了模型并行技术来优化性能。