One-sentence Summary

NVIDIA researchers introduce ChatQA 2, an Llama 3.0-based model with a $128\text{K}$128K context window, and propose a continued training recipe alongside three-stage instruction tuning that extends the context from $8\text{K}$8K to $128\text{K}$128K tokens, achieving superior performance over GPT-4-Turbo-2024-04-09 and Qwen2-72B-Instruct on ultra-long tasks and RAG benchmarks, with RAG outperforming direct long-context solutions when retrieving more chunks.

Key Contributions

• Introduces Llama3-ChatQA-2-70B, an open-source long-context model with a 128K context window, along with a reproducible continued training recipe that extends Llama 3 70B from 8K to 128K tokens and a three-stage instruction tuning process. The model reaches GPT-4-Turbo-2024-04-09 level accuracy, outperforming GPT-4-Turbo-2024-04-09, Qwen2-72B-Instruct, and Llama3.1-70B-Instruct on ultra-long tasks beyond 100K tokens and on the ChatRAG benchmark within a 4K context window.

• Conducts an extensive comparison between direct long-context and RAG solutions using the same state-of-the-art long-context LLMs, showing that RAG accuracy consistently improves as the number of retrieved top-k chunks grows. With a sufficiently large top-k, RAG outperforms direct long-context on both 32K and 128K benchmarks using the same underlying models.

• Shows that RAG with a long-context retriever remains robust to variations in chunk size and that accuracy improves as the total number of retrieved tokens increases. Releases model weights, training data, and evaluation setup for the research community.

Introduction

The open LLM community has produced capable open-access models such as Llama-3-70B-Instruct and Qwen2-72B-Instruct, yet these still lag behind frontier proprietary systems like GPT-4-Turbo across many domains. Extending context windows has become a central trend, with proprietary models supporting 128K to 10M tokens and open models reaching up to 200K, but prior open attempts face key limitations: training data and reproduction recipes are often unavailable, making replication difficult, and evaluations rely mostly on synthetic tests like Needle in a Haystack rather than real-world tasks. Long-context processing and retrieval-augmented generation (RAG) are frequently framed as competing approaches, but they complement each other: RAG provides efficiency for query-based tasks over billions of tokens, while long-context models handle whole-document summarization. Earlier work by Xu et al. extended Llama2 to only 16K and 32K tokens, reaching just GPT-3.5-turbo-level capability.

The authors present ChatQA 2, which pushes open long-context LLMs to GPT-4-Turbo-level capability with a 128K context window and pairs this with a state-of-the-art long-context retriever for RAG. Their main technical contribution is a two-step recipe: first extending Llama3-70B's context from 8K to 128K via continued pretraining on SlimPajama with upsampled long sequences, then applying a three-stage instruction tuning process that separately enhances instruction-following, RAG, and long-context understanding. The resulting Llama3-ChatQA-2-70B outperforms GPT-4-Turbo-2024-04-09, Qwen2-72B-Instruct, and Llama3.1-70B-Instruct on long-context tasks beyond 100K tokens and on RAG benchmarks within 4K tokens. The authors also show that their RAG accuracy remains robust to chunk size variations and continues to improve with more retrieved tokens, while demonstrating that RAG can still beat long-context-only models given sufficient top-k chunks. Training data and the full reproduction recipe are publicly released.

Dataset

Dataset Description

The authors construct a suite of datasets to extend the Llama3 base model's context window from 8K to 128K tokens, followed by a three-stage instruction-tuning process. The data is divided into a pretraining corpus and a long-context supervised fine-tuning (SFT) dataset.

• Long Context Pretraining Corpus

  • Source: Curated from the SlimPajama corpus.
  • Composition: Long-context documents are upsampled with a hyperparameter set to 0.1 to create a corpus of 10 billion tokens, formatted into sequences of 128K tokens.
  • Processing: To help the model adapt to longer inputs, the authors separate different documents with the special character <s> instead of the reserved <BOS> and <EOS> tokens, hypothesizing that the latter signal the model to ignore previous context.
  • Usage: Used for a continued pretraining phase with a batch size of 32 (4 million tokens) and a learning rate of 3e-5 for 2000 steps, totaling 8 billion tokens.

