Breaking the Stage Barrier: A Novel Single-Stage Approach to Long Context Extension for Large Language Models

TL;DR

This paper introduces HARPE, a single-stage continual pretraining method for large language models that enhances their ability to handle long contexts efficiently, outperforming multi-stage approaches.

Contribution

The paper proposes a novel single-stage pretraining technique, HARPE, which simplifies long context training for LLMs by using head-adaptive rotary position encoding.

Findings

HARPE outperforms multi-stage methods on 4 language modeling benchmarks.

HARPE effectively models long contexts with a single training stage.

The approach simplifies training while maintaining or improving performance.

Abstract

Recently, Large language models (LLMs) have revolutionized Natural Language Processing (NLP). Pretrained LLMs, due to limited training context size, struggle with handling long token sequences, limiting their performance on various downstream tasks. Current solutions toward long context modeling often employ multi-stage continual pertaining, which progressively increases the effective context length through several continual pretraining stages. However, those approaches require extensive manual tuning and human expertise. In this paper, we introduce a novel single-stage continual pretraining method, Head-Adaptive Rotary Position Encoding (HARPE), to equip LLMs with long context modeling capabilities while simplifying the training process. Our HARPE leverages different Rotary Position Encoding (RoPE) base frequency values across different attention heads and directly trains LLMs on the target context length. Extensive experiments on 4 language modeling benchmarks, including the latest RULER benchmark, demonstrate that HARPE excels in understanding and integrating long-context tasks with single-stage training, matching and even outperforming existing multi-stage methods. Our results highlight that HARPE successfully breaks the stage barrier for training LLMs with long context modeling capabilities.