Updates

🔥🔥 SkyLadder is accepted to NeurIPS 2025! See you in San Diego!

Introduction

Recent advancements in LLM pretraining have featured ever-expanding context windows to process longer sequences. However, our pilot study reveals that models pretrained with shorter context windows consistently outperform their long-context counterparts under a fixed token budget. This finding motivates us to explore an optimal context window scheduling strategy to better balance long-context capability with pretraining efficiency. To this end, we propose SkyLadder, a simple yet effective approach that implements a short-to-long context window transition. SkyLadder preserves strong standard benchmark performance, while matching or exceeding baseline results on long-context tasks. Through extensive experiments, we pre-train 1B-parameter models (up to 32K context) and 3B-parameter models (8K context) on 100B tokens, demonstrating that SkyLadder yields consistent gains of up to 3.7% on common benchmarks, while achieving up to 22% faster training speeds compared to baselines

Quick Start

This project is based on the TinyLlama project. It has been adapted to support pretraining with context window scheduling, intra-document masking, etc.

Installation

If you already an environment built for TinyLlama, you can directly use it. Otherwise, please use the following commands to build a new environment. Here, we expect a CUDA version of 11.8

conda create -n ladder-pretrain python=3.8
conda activate ladder-pretrain
# install the latest compatible version of torch and xformers, this should install torch 2.4.1
pip install ninja
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu118
# install flash attention
git clone --branch v2.3.3 --depth 1 https://github.com/Dao-AILab/flash-attention.git
cd flash-attention
python setup.py install
cd csrc/rotary && pip install .
cd ../layer_norm && pip install .
cd ../xentropy && pip install .
cd ../.. && rm -rf flash-attention
# install other dependencies 
pip install -r requirements.txt

If you wish to use Docker, you could follow the instruction here for the docker image name siviltaramqian/tinyllama:latest

Data preparation

The data preparation process is the same as the original tinyllama project. First make sure that your data is in jsonl format in one directory of the following structure:

TEXT_DIR
├── cc
│   ├── train
│   │   ├── 0.jsonl
│   │   ├── 1.jsonl
│   │   └── ...
│   └── valid
│       ├── 0.jsonl
│       ├── 1.jsonl
│       └── ...
└── ...

You can download the 30B corpus of CommonCrawl (from SlimPajama), and the high-quality FineWeb-pro dataset.

Then run the following:

export TEXT_DIR=<YOUR_TEXT_DIR>
export BINS_ROOT=<YOUR_BIN_DIR> # where to store the processed chunks
bash scripts/pajama_processing.sh cc 8k

where cc is the dataset name and 8k is the sequence length (supporting from 512 to 16k). The TEXT_DIR is the directory where the text data is stored and the BIN_DIR is the directory where the processed data will be stored. After this step, you will have the data in the following structure:

BINS_ROOT
├── cc_8k
│   ├── train_0.bin
│   ├── train_1.bin
│   ├── ...
│   ├── valid_0.bin
│   ├── valid_1.bin
│   └── ...
└── ...

Pretraining

Next, you can start pretraining by running the following:

export WANDB_API_KEY=<YOUR_WANDB_API_KEY> # if you want to log into wandb
export BINS_ROOT=<YOUR_BIN_DIR> # from the previous data preparation step
bash scripts/pretraining.sh tiny_LLaMA_1b_8k cc_8k cc_8k # replace 1b_8k with 120M_8k or 360M_8k for smaller models

The general usage of pretraining.sh is bash scripts/pretraining.sh model_config_name train_dataset_name eval_dataset_name. For instance, tiny_LLaMA_1b_8k is the model config name, cc_8k is the training dataset name, and cc_8k is the evaluation dataset name. The script will look for bins created in the previous step. Those with a train_* prefix are used for training and those with a valid_* prefix are used for evaluation.

You can simply replace the model config name to get different models:

bash scripts/pretraining.sh tiny_LLaMA_1b_8k cc_8k cc_8k # baseline with standard causal attention
bash scripts/pretraining.sh tiny_LLaMA_1b_8k_intramask cc_8k cc_8k # intradocument masking
bash scripts/pretraining.sh tiny_LLaMA_1b_8k_dm8 cc_8k cc_8k # skyladder with alpha=1/8
bash scripts/pretraining.sh tiny_LLaMA_1b_8k_intradm8 cc_8k cc_8k # intradocument masking + skyladder with alpha=1/8

Here, dm8 means that