PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, Jiaya Jia

arXiv:2309.12307v3cs.CLcs.AIcs.LG

TL;DR

Extending pretrained LLMs to long contexts is computationally expensive. LongLoRA combines shifted sparse attention with LoRA that trains embeddings and normalization, achieving performance comparable to full-attention, fully fine-tuned results at lower computational cost while extending Llama2 7B to 100k or 70B to 32k context on one 8× A100 machine.

  • Problem

    Extending LLMs to long contexts is computationally expensive because self-attention costs grow quadratically with sequence length.

    Source: §1 · p. 2, §3.1 · p. 4

  • Method

    LongLoRA combines shifted sparse attention during training with LoRA that trains embedding and normalization layers for long-context adaptation.

    Source: §1 · p. 2, §3.3 · p. 6

  • Results

    LongLoRA achieves performance comparable to full-attention, fully fine-tuned results at much lower computational cost, extending Llama2 7B to 100k or 70B to 32k context on one 8× A100 machine.

    Source: §1 · p. 3

  • Takeaways & Limitations

    LongLoRA models retain standard attention at inference, allowing reuse of existing infrastructure and optimizations.

    Source: §5 · p. 9

    LoRA that adapts only attention weights does not work for long-context extension.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

We present LongLoRA, an efficient fine-tuning approach that extends the context sizes of pre-trained large language models (LLMs), with limited computation cost. Typically, training LLMs with long context sizes is computationally expensive, requiring extensive training hours and GPU resources. For example, training on the context length of 8192 needs 16x computational costs in self-attention layers as that of 2048. In this paper, we speed up the context extension of LLMs in two aspects. On the one hand, although dense global attention is needed during inference, fine-tuning the model can be effectively and efficiently done by sparse local attention. The proposed shifted sparse attention effectively enables context extension, leading to non-trivial computation saving with similar performance to fine-tuning with vanilla attention. Particularly, it can be implemented with only two lines of code in training, while being optional in inference. On the other hand, we revisit the parameter-efficient fine-tuning regime for context expansion. Notably, we find that LoRA for context extension works well under the premise of trainable embedding and normalization. LongLoRA combines this improved LoRA with S^2-Attn. LongLoRA demonstrates strong empirical results on various tasks on Llama2 models from 7B/13B to 70B. LongLoRA extends Llama2 7B from 4k context to 100k, or Llama2 70B to 32k on a single 8x A100 machine. LongLoRA extends models' context while retaining their original architectures, and is compatible with most existing techniques, like Flash-Attention2. In addition, we further conduct supervised fine-tuning with LongLoRA and our long instruction-following LongAlpaca dataset.

1 INTRODUCTION

LongLoRA extends pre-trained LLM context by combining shifted sparse attention during fine-tuning with trainable embeddings and normalization; models retain standard attention at inference.

1 INTRODUCTION

  • Pre-defined context windows limit applications such as long-document summarization, while training from scratch or fine-tuning for longer contexts is computationally expensive.The paper notes that standard self-attention costs rise sharply as context expands, even when using LoRA.Source: §1 · p. 2
  • Plain LoRA yields high long-context perplexity even at rank 256; LongLoRA trains embeddings and normalization, which comprise <2% and ≤0.004% of Llama2 7B parameters, respectively.The paper identifies these trainable layers as key to effective context-extension LoRA fine-tuning.Source: §1 · pp. 2–3, Figure 2 · p. 2
  • S2-Attn splits tokens into groups and shifts half the heads by half a group, enabling information flow between groups; group size 2048 approximates 8192-token training.Attention is computed within each group during fine-tuning.Source: §1 · p. 2, Figure 3 · p. 3
  • LongLoRA retains standard attention at inference, allowing existing LLM optimizations and infrastructure; Flash-Attention2 is compatible in both training and inference.The authors attribute compatibility to short attention resembling the attention scheme used during pre-training.Source: §1 · p. 2, Figure 2 · p. 2
  • LongLoRA extends Llama2 7B to 100k context and 70B to 32k on one 8× A100 machine, with performance comparable to full-attention, fully fine-tuned results at much lower computational cost.The reported experiments also cover Llama2 13B.Source: §1 · p. 3
  • The authors also introduce LongAlpaca, a self-collected long instruction-following dataset, and use it to further fine-tune LongLoRA models on long questions and answers.Its questions cover technical papers, science fiction, and other books; the paper presents SFT as important for chat ability.Source: §1 · p. 3

2 RELATED WORK

