Seq1F1B: Efficient Sequence-Level Pipeline Parallelism for Large Language Model Training

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Sun, Ao, Zhao, Weilin, Han, Xu, Yang, Cheng, Zhang, Xinrong, Liu, Zhiyuan, Shi, Chuan, Sun, Maosong
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929584901980160
author Sun, Ao
Zhao, Weilin
Han, Xu
Yang, Cheng
Zhang, Xinrong
Liu, Zhiyuan
Shi, Chuan
Sun, Maosong
author_facet Sun, Ao
Zhao, Weilin
Han, Xu
Yang, Cheng
Zhang, Xinrong
Liu, Zhiyuan
Shi, Chuan
Sun, Maosong
contents The emergence of large language models (LLMs) relies heavily on distributed training strategies, among which pipeline parallelism plays a crucial role. As LLMs' training sequence length extends to 32k or even 128k, the current pipeline parallel methods face severe bottlenecks, including high memory footprints and substantial pipeline bubbles, greatly hindering model scalability and training throughput. To enhance memory efficiency and training throughput, in this work, we introduce an efficient sequence-level one-forward-one-backward (1F1B) pipeline scheduling method tailored for training LLMs on long sequences named Seq1F1B. Seq1F1B decomposes batch-level schedulable units into finer sequence-level units, reducing bubble size and memory footprint. Considering that Seq1F1B may produce slight extra bubbles if sequences are split evenly, we design a computation-wise strategy to partition input sequences and mitigate this side effect. Compared to competitive pipeline baseline methods such as Megatron 1F1B pipeline parallelism, our method achieves higher training throughput with less memory footprint. Notably, Seq1F1B efficiently trains a LLM with 30B parameters on sequences up to 64k using 64 NVIDIA A100 GPUs without recomputation strategies, a feat unachievable with existing methods. Our source code is based on Megatron-LM, and now is avaiable at: https://github.com/MayDomine/Seq1F1B.git.
format Preprint
id arxiv_https___arxiv_org_abs_2406_03488
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Seq1F1B: Efficient Sequence-Level Pipeline Parallelism for Large Language Model Training
Sun, Ao
Zhao, Weilin
Han, Xu
Yang, Cheng
Zhang, Xinrong
Liu, Zhiyuan
Shi, Chuan
Sun, Maosong
Distributed, Parallel, and Cluster Computing
The emergence of large language models (LLMs) relies heavily on distributed training strategies, among which pipeline parallelism plays a crucial role. As LLMs' training sequence length extends to 32k or even 128k, the current pipeline parallel methods face severe bottlenecks, including high memory footprints and substantial pipeline bubbles, greatly hindering model scalability and training throughput. To enhance memory efficiency and training throughput, in this work, we introduce an efficient sequence-level one-forward-one-backward (1F1B) pipeline scheduling method tailored for training LLMs on long sequences named Seq1F1B. Seq1F1B decomposes batch-level schedulable units into finer sequence-level units, reducing bubble size and memory footprint. Considering that Seq1F1B may produce slight extra bubbles if sequences are split evenly, we design a computation-wise strategy to partition input sequences and mitigate this side effect. Compared to competitive pipeline baseline methods such as Megatron 1F1B pipeline parallelism, our method achieves higher training throughput with less memory footprint. Notably, Seq1F1B efficiently trains a LLM with 30B parameters on sequences up to 64k using 64 NVIDIA A100 GPUs without recomputation strategies, a feat unachievable with existing methods. Our source code is based on Megatron-LM, and now is avaiable at: https://github.com/MayDomine/Seq1F1B.git.
title Seq1F1B: Efficient Sequence-Level Pipeline Parallelism for Large Language Model Training
topic Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2406.03488