How ChunkFlow Boosts Long-Context Model Training Up to 4.5× Faster

Research Background

Long‑text capability is a core strength of large language models (LLMs) and is essential for many downstream tasks. Continue pre‑training and long‑context fine‑tuning are key to extending this ability, but real‑world datasets exhibit a long‑tail distribution where short samples dominate and a few ultra‑long sequences cause severe performance bottlenecks.

Existing Issues

1. Fixed memory and parallel strategies designed for the longest sequence conflict with variable‑length data, leading to load imbalance and pipeline bubbles that degrade training efficiency.

2. Ultra‑long sequences generate massive activation memory, forcing recompute or offload techniques and further worsening memory imbalance, especially in pipeline parallelism.

ChunkFlow Solution

ChunkFlow reorganizes training data into chunks of a predefined ChunkSize. Short sequences are concatenated, long sequences are split, and the resulting chunks are scheduled with a state‑aware mechanism that preserves computational correctness via attention masks and handles dependent chunks through recompute and offload.

The approach ensures that memory consumption scales with the fixed ChunkSize rather than the maximum sequence length, providing predictable GPU usage.

Experimental Results

Benchmarks on various Qwen 2.5 models with 32K and 256K context lengths show that ChunkFlow outperforms Megatron‑LM, delivering up to 4.53× end‑to‑end training speedup. Varying ChunkSize demonstrates controllable peak memory, improving robustness across different workloads.

ChunkFlow now powers all Qwen series for supervised fine‑tuning (SFT) and long‑sequence CPT tasks, consistently yielding more than 2× performance gains and substantial GPU cost savings.

Paper Information

Title: Efficient Long Context Fine-tuning with Chunk Flow

Authors: Xiulong Yuan, Hongtao Xu, Wenting Shen, Ang Wang, Xiafei Qiu, Jie Zhang, Yuqiong Liu, Bowen Yu, Junyang Lin, Mingzhen Li, Weile Jia, Yong Li, Wei Lin

Link: https://arxiv.org/pdf/2503.02356

Illustrations