论文信息
➤ 论文地址
🔗https://arxiv.org/abs/2406.03488
➤ 开源地址
🔗https://github.com/thunlp/Seq1F1B
Seq1F1B 显存占用和训练吞吐
实验设置
显存占用对比
2.7B、7B、13B和30B显存占用对比（左右滑动）
研究背景
主要创新
我们提出了序列级别的 1F1B 流水线技术，解决了以往 1F1B 流水线在长文本训练中面临的显存瓶颈。 提出了根据计算量对序列切分的负载均衡技术，使得序列切分后的流水线调度更加高效。 通过实验证明我们方法的有效性，大大降低了 1F1B 的显存占用，并实现了更高的吞吐。 基于 Nvidia/Megatron-LM 进行开发，保证可复现性和兼容性。
具体方法
➤ 1F1B 流水的原理
传统的 1F1B 流水线并行包含三个阶段：Warmup，Steady，Cooling down。
1F1B时序图
上图为 1F1B 的时序图，在 Warmup 阶段，前面的流水线阶段要比后面的多一个前向计算，而最后面的阶段在完成一个前向计算后，要马上进行该前向计算对应的反向计算，并在 Steady 阶段进行前向和反向的交替计算，这一过程被称为 one-forward-one-backward，也是 1F1B 名字的由来。长文本训练中，1F1B 每个批次的前向计算和反向计算时间会变长，因此流水线气泡（图中的白色部分，即流水线并行中设备的空闲时间）也会随之变长，这也就会导致更长的流水线空占时间，最后的结果在时序图上表现为总时间变长。
➤ Seq1F1B 训练技术解析
Seq1F1B时序图
我们提出的 Seq1F1B训练技术，解决了 Attention 模块中序列之间的依赖问题的同时，通过序列切分进一步减少了气泡时间。由上图可以看到，Sequence 2的前向计算在Sequence 1后面，而其 Sequence 2 的反向计算在 Sequence 1 的前面。通过 Seq1F1B 设计，流水线的负载变得更加紧凑，不同流水线阶段间的空置时间变得更少了。除此之外，由于前向计算和反向计算的单位变得更小了，因此在前向阶段和反向阶段中间，模型需要存储的激活值也变小了，极大降低了激活值带来的显存占用。
尽管更细粒度的流水线调度使得流水线气泡减少，但是由于模型的自回归特性，如果根据序列长度直接均分序列，不同序列之间的前向 / 反向计算量不相等。例如，在上面的例子中，Sequence 2 的前向 / 反向计算时间大于 Sequence 1，从而导致多余的空闲气泡。为了解决这一问题，我们根据根据序列对应的计算量来切分序列，而不是简单的均分序列，这一改进在试验中大约有 1.2 倍加速比，充分证明了该切分方法的有效性。
序列切分策略加速比
其它创新点
除了论文提出的主要策略以外，为了和主流的流水线优化技术结合，Seq1F1B还有其它创新点。
➤ Seq1F1B-I （Interleaved version）
Seq1F1B-I 时序图
➤ Seq1F1B integrated with ZeroBubble Pipeline
Seq1F1B 和ZB-H1结合时序图
上图是 ZB-H1 的时序图，其作为 Zero-Bubble 流水线的简化形式，通过简单将反传的梯度计算分为权重梯度计算和激活梯度计算两部分，实现了接近零气泡的流水线调度，而在 ZB-H1 基础上，我们结合 Seq1F1B 技术，提出了 Seq1F1B integrated with ZeroBubble Pipeline，进一步减少了空占气泡，提升了流水线效率。
关于我们
清华大学自然语言处理与社会人文计算实验室（TsinghuaNLP）成立于 20 世纪七十年代末，是国内开展自然语言处理研究最早、深具影响力的科研单位，也是中国中文信息学会计算语言学专业委员会及中国人工智能学会因果与不确定性人工智能专业委员会的挂靠单位。实验室在学术总体带头人孙茂松教授及学术方向带头人刘洋教授、刘知远副教授的带领下，围绕以中文为核心的自然语言处理，在语言大模型、跨模态大模型、中文信息处理、机器翻译、知识图谱、智慧教育、社会人文和艺术计算等方面开展系统深入的研究，在国内外具有较大的学术影响。近年来，实验室承担了国家 973 项目、国家重点研发项目、国家社会科学基金重大项目等多项重要研究任务，并与腾讯、华为、搜狗等企业建立密切的学术合作关系。
Website：http://nlp.csai.tsinghua.edu.cn/
Email：thunlp@163.com
GitHub：https://github.com/thunlp
https://github.com/OpenBMB
Twitter：https://twitter.com/TsinghuaNLP
THUNLP 长期开放招聘
