LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models

We are truly appreciated for your valuable comments. In the following, we provide responses to the concerns.

Q1: “Other generative tasks that require longer context.”

A: Thanks for your suggestion! We have included the evaluation on LongBench[1] and L-Eval [2], as shown in Table 9 and Table 10 in the revision. We fine-tuned Llama2 7B model with our long QA data. In these benchmarks, there are comprehensive generative tasks, including document QA, summarization, few-shot learning, code completion synthetic tasks, and other open-ended tasks in L-Eval. Our model presents comparable or even better performance than other counterparts, with about 4 hours for fine-tuning on 8 A100 GPUs. It takes 60 million tokens per epoch, 5 epochs, and 0.3 billion tokens in total for the supervised fine-tuning.

Table 1 - Evaluation on LongBench English tasks.

Table 2 - Evaluation on L-Eval open-ended tasks, i.e., comparing models to GPT-3.5-Turbo and judging win rates via GPT-4.

Q2: “The improvement on perplexity in Table 4.”

A: Sorry for the confusion caused. The original Table 4 (Table 3 in the revision) is not designed to show perplexity improvements, which we never claim. The goal of LongLoRA is to significantly improve the training efficiency (in terms of both memory consumption and training speed) without compromising the perplexity achieved through finetuning the full model. We have clarified this point in the caption of Table 3 In the revision.

[1] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, Juanzi Li: LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. CoRR abs/2308.14508 (2023)

[2] Chenxin An, Shansan Gong, Ming Zhong, Mukai Li, Jun Zhang, Lingpeng Kong, Xipeng Qiu: L-Eval: Instituting Standardized Evaluation for Long Context Language Models. CoRR abs/2307.11088 (2023)

We are truly appreciated for your valuable comments. In the following, we provide responses to the concerns.

Q1: “Comparisons to methods that train long-context from scratch.”

A: Thanks for this question. We agree with you that training long-context from scratch could be an interesting direction. Nevertheless, training long-context LLMs from scratch is too computationally expensive and unaffordable for most researchers. LongLoRA provides an efficiency advantage of orders of magnitude compared to training long-context LLMs from scratch. For instance, Llama-2 models require training with 2 trillion tokens across hundreds of GPUs, whereas LongLoRA models are finetuned on about 2 billion tokens for 32k context length using just 8 A100 GPUs.

Recently, Meta published a concurrent research paper, LLama2-Long [1]. In this paper, the authors empirically verified that long-context continual training on short-context models is more efficient and similarly effective compared to pretraining from scratch with long sequences. This conclusion could further magnify the importance of LongLoRA.

Q2: “Regular LoRA and LongLoRA use the same amount of memory in Table 11. $S^2$-Attn for memory, or only flops?”

A: The reason why $S^2$-Attn does not reduce the memory usage on top of LoRA (LoRA+) in the original Table 11 is that we adopt FlashAttention-2 [2] for the implementation of self-attention layers. This library fuses all operators in multi-head self-attention (MHSA) into a single CUDA kernel and thereby avoids writing the attention score matrix to DRAM. So the memory space required for MHSA is around $2 * B * N * C$, where B is the batch size, N is the sequence length and C is the number of channels. The multiplier 2 corresponds to input and output activations.

Nevertheless, if FlashAttention-2 is not used, the vanilla dense attention requires $2 * B * N * C + B * N^2 * C$ memory space. Here, $B * N^2 * C$ corresponds to attention scores. In comparison, our $S^2$-Attn requires only $2 * B * N * C + 4 * B * (N/4)^2 * C = 2 * B * N * C + 1/4 * B * N^2 * C$ memory space. Empirically, we include an additional comparison that disables the FlashAttention-2 in the table below. Under a context length of 8192, $S^2$-Attn improves the training speed by 2.1x and memory usage by 1.8x. With a 16k context length, the original dense attention runs out of memory during training while $S^2$-Attn is still feasible. We have included this comparison in the Table 13 of the revision.

[1] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, Hao Ma: Effective Long-Context Scaling of Foundation Models. CoRR abs/2309.16039 (2023)

[2] Tri Dao: FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. CoRR abs/2307.08691 (2023)

We are truly appreciated for your valuable comments. In the following, we provide responses to the concerns.

