一、背景

我們之前已經分享過幾篇關于 Sample Packing 相關的文章，也提到了其中的性能優化問題。最近今天又看到一篇新的論文，這里進行簡單介紹。

對應的論文為：[2410.08081] Packing Analysis: Packing Is More Appropriate for Large Models or Datasets in Supervised Fine-tuning

相關工作可以參考我們之前的文章：

• Sample Packing：長序列 LLM 訓練的 Attention 問題及優化
• Sample Packing 綜述：LLM 效果與效率的 Tradeoff
• Binary Block Masking：加快稀疏 Attention 的一種新方法?

二、摘要

Sample Packing 技術最初應用于大模型預訓練階段，旨在通過合并不同的訓練樣本來適應模型的最大輸入長度，從而最大化硬件利用率。盡管在預訓練中已展示了有效性，但對監督微調（SFT）階段還缺乏全面分析：

• Packing 能否在保持性能的同時有效提升訓練效率。
• 適合采用 Packing 技術進行微調的模型和數據集的規模。
• Packing 無關或者相關的訓練樣本是否可能導致模型過度忽略或依賴上下文。

本文中，作者對使用 Padding 和 Packing 的 SFT 方案進行了對比，涵蓋從 69K 到 1.2M 的 SFT 數據集，以及 8B 和 70B 規模的模型。作者首次提供了關于 Padding 和 Packing 優缺點的全面分析，以及在各種訓練場景中實施 Packing 的考量。作者的分析涉及多個基準，包括知識、推理和編碼，以及基于 GPT 的評估、訓練效率等。此外，作者還開源了相應微調和評估代碼，并提供了在不同規模數據集上微調的 Checkpoint，旨在推動未來 Packing 技術的研究。

PS：個人感覺這篇文章還有很多可以改善的地方，比如說：

• 沒有考慮 Document Level Mask 的問題。
• 沒有考慮長序列場景。
• 可以提供更多的實驗證明其觀點：“特殊 Token [EOS] 可以有效地幫助模型區分兩個相鄰的訓練樣本”。

三、方法

3.1 方法概述

如下圖 Figure 1 所示為對應的 Padding 和 Packing 方案，其中 SFT 數據集包含單輪對話（一個樣本只包含一個輸入和對應的一個輸出）和多輪對話（包含多個輸入-輸出對）。

• Padding：

每個序列中只包含一個樣本，如果不超過 Max Sequence Length 則補充[PAD] 占位 Token，到滿足 Max Sequence Length。

如果是多輪對話，每個輸入-輸出對之后也要添加[EOS] 結束 Token。

• Greedy Packing：
• 首先，將多輪對話的多個輸入-輸出對合并成單個樣本，用 [EOS] 間隔。
• 然后，將所有的數據按照長度排序。
• 最后，從最長的樣本開始 Packing，將連續的幾個樣本打包，如果加上某個樣本后總長度超過 Max Sequence Length，則放棄這個樣本，補充 [PAD] Token 并開始一個新的序列。
• Random Packing：
• 首先，將所有訓練樣本拼接為一個單一的序列，同樣會使用 [EOS] Token 間隔。
• 然后，按照 Max Sequence Length 將序列截斷。?

3.2 優劣

Padding 方案的優勢和劣勢：

• 優勢：

實現簡單，任何框架都比較容易實現。

簡單的 Padding 和截斷可以避免變長序列引申出的各種計算錯誤或性能問題。

• 劣勢：
• 如果 Padding 太多，會導致計算量明顯增加，并且對訓練沒有任何幫助。（PS：其實很多框架都支持變長計算，比如 FlashAttention，識別出 Padding 的 Token 可以不參與計算）

Random Packing 方案的優勢和劣勢：

• 優勢：

可以最大限度減少未使用的空間（Token），從而提升計算效率。

可以提供更廣泛的上下文組合，從而提供更多元化的場景，有可能幫助提升模型的泛化能力。（也就是說不會使用 Document Level 的 Mask？）

• 劣勢：
• 可能導致兩個不同或相似的樣本串聯，可能導致模型過度忽略或依賴上下文。不過作者提到這并不會導致模型過度關注 Packing 的上下文，主要是兩個原因：

Random Packing 導致兩個相似訓練樣本打包在一起的概率非常低。

特殊 Token [EOS] 可以有效地幫助模型區分兩個相鄰的訓練樣本。（PS：是否可以通過驗證 Attention Score 來判斷是否真的如此？）

• 這種方法也可能導致單個對話中的不同輪次被劃分到了不同的子序列中，比如 Figure 1 中的 (Instruction 1 + Answer 1) 和 (Instruction 2 + Answer 2) 被劃分到不同序列中

Greedy Packing 方案的優勢和劣勢：

• 優勢：

可以降低將一個多輪對話樣本劃分到不同序列的風險。

• 劣勢：
• 需要對整個樣本排序，會打亂樣本分布的隨機性，也會影響一個 Batch 中樣本的多樣性。

PS：上述 Greedy Packing 方案也有可以改進的地方，作者的 Packing 方法中會優先使用當前更長的序列來拼接，如果超過 Max Sequence Length 就直接結束，這也就導致更多的短上下文樣本無法利用這一部分空間。比如說，Max Sequence Length 為 2048，當前已拼接的序列長度為 1800，需要 Padding 248 個 Token，假設下一個樣本的長度為 300，則會直接結束；然而，當前還有大量長度小于 248 的 Token 可以被拼接。

四、實驗

4.1 實驗配置

包含 4 個 SFT 數據集：

• WildChat(GPT-4)：WildChat 是用戶訪問 ChatGPT 或者 GPT-4 生成的數據。這里作者只選擇了 GPT-4 相關數據，包含 69K 個真實世界的數據。
• TULU：包含 326K 個對話數據集，也是來自真實世界。
• WildChat(Full)：來自 WildChat 的全量 652K 數據。
• Open-Source 1M：總共包含 1.2M 對話數據，來自多個數據源。如下圖 Table 1 所示。?

訓練配置如下圖 Table 2 所示，主要是基于 LLaMA-3-8B 和 LLaMA-3-70B 微調，訓練是在 4 臺機器，每臺集群 8 個 A800 GPU：

4.2 實驗結果

如下圖 Table 3 和 Table 4 所示，作者使用各種模型和數據集驗證，Packing 的平均性能會比 Padding 的更好。比如 LLaMA-3-70B 上微調 WildChat（GPT-4）數據集，Padding 對應平均性能為 61.96，而 Packing 對應為 65.97（Random Packing）。

PS：上述的結論有點奇怪，智譜 AI 在 [2401.18058] LongAlign: A Recipe for Long Context Alignment of Large Language Models 中的實驗表明，短序列任務上的差異并不明顯，而在長序列上甚至 Packing 會導致指標下降。當然，LongAlign 使用了 Document Level 的 Mask，實際上 Meta LLaMA 等論文也建議在微調時使用 Document Level Mask。

如下圖 Table 5 所示，使用 Packing 的方案確實可以有效提升訓練速度，可以實現 3x 左右的加速，主要是 Packing 后 Padding 少了很多，訓練的 Step 數少了很多：

上述結論也與智譜 AI 在 [2401.18058] LongAlign: A Recipe for Long Context Alignment of Large Language Models 中的結論類似，如下圖 Figure 5 所示：

五、參考鏈接

1. https://arxiv.org/abs/2410.08081
2. https://arxiv.org/abs/2401.18058

本文轉載自 ??AI閑談??，作者： AI閑談