princeton-nlp/prolong-data-64K

[Paper] [HF Collection] [Code]

ProLong (Princeton long-context language models) is a family of long-context models that are continued trained and supervised fine-tuned from Llama-3-8B, with a maximum context window of 512K tokens. Our main ProLong model is one of the best-performing long-context models at the 10B scale (evaluated by HELMET).

To train this strong long-context model, we conduct thorough ablations on the long-context pre-training data, SFT data, and numerous other design choices. We demonstrate our findings in our paper, How to Train Long-Context Language Models (Effectively).

Authors: Tianyu Gao*, Alexander Wettig*, Howard Yen, Danqi Chen (* equal contribution)

Contact: {tianyug, awettig}@princeton.edu

Dataset Loading

This dataset contains 31B tokens, tokenzized with the Llama-3 tokenizer and packed to sequences of 65,536 tokens. The data is stored as MDS (Mosaic Data Shard) and requires mosaicml-streaming to be loaded. Instead of datasets.load_dataset, download the data by cloning the repository or the huggingface_hub.snapshot_download function.

When loading the datasets with mosaicml-streaming, each entry has the following fields:

• input_ids: a 1-dimensional array of length 65,536 containing the token ids
• indices: a list of (start_index, end_index) tuples that identify the subsequences in input_ids of separate documents. This is particularly important for short-context datasets that are packed to 524,288 sequence length
• domain: (optional) string of the dataset split

This dataset contains the following subsets as folders:

The ProLong Models

• princeton_nlp/Llama-3-8B-ProLong-64k-Base
• princeton_nlp/Llama-3-8B-ProLong-64k-Instruct
• princeton_nlp/Llama-3-8B-ProLong-512k-Base
• ⭐ princeton_nlp/Llama-3-8B-ProLong-512k-Instruct

The ProLong Data

• Stage 1 64K training: princeton-nlp/prolong-data-64K ← you are here!
• Stage 2 128K training: princeton-nlp/prolong-data-512K

Data Compositions

ProLong training data and recipe.

Citation

@article{gao2024prolong,
    title={Enabling Large Language Models to Generate Text with Citations},
    author={Gao, Tianyu and Wettig, Alexander and Yen, Howard and Chen, Danqi},
    year={2024},
}