PKU-DAIR實驗室兩項成果被SIGMOD 2025錄用

SIGMOD是計算機科學領域的頂級國際學術會議之一,，專注于數據庫系統(tǒng)和數據管理技術的研究。作為中國計算機學會（CCF）推薦的A類會議，SIGMOD與VLDB,、ICDE并稱為數據庫領域的三大頂會,。PKU-DAIR實驗室《Malleus: Straggler-Resilient Hybrid Parallel Training of Large-scale Models via Malleable Data and Model Parallelization》和《PQCache: Product Quantization-based KVCache for Long Context LLM Inference》兩篇論文在SIGMOD第四輪評審中被成功錄用。

1. Malleus: Straggler-Resilient Hybrid Parallel Training of Large-scale Models via Malleable Data and Model Parallelization

作者：Haoyang Li, Fangcheng Fu, Hao Ge, Sheng Lin, Xuanyu Wang, Jiawen Niu, Yujie Wang, Hailin Zhang, Xiaonan Nie, Bin Cui

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

Github鏈接：https://github.com/PKU-DAIR/Hetu

圖1. Malleus概覽

在大規(guī)模模型訓練過程中,，個別GPU或機器出現(xiàn)性能波動（稱為“掉隊者”,，straggler）會顯著影響整個集群的性能。這些straggler的產生可能由多種因素引起,，例如GPU過熱,、后臺進程干擾、資源競爭或未知故障等,。目前,，常見的解決方案是將掉隊者所在的整個數據并行組（replica）移除，但這種做法會導致大量正常節(jié)點無法得到充分利用,。我們意識到,，一方面，性能下降的straggler仍然可以被有效利用,；另一方面,，對于性能嚴重下降的straggler，僅需移除其自身,，而不必移除整個replica,。

為此，我們通過引入四個維度的異構（data,、layer,、device和stage），設計了一個能夠感知straggler并細粒度調整異構訓練策略的系統(tǒng)Malleus,。實驗表明,，在多種straggler場景下,，與Megatron-LM和DeepSpeed相比,，我們的系統(tǒng)能夠實現(xiàn)2.63-5.28倍的加速比，并具備更快的恢復速度,。

2. PQCache: Product Quantization-based KVCache for Long Context LLM Inference

作者：Hailin Zhang, Xiaodong Ji, Yilin Chen, Fangcheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, Bin Cui

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

Github鏈接：https://github.com/HugoZHL/PQCache

圖2. PQCache概覽

近年來,，大型語言模型 (LLM) 的上下文窗口不斷擴展，從最初的幾千token發(fā)展到如今的百萬token級別,。然而,，GPU顯存的限制使得LLM推理過程中token的中間表示——鍵值緩存KVCache成為主要的內存瓶頸，嚴重制約了模型的實際應用?，F(xiàn)有方法嘗試通過選擇性保留部分tokens的KVCache來緩解這一問題,，但往往面臨模型質量下降或服務延遲增加的困境。

借鑒數據管理領域的信息檢索技術，我們創(chuàng)新性地將KVCache的管理問題轉化為嵌入檢索任務,。我們提出了PQCache方法,，基于乘積量化（PQ）來管理KVCache，在確保低服務延遲的同時保持模型質量,。在預填充階段,，我們將PQ應用于每個LLM層中自注意力頭的所有token的鍵。在自回歸解碼階段,，我們使用PQ編碼和聚類中心來近似檢索重要的token,，然后獲取相應的鍵值對進行自注意計算。通過精心設計重疊和緩存,，我們最大限度地減少了兩個階段的任何額外計算和通信開銷,。大量實驗表明，PQCache 兼具有效性和效率,，在多個長文本benchmark和任務上表現(xiàn)優(yōu)異,，并且在預填充和解碼階段都具有較低的系統(tǒng)延遲。

實驗室簡介

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