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
Venue: ICML 2025 poster (2025)
Keywords: Long context fine-tuning
Large language model
Primary area: deep_learning->large_language_models
TL;DR: Efficient Long Context Fine-tuning with Chunk Flow

Abstract: 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.

Efficient Long Context Fine-tuning with Chunk Flow
Xiulong Yuan
Hongtao Xu 211 Wenting Shen'
Ang Wang' Xiafei Qiu' Jie Zhang' Yugiong Liu'
Bowen Yu! Junyang Lin Mingzhen Li
Weile Jia
Yong Li
Wei Lin
Abstract
1. Introduction
Long context fine-tuning of large language mod- els(LLMs) involves training on datasets that are predominantly composed of short sequences and small proportion of longer sequences. However, existing approaches overlook this long-tail tail distri- bution and employ training strategies designed specifically for long sequences. Moreover, these approaches also fail address the challenges posed by variable sequence lengths during dis- tributed training, such as load imbalance in data parallelism and severe pipeline bubbles in pipeline parallelism. These issues lead to suboptimal train- ing performance and poor GPU resource utiliza- tion. To tackle these problems, we propose a chunk-centric training method named ChunkFlow. ChunkFlow reorganizes input sequences into uni- formly sized chunks by consolidating short se- quences and splitting longer ones. This approach achieves optimal computational efficiency and bal- ance among training inputs. Additionally, Chunk- Flow incorporates a state-aware chunk scheduling mechanism to ensure that the peak memory us- age during training is primarily determined by the chunk size rather than the maximum sequence length in the dataset. Integrating this schedul- ing mechanism with existing pipeline scheduling algorithms further enhances the performance of distributed training. Experimental results demon- strate that, compared with Megatron-LM, Chunk- Flow can be up 4.53x faster in the long context fine-ture tine-tuning of LLMs. Furthermore, that ChunkFlow serves as an effective solution for broader range of scenarios, such as long con- text continual pre-traine pre-training, where datasets contain variable-length sequences.
The ability of large language models (L.LM)) to handle long contexts is crucial for tasks involving extensive inputs, such as comprehensive document analysis, extensive dialogue management, code generation from complex specifications, and detailed question answering over lengthy texts (Bai et al., 2024; OpenAI, 2024; Anthropic, 2025). Mainstream models like Llama can support context lengths of up 128K tokens, while Google's Gemini can manage sequences up to 1M tokens(Gemini Team, 2024; Meta, 2024). However, given the significant costs associated with directly training models on long contexts, current practices typically involve an initial pre-training phase using shorter contexts (e.g., Llama with a context length of 8K), followed by a continual training stage to progressively extend their context lengths. Finally, Long context fine-tuning(Long SFT), which the focus of this paper, is employed further enhance the model's effectiveness in processing long contexts (Alibaba, 2025b; Meta, 2024).
To improve models' ability to handle long sequences with- out compromising their performance on short ones, Long SFT is typically conducted on datasets predominantly com- prising short sequences with small percentage of long ones (Xiong et al., 2024). For instance, Llama3 was fine- tuned using 99.89% short sequences (averaging under K tokens) and 0.11% long sequences (averaging around 37K tokens) (Meta, 2024). Current training methods, however, often prioritize long sequences at the expense of the unique characteristics of the long-tail distribution. This approach results in suboptimal training efficiency and inefficient GPU resource utilization (Zhao et al., 2024). The micro-batch size for each training step is usually determined by the memory requirements of long sequences, leading to significant under- utilization of GPU memory when processing mostly shorter sequences. Additionally, handling long sequences typically demands more GPUs than processing primarily short se- quences. Existing methods allocate GPU resources based solely the needs of long sequences, resulting in finer partitioning of computations. Consequently, this not only di- minishes the training efficiency for short sequences but also exacerbates resource utilization issues. Moreover, variable sequence lengths in fne-tuning datasets also pose signifi- cant challenges in distributed training. For example, when
believe
Equal contribution Alibaba Group School of Advanced Inter- disciplinary Sciences, University of Chinese Academy of Sciences 2 State Key Lab of Processors, Institute of Computing Technol- ogy, CAS. Correspondence to: Yong Li -jjufenn.l@@alibbbb @aalibaba- inc.com>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Efficient Long Context Fine-tuning with Chunk Flow
employing data parallelism, it causes load imbalance among data parallel ranks. When pipeline parallelism is used, it leads to severe pipeline bubbles. Although some studies aim to address the load imbalance issue (Bai et al., 2024), existing approaches fail to effectively tackle the pipeline bubble problem.
mechanisms capture relationships between elements in sequence. However, while self-attention is powerful, it inherently allows each token to attend to all other tokens in the sequence, including future ones. This characteristic is problematic for tasks like language modeling (OpenAl, 2024; Meta, 2024; Alibaba, 2025b), where predictions must be made sequentially-one token at a time-without access to future information. To address this limitation, causal self-attention introduces causal mask, enforcing an auto- regressive property that ensures each token can only attend to previous tokens in the sequence. In this work, we leverage this causal property of LLMs to process long sequences in a chunk-by-chund manner.
In this work, we introduce ChunkFlow to address the aforementioned issues. During each training step, given sampled batch of sequences, ChunkFlow first reorganizes these sequences into new chunks. Specifically, it packs multiple short sequences into single chunks and divides longer sequences into several smaller ones, ensuring that the total sequence length of each chunk is approximately equal and does not exceed predefined ChunkSize. The ChunkSSze is determined based on the training configu- ration and GPU memory constraints, aiming to strike optimal balance between GPU computational efficiency and memory consumption. Next, ChunkFlow employs state-aware chunk scheduling approach to manage these chunks during forward and backward passes. This schedul- ing method not only correctly handles computational depen- dencies among chunks from the same long sequence but also ensures nearly constant memory consumption mainly reg- ulated by ChunkiSize, regardless of the original sequence lengths. Furthermore, we integrate our aware schedul- ing algorithm with existing pipeline scheduling algorithms to significantly boost the performance of distributed training on sequences of variable lengths. In this paper, we imple- ment a aware IFIB (Shoeybi et al.) chunk schedul- ing method, reserving the integration with more advanced pipeline scheduling algorithms for future work. By employ- ing this chunk-centered training method, ChunkFlow signif- icantly improves end-to-end training performance for long context fine-tuning. Extensive experiments demonstrate that, compared with Megatron-LM, ChunkFlow achieves up to 4.53x speedup in training performance for fine-tuning various sizes of Qwen2.5-series LLMs with different levels of context lengths. Moreover, we believe that ChunkFlow serves as an effective solution across wide range of con- texts, especially when dealing with training datasets contain- ing sequences of variable lengths. This not only highlights its efficiency but also underscores its adaptability across diverse scenarios, showcasing ChunkFlow's 's versatility and broad applicability.
2.2. Sequence Packing
Sequence packing (Kosec et al.. 2021) is widely-used training technique for LLMS when dealing with datasets that contain variable-length sequences. This method in- volves concatenating multiple sequences within batch into single sequence, therehy eliminating the need for padding. As result, it reduces redundant computations and mem- ory usage. In this paper, we adopt sequence packing as the default approach.
2.3. Distributed Parallelism Strategies
In this section, we present commonly employed paralleliza- tion strategies for training LLMs These strategies are fre- quently combined to achieve optimal training performance.
Data Parallelism (DP): In data parallelism(Li et al., b; Rajbhandari et al.; Zhao al.), the dataset is split into smaller subsets, and each subset is processed independently by replica of the model running on different devices.
Model Parallelism (MP): Model Parallelism is a critical technique for training large-scale deep learning models that exceed the memory capacity of a single device. It involves splitting a model into smaller components and distributing them across multiple devices, enabling the training of mod- els with billions or even trillions of parameters. Model paral- lelism can be broadly categorized into two main approaches: Tensor (rrallelssm(P))Sheeyb et al.) and Pipeline Par- allelismPPP)(Hunnn et al.: Narayanan et al., 2021). Ten- sor Parallclism(TP) divides individual layers or operations within a model across devices, ensuring that each device handles a portion of the tensor operations. Pipeline Paral- lelism(PP).. on the other hand, splits the model into sequen- tial stages, where each stage consists of one more layers and is assigned to different device(Fan et al.; Qi et al.). During training, the input data is processed stage-by-stage, with intermediate activations passed between devices. How- ever, PP introduces inefficiencies known as pipeline bubbles, which occur when some devices remain idle during certain computational phases(Qi al.). These bubbles reduce hard-
2. Preliminaries
In this section, we briefly walk through the preliminary literature related to this work.
2.1. Transformer and Causal Self-Attention
The Transformer architecture (Vaswani et al., 2017) has revolutionized deep learning and relies on self-attention
2