Q1: About Figure 3

A: Thanks for your constructive suggestion. We have updated Figure 3 in the revision, with indices and annotations to the matrices, and the visual connections between matrices and attention patterns.

Q2: “Be more consistent about $S^2$-Attn, shift short attention, LoRA+, and LongLoRA.”

A: We have updated these terms to be more consistent in the revision. We make more explicit that LongLoRA is the combination of $S^2$-Attn and LoRA+ in the abstract.

Q3: “Table 7 is a great candidate for a line plot.”

A: We have transformed the original Table 7 into a line plot. Please refer to Figure 5 in the revision.

Q4: “Reference a specific section in the Appendix”

A: In the revision, we have added specific section numbers in the appendix when we reference them in the paper.

Q5: “The attention patterns ablation feels repetitive to Section 3.2”

A: Thanks for the reminder! To avoid repetition, we have relocated the original Table 2 to Table 6 (in Section 4.2) and shortened relevant descriptions in Section 3.2. In Section 3.2, the objective is to show that we can “train sparse, test dense” with $S^2$-Attn. In Section 4.2, we further demonstrate that $S^2$-Attn achieves superior performance to existing sparse attention patterns.

Q6: “Retrieval-based evaluation - able to handle longer context.”

A: Sorry for the confusion caused. We have removed the word “somehow” in the section you mentioned. Retrieval-based evaluation shows that LongLoRA extends the context window of a pre-trained LLM. Specifically, Llama-2-7b sharply fails to retrieve the passkey when the context window exceeds 4k, while our method maintains a high retrieval accuracy (60-90%) even at a much longer context length of 33k-45k.

Q7: “The original standard self-attention at inference time.”

A: In Table 6 of the revision (the original Table 2), for each attention pattern, we evaluate its performance under two protocols. In the first row, we use sparse attention in both training and testing. In the second row, we use sparse attention only for training and the standard full attention for testing. For our $S^2$-Attn, it achieves the best perplexity with the original full self-attention at inference time. We have elaborated on these details in the table caption in the revision.

Q8: “An additional baseline that trains LoRA + embeddings.”

A: Thanks for your reminder. We have conducted this ablation and included it in Table 2 in the revision. Finetuning the embedding layer brings larger benefits than normalization layers.

We are truly appreciated for your valuable comments. In the following, we provide responses to the concerns.

Q1: “Evaluation only on perplexity and retrieval setting.”

A: We include additional comparisons on LongBench [1] and L-Eval [2] benchmarks. We fine-tune Llama2 7B with our method on our long QA data. We compare our model with GPT-3.5-Turbo and other Llama2-based long-context models, Vicuna and LongChat models, in the tables below. It shows that our 7B model presents comparable or even better performance than these Llama2-based long-context models, while ours only takes about 4 hours, about 0.3 billion tokens, for supervised fine-tuning on a single 8x A100 machine. We have included these evaluations in Table 9 and Table 10 in the revision.

Table 1 - Evaluation on LongBench English tasks

Table 2 - Evaluation on L-Eval open-ended tasks, i.e., comparing models to GPT-3.5-Turbo and judging win rates via GPT-4

Q2: “Ablation on group size for longer sequence lengths”

A: Thanks for your suggestion! We perform additional ablation experiments on group size using a context length of 16384, as depicted in the table below. The results indicate that setting the group size to 1/4 of the context length is optimal in terms of efficiency-accuracy tradeoff. We have included these results in Table 7 in the revision.

Q3: “The method to estimate model FLOPs.”

A: We profile the context stage FLOPs of Llama2-7B using a batch size of 1 and various context lengths using a third-party tool, torchprofile [3]. The tool traces the computation graph and sums up the FLOPs of each node in the graph (e.g. Q/K/V/O projections, multi-head self-attention, fully-connected layers, and normalization layers).

[1] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, Juanzi Li: LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. CoRR abs/2308.14508 (2023)

[2] Chenxin An, Shansan Gong, Ming Zhong, Mukai Li, Jun Zhang, Lingpeng Kong, Xipeng Qiu: L-Eval: Instituting Standardized Evaluation for Long Context Language Models. CoRR abs/2307.11088 (2023)

[3] torchprofile. https://github.com/zhijian-liu/torchprofile