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 pre-trained LLMs to long contexts is computationally expensive, and standard LoRA alone performs poorly for this setting. LongLoRA combines shifted sparse attention with trainable embedding and normalization layers, achieving long-context extensions while retaining standard inference attention and supporting long instruction fine-tuning.

  • Problem

    Extending pre-trained LLMs to long contexts requires substantial computation, while LoRA alone leaves an increasing performance gap from full fine-tuning as target context length grows.

  • Method

    LongLoRA combines shifted sparse attention during fine-tuning with LoRA that also trains embedding and normalization layers.

  • Results

    LongLoRA achieves comparable performance to full-attention and fully fine-tuned results while extending Llama2 7B to 100k context and Llama2 70B to 32k on one 8× A100 machine.

  • Takeaways & Limitations

    LongLoRA extends context windows while retaining standard attention at inference and enabling reuse of existing optimization and infrastructure.

  • Takeaways & Limitations

    Plain LoRA with adaptation restricted to attention weights does not work sufficiently for long-context extension.

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 addresses the high cost of extending pre-trained LLMs to longer contexts by combining shifted sparse attention during fine-tuning with an improved LoRA setup. It retains standard attention at inference, achieves long-context extensions with reduced computation, and supports long instruction fine-tuning.

  • Motivation: Long-context fine-tuning is computationally expensive, with prior approaches requiring tens to hundreds of accelerators.Position Interpolation used 32 A100 GPUs for 2k-to-8k extension and 128 A100 GPUs for longer-context fine-tuning; FOT used 32 or 128 TPUs.
  • Method: LongLoRA uses shifted sparse attention during fine-tuning to approximate long-context self-attention through local groups and cross-group token shifts.Attention is computed within groups, while shifting half the attention heads by half a group size enables information flow between neighboring groups.
  • Method: The trained model retains standard self-attention at inference, preserving compatibility with existing optimization and infrastructure such as Flash-Attention2.The inference architecture is unchanged even though sparse attention is used during training.
  • Method: LongLoRA makes embedding and normalization layers trainable alongside LoRA weights, which the authors identify as pivotal for long-context adaptation.Embedding and normalization parameters represent less than 2% and at most 0.004% of Llama2 7B parameters, respectively.
  • Results: LongLoRA achieves comparable performance to full-attention and fully fine-tuned models at substantially lower computational cost.It extends Llama2 7B to 100k context or Llama2 70B to 32k context on a single 8× A100 machine.
  • Results: The authors also provide supervised fine-tuning with LongAlpaca, a self-collected long instruction-following dataset.The dataset includes questions for technical papers, science fiction, and other books.

2 RELATED WORK

Prior long-context work uses retrieval, modified attention, position-embedding changes, or expensive full fine-tuning. LongLoRA instead focuses on efficient fine-tuning while preserving the original attention architecture during inference.

  • Long-context Transformers: Retrieval-based methods augment language models by fetching related documents and adding retrieved results to the context.LongLoRA is described as complementary because its attention mechanism remains unmodified during inference.
  • Long-context Transformers: Several long-context approaches approximate multi-head attention, whereas LongLoRA preserves full access to the input through unmodified inference attention.The paper contrasts its approach with compressed retrieved-token methods that can be lossy.
  • Long-context LLMs: Earlier long-context LLM methods relying on full fine-tuning required 128 A100 GPUs or 128 TPUv3 devices.The cited examples include Position Interpolation and Focused Transformer.
  • Long-context LLMs: Position-embedding methods extend context by modifying rotary embeddings or related positional schemes, while LongLoRA targets efficient fine-tuning.The paper presents these directions as orthogonal to its focus on retaining the original inference architecture.
  • Efficient Fine-tuning: Parameter-efficient fine-tuning includes LoRA, prompt tuning, prefix tuning, hidden-state tuning, bias tuning, masked weight learning, and input tuning.LongLoRA trains input embedding layers but reports that this alone is insufficient for long-context extension.

3 LONGLORA