---

Efficient Long Context Fine-tuning with Chunk Flow
ware utilization and hinder training efficiency, making it essential to minimize pipeline bubbles for achieving high training performance and efficient utilization of distributed resources.
long-tail distribution of sequence lengths, posing signifi- cant challenges for efficient training and resource utilization. Notably, this distribution pattern has also been observed by Meta(Meta, 2024) as well as in our in-house propri- etary training dataset, which is specifically collected for fine-tuning LLMs with context length over 256K.
Sequence Parallelism (SP): Sequence Parallelism(Li et al., a; Korthikanti et al.; Liu et al.; Jacobs et al.) is a technique used to distribute the computation of long sequences across multiple devices in LLMs. Instead of processing an entire sequence on single device, the sequence is divided into smaller chunks, and each chunk is processed independently on different devices. This approach reduces memory usage per device and enables efficient training of models with long input sequences.
Sequence Length
Proportion of Sequences
< 1K
90,499%
< 4K
99.539%
< 8K
99.908% %
< 32K
99.987%
< 128K
99.9969
Longest
303K
Token-Level Pipeline Parallelism: Token-Level Pipeline Parallelism(Li al., c; Ao et al., 2025; Ma et al.) is novel approach to parallelizing the computation of LLMs by ploiting the unique properties of causal attention In this method, the input sequence is divided into smaller chunks at the token level, and these chunks are processed sequen- tially across multiple devices in a pipeline fashion. The key innovation lies in leveraging the causal masking property of causal attention, which ensures that each token only depends on its preceding tokens and not on future ones. This allows different pipeline stages to process distinct parts of the se- quence independently, without violating the auto- rugo-reegeesive nature of the model. By assigning different token chunks to different pipeline stages, token-level pipeline parallelism reduces the memory burden on individual devices and im- proves hardware utilization. Furthermore, it reduces idle time (pipeline bubbles) by carefully scheduling computa- tions and overlapping communication with computation. Overall, this approach enables efficient training of long sequences while maintaining the sequential dependencies required by causal attention.
Table 1. Sequence Length Distribution in LMSysChatIM
Apply Training Strategies Tailored For Long Sequences Causes Severe GPU Resource Un- derutilization and Poor Fine-tuning Efficiency. When fine-tuning the Qwen2.5-7B model on the LMSysChatIM dataset with a context length of 32K (excluding sequences longer than 32K), we use GPUs and set the global batch size to 64. If we directly apply a training strategy designed for long sequences, the micro-bated size has to be set to 1, and gradient accumulation over 64 micro-steps is required to prevent Out-Of- sequences of 32K length. However, nearly 90% of the se- quences in the dataset are shorter than IK tokens. Given that memory consumption is proportional to sequence length, this approach results in significant underutilization of GPU resources during most micro-steps. Figure 1 depicts the memory usage across randomly selected set of 1000 con- secutive training micro-steps. The results reveal that while peak memory usage can soar up to 75GB, a staggering 97.7% of these micro-steps consume less than 45GB of memory. This clearly indicates suboptimal utilization of GPU resources.
Observation 2:
(OOM) errors when processing
3. Observations
Before delving into the design of ChunkFlow, we first present three key observations that motivate our work. To better illustrate these points, we use the example of train- ing the Qwen2.5-7B model on the popular dataset LM- SysChatIM using Megatron-LM. However, it is crucial to emphasize that these insights are not limited to this spe- cific scenario; rather, they are applicable to all long-context fine-tuning tasks.
Mernary Footprint of Gifferent Iterations
Device
70
Manoory Foettrent 5010
Observation 1: Extremely Long-tailed Sequence Length Distribution Characteristics in Datasets.
Table 1 shows sequences distribution statistics in LM- S M dataset. It can be observed that over 99% of the sequences in the dataset are shorter than 4K tokens,
200
420
820
850
1000
Keration Stags
Figure I. Memory Footprints Different Iterations
while the longest sequence extends approximately 300K tokens. This pronounced disparity results in an extremely
Worse still, if we fine-tune the model with a 256K context
3

---

Efficient Long Context Fine-tuning with Chunk Flow
length (excluding sequences longer than 256K), we must allocate 16 GPUs avoid OOM errors when handling se- quences over 32K, even though these long sequences make up only 0.013% of the dataset. This approach requires dis- tributing computations across 16 GPUs, significantly degrad- ing the training performance for sequences shorter than 32K. Since sequences under 32K tokens dominate the dataset, us- ing 16 GPUs not only reduces overall end-to-end training efficiency but also worsens resource utilization. Experi- ments show that partitioning computations across 16 GPUs instead of 4 GPUs results in roughly 65% drop in training performance for sequences under 32K tokens.
bubble ratio should be ( et al.; Narayanan al., 2021). This discrepancy highlights that variable sequence lengths significantly exacerbate the pipeline bubble prob- lem.
4. ChunkFlow
besstion Start Semple Batch Secuensce
Chunk Construction Algonthi
Commrnust Churiz From irpul Sacusercss
ChunkSize
Observation 3: Variable Sequence Length Cause Severe Pipeline Bubbles. The variation in sequence lengths within fine-tuning datasets also poses significant challenges for efficient distributed training, such as load imbalance among data parallel ranks in data parallelism and severe pipeline bubbles in pipeline parallelism. In this section, we primarily focus on the pipeline bubble problem.
Coria uttid Churis
State Hate--wware Scheduling Algonthm
Scheotiing Chunis Traiing I Schwitslin Ruach Esccation Itanatio End
Figure 3. Overall Workflow in ChunkFlow
We illustrate the overall workflow of ChunkFlow in Fig- ure 3. At each training step, given a sampled batch of se- quences, ChunkFlow first reorganizes these sequences into a new set of chunks using a heuristic algorithm. This pro- cess involves merging multiple short sequences into single chunks and splitting longer sequences into several chunks as needed. Subsequently, a state-aware scheduling algorithm is employed to schedule these chunks for forward and back- ward passes. This scheduling approach not only properly manages computational dependencies between chunks from the same long sequence but also ensures nearly constant peak memory usage dictated by K * ChunkSize. Here, CleunkSize is the length limit of the constructed chunks, while к specifies how many chunks' activations should be saved by the scheduler. Additionally, we integrate Chunk- Flow's ss scheduling algorithm with existing pipeline schedul- ing methods, such as 1F1B, and implement a state-aware IFIB chunk scheduling method. This method can effec- tively enhance distributed training over variable-length se- quences. We plan to incorporate ChunkFlow's idea into more advanced pipeline scheduling algorithms in future work. What's 's more, it should be noted that gradients from each chunk are accumulated to ensure mathematical equiva- lence with existing training methods.
Figure 2. (a)
Sequences
with
Different with
Lengths:(b) Standard
IFIB Scheduling
Result
Pipeline Paraliel World Size =
We use the sequences in Figure 2(a) to explain pipeline parallel execution. The batch contains four sequences: two with Unit tokens and two with 2 * Unit and 4 * Unit to- kens, respectively. In addition to the commonly adopted assumption that the backward pass takes twice as long as the forward pass (Shoeybi et al.: Qi al.), we introduce another reasonable assumption: the execution time of dif- ferent sequences is proportional to their lengths. We also employ the bubble ratio, calculated according to Equation I, to quantify the proportion of GPU time wasted by pipeline bubbles during the entire batch execution.
Total Bubble Time
Bubble Ratio-
Total Execution Time
4.1. Chunk Construction
Figure 2(b) illustrates the scheduling results achieved by directly applying the standard 1FIB method to these four sequences with Pipeline Parallel W orld Size of 4. It can be observed that bubbles account for 57.14% of the total execution time. However, theoretically, when scheduling four sequences of equal length under this configuration, the
For a given batch of sequences, we employ the heuristic al- gorithm described in Algorithm 1 to reorganize them into a new set of chunks. Sequences that exceed the ShankSize are divided into multiple chunks. For the remaining se- quences shorter than the ChunkSSze, the algorithm treats chunk construction as a bin packing problem with two con-
4

