2025年12月，第39屆神經信息處理系統大會（NeurIPS：Annual Conference on Neural Information Processing System）在美國加利福尼亞州聖迭戈順利召開。NeurIPS是機器學習領域的頂級會議，與ICML、ICLR並稱為機器學習領域三大會議。阿里雲 PAI 團隊與中國科學院大學前沿交叉科學學院等單位合作的研究成果——輕量級動態數據調度方案 Skrull，論文被 NeurIPS2025 會議接收。

長上下文微調（Long-SFT）對提升大模型處理長文本的能力至關重要，但混合長短序列的訓練數據給現有系統帶來效率瓶頸。Skrull 通過在線平衡長短序列的計算負載，在幾乎零調度開銷下顯著提升 Long-SFT 的訓練效率。實測表明，Skrull 相比基線平均提速 3.76 倍，最高達 7.54 倍，為高效長上下文訓練提供了實用的系統優化思路。

一、研究背景

長文本能力是語言模型的核心能力之一，對諸多下游任務都至關重要。續訓練（Continue Pre-Training）和長文本微調（Long Context Fine-Tuning）是擴展大語言模型長文本能力的重要一環。通常情況下，這些訓練場景通常會在精心挑選的數據集中進行，在數據長度分佈上會有顯著的特點，如展現出極度的長尾效應（數據集中短數據佔絕大多數，同時存在超長的訓練數據）或者是雙峯分佈的特徵（短序列和長序列同時在數據集中佔大多數）。這種特殊的數據分佈特徵給現有的訓練系統帶來了廣泛的性能問題。繼 PAI 團隊論文【ICML 2025】Chunkflow後，Skrull在上下文並行（Context Parallelism）以及負載均衡的角度繼續優化系統訓練性能。

二、論文思路

針對這種特殊數據集長度分佈的訓練數據集，原始的序列並行方案很難達到最優性能。首先，訓練數據中長度值差異顯著。單一的上下文並行方案在處理這類場景面臨困難。長序列需要更大的上下文並行維度以減少顯存壓力，但會給短序列處理時帶來更多的通信代價以及性能劣化。特別地，在長文本微調場景中，訓練數據中的短文本通常是佔絕大多數的。因此，如何在維持長文本處理能力的同時，高效地處理較短數據成為了提升該場景訓練系統性能的關鍵問題。Skrull論文中提供了一個高效且魯棒的解決方案。

為了保持長文本處理能力同時提升短文本的訓練效率，Skrull在每個iteration動態地將訓練數據分為兩組（分佈式計算的數據組和局部計算的數據組）。分佈式計算組如同上下文並行的機制一樣，將訓練數據切分到不同的GPU上計算，並通過通信傳輸attention計算所需的Key/Value Cache。局部計算的數據將被完整分配到上下文並行組的某個GPU上，以避免額外的通信和提升計算效率。於此同時，由於兩組計算沒有依賴性，分佈式計算的通信時間可以與局部計算重疊，進一步提升性能。

同時，負載均衡成為提升系統性能的重要環節。局部計算的數據數量和長度同樣表現出顯著差異。尤其是attention機制中，計算量（FLOPs）與數據長度的二次方增長的趨勢與顯存佔用的一次方增長趨勢的差異，使得在追求負載均衡的同時難以對峯值顯存做出有效控制，增加了顯存溢出的風險。

因此，為了拿到最大收益，我們需要規劃出最高效的數據分組以及數據分配方案。理論上，我們可以根據性能建模將該問題形式化成優化問題。但是為了實際的效果以及訓練時表現的魯棒性，Skrull系統使用啓發式的方案來完成上述數據的分組與分配。我們觀察到，儘可能多的將訓練數據用作局部計算能減少通信量和提升運算效率，但不恰當的分組也增大了顯存溢出的風險。同時，我們需要時刻保持計算的負載均衡。我們可以通過統計每個GPU實際負載（FLOPs）來判斷負載均衡情況，從而指導局部計算數據的分配。前兩點設計，雖然最大化了性能收益，但都共同增加了顯存溢出風險。因此，我們設計了回滾機制來排除這種風險。因為訓練顯存佔用與序列長度的線性關係，我們在確定模型和訓練策略的基礎上，很容易就可以推算出單個GPU最長可容納的序列總長度，即為BucketSize。我們將BucketSize作為數據分配的硬約束，當分配序列超出時，我們將會強制回滾操作，保證了訓練的穩定性。

我們上述的優化都是在一個微批次中進行。事實上，我們可以在Global batch內就做這種數據調度以獲取更大的性能提升空間，同時不影響模型訓練的優化軌跡。同理，我們通過排序並間隔取長短序列的方式，使得其在數據並行維度更加負載均衡、並將長短序列均勻分配到不同微批次中。

三、實驗數據

在多種尺寸的Qwen系列模型中驗證系統收益。選取了三個數據集，分別代表常見的長尾和雙峯分佈。（注：前兩個數據集不是專用於長文本微調場景但是其數據分佈與該場景極為相似）。下圖展示了在不同配置下Qwen-0.5B和Qwen-7B相對於DeepSpeed（Zero-2）和簡單排序（sorted batching）均取得了顯著的加速。

同時，我們測試了不同BatchSize和BucketSize設定對於性能的影響、更大尺寸模型以及高效微調方法Lora的兼容性（如下圖所示）。進一步的消融實驗和分析（如下表所示）展示出Skrull的啓發式策略以及回滾機制對於性能提升的重要性。

四、更多論文相關信息

論文標題：
Skrull: Towards Efficient Long Context Fine-tuning through Dynamic Data Scheduling

論文作者：
Hongtao Xu，Wenting Shen，Yuanxin Wei，Ang Wang，Guo Runfan，TianxingWang，Yong Li，Mingzhen Li，Weile Jia

論文鏈接：
https://arxiv.org/abs/2505.19609