Efficient Long Context Fine-tuning with Chunk Flow

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.

(Q1) The cost of grid search compares to training cost itself

We 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. Below, 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.

From 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. Moreover, 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.

(Q2): How and why many of the design choices/intuitions in ChunkFlow could help improve the computational/memory efficiency of pre-training

Although 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.

As 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.

Regarding 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.

[1] Meta(2024). The Llama 3 Herd of Models. ArXiv, abs/2407.21783.

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.

Q1: Overhead of keeping dependent chunks when N > K

We 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.

Since 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.

When 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.

Q2: Components retained for dependent sequences in causal attention

For 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:

1. Forward Dependency : Subsequent chunks depend on K/V tensors from prior chunks to maintain causal masking (Section 4.2).

2. Backward Dependency : Gradients for K/V tensors in earlier chunks require activations from later chunks(Figure 5).

Q3: Efficacy in datasets dominated by long sequences

Even 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. 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.

[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.

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.

(Q1): The paper does not discuss other efficient training methods like LoRA, or other long-sequence papers, like RingAttention.

We 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.

It is worth to highlight that ChunkFlow is orthogonal to LoRA and Ring-attention. We conduct a discussion here to provide further explanations.

● 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.

● 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.

(Q2): This paper did not compare with any other long context models, on any long context benchmark, like LongBench.

We 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.

It 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.

However, 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.

(Q3): How does ChunkFlow perform on datasets with even longer sequences (e.g., >1M tokens)? Could the method be extended to non-causal models?

As 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.

For 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.

Due 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.

[1] Meta(2024). The Llama 3 Herd of Models. ArXiv, abs/2407.21783.