Prior long-context work uses retrieval, approximate attention, memory compression, or position-encoding changes; LongLoRA instead targets efficient fine-tuning while retaining unmodified full attention at inference. It builds on parameter-efficient fine-tuning while noting that trainable embeddings alone are insufficient for context extension.

Long-context Transformers

Long-context LLMs

  • Prior LLM context-extension approaches include full-finetuning methods Position Interpolation (Chen et al., 2023) and Focused Transformer (Tworkowski et al., 2023), while Landmark attention (Mohtashami & Jaggi, 2023) is efficient but lossy.Position Interpolation and Focused Transformer require 128 A100 GPUs or 128 TPUv3 for training; Landmark attention compresses long-context inputs into retrieved tokens, whereas LongLoRA saves fine-tuning costs while preserving attention quality.Source: §2 · pp. 3–4
  • Position-embedding approaches include Position Interpolation (Chen et al., 2023), NTK-aware (ntk, 2023), Yarn (Peng et al., 2023), positional Skipping, and out-of-distribution methods.LongLoRA focuses on efficient fine-tuning while retaining the original inference architecture, making it orthogonal to these position-embedding methods.Source: §2 · p. 4

Efficient Fine-tuning

3 LONGLORA

The section tackles the cost of long-sequence attention and the limits of standard LoRA for context extension. It presents shifted sparse attention and trainable embedding and normalization layers as ways to make long-context fine-tuning more efficient and effective.

3.1 BACKGROUND

  • Self-attention has quadratic computation cost in sequence length, and standard LoRA’s attention-only adaptation does not work for long-context extension.Standard LoRA freezes other layers, including MLP and normalization layers.Source: §3.1 · p. 4

3.2 SHIFTED SPARSE ATTENTION

  • Shifting the group partition by half a group in half the attention heads enables information flow between groups without extra computation, bringing results close to the standard-attention baseline.Short attention alone lacks communication between groups, and its perplexity worsens as context length increases.Source: §3.2 · p. 5
  • S2-Attn supports full-attention testing, while other efficient attention designs discussed require testing with the same attention used during fine-tuning; shifting helps avoid overfitting to specific patterns.The passage attributes gaps from standard full attention to designs developed for training from scratch.Source: §3.2 · p. 5
  • S2-Attn requires only two steps—shifting tokens in half the attention heads and transposing features from token to batch dimension—and can be implemented in two lines of code.The paper provides PyTorch-style pseudocode for the implementation.Source: §3.2 · p. 5, §3.1 · p. 5

3.3 IMPROVED LORA FOR LONG CONTEXT

  • On proof-pile, the method achieves performance comparable to full-attention or full-fine-tuning baselines with better efficiency for Llama2 7B and 13B at 8k–32k context lengths.The models are fine-tuned on RedPajama; the cited passage provides no specific perplexity values.Source: Table 3 · p. 6
  • Opening embedding and normalization layers for training yields LoRA+, addressing the growing gap between standard LoRA and full fine-tuning that larger LoRA ranks do not close.Normalization parameters are only 0.004% of Llama2 7B parameters.Source: §3.3 · pp. 5–6, Table 2 · p. 6

4 EXPERIMENT

The experiments evaluate LongLoRA on long-context language modeling, retrieval, and attention-pattern ablations. Results show improved perplexity with longer contexts, context extension up to 100k, and competitive topic and passkey retrieval.

4.1 EXPERIMENTAL SETTINGS

  • The experiments extend pre-trained Llama2 7B, 13B, and 70B models with Position Interpolation and train on Redpajama, evaluating long-sequence modeling on PG19 and proof-pile.Perplexity evaluation uses a sliding window with S = 256.Source: §4.1 · pp. 6–7

4.2 MAIN RESULTS

  • On a single 8× A100 machine, LongLoRA fine-tuned Llama2 7B to 100k, 13B to 65536, and 70B to 32768 context lengths.The extended models showed some perplexity degradation at small context sizes, a known limitation of Position Interpolation.Source: §4.1 · p. 6, Table 4 · p. 7, §4.2 · p. 7
  • For Llama2 7B, perplexity fell from 2.72 to 2.50 (-0.22) as context increased from 8192 to 32768; for 13B, it decreased by -0.28.For the same training and evaluation context lengths, perplexity decreased as context size increased.Source: §4.2 · p. 7
  • On LongChat topic retrieval, the Llama2 13B model achieved comparable performance to LongChat-13B and slightly outperformed it at 16k evaluation.It was fine-tuned on Redpajama with an 18k context length, at a cost similar to training at 16k; the task tests conversations from 3k to 16k.Source: §4.2 · pp. 7–8, Table 5 · p. 7
  • 33k–34k: the 32k-fine-tuned LongLoRA 7B model retained reasonable passkey retrieval accuracy, while extended position interpolation reached 60%-90% accuracy from 33k to 45k.Without further fine-tuning, extending position interpolation to 48k enabled longer-document retrieval; the original Llama2 7B sharply degraded after 4k.Source: §4.2 · p. 8, Figure 4 · p. 8

