SIGMOD是计算机科学领域的顶级国际学术会议之一，专注于数据库系统和数据管理技术的研究。作为中国计算机学会（CCF）推荐的A类会议，SIGMOD与VLDB、ICDE并称为数据库领域的三大顶会，近年来论文接收率维持在17%-20%左右。PKU-DAIR实验室的论文《Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment》在SIGMOD第二轮评审中成功被录用。

Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment

作者：Haoyang Li, Fangcheng Fu, Sheng Lin, Hao Ge, Xuanyu Wang, Jiawen Niu, Jinbao Xue, Yangyu Tao, Di Wang, Jie Jiang, Bin Cui

代码链接：https://github.com/PKU-DAIR/Hetu

背景

在大规模Transformer模型训练中，输入序列往往存在显著的长度差异。传统做法是通过padding将较短的序列补齐到批次中的最大长度，以便统一并行计算，但这种方式会引入对填充token的冗余计算，造成算力浪费。为提升效率，现代训练系统广泛采用packing技术，即将多个不同长度的序列拼接成一个较长序列，再通过block-diagonal attention等机制保证语义隔离，从而减少无效计算。

目前，主流的训练系统通常采用两类方法：一方面，并行策略上采用静态、同构的配置，即所有训练流水线（replicas）在整个训练过程中都保持相同的并行度组合（如固定的tensor parallel与pipeline parallel度数）；另一方面，数据管理上使用max-length packing，即先设定一个最大序列长度（context length），再将不同长度的原始序列packing成不超过该阈值的长序列，并平均分配到各设备上。这种“静态同构并行+最大长度packing”的组合在实现上简洁，但却存在四种负载不均：

1. 迭代内采样不均衡（intra-iteration sampling imbalance）：同一批次中，短序列占主导地位，而长序列稀少。这种长度差异使得同构并行策略为了给长序列预留显存，不得不采取低效的资源分配方案，从而拖累了短序列的训练效率。

2. 迭代间采样不均衡（inter-iteration sampling imbalance）：不同批次的序列长度分布差异显著，最大序列长度波动剧烈。静态的并行策略通常按最长序列配置，结果在大多数短序列批次中浪费资源。

图1：迭代内与迭代间采样不均衡

     3. 流水线内 packing 不均衡（intra-pipeline packing imbalance）：packing虽然平衡了显存占用，但注意力机制的计算复杂度随序列长度平方增长，导致不同packed序列的计算时间差异极大，增加流水线气泡。

    4. 流水线间 packing 不均衡（inter-pipeline packing imbalance）：当不同流水线处理的packed序列复杂度差异较大时，最慢的流水线决定了整体迭代速度，从而拉低集群效率。

图2：流水线内与流水线间packing不均衡

这四类不均衡使得现有依赖静态同构并行与最大长度packing的方法在实际训练中效率低下。因此，需要在并行策略与数据管理层面协同优化，以适应真实世界中变长序列的特性。

方法

我们提出Hydraulis系统，从并行策略与数据分配两个维度协同设计，核心方法包括：

   1. 动态异构并行策略：针对采样不均衡，Hydraulis支持在不同迭代动态选择合适的并行策略；在同一迭代内，不同流水线可采用异构并行配置（如tensor/pipeline/context parallel度数不同），使长短序列各自匹配最优方案，从而同时缓解intra-iteration与inter-iteration的不均衡。

2.优化—传播解耦与子图机制：Hydraulis 通过将传播阶段（前向/反向计算）与优化阶段（梯度同步与更新）解耦，保持优化阶段的参数分片固定不变，同时允许传播阶段灵活切换异构并行配置。为进一步解决两者之间的异构通信问题，我们引入了 pull（拉取参数）和 push（同步梯度）操作。基于子图抽象，我们进一步统一表达了不同并行策略下的通信模式，实现了异构并行策略之间的无缝切换。

3. 两阶段序列分配（two-stage sequence assignment）：针对packing不均衡，Hydraulis先在流水线间分配原始序列以保证负载均衡（解决inter-pipeline imbalance），再在流水线内部通过基于整数线性规划的优化方法进行packing，使不同微批次的执行时间更加接近（解决intra-pipeline imbalance）。

4.数据分布感知的策略候选生成：结合数据集的全局序列长度分布，Hydraulis使用动态规划生成一组高质量的候选策略，训练时快速评估与选择，避免了迭代中穷举搜索的开销。

图3：Hydraulis系统框架

实验

在LLaMA2 7B、13B和32B模型上，以CommonCrawl和GitHub两个大规模数据集为基准，我们在64张Nvidia A800 GPU的集群中开展实验。结果显示，Hydraulis相较于Megatron-LM、DeepSpeed和HotSPa，在端到端性能上提升1.32–2.66倍。进一步实验表明，在 GPU 数量扩展、最大序列长度增大和batch size调整等不同条件下，Hydraulis均展现出良好的扩展性和稳定性。

图4：端到端实验对比

总  结

本文首次系统化地提出通过并行策略与数据管理的协同优化来应对大模型训练中的四类不均衡问题。Hydraulis通过动态异构并行以及合理的序列分配，有效提升了对变长序列训练的适应性，大幅改善了负载均衡。实验验证了其在真实大规模集群上的显著优势，为未来大模型训练在变长数据场景下的系统设计提供了新的方向。

实验室简介

北京大学数据与智能实验室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR实验室）由北京大学计算机学院崔斌教授领导，长期从事数据库系统、大数据管理与分析、人工智能等领域的前沿研究，在理论和技术创新以及系统研发上取得多项成果，已在国际顶级学术会议和期刊发表学术论文200余篇，发布多个开源项目。课题组同学曾数十次获得包括CCF优博、ACM中国优博、北大优博、微软学者、苹果奖学金、谷歌奖学金等荣誉。PKU-DAIR实验室持续与工业界展开卓有成效的合作，与腾讯、阿里巴巴、苹果、微软、百度、快手、中兴通讯等多家知名企业开展项目合作和前沿探索，解决实际问题，进行科研成果的转化落地。