本文中我們通過幾篇論文來具體介紹 Sample Packing 相關的方案和對應的各種問題，比如 GraphCore 的PackedBert、Meta 的 In-Context-Pretraining、智譜 AI 的 LongAlign、Amazon 的 Fewer Truncations 以及 IBM 的 Packing with FlashAttention。

一、背景

上一篇文章（???Sample Packing：長序列 LLM 訓練的 Attention 問題及優(yōu)化??）中我們簡單介紹了 Sample Packing 相關的問題和部分簡單實驗。本文中我們通過幾篇論文來具體介紹 Sample Packing 相關的方案和對應的各種問題，比如 GraphCore 的PackedBert、Meta 的 In-Context-Pretraining、智譜 AI 的 LongAlign、Amazon 的 Fewer Truncations 以及 IBM 的 Packing with FlashAttention。

二、方法

Sample Packing 可以看成是一個經(jīng)典的 Bin Packing 組合優(yōu)化（Combinatorial Optimization）問題，核心思想是將一組不同體積的物品放入容量固定的箱子中，目標是最小化所需箱子數(shù)量。

Bin Packing 是一個典型的 NP-Hard 問題，因此往往需要關注其效率。對于 N 個 Sample 的數(shù)據(jù)集，通常需要先排序，對應 O(N*logN) 的時間復雜度；然后 Packing，需要 O(N*logN) 的時間復雜度。不過由于 Sample 的長度通常比較有限，可以采用計數(shù)排序的方式進行優(yōu)化；此外，整個訓練集通常有多個數(shù)據(jù)子集組成，為了保證每個數(shù)據(jù)子集具有不同的采樣權重，往往一個 Sequence 中的 Sample 來自同一個數(shù)據(jù)子集，一個大的 Bin Packing 問題也變成了多個小的 Bin Packing 問題，可以進一步降低復雜度。本文中我們就不再具體介紹 Bin Packing 的優(yōu)化問題。

如果從物品順序角度考慮，Bin Packing 可以分為兩類：

• Online Bin Packing：物品按順序到達，必須立即決定放入哪個箱子，無法預知后續(xù)物品的大小。對于 LLM 訓練而言，Online 方式依然可以實現(xiàn) Batch 內(nèi)（窗口內(nèi)）的亂序，也可以通過梯度累加增加 Batch 的大小。
• Offline Bin Packing：預先知道所有物品大小，可以全局排序。對于 LLM 訓練而言，相當于訓練前就預先知道所有序列的長度，并對所有 Sample 打包。

當然，Bin Packing 方案也可能有不同的約束條件：

• 一次性 Packing：每個物品只能放進一個箱子，不允許拆分物品。對于 LLM 而言，可以理解為 Sample 不能截斷。
• 多維 Packing：物品和箱子不僅有體積，還有其他屬性，比如重量、形狀，需要同時滿足多維約束。對于 LLM，可以約束同一個 Batch 內(nèi)不同 Sequence 的計算量盡量類似，以實現(xiàn)更好的負載均衡。

整體來說，對于 LLM 訓練可以從數(shù)據(jù)分布和計算效率的角度考慮相關的方案，數(shù)據(jù)分布可能影響模型訓練的效果，計算效率會影響訓練的速度，往往需要綜合考量。

影響數(shù)據(jù)分布的幾個因素：

• 是否隨機采樣，比如排序機制可能引入分布的不一致。
• 是否交叉污染，比如 Packing 是否采用 Document Level 的 Mask（Block Diagonal Mask，也可以通過 Position ID 區(qū)分）。
• 是否有 Sample 的截斷。

影響計算效率的幾個因素：

• 是否 Padding，Padding 是否參與計算。
• Packing 是否采用 Document Level Mask。
• 是否存在負載不均衡問題（主要是指稀疏度不同，計算量不同）。

三、GraphCore PackedBert

3.1 概述

GraphCore 在 2021 年 07 月的 [2107.02027] Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance 中已經(jīng)討論了 Sample Packing 導致的交叉污染（Cross Contamination）問題。作者研究了新的 Packing 算法，并通過修改 Attention Mask 和 Position ID 來避免交叉污染，提升 Bert 模型的訓練效率。

如下圖 Table 1 所示，不同的 Packing 方式有不同的 EFF（有效率），其中 Baseline 的 None 表示有 50% 左右的 Padding Token，而 SPFHP 和 NNLSHP（本文提出） 能獲得相對比較高的 EFF。

3.2 結果

如下圖 Figure 4 所示可以看出，需要相應修改 Attention Mask 才能保證精度：

如下圖 Figure 5 所示，隨著 Accelerator 的增加，本文的方案能獲得比較穩(wěn)定的加速，非常接近理論速度（最上的藍線）；而非 Padding 的方案可能隨著 Accelerator 的增加出現(xiàn)明顯的降速。

