PKU-DAIR实验室两项成果被ASPLOS 2025录用
ASPLOS (ACM International Conference on Architectural Support for Programming Languages and Operating Systems) 是计算机科学领域顶级的国际学术会议之一，专注于计算机体系结构、编程语言与操作系统等领域。ASPLOS是计算机系统领域最具影响力的会议之一，属于计算机科学领域的A类会议（中国计算机学会CCF评定的A类会议），录用的难度相对较大，接受率通常较低，约为15%到20%之间，每年接收的论文约为100~150篇。PKU-DAIR实验室《Spindle: Efficient Distributed Training of Multi-Task Large Models via Wavefront Scheduling》和《FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism》两篇论文被ASPLOS 2025录用！论文核心作者包括北大PKU-DAIR实验室四年级博士生王驭捷、三年级硕士生竺沈涵、北航本科四年级科研实习生王士举等。欢迎对分布式深度学习系统、高效大模型框架感兴趣的学界业界人士关注我们的工作！
论文接收信息如下：
1. Spindle: Efficient Distributed Training of Multi-Task Large Models via Wavefront Scheduling
作者：Yujie Wang, Shenhan Zhu, Fangcheng Fu, Xupeng Miao, Jie Zhang, Juan Zhu, Fan Hong, Yong Li, Bin Cui
论文链接：https://arxiv.org/abs/2409.03365
近期最先进的大规模基座模型具备了同时理解和处理多个任务、多种数据模态的强大能力，例如语言任务、图像任务、视频任务、语音任务等。这类模型通常由一个统一的基础模型结构和多个专门化的模型组件组成，其模型结构复杂，图1展示了多任务多模态大模型的复杂结构。现有的大模型训练系统均主要针对单一任务、单一模态的模型而设计，而如此复杂的多任务多模态大模型会对高效训练系统的设计带来巨大挑战，一方面其复杂模型结构和多任务多模态特性使其存在严重的负载异构性，导致现有系统训练这类模型容易造成资源浪费和低效，另一方面复杂的模型结构同时也带来复杂的执行依赖，而现有系统无法高效处理这种依赖。
图1：多任务多模态大模型的复杂模型结构
图2：现有系统训练多任务多模态大模型的资源低效和浪费
为此，本工作设计了一个全新的训练系统——Spindle，旨在通过波面调度（Wavefront Scheduling）的方法实现多任务多模态大模型高效分布式训练。Spindle的核心思想是将负载异构和执行依赖的复杂模型分解为多个顺序执行的波（Wave）。具体而言，波（Wave）是Spindle中最小粒度的执行调度单位，在每个波内，集群中会存在多个算子并行执行，这些算子各自占据一部分集群资源并具有相近的执行时间开销，一个波代表着一种集群资源分配模式。Spindle将复杂的模型执行分解为多个波，并顺序地解决联合优化问题，包括异构负载感知的并行化和依赖驱动的执行调度，从而实现多任务大模型的高效训练。
图3：Spindle系统框架概览
图4：Spindle基于波面调度的训练执行规划示意图
图3展示了Spindle的系统框架架构，其由执行规划器（Execution Planner）和训练框架（Training Framework）组成，执行规划器分为五个系统组件，包括：图收缩组件（Graph Contraction）负责对庞大的计算图进行收缩简化，可扩展性估计器（Scalability Estimator）负责对负载异构的算子进行执行开销和可扩展性的准确估计，资源分配器（Resource Allocator）负责对算子分配合适的计算资源量，波面调度器（Wavefront Scheduler）负责将算子组织成多个顺序执行的波并生成波面调度方案，设备放置器（Device Placement）负责将各个算子放置到GPU设备上并生成最终的执行规划，并交由运行时引擎（Runtime Engine）进行高效的多任务大模型训练。图4给出了一个Spindle基于波面调度的训练执行规划的示意图，其包含了六个顺序执行的波。
图5：Spindle相比于现有系统的端到端性能对比实验
我们实现并构建了该系统，并在多种多任务多模态模型上进行了评估。图5展示了Spindle相比于现有系统的端到端性能对比实验，实验结果表明，Spindle在性能和效率方面优于现有的训练系统（例如Megatron-LM和DeepSpeed），加速比最高可达71%。
2. FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism
作者：Yujie Wang, Shiju Wang, Shenhan Zhu, Fangcheng Fu, Xinyi Liu, Xuefeng Xiao, Huixia Li, Jiashi Li, Faming Wu, Bin Cui
论文链接：https://arxiv.org/abs/2412.01523
代码链接：https://github.com/PKU-DAIR/Hetu-Galvatron
随着大语言模型（LLMs）的快速发展和取得的巨大成就，扩展其上下文长度（即最大支持的序列长度）成为了一个迫切需求。为支持LLM的长上下文训练，序列并行作为关键技术应运而生，它将每个输入序列切分并分散到多个设备上，并通过必要的通信来处理这些序列。然而，现有的序列并行方法假设输入序列是同质的（即所有序列长度一致），并采用单一静态并行策略，这在实际应用中是低效的。
图1：大语言模型训练语料库中序列长度的长尾分布
图2：不同序列并行度下不同序列的训练时间和通信时间占比
具体而言，在真实世界的大语言模型训练语料库中，序列长度表现出显著的差异，通常呈长尾分布（如图1所示），短序列占绝大部分，而长序列较少，这种差异性导致训练负载的异构性。长序列需要更大的并行度来分摊显存开销，但这会受到跨机通信带宽的限制，导致较高的额外通信开销（如图2展示）。相对而言，低并行度策略更为高效。然而，现有同构系统在训练中只能使用单一静态的策略，这使得短序列也不得不采用低效的高并行度策略，而短序列又占据了训练数据集的绝大部分，这导致现有同构系统处理真实异构数据时效率很低。
图3：FlexSP异构自适应序列并行的例子
为了解决这一问题，我们首次提出了一个基于异构理念设计的训练系统FlexSP，其使用异质自适应的序列并行方法来处理不同序列之间的异质工作负载。图3展示了一个简单的例子，表明对于不同长度的序列采用合适的并行策略能够提高训练效率。我们的系统会在每轮训练迭代中捕捉序列长度的异质性，并根据工作负载特征分配最优的异构序列并行策略组合。这种基于异构理念设计的并行范式和训练系统能够天然捕捉异构负载的特征，并实现高效训练。
图4：FlexSP的系统架构图
图4展示了我们系统FlexSP的系统架构图，其由求解器和执行器组成，求解器包含：并行规划器（Parallelism Planner）负责将异构序列并行的问题建模为线性规划优化问题并高效地求解，其包括了异构序列并行组的构建和时间均衡的并行组序列分配；序列分组器（Sequence Blaster）负责将输入的数据批次（Batch），切分为若干个数据微批次（Micro-batch），使得集群的显存能够容纳每一个微批次，并将每个微批次交由并行规划器进行异构策略的求解。求解器生成的最优规划会交由执行器进行训练。值得一提的是，FlexSP通过数据预取和求解-训练解耦的方式，将求解开销（Overlap）完全隐藏于训练时间中，因此不会引起任何额外开销。
图5：FlexSP相比于SOTA系统的端到端性能实验对比
我们构建了基于异构理念设计的LLM训练系统FlexSP，并在多个数据集、多种大小模型上进行了实验，图5展示了FlexSP相比于SOTA训练系统（例如Megatron-LM, DeepSpeed）的性能对比，实验结果表明，我们的系统在性能上比现有的最先进训练框架提高了最多1.98倍。
实验室简介
北京大学数据与智能实验室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR实验室）由北京大学计算机学院崔斌教授领导，长期从事数据库系统、大数据管理与分析、人工智能等领域的前沿研究，在理论和技术创新以及系统研发上取得多项成果，已在国际顶级学术会议和期刊发表学术论文200余篇，发布多个开源项目。课题组同学曾数十次获得包括CCF优博、ACM中国优博、北大优博、微软学者、苹果奖学金、谷歌奖学金等荣誉。PKU-DAIR实验室持续与工业界展开卓有成效的合作，与腾讯、阿里巴巴、苹果、微软、百度、快手、中兴通讯等多家知名企业开展项目合作和前沿探索，解决实际问题，进行科研成果的转化落地。
168
2025-02-21
