How to Train Long-Context Language Models (Effectively)

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., 2023; Peng et al., 2024), Fu et al. (2024) find that the above methods cannot even perform the simple needle-in-a-haystack (NIAH; Kamradt, 2024) 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., 2024) and Jamba (Lenz et al., 2024), 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., 2024) 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., 2024). ProLong has a maximum context length of 512K tokens, making it one of the longest-context LMs available.Throughout the paper, we use binary prefixes K$=2^{10}$, M=$2^{20}$, and B=$2^{30}$.

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, 2024) and RULER (Hsieh et al., 2024) 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.

We first make the decision to use HELMET (Yen et al., 2024b) 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. (2024b) 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. (2024) 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.

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.

Besides synthetic recall tasks, many previous works rely on perplexity (PPL) for evaluating long-context extensions of LMs (Chen et al., 2023; Fu et al., 2024; Lu et al., 2024), which is commonly measured on the PG19 books dataset (Rae et al., 2020). 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 Evaluate after supervised fine-tuning

Supervised fine-tuning (SFT; Ouyang et al., 2022) 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. (2024). The mix is based on SlimPajama (Soboleva et al., 2023) 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., 2023).

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.

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.

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., 2023): HellaSwag (Zellers et al., 2019), MMLU (Hendrycks et al., 2021), ARC-challenge (Clark et al., 2018), WinoGrande (Sakaguchi et al., 2021), and GSM8K (Cobbe et al., 2021). 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.

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., 2024) 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.

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., 2023; Wettig et al., 2024; Li et al., 2024a) 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.

We analyze the quantity of long data in SlimPajama (SP; Soboleva et al., 2023). 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.

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

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.

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., 2024) or adds some short training data (Yen et al., 2024a; Fu et al., 2024). 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.

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., 2023), FineWeb-Edu (Penedo et al., 2024), DCLM-Baseline (Li et al., 2024a), and our own ProLong ShortMix. Our ShortMix is inspired by the “stage 2 training” in MiniCPM (Hu et al., 2024a) and Dolma-1.7 (Soldaini et al., 2024), which use more knowledge-intensive, downstream-related data at the end of pre-training. Table 5 shows the composition of our ShortMix. 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 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.

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.

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.

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., 2022; Xiao et al., 2024b; a; Yen et al., 2024a; Chen et al., 2023). 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.While 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.

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., 2022) 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., 2023) conversations have 1.2K tokens on average and 4.1K tokens maximum. To bridge this gap, several works (Xiong et al., 2023; Dubey et al., 2024; Xiong et al., 2024) have proposed to generate long instruction data synthetically.

We consider three popular SFT datasets—UltraChat (Ding et al., 2023), Tulu-v2 (Ivison et al., 2023), ShareGPThttps://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., 2021). 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.

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.

Why do our conclusions about synthetic data differ from previous work? We offer the following hypotheses: (1) Previous work like Xiong et al. (2024); Bai et al. (2024a) 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., 2024)—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.

The ProLong Model: Recipe and Results

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., 2023). 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., 2024), 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.

We find that changing the RoPE (Su et al., 2021) frequency base to achieve position extrapolation (Xiong et al., 2023; emozilla, 2023) 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.

Ding et al. (2024a) show that masking out attention across document boundaries improve model performance and this was also used during Llama-3 pre-training (Dubey et al., 2024). 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.

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.

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 MegaBeamhttps://huggingface.co/aws-prototyping/MegaBeam-Mistral-7B-512k., Llama-3.1 (Dubey et al., 2024), Qwen2 (Yang et al., 2024a), Phi-3 (Abdin et al., 2024), Mistral-Nemohttps://huggingface.co/mistralai/Mistral-Nemo-Instruct-2407., Jamba-1.5 (Lenz et al., 2024), Claude-3.5-Sonnet (Anthropic, 2024), Gemini-1.5 (Reid et al., 2024), and GPT-4o (Achiam et al., 2023).

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.

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 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.

Besides HELMET, we also evaluate our models on NoCha (Karpinska et al., 2024)—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.

Related Work

Many works explore extending the LM context windows with minimal training, either by position extrapolation (Chen et al., 2023; Peng et al., 2024; Chen et al., 2024; Ding et al., 2024b; Liu et al., 2024a; Zhang et al., 2024b; Zhu et al., 2024; Zhao et al., 2024; Wu et al., 2024; Hu et al., 2024b) or manipulating the attention patterns (Chen et al., 2024; Xiao et al., 2024b; a; Bertsch et al., 2023; Jin et al., 2024). Yoshida et al. (2020); Choromanski et al. (2021); Chevalier et al. (2023) instead explore the idea of compressing the long contexts into shorter forms. However, Fu et al. (2024); Lu et al. (2024) 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., 2024) and Jamba (Lieber et al., 2024) 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. (2024) 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. (2023); Dubey et al. (2024); Lieber et al. (2024); Xiong et al. (2024); An et al. (2024b); Bai et al. (2024a) 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.

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

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

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.

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.

References

Appendix A Experiment Details

Table 13 shows all the datasets we used for the long-context evaluation from HELMET (Yen et al., 2024b). 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, 2024) 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. (2024b) 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., 2023), 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, 2021; Deutsch et al., 2022; Goyal et al., 2023; Chang et al., 2024), 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 (2021); Gao et al. (2023): 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. (2024b) show that the model-based evaluation correlates with human perceptions significantly better than traditional metrics.

A.2 Data processing

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

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.

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. (2024). 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.

A.3 Implementation details

We use various open-source packages and tools for the ProLong training and evaluation. We use PyTorch (Paszke et al., 2019) and Hugging Face transformers (Wolf et al., 2020) for the model training. We use mosaic-streaming (Mosaic ML, 2022) for loading and mixing the data and FlashAttention 2 (Dao, 2024) for efficient attention implementation. We implement sequence parallelism based on DeepSpeed-Ulysses (Jacobs et al., 2023) 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., 2024b) and for short-context evaluation, we use lm-eval-harness (Gao et al., 2021).

Since we do document masking in attention (§6), we use the variable-length attention implementation from FlashAttention 2 (Dao, 2024) 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.

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. (2024a) 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.

Appendix B More Ablations

Xiong et al. (2023); emozilla (2023) 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 (2023) 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.

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.

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.

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

Appendix D Examples

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