论文分析｜高效长文本生成的技术与应用

Preface 前言

目前大模型公司很多在追求长文a本， 对算力需求极大，如何能够现实地处理该问题很重要。特别是随着Transformer模型尺寸和复杂性的增长，它们在训练期间的内存需求呈指数级增加。

语言模型训练的瓶颈在于显存占用非常大，这需要创新的解决方案来优化内存使用，同时保持性能。

本次将介绍一篇最新前沿论文，提出了一种名为MINI-SEQUENCE TRANSFORMER（MST）的技术，用于高效准确地训练极长序列的大型语言模型（LLMs），通过减少中间内存使用，实现了显著的内存节省，而不影响模型性能或训练收敛速度。MST方法通用、易于集成，并且支持分布式训练。

论文链接 🔗

https://arxiv.org/abs/2407.15892

贡献者：

Cheng Luo, Jiawei Zhao, Zhuoming Chen, 

Beidi Chen, Anima Anandkumar

Background 前沿背景

训练时必须需要在显存里存储以下内容：

1. 权重：模型的参数，包括所有层的权重矩阵，需要在训练前加载到显存中。

2. 正向传播时需要储存：

        a. 激活值（Activations）：计算并存储每层的激活，这些值需要保存以便于在反向传播时计算梯度。

        b. 中间值（Intermediate Values）——计算时每一层时都需要储存：在模型的不同层，特别是多头自注意力（Multi-Head Attention）层和多层感知器（MLP）层中，计算过程中会产生中间值，如Q（Query）、K（Key）、V（Value）张量，以及MLP层的中间线性变换结果。

        拆解开来主要是：

• Transformer 中的 Attention 层，计算 QKV 和注意力矩阵；
• Transformer 中的 MLP 层，将序列嵌入维度放大再缩小；

• LM Head （Language Modeling Head）将嵌入映射为 logits，之后计算 loss / 下一个 token。

3. 反向传播时需要计算并存储：梯度（Gradients）——在反向传播过程中计算得到的模型权重的梯度，用于更新模型参数，保持优化器的状态等等。

Innovation 技术创新

文章提出了MST方法：多个小序列（Mini-sequences）迭代处理

核心思路：通过将输入序列划分为多个小序列（mini-sequences），并迭代处理这些小序列，从而减少了中间内存的使用。

MST方法关注的是计算过程中间状态，因为Transformer 在计算过程中会产生非常巨大的中间状态：

• 直接计算 Attention 会产生 NxN 大小的注意力矩阵

• MLP 层也需要将嵌入维度放大四倍左右

• LM Head 的映射会产生几万到几十万维度的矩阵，因为要计算整个词表中每个单词的概率

但与此相对的是 GPU 的显存是有限的，而且 GPU 的显存分了不同层次，速度快的 HBM 非常小，搬运到低层次的 SRAM 很费时间。因此降低中间状态的大小就非常重要。所幸语言模型的计算在token 之间依赖比较小，因此可以并行操作，这就导向一个关键技巧，也就是分块处理。

假设有 16 个 token 要处理，每个 token 处理时需要X的内存，如果一次性处理，那就需要 16X 的内存；但如果分两块处理，那么整个过程最多只需要 8X 的内存；第一次计算完成后，相应内存就可以释放了。当然，不是分块越多越好，分块越多，各种搬运记录计算结果，汇总结果所需的时间也会越长。

FlashAttention 其实就在使用这个技巧，因为 Attention 的中间状态太大了，随着序列长度二次方增长谁也受不了。随着 Attention 的中间状态被 FlashAttention 和 Ulysses 打下来，我们自然就盯上其他中间状态了。

本文就在讨论分块解决 Llama 3 中 MLP 层和 LM Head 的中间状态。

分解过程都是类似的，都是 分解、计算、汇总 。设序列长度为 N, 嵌入维度为 d，序列嵌入 X∈R^(N*d)那么就将 X 沿着分成 M 块，每块计算结果，然后汇总。

MLP层：Llama3 的 MLP 层有三个线性层 Wgate,Wup,Wdown。前向计算过程为：

• 首先将 N*d 大小的序列矩阵通过 Wgate, Wup 放大为两个 N*l 的矩阵

• 然后两个矩阵逐元素点乘，最后通过 Wdown 缩小为原来的 N*d 大小。

汇总的过程就很简单，直接在序列维度拼接就可以。

LM Head 计算 ： 计算loss 的过程是将序列矩阵右乘一个矩阵映射到词表，也就是得到 NxV 的矩阵，然后算交叉熵。此时需要检查标签是否在块的范围内。汇总的过程是将 loss 加和再除以 M 取平均。

