Some of these approaches are retrieval-based, which augment language models via fetching related documents and including the retrieved results into contexts. Our work is complementary to these works, as our attention mechanism is unmodified during inference. Many works modify multi-head attention to be approximated ones. They alleviate the quadratic complexity of the self-attention computation. For example, Longformer and BigBird use sparse attention to handle long sequences. Other works utilize memory mechanisms as a compression on past inputs, to look up relevant tokens. One limitation of these works is that these compressions have a large gap to full attention, making it infeasible to fine-tune pre-trained LLMs. Although our work also involves an approximation of attention mechanism, it has a similar shape and a small gap to standard attention. This enables fine-tuning pre-trained LLMs on S2-Attn and maintain full attention during inference. Long-context LLMs. LLMs are typically pre-trained with a pre-defined context length, such as 2048 for LLaMA and 4096 for Llama2. Training LLMs with long context from scratch is prohibitively expensive for most researchers. Recently, several works have tried to extend the context length of LLMs via fine-tuning. Position Interpolation modifies rotary position encoding and extends the context length of LLaMA to 32768. Focused Transformer utilizes contrastive learning to train LongLLaMA. Both of them rely on full fine-tuning, which is computationally expensive (128 A100 GPUs / 128 TPUv3 for training). Landmark attention is an efficient approach, but somewhat lossy. It compresses long context inputs into retrieved tokens. Our method saves substantial fine-tuning costs, while preserving the quality of the original attention. Ours maintain full access to the entire input via unmodified attention during inference. Some literature focuses on the position embedding modification of LLMs for long context extension, including Position Interpolation, NTK-aware, Yarn, positional Skipping, and methods based on out-of-distribution analysis. Our method focuses on efficient fine-tuning and retaining the original architecture during inference, which is orthogonal to these position embedding methods. Our models apply the Position Interpolation in experiments. Efficient Fine-tuning. This work is based on LoRA, a classical efficient fine-tuning approach. In addition to LoRA, there are many other parameter-efficient fine-tuning methods, including prompt tuning, prefix tuning, hidden state tuning, bias tuning, and masked weight learning. Input-tuning introduces an adapter to tune input embedding. Although the input embedding layers are also trainable in ours, this is not enough for long context extension. We make a comprehensive analysis on layer types in experiments. Existing work shows sparse masks can effectively save training costs and avoid performance drops. Section 3: LongLoRA  Subsection 3.1: Background  Transformer. LLMs are typically built with transformers. Taking Llama2 for example, as shown in Figure 2, an LLM model consists of an embedding input layer and a number of decoder layers. Each decoder layer comprises a self-attention module. It maps input features into a set of queries, keys, and values {q, k, v}, via linear projection layers with weight matrices {Wq, Wk, Wv}. Given {q, k, v}, it computes the outputs by taking the softmax of the query-key dot product and multiplying it by the value. The outputs are then projected by a linear layer with a weight matrix Wo. And MLP layers are followed. Before and after self-attention modules, layer normalization is applied. A final normalization is conducted after all decoder layers. For long sequences, self-attention struggles with computation cost, which is quadratic to the sequence length. This dramatically slows down the training procedure and increases GPU memory costs. Low-rank Adaptation. LoRA hypothesizes that the weight updates in pre-trained models have a low intrinsic rank during adaptation. For a pre-trained weight matrix W, it is updated with a low-rank decomposition where the change in weight is represented as the product of two smaller matrices, B and A. The rank r is much smaller than the dimensions of the original matrix. During training, W is frozen with no gradient updates, while A and B are trainable. This is the reason why LoRA training is much more efficient than full fine-tuning. In the Transformer structure, LoRA only adapts the attention weights (Wq, Wk, Wv, Wo) and freezes all other layers, including MLP and normalization layers. This manner is simple and parameter-efficient. However, we empirically show that only low-rank adaptation in attention weights does not work for long context extension. Subsection 3.2: Shifted Sparse Attention Standard self-attention costs are quadratic in sequence length, making LLMs on long sequences high memory cost and slow. To avoid this issue during training, we propose Shifted Sparse Attention (S2-Attn), as shown in Figure 2. In the following, we make a pilot study and explain our design step by step. Pilot Study. We build up a standard baseline that is trained and tested with full attention and fine-tuning, which presents consistently good quality in various context lengths. The first trial is to train with short attention, only pattern 1 in Figure 2. As we know for a long context, the high cost mainly comes from self-attention modules. Thus, in this trial, since the input is long, we split into several groups in self-attention. For example, the model takes 8192 tokens as input in both the training and testing stages, but self-attention is conducted in each group with a 2048 size. The group number is 4, as ablated in Section B.2 in the appendix. This pattern is efficient but still does not work in a very long context. The perplexity becomes larger as the context length increases. The reason behind this is that there is no information exchange between different groups. To introduce communication between groups, we include a shifted pattern, as shown in Figure 2. We shift the group partition by half group size in half attention heads. Taking the overall 8192 context length for example, in pattern 1, the first group conducts self-attention from 1st to 2048th tokens. In Pattern 2, the group partition is shifted by 1024. The first attention group begins from 1025th and ends at 3072th tokens, while the first and the last 1024 tokens belong to the same group. We use patterns 1 and 2 in each half self-attention heads respectively. This manner does not increase additional computation costs but enables the information flow between different groups. We show that it gets close to the standard attention baseline. Consistency to Full Attention. Existing efficient attention designs can also improve the efficiency of long-context LLMs. However, most of them are not suitable for long-context fine-tuning. Because, these transformers, designed for training from scratch, have gaps to the standard full attention, which is used in pre-training. We show that S2-Attn not only enables efficient fine-tuning but also supports full attention testing. Although other attentions can also be used in long context fine-tuning, models must be tested with the attention used during fine-tuning. Shifting prevents models from being over-fitted to specific attention patterns. Easy Implementation. S2-Attn is easy to implement. It involves only two steps: (1) shifting tokens in half attention heads, and (2) transposing features from token dimension to batch dimension. Two lines of code are enough. Subsection 3.3: Improved LoRA for Long Context LoRA is an efficient and popular manner for adapting LLMs to other datasets. It saves much trainable parameters and memory cost, compared to full fine-tuning. However, adapting LLMs from short context length to long is not easy. We empirically observe an obvious gap between LoRA and full fine-tuning. The gap between LoRA and full fine-tuning grows as the target context length becomes larger. And LoRA with larger ranks cannot reduce the gap. To bridge this gap, we open embedding and normalization layers for training. They occupy limited parameters but make effects for long context adaptation. Especially for normalization layers, the parameters are only 0.004 percent in the whole Llama2 7B. We denote this improved version of LoRA as LoRA+ in experiments. Section 4: Experiment Subsection 4.1: Experimental Settings  Models We extend the pre-trained 7B, 13B, and 70B Llama2 models. The maximum extended context window sizes are up to 100k for 7B models, 65536 for 13B models, and 32768 for 70B models. The position indices for these models are re-scaled with Position Interpolation. Training Procedure We follow most training hyper-parameters in Position Interpolation, except that our batch size is smaller as we use a single 8 by A100 GPUs machine in some cases. All models are fine-tuned via the next token prediction objective. We use AdamW with beta 1 = 0.9 and beta 2 = 0.95. The learning rate is set to 2 times 10 to the power of negative 5 for 7B and 13B models, and 10 to the power of negative 5 for 70B models. We also use a linear learning rate warmup. The weight decay is zero. We set the per-device batch size as 1 and gradient accumulation steps as 8, which means that the global batch size equals 64, using 8 GPUs. We train our models for 1000 steps. Datasets We use the Redpajama dataset for training. We evaluate the long-sequence language modeling performance of our fine-tuned models on the book corpus dataset PG19 and the cleaned Arxiv Math proof-pile dataset. We use the test split of PG19, consisting of 100 documents. For the proof-pile dataset, we also use the test split of it for evaluation. We follow Position Interpolation for proof-pile data processing. We evaluate perplexity by using a sliding window approach with S=256. Subsection 4.2: Main Results  Long-sequence Language Modeling. We report the perplexity for our models and baseline on proof-pile and PG19 datasets. Under certain training context lengths, our models achieve better perplexity with longer context sizes. This indicates the effectiveness of our efficient fine-tuning method. For the same training and evaluation context length cases, the perplexity decreases as the context size increases. By increasing the context window size from 8192 to 32768, for the Llama2 7B model, we observe that the perplexity gets better from 2.72 to 2.50 by -0.22. For Llama2 13B model, we observe that the perplexity reduces by -0.28. We further examine the maximum context length that we can fine-tune on a single 8 by A100 machine. We extend Llama2 7B, 13B, and 70B to 100k, 65536, and 32768 context length respectively. LongLoRA achieves promising results on these extremely large settings. In addition, we find some perplexity degradation on small context sizes for the extended models. This is a known limitation of Position Interpolation. Retrieval-based Evaluation. We conduct experiments on retrieval in long contexts. We compare our model with other open LLMs on the topic retrieval task introduced in LongChat. This task is to retrieve the target topic from a very long conversation, with lengths varying from 3k, 6k, 10k, 13k, to 16k. As some questions in LongChat are longer than 16k, we fine-tuned Llama2 13B with a context length of 18k. The training cost is similar to that for 16k. Our model achieves comparable performance to LongChat-13B, the state-of-the-art model in this task. Unlike LongChat-13B, which is fully fine-tuned on self-collected long context conversation text, our model is efficiently adapted on RedPajama via next-token generation. Our model even slightly outperforms LongChat-13B in the 16k evaluation. We present the passkey retrieval accuracy of our model, following Landmark Attention. This task has also been adopted by other literature. In this task, the models need to find a random passkey hidden in a long document. We show the document format is in Section B.1 in the appendix. We study Llama2 7B and our LongLoRA model which fine-tunes Llama2 7B with 32768 context length. We test the passkey retrieval accuracy from 1k to 34k, with an interval of roughly 1k (as the sentence length can not be precisely controlled). For each document length, we test the model 10 times with different random passkey values.