解决长序列Transformer训练效率低问题，提升大规模模型训练速度。
HelixPipe: Efficient Distributed Training of Long Sequence Transformers with Attention Parallel Pipeline Parallelism
- 分阶段并行计算注意力，减少流水线空档
- 在64张H20 GPU上训练128k序列时提速26%
- 适合超长序列大模型分布式训练场景
随着Transformer序列长度增长，现有流水线并行方法因注意力计算的二次复杂度和高内存开销而表现不佳。为此，我们提出HelixPipe，一种面向长序列Transformer训练的新流水线并行方案。首先，引入注意力并行分区，将不同微批次的注意力计算分配到不同流水线阶段并行执行，降低流水线空档。其次，采用两阶段先进后出微批次调度策略，在平衡内存使用的同时实现通信与计算重叠。此外，通过不带注意力的重新计算和分块MLP缓解内存碎片，支持更长序列。实验表明，随着序列长度增加，HelixPipe优势愈发显著，在不同流水线规模、模型大小和集群配置下均展现出更高吞吐量与可扩展性。特别地，在64张H20 GPU上训练70亿参数模型且序列长度达128k时，相比基线方法提升26%性能。代码已开源。
原文摘要 · Abstract (English)
As transformer sequence lengths grow, existing pipeline parallelisms incur suboptimal performance due to the quadratic attention computation and the substantial memory overhead. To relieve these challenges, we propose HelixPipe, a novel pipeline parallelism for long sequence transformer training. First, HelixPipe introduces attention parallel partition, which schedules attention computations of different micro batches across different pipeline stages in parallel, reducing pipeline bubbles. Second, it employs a two-fold first-in-last-out micro batch schedule to balance memory usage and overlap communication with computation. Additionally, HelixPipe utilizes recomputation without attention and chunked MLP to mitigate fragmentation and enable longer sequences. Experiments demonstrate that HelixPipe gains increasing advantages with longer sequence lengths, and outperforms existing methods in throughput and scalability across varying pipeline sizes, model sizes, and cluster configurations. Notably, it achieves a 26\% speedup over baseline methods when training a 7B model with 128k sequence length on 64 H20 GPUs. Code is available at https://github.com/code-tunnel/Megatron-LM/tree/dev.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。