• Long Context Instruction-Tuning Dataset (Long SFT)

  • Composition: A blend of data from two categories based on sequence length.
    • For sequences under 32K: Sourced from existing long-context datasets including LongAlpaca12k, GPT-4 samples from Open Orca, and Long Data Collections.
    • For sequences from 32K to 128K: Synthetic data generated from the NarrativeQA dataset, which contains human-generated summaries, questions, and answers based on long source web pages.
  • Processing: To create longer contexts, the authors insert a summary into its corresponding long web page document at a random location, ensuring the sentence structure is not disrupted. This preserves the grounding of the original question-answer pairs.
  • Usage: This long SFT dataset is blended with the short SFT dataset from the first two instruction-tuning stages for training. The authors intentionally exclude NarrativeQA from all evaluation benchmarks to prevent data contamination.

• Evaluation Benchmarks

  • Ultra-Long Context (over 100K): Four tasks from InfiniteBench are used: En.Sum (ROUGE-L-Sum score), En.QA (F1 score), En.MC (Exact Match score), and En.Dia (Exact Match score).
  • Long Context (within 32K): Six datasets are used from Xu et al. (2024), excluding NarrativeQA. These include QMSum (geometric mean of ROUGE scores), Qasper and MultiFieldQA-en (F1 scores), QuALITY (Exact Match score), and HotpotQA and MuSiQue (F1 scores).
  • Short Context (within 4K): ChatRAG Bench is used, excluding the HDial dataset. For Doc2Dial, QuAC, and QReCC, documents are segmented into roughly 300-word chunks with the top 5 retrieved. For TopiOCQA and INSCIT, the top-20 chunks are retrieved. The remaining datasets are CoQA, DoQA, ConvFinQA, and SQA. The average F1 score across these datasets is reported.

Method

The authors present a systematic approach to adapting Llama3 for long context understanding and retrieval-augmented generation. The method proceeds in three stages: extending the pretrained context window from 8K to 128K, instruction-tuning with long-context supervision, and pairing the resulting model with a long-context retriever.

Extending Context Window to 128K

To extend the context window, the authors prepare a long-context pretraining corpus derived from Slimpajama. They upsample long-context documents with a hyperparameter set to $0.1$0.1 to produce 10 billion tokens with a sequence length of 128k. Since Llama3 is pretrained with a much higher RoPE base frequency of $500{,}000$500,000 compared to Llama2, the authors increase the RoPE base frequency to $150$150M accordingly to accommodate longer sequences. Training uses a batch size of 32, giving 4 million tokens per batch, a learning rate of $3 \times 10^{-5}$3×10−5, and 2000 steps totaling 8 billion tokens.

An interesting finding is that separating different documents using special characters, such as "<s>", proves more effective than using the reserved beginning and ending tokens <BOS> and <EOS>. The authors hypothesize that these tokens signal the model to ignore previous chunks of text after pretraining, which is counterproductive for adapting the LLM to longer context inputs.

Instruction-Tuning with Long Context Data

To enhance both long-context understanding capability and RAG performance, the authors implement three stages of instruction-tuning. The first two stages follow ChatQA 1.5, where the model is initially trained on a high-quality instruction-following dataset, and then trained on a blend of conversational QA data with provided context. These stages involve relatively short contexts with a maximum sequence length of only 4K tokens.

To strengthen the model's ability to handle very long sequences up to 128K tokens, the authors collect a long SFT dataset covering two categories. For sequences less than 32K, they leverage existing long-context datasets, including LongAlpaca12k, GPT-4 samples from Open Orca, and Long Data Collections. For sequences between 32K and 128K, they rely on synthetic datasets built from NarrativeQA, which contains summary paragraphs, questions, answers, and source long web pages. The summaries are human-generated based on the source web pages, while question-answer pairs are human-annotated using the summaries. To extend the context length, a summary is inserted into the corresponding long web page document at a random location while ensuring the sentence structure remains intact. This approach maintains the grounding of the question-answer pairs within the augmented long documents. Since NarrativeQA is used for synthetic data generation, it is intentionally excluded from evaluation benchmarks to avoid potential data contamination. Both the full long SFT dataset and the short SFT dataset from the first two stages are blended for training, with a learning rate of $3 \times 10^{-5}$3×10−5 and a batch size of 32.

Long Context Retriever Meets Long Context LLM

The authors identify two issues in current RAG pipelines. First, top-k chunk-wise retrieval introduces non-negligible fragmentation of context for generating accurate answers, since previous state-of-the-art dense-embedding based retrievers only support 512 tokens. Second, small top-k values, such as 5 or 10, usually lead to relatively low recall, while much larger k values, such as 100, can lead to worse generation as prior LLMs cannot effectively utilize many chunked contexts. To address these issues, the authors propose using a recent long-context retriever that supports thousands of tokens. Specifically, they employ the E5-mistral embedding model as the retriever. The input to this model is a chunk of text represented by a sequence of subword tokens, and the output is an embedding vector. Long documents or corpus are chunked and embedded into a set of dense embedding vectors, which are indexed and retrieved via k-nearest-neighbor search at inference time, where the embedding of the prompt or question serves as the query.

