[论文解读] TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer

Zhen Qin, Li Dong+9 more|arXiv (Cornell University)|Jul 27, 2023

Topic Modeling被引用 5

一句话总结

TransNormerLLM 是一种基于线性注意力机制的大规模语言模型，通过集成先进的技术（如具有指数衰减的 LRPE、Lightning Attention 加速训练、门控机制、张量归一化以及稳健的推理算法），在准确率和效率方面均优于传统的 Transformer 模型。该模型在 385M 到 175B 参数规模下实现了训练速度提升超过 2 倍、显存使用减少 4 倍，同时保持了高性能表现。

ABSTRACT

We present TransNormerLLM, the first linear attention-based Large Language Model (LLM) that outperforms conventional softmax attention-based models in terms of both accuracy and efficiency. TransNormerLLM evolves from the previous linear attention architecture TransNormer by making advanced modifications that include positional embedding, linear attention acceleration, gating mechanisms, tensor normalization, and inference acceleration and stabilization. Specifically, we use LRPE together with an exponential decay to avoid attention dilution issues while allowing the model to retain global interactions between tokens. Additionally, we propose Lightning Attention, a cutting-edge technique that accelerates linear attention by more than twice in runtime and reduces memory usage by a remarkable four times. To further enhance the performance of TransNormer, we leverage a gating mechanism for smooth training and a new tensor normalization scheme to accelerate the model, resulting in an impressive acceleration of over $20\%$. Furthermore, we develop a robust inference algorithm that ensures numerical stability and consistent inference speed, regardless of the sequence length, showcasing superior efficiency during both training and inference stages. We also implement an efficient model parallel schema for TransNormerLLM, enabling seamless deployment on large-scale clusters and facilitating expansion to even more extensive models, i.e., LLMs with 175B parameters. We validate our model design through a series of ablations and train models with sizes of 385M, 1B, and 7B on our self-collected corpus. Benchmark results demonstrate that our models not only match the performance of state-of-the-art LLMs with Transformer but are also significantly faster. Code is released at: https://github.com/OpenNLPLab/TransnormerLLM.

研究动机与目标

• 开发一种在准确率和推理效率方面均超越传统基于 softmax 注意力机制的 Transformer 模型的大规模语言模型。
• 通过架构和优化创新，解决现有线性注意力机制存在的注意力稀释和实际性能不佳等问题。
• 通过高效的模型并行和内存优化技术，实现超大规模模型（最高达 175B 参数）的可扩展训练与部署。
• 确保在不同序列长度下具备数值稳定性和一致的推理速度，提升实际部署的可行性。
• 在包含超过 2 万亿 token 的大规模自收集语料上验证模型，证明其具备强大的泛化能力和可扩展性。

提出的方法

• 将 TransNormer 中的 DiagAttention 替换为线性注意力，实现线性复杂度并提升全局 token 交互能力。
• 引入 LRPE（学习型相对位置编码）并采用指数衰减策略，缓解注意力稀释问题，同时保留长距离依赖关系。
• 开发 Lightning Attention，一种新型 I/O 友好的技术，使线性注意力的训练速度提升 2 倍以上，显存使用减少 4 倍。
• 应用门控机制以稳定训练过程，并引入新型张量归一化方案，加速推理，实现超过 20% 的速度提升。
• 设计一种稳健的推理算法，确保在不同序列长度下均具备数值稳定性和恒定的推理速度。
• 实现一种高效的模型并行方案，支持最多 8 路并行，可在大规模集群上部署，并扩展至 175B 参数模型。

实验结果

研究问题

• RQ1 基于线性注意力机制的大规模语言模型能否在准确率和效率方面均优于标准 Transformer 模型？
• RQ2 如何在保留长序列建模能力的前提下，缓解线性注意力机制中的注意力稀释问题？
• RQ3 像 Lightning Attention 这类新型加速技术，能在多大程度上减少大规模 LLM 的训练时间和显存消耗？
• RQ4 统一的推理算法能否在不同序列长度下保持一致的推理速度和数值稳定性？
• RQ5 当在大规模数据集上训练并分布于大规模 GPU 集群时，基于线性注意力机制的大规模语言模型的可扩展性上限是什么？

主要发现

• 在相同条件下，TransNormerLLM-7B 的推理速度达到 4081.0 tokens/sec/GPU，优于 Transformer-7B 的 3362.7 tokens/sec/GPU。
• 在 8 路模型并行配置下，TransNormerLLM-7B 的 GPU 显存使用量降低至 24.1GB，相比 1× 配置减少了 62.3%。
• 在 175B 参数模型上，TransNormerLLM 在相同上下文长度下相比 Transformer 基线模型实现了 1.35 倍的相对训练速度提升。
• 由于其稳健的推理算法，该模型在不同序列长度下均保持一致的推理速度，确保了可预测的性能表现。
• Lightning Attention 实现了 4 倍显存减少和 2 倍以上的训练加速，显著提升了真实训练场景下的效率。
• 该模型在 6TB、2 万亿 token 的自收集语料上实现了具有竞争力的性能，达到或超越了同类规模模型的最先进水平。