MST 没有改变 FLOP，但会增加 HBM 访问次数，标准 MLP 的访问次数为   Θ (Nd+NI+dI)，而 MST 的访问次数为 Θ (Nd+NI+dIM)。标准 LM Head 的次数是 Θ (Nd+NI+dV)，而 MST 下这个数字是 Θ (Nd+NI+dVM)。如果序列长度很短，中间状态占主导， dI,dV 占主导，MST会降低 GPU 显存的访问量。

Result 研究结果

论文通过实验评估了MST在不同模型上的性能，包括Llama2和Llama3，并展示了在不同序列长度和批处理大小下的训练时间和内存开销。

1. Loss 基本没有损失，Llama 3 8B 上最大序列长度是基本实现的12-20倍，激活重计算的1.8-4倍；

   

2. 计算速度有些许损失；
   
   

3. 峰值内存下降了很多。

结果分析：

• MST通用性和易用性：MST是完全通用的，并且可以轻松集成到现有的LLM训练框架中，只需最小的代码更改。

• 分布式训练支持：MST可以与DeepSpeed-Ulysses等技术一起工作，支持按GPU数量线性扩展序列长度。

• MST的局限性和未来的发展方向：例如将方法编译到CUDA以提高性能和内存效率，以及在小序列训练中可能出现的性能下降问题。

Thinking 后续思考

• 模型规模的增大和词表的扩展（大词表）：传统的串行计算方法可能无法满足效率需求，长序列下模型训练和推理会出现新的瓶颈，需要从硬件出发设计算法进行并行优化，如利用现代硬件（如GPU、TPU）的并行处理能力，可以显著加速模型的训练和推理过程。

• 分块计算：减少中间状态的思想非常简单但也非常好用，可以优化优化器的中间状态，值得继续推广

💡一点思考：长文本的能力是否值得长期探索？

1. 一方面，长文本本身的应用可能并没有那么多或者需求并不强烈，如debug 代码库里的代码，但是在多模态的发展上，处理long sequence确实十分重要。

2. 另一方面，长文本的能力在一定程度上反应智能的水平。人脑限制在20瓦，在这个限制下自然会放弃对一些记忆细节的把握，但是机器学习模型没有这方面的限制，因此没有必要过于参考人类生理进化的过程。eg：如果大模型能够同时记住所有学科的很多细节，那么很可能在科学研究上的突破性贡献要强于人类。

本周日，我们有幸邀请到了本文的作者之一Beidi Chen来与我们共同探讨高效长文本生成的技术与应用。欢迎大家来参加我们的活动。

活动介绍【奇绩潜空间】

奇绩潜空间第三季第二期邀请到的嘉宾是卡内基梅隆大学助理教授——陈贝迪，她将为我们带来关于 Efficient Long-context Generation 的分享。

为更好的帮助来到线下的创业者进行交流，我们全新升级【Co-Founder Matching】环节，你可以在这里更高效地与 GenAI 时代最有活力的创业者和研究者面对面交流（往期参与者大多是来自 BAT、字节、京东、美团等大厂的算法工程师、架构师、AI产品经理、创新战略部经理及大模型创业者们），本期我们也邀请到奇绩校友「明之几何」创始人甘少多一起参与关于机器学习与创业实践的讨论，你将有机在这里找到志同道合、有共同创业梦想的小伙伴。

此外，我们创建了「高效长文本生成技术交流群」，欢迎加企微进群交流互动。

【分享嘉宾介绍】

主要研究方向：

Efficient AI，特别关注基于目前硬件设计及优化算法，以现实问题为目标加速大规模机器学习系统。代表作包括H2O, Streaming LLM, Monarch, Pixelated Butterfly等。

学术成就：

陈教授在机器学习领域的成就非常显著，她在2022年和2023年的ICML会议上分别发表了多篇有影响力的论文，包括FlexGen、StreamingLLM和DejaVu等，这些研究成果在大型语言模型的高效推理和训练方面取得了显著进展。

她还曾获得ICML 2022的杰出论文奖，是对Monarch矩阵设计的贡献的认可。

【Co-Founder Matching 嘉宾介绍】

甘少多，明之几何（Geoming AI）CEO，苏黎世联邦理工学院（ETH Zurich）博士学位，剑桥大学访问学者。研究方向聚焦大规模深度学习系统、大模型训推优化和 Agent 系统框架研发，其学术成果成功落地于微软 Deepspeed、PyTorch Lightning、快手等公司的业务。明之几何致力于实现人力与算力协同工作的最佳配比，让真正好用的 AI Agent 成为每个企业的标配。目前已帮助用户在营销、客服、社媒运营、企业知识等场景实现大模型技术落地。

【活动时间】

北京时间：8月11日 11:00-13:30

【活动地点】

北京线下 / 上海线下 / 线上 均有场次

【活动报名】

扫描下方二维码，填写信息完成报名
（找到奇绩校友帮忙推荐，可增加报名成功率）

（提醒：报名的小伙伴请联系小助手获取往期潜空间嘉宾分享笔记）