---

Efficient Long Context Fine-tuning with Chunk Flow
Algorithm 1 ChankCoonstrccionnlgggittn
thus rely on other chunks' states (primarily key/value ten- sors and their corresponding gradients in the causal attention modules) from the same long sequence to ensure correct computation. For example, in Figure 4, Chunk 1, Chunk 2, and Chunk 3 are standalone chunks, whereas Chunk 4 to Chunk 7 are dependent chunks. During training, standalone chunks do not require special scheduling and naturally con- sume memory according to the ChunkSize. However, de- pendent chunks must be carefully scheduled to ensure that peak memory usage remains within the limits defined by the ChunkSize. Given chunks extracted from the same long sequence (indexed from I to N), these chunks need to be processed in ascending order for forward passes and in descending order for backward passes. This because, dur- ing forward propagation, subsequent chunks depend on the key/value tensors in the causal attention modules of preced- ing chunks, while during backward propagation, preceding chunks rely on the gradients of the key/value tensors from subsequent chunks. A naive scheduling approach would cause memory consumption to scale linearly with the length of the original sequence(Li et al., c; Ao et al., 2025).
1: Given ChunkSSzee List[sequence).
2: Return ResultChumks List [Chankk] as result.
3: LongSequences + Select sequences longer than ChuukSize.
4: ShortSequences + Select sequences shorter than ChunkSize.
5: for Sequence E LongSequences do
6: Divide Sequence by ChunkSize into multiple chunks and append chunks them to ResultChunks
7: end for
do 9: ResultBins + Try binpacking ShortSequence
8: for BinCnl = 1,...,
into BinCnt bins with ChunkSize as bin's max weight limit. If failed, continue trying for next BinCnt,, otherwise, we take the binpacked result
10: end for
11: for Bin RResultBins do
12: Pack sequences in Bin into single chunk and add
it to ResultChunks
13: end for
To address this challenge, ChunkFlow introduces a state- aware chunk scheduling method detailed in Algorithm 2. Given a list of dependent chunks (where standalone chunks can be considered special case with list size of 1) and a parameter K, our scheduling approach ensures that memory consumption scales with K * Chhnk (with K default- ing to 1) instead of the original sequence length. For scenar- ios where N = sizeof DependentChunkss and N > K, the forward passes of the first (N K) chunks are executed twice. During the initial forward pass, the activations of these chunks are discarded, while the causal attention mod- ules' key/value tensors are stored as state for subsequent reuse during the second forward pass. Applying this chunk scheduling algorithm to all dependent chunk groups yields the final chunk execution order result. Figure 5 illustrates the scheduling results for the chunks constructed in Figure 4, with K values set to 1 and 2, respectively. We observe that when K=1, Chunk 3 is executed twice, and at any given time during the entire execution, at most one chunk's 's ac- tivation is stored. However, when K=2, the activations of two chunks are retained, which results in higher memory consumption but also leads to improved end-to-end perfor- mance. It is worth noting that since Group-Query Atten- tion (Anniiiiiis et al., 2023) has been widely adopted in mainstream large language models (LLMs) such as Llama, Qwen, and Gemini, the storage of the corresponding atten- tion's 's key/value tensors and their gradients does not impose a significant memory overhead.
straints: the number of bins and the maximum weight limit per bin. The algorithm prioritizes minimizing the number of bins to maximize GPU computation efficiency.
Figure 4 illustrates an example of the chunk construction result from a batch of 16 input sequences. It shows that Sequence 6 is split into four chunks (Chunk 4 to Chunk 7), while the other shorter sequences are grouped into three chunks (Chunk 1 to Chunk 3).
2 & 0X
Coneruct Chunls from Sa mampled Boqueress
Churksial
Figure 4. Chunk Construction Result from Batch of 16 Sequences
4.2. State-Aware Chunk Scheduling
Given the chunks generated by Algorithm 1, we can observe two types of chunks: standalone chunks and dependent chunks. Standalone chunks contain complete sequences and can be processed independently. In contrast, dependent chunks contain segments of original long sequences and
In summary, this scheduling method ensures predictable memory usage pattern irrespective of the sequence length, while also achieving optimal computational efficiency.
5

---

Efficient Long Context Fine-tuning with Chunk Flow
Algorithm 2 ChunkSchhcduinnggAgorrnnn
Given DependentChunks List List(Chunkk) K.
2: StateStore for sharing states across Chunk's execu- tion
(8)
3: List for saving losses for each Chunk's execution <= K then
4: if size_of (Dependent
5: for Chank E DependentChuunks do
6: Loss = model model(Chunk;; StateStore)
K-2
7: Append Loss to LossTist
8: end for
9: for Loss reversed do Loss))
Figure 5. (a) Chunk Scheduling Result with 1;(b) Chunk
10:
Scheduling Result with K
11: end for
12:
else
13: for C'uunk E DependentChunns do
bubble ratio 47.8% resulting in 12% enhancement in end-to-end efficiency.
Loss model(Chunk, StateStore)
if Chunk.Idæ K then
16:
Append Loss to LissList else
17:
18:
Discard activations for Chunk end if
03130 0
6AR
1000
1000
Chand and BA
19:
20: end for
for Loss E Lousrsed(LossList) do wtw.:rddwwtt..rrddeet.ccocaaaaie end for
22:
001AR 282
2 22
BARRAR
23:
24:
for C'hunk E DependentChauks do
Figure 6. (a) State-Aware IFIB Scheduling Result with
25:
if Chunk..Iæ < K then
ChankSize = * Unit and K = 1;(b) State-Aware IF1B
26:
Loss = model (Channk; StateStore)
Scheduling Result with ChankSize 2 Unit and K
27:
backward with. gradient accumulation(los)
28:
end if
end for
5. Determining ChunkSize and K
30:
end if
The ChunkSize and K are primarily determined based on the training configuration and memory constraints, and they significantly impact end-to-end training performance. When pipeline parallelism is not utilized, K should always be set to 1, and ChunkSize should be maximized within memory constraints. This configuration aims to achieve maximum GPU utilization, thereby providing optimal training perfor- mance.
4.3. Incorperation with Pipeline Parallism: 1F1B Algorithm
State-Aware
As mentioned in Section 3, directly applying the standard pipeline scheduling method to variable-length sequences can result a relatively high bubble ratio, indicating wasted GPU idle time and poor end-to-end efficiency for distributed training. We incorporate ChunkFlow's 's idea with pipeline parallelism solve this problem. Specifically, in this paper, we combine our state-aware scheduling mechanism with the standard IFIB scheduling algorithm and implement a state-aware FFIB scheduling method for efficient pipeline parallel execution over variable-length sequences.
However, when pipeline parallelism is employed, must carefully consider these two key parameters. A too large ChunkSize results in fewer chunks, which can lead to more pipeline bubbles and degraded training perfor- mance. Conversely, a too small ChunkSize, while reduc- ing pipeline bubbles, may fail to fully utilize the GPU's computational efficiency, leading to suboptimal overall per- formance. For example, setting ChunkSize and K for the sequences in Figure 2 produces only 2 chunks. Scheduling them increases the bubble ratio to 60% and leads to a 15% performance degradation (as shown in Figure 7) compared with directly scheduling these four sequences using the standard IFIB method. For given training configuration, we employ grid search approach
For example, given the sequences in Figure our state- aware IFIB method operates as illustrated in Figure 6. When setting ChunkSize = 2, we obtain 4 chunks. Com- pared to the baseline method described in Section 3, our new scheduling approach achieves a reduced bubble ratio of 54.1% % and an approximately 8% improvement in efficiency with By increasing to 2, we further decrease the
Unit
6

