一句话：提出 AdaPipe，通过自适应重计算与灵活阶段划分，优化流水线并行的内存与计算均衡，加速大模型训练。
大语言模型规模与序列长度增长带来巨大内存需求，需要更大规模并行系统。流水线并行虽被广泛采用，却引入更重且不均衡的内存消耗；重计算可缓解内存压力，但会带来额外计算开销。AdaPipe 旨在寻找优化的重计算与流水线阶段划分策略：它采用自适应重计算以最大化内存利用率并降低各流水线阶段的计算代价，同时用灵活的阶段划分算法平衡阶段间计算。作者在 NVIDIA GPU 与昇腾 NPU 集群上训练 GPT-3 175B 和 Llama 2 70B，分别取得最高 1.32× 与 1.22× 加速。
Large language models (LLMs) have demonstrated powerful capabilities, requiring huge memory with their increasing sizes and sequence lengths, thus demanding larger parallel systems. The broadly adopted pipeline parallelism introduces even heavier and unbalanced memory consumption. Recomputation is a widely employed technique to mitigate the problem but introduces extra computation overhead. This paper proposes AdaPipe, which aims to find the optimized recomputation and pipeline stage partitioning strategy. AdaPipe employs adaptive recomputation to maximize memory utilization and reduce the computation cost of each pipeline stage. A flexible stage partitioning algorithm is also adopted to balance the computation between different stages. We evaluate AdaPipe by training two representative models, GPT-3 (175B) and Llama 2 (70B), achieving up to 1.32× and 1.22× speedup on clusters with NVIDIA GPUs and Ascend NPUs respectively.
@inproceedings{SunCWF0WC24,
title = {AdaPipe: Optimizing Pipeline Parallelism with Adaptive Recomputation and Partitioning},
author = {Zhenbo Sun and Huanqi Cao and Yuanwei Wang and Guanyu Feng and Shengqi Chen and Haojie Wang and Wenguang Chen},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620666.3651359}
}