How to Train Long-Context Language Models (Effectively)

## Abstract

Abstract We study continued training and supervised fine-tuning (SFT) of a language model (LM) to make effective use of long-context information.
We first establish a reliable evaluation protocol to guide model development—instead of perplexity or simple needle-in-a-haystack (NIAH) tests,
we use a broad set of long-context tasks, and we evaluate models after SFT with instruction data as this better reveals long-context abilities.
Supported by our robust evaluations, we run thorough experiments to decide the data mix for continued pre-training, the instruction tuning dataset, and many other design choices. We find that
(1) code repositories and books are excellent sources of long data, but it is crucial to combine them with high-quality short data;
(2) training with a sequence length beyond the evaluation length boosts long-context performance;
(3) for SFT, using only short instruction datasets yields strong performance on long-context tasks.
Our final model, ProLong-8B , which is initialized from Llama-3 and trained on 40B tokens,
demonstrates state-of-the-art long-context performance among similarly sized models
at a length of 128K. ProLong outperforms Llama-3.1-8B-Instruct on the majority of long-context tasks despite having seen only 5% as many tokens during long-context training.
Additionally, ProLong can effectively process up to 512K tokens, one of the longest context windows of publicly available LMs. 1 1 1 Our code, data, and models are available at https://github.com/princeton-nlp/ProLong .

## 1 Introduction

The ability of language models (LMs) to process extremely long inputs (for example, 128K tokens) has enabled new applications, such as book summarization or learning new tasks on the fly from many examples. However, adapting
LMs to process long contexts is challenging from an infrastructure and data perspective, and many design decisions are not well understood by open-source practitioners.

While many works have focused on extending the context length of pre-trained LMs with minimal training (Chen et al., ; Peng et al., ),
Fu et al. () find that
the above methods cannot even perform the simple needle-in-a-haystack (NIAH; Kamradt, ) task
and it is necessary to continually train the LM on long documents for billions of tokens. Frontier open-source models, such as Llama-3.1 (Dubey et al., ) and Jamba (Lenz et al., ), also employ a long-context continued training stage, followed by supervised fine-tuning (SFT) on instruction data.
We adopt the same setting and study continued training and SFT of a pre-trained LM for effective long-context use.

We first establish a reliable evaluation protocol to provide a meaningful signal for model development.
Most existing works rely on either perplexity or NIAH for ablating training recipes.
We demonstrate that neither is robust for guiding the development
and opt for a broad range of downstream applications, such as retrieval-augmented generation (RAG), long-document summarization, and many-shot in-context learning (ICL).
Importantly, we also conduct our evaluations after performing SFT,
even for all our ablation runs on continued pre-training.
We observe that, on some long-context tasks, performance gains only emerge after SFT, which means that best design choices can differ before and after SFT.

Guided by our evaluation protocol, we run comprehensive experiments with Llama-3-8B (8K original context window; Dubey et al., ) to study each component of long-context continued training, including data mixture, data and length scaling, supervised fine-tuning, and many other design choices such as cross-document attention masking and position extrapolation.
Many of our findings are surprising or contradictory to existing claims, for example, (1) training only on long data hurts long-context performance,
(2) training on longer sequences than the evaluation length helps, and (3) SFT on only short instruction data is sufficient for good long-context performance.
We outline our main takeaways and the structure of the paper in the takeaway box at the beginning of this section.

Our final model, ProLong, achieves the best performance at a 128K context length among 10B-parameter models, while taking only $5\%$ of the data budget compared to Llama-3.1’s long-context training (Dubey et al., ).
ProLong has a maximum context length of 512K tokens, making it one of the longest-context LMs available.(^2^22Throughout the paper, we use binary prefixes K$=2^{10}$, M=$2^{20}$, and B=$2^{30}$.)

## 2 Guide Model Development With Meaningful Evaluations

A pre-requisite for training a strong LM is having a robust evaluation suite that can guide model development while tracking its utility in real-world applications.
While synthetic benchmarks like needle-in-a-haystack (NIAH; Kamradt, ) and RULER (Hsieh et al., ) have gained much popularity due to their simplicity and controllability, we are interested in a wider range of tasks that reflect practical usage, such as the ability to reason over the whole document.
In the following, we describe our evaluation protocols and
showcase why they are critical to our model development.

### 2.1 Evaluate on diverse and realistic tasks

We first make the decision to use HELMET (Yen et al., ) as our evaluation suite,
as it is one of the most comprehensive long-context benchmarks, covering the following tasks:

- •
Recall: Given a JSON file with random key-values pairs, retrieve the value for a key.
- •
RAG: Answer a question given retrieved Wikipedia documents (NQ, HotPotQA, PopQA).
- •
Re-ranking: Produce top-10 rankings from a shuffled list of documents (MSMARCO).
- •
ICL: Learn classification tasks from many in-context examples, where the #classes ranges from 6 to 151; average of 5 datasets (TREC coarse/fine, NLU, Banking77, Clinc-150).
- •
QA: Answer a question given a full-length book (NarrativeQA).
- •
Summarization: Summarize long legal documents (Multi-LexSum).

Overall, these diverse tasks reflect a range of long-context abilities including recall, reasoning, learning from context, and robustness to noisy inputs.
Yen et al. () also show that HELMET produces model performance trends that are more consistent with human perceptions unlike other long-context benchmarks.

We showcase the importance of a robust evaluation suite in Table 1.
As a predecessor of our work, Fu et al. () only consider needle-in-a-haystack (NIAH) and perplexity during model development; evaluations on 3 tasks from HELMET reveal major short-comings of their models despite perfect NIAH scores.
We also see how NIAH and even the HELMET recall task become saturated for strong models (Llama-3.1-8B vs. 70B) while other task categories continue to detect differences in their long-context abilities.

Table: Table 1: HELMET offers a more holistic long-context evaluation. We reproduce Fu et al. () on Llama-3-8B with SFT. We report the instruct Llama versions.

|  |  | HELMET |  |  |
| --- | --- | --- | --- | --- |
| Models | NIAH | Recall | RAG | Re-rank |
| Fu et al. () | 100 | 95.8 | 52.1 | 23.1 |
| Llama-3.1-8B | 100 | 99.4 | 56.3 | 37.0 |
| Llama-3.1-70B | 100 | 100 | 62.1 | 49.2 |

We offer more details about the HELMET evaluation, including its careful choice of metrics, in §A.1.
We did not use all tasks from HELMET for efficiency reasons and we also do not want to overfit to HELMET.
If not otherwise specified, we average the performance for each category over all datasets and over evaluation lengths of 32K and 64K;
for the final long-context score, we macro-average all categories.

Figure: Figure 1: Making design decisions based on perplexity (PPL) is not optimal for long-context downstream tasks.

##### Why not perplexity?

Besides synthetic recall tasks, many previous works rely on perplexity (PPL) for evaluating long-context extensions of LMs (Chen et al., ; Fu et al., ; Lu et al., ), which is commonly measured on the PG19 books dataset (Rae et al., ).
We use the ablation experiment from §3.2 to showcase why perplexity is not an indicative metric for developing long-context models.
The experiment studies how the ratio of long documents affects the performance. We report both our evaluation and the perplexity measured on the last 32K tokens of 64K-length documents from PG19.
As shown in Figure 1, while using more long data continues to improve PPL,
it is clear that using 100% long data significantly hurts downstream long-context performance.

### 2.2 Evaluate after supervised fine-tuning

Supervised fine-tuning (SFT; Ouyang et al., ) is an additional training stage that fine-tunes the model on a small amount of natural-language instructions and corresponding responses; it enables a base LM to address user queries in a chat format and has become a standard step for producing frontier LMs.
Here, we consider the difference between evaluating a model before or after SFT.

In preliminary experiments, we continue training Llama-3-8B-Base on 5B-token subsets from the data mix by Fu et al. (). The mix is based on SlimPajama (Soboleva et al., ) and upsamples long documents to constitute roughly 70% of tokens, while retaining the original domain proportions.
Then we conduct SFT on several intermediate checkpoints with UltraChat (Ding et al., ).

We show the benchmarking results before and after SFT in Figure 2.
Long-context evaluation shows clearer signals when it is conducted after SFT:
(1) SFT shows that the model continues to improve with more training tokens on RAG and re-ranking, while the improvement is less clear or does not exist when evaluated before SFT.
(2) SFT enables evaluation on realistic applications like QA and summarization, which require instruction following and have low performance before SFT.
We also note that the variance from two random training runs is not substantially higher after the additional SFT phase. Therefore, unless otherwise specified, we report the long-context performance after SFT.

Figure: Figure 2: Improvements on RAG and re-ranking tasks are only observed when evaluating models after a supervised fine-tuning (SFT) phase on instruction data. The models are trained on the pre-training data mix by Fu et al. (). We report the mean and standard deviations over two training runs.
Refer to caption: x2.png