四、Meta In-Context-Pretraining

4.1 概述

Meta 作者在 [2310.10638] In-context Pretraining: Language Modeling Beyond Document Boundaries 也關注了 Sample Packing 的問題。作者指出，之前的預訓練流程在訓練時將隨機的短文檔拼接起來形成輸入上下文，這些文檔之間沒有提供預測下一個文檔的信號，導致計算效率不高。因此作者提出了新的方案：In-context Pretraining，通過改變文檔的順序，使得每個上下文包含相關的文檔，從而明確鼓勵模型跨文檔邊界進行閱讀和推理。

4.2 方案

如下圖所示，在預訓練前會計算文檔的相似性，在 Packing 時利用上這種相似性，保證 Sequence 中文檔盡可能相關。

4.3 結果

作者使用 CommonCrawl 數(shù)據(jù)集，預訓練了 0.3B 到 7B 參數(shù)量的多個模型，并在多種任務上評估，包括上下文學習、閱讀理解、對先前上下文的忠實度、長上下文推理和檢索增強等。與使用標準方法預訓練的模型相比，In-context Pretraining 方法訓練出的模型（ICLM）顯示出顯著的性能提升。

如下圖 Figure 3 所示，本文 ICLM 訓練的模型獲得了更低的困惑度：

如下圖 Table 1 所示，基于 ICLM 訓練的模型在下游 In-context Learning 任務上也獲得了更好的效果：

五、智譜 LongAlign

5.1 概述

智譜 AI 在 [2401.18058] LongAlign: A Recipe for Long Context Alignment of Large Language Models 中討論了部分 Sample Packing 相關問題。如下圖 Figure 3 左圖所示，Sequence 的長度各不相同，從 0 - 60K，如果采用 Naive Batching 方式，會導致明顯的 Bubble 問題（雖然 NoPadding 技術可以避免重復計算，但是如果采用 Data Parallelism 方式，比較快的設備需要等待比較慢的設備計算完成）。為了解決效率和效果問題，作者提出了 3 種解決方案：Packing、Loss Weighting 和 Sorted Batching。

5.2 Packing

如 Figure 3 右上圖所示，就是我們之前介紹的 Sample Packing：將不同的 Sample 拼接在一個 Sequence 里，并且保證盡可能接近 Max Sequence Length，末尾的部分 Token 進行 Padding。然后通過 Block Diagonal Attention Mask 來區(qū)別不同的 Sample，以避免 Sample 之間的交叉污染，也就是 Document Level Attention。

PS：作者介紹，這里同樣是使用了 FlashAttention2 的 Varlen 特性。

5.3 Loss Weighting

假設訓練時的 Batch Size 為 K，總共包含 M 個 Sample，第 i 個 Sample 的 Token 數(shù)為 Ni，則對應的 Loss 如下圖所示：

然而，增加 Sample Packing 之后會引入一個問題，如下圖所示，一個 Sequence 中的不同 Sample 會被看成一個 Sample 來計算損失。當有些 Sample 比較長，其對應的 Token 很多，那么這個 Sample 對 Loss 的貢獻就更大，模型可能會在訓練時更傾向于優(yōu)化長 Sample 的表現(xiàn)，進而可能會導致對短 Sample 的學習有所欠缺。

為了解決這個問題，作者提出 Loss Weighting，也就是對不同 Sample 的 Loss 加權。如下圖所示，保證其和上述公式（2）等價。作者聲稱可以在下游任務上帶來 10% 左右的效果提升。

PS：不過這里還會引入另外一個問題，因為采用了 Sample Packing，那么實際上不同 Step 中 Sample 的個數(shù)在不斷變化。比如一個 Batch 里都是短 Sample，那么對應的 M 會比較大；如果一個 Batch 里都是長 Sample，相應的 M 會比較小。這樣可能引入兩個問題，1. 相當于一個 Batch Size 在不斷變化；2. 同樣一個 Sample 可能會因為順序的原因被賦予不同的權重。因此需要盡量保證Batch 中 Sample 的平均個數(shù)比較穩(wěn)定。 

5.4 Sorted Batching

如上圖 Figure 3 右下圖所示，可以將所有 Sample 進行排序，在組 Batch 時盡量保證一個 Batch 中的 Sample 長度相同（沒有 Packing）。這樣可以保證不同設備的計算盡可能的均衡。然而，這種方式也不可避免地引入不同 Batch 的數(shù)據(jù)分布的偏差，有些 Batch 都是長序列，有些 Batch 都是短序列，對于 SGD 優(yōu)化來說可能并不友好。不過作者發(fā)現(xiàn)這種方式可以顯著加快訓練速度，而不會對效果產(chǎn)生明顯的負面影響。這可能是因為使用了大的梯度累加（Micro Batch 中長度類似，但整個 Batch 中包含各種長度的 Sample）。

