训练的 LLM Packing Sample Attention 长序列 疑问及优化

一、背景

之前看过局部 Megatron-LM 的源码，也详细剖析过对应的>Sample Packing 中有很多可以探讨的技术点，比如 Attention 的成功和优化，Sample 的组合及负载平衡疑问（有点相似调度疑问）以及不同打算对成果的影响等。咱们这里只是先便捷引见一下关系疑问和试验，后续会进一步探求更多上班，比如 Document Level 的 Mask 究竟对预训练成果影响有多大，对 Attention 启动优化还能带来多少优化，如何设计一个比拟好的 Packing 战略等？

关系上班可以参考咱们之前的文章：

二、Dataset +>之前的文章（​​ ​LLM 预训练语料、预处置和数据集索引、加载总结​ ​）中详细引见过 Megatron-LM（DeepSpeed-Megatron）中预训练>如下图所示，便捷展现了 Megatron-LM 中如何 Packing 多个 Document，实践上就是一个多级的索引。须要说明的是，这里其实会引入很多随机读操作，会极大影响读的性能。不过普通 LLM 计算代价都很高，这里也往往不会造成瓶颈。

三、Attention Mask

关于单个 Document 而言，Decoder Only 的 GPT 模型具备 Causal 特性，也就是每个 Token 不能看到之后的 Token，因此在实践训练中须要减少 Attention Mask。如下图所示，这种状况下 Attention Mask 是一个规范的下三角矩阵（Causal Mask），也就是绿色局部为 1，其余局部为 0：

假设一个 Sample 里蕴含多个样本，则 Attention Mask 矩阵须要变成如下图所示的块对角矩阵方式（Block Diagonal Mask）。比如 Sequence Length 为 16，4 个 Document 的长度区分为 3，4，5，4，则对应 Attention Mask 矩阵如下图所示，对角线上的 4 个矩阵（红框）都是规范的下三角矩阵。依照这种方式可以保障和 4 个 Document 独自作为 Sample 训练是等价的：

四、Reset Attention Mask

4.1 能否须要

那么在实践经常使用中能否须要严厉依照 Block Diagonal Mask 的方式经常使用呢？答案能否认的，比如 Megatron-LM 可以经过 reset_attention_mask 来控制是经常使用 Block Diagonal Mask 还是规范的 Causal Mask，自动值为 False。很多模型在预训练时也会驳回自动性能，即使用 Causal Mask。

在浪潮的 Yuan-1.0 报告（“源1.0”大模型技术白皮书）中有提到，为了防止不同 Document 之间的相互搅扰而将 reset_attention_mask 设置为 True，也就是 Block Diagonal Mask：

在 Meta 的 LLaMA 3.1 技术报告（[2407.21783] The Llama 3 Herd of Models）中也提到，在 LLaMA 3.1 模型的预训练中会关上这特性能。不过作者也做了说明，关于 8K Sequence Length 的预训练而言，对模型最终的成果影响不大，对长序列的 Continuous PreTraining 影响比拟大：

在 [2402.08268] World Model on Million-Length Video And Language With Blockwise RingAttention 中作者提出了“环球模型”，为了优化超长序列的训练效率，作者驳回了 Sample Packing 的战略，并且做了关系消融试验。如下图 Table 10 所示，驳回 Naive Packing（不对 Attention Mask 不凡处置）相比经常使用了 Block Diagonal Mask 的 LWM 的性能会差很多：

PS：当然，目前还没有更多无关预训练中能否 reset_attention_mask 的消融试验，咱们后续会启动关系测试。此外，假设驳回相对位置编码，Position-id 也须要相应的调整，在 Megatron-LM 中对应 reset_position_id 选项。

4.2 性能疑问

如下图为 Megatron-LM/megatron/core/datasets/gpt_dataset.py 中 reset_attention_mask 的成功方式，首先会将 attention_mask 初始化为规范的 Causal Mask 方式，而后从第二个 Document 开局，将之前的 mask 置为 0：

详细来说如下图所示，初始是一个规范的 Causal Mask 矩阵，而后会将 4x3、5x(3+4) 和 4x(3+4+5) 的区域依次置为 0，之后会变成 Block Diagonal Mask：

实践上咱们曾经知道这里是规范的 Block Diagonal Mask，可以经常使用 torch.block_diag() 极速创立。实测当序列比拟长时（比如 32K），两种方式速度或许会差几十倍，造成 reset_attention_mask 或许成为训练瓶颈：

除此之外，当序列十分长时，Attention Mask 也会占据很大的存储空间，为了计算效率，往往会经常使用整型而不是 Bool 类型。假定以 int8 存储，32K 序列长度对应的 Mask 大小为 32K * 32K = 1GB，128K 时更是高达 16GB。为了防止显存糜费，其实不用将其拼成大的 Block Diagonal Mask，而保管几个小的 Causal Mask 即可。

五、Attention 优化

5.1 FlashAttention

LLM 预训练基本都会经常使用 FlashAttention，其对 Casual Mask 的方式启动了优化，如下图所示，假定 16x16 的 Attention Mask，在计算时依照 4x4 分块，则可以将其分为 3 种状况：