---

Efficient Long Context Fine-tuning with Chunk Flow
Model
32K
256K
7B
selective selectiv >
< 4,, 4,4, full >
11000PA0
10 1 8 3 8
14B
< [, 4, 4, selective >
< 4, 4,4, full
32B
< 4, 4, selective
< 4, 4,4, full
HAN
48
72B
< 8, 3,8, selective
>
<. 8,8, 4, selective >
Table 3. Parallel Strategies for Training Different els with
Mod- Formatted in
Figure 7. Unsuitable ChankSize and K Leads to Performance Degradation
Different Context Length, TP,SP, Recompute Granularity >
<
Model Size
32K
256K
tune the parameters ChunkSize and K. The additional overhead introduced by this process accounts for less than 0.5% of the total training time in our practice, and the result- ing best combination is selected to maximize performance.
7B
(32K, 1)
(8K, 16)
14B
(8K,8)
(8K, 8)
32B
(8K, 6)
(8K, 6)
72B
(8K, 16)
(8K, 16)
6. Evaluation
Table 4. Parameter Setting (ChunkSize, K)
ChunkFlow,
Formatted in
6.1. Evaluation Setup
We evaluate ChunkFlow on Qwen2.5-serie LLMs and choose Megatron-LM the baseline. This choice is due to the fact that Megatron-LM represents the state-of-the -ftee---eeee-t-- art in LLM training, and ChunkFlow is built on top of it. Regarding the training dataset, we construct an eval- uation dataset that better aligns with the real-world distri- bution characteristics for long context fine-tuning based on the methods provided in (Fu et al., 2024; ChatGLM, 2025). The sequence length distribution in the evaluation dataset is shown in Table 2 It can be observed that this evaluation dataset shares similar distribution characteris- tics with LMSYsChatIM (Zheng et al., 2023), but has a slightly higher proportion of sequences longer than 32K and those shorter than IK. All experiments are conducted using Alibaba Cloud ml 8xlarge gungge-glliie 100 instances (Alibaba, 2025a), with global batch size of 256 and a micro-batch size of 1.
and compare the performance with Megatron-LM. The configurations used for training different-sized models in Mcgatron-LM with various context lengths are shown in Table 3. These configurations achieve the best performance in Megatron-LNN while ensuring no OOM errors occur.
ChunkFlow adopts the same parallel strategies as the base- line (i.e., the combination of < TP,, SP, PP >) but incor- porates the selective recomputation strategy across all exper- iments. This is because ChunkFlow's 's memory consumption is primarily determined by the ChunksSize rather than the length of the longest sequence, allowing it to avoid OOM is- sues without requiring full recomputation strategy. Addition- ally, we list the best-performance configurations for Chunk- Flow, obtained through a grid search over Chank:Size and K, in Table 4.
Figure 8 illustrates the performance results of ChunkFlow and Megatron-LMM, with performance measured by aver- age iteration time. To facilitate a clearer comparison, we normalize Megatron-LM'ss results relative to those of Chunk- Flow. The results indicate that ChunkFlow can accelerate long context fine-tuning by up to 4.53x. This significant improvement in ChunkFlow can be attributed to two key design innovations. First, ChunkFlow consolidates short se- quences into single chunks, which greatly enhances compu- tational efficiency. Second, it's state-aware IFIB scheduling mechanism reduces pipeline bubbles and thus further boosts overall performance. Additionally, unlike Megatron-LM, ChunkFlow avoids memory bottlenecks, eliminating the need for full recomputation when fine-tuning models with 256K context lengths, such as 7B, 14B, and 32B models.
Sequence Length
Proportion of Sequences
< 1K
08.17%
< 4K
99.72%
< 8K
99.83%%
< 32K
99. 99.92%
< 128K
99.98%
Longest
256K
Table Sequence Length Distribution in Evaluation Dataset
6.2. End-to-end Evaluation
We first evaluate the end-to-end training efficiency of us- ing ChunkFlow for fine-tuning Qwen2.5 models listed in Table on the evaluation dataset with context lengths of 32K and 256K, respectively. For each experiment, we ex- clude sequences exceeding the context length in the dataset
Notably, even when using the same recomputation strategy, experiments on fine-tuning muning the 72B model demonstrate that ChunkFlow still achieves substantial performance gains, underscoring its superior efficiency and scalability.
7

---

Efficient Long Context Fine-tuning with Chunk Flow
Normalieed Assrage hand Time Bealls
for dependent chunk group. The results in Table 6 indi- cate that both ChunkSize and K significantly influence overall performance. The (8K, 4) configuration achieves optimal performance. In contrast, the (2K, 16) configura- tion leads to suboptimal performance since it yields chunks that are too small to fully exploit the GPU's computational efficiency. Additionally, the (32K, 1) configuration results in fewer chunks, leading increased pipeline bubbles and consequently reducing overall performance.
24
230
Figure 8. Normalized End-to-End Training Performance Results
(ChunkSize K)
Training Performance(ms)
(2K, 16)
29810
6.3. Case Study
(SK, 4)
23774
(32K, 1)
28942
In this section, we conduct experiments to explore the mem- ory consumption characteristic of ChunkFlow as well as the performance impact of the ChunkSize and K.
Table 6. Impacts of ChunkSize and K on Training Efficiency
6.3.1. MEMORY CONSUMPTION CHARACTERISTIC IN CHUNKFLOW
7. Related Works
A significant portion of research in the field of long context fine-tuning primarily focuses on enhancing model perfor- mance by meticulously constructing datasets (Meta, 2024; Alibaba, 2025b; Bai et al., 2024; Zhao et al., 2024). How- ever, the training efficiency of this task has not been thor- oughly explored. As the importance of long context fine- tuning grows, there is increasing emphasis on improving end-to-end training performance by leveraging the unique characteristics of this workload, particularly the variabil- ity in dataset lengths. Sequence packing is proposed as a method to eliminate padding when processing batches with varying sample lengths, thereby reducing unnecessary mem- ory usage and computational demands (Kosec et al., 2021). Additionally, smart batching (Bai et al., 2024), which sorts mini-batches during training steps, is employed to achieve better load balancing among data parallel ranks. Further- more, several studies that do not directly address the issue of long context fine-tuning have contributed solutions to related challenges within this domain. HotSPa (Ge et al., 2024) introduces multiple sequence parallel strategies tai- lored for varying sequence lengths in the training datasets. There is also line of research, such LoRA(Hu al., 2022), that focuses parameter-efficient fine-tuning im- prove training efficiency. These approaches are orthogonal to ChunkFlow and can be seamlessly integrated with it.
Table 5 demonstrates the peak memory usage when fine- tuning the 7B model with different context lengths and vary- ing ChunkSize. All experiments share the same training configuration of < 4, 4, 1, selective >, and K is set to The results indicate that ChunkFlow's memory consumption is primarily determined by the ChunkSize regardless of the maximum sequence length in datasets. However, it can also be observed that training with 256K context length consumes slightly more memory compared with training with a 32K context length using the same ChunkSize. This because, in our current implementation, we directly save all key/value tensors in memory without further offloading optimizations. We leave this optimization for future work.
Context Length
Chunk
Peak Memory
32K
2K
41.6 GiB
256K
2K
45.6 GiB
32K
4K
47.5 GiB
256K
4K
50.8 GiB
32K
8K
59.3 GiB
256K
8K
63.8 GiB
Table Memory Consumption Characteristic in ChunkFlow
6.3.2. PERFORMANCE IMPACT OF Siznk AND K
8. Conclusion
We also investigate the performance impacts of ChunkSize and K by fine-tuning the 7B model with a context length of 256K and various combinations of these two parameters. All experiments are conducted using the < 4, 4, 4, seleclive >> training strategy, and training performance is measured in terms of average iteration time. We maintain ChunkSize * K constant across different settings to ensure that all ex- periments save the same total amount of chunk activations
In this work, we propose a chunk-centric training method called ChunkFlow to effectively address the challenges posed by the unique distribution characteristics of long- context fine-tuning datasets for LLMs. ChunkFlow reor- ganizes input sequences into uniformly sized chunks by consolidating short sequences and splitting long ones. also introduces a novel state-aware chunk scheduling algo-
8

---

Efficient Long Context Fine-tuning with Chunk Flow
rithm manage these chunks for forwards and backwards. This scheduling algorithm ensures nearly constant memory consumption, which is governed by the chunk size. Addi- tionally, we integrate ChunkFlow's 's approach with existing pipeline scheduling methods and implement a state-aware IFIB scheduling technique, further enhancing distributed training performance on variable-length sequences. Evalua- tion results demonstrate that ChunkFlow achieves up 4.53x speedup in performance compared to the state-of-the- art training system, Megatron-LM. Furthermore, we believe that ChunkFlow can serve as an effective solution for broader range of scenarios involving the training of LLMs on variable-length sequences.
Alibaba. Qwen2.5 technical report, 2025b. URL https:
Anthropic.
Introducing 100k context windows.
https: //wp:::/wwww.. anthropic cttthroppc.com/eesss....... com/news/
100x-context-windde. 2025. 2025-01-20.
Accessed:
Ao, S., Zhao. W., Han, X., Yang, C., Zhang, X., Liu, Z., Shi. C., and Sun, M. Seq1FIB: Efficient sequence-level pipeline parallelism for large language model training. In Chiruzzo, L.,, Ritter, A., and Wang, L. (eds.), Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguis- tics: Human Language Technologies (Volume I: Long Pa- pers), pp. 1998-90088 Albuquerque, New Mexico, April 2025. Association for Computational Linguistics. ISBN 979-8-89 18-899976---966 URL https: / aac.aantnnoooo....o o org /zgg2025..0 haw.ssnnacclloooo
Impact Statement
This paper presents work whose goal is to advance the sys- tem efficiency of long-context fine-tuning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
454/.
Bai, Y., Lv, X., Zhang, J., He, Y., Qi, J., Hou, L.,, Tang, J., Dong, Y., and Li, J. LongAlign: A recipe for long context alignment of large language models. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Findings of the Association for Computational Linguis- tics: EMNLP 2024, pp. 1376-1395, Miami, Florida, USA, November 2024. Association for Computational Linguisties. doi: 10.18866/w122224.miiinn-e-e 74. URL https: http://alaaaaahcccc wwwp:gg...rrr////e org/2024. f findings-emm1p.275
Acknowledgements
This work is supported by the following funding: Na- tional Science Foundation of China (92270206, 62372435), China National Postdoctoral Program for Innovative Tal- ents (BX20240383) and Beijing Natural Science Foundation (4254087).
References
ChatGLM. Chatglm with Im context length. https://
medium. Soeeuumm.com/Chhaatst//gg/// long- lon---calnnggggggg scalings pre-traiined-modee--cttttttt -to-millions-cssseeeee 2025-01-20.
Ainslie, J., Lee- Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebron, F., and Sanghai, S. GQA: Training generalized multi-query transformer models from multi-head check- points. In Bouamor, H., Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference Empirical Meth- ods in Natural Language Processing, pp. 4895-4901, Singapore, December 2023. Association for Computa- tional Linguistics. doi: 10.18655/11/002.emmllpp main. 298. URL ht https: htw.aaanttoooy.oorg//000 emn 1 p-ma in. 298/.
2025. Accessed:
Fan, S., Rong, Y., Meng, C., Cao, 7., Wang, S., Zheng, Z., Wu, C., Long, G., Yang, J., Xia, L., Diao, L., Liu, X., and Lin, W. DAPPLE: a pipelined data parallel approach for training large models. In Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, PPoPP '21, pp. 431-445, As- sociation for Computing Machinery. ISBN 978- 1-450- 8294-6. doi: 10. 10.1145/3437800
Alibaba. Alibaba cloud pai dlc gpu instance.
URL https: 111000..44193.. event-
https://bellp.aaiyynccco///abba
/ / //doioooo org. 10.11000111 . place: Virtual Event, Republic of Korea.
Fu, Y., Panda, R., Niu, X., Yue, X., Hajishirzi, H., Kim, Y., and Peng. H. Data engineering for scaling language models to 128k context, 2024. URL https: /arxivv. org hatp/ab///a 24002 10171.
helpe
mlgerhefbbxlarrge--A------1----0
Ge, H., Fu, F., Li, H., Wang, X., Lin, S., Wang, Y., Nie, X., Zhang, H., Miao, X., and Cui, B. Enabling paral- lelism hot switching for efficient training of large lan- guage models. In Proceedings of the ACM SIGOPS
211280011305532215557006888888
4-RE new4-PO_0-P1 2025a. Accessed:
2025-01-20.
9

---

Efficient Long Context Fine-tuning with Chunk Flow
30th Symposium on Operating Systems Principles, SOSP 224, pp. 178-194, New York, NY, USA, 2024. Associa- tion for Computing Machinery. ISBN 9798400712517. doi: 10. 0001144//6644155369996666 URL hatpps://do org/10. 1 45 /1..//6999711....59..
Li, Z., Zhuang. S., Guo, S., Zhuo D., Zhang, H., Song, D. X., and Stoica, I. TeraPipe: Token-level pipeline parallelism for training large-scale language models. In International Conference on Machine Learn- ing, C. URL https: /api. ippi.... semant cscholar. org http/CorpusD:::20 :ttrppuII::21194422 :21::222113222
Gemini Team, G. Gemini 1.5:
Unlocking multimodal un- derstanding across millions of tokens context, 2024. URL https::/p//.......t....tt / //a//aaaaaa /arxivv.... orgg abs 2403. 05530.
Liu, H., Zaharia, M., and Abbeel, P. RingAttention with blockwise transformers for near-infinite context. In The Twelfth International Conference on Learning Rep- resentations. URL https: / wpwp//veeweeeeeeeeeeeeeeeeeee.em net// forum?id S
Hu, E. J., Shen, Y., Wallis, P., Allen-llnnitee Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adaptation of large language models. In ICLR 2022, April 2022.
Ma, R., Yang, X., Wang. J., Qi, Q., Sun, H., Wang, J., Zhuang, Z., and Liao, J. HPipe: Large language model pipeline parallelism for long context on heterogeneous cost-effective devices. In Yang. Y., Davani, A., Sil, A., and Kumar, A. (eds.), Proceedings of the 2024 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies ( (Volume 6: Industry Track), pp. 1-9, Association for Computational Linguistics. doi: 10. 186533/1122024. naacl-industry.. URI. https:: //ac lnthholoy.o org/202 naacl innustry.....
Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, M. X., Chen, D., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., and Chen, Z. GPipe: efficient training of giant neural networks using pipeline parallelism. In Proceedings of the 33rd Inter- national Conference on Neural Information Processing Systems. Curran Associates Inc.
Jacobs, S. A., Tanaka, M., Zhang, C.,, Zhang, M., Am- inadabi, R. Y., Song, S. L., Rajbhandari, S., and He, Y. System optimizations for enabling training of ex- treme long sequence transformer models. In Proceed- ings of the 43rd ACM Symposium on Principles of Dis- tributed Computing, PODC '24, pp. 121-130. Associ- ation for Computing Machinery. ISBN 979-8-4007- URL https:
Meta. The Ilama 3 herd of models, 2024. URL https:
(laaxiv.org/aae abs 22477.2077773
0668-4. doi: 1411144661118886666666888 //doi. aoooioorg/10 www/oorg.0... /10. 1t..11..1111/1 place: Nantes, France.
Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Pat- wary, M., Korthikanti, V., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Zaharia, M. Efficient large-scale language model train- ing on gpu clusters using megatron-Im. In Proceed- ings of the International Conference for High Perfor- mance Computing, Networking, Storage and Analysis, SC '21, New York, NY, USA, 2021. Association for Computing Machinery. ISBN 9781450384421, doi: 10. 1145/345881173476200, URL https://doi.crs 10 N0.1145/3458881 34762299..77....
366288 event-
Korthikanti, V. A., Casper, J., Lym, S., McAfee, L., An- dersch, M., Shoeybi, M., and Catanzaro, B. Reducing activation recomputation in large transformer models. In MLSys '23.
Kosec, M., Fu, S., and Krell, M. M. Packing: Towards 2x NLP BERT acceleration. CoRR, 0atp22107.000,,, 2021. URL
http://axxxvv..g/aaa/.
02027.
Li, S., Xue, F., Baranwal, C., Li, Y., and You, Y. Sequence parallelism: Long sequence training from system per- spective. In Rogers, A., Boyd-Graber J., and Okazaki, N. (eds.), Proceedings of the ó1 Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers). pp. 2391-2404. Association for Compu- tational Linguistics, a. doi: 10.18663//1202233cccooo
OpenAL. Gpt-4 technical report, 2024. URL https://
arxiv. 2203. 08774.
Qi, P., Wan, X., Huang, G., and Lin, M. Zero bubble (al- most) pipeline parallelism. In The Tweifth International Conference on Learning Representations. URL https: / //openevieww net/ www.mpdd=ttutTttCCCC.
134. URL https: http::tttttaaaaatiiiete http::::///acaahooooo.y/ org..o20....g/22...... acl-long: 134/.
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. ZeRO: memory optimizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC '20. IEEE Press. ISBN 978-1-7281- 9998-6. Place: Atlanta, Georgia.
Li, S., Zhao, Y., Varma, R., Salpekar, O., Noordhuis, P., Li, T., Paszke, A., Smith, J., Vaughan, B., Dama- nia, P., and Chintala, S. PyTorch distributed: experi- ences on accelerating data parallel training. 13(12):3005- 3018, b. ISSN 2150-8097. doi: 10.14778/3415478.
3415530. URL hppp:://doi.or///100/ 14778/ 3415478. 3415530. Publisher: VLDB Endowment.
Shocybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LMM, Training multi-
10

---

Efficient Long Context Fine-tuning with Chunk Flow
billion parameter language models using model par- allelism. URL https: wtt.aaivvoorrr.............. 9999 199.
08053. -eprint: 1909.08053.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u. and Polosukhin, Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.). Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL
neuripss cc cc/papee_fille/sseee////1/////// #ffa245557aee31foo03ccccaaaaaaa. pdf.
Xiong, W., Liu, J., Molybog, I., Zhang, H., Bhargava, P., Hou, R., Martin, L., Rungta, R., Sankararaman, K. A., Oguz, B., Khabsa, M., Fang, H., Mehdad, Y., Narang, S., Malik, K., Fan, A., Bhosale, S., Edunov, S., Lewis, M.. Wang, S., and Ma, H. Effective long-context scal- ing of foundation models. In Duh, K., Gomez, H., and Bethard, S. (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies (Volume 1: Long Papers), pp. (643-4663, Mex- City, Mexico, June 2024. Association for Compu- tational Linguistics. doi: 10. 10.186553w12022..aaadd 260. URL https: /al/aalaann //aaaaaaaaoolooooo.oooo.. org//024, baacl-lons 11:11-1-0 260/.
Zhao, L., Wei. T., Zeng. L., Cheng, C., Yang, L.,, Cheng, P., Wang. L., Li, C., Wu, X., Zhu, B., Gan, Y., Hu, R., Yan, S., Fang. H., and Zhou, Y. Longskywork: A training recipe for efficiently extending context length in large language models, 2024. htR..ttps:/ /tattps:///riiii /wwaaaixvvvorooo abs/2406. 00605.
Zhao, Y., Gu, A., Varma, R., Luo, L., Huang, C.-C., Xu, M.. Wright, L., Shojanazeri, H., Ott, M., Shleifer, S., Desmaison, A., Balioglu, C., Damania, P., Nguyen, B., Chauhan, G., Hao, Y., Mathews, A., and Li, S. PyTorch FSDP: Experiences on scaling fully sharded data par- allel. (6(12):388838866 3860. ISSN 2150-8097. doi: 10. 14778/361 15t0:36111500.1111 URL https: //o.ioorrgg.oooooo 10. Endowment.
Publisher: VLDB
Zheng. L., Chiang. W.-L,, Sheng, Y. Li, T., Zhuang, S., Wu, Z., Zhuang, Y., Li, Z., Lin, Z.,, Xing, E. P., Gonzalez, J. E., Stoica, I., and Zhang, H. Lmsys-chat- lm: A large-scale real-world Ilm conversation dataset, 2023.


Reviews:
[{"review_id": "YS8XZe8JaZ", "reviewer": "Reviewer_HVAi", "summary": "This paper introduces ChunkFlow, a novel method for efficient long-context fine-tuning of large language models (LLMs). ChunkFlow addresses the challenges of variable sequence lengths in training datasets by reorganizing input sequences into uniformly sized chunks, merging short sequences and splitting long ones. It employs a state-aware chunk scheduling mechanism to manage computational dependencies and ensure consistent memory usage, primarily determined by the chunk size. The method also reduces pipeline bubbles and improving distributed training efficiency.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer **HVAi**'s  insightful feedback. Below we clarify the concerns regarding related work discussion and broader evaluations, with references to our methodology and results in the paper.\n\n***(Q1): The paper does not discuss other efficient training methods like LoRA, or other long-sequence papers, like RingAttention.***\n\nWe would like to clarify that ChunkFlow emerged in response to the practical need to accelerate the fine-tuning of LLMs on datasets with sequences of varying lengths. Within this context, we discuss several relavant directions in Section 2(Premilinaries) and Section 7(Related Works). We also give a brief introduction about long sequence training methods such as sequence parallelism and token-level pipeline parallelism in Section 2 and RingAttention is also referenced in this section. \n\n**It is worth to highlight that ChunkFlow is orthogonal to LoRA and Ring-attention**. We conduct a discussion here to provide further explanations. \n\n● LoRA focuses on Parameter-Efficient Fine-Tuning by substantially reducing the number of training parameters. Conversely, ChunkFlow aims at enhancing the training efficiency of LLMs on variable-length datasets. Moreover, integrating ChunkFlow into the LoRA training process can effectively boost the training efficiency of LoRA.\n\n● RingAttention is an effective approach for handling extremely long sequences. However, it does not address training issues such as low resource utilization and pipeline bubbles caused by the variability in sequence lengths. When performing long context finetuning at scale, ChunkFlow and the RingAttention can complement each other. The RingAttention mechanism offers a method for distributed attention computation. ChunkFlow, through its unified chunking and state-aware scheduling strategy, reduces pipeline bubbles and enchances computation efficiency (as depicted in Figure 6), thereby improving training performance. We will incorporate the above-mentioned suggestions into the final version.\n\n***(Q2): This paper did not compare with any other long context models, on any long context benchmark, like LongBench.***\n\nWe sincerely appreciate the valuable suggestions put forward by reviewer HVAi and we conduct more experiments using the ***Llama3-8B*** model on LongBench dataset. The results and analysis will also be reflected in our final version.\n\nIt is worth to clarify that the LongBench dataset is primarily employed to evaluate the long context understanding capabilities of LLMs, rather than for long context fine-tuning.  The table below presents the distribution characteristics of sequence lengths in LongBench dataset and demonstrate significantly difference distribution with our previously mentioned SFT dataset, as well as Llama SFT dataset[1]. ***Those distinctions helps explain why LongBench wasn't initially included in our experiment datasets.*** \n\n| Sequence Length | Propotion Of Sequences in LongBench| Propotion Of Sequences in Meta |\n|:-------|:--------:|-------:|\n| < 1K  | 0.26%    |  ~ 99% |\n| < 4K  | 22.65%   | - |\n| < 8K  | 48.4%| -  |\n| < 16K  | 82.59%   | -  |\n| < 32K  | 98.06%   | -  |\n| Longest| 64K   | 128K  |\n\nHowever, we agree that benchmarking on LongBench provides important validation of ChunkFlow's effectiveness, and we conduct the experiments using llama3-8B on longbench to further highlight our contributions. For ChunkFlow, we configure the `ChunkSize=16384, K=2`. Both experiments adopt the identical `<TP=2, SP=2, PP=4>` parallelization strategy. Under these circumstances, ChunkFlow demonstrates a **1.7x** speedup compared to Megatron-LM. Evidently, this showcases the superiority of our design.\n\n***(Q3): How does ChunkFlow perform on datasets with even longer sequences (e.g., >1M tokens)? Could the method be extended to non-causal models?*** \n\nAs shown in Figure 8, ChunkFlow shows greater performance improvement when fine-tuning models with longer context lengths. As the fine-tuning context length increases, the sequence length distribution widens. Training strategies for the longest sequences perform more poorly for most short ones. This implies that for models with even longer contexts (e.g., >1M tokens), ChunkFlow can achieve far better performance than baselines.\n\nFor applying ChunkFlow to non-causal model, it has two chunk-construction methods: short-sequence consolidation and long-sequence splitting. The former can be directly used in non-causal model training. The latter depends on causal-attention. So, when using ChunkFlow for non-causal models, setting ChunkSize to the dataset's maximum sequence length can effectively improve training efficiency.\n\nDue to limitations in computational resources and the timeline for this submission, we regret that we were unable to include these specific (>1M, speedups on LoRA and non-causal) experiments. We will incorporate the above-mentioned experiments in the final version.\n\n[1] Meta(2024). The Llama 3 Herd of Models. ArXiv, abs/2407.21783.", "claims_and_evidence": "The paper claims that ChunkFlow improves training efficiency for long-context fine-tuning by reorganizing sequences into uniform chunks and using state-aware scheduling. Evidence in the experiments in Section 6 includes the comparison to Megatron-LM, with Peak Memory and training performance.", "methods_and_evaluation_criteria": "ChunkFlow uses chunk construction and state-aware scheduling to manage variable-length sequences. Evaluation is based on training performance metrics like iteration time and memory usage.", "theoretical_claims": "The paper argues that chunk-based training and state-aware scheduling can optimize GPU utilization and reduce pipeline bubbles in distributed training. This is the claim, but no so theoretical.", "experimental_designs_or_analyses": "Experiments involve fine-tuning Qwen2.5-series LLMs with varying context lengths, comparing ChunkFlow to Megatron-LM on metrics like memory consumption and training speed.", "relation_to_broader_scientific_literature": "ChunkFlow builds on existing work in sequence packing and pipeline parallelism.", "essential_references_not_discussed": "The paper does not discuss other efficient training methods like LoRA, or other  long-sequence papers, like  RingAttention.", "other_strengths_and_weaknesses": "This paper did not compare with any other long-context models, on any long-context benchmark, like LongBench.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "How does ChunkFlow perform on datasets with even longer sequences (e.g., >1M tokens)? \nCould the method be extended to non-causal models?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IU6U3gP9Sv", "reviewer": "Reviewer_xBtd", "summary": "This paper propose Chunk Flow, a novel chunking and scheduling method for pipeline parallel long sequence training. It first discussed the long-tail phenomena of long-context LLMs training, and the potential issues result from it, such as underutilize of GPU memory and pipeline bubbles. Then the authors propose to chunk the longer sequence to a max chunk size, with a scheduler handling the defendency of the causal mask of long sequence. The author compares chunk flow against the state of the art pre-training framework megatron and find chunk flow well maintained the peak memory across various length long sequence during the training process, while achieve a 4x acceleration compared to megatron.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank **xBtd** for their very positive review of our work, noting that our paper is \"a solid study and have provide well explaination for the majority of the paper.\" and emphasizing strongly the importance of the research problem. In accordance with reviewer xBtd’s suggestions, we will redo all figures to enhance their readability in the final release. Additionally, we provide detailed responses to reviewer xBtd’s other concerns. \n\n***Q1: Overhead of keeping dependent chunks when N > K***\n\nWe sincerely appreciate the reviewer’s attention to the issue of computational overhead when the number of dependent chunks (N) exceeds K. In ChunkFlow, when N>K, the forward passes of the first `N-K` chunks are executed twice. After the first-time forward passes of these `N-K` chunks, **we discard their activation values**. Instead, **we only retain the key/value tensors of their attention mechanisms as states for subsequent reuse**.\n \nSince Multi-Query Attention (MQA) and Group-Query Attention (GQA) are widely adopted in mainstream large language models (LLMs) like LLama, Qwen, and Gemini, **storing these attention key/value tensors does not impose a substantial burden**. As can be seen from Table 5 in the paper, compared to processing 32K-length sequences, processing 256K-length sequences only consumes approximately 4GB more memory(Still much less memory footprint compared to Megatron-LM baseline). This clearly demonstrates that retaining these states in memory does not pose a severe problem.\n\nWhen training models with longer contexts, such as 2M token context models, retaining these states could ultimately lead to excessive memory consumption. As noted in Section 6.3.1, we plan to optimize this memory consumption through carefully designed memory offloading strategies.\n\n***Q2: Components retained for dependent sequences in causal attention***\n\nFor dependent chunks (split from long sequences), ChunkFlow retains key/value (K/V) tensors and their gradients in causal attention layers. These components are critical for:\n\n1. Forward Dependency : Subsequent chunks depend on K/V tensors from prior chunks to maintain causal masking (Section 4.2).\n\n2. Backward Dependency : Gradients for K/V tensors in earlier chunks require activations from later chunks(Figure 5).\n\n***Q3: Efficacy in datasets dominated by long sequences***\n\nEven when long sequences are more prevalent in datasets, ChunkFlow continues to demonstrate substantial advantages. However, the extent of performance improvement hinges on the distribution of data. For scenario analysis, we classify sequences as long if they exceed ChunkSize; otherwise, they are categorized as short. Whenever short sequences are present, we can consolidate them into a single chunk, significantly enhancing the utilization of GPU resources. As demonstrated in Figure 6, for long sequences, ChunkFlow splits them into uniformly sized chunks, effectively reducing pipeline bubbles. Subsequently, we carried out an experiment on the LongBench dataset [1]. The dataset exhibits a sequence-length distribution pattern consistent with what the reviewer envisioned.\n LongBench dataset is used for bilingual, multitask, and comprehensive assessment of long context understanding capabilities of LLMs and the table below demonstrates sequence length distribution in the dataset. When setting ChunkSize=8K, we can see that more than 50% sequences are long sequences and ChunkFlow achieve **1.7X** speeup over Megatron-LM, showing its effectiveness in a broader scenario.  If we further filter the LongBench dataset by removing sequences with fewer than 8K tokens, we can obtain a new dataset that **solely consists of long sequences**. Thanks to the Chunking mechanism of ChunkFlow, which can reduce pipeline bubbles, ChunkFlow still achieves a **1.4X** speedup compared to Megatron-LM.\n\n| Sequence Length| Propotion Of Sequences |\n|---------|---------|\n| < 1K | 0.26% |\n| < 4K | 22.65% |\n| < 8K | 48.4% |\n| < 16K | 82.59% |\n| < 32K | 98.06% |\n| Longest | 64K |\n\n[1] Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., & Li, J. (2023). LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. ArXiv, abs/2308.14508.", "claims_and_evidence": "-", "methods_and_evaluation_criteria": "The strength of this paper:\n1. Addressing a critical real world problem. Training long sequence model is very expensive and many existing frameworks have very limited optimization on long sequence pre-training. Chunk Flow ensures a upper bound of the peak memory and reduce pipeline bubble which can help the scaling of sequence length and size of LLMs.\n\n2. Good empirical results. The proposed method achieves good peak memory optimization, enables lower 60GB memory with 8k chunk size, which is quite impressive. The end2end speedup is 4x faster against megatron, which can significantly reduce the cost to train long sequence LLMs.\n\nThe limitation and questions for the authors:\n1. When the N > K, ie, the longest sequence exceed the num of chunks, it did duplicate computations for the overlapping chunks, and it keeps dependent chunk on GPU, which can leads to significant overhead.\n\n2. It is unclear that what components are actually retained for dependent sequence, especially within causal attention.\n\n3. What is the training data has more longer sequence than shorter sequence, how much is the benefit to use chunk flow in this scenarios.", "theoretical_claims": "-", "experimental_designs_or_analyses": "-", "relation_to_broader_scientific_literature": "-", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "-", "other_comments_or_suggestions": "This is overall a solid study and have provide well explaination for the majority of the paper. Good experimental results. But the author should redo the figures as they are not quite readable.", "questions_for_authors": "-", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VxhzSxm374", "reviewer": "Reviewer_dcDa", "summary": "This paper introduces ChunkFlow, an LLM training (fine-tuning) method that aims to improve the computational as well as memory efficiency of long-context training / fine-tuning. The authors start from three empirical observations in long-context fine-tuning, point out existing efficiency bottlenecks, and design ChunkFlow based on them. Given a batch of input sequence, ChunkFlow reorganizes it to a new set of chunks based on heuristics, so that the size of each chunk does not exceed a pre-defined chink size. The method also incorporates state-aware chunk scheduling and state-aware 1F1B for better training efficiency. Evaluations show that as compared to Megatron-LM, ChunkFlow accelerates long-context fine-tuning by up to 4.53x.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are sincerely grateful to Reviewer **dcDa** for devoting time to review our work and providing invaluable feedback. Regarding the clarity of figures and the impact statement, we will, in strict compliance with Reviewer dcDa’s suggestions, **incorporate an impact statement and refine all figures—with particular focus on Figures 2, 3, 5, 6, and 7—in the final version of our release**. Moreover, we offer a comprehensive response below to other concerns raised by Reviewer dcDa.\n\n***(Q1) The cost of grid search compares to training cost itself***\n\nWe sincerely appreciate reviewer dcDa's meticulous review of our paper and point out a key implementation detail regrading to grid search and allows us to furtherly clarify the whole training process of ChunkFlow. When compared to the remarkable benefits of identifying the optimal `(ChunkSize, K)`, the overhead of grid search is almost negligible.\nBelow, we use the grid search process presented in Table 6 to illustrate this. In our research on the 72B model (relevant data can be found in Table 6), we evaluated several configurations, namely `[(32K, 1), (16K, 2), (8K, 4)]` . To accurately estimate training performance, each `(ChunkSize, K)` combination had to undergo 10 training steps. The entire grid search process accumulated 15 min. Significantly, training with a full 256K-context using ChunkFlow takes approximately multiple days or even weeks(~70 hours in our case). **Evidently, the overhead of grid search accounted for less than 0.3% of the total training time**.\n\nFrom a cost-benefit perspective, the optimal configuration `(8K, 4)` achieved a **1.21x** speed-up compared to the worst-performing configuration(i.e.`(2K, 16)` ). As a result, approximately **12 hours of training time are saved** over the entire training process in our case. \nMoreover, these validated optimal parameters can be reused across repeated training runs, such as training new 72B model versions with new algorithms or training on updated datasets. This effectively amortizes the one-time cost of grid search.\n\n***(Q2): How and why many of the design choices/intuitions in ChunkFlow could help improve the computational/memory efficiency of pre-training***\n\nAlthough ChunkFlow is mainly designed to boost training efficiency for LLM's post-training phase like SFT and RLHF which involve training on variable-length sequence datasets. It's design can also be applied to other training stages like pretraining and long context continual pretraining.\n\nAs mentioned in Meta's tech report[1], there are three main stages in developing a new LLM model: pretraining, long context continual pretraining (for context length extension), and post-training (for human preference alignment). In Section 1(Introduction), we claim that ChunkFlow can enhance long context continual pretraining efficiency because this phase also involves training on datasets with variable-length sequences[1] and exhibits efficiency challenges similar to those in the SFT stage (e.g., pipeline bubbles and resource underutilization). **Consequently, ChunkFlow’s design naturally improves computational and memory efficiency in long context continual pretraining scenario**.\n\nRegarding pretraining, even though it typically operates on uniformly sized sequences, splitting sequences into chunks and utilizing ChunkFlow's state-ware scheduler to process them sequentially can still reduce pipeline bubbles(due to increased number of micro-batches)and alleviate memory pressure. **This suggests that ChunkFlow’s chunk-centric approach retains value even in scenarios with uniform sequence lengths**.\n\n[1] Meta(2024). The Llama 3 Herd of Models. ArXiv, abs/2407.21783.", "claims_and_evidence": "The claims made by the authors in this paper are supported by citations or experiment results.", "methods_and_evaluation_criteria": "No particular flaw in evaluation design. See issues/concerns in \"Questions For Authors\".", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "No particular flaw in experiment design. See issues/concerns in \"Questions For Authors\".", "relation_to_broader_scientific_literature": "This paper lies within the broader area of LLM fine-tuning systems. In particular, it addresses the computational and memory efficiency of long-context fine-tuning (a special case of LLM fine-tuning).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Issues not mentioned in prior sections:\n- Missing an impact statement.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Thank you for submitting this paper to NeurIPS. This is a well-motivated paper, drawing inspirations from empirical observations and presenting a concrete solution that is well-evaluated. \n\nBelow are a few questions and concerns:\n- The figures do a good job in illustrating different parts of the system (esp. reorganizing the chunks) but really need to be fixed. Fonts are too small to be read clearly (esp. figures 2, 3, 5, 6, 7). \n- You mention that you use grid search to find best K and chunk size --- I am curious what the cost is for the grid search process, before fine-tuning begins. And, in particular, how that cost compares to the training cost itself, e.g. the cost in terms of latency in table 6.\n- How does ChunkFlow generalize to pre-training? I understand that pre-training experiments might be too hard to implement, but I am curious how and why many of the design choices/intuitions in ChunkFlow could help improve the computational/memory efficiency of pre-training.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces ChunkFlow, a method designed to enhance the efficiency of fine-tuning LLMs on long context data by reorganizing input sequences into uniformly sized chunks and employing a state-aware chunk scheduling mechanism. This approach aims to address common issues in distributed training, such as load imbalance and pipeline bubbles, which can lead to suboptimal training performance and underutilized GPU resources. By ensuring optimal computational efficiency and balance, ChunkFlow claims to accelerate training performance by up to 4.53 times compared to the Megatron-LM baseline, with applications beyond fine-tuning, such as continual pre-training and pre-training stages.
It addresses a critical real-world problem in the training of LLMs, particularly in terms of memory and computational efficiency. 
In the rebuttals the authors addressed the reviewers' concern well.
Decision: Accept (poster)