SIGMOD 2025 | PKU-DAIR實驗室論文

被SIGMOD 2025錄用

ACM SIGMOD是數據庫領域影響力最高的國際學術會議之一，也是CCF推薦的A類國際學術會議之一，近年來接收率17%-20%左右。SIGMOD2025將于2025年6月22日-27日在德國柏林舉行。

PKU-DAIR實驗室論文《MEMO:Fine-grained Tensor Management For Ultra-long ContextLLM Training》被數據庫領域頂級會議SIGMOD 2025錄用。

  面向長序列大模型訓練的細粒度張量管理

作者：Pinxue Zhao, Hailin Zhang, Fangcheng Fu, Xiaonan Nie, Qibin Liu, Fang Yang, Yuanbo Peng, Dian Jiao, Shuaipeng Li, Jinbao Xue, Yangyu Tao, Bin Cui

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

1. 引言

在預訓練階段使用長文本是增強大模型長文本能力的關鍵方法。然而，在實際中，長文本訓練面臨以下嚴重挑戰：

（1）重計算開銷大。在長文本訓練中，激活值顯存隨文本長度線性增大。由于GPU顯存限制，實際訓練中通常需開啟全重計算。重計算會導致約20%的額外端到端計算開銷。

（2）顯存碎片化導致Pytorch顯存分配器經常調用``cudaFree''和``cudaMalloc''兩個同步操作重新整理顯存，嚴重拖慢訓練效率。

2. 方法

（1）針對重計算開銷，我們注意到，由于Transformer層的計算時間隨序列長度平方增長，而層內激活值隨序列長度線性增長，我們提出使用CPU Swap替代（或部分代替）重計算，從而減小重計算開銷。具體而言，為確保CPU-GPU數據傳輸被單個Transformer層的計算覆蓋，以及確保系統CPU內存不被耗盡，我們提出了細粒度的Swap策略。因為在長序列場景下，FlashAttention計算成為整個Transformer層計算的瓶頸，我們將每個Transformer層的輸入和FlashAttention的輸出完整Swap，而對其余激活值，只Swap其一部分。通過這樣的設計，我們盡可能減小了重計算開銷，充分利用了空閑的CPU內存和PCIe帶寬。

圖1. 細粒度swapping和重計算示意

  （2）針對顯存碎片問題，我們利用大語言模型由多個相同的Transformer層堆疊而成的性質，提出一種雙層的混合整數線性規劃算法解決這一問題。具體而言，我們首先執行一輪訓練，獲得張量顯存的存取序列。基于這一序列，我們使用混合整數線性規劃算法確定每個張量的地址，最小化顯存峰值和碎片。進而，在正式訓練時，每個激活值張量的地址都已預先確定，從而避免了Pytorch顯存分配器頻繁進行顯存重整。

圖2. 雙層MIP算法示意

3. 實驗結果

我們基于Megatron-LM和TransformerEngine實現了我們的系統，MEMO。端到端實驗結果表明，我們通過減少顯存碎片以及減少重計算開銷，可以在單機8卡A800上支持7B GPT模型訓練1 Million長度的文本，并且達到超50%的MFU。更多實驗數據可以參考論文。

下圖展示了MEMO的擴展性。我們可以發現MEMO具有線性擴展性，在8機64卡上支持7B模型，8 Million序列長度的訓練。其余實驗可參考論文。

圖3. (a) 支持的最大序列長度 (b)在最大序列長度下的MFU (c)不同序列長度下的MFU

4. 總結

本工作提出面向長序列語言模型訓練的一種細粒度張量管理方案--MEMO。通過結合細粒度的CPU Swapping和重計算技術，以及顯存規劃，MEMO成功支持更長序列的訓練，并且實現了更高的訓練效率，以及擁有較好的擴展性。

實驗室簡介

北京大學數據與智能實驗室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR實驗室）由北京大學計算機學院崔斌教授領導，長期從事數據庫系統、大數據管理與分析、人工智能等領域的前沿研究，在理論和技術創新以及系統研發上取得多項成果，已在國際頂級學術會議和期刊發表學術論文100余篇，發布多個開源項目。課題組同學曾數十次獲得包括CCF優博、ACM中國優博、北大優博、微軟學者、蘋果獎學金、谷歌獎學金等榮譽。PKU-DAIR實驗室持續與工業界展開卓有成效的合作，與騰訊、阿里巴巴、蘋果、微軟、百度、快手、中興通訊等多家知名企業開展項目合作和前沿探索，解決實際問題，進行科研成果的轉化落地。