关于上述 Block Diagonal Mask，依然可以经常使用 Causal Mask 的方式计算，不过会造成少量的有效计算。幸运的是，FlashAttention V2 允许可变序列长度（Varlen）的 Batching Attention 计算，可以防止 Padding 造成的有效计算。因此也就可以借用这种机制来对 Block Diagonal Mask 启动解构，从新合成为多个 Causal Mask 区分计算，可以防止很多有效计算。如下图所示，可以将其看成 4 个独立的 Attention 计算，详细可以参考 FlashAttention Github 上的关系探讨：How to implement example packing with flash_attn v2? · Issue #654 · Dao-AILab/flash-attention · GitHub 和 Will attention_mask be extended to 3D? (concatenate short samples for efficient training) · Issue #432 · Dao-AILab/flash-attention · GitHub。

在 GLM-4（[2406.12793] ChatGLM: A Family of Large Language Models from GLM-130B to GLM-4 All Tools）中也运行了 Sample Packing 打算，并且雷同经常使用了 Block Diagonal Mask 机制来区分不同的 Document。并且作者也是基于 FlashAttention 的 Varlen 性能来成功。

5.2 Pytorch FlashAttention

Pytorch 的 scaled_dot_product_attention 提供了高效的 Attention 成功，也集成了 FlashAttention2 的成功，但是其不允许上述的可变序列长度的性能，造成针对 Block Diagonal Mask 场景时会存在少量的重复计算。

此外，咱们在之前的文章中也屡次提到，当序列比拟短时，Attention 局部计算的占比并不是特意大，因此其中的冗余计算或许对全体训练速度影响不大；但当序列比拟长时，Attention 局部计算的占比会越来越大，冗余计算或许会对训练速度有比拟大的影响，也就须要对其启动优化。

5.3 FlexAttention

Pytorch 在 2.5.0 版本引入了 FlexAttention（FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention），可以很容易允许各种 Attention Mask 变种，比如规范 Causal Mask、Sliding Window + Causal、Prefix Mask 以及 Document Mask（Block Diagonal Mask）等，相比 FlashAttention 也愈加的灵敏。

咱们基于 FlexAttention 启动了关系测试，以验证经常使用 Block Diagonal Mask 的性能长处。首先以两个 16K Document 拼接为一个 32K Sample 为例，Attention Mask 大略是如下图所示方式，对应的稠密度为 74.80%（整个 Mask 中 0 的占比）：

如下图所示咱们在 H100 GPU 上启动的 Attention 关系性能测试。可以看出， Pytorch 的 Causal + FlashAttention2 方式确实可以到达十分高的 TFLOPS，显著高于 FlexAttention。但是，由于 FlexAttention 中防止了很多有效计算，实践的 Forward 和 Backward 期间反而更短：

当然，也并不异常着 FlexAttention 总是更优的，还和 Sample 中 Document 长度无关。如下图所示为相应测试结果，32K 示意 Sample 中只要一个 Document，2K + 30K 示意 Sample 中有 2 个 Document，一个长度 2K，一个长度 30K。从下图基本上可以得出这样一个论断：当 Sample 中最长的 Document 的长度 <= Sequence Length/2 时，经常使用 FlexAttention 或许会带来更大的收益：

那么为什么“最长的 Document 的长度 <= Sequence Length/2”时会有收益呢？其实可以便捷从稠密度的角度思考：假定 a1 + a2 + a3 + ... + an = S，并且 0 < a1 <= a2 <= a3 <= ... <= an <= S/2，那么可以用数学演绎法得出 (a1)^2 + (a2)^2 + (a3)^2 + ... + (an)^2 <= S^2/2。也就是说，最长的 Document 的长度 <= Sequence Length/2 时，稠密度会 >= 75%（还要思考 Causal 特性），相应的 FlashAttention 中至少有一半的冗余计算。

因此，咱们也须要充沛思考在长文本训练环节中短文本的占比，极其状况下训练数据所有是超长文本，每个 Sample 中都只要一个 Document，Block Diagonal Mask 会退步为 Causal Mask。不过有些时刻为了防止模型产生劫难性忘记，也会混合一些短文本数据，或许高品质的预训练数据，无法防止的会产生冗余计算的疑问。

5.4 Sequence Parallel

咱们在之前的序列并行文章（​​ ​大规模散布式 AI 模型训练系列——序列并行​ ​）中也提到过，针对长序列场景理论会驳回 RingAttention 和 USP 等，但是不论是 RingAttention 还是其 LoadBalance 版本（如下图 Figure 3 所示）等都没有太多探讨 Sample Packing 的状况。关于 Block Diagonal Mask 场景，其相应的优化，LoadBalance 战略也或许须要对应调整：

在 [2402.08268] World Model on Million-Length Video And Language With Blockwise RingAttention 中作者（也是 RingAttention 的作者）宣称针对 Block Diagonal Mask 场景对 RingAttention 启动关系优化，但并没有对比优化前后训练速度的优化。

PS：全体来说，在各种序列并行技术中更好的兼容 Block Diagonal Mask 场景又会有更多的应战，咱们留作后续引见。

六、参考链接

本文转载自​​，作者：