PS：這種方式也就對應 Transformer 等工作中的 LengthGroupedSampler，如下圖所示，排序后可以有效降低無效計算（圖片參考 數(shù)據(jù)分組— XTuner 0.1.23 文檔）。

5.5 結果

如下圖 Figure 5 所示，作者在 8xA800 GPU 上進行速度對比，可以看成，Packing 和 Sorted Batching 相比 Naive Batching 都有 2x-3x 的加速：

如下圖 Table 3 所示，作者基于 ChatGLM3-6B-64K 和 LLaMA-2-7B-64K 進行了相關效果驗證。可以看出，Loss Weighting 在 LongBench-Chat 上能帶來 5%-10% 的提升，但是在其他任務上并不明顯，并且這些方法看著都不是特別魯棒。

5.6 Packing 負載均衡

智譜 AI 在訓練 GLM4（GLM Long: Scaling Pre-trained Model Contexts to Millions | by ChatGLM | Medium） 模型時進一步解決了 Packing 帶來的負載均衡問題。如下圖所示，雖然都 Packing 到了相同的長度，但是由于其中的 Sample 個數(shù)、長度不同，導致其稀疏度差距很大，計算量也相應差距很大。如果它們在不同的設備上執(zhí)行，同樣會存在計算不均導致的 Bubble 問題。

如下表所示，我們在上一篇文章中的實驗也能說明這個問題，隨著 Sequence 中 Sample 分布的不同，計算的耗時甚至可能差 10x：

如下圖所示，作者發(fā)現(xiàn)訓練中每個 Step 的時間存在較大幅度的波動，這種現(xiàn)象在短文本的 Packing SFT 中并不明顯。這是因為短文本時 Attention 的計算占比并不高，而超長文本訓練中會尤其明顯。

為了解決上述問題，作者進一步提出了 Sorted Packing。具體來說，作者在構建 Batch 數(shù)據(jù)時考慮了計算復雜度，以確保每個 Batch 中計算復雜度相似，從而減少 Bubble 時間。（PS：這里需要注意，計算復雜度不等于執(zhí)行速度，如果能針對預估計算速度來打包也許能獲得更優(yōu)的效果）

作者指出也使用了 layer accumulation 技術（PS：上述介紹的梯度累加？）來避免排序?qū)е碌钠脝栴}。

六、Amazon Fewer Truncations

6.1 概述

在 [2404.10830] Fewer Truncations Improve Language Modeling 中，作者探討了數(shù)據(jù)截斷問題對模型效果的影響。作者指出，截斷會損害數(shù)據(jù)完整性，從而阻礙模型學習基于完整上下文撰寫邏輯連貫且事實一致的內(nèi)容的能力。

為了解決這個問題，作者提出了 Best-fit Packing，通過長度感知組合優(yōu)化來 Packing，可以完全消除不必要的截斷，同時基本不影響訓練效率。通過文本和代碼預訓練的實驗結果表明，提出的方法可以在閱讀理解上相對提升 4.7%，上下文跟隨提升 16.8%，程序生成提升 9.2%。此外，也可以有效減少 58.3% 的封閉域幻覺問題。

PS：論文中對比實驗時 Baseline 的 Concatenation 方案中有截斷，并且沒有使用 Block Diagonal Mask；而 Best-Fit Packing 使用了 Block Diagonal Mask，且沒有截斷。

6.2 方案

如下圖 Figure 1 所示為本文 Best-Fit Packing 與傳統(tǒng)方案的對比。

• 右圖所示為傳統(tǒng)方案：可以理解為將所有樣本排成一行，然后按照 Max Sequence Length 進行切割，會導致大量 Sample 被截斷。
• 左圖為本文的方案：首先將所有 Document 按照 Max Sequence Length 截斷，然后使用 Best-Fit Decreasing 算法來進行組合優(yōu)化（Bin Packing 優(yōu)化）。?

如下圖 Table 2 所示，使用本文的 Best-Fit Packing 可以保證只增加不到 0.003% Sequence 數(shù)量，對訓練效率的影響也就微乎其微。

6.3 結果

如下圖 Table 3 所示，作者訓練了 3 個不同規(guī)模、序列長度的模型：

如下圖 Table 4 所示，提出方案訓練出的模型的閱讀理解能力有明顯提升：

如下圖 Table 9 所示，幻覺問題也可以明顯降低：