We dive deeper into supervised fine-tuning in §5 and explore different training datasets, as well as the use of synthetic long instruction data. However, we find that simply fine-tuning on UltraChat remains a surprisingly competitive choice.

### 2.3 Check that short-context performance is preserved

Long-context abilities should not come at the expense of short-context performance, particularly since short-context evaluations cover a wider range of capabilities, e.g., world knowledge, commonsense, and mathematical reasoning.
However, short-context evaluation has largely been neglected by previous long-context research.
We report on 5 tasks from the the Open LLM Leaderboard (Beeching et al., ): HellaSwag (Zellers et al., ), MMLU (Hendrycks et al., ), ARC-challenge (Clark et al., ), WinoGrande (Sakaguchi et al., ), and GSM8K (Cobbe et al., ).
We evaluate short-context performance before SFT, since this allows for a direct comparison to the base model which was used as initialization for the long-context training.

Table: Table 2: Applying position extrapolation (PE) to Llama-3-8B by changing the RoPE frequency base (§B.1) or fine-tuning it on a long-context SlimPajama mixture (Fu et al., ; Soboleva et al., ) deteriorates the performance of this top-shelf pre-trained LM on short-context tasks.

|  | HSwag | MMLU | ARC-c | WG | GSM8K |
| --- | --- | --- | --- | --- | --- |
| Llama-3-8B | 82.1 | 66.5 | 59.4 | 77.1 | 44.7 |
| + PE | 81.5 | 64.7 | 58.1 | 75.5 | 40.1 |
| + SlimPajama | 81.0 | 63.1 | 57.8 | 75.1 | 40.6 |

##### Previous techniques deteriorate short-context performance.

We show in Table 2 that both training-free position extrapolation, as well as fine-tuning with an existing long data mixture (Fu et al., )
do not preserve the strong performance of Llama-3-8B on standard short-context tasks.
This motivates us to find data sources which retain the initial model’s strong short-context performance.

## 3 Long-Context Data Curation

The quality and composition of training data has been found to be the most important factor for LM pre-training (Penedo et al., ; Wettig et al., ; Li et al., ) and is therefore
a primary focus of our study. To make data decisions, we perform ablation experiments:
we continue to train Llama-3-8B-Base for 5B tokens with a maximum length of 64K tokens and evaluate according to §2. See §A.4 for more details of our ablation setting.

We aim to boost the long-context task performance while preserving the short-context performance of the original model. Starting from the intuition that the data should be a mixture of long and short documents, we study these choices separately.
In our ablations, the long data is comprised of single-document chunks of 64K tokens, whereas for the short data, we construct batches by packing documents until we reach 64K tokens per sequence.

### 3.1 Code repositories and books are good sources of long-context data

Table: Table 3: Long text documents ($\geq$64K tokens) by data sources.

| Data | #Long |
| --- | --- |
| tokens |  |
| Code Repos | 98.8B |
| SP/Books | 33.2B |
| SP/CC | 15.3B |
| SP/Arxiv | 5.2B |
| SP/GitHub | 2.8B |
| SP/Wiki | 0.1B |
| SP/StackEx | $<$0.1B |
| SP/C4 | $<$0.1B |

##### SlimPajama.

We analyze the quantity of long data in SlimPajama (SP; Soboleva et al., ). Table 3 shows that books account for the majority of long-context tokens.
When inspecting the long data in CommonCrawl (CC), we observe that though varied in quality, it also contains some book-like content, which future work could identify via data selection methods.

##### Code repositories.

While only few files from GitHub reach a very long length (which also tend to be lower quality as suggested by Singh et al., ), we construct an abundant source of long-context data from the Stack (Kocetkov et al., ) by concatenating all files from a repository to form a single document. Unlike Guo et al. (), we do not order the files based on dependencies, which should increase the distance between dependent files and reduce recency bias.

##### Data mixture.

We train models with 60% of long-context data and 40% of our ShortMix (§3.3).
Table 4 shows that using code repositories alone performs the best on stress-test recall tasks. Meanwhile, books are more broadly beneficial for in-context learning, summarization and re-ranking.
An equal mix of books and code repositories achieves the best overall performance.
Note that short-context task performance remains consistent due to our high-quality short data mix.

**Table 4: Impact of different long data sources, while keeping the 40% short data component fixed. Code repositories particularly helps recall, while books are more effective on re-ranking, ICL, and summarization. Mixing the two achieves the overall best performance.**
| Long Data (60%) | Long-Context | Short-Context |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Recall | RAG | Re-rank | ICL | QA | Summ. | Avg. | Avg. |  |
| CommonCrawl | 84.1 | 53.3 | 28.1 | 67.5 | 35.2 | 37.0 | 50.9 | 66.5 |
| Books | 94.9 | 53.9 | 30.7 | 72.2 | 33.2 | 37.7 | 53.8 | 65.5 |
| Code Repos | 99.2 | 53.8 | 29.0 | 61.2 | 34.7 | 36.2 | 52.3 | 65.9 |
| Books/Repos 1:1 | 96.0 | 54.9 | 29.4 | 73.9 | 35.7 | 37.9 | 54.6 | 65.5 |

### 3.2 Training only on long data hurts long-context performance

The ratio between short/long data is another crucial factor for downstream performance.
Prior work either trains only on long data (Peng et al., ) or adds some short training data (Yen et al., ; Fu et al., ).
However, we are the first to systematically study the impact of short/long ratio.

Figure 3 shows that
short task performance monotonically decreases as the long data increases.
The trends for long-context vary by tasks and are further complicated by SFT:
On tasks like recall and RAG, the performance before SFT prefers high proportions of long data, while the performance after SFT drastically deteriorates with more long data.
We hypothesize that specializing the model only on long data makes it a poor initialization for generic SFT—highlighting the importance of evaluating checkpoints after SFT (§2.2).
While some long-context tasks benefit from more long data consistently (ICL) or show no clear pattern (re-ranking), the best average performance is achieved at 60% long data and 40% short data, which we adopt for our final ProLong model.

Figure: Figure 3: Impact of short/long data ratio. All models are trained on books/repos long data and our ShortMix for 5B tokens. More long data initially improves long-context performance, but then becomes impairing. More long data also consistently degrades the short-context performance.
Refer to caption: x3.png

### 3.3 Choosing a high-quality short-context mix is important

We saw in §2.3 that it is difficult to preserve the strong performance of Llama-3-8B on short-context tasks during long-context fine-tuning.
We adopt our best long-context settings (Book/repo data and 60% long/40% short) and study the impact of different short-context training mixes. We experiment with SlimPajama (Soboleva et al., ), FineWeb-Edu (Penedo et al., ), DCLM-Baseline (Li et al., ), and our own ProLong ShortMix.
Our ShortMix is inspired by the “stage 2 training” in MiniCPM (Hu et al., ) and Dolma-1.7 (Soldaini et al., ), which use more knowledge-intensive, downstream-related data at the end of pre-training.
Table 5 shows the composition of our ShortMix.(^3^33 Since we do not truncate documents in the short data component unnecessarily, it includes a small percentage of documents longer than 8K. See Table 14 in the appendix for the dataset length statistics.)

Table: Table 5: Our ShortMix.

| Components | % |
| --- | --- |
| FineWeb | 25 |
| FineWeb-Edu | 25 |
| Wikipedia | 10 |
| Tulu-v2 | 10 |
| StackExchange | 10 |
| ArXiv | 10 |
| OpenWebMath | 10 |

**Table 6: Impact of different short data sources. The long-context performance is the average of 6 categories at the lengths of 32K and 64K.**
| Short Data (40%) | Long-Context | Short-Context |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Avg. | HellaS. | MMLU | ARC-c | WG | GSM8K | Avg. |  |
| Original model (Llama-3-8B) | - | 82.1 | 66.5 | 59.4 | 77.1 | 44.7 | 66.0 |
| SlimPajama | 52.9 | 81.2 | 63.0 | 58.5 | 76.2 | 41.9 | 64.2 |
| FineWeb-Edu | 53.0 | 81.0 | 62.6 | 57.7 | 74.4 | 39.4 | 63.0 |
| DCLM-Baseline | 52.0 | 82.0 | 65.6 | 59.6 | 77.4 | 39.4 | 64.8 |
| ProLong ShortMix | 54.6 | 81.6 | 65.3 | 58.0 | 76.2 | 46.6 | 65.5 |

Table 6 demonstrates that the short data component has a substantial impact on both short-context and long-context downstream performance.
Our curated ShortMix outperforms other short data sources on both short and long-context tasks and our data domains are particularly important for retaining Llama-3-8B’s performance on mathematical reasoning.
Surprisingly, we find that fine-tuning only using FineWeb-Edu—a dataset that is curated to help with knowledge-intensive tasks like MMLU—performs poorly as a short-context component, and we combine it with more diverse data sources in our ShortMix.
DCLM-Baseline performs well on all short-context tasks except for GSM8K.
This can likely be improved by combining with math-related datasets, but as we added the DCLM-baseline ablation at the conclusion of the project, we leave this exploration to future work.

