Research map / Large language models / Instruction and fine-tuning
Efficient Long Context Fine-tuning with Chunk Flow
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ChunkFlow, a chunk-centric training system that repacks variable-length sequences into uniform chunks and uses state-aware scheduling plus a state-aware 1F1B pipeline schedule to speed up long-context fine-tuning.
Problem
Long context fine-tuning uses datasets dominated by short sequences with a small long tail; existing strategies tuned for long sequences cause GPU memory underutilization, excess parallelism for short data, load imbalance in data parallelism, and pipeline bubbles.
Approach
Reorganize each batch into uniformly sized chunks (packing short sequences together, splitting long ones) and schedule dependent chunks with a state-aware scheme that stores only K chunks' activations, so peak memory is governed by ChunkSize*K rather than the longest sequence; combine this with 1F1B pipeline scheduling.
Key results
- ChunkFlow achieves up to 4.53x speedup over Megatron-LM when fine-tuning Qwen2.5 models (7B to 72B) at 32K and 256K context lengths.
- Peak memory for 7B at 32K context rises from 41.6 GiB (ChunkSize 2K) to 59.3 GiB (ChunkSize 8K), largely independent of maximum sequence length.
- State-aware 1F1B reduces the pipeline bubble ratio from 57.14% (standard 1F1B) to 54.1% with K=1 and 47.8% with K=2 in the toy example.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ChunkFlow chunk construction and state-aware scheduling instead of Megatron-LM | Qwen2.5-series long-context fine-tuning · Long-context fine-tuning evaluation dataset | training speed: up to 4.53x | Helped |
| Use state-aware 1F1B chunk scheduling with ChunkSize 2*Unit and K 1 instead of standard 1F1B | Variable-length pipeline scheduling · Four sequences with lengths Unit, Unit, 2*Unit, and 4*Unit | bubble ratio: 57.14% -> 54.1% | Helped |
| Use state-aware 1F1B chunk scheduling with ChunkSize 2*Unit and K 1 instead of standard 1F1B | Variable-length pipeline scheduling · Four sequences with lengths Unit, Unit, 2*Unit, and 4*Unit | efficiency: approximately 8% | Helped |
| Increase K from 1 to 2 in state-aware 1F1B chunk scheduling (ablation) | Variable-length pipeline scheduling · Four sequences with lengths Unit, Unit, 2*Unit, and 4*Unit | bubble ratio: 54.1% -> 47.8% | Helped |
| Increase K from 1 to 2 in state-aware 1F1B chunk scheduling (ablation) | Variable-length pipeline scheduling · Four sequences with lengths Unit, Unit, 2*Unit, and 4*Unit | end-to-end efficiency: 12% | Helped |
| Set ChunkSize to 4*Unit and K to 1, yielding only 2 chunks (ablation) | ChunkFlow state-aware pipeline scheduling · Four-sequence pipeline scheduling example | training performance: 15% | Hurt |
| Set ChunkSize to 4*Unit and K to 1, yielding only 2 chunks (ablation) | ChunkFlow state-aware pipeline scheduling · Four-sequence pipeline scheduling example | bubble ratio: 60% | Hurt |
| Use ChunkSize 2K and K 16 instead of ChunkSize 8K and K 4 (ablation) | ChunkFlow on Qwen2.5-7B · 256K context | average iteration time: 23774 -> 29810 | Hurt |
| Use ChunkSize 32K and K 1 instead of ChunkSize 8K and K 4 (ablation) | ChunkFlow on Qwen2.5-7B · 256K context | average iteration time: 23774 -> 28942 | Hurt |
| Increase ChunkSize from 2K to 4K to 8K (ablation) | ChunkFlow on Qwen2.5-7B · 32K context | peak memory: 41.6 GiB; 47.5 GiB; 59.3 GiB | Hurt |
| Increase ChunkSize from 2K to 4K to 8K (ablation) | ChunkFlow on Qwen2.5-7B · 256K context | peak memory: 45.6 GiB; 50.8 GiB; 63.8 GiB | Hurt |
Surprising findings
- Setting ChunkSize=4*Unit with K=1 in the toy example produced only 2 chunks and degraded performance by 15% relative to standard 1F1B, showing that too-large chunks can hurt pipelining.
Limitations
- Pipeline integration implemented only for standard 1F1B; advanced pipeline schedules left for future work.
- K/V tensors are kept in memory without offloading, so long-context memory grows somewhat with max sequence length.
- ChunkSize and K require grid search per training configuration.
Experimental setup
Datasets: LMSysChat1M constructed long-context evaluation dataset (context lengths 32K and 256K)
Metrics: average iteration time (training performance) peak memory usage pipeline bubble ratio speedup over Megatron-LM
Baselines: Megatron-LM
Builds on: Megatron-LM 1F1B pipeline scheduling sequence packing token-level pipeline parallelism
Closest papers by meaning
- Skrull: Towards Efficient Long Context Fine-tuning through Dynamic Data Scheduling (NeurIPS 2025 · similarity 0.91)
- Training-Free Long-Context Scaling of Large Language Models (ICML 2024 · similarity 0.91)
- Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum (NeurIPS 2024 · similarity 0.90)
- A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts (ICLR 2025 · similarity 0.89)
- LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models (ICLR 2024 · similarity 0.89)
- InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.