LongLoRA combines shifted sparse attention with an improved LoRA setup to make long-context fine-tuning more efficient while retaining compatibility with full-attention evaluation.

  • Background: O(n^2) self-attention makes long-sequence training slow and memory-intensive.The computational burden grows quadratically with sequence length.
  • Shifted Sparse Attention: S2-Attn splits long sequences into groups and shifts the partition in half the attention heads to communicate across groups.The implementation uses token shifting and feature transposition, with two code lines sufficient for the procedure.
  • Shifted Sparse Attention: S2-Attn supports efficient fine-tuning while allowing full attention during testing, unlike alternatives that require matching training and testing attention patterns.The shifting is intended to prevent overfitting to a specific attention pattern.
  • Improved LoRA for Long Context: Standard LoRA shows an increasing gap from full fine-tuning as the target context length grows, and larger ranks do not close it.This limitation motivates changes beyond adapting only attention weights.
  • Experiments: LongLoRA achieves performance comparable to full-attention or fully fine-tuned baselines with lower efficiency costs across Llama2 7B and 13B experiments.The method combines S2-Attn and LoRA+ for long-context fine-tuning.

4 EXPERIMENT

Experiments evaluate LongLoRA across model sizes, language modeling, retrieval, and attention ablations. The method reaches very long context lengths with comparable performance and reduced fine-tuning cost.

  • Experimental settings: LongLoRA fine-tunes Llama2 7B, 13B, and 70B to maximum context lengths of 100k, 64k, and 32k, respectively, on a single 8× A100 machine.The reported settings use Flash-Attention2 and DeepSpeed stage 3 during fine-tuning.
  • Retrieval-based evaluation: LongLoRA achieves comparable performance to LongChat-13B on topic retrieval across conversations ranging from 3k to 16k context lengths, with lower fine-tuning cost.The model is fine-tuned on Llama2 13B because some evaluation questions exceed 16k tokens.
  • Long-sequence language modeling: For long-sequence language modeling, perplexity improves as evaluation context length increases under the reported training settings.The experiments use RedPajama for training and PG19 and proof-pile datasets for evaluation.
  • Retrieval-based evaluation: 33k or 34k: the 32k-fine-tuned 7B model retains reasonable passkey retrieval accuracy without further fine-tuning.Extending position interpolation to 48k further supports longer-document retrieval.
  • Ablation study: LongLoRA’s S2-Attn is evaluated against shift, dilated, block-sparse, and stride-sparse attention patterns for 32k context fine-tuning.The experiments vary shifting across attention heads or sequential layers and compare alternative efficient attention designs.
  • Ablation study: 15.82 perplexity: without fine-tuning, the 7B model has limited long-context capability, while perplexity drops quickly during training.Full fine-tuning converges faster initially, but the methods approach each other after 200 steps without a large final gap.

5 CONCLUSION

The conclusion presents LongLoRA as an efficient approach for substantially extending LLM context lengths while preserving the standard inference architecture. It combines S2-Attn, trainable embeddings and normalization, and LongAlpaca supervised fine-tuning.

  • Conclusion: LongLoRA reduces GPU memory cost and training time relative to standard full fine-tuning, with minimal accuracy compromise.Its S2-Attn approximates standard self-attention during training and requires only two lines of code.
  • Conclusion: Models trained with S2-Attn retain standard attention during inference, allowing existing infrastructure and optimization techniques to be reused.The conclusion specifically notes compatibility with pre-existing optimization and infrastructure.
  • Conclusion: LongLoRA extends Llama2 7B to 100k context and Llama2 70B to 32k context on a single 8× A100 machine.The method bridges LoRA and full fine-tuning through trainable normalization and embedding layers.
  • Conclusion: The authors additionally introduce LongAlpaca and conduct supervised fine-tuning with LongLoRA for long instruction-following.The dataset contains long questions and corresponding answers across technical papers, science fiction, and other books.

B.1 EVALUATION PERPLEXITY ON PG19 TEST SPLIT.