## 4 Scaling the Size and Length of the Training Data

Training for more steps is well-known to improve downstream tasks in regular pre-training, but little analysis has been done in the context of long-context continued training.
We incorporate the lessons from our ablation experiments and arrive at the ProLong recipe, which we describe in detail in §6. Notably, we scale up the training budget to longer sequences (up to 512K) and more tokens (20B tokens at a maximum sequence length of 64K and an additional 20B tokens at 512K). We reset the learning rate schedule and increase the RoPE frequency base when switching from 64K to 512K context lengths.
In this section, we analyze the impact of these decisions.

##### Increasing the number of steps helps.

In Figure 4, we plot the downstream performance of intermediate checkpoints of our 40B-token runs.
While the long-context performance fluctuates throughout training, we observe positive trends on recall, RAG, re-ranking, and summarization.
For short-context tasks, we observe the average performance initially drops from the initialization, but gradually recovers. Performance again drops when switching from 64K to 512K sequence length, but also recovers with additional training.

Figure: Figure 4: Performance (avg. of 32K and 64K) of our ProLong model throughout training.
Refer to caption: x4.png

##### Increasing the training length beyond the evaluation length helps.

One might assume that we should train long-context models on the maximum sequence length that we want the model to support. Many works even emphasize extrapolation to even longer sequences at inference time (Press et al., ; Xiao et al., ; ; Yen et al., ; Chen et al., ).
In contrast, we observe that training on a longer sequence length (512K tokens) substantially improves the long-context performance at a shorter evaluation length (64K tokens).

We establish this by initializing
with a model that was trained for 20B tokens at 64K and either (1) continuing training at 64K, or (2) switching to the 512K training.
We use the same hyperparameters and data mixtures in either experiment.
We evaluate a checkpoint after 4B training tokens
at a evaluation length of 64K.
Comparing the two runs
in Table 7,
we see consistent gains from switching to the 512K training length.(^4^44While we demonstrate the benefit of longer data, we note that training with longer sequences is more expensive, and may therefore not be the computationally optimal choice.)

**Table 7: Impact of training models on different sequence lengths. All the results are evaluated at a sequence length of 64K. We see that training at a maximum length beyond the evaluation context window consistently improves the long-context performance.**
| Max Seq. Length | Recall | RAG | Re-rank | ICL |
| --- | --- | --- | --- | --- |
| ProLong 64K training (20B) | 96.5 | 52.7 | 22.8 | 70.6 |
| +4B 64K training | 95.0 | 56.4 | 28.0 | 78.8 |
| +4B 512K training | 98.5 | 56.9 | 32.9 | 79.2 |

## 5 Supervised Fine-Tuning for Long-Context LMs

In this section, we study how to best enable long-context language models to follow instructions.
We focus on supervised fine-tuning on instruction datasets (Ouyang et al., ) and
leave reinforcement learning and preference optimization for future work.

All our experiments in this section use the ProLong base model, which was
trained for
40B tokens at a maximum sequence length of 512K.
In comparison, open-source instruction data are very short, e.g.,
UltraChat (Ding et al., ) conversations have 1.2K tokens on average and 4.1K tokens maximum.
To bridge this gap, several works (Xiong et al., ; Dubey et al., ; Xiong et al., )
have proposed to generate long instruction data synthetically.

