hazyresearch / based Goto Github PK

Simple linear attention language models balance the recall-throughput tradeoff.

Based is an efficient architecture inspired by recovering attention-like capabilities (i.e., recall). We do so by combining 2 simple ideas:

1. Short sliding window attention (e.g., window size 64), to model fine-grained local dependencies
2. "Dense" and global linear attention, to model long-range dependencies

In this way, we aim to capture the same dependencies as Transformers in a 100% subquadratic model, with exact softmax attention locally and a softmax-approximating linear attention for all other tokens.

We find this helps close many of the performance gaps between Transformers and recent subquadratic alternatives (matching perplexity is not all you need? [1, 2, 3]).

In this repo, please find code to (1) train new models and (2) evaluate existing checkpoints on downstream tasks.

Note. The code in this repository is tested on python=3.8.18 and torch=2.1.2. We recommend using these versions in a clean environment.

# clone the repository
git clone [email protected]:HazyResearch/based.git
cd based

# install torch
pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/cu118 # due to observed causal-conv1d dependency

# install based package
pip install -e .

# Note that sometimes the causal-conv1d interface changes (https://github.com/state-spaces/mamba/pull/168) in case you run into an error.

We are releasing the following checkpoints for research, trained at the 360M and 1.3B parameter scales. Each checkpoint is trained on the same 10B to 50B tokens (specified below) of the Pile corpus, using the same data order. The checkpoints are trained using the same code and infrastructure. A quick start notebook is provided at notebooks/03-24-quick-start.ipynb and further details are below:

Use the code below to load the Based checkpoints:

import torch
from transformers import AutoTokenizer
from based.models.gpt import GPTLMHeadModel

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = GPTLMHeadModel.from_pretrained_hf("hazyresearch/based-360m")

Warning. We are releasing these models for the purpose of efficient architecture research. Because they have not been instruction fine-tuned or audited, they are not intended for use in any downstream applications.

The following code will run text generation for a prompt and print out the response.

input = tokenizer.encode("If I take one more step, it will be", return_tensors="pt").to("cuda")
output = model.generate(input, max_length=20)
print(tokenizer.decode(output[0]))

Note. For the checkpoints from other models, you will need to install other dependencies and use slightly different code.

To load the Attention models, use the following code:

import torch
from transformers import AutoTokenizer
from based.models.transformer.gpt import GPTLMHeadModel

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = GPTLMHeadModel.from_pretrained_hf("hazyresearch/attn-360m").to("cuda")

To use the Mamba checkpoints, first run pip install mamba-ssm and then use the following code:

import torch
from transformers import AutoTokenizer
from based.models.mamba import MambaLMHeadModel

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = MambaLMHeadModel.from_pretrained_hf("hazyresearch/mamba-360m").to("cuda")

In order to train a new model with our code, you'll need to do a bit of additional setup:

# install train extra dependencies
pip install -e .[train]

# install apex (if you run into issues, likely candidates are torch or pip version issues; if using torch 2.0.1, this may help https://github.com/NVIDIA/apex/issues/1735)
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
cd ..

We breakdown this section into three parts: 1) how to set up a training config and launch; 2) how to set up fast training kernels, and 3) how to install extra optimizations for training.

To train a new model, construct a config.yaml file at train/configs/experiment/. We are including the configs used to produce the pretrained checkpoints for the paper (released on HF below) at train/configs/experiment/reference/.

You can launch a training job using the following command from the train/ directory, where you can modify the config name and number of GPUs (trainer.devices):

cd train/
python run.py experiment=reference/based-1b trainer.devices=8

In our paper, we evaluated on the Pile corpus, which is no longer available online, so the train/configs/experiment/reference/ configs are unfortunately not directly runnable. For your use, we are including an example config that would train on the WikiText103 language modeling data. You can launch using the following script:

cd train/
python run.py experiment=example/based-360m trainer.devices=8

You can adapt the training dataset by adding a new dataset config file under train/configs/datamodule/. Follow the examples in wikitext103.yaml. Once you've constructed the yaml file for your new dataset, go to the experiment config (e.g. train/configs/experiment/example/based-360m.yaml) and update the name of the datamodule under override datamodule to the filename of your new dataset yaml file.

Be sure to update the checkpointing directory in the config prior to launching training.

We support a few different training views in this repo. The choice of parallel_implementation in your training config determines which training view gets used:

• Option 1 (parallel_implementation = "quadratic"): default, quadratic PyTorch view.
• Option 2 (parallel_implementation = "fla_parallel"): Flash linear attention kernel. Use the following to install:

pip install triton==2.2.0
pip install -U git+https://github.com/sustcsonglin/flash-linear-attention

• Option 3 (parallel_implementation = "linear"): Fast transformers linear attention kernel. Use the following to install:

cd train/csrc/causal_dot_prod/
python setup.py install

We have provided benchmarking plots for different kernels in the benchmark/examples/linear_attention_forward/ folder. We are providing WandB training curves here showing how training using the fla-parallel mode allows Based to train fast with strong quality!