PG19 test-split evaluations show that perplexity changes with evaluation context length and differs from proof-pile results because the datasets have different writing styles.

  • PG19 evaluation: For models trained at a certain context length, perplexity improves as the evaluation context length increases on the PG19 test split.The reported evaluation uses the same settings as the proof-pile evaluation in the paper.
  • PG19 evaluation: PG19 perplexity is higher than proof-pile perplexity because PG19 contains substantially different writing styles.The comparison concerns evaluation results reported in Tables 14 and 15.

B.2 ABLATION ON GROUP SIZES.

The group-size ablation evaluates S2-Attn at 8192 and 16384 context lengths against full attention. Group sizes of one-half or one-quarter of the target context length remain close to full-attention fine-tuning.

  • Group-size ablation: 1/2 and 1/4: these S2-Attn group-size settings have minor gaps to full-attention fine-tuning on PG19 validation.The ablation varies group size across {1/2, 1/4, 1/6, 1/8} of the target context length.
  • Group-size ablation: Group sizes smaller than 1/4 of the target context length are not good enough in the reported experiments.The authors set the group size to 1/4 of the context length by default.

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

The study evaluates S2-Attn variants and compares LongLoRA with other long-context models on LongBench and LEval benchmarks.

  • S2-Attn variants: S2-Attn variants alter shifting direction, grouping, or token swapping to test design choices.
  • Long-context benchmarks: LongBench evaluation compares models using the benchmark’s reported per-column rankings.
  • Long-context benchmarks: LEval evaluation compares various models with GPT-3.5-Turbo using GPT-4-judged win rates.

B.4 EVALUATION ON LONG-CONTEXT BENCHMARKS.

LongLoRA is evaluated on LongBench and LEval after supervised fine-tuning Llama2 7B to a 16,384-token context, with performance comparable to or better than several baselines.

  • 16,384-token context: the Llama2 7B model is fine-tuned with supervised fine-tuning and the LongAlpaca data introduced in Section B.6.
  • LongLoRA is compared with GPT-3.5-Turbo and Llama2-based long-context models including Vicuna and LongChat.
  • The 7B model shows comparable or better performance than the compared long-context models on LongBench and LEval.

B.5 EFFICIENCY ANALYSIS.

The efficiency analysis examines computation, memory, and training time, showing that S2-Attn reduces long-context costs while preserving the original attention architecture for inference.

  • FLOPs analysis: The FLOPs profiling separates FFN, projection, attention, and other layers to track how attention costs change with context length.
  • S2-Attn variants: S2-Attn modifies shifting variants while retaining the method’s sparse-attention design.
  • FLOPs analysis: 72.2%: self-attention accounts for this share of total FLOPs at 65,536 context with full attention, versus 39.4% with S2-Attn.
  • Training efficiency: LongLoRA requires fewer training hours than LoRA and full fine-tuning, although peak-memory differences from LoRA are limited with Flash-Attention2.
  • Training efficiency: 56.6%: LongLoRA uses this fraction of LoRA’s training hours at 65,536 context length.
  • Training efficiency: Without Flash-Attention2, S2-Attn provides greater speedup, while full attention reaches out-of-memory at 16,384 context on an 8× A100 machine.

B.6 SUPERVISED FINE-TUNING.

The supervised fine-tuning setup targets long-document question answering with LongAlpaca, using varied materials and questions, while efficiency comparisons and examples support the broader training approach.

  • Dataset: The dataset covers technical papers, science fiction, and books, with questions involving summarization, relationships, characters, and limitations.
  • Efficiency comparison: Without Flash-Attention2, S2-Attn improves training speed by 2.1× and GPU memory cost by 1.8× at 8,192 context length.
  • Prompt format: The prompt places material before a question and asks the model to memorize the content before answering.
  • Dataset: LongAlpaca-12k contains 9,000 long-context question-answer pairs and 3,000 short pairs sampled from Alpaca.
  • Training setup: Supervised fine-tuning uses five training epochs and the same learning rate, weight decay, and batch sizes as context extension.
  • Examples: In a book-question example, the model gives a structured answer identifying Pettigrew as Sirius Black’s intended target rather than Harry.
  • Examples: The examples include book-related questions and questions about papers’ contributions, limitations, and summarizations.
Loading 2309.12307v3…