如下圖 Table 10 所示，作者也針對 Attention Mask 進行了相關消融實驗，可以看出原始 Concatenation 方案加入 Block Diagonal Mask 后也有一定的提升。可以證明 Attention Mask 和截斷都會對效果有一定影響，不過 Packing（避免截斷） 的影響似乎更大一些。

七、IBM Packing with FlashAttention2

7.1 概述

在 [2407.09105] Enhancing Training Efficiency Using Packing with Flash Attention 中，作者總結了不同 Packing 策略、Mask 方式及與 FlashAttention 結合的優(yōu)勢。此外，作者也將相關工作提交到了 Huggingface Transformer 中，提供了新的 DataCollatorWithFlattening，具體可以參考：通過打包Flash Attention 來提升Hugging Face 訓練效率。

7.2 相關方案

如下圖 Table 1 所示，作者分析了不同的 Packing 方案以及它們的影響，具體包含如下幾種方式：

• RandomSampling + Padding：最傳統(tǒng)的隨機采樣，然后 Padding 的方式。存在冗余計算，并且占比很高。
• GroupByLength+Padding：先排序，然后盡量保證每個 Batch 中的序列長度接近。可以減少 Padding 的占比。
• RandomSampling + PosID：隨機采樣，但是不 Padding，而是通過 PosID 支持變長序列。幾乎沒有冗余計算，但可能存在明顯的負載不均衡（計算量）。
• FixedLengthPacking：隨機采樣，隨機 Packing，并且最后一個Sample 可能截斷，保證填滿 Max Sequence Length。沒有區(qū)分不同 Sample，也就是 Causal Mask，沒有冗余計算，并且負載很均衡。
• FixedLengthPacking + PosID：相比FixedLengthPacking多了 PosID，也就是可以區(qū)分不同 Sample，對應 Block Diagonal Mask。但依然會存在末尾截斷，并且可能負載不均衡。
• MultiPack + PosID：使 Sequence 中的數(shù)據(jù)盡量接近 Batch 的 Max Sequence Length，降低 Sequence 中的長度不均衡，可以參考GitHub - imoneoi/multipack_sampler: Multipack distributed sampler for fast padding-free training of LLMs。需要對數(shù)據(jù)進行排序。
• SortedPacking + PosID：通過排序，使同一個 Batch 中的計算復雜度盡量接近。可以盡可能降低計算負載不均衡問題。
• RandomPacking + PosID：與FixedLengthPacking + PosID相比主要的區(qū)別就是最后一個 Sample不截斷，可能存在部分 Bubble。?

7.3 結果

作者通過微調(diào)任務對比了一系列模型使用不同方案的效果和速度，其中 Max Sequence Length（msl）為 4096，每個 GPU 的 Mini-Batch Size 為 4。對應配置如下：

• no：表示最原始的RandomSampling + Padding，可以作為基線，冗余計算比較多，速度最慢，但是效果有保障。
• yes：表示FixedLengthPacking，存在交叉污染。
• flat：表示訓練前 Offline Packing 好，也就是引入了排序，并且使用 PosID 實現(xiàn) Document Level Mask。
• mini：表示訓練中 mini batch 的 Online Packing，和 Random 類似，并且使用 PosID 實現(xiàn) Document Level Mask。

如下圖所示（PS：這里只是部分結果，全量請參考論文，結論基本一致，可以看出 yes 和 flat 都會對精度（VLoss）有比較大的影響，但速度（Tok/s）確實快了很多，可以達到 Baseline 的 3x-4x；而 mini 可以在保證精度的情況實現(xiàn)實現(xiàn) 2x 左右加速。

作者使用 FLAN_20k 數(shù)據(jù)集在 Mistral-7B 上針對之前提到的幾種方案做了更多實驗（PS：這里只是部分，gas 表示梯度累加次數(shù)），看起來 Packing 的方案都能獲得不錯的速度，但是精度的影響因素就比較多。但整體來說 RandomPacking + PosID 的方案還不錯，看著似乎也不能有太多的梯度累加。

八、參考鏈接

1. ??https://arxiv.org/abs/2107.02027??
2. ??https://arxiv.org/abs/2310.10638??
3. ??https://arxiv.org/abs/2401.18058??
4. ??https://xtuner.readthedocs.io/zh-cn/latest/acceleration/length_grouped_sampler.html#length-grouped-sampler??
5. ??https://medium.com/@ChatGLM/glm-long-scaling-pre-trained-model-contexts-to-millions-caa3c48dea85??
6. ??https://arxiv.org/abs/2404.10830??
7. ??https://arxiv.org/abs/2407.09105??
8. ??https://huggingface.co/blog/zh/packing-with-FA2??
9. ??https://github.com/imoneoi/multipack_sampler??

本文轉(zhuǎn)載自 ??AI閑談??，作者： AI閑談