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实验室持续与工业界展开卓有成效的合作，与腾讯、阿里巴巴、苹果、微软、百度、快手、中兴通讯等多家知名企业开展项目合作和前沿探索，解决实际问题，进行科研成果的转化落地。