• Kernels for other fused operations: The config defaults will use fused kernels from the Flash Attention repo, which can all be installed by cloning the Flash Attention repo and python setup.py install the relevant kernels here. In particular, the fused_dense_lib, layer_norm, rotary, and xentropy kernels. Alternatively, you can change the codepaths to avoid the use of these kernels -- for instance by specifying fused_dense False in the experiment config, or by replacing the RMSNorm import in based/models/gpt.py to import from based/ops/triton/layer_norm.
• Decay If you want to explore the optional decay strategy discussed in the Based paper, you can checkout the notebooks/03-31-decay.ipynb notebook.
• References Note that this training code is from: https://github.com/Dao-AILab/flash-attention/tree/main/training, the Flash Linear Attention kernel is from https://github.com/sustcsonglin/flash-linear-attention, and the Fast Transformers kernel is from https://github.com/idiap/fast-transformers. Please cite them if you use their work!

In our paper, we evaluate pretrained language models on a standard suite of benchmarks from the LM Evaluation Harness, as well as a suite of three recall-intensive tasks:

• SWDE (Info. extraction). A popular information extraction benchmark for semi-structured data. SWDE includes raw HTML documents from 8 Movie and 5 University websites (e.g.IMDB, US News) and annotations for 8-274 attributes per website (e.g., Movie runtime). HuggingFace: hazyresearch/based-swde
• FDA (Info. extraction). A popular information extraction benchmark for unstructured data. The FDA setting contains 16 gold attributes and 100 PDF documents, which are up to 20 pages long, randomly sampled from FDA 510(k). HuggingFace: hazyresearch/based-fda
• SQUAD-Completion (Document-QA). We find that original SQUAD dataset is challenging for our models without instruction fine-tuning. So we introduce a modified version of SQUAD where questions are reworded as next-token prediction tasks. For example, "What is the capital of France?" becomes "The capital of France is". HuggingFace: hazyresearch/based-squad

Under evaluate, we have a clone of EleutherAI's lm-evaluation-harness that includes these new tasks and provides scripts for running all the evaluations from the paper. The following instructions can be used to reproduce our results on the LM-Eval harness using the pretrained checkpoints.

cd evaluate 

# init the submodule and install
git submodule init
git submodule update
pip install -e . 

We provide a script evaluate/launch.py that launch evaluations on the checkpoints we've released.

For example, running the following from the evaluate folder will evaluate the 360M Based, Mamba, and Attention models on the SWDE dataset.

You can set your huggingface cache directory to a location with sufficient space (export TRANSFORMERS_CACHE, export HF_HOME).

python launch.py \
    --task swde  --task fda --task squad_completion \
    --model "hazyresearch/based-360m" \
    --model "hazyresearch/mamba-360m" \
    --model "hazyresearch/attn-360m" \
    --model "hazyresearch/based-1b" \
    --model "hazyresearch/mamba-1b" \
    --model "hazyresearch/attn-1b"

Optionally, if you have access to multiple GPUs, you can pass the -p flag to run each evaluation on a different GPU. To run a limited number of samples for each task (e.g. 100), use the --limit=100 option.

Below we include the results produced from running the command above. Note: the results below are on the new models trained and evaluated with the cleaned-up code in this repository. As a result, the results reported in our paper differ slightly, however the trends and conclusions remain the same.

Note that the results shown may differ slightly if the Flash-Attention kernels are not used during inference.

In our paper, we demonstrate the recall-throughput tradeoff using a synthetic associative recall task (see Figure 2, below, and Figure 3 in the paper).

The code for reproducing these figures is provided in a separate repository: HazyResearch/zoology. Follow the setup instruction in the Zoology README. The instructions for reproducing the are provided in zoology/experiments. For example, you can create the figure above using.

python -m zoology.launch zoology/experiments/arxiv24_based_figure2/configs.py -p

We include the kernels evaluated in the Based paper under based/benchmarking/. We provide additional details on the CUDA releases in the README in this folder. Stay tuned!

This repo contains work based on the following papers. Please consider citing if you found the work or code useful:

# Based
@article{arora2024simple,
  title={Simple linear attention language models balance the recall-throughput tradeoff},
  author={Arora, Simran and Eyuboglu, Sabri and Zhang, Michael and Timalsina, Aman and Alberti, Silas and Zinsley, Dylan and Zou, James and Rudra, Atri and Ré, Christopher},
  journal={arXiv:2402.18668},
  year={2024}
}

# Hedgehog (Linear attention)
@article{zhang2024hedgehog,
  title={The Hedgehog \& the Porcupine: Expressive Linear Attentions with Softmax Mimicry},
  author={Zhang, Michael and Bhatia, Kush and Kumbong, Hermann and R{\'e}, Christopher},
  journal={arXiv preprint arXiv:2402.04347},
  year={2024}
}

# Zoology (BaseConv, Synthetics, Recall Problem)
@article{arora2023zoology,
  title={Zoology: Measuring and Improving Recall in Efficient Language Models},
  author={Arora, Simran and Eyuboglu, Sabri and Timalsina, Aman and Johnson, Isys and Poli, Michael and Zou, James and Rudra, Atri and Ré, Christopher},
  journal={arXiv:2312.04927},
  year={2023}
}

This project was made possible by a number of other open source projects; please cite if you use their work! Notably:

• Our training code and sliding window implementation are based on Tri Dao's FlashAttention.
• We use EleutherAI's lm-evaluation-harness for evaluation.
• We use the conv1d kernel from Mamba.
• We integrated the causal dot product kernel from Fast Transformers.
• We integrated the based kernels from Flash Linear Attention.

Models in this project were trained using compute provided by:

• Together.ai
• Google Cloud Platform through Stanford HAI

Please reach out with feedback and questions!