Experiment

The evaluation compares Llama3-ChatQA-2 models against SOTA long-context LLMs (including GPT-4-Turbo, Qwen2, and Llama-3.1) across three benchmark categories: ultra-long contexts beyond 100K tokens, long contexts within 32K, and short contexts within 4K, with RAG baselines using the E5-mistral retriever. The proposed model achieves 100% accuracy on the Needle In A Haystack test, outperforms prior SOTA on ultra-long InfiniteBench tasks (particularly QA, with weaker summarization due to limited training data), and leads on short-context ChatRAG Bench despite slight gaps on 32K tasks attributed to reduced pretraining scale. Ablation studies show that a three-stage training recipe beats all-in-one training, and that RAG can outperform full long-context solutions when more retrieved chunks are used, with accuracy improving as token count increases.

The model achieves the top average scores on both ultra-long (>100K) and short (4K) context tasks among the compared models, while lagging slightly on mid-length (32K) tasks. It also stands out for open-sourcing its training data and reproduction recipe, which are not available for other open-access 128K models. Highest average score on real-world ultra-long tasks beyond 100K tokens among all compared models. Leads on short-context (4K) RAG tasks, outperforming both proprietary and open-access counterparts. Slightly lower performance on 32K tasks relative to proprietary models and some open-access models. Open-sources training data and recipe for extending 8K-context base models to 128K, unlike other open-access models.

The model achieves a higher average score than several strong baselines on the InfiniteBench tasks, with particularly strong performance on QA and multiple choice tasks but lower scores on summarization. Using retrieval-augmented generation (RAG) with top-5 chunks yields worse results than processing the full long context directly. Outperforms GPT-4-Turbo, Claude 2, and other large long-context models on average. Shows excellent performance on QA and multiple choice tasks, while summarization scores are relatively low, likely due to limited summarization data in training. RAG-based evaluation scores are lower than those from direct long-context processing.

On long context benchmarks within 32K tokens, GPT-4-Turbo leads all models, while the proposed model performs competitively but falls slightly behind two other 70B models. A default RAG setup with top-5 chunks of 1200 tokens consistently yields lower average scores than direct long-context processing across all evaluated models. GPT-4-Turbo achieves the highest average score, with the proposed model trailing Qwen2-72B-Instruct and Llama3.1-70B-Instruct but ahead of Llama-3-70B-Instruct-Gradient-262k. Using a default RAG configuration with top-5 chunks of 1200 tokens results in lower average scores than the direct long-context approach for every model.

The proposed model achieves strong average performance on CHATRAG BENCH short-context tasks, surpassing several large-context models while slightly trailing a dedicated short-context model. This indicates a trade-off when extending context windows to longer lengths. Outperforms GPT-4-Turbo, Qwen2-72B-Instruct, and Llama3.1-70B-Instruct on average score despite those models supporting 128K context. Slightly behind Llama3-ChatQA-1.5-70B, which only handles 4K context, suggesting that extending context length comes with modest short-task performance costs.

The comparison shows that for tasks beyond 100K tokens, retrieval-augmented generation (RAG) achieves slightly higher average accuracy than direct long-context evaluation for both models tested. This indicates that RAG remains competitive or better even at very long context lengths, while also offering lower inference cost. RAG with top-k retrieval outperforms direct long-context evaluation for Llama3-ChatQA-2-70B. The same advantage holds for Qwen2-72B-Instruct, where RAG also edges out the full long-context approach.

The evaluation spans ultra-long (>100K), mid-length (32K), and short (4K) context benchmarks, where the proposed model ranks best on ultra-long and short-context RAG tasks, falls slightly behind on 32K tasks, and uniquely open-sources its training data and reproduction recipe. It excels at QA and multiple choice but shows weaker summarization performance. Direct long-context processing generally beats RAG with a default top-5 chunk setup within 128K contexts, yet for tasks beyond 100K tokens, RAG edges out direct evaluation while also reducing inference cost. Extending context length carries a modest trade-off, as the model trails a dedicated short-context model on the ChatRAG Bench despite surpassing several long-context counterparts.