4.3 ABLATION STUDY

  • 15.82 perplexity at step 0 fell quickly; full fine-tuning converged faster than low-rank training, but their performance gap was small by the end.The curves came closer after 200 steps for Llama2 7B extending to 8192 context on PG19 validation.Source: §4.3 · pp. 8–9, Figure 5 · p. 8
  • Among shift settings, shifting between layers was acceptable but not best; using only pattern 1 or pattern 2 did not work, and shift direction had little effect.The comparison fine-tuned Llama2 7B to 32768 context with improved LoRA and evaluated PG19 perplexity; it also tested shift, dilated, block-sparse, and stride-sparse attention patterns.Source: §4.3 · p. 9, Table 6 · p. 9
  • Dilated attention performed well with full fine-tuning but poorly with low-rank adaptation, while stride sparse attention was harmful and lagged full attention.The tested alternatives included dilated, block sparse, and stride sparse attention, with similar computational costs.Source: §4.3 · p. 9, Table 6 · p. 9

5 CONCLUSION … A.2 FORMAT OF PASSKEY RETRIEVAL

LongLoRA combines sparse training attention with trainable embeddings and normalization to extend LLM context efficiently while retaining standard attention at inference. The paper also describes its experimental environment and passkey-retrieval input format.

5 CONCLUSION

  • LongLoRA extends context with less GPU memory and training time than standard full fine-tuning, with minimal accuracy compromise.Source: §5 · p. 9
  • S2-Attn approximates standard self-attention during training, takes two lines of code to implement, and preserves standard attention at inference.This allows reuse of most pre-existing infrastructure and optimization.Source: §5 · p. 9
  • Trainable normalization and embeddings bridge the gap between LoRA and full fine-tuning for context extension.Source: §5 · p. 9
  • 100k context for Llama2 7B and 32k for Llama2 70B are achieved on a single 8× A100 machine; the work also introduces LongAlpaca and uses it for supervised fine-tuning.Source: §5 · p. 9

A.1 ENVIRONMENTS

  • Experiments use 8× A100 GPUs, PyTorch (Paszke et al., 2019), DeepSpeed, and Flash-Attention2 (Dao, 2023).DeepSpeed stage 2 is the default, stage 3 is used for maximum-context experiments, and gradient checkpointing is enabled by default; 3090 Ti GPUs can suffice for 7B fine-tuning to 8192 context.Source: §A.1 · p. 13

A.2 FORMAT OF PASSKEY RETRIEVAL

  • Passkey retrieval places a randomly sampled passkey amid repeated irrelevant text, with M and N determining document length.The input asks the model to retrieve the passkey after the surrounding text.Source: §A.2 · p. 13

B EXPERIMENTS · B.1 EVALUATION PERPLEXITY ON PG19 TEST SPLIT.

On the PG19 test split, models achieve better perplexity as evaluation context length increases, while PG19 perplexity is higher than on proof-pile, which the paper attributes to different writing styles.

B.1 EVALUATION PERPLEXITY ON PG19 TEST SPLIT.

  • As evaluation context length increases, models trained for a given context length achieve better perplexity on PG19.The evaluation uses the same settings as the proof-pile evaluation.Source: §B.1 · p. 13
  • PG19 perplexity is higher than proof-pile perplexity, which the paper attributes to PG19’s different writing styles (Rae et al., 2020).Source: §B.1 · p. 13

B.2 ABLATION ON GROUP SIZES.

The ablation finds that S2-Attn group sizes of 1/2 and 1/4 of the target context length have only minor gaps to full-attention fine-tuning, while smaller groups are not good enough; LongLoRA defaults to 1/4.

B.2 ABLATION ON GROUP SIZES.

  • Group sizes of 1/2 and 1/4 have minor gaps to full-attention fine-tuning, whereas sizes below 1/4 are not good enough.The study varies S2-Attn group size across 1/2, 1/4, 1/6, and 1/8 of the target context length.Source: §B.2 · p. 13, Table 7 · p. 13
  • LongLoRA uses a group size of 1/4 of the context length by default.For example, 1/8 of an 8192-token context corresponds to a group size of 1024.Source: §B.2 · p. 13
  • The ablation fine-tunes Llama2 7B to 8192- and 16384-token contexts and evaluates on the PG19 validation set.It compares group sizes from 1/2 to 1/8 of the target context length against standard full attention.Source: Table 7 · p. 13, §B.2 · p. 13