We consider three popular SFT datasets—UltraChat (Ding et al., ), Tulu-v2 (Ivison et al., ), ShareGPT(^5^55[https://huggingface.co/datasets/RyokoAI/ShareGPT52K](https://huggingface.co/datasets/RyokoAI/ShareGPT52K).)—and three sources of synthetic data: For synthetic QA, we prompt Llama-3-8B-Instruct
to generate a question-and-answer pair given a random chunk from a long document; we reuse the QA pairs for synthetic RAG but we present
a random list of chunks from the document to mimic retrieved passages; for synthetic summarization, we generate summaries for long books via recursive summarization (Wu et al., ).
For all synthetic data, we write several templates, which we sample at random to increase diversity. More details can be found in §A.5.
We always use a combination of 40% synthetic QA, 30% synthetic RAG, and 30% synthetic summarization in our synthetic instruction dataset.
The hyperparameters for the instruction tuning experiments can be found in Table 9.

##### Short-context instruction data yields strong long-context results.

We first establish that UltraChat outperforms Tulu-v2 and ShareGPT
in Table 22.
We therefore use it when studying the ratio of synthetic long-context instruction data in Table 8.
Surprisingly, we find that adding synthetic data does not improve the performance on these very long-context tasks, and adding even as little as 1% synthetic data hurts the performance in our setting.
Therefore, we use only short-context UltraChat data for SFT of our final ProLong model.

**Table 8: Effect of different ratios of synthetic SFT data (mixed with UltraChat). We report the 32K-and-64K-averaged performance except tasks marked with ^†, which are evaluated at 512K for stress testing. The number of percentage is based on #tokens, not #samples.**
| % Synthetic Data | JsonKV^† | RAG | Re-rank | ICL | QA^† | Summ.^† | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 0% | 65.7 | 58.1 | 38.5 | 80.3 | 49.7 | 42.1 | 55.7 |
| 1% | 61.5 | 57.0 | 38.3 | 80.8 | 45.3 | 41.5 | 54.1 |
| 3% | 62.0 | 56.4 | 37.9 | 80.6 | 44.8 | 39.5 | 53.5 |
| 10% | 70.3 | 55.5 | 36.1 | 80.6 | 41.7 | 39.4 | 53.9 |
| 50% | 45.8 | 48.8 | 18.8 | 70.5 | 42.3 | 33.3 | 43.3 |

Why do our conclusions about synthetic data differ from previous work?
We offer the following hypotheses:
(1) Previous work like Xiong et al. (); Bai et al. () may have insufficient long-context training and the synthetic data acts as additional long-context training data.
(2) Our instruction dataset is much smaller compared to the private instruction data used for Llama-3.1 (Dubey et al., )—it is possible that when using an extensive short instruction dataset, mixing in synthetic long data avoids the model from degenerating on long-context tasks.

**Table 9: The training recipe for ProLong.**
| Continued Long-context Training |  |  |
| --- | --- | --- |
| Data | 30% code repos, 30% books, 3% textbooks, 37% ShortMix |  |
|  | ShortMix: | 27% FineWeb-Edu, 27% FineWeb, |
|  |  | 11% Tulu-v2, 11% StackExchange, |
|  |  | 8% Wikipedia, 8% OpenWebMath, 8% ArXiv |
| Length<br>Curriculum | Stage 1 (64K): | Code repos, books, and textbooks at length 64K |
| Stage 2 (512K): | Code repos: 50% at length 512K, 50% at length 64K |  |
|  |  | Books: 17% at length 512K, 83% at length 64K |
|  |  | Textbooks at length 512K |
| Steps | Stage 1: 20B tokens (2.2K H100 hours),   Stage 2: 20B tokens (12.2K H100 hours) |  |
| Model | Initialization: | Llama-3-8B-Instruct (original RoPE base freq. $5\times 10^{5}$) |
|  | RoPE: | Stage 1: $8\times 10^{6}$, Stage 2:   $1.28\times 10^{8}$ |
|  | Attention: | Full attention with cross-document attention masking |
| Optim. | AdamW (weight decay = $0.1$, $\beta_{1}=0.9$, $\beta_{2}=0.95$) |  |
|  | LR: | $1e-5$ with $10\%$ warmup and cosine decay to $1e-6$, each stage |
|  | Batch size: | 4M tokens for stage 1, 8M tokens for stage 2 |
| Supervised Fine-tuning (SFT) |  |  |
| Data | UltraChat |  |
| Steps | 1B tokens |  |
| Optim. | AdamW (weight decay = $0.1$, $\beta_{1}=0.9$, $\beta_{2}=0.95$) |  |
|  | LR = $2e-5$ (cosine decay to $2e-6$), warmup = $5\%$ |  |
|  | Batch size = 4M tokens |  |

## 6 The ProLong Model: Recipe and Results

### 6.1 Final recipe

We summarize the training recipe for ProLong in Table 9.
Our final model starts from the Llama-3-8B-Instruct model and is trained on 64K sequence length for 20B tokens.
It is then further trained on 512K sequence length for 20B tokens (ProLong base), which we achieve using sequence parallelism (Li et al., ).
We obtain the final ProLong model via SFT of the base model on UltraChat.
One small difference on the data mixture between our ablations and the final model is that we mix in 3% high-quality textbooks (Chevalier et al., ), as book-like data are shown to be beneficial for long-context (§3.1) and textbooks are highly educational. This also slightly changes the proportions of ShortMix.
You can find more details about our data processing (§A.2) and the training stack (§A.3) in the appendix.

In the following, we elaborate on several carefully ablated design choices in our recipe.

##### RoPE frequency base tuning.

We find that changing the RoPE (Su et al., ) frequency base to achieve position extrapolation (Xiong et al., ; emozilla, ) significantly improves long-context performance, even with a significant amount of training. §B.1 shows our ablation on the best RoPE base to use. While the original Llama models use a RoPE base of $10^{5}$, we use a base of $8\times 10^{6}$ for the 64K setting and $1.28\times 10^{8}$ for the 512K setting.

##### Disabling cross-document attention.

Ding et al. () show that masking out attention across document boundaries improve model performance and this was also used during Llama-3 pre-training (Dubey et al., ).
In §B.2, we show that disabling cross-document attention in continued training benefits both the short and long-context performance.
Disabling cross-document attention can also result in higher training throughput, which we describe in more detail in §A.3.

##### Starting from Llama-3-8B-Instruct.

While we conduct all our long-context training ablations with the base model of Llama-3-8B,
we use Llama-3-8B-Instruct as the initialization for the final ProLong model.
§B.3 shows that while slightly improving the long-context performance, Llama-3-8B-Instruct significantly enhances the short-context performance.

### 6.2 ProLong performance

We present the final HELMET evaluation results of ProLong in Table 10.
We compare to a number of frontier long-context LMs, namely MegaBeam(^6^66[https://huggingface.co/aws-prototyping/MegaBeam-Mistral-7B-512k](https://huggingface.co/aws-prototyping/MegaBeam-Mistral-7B-512k).), Llama-3.1 (Dubey et al., ), Qwen2 (Yang et al., ), Phi-3 (Abdin et al., ), Mistral-Nemo(^7^77[https://huggingface.co/mistralai/Mistral-Nemo-Instruct-2407](https://huggingface.co/mistralai/Mistral-Nemo-Instruct-2407).), Jamba-1.5 (Lenz et al., ), Claude-3.5-Sonnet (Anthropic, ), Gemini-1.5 (Reid et al., ), and GPT-4o (Achiam et al., ).

ProLong outperforms all 10B-scale models on our long-context evaluation. Notably, ProLong outperforms Llama-3.1-8B-Instruct
on all categories except summarization.
ProLong achieves this with only 5% of Llama-3.1’s long-context data budget (40B vs. 800B tokens).
We also showcase the strength of ProLong with several QA examples in Table 23.

**Table 10: Our main evaluation results on HELMET (Yen et al., ; details in §A.1). All results are averaged over sequence lengths of 32K, 64K, and 128K. For all models, we use the corresponding instruction version. ProLong is one of the best performing 10B-scale LMs. The complete set of results can be found in §C.**
| Model | Max Len. | Recall | RAG | ICL | Re-rank | QA | Summ. | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| ProLong (8B) | 512K | 99.4 | 66.0 | 81.1 | 33.2 | 40.8 | 40.5 | 60.2 |
| MegaBeam-Mistral (7B) | 512K | 99.4 | 58.1 | 82.1 | 22.1 | 33.7 | 43.6 | 56.5 |
| Meta-Llama-3.1 (8B) | 128K | 98.7 | 62.8 | 79.7 | 26.6 | 40.4 | 46.1 | 59.0 |
| Qwen2 (7B) | 128K | 34.4 | 43.4 | 54.8 | 4.6 | 23.3 | 38.5 | 33.2 |
| Phi-3-small (7B) | 128K | 74.8 | 60.6 | 82.0 | 18.5 | 34.1 | 42.4 | 52.1 |
| Mistral-Nemo (12B) | 128K | 24.9 | 48.1 | 82.0 | 4.7 | 37.7 | 37.0 | 39.1 |
| Jamba-1.5-Mini (12B/52B) | 256K | 87.7 | 61.3 | 88.4 | 25.9 | 42.0 | 38.6 | 57.3 |
| Meta-Llama-3.1 (70B) | 128K | 98.5 | 65.9 | 80.0 | 39.4 | 47.2 | 51.1 | 63.7 |
| Claude-3.5-Sonnet | 200K | 99.4 | 44.0 | 79.3 | 19.9 | 38.1 | 49.2 | 55.0 |
| Gemini-1.5-Pro | 2M | 94.2 | 71.4 | 78.9 | 65.3 | 44.4 | 56.2 | 68.4 |
| GPT-4o | 128K | 99.9 | 71.5 | 86.7 | 59.6 | 47.0 | 55.7 | 70.1 |

Since most existing models do not support more than 128K tokens,
to showcase ProLong’s 512K context length,
we stress test ProLong on the QA and summarization tasks from 32K to 512K(^8^88 In QA and summarization, we truncate the documents at the evaluation length; hence an effective long-context model should demonstrate better performance on longer lengths.). Table 11 shows that ProLong continues to improve at a longer context window.

Table: Table 11: ProLong at 512K.

|  | 32K | 64K | 128K | 512K |
| --- | --- | --- | --- | --- |
| QA | 31.7 | 43.7 | 46.7 | 49.7 |
| Summ | 40.4 | 39.8 | 41.5 | 42.1 |

**Table 12: Results on the NoCha benchmark (Karpinska et al., ).(^10^1010[https://github.com/marzenakrp/nocha](https://github.com/marzenakrp/nocha). NoCha has a private test set and all evaluation is done by the NoCha authors. Hence, we report models from Table 10 that are also on the NoCha leaderboard.) ProLong is the only model that achieves above-random performance in the $<$75K category and it consistently beats Llama-3.1. Different from the original NoCha leaderboard, we report the average accuracy over all test instances without filtering the test examples based on the model’s context window lengths.**
| Model | Max Len. | $<$75K | 75K-127K | 127K-180K | $>$180K |
| --- | --- | --- | --- | --- | --- |
| ProLong (8B) | 512K | 28.4 | 17.0 | 13.1 | 20.3 |
| MegaBeam-Mistral (7B) | 512K | 19.8 | 18.3 | 17.5 | 15.6 |
| Meta-Llama-3.1 (8B) | 128K | 17.3 | 16.4 | 0.0 | 0.0 |
| Mistral-Nemo (12B) | 128K | 13.6 | 0.4 | 0.0 | 0.0 |
| Jamba-1.5-Mini (12B/52B) | 256K | 27.2 | 28.0 | 24.4 | 6.2 |
| Meta-Llama-3.1 (70B) | 128K | 42.0 | 25.0 | 0.0 | 0.0 |
| Gemini-1.5-Pro | 2M | 24.7 | 38.8 | 35.3 | 46.9 |
| GPT-4o | 128K | 55.6 | 58.4 | 0.0 | 0.0 |

Besides HELMET, we also evaluate our models on NoCha (Karpinska et al., )—a claim verification dataset on 67 recently published English fictional books.
We chose this dataset because (1) it minimizes the data contamination problem as all the books are unlikely to exist in the model pre-training data; (2) all the claims are written by human readers and require global reasoning. Each test instance contains two contradictory claims, and the models must correctly judge both to pass.

footnote 10 demonstrates the NoCha evaluation results.
Among 10B-scale models, ProLong achieves the best accuracy on the extremely long test instances ($>$180K); on test instances $<$75K tokens, ProLong significantly outperforms other models and is the only model that is better than random guessing (25%).
This further showcases the strength of our training recipe and the ProLong model.

## 7 Related Work

##### Adapting existing LMs for long contexts.

Many works explore
extending the LM context windows with minimal training,
either by
position extrapolation (Chen et al., ; Peng et al., ; Chen et al., ; Ding et al., ; Liu et al., ; Zhang et al., ; Zhu et al., ; Zhao et al., ; Wu et al., ; Hu et al., )
or
manipulating the attention patterns (Chen et al., ; Xiao et al., ; ; Bertsch et al., ; Jin et al., ).
Yoshida et al. (); Choromanski et al. (); Chevalier et al. ()
instead explore the idea of compressing the long contexts into shorter forms.
However, Fu et al. (); Lu et al. () show that using full attention, applying simple position extrapolation, and fine-tuning the model on long documents reach much stronger results.

Llama 3.1 (Dubey et al., ) and Jamba (Lieber et al., ) achieve long-context capabilities by adding a long-context continued training stage between standard pre-training and supervised fine-tuning, which is the setting we follow.
Fu et al. () study the data engineering for this setting and argue that 0.5B tokens of domain-balanced, length-upsampled data is sufficient for acquiring the long-context recall ability—which we show is not sufficient if a more holistic evaluation is taken.
Xiong et al. (); Dubey et al. (); Lieber et al. (); Xiong et al. (); An et al. (); Bai et al. () also adopt synthetically-generated long data in the SFT stage; however,
we find that
using standard, short-context instruction data achieves the best long-context results in our setting.

##### Efficient long-context architectures.

There have been many efforts in designing more efficient architectures,
for example,
linear attention/RNNs (Gu & Dao, ; Dao & Gu, ; Ma et al., ; Sun et al., ; Peng et al., ; Yang et al., ),
and
alternative attention architectures (Rubin & Berant, ; Sun et al., ; Yen et al., ).
However,
they often require training from scratch and many have the inherent limitations in terms of long-context recall (Jelassi et al., ; Arora et al., ).
Recent works explore hybrid models (Waleffe et al., ; Lieber et al., )) or distilling existing LMs into hybrid models (Wang et al., ) and show promising results.

##### Long-context evaluation.

Many benchmarks have been proposed for long-context evaluation
(Shaham et al., ; Hsieh et al., ; Krishna et al., ; Zhang et al., ; An et al., ; Bai et al., )
There are works studying particular aspects of long-context LMs as well,
such as positional bias (Liu et al., ),
in-context learning (Bertsch et al., ; Li et al., ), and book-length summarization (Kim et al., ).
In this work, we follow Yen et al. ()
for its diverse application coverage and reliable evaluations.

## 8 Conclusion

We study the problem of given a short-context pre-trained LM, how to
most effectively continually pre-train and SFT the model to be long-context.
We conduct thorough ablations on each component and
many of our findings contradict existing practices or beliefs.
We use all the findings to produce ProLong, a new state-of-the-art long-context LM. We release all our code, data, and models publicly
and hope that our findings will
boost research and applications of long-context LMs.

#### Limitations

Although we try to ablate the major components of our training recipe,
due to resource limitations, we cannot exhaust all aspects, such as the optimization hyperparameters and additional data mixtures.
We also limit ourselves to the 10B-scale regime and the Llama-3 models, which may limit the generalizability of our findings and recipe.
Another concern is that we are overfitting to the tasks chosen for model development—however, we do not directly train on those datasets and guiding model development with benchmark tasks has become a common practice in pre-trained LM development. We also show that our final recipe and model perform well on an additional evaluation dataset, NoCha.

#### Acknowledgments

We acknowledge Mengzhou Xia, Zexuan Zhong, Samyak Gupta, Dan Friedman, Yihe Dong, Abhishek Panigrah, Adithya Bhaskar, Colin Wang, Carlos Jimenez, and other members of Princeton Language and Intelligence for their helpful feedback and discussion. We also thank Luca Soldaini for providing comments on a draft.
We thank Marzena Karpinska, Tanya Goyal, and Mohit Iyyer for their help with the NoCha evaluation.
Tianyu Gao is supported by an IBM PhD Fellowship.
This work is gratefully supported by an NSF CAREER award (IIS-2239290), a grant from Intel, and Microsoft Azure credits through the “Accelerate Foundation Models Academic Research” Initiative.

## Appendix A Experiment Details

### A.1 Evaluation

**Table 13: The details for our long-context evaluation following HELMET (Yen et al., ).**
| Category | Metrics | Tasks and Datasets |
| --- | --- | --- |
| Recall | SubEM | Given a randomly-generated long JSON file and a key, retrieve the corresponding value (Liu et al., ). |
| RAG | SubEM | Given a question and many retrieved Wikipedia documents (shuffled), answer the question (Liu et al., ). Datasets: *NaturalQuestion* (Kwiatkowski et al., ), *HotpotQA* (Yang et al., ), and *PopQA* (Mallen et al., ). |
| Re-rank | nDCG@10 | Given a query and many retrieved documents (shuffled), re-rank the top-10 documents. Datasets: *MSMARCO* (Bajaj et al., ). |
| ICL | Accuracy | Datasets selected from Bertsch et al. (): *TREC coarse*, *TREC fine* (Hovy et al., ), *NLU* (Liu et al., ), *Banking77* (Casanueva et al., ), and *Clinc-150* (Larson et al., ). |
| QA | GPT-4o score | Given a book, answer the question. Datasets (# tokens): *NarrativeQA* (medium: 73K; max: 518K; Kočiský et al., ). |
| Summ. | GPT-4o score | Summarize a given legal document. Datasets (# tokens): *Multi-LexSum* (medium: 90K; max: 5M; Shen et al., ) |

Table 13 shows all the datasets we used for the long-context evaluation from HELMET (Yen et al., ).
Note that we did not use all the datasets from HELMET for efficiency reasons and we also do not want to overfit to HELMET.
We highlight some of the evaluation protocol improvements that HELMET implemented compared to previous benchmarks here:

- •
Sufficient context lengths and fine-grained control. HELMET can evaluate models at a context length of 128K tokens and beyond.
The evaluation protocol also allows for reporting results at different lengths, giving developers fine-trained controls for different needs of long contexts.
- •
Better synthetic recall tasks. As shown in HELMET, needle-in-a-haystack (Kamradt, ) is mostly saturated because of its simplicity—the model only needs to find a needle in some irrelevant context. We instead use the more challenging JSON KV task, first proposed in Liu et al. () and included in HELMET, where the model is required to find the corresponding value to a given key among a large JSON file.
- •
Using class-balanced demonstrations and abstract labels for ICL. To disentangle models’ ability of learning from demonstrations from their pre-training bias of the task or the dataset label distribution (Pan et al., ), HELMET samples the same number of demonstrations for each class and uses number labels (*1*, *2*, …) instead of natural-language labels (e.g., *location*, *description*, …).
- •
Model-based evaluation for long-context QA and summarization. Instead of using traditional metrics like ROUGE (which has shown to be poorly indicative of the real model performance: Deutsch & Roth, ; Deutsch et al., ; Goyal et al., ; Chang et al., ), HELMET uses model-based evaluations to compare the reference answer and the model output. For QA, HELMET uses GPT-4o to score the model output given the question and the reference answer at a 0-3 scale.
For summarization, HELMET takes a similar approach as Zhang & Bansal (); Gao et al. (): it first uses GPT-4o to decompose the reference summary into atomic claims; then it uses GPT-4o to check whether each reference atomic claim is covered by the model output (recall) and whether each sentence in the model output is covered by the reference summary (precision). Yen et al. () show that the model-based evaluation correlates with human perceptions significantly better than traditional metrics.

### A.2 Data processing

##### Data sources.

We list all the data sources we have explored in our ablations and main experiments here:
the Stack (Kocetkov et al., ),
SlimPajama (Together, ; Soboleva et al., ),
FineWeb (we use the 2023-50 snapshot), FineWeb-Edu (we use a random sample) (Penedo et al., ),
Tulu-v2 (Ivison et al., ),
OpenWebMath (Paster et al., ),
textbooks (Chevalier et al., ),
and Dolma (Soldaini et al., ).
The Books, StackExchange, and ArXiv data are from SlimPajama. The Wikipedia data are from Dolma.

##### Data filtering and packing.

For the short training data and the SFT data, we randomly sample and concatenate the documents or conversations into 64K chunks. The last document for each chunk is truncated. The truncated part is used as the beginning for the next chunk for the short training data but is discarded for the SFT data.
For the long-context training data, we filter out the documents that are shorter than 64K; we do the same for the 512K setting, while making sure that the 64K documents packed to 512K length are distinct from the 512K documents.

##### Final data mixture.

For 512K length, we use a mix of 64K and 512K long data. For the ratio of 64K/512K data, we choose 50%/50% for code and 83%/17%, which are roughly chosen according to the natural availability of very long data, i.e., there are relatively fewer books of length 512K than code repositories. One benefit of retaining 64K-long documents is that we can process these without sequence parallelism and the associated communication overhead.
We use a slightly different long data mixture in our ablations (Table 5) and our main ProLong experiment (Table 9).
For the final model, we mix 3% textbooks into the long-context training data. The textbooks are open-source resources from libretexts.org, collected and made available by Chevalier et al. ().
We pre-process the data by concatenating chapters from the same text books, as well as books from the same subject areas. This results in extremely long sequences which we pack into contexts of either 64K or 512K tokens.
Though we do not have an ablation for adding this data due to limited resources, we believe that it should have a slight positive effect to the final model performance as textbooks are highly educational long-context data.

**Table 14: % Proportion of long documents for the short data components used in Table 6. These statistics are computed after packing and truncation and therefore correspond to the document lengths as seen by the model. We highlight that the proportion of documents beyond 32K is below 1% for ShortMix.**
|  | $>$4K | $>$8K | $>$16K | $>$32K |
| --- | --- | --- | --- | --- |
| FineWeb | 1.4 | 0.3 | 0.1 | 0.0 |
| FineWeb-Edu | 2.8 | 0.8 | 0.2 | 0.0 |
| Wikipedia | 1.6 | 0.4 | 0.0 | 0.0 |
| Tulu-v2 | 0.0 | 0.0 | 0.0 | 0.0 |
| StackExchange | 0.6 | 0.1 | 0.0 | 0.0 |
| ArXiv | 85.7 | 64.0 | 30.3 | 7.6 |
| OpenWebMath | 11.1 | 4.3 | 1.2 | 0.3 |
| ShortMix | 10.9 | 7.2 | 3.2 | 0.8 |
| SlimPajama | 11.3 | 7.4 | 4.9 | 3.2 |
| FineWeb-Edu | 2.8 | 0.8 | 0.2 | 0.0 |
| DCLM-Baseline | 4.9 | 1.7 | 0.4 | 0.1 |

### A.3 Implementation details

##### Technical stack.

We use various open-source packages and tools for the ProLong training and evaluation.
We use PyTorch (Paszke et al., ) and Hugging Face transformers (Wolf et al., ) for the model training.
We use mosaic-streaming (Mosaic ML, ) for loading and mixing the data
and FlashAttention 2 (Dao, ) for efficient attention implementation.
We implement sequence parallelism based on DeepSpeed-Ulysses (Jacobs et al., ) across groups of 8 GPUs on the same node. We only perform distributed attention if it is necessary, i.e., only on sequences of 512K length.
For long-context evaluation, we use HELMET (Yen et al., ) and for short-context evaluation, we use lm-eval-harness (Gao et al., ).

##### Attention and batching.

Since we do document masking in attention (§6),
we use the variable-length attention implementation from FlashAttention 2 (Dao, ) to speed up long-context training:
for sequences that are concatenations of multiple short documents,
instead of computing the full attention with masking,
we instead compute the attention for each individual document.
Since the complexity of attention is quadratic to the sequence length, this improves the training speed.
However, the improvement is negligible in a distributed training setting with FSDP, since GPUs processing short sequence batches have to wait on other GPUs processing long sequences.
We therefore implement a smart batching algorithm:
In our setting, a gradient step usually consists of multiple gradient accumulation steps, where each device processes a smaller minibatch.
We sort all the minibatches per training step by the sum of the squared lengths of documents in the sequence.
This leads to more balanced sequence lengths across the GPUs and effective speedups, as can be seen in Table 15, without affecting the gradient updates or loss during training. However, the efficiency gains are diminished when training with more GPUs, as this reduces the number of gradient accumulation steps.

**Table 15: Throughput per device of our ablation runs from Table 20, when training with 8 Nvidia H100 GPUs with FSDP. Our strategy of reordering minibatches is important for realizing the speed benefits from variable-length attention.**
|  | Throughput |
| --- | --- |
|  | (tokens/s/GPU) |
| 64K full attention | 2770 |
| Variable-length attention | 2780(+0.4%) |
| + Minibatch reordering | 3095(+11.7%) |

##### Token-averaged loss.

We found that in the SFT stage,
the distribution of the training tokens (in SFT, the tokens from the instructions are masked out and the models are only trained on the responses) on each GPU device can be extremely imbalanced, especially when there is synthetic data (most tokens in a synthetic data instance are from the instruction).
Conventional all-reduce loss in distributed training averages over the sequences instead of valid tokens, which skews the optimization and also our control over the domain proportions.
Instead, we change the all-reduce loss to be the average over all valid training tokens.
Bai et al. () implements their SFT loss in a similar way.

### A.4 The ablation setting

For all our ablations, unless specified, we train the base model of Llama-3-8B (instead of Instruct) on a 64K sequence length for 5B tokens, with the same hyperparameters as specified in Table 9.
We choose this context length, as it is the highest power of 2 value for which we can train without sequence parallelism.
By default, we use the same training data as the 64K ProLong setting, except that we remove the textbooks and use the ShortMix proportions in Table 5.
For SFT, we use the same settings as specified in Table 9.

### A.5 Generating synthetic SFT data

We prompt Llama-3-8B-Instruct to generate the synthetic data and
Table 16 shows the prompt we used for generating the synthetic QA data for books.
We also write predefined templates and randomly sample one for each synthetic instance to increase the diversity, and Table 17 provides some examples.

**Table 16: Prompts for generating synthetic QA data.**
| Given the following snippet of a book, ask a relevant question and provide the answer. The question and the answer should follow the following rules: |
| --- |
|  |
| (1) The question should be specific enough that it can only be answered with the snippet. The question should also be interesting and intellectual enough that a curious reader of the book would ask about it. |
| (2) The question and the answer should be comprehensible given just the whole book without highlighting the snippet. With that being said, the question should NOT refer to the snippet directly (e.g., do NOT say things like "Question: given the conversation in the snippet, what ..."). The answer also should not mention "the snippet …" explicitly (assuming that the snippet is never provided), but it can copy the snippet content as a reference when answering the question. |
| (3) The answer should be concise but also should provide references to the book when needed. For example, “Wellington Yueh betrayed the Atreides, as the book mentioned, ’...’". |
|  |
| *** Start of the snippet *** |
|  |
| {sampled snippet} |
|  |
| *** End of the snippet *** |
|  |
| Before generating the question and the answer, first reason about what this snippet is about. In your generation, stick to the following format: |
|  |
| Reasoning: this snippet is about ... |
| Question: ... |
| Answer: ... |

**Table 17: Examples for question prompts and templates used for generating diverse synthetic QA data. We sample one question prompt and one template each time and combine them with the documents and the generated QA pairs to form a synthetic training example.**
| Example question prompts for synthetic QA data |
| --- |
| Given the document, please answer the question. |
| Here is a piece of text; answer the following question based on it. |
| Please answer the question using the provided content. |
| Based on the given passage, respond to the question. |
| Read the snippet and answer the question that follows. |
| Using the provided text, answer the following question. |
| Example templates for combining questions, answers, and contexts for synthetic QA data |
| {prompt}\n\n{documents}\n\nQuestion: {question} |
| {prompt}\n\n==== document starts ====\n{documents}\n==== document ends ====\n\nQuestion: {question} |
| {prompt}\n\n{documents}\n\n{question} |
| {prompt} Question: {question}\n\n{documents} |
| {prompt} {question}\n\n{documents} |
| {prompt}\n\n{question}\n\n{documents} |

## Appendix B More Ablations

### B.1 Position extrapolation

Xiong et al. (); emozilla () show that
changing the RoPE frequency base to a larger value in
continual long-context pre-training or in inference time can improve the long-context performance.
emozilla () suggests that one should scale the frequency base by a factor of $t^{\frac{d}{d-2}}$, where $t$ is the ratio between the target sequence length and the original LM length, and $d$ is the attention head dimension.

We conduct ablation studies, at both 64K (same as our standard ablation setting as specified in §A.4) and 512K (starting from ProLong-64K and training with the 512K data mixture for 5B tokens) sequence lengths, on what frequency bases we should use.
§B.1 and §B.1 show the results.
We first see that using the original 500,000 frequency base from Llama-3 leads to significant performance degradation. While dynamic NTK suggests $4\times 10^{6}$, we find that further scaling it to $8\times 10^{6}$ leads to better performance. Similar, we see that when scaling the 64K model to 512K, while dynamic NTK suggests a $64\times 10^{6}$ frequency base, much larger frequency bases ($128\times 10^{6}$ and $256\times 10^{6}$) lead to better performance.
We use $8\times 10^{6}$ for 64K and $128\times 10^{6}$ for 512K for our final ProLong models.

**Table 18: Ablation study on RoPE frequency base at a maximum training length of 64K. Dynamic NTK (emozilla, ) roughly suggests to use 4m as the frequency base.**
| RoPE Base<br>($\times 10^{6}$) | Long-Context | Short-Context |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Recall | RAG | Re-rank | ICL | QA | Summ. | Avg. | Avg. |  |
| 0.5 | 25.8 | 37.0 | 4.4 | 73.8 | 17.5 | 16.3 | 29.1 | 65.0 |
| 4.0 | 81.3 | 47.8 | 18.2 | 76.5 | 31.8 | 36.3 | 48.7 | 65.3 |
| 8.0 | 96.0 | 54.9 | 29.4 | 73.9 | 35.7 | 37.9 | 54.6 | 65.5 |

### B.2 Document masks

We experiment whether to use document masks in attention in Table 20. Standard training concatenates multiple short documents into a single sequence (in our case, a 64K sequence), uses a special token to separate documents, and performs full attention over the whole sequence. When the document masks are used, we do not allow the attention to cross the document boundaries.
We find that using document masks in continual long-context training leads to both better long-context results and short-context performance.
For all our other ablations and the main experiment, we use document masks.

**Table 20: Impact of using document masks in attention.**
| Attention | Long-Context | Short-Context |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Recall | RAG | Re-rank | ICL | QA | Summ. | Avg. | Avg. |  |
| No doc masks | 97.4 | 53.6 | 20.4 | 76.6 | 37.2 | 36.3 | 53.6 | 64.9 |
| Document masks | 96.0 | 54.9 | 29.4 | 73.9 | 35.7 | 37.9 | 54.6 | 65.5 |

### B.3 Initialization

We use the base model for Llama-3-8B as the initialization for all our ablations to make sure the findings are generalizable and are not confounded by the Llama instruction tuning.
However, for our final ProLong model, we use Llama-3-8B-Instruct as the initialization to achieve the best performance.
We see in Table 21 (using the ablation setting from §A.4) that using Llama-3-8B-Instruct as the initialization achieves slightly better long-context performance and much stronger short-context performance.

**Table 21: Differences of using the base Llama-3-8B model vs. Llama-3-8B-Instruct.**
| Base Model | Long-Context | Short-Context |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Avg. | HellaS. | MMLU | ARC-c | WG | GSM8K | Avg. |  |
| Llama-3-8B-Base | 54.6 | 81.6 | 65.3 | 58.0 | 76.2 | 46.6 | 65.5 |
| Llama-3-8B-Instruct | 55.0 | 80.8 | 66.1 | 58.5 | 75.6 | 57.7 | 67.7 |

### B.4 Instruction-tuning datasets

Initialized from the ProLong base model, we experiment with different public, short-context SFT datasets.
All runs use the same SFT hyperparameters as specified in Table 9.
Table 22 shows that using UltraChat leads to the best overall results.
Note that this does not necessarily mean that UltraChat is the best SFT dataset for all base models or applications.

**Table 22: Ablations on using different short-context SFT datasets. We report the 32K-and-64K-averaged performance except tasks marked with ^†, which are evaluated at 512K for stress testing.**
| SFT Data | Long-Context |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Recall^† | RAG | Re-rank | ICL | QA^† | Summ.^† | Avg. |  |
| UltraChat | 65.7 | 58.1 | 38.5 | 80.3 | 49.7 | 42.1 | 55.7 |
| Tulu v2 | 61.5 | 45.4 | 25.1 | 81.8 | 40.4 | 40.3 | 49.1 |
| ShareGPT | 40.5 | 47.5 | 26.7 | 79.6 | 42.7 | 34.4 | 45.2 |

## Appendix C Full Evaluation Results

We provide all the evaluation results, including specific numbers for each dataset and evaluation length, at [https://shorturl.at/pULbx](https://shorturl.at/pULbx).

## Appendix D Examples

Table 23 shows some long-document QA examples with ProLong, Llama-3.1, and GPT-4o.

**Table 23: Long-document QA examples with ProLong (512K), Llama-3.1-8B-Instruct (128K), and GPT-4o (via the ChatGPT interface). Note that it is very likely that the ChatGPT-version GPT-4o uses a retrieval component for long documents, hence it can process contexts beyond its 128K context window. We see that ProLong can answer questions effectively with a 300K document while both Llama-3.1 and GPT-4o fail. On a shorter document (Llama-3.1 paper), ProLong can understand very technical questions and answer them accurately.**
| Instacart S1 file (300K tokens) |  |
| --- | --- |
| Q: How many shares and at what price did instacard sell their series G redeemable convertible preferred stock in June 2020 and July 2020? | ProLong: Instacart sold an aggregate of 6,757,893 shares of their … at a purchase price of $48.0919 … |
| llama-3.1-8B: The following is a II/5 [collapsed] |  |
| GPT-4o: … The company sold a total of 6,758,000 shares at a price of $48.0919 per share during this period. |  |
| Llama-3.1 paper (94K tokens) |  |
| Q: During training of the llama3 model, how many unexpected interruptions are caused by hardware related reasons? | ProLong: Approximately 78% of unexpected … |
| llama-3.1-8B: According to Table 5, approximately 58.7% … |  |
| GPT-4o: … approximately 78% of unexpected interruptions … |  |
| Q: Is Mark Zuckerberg a contributor to the paper? | ProLong: There is no mention of Mark Zuckerberg as a contributor to the paper. |
| llama-3.1-8B: Yes, Mark Zuckerberg is listed as a contributor … |  |
| GPT-4o: Mark Zuckerberg is acknowledged in the paper for …, but he is not listed as a core contributor … |  |
| Q: What ratio did the llama 3 team use to mix synthetic long-context data with short data in SFT? | ProLong: The llama 3 team used a ratio of 0.1% synthetic … |
| llama-3.1-8B: 0.10% |  |
| GPT-4o: The Llama 3 team used a mix ratio of 30% synthetic … |  |
| Q: What benchmark did they evaluate on for ”instruction-following” ability? Provide references too. | ProLong: They evaluated on IFEval (Zhou et al., 2023) for instruction-following ability. Reference: Jeffrey Zhou, … Instruction-following evaluation for large language models. arXiv preprint … |
| llama-3.1-8B: The benchmark they evaluated on for ”instruction-following” ability is IFEval. Reference: Section 5.2.2, ”General Knowledge and Instruction-Following Benchmarks” |  |
| GPT-4o: The benchmark they used to evaluate ”instruction- following” ability is IFEval. This is mentioned in the context of assessing how Direct Preference Optimization (DPO) … |  |

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: * How does this recipe scale? Would this recipe work for larger models? 
* Since the model is trained for 512k context length, it would be very helpful to understand the performance of the model at 512k context length tasks. Do any of the evaluations have such long context tasks? 
* The authors mention that the performance on short context tasks is important but don't report performance on these tasks for their final ProLong model. How does ProLong perform on short context tasks?
soundness: 3
strengths: * Careful examination of individual changes (data mix, sft etc.)
* Final recipe does provide a model that outperforms some of the existing models
summary: The main contributions of the paper are related to evaluations, data mix, and training recipe. 

Evaluation related contribution
* Using HELMET (an existing benchmark) instead of perplexity and needle in the haystack for evaluating long context expansion.
* The authors suggest evaluating LLMs for long context after SFT instead of after pretraining.

Data mix related contributions

* Carefully designed data mixtures like high quality code and books are important. 
* Additionally, the authors ablate the ratio of short context and long context data in the data mix to show that a very high fraction of long context data is detrimental to both long context and short context tasks.

Training recipe: 
* Training on more tokens improves performance of long context tasks
* Increasing pretraining context length beyond target context length can also help long context tasks.
* They also show that synthetic data during SFT can be harmful to long context benchmarks

They combine these improvements to develop a new model called ProLong which achieves better results than prior models on HELMET and NoCha (another long context benchmark).
weaknesses: The paper contributions are somewhat limited. The main contributions are around datamixing and training recipe. The authors make some interesting claims but unfortunately don't investigate these claims further. 

* The results on perplexity are a bit surprising. The authors show that increasing long context data results in perplexity improvements but long context tasks performance is lower. It would be really interesting to understand this further. Is the quality degraded due to the data mix? Or is the target eval task not correlated to the PG19 books dataset?
* The authors show that using synthetic data during SFT can lead to poor performance. It would have been helpful if the authors had verified that the generated data is correct prior to dismissing it. 
* Additionally, using a large model (say Llama 70B) for synthetic data generation could have been more interesting. Perhaps, this larger model would have generated data that can be used for SFT.
* The trends between SFT evaluation and pretrained model evaluation are not clear. In Figure 2, many of the benchmarks show similar trends before and after SFT, so it's hard to conclude that it is necessary to measure performance after SFT as the authors claim. 

The paper would also be stronger if the authors show that their methods generalize beyond just Llama 3 8b. Would the same set of techniques also expand context lengths effectively for other open source models?

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: - Why training on all long context data will harm the downstream long-context tasks? Is it because (1) the proposed ShortMix data actually contains lots of QA data from Fineweb-Edu and Stack Exchange which can improve instruction following (2) the curated long context data (code repos+books) does not have QA data, and training purely on it will harm the instruction following ability?
In Line 351, why use the llama-3-8B-instruct model to synthesize long context SFT data, instead of more powerful models such as lama-3-70B-instruct? I would assume an 8B model cannot accomplish these difficult tasks and can provide low-quality long-context SFT data.

- Could you also evaluate ProLong on some popular synthetic long-context benchmarks such as RULER? In this way, we can have a better understanding of the limitations of ProLong.

- The stage 2 training also doubled the batch size from 4M to 8M. Do you have an ablation study for it? Batch size ramp-up is believed to provide more accurate optimization trajectory and boost model convergences, therefore the performance gains from 512K length training may actually be from the batch size ramp-up instead of longer context.

- Could you also evaluate the final ProLong model on short-context tasks, and add comparisons with a llama3-8B-instruct model SFTed with UltraChat?
soundness: 3
strengths: - The model trained with the final recipe (ProLong) shows strong performance on long context benchmarks.

- The proposed data mixture can be a valuable contribution for the community.
summary: The paper proposes a continual training recipe, including data curation and multi-stage context extension, that successfully equips the llama-3-8B-instruct model with long context ability up to 512K tokens.
weaknesses: - Missing crucial baselines: The author should at least include a very baseline that applies the recipe of Fu et al. (2024) on the llama-3-8B model with 40B tokens to provide a fair comparison between the ProLong recipe and the Fu et al. (2024) recipe on both short and long context tasks. The Table 10 and 12 is not very informative because of lacking both fair comparisons and the short-context results.

- Misleading evaluation: As in L174, the short-context performance is evaluated before SFT, while the long context performance is after SFT (As in L149). However, what people really care about is if the final model after SFT can have both good short and long-context performances.

- The paper claims counterintuitive conclusions but lacks in-depth analyses explaining why they happen. The author claims that training on longer than evaluation context can bring performance benefits (in L331), but it can be simply due to the fact of training on 4B more tokens, which can improve the instruction following ability because of the short data mixture. The author claims that the long synthetic SFT data does not improve performance but it can be simply because of using a weak 8B model for data generation. See more details in the Question section.

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
questions: See weaknesses above.
soundness: 2
strengths: The paper gives a thorough investigation of data engineering for long context training across various aspects, covering long-short performance balance, SFT, and evaluation. This paper can potentially serve as a handbook for long context training. The paper is well-organized and easy to understand.
summary: This paper investigates the optimal recipe of training a long-context LLM from the data engineering perspective. The authors investigate the effects of data source, mixing proportion of long and short data, etc, and evaluate the trained models on realistic tasks beyond perplexity.
weaknesses: - The evaluation beyond ppl is good and there are some previous studies [1] proprose ppl is not proper for evaluating long-context performance. However, it is kind of confused if the evaluation can serve as a "contribution" of this paper? It seems the paper just adopt the HELMET benchmark (I don't think the decision is a contribution as most long-context works have adopted some realistic benchmarks like LongBench [2] and InfiniteBench [3]).

- The paper proposed short-context instruction tuning is sufficient for achieving good performance. The study is conducted by mixing partial long synthetic instruction data. The authors found that the incorporation of long synthetic data would degrade the long-context performance.  As our understanding in short-context scenario, data quality is (most) essential for SFT [4]. However, the synthetic long-context data (generated by LLaMA-3-8B) is probably of low quality. So I think this study only demonstrates that long synthetic data is not good enough for improving long-context performance. BTW, have you compared the long-context performance before and after SFT? LLaMA 3.1 technical report [5] proposed the long-context performance would degenerate with short instruction data only after SFT. I wonder if this point holds here.

- I believe summarization is a vital and basic long-context ability, however, there is a significant performance drop in the summarization task after ProLong's training (compared to naive LLaMA-3.1). Is this due to the lack of some long data types/sources in training data (books and code only)? The information in long books and codes may be distributed without summarizing. I feel like arxiv/wikipedia  (which you only use in shortmix) are good data sources for summarization, as they naively contain an abstract section. The authors can consider including more data sources with ablations (as that in Table 4).

- Would you consider reporting results on more benchmarks like RULER [6] and InfiniteBench [3]? HELMET is a good but new benchmark while some experiences (e.g., the degeneration of long-context performance when using short instruction data only.) are acquired from evaluation on some older benchmark. More results would help us align the findings.


[1] Can Perplexity Reflect Large Language Model’s Ability in Long Text Understanding?

[2] LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding

[3] ∞Bench: Extending Long Context Evaluation Beyond 100K Tokens

[4] LIMA: Less Is More for Alignment

[5] The Llama 3 Herd of Models

[6] RULER: What's the Real Context Size of Your Long-Context Language Models?

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 4
questions: See Weakness above.
soundness: 3
strengths: 1. The paper is very well structured and written, and is easy to follow.
2. The paper thoroughly explored important components and provide useful observations for long-context language models development, including training recipe and evaluation protocol. From the evaluation perspective, the authors underscore the importance of using diverse and realistic long-context tasks rather than relying on conventional perplexity-based and needle-in-a-haystack-like benchmarks. The preservation of short-context performance is also highlighted. From the training perspective, the authors conducted experiments and found empirically effective choices of data mixture, data ratio, and data quality.
3. The ultimately trained model, ProLong, achieves strong long context performance at 10B parameter level.
4. The assets associated with this paper could serve as useful resources for the research community and are open-sourced.
summary: This paper study the problem of how to effectively adapting a short-context language models to be long-context. The paper first starts with identifying deficiency in widely used perplexity and Needle-In-A-Heystack test and establishing a reliable evaluation protocol for long-context LLM by adopting HELMET, which covers diverse range of realistic long-context applications. 
The authors also advocate for (1) checking model performance after supervised finetuning (2) ensuring that model performance on short-context tasks is not compromised after long context training. For long context data curation, the authors provide best practices for (1) data mixture in long-context data (2) quantity ratio between long and short context data (3) choosing high-quality short-context data. 
Next, the paper discuss the impact of scaling training tokens and context length on both long and short task performance. Finally, the recipe ends with exploration on the choice of SFT data, pointing out that short-context instruction data can yield strong long-context result. The resulting long-context model ProLong as well as associated code and data are open-sourced.
weaknesses: 1. The choice of long-context evaluation benchmark(HELMET) seems arbitrary. The authors should discuss the rationale behind selecting HELMET instead of other long-context benchmarks, such as LongBench, Infinite Bench, etc.
2. The training-free length extrapolation methods compared in the paper is limited. The experiments would be more solid if more advanced training-free extension methods are incorporated, e.g., SelfExtend and ChunkLLaMa.
3. There exists other community-tuned long-context LLMs initialized from LLaMa-3-8B-Instruct, e.g., gradientai/Llama-3-8B-Instruct-Gradient-1048k. As they claim their training process only consume 1.4B tokens, it would be more convincing to include such models as strong baselines.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 4
questions: The paper's contributions are not clearly defined in terms of their nature, like technical report？ If the paper aims to contribute as a methodological study on data engineering, the novelty seems insufficient, as it mainly provides data mixture ratios without in-depth justification or exploration. On the other hand, if it aims to be an experimental study, the lack of diverse experiments and limited comparisons with other training approaches weakens its claims of novelty.
soundness: 1
strengths: 1.  The study highlights that training on a mix of both long and short-context data sources is key to achieving good performance across tasks. The authors show that using long-context data like code repositories and books, balanced with high-quality short-context data, preserves both long and short-context abilities of the model.
2.  The extensive ablation experiments provide meaningful insights into effective training practices. For instance, the paper reveals that training only on long-context data negatively impacts overall performance and that including high-quality short-context data is essential for optimal results.
3. The resulting ProLong-8B model achieves strong performance on long-context tasks compared to similarly sized models
summary: The paper addresses the challenges involved in training language models (LMs) to make effective use of long-context information. It proposes a systematic methodology that goes beyond traditional evaluation metrics, such as perplexity, to train models capable of handling long-context inputs more effectively. The authors introduce ProLong-8B, a model initialized from Llama-3 and trained on 40 billion tokens, demonstrating state-of-the-art performance for long-context tasks.
weaknesses: 1. **Misalignment Between Title and Content**: A significant portion of the paper focuses on data engineering rather than providing comprehensive training methodologies for long-context language models, which does not fully align with the promise of the title. To better align with the title, the paper could have included more discussions on training techniques, such as **pose: efficient context window extension of LLMs via positional skip-wise training** and **CLEX: Continuous Length Extrapolation for Large Language Models**. Additionally, incorporating experiments involving different rotary position encoding methods and synthetic data, such as **LongSkywork: A Training Recipe for Efficiently Extending Context Length in Large Language Models**, would provide a more thorough exploration of how to train long-context LLMs.

2. **Benchmark Justification**: In Section 2, the necessity of introducing a new benchmark is not clearly justified, especially considering that many of the tasks overlap with existing benchmarks like **RULER**. Providing more rationale for the benchmark's uniqueness？

3. **Lack of Explanation**: In Section 3.2, the authors note that previous techniques (using long text) deteriorate short-context performance, but they do not provide a satisfactory explanation for this phenomenon. Without a clear rationale, the motivation for mixing in short-context training data remains weak.

4. **Insufficient Explanation of Data Mixture Ratios**: The paper proposes two mixture ratios—one for combining long-context data sources (books and repositories) and another for mixing long and short-context data. While the experiments show the effectiveness of these ratios, there is no clear explanation of why these particular proportions work. 

5. **Lack of Comparison with Short-Context Performance Task**: The final experiments lack a comparison with short-context performance tasks. Including such comparisons would provide a more holistic evaluation of the model's performance.

Meta-review: The paper presents ProLong-8B, a language model fine-tuned to handle extended context lengths up to 512K tokens. Key findings include the effectiveness of combining code repositories and books with high-quality short data, the benefits of training with sequence lengths exceeding evaluation lengths, and the strong performance achieved using only short instruction datasets. However, there are major concerns regarding the experiments and their results i.e. how these results can be interpreted as valuable insights for the community, the intuitions behind some observations are not clearly justified, and comparison with relevant baselines. I encourage the authors to address these issues, focusing on how their results can be translated as more generalisable insights for the community.

Decision: Reject