Efficient Long Context Fine-tuning with Chunk Flow

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Yuan, Xiulong, Xu, Hongtao, Shen, Wenting, Wang, Ang, Qiu, Xiafei, Zhang, Jie, Liu, Yuqiong, Yu, Bowen, Lin, Junyang, Li, Mingzhen, Jia, Weile, Li, Yong, Lin, Wei
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866909683517751296
author Yuan, Xiulong
Xu, Hongtao
Shen, Wenting
Wang, Ang
Qiu, Xiafei
Zhang, Jie
Liu, Yuqiong
Yu, Bowen
Lin, Junyang
Li, Mingzhen
Jia, Weile
Li, Yong
Lin, Wei
author_facet Yuan, Xiulong
Xu, Hongtao
Shen, Wenting
Wang, Ang
Qiu, Xiafei
Zhang, Jie
Liu, Yuqiong
Yu, Bowen
Lin, Junyang
Li, Mingzhen
Jia, Weile
Li, Yong
Lin, Wei
contents Long context fine-tuning of large language models(LLMs) involves training on datasets that are predominantly composed of short sequences and a small proportion of longer sequences. However, existing approaches overlook this long-tail distribution and employ training strategies designed specifically for long sequences. Moreover, these approaches also fail to address the challenges posed by variable sequence lengths during distributed training, such as load imbalance in data parallelism and severe pipeline bubbles in pipeline parallelism. These issues lead to suboptimal training performance and poor GPU resource utilization. To tackle these problems, we propose a chunk-centric training method named ChunkFlow. ChunkFlow reorganizes input sequences into uniformly sized chunks by consolidating short sequences and splitting longer ones. This approach achieves optimal computational efficiency and balance among training inputs. Additionally, ChunkFlow incorporates a state-aware chunk scheduling mechanism to ensure that the peak memory usage during training is primarily determined by the chunk size rather than the maximum sequence length in the dataset. Integrating this scheduling mechanism with existing pipeline scheduling algorithms further enhances the performance of distributed training. Experimental results demonstrate that, compared with Megatron-LM, ChunkFlow can be up to 4.53x faster in the long context fine-tuning of LLMs. Furthermore, we believe that ChunkFlow serves as an effective solution for a broader range of scenarios, such as long context continual pre-training, where datasets contain variable-length sequences.
format Preprint
id arxiv_https___arxiv_org_abs_2503_02356
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Efficient Long Context Fine-tuning with Chunk Flow
Yuan, Xiulong
Xu, Hongtao
Shen, Wenting
Wang, Ang
Qiu, Xiafei
Zhang, Jie
Liu, Yuqiong
Yu, Bowen
Lin, Junyang
Li, Mingzhen
Jia, Weile
Li, Yong
Lin, Wei
Distributed, Parallel, and Cluster Computing
Long context fine-tuning of large language models(LLMs) involves training on datasets that are predominantly composed of short sequences and a small proportion of longer sequences. However, existing approaches overlook this long-tail distribution and employ training strategies designed specifically for long sequences. Moreover, these approaches also fail to address the challenges posed by variable sequence lengths during distributed training, such as load imbalance in data parallelism and severe pipeline bubbles in pipeline parallelism. These issues lead to suboptimal training performance and poor GPU resource utilization. To tackle these problems, we propose a chunk-centric training method named ChunkFlow. ChunkFlow reorganizes input sequences into uniformly sized chunks by consolidating short sequences and splitting longer ones. This approach achieves optimal computational efficiency and balance among training inputs. Additionally, ChunkFlow incorporates a state-aware chunk scheduling mechanism to ensure that the peak memory usage during training is primarily determined by the chunk size rather than the maximum sequence length in the dataset. Integrating this scheduling mechanism with existing pipeline scheduling algorithms further enhances the performance of distributed training. Experimental results demonstrate that, compared with Megatron-LM, ChunkFlow can be up to 4.53x faster in the long context fine-tuning of LLMs. Furthermore, we believe that ChunkFlow serves as an effective solution for a broader range of scenarios, such as long context continual pre-training, where datasets contain variable-length sequences.
title Efficient Long Context Fine-tuning with Chunk Flow
topic Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2503.02356