B.3 ABLATION ON THE VARIANTS OF S2-ATTN.

S2-Attn variants—including reversing the shift and changing how shifted tokens are grouped or ordered—show similar perplexity to the original, while the latter variants require additional steps.

B.3 ABLATION ON THE VARIANTS OF S2-ATTN.

  • Reversing the shift direction has no effect on perplexity, and variants using separate groups or swapping shifted and front tokens achieve similar perplexity to the original.The variants were evaluated by fine-tuning Llama2 7B to 8192 context and testing on the PG19 validation set.Source: §B.3 · p. 14, Table 8 · p. 14
  • Separate groups for shifted tokens and swapping shifted with front tokens address the concern that moving last tokens to the front may conflict with causal masks.The authors suggest communications among front and last tokens are limited within the local group because those tokens were originally far from others.Source: §B.3 · p. 14
  • S2-Attn is used only during fine-tuning; inference uses standard causal masks and full attention.Variants 2 and 3 also work well, but involve additional steps compared with the original S2-Attn.Source: §B.3 · p. 14

B.4 EVALUATION ON LONG-CONTEXT BENCHMARKS.

On LongBench and LEval, LongLoRA’s Llama2 7B model at 16,384-token context achieved comparable or better performance than other Llama2-based long-context models, with about four hours of fine-tuning on one 8× A100 machine.

B.4 EVALUATION ON LONG-CONTEXT BENCHMARKS.

  • About 4 hours and 0.3 billion tokens on one 8× A100 machine produced a 16,384-context Llama2 7B model with comparable or better LongBench and LEval performance than other Llama2-based long-context models.The evaluation used the supervised fine-tuning method and data introduced in Section B.6; GPT-3.5-Turbo was also included in the comparison.Source: §B.4 · p. 14

B.5 EFFICIENCY ANALYSIS.

LongLoRA’s efficiency gains come from S2-Attn reducing attention’s growing share of FLOPs and lowering training time. With Flash-Attention2, peak memory differs little from LoRA; without it, S2-Attn avoids an out-of-memory limit at 16,384 context.

B.5 EFFICIENCY ANALYSIS.

  • Attention accounts for 24.5% of total FLOPs at 8,192 context and 72.2% at 65,536, falling to 39.4% with S2-Attn.Source: Table 11 · p. 15, §B.5 · p. 15
  • At 65,536 context, LongLoRA uses 56.6% of plain LoRA’s (Hu et al., 2022) training hours, while peak memory differs little with Flash-Attention2.The authors attribute the limited peak-memory difference to highly optimized Flash-Attention2 (Dao, 2023).Source: §B.5 · p. 15
  • Without Flash-Attention2 (Dao, 2023), S2-Attn provides greater speedup, including with LoRA+, and supports 16,384-context fine-tuning on 8× A100 where full attention runs out of memory.Source: §B.5 · p. 15

B.6 SUPERVISED FINE-TUNING.

The section describes LongAlpaca-12k and supervised fine-tuning to improve question-answering with long-context material, then presents example model answers.

B.6 SUPERVISED FINE-TUNING.

  • To improve QA ability where RedPajama-fine-tuned models had good perplexity but limited chat ability, the authors created LongAlpaca-12k with 9k long-context QAs and 3k short Alpaca QAs.The short QAs were sampled from the original Alpaca data.Source: §B.6 · pp. 15–16
  • LongAlpaca-12k covers technical papers, science fiction, books, and other materials; potentially harmful content was filtered, and questions address summarization, relationships, characters, and limitations.The prompt supplies material type and content before asking a question designed for that material.Source: §B.6 · p. 16
  • For SFT, the authors trained for 5 epochs using the same learning rate, weight decay, and batch sizes as the context-extension step.The example questions shown in Figures 8 and 9 were not in the training set.Source: §B.6 · p. 16
  • On book-related examples, the SFT 13B model gives more structured, concise answers grounded in detailed context than chat Llama2 13B without long documents; paper examples address contributions, limitations, and summaries.The book examples provide the related sections to the SFT model and probe the chat model using book title and question.Source: Figure 8 · p. 18, Figure 9 · p. 19

Related papers

Loading 2309.12307v3…