Minimal GPT-NeoX-20B

This is a fairly minimal implementation of GPT-NeoX-20B in PyTorch. It is meant primarily as an educational/reference implementation, rather than an optimized or feature-full implementation.

GPT-NeoX-20B is a 20B-parameter autoregressive Transformer model developed by EleutherAI with the support of CoreWeave, trained using the GPT-NeoX library.

Some notes about the model:

• The model weights and activations come in half-precision (fp16).
• In fp16, loading the model weights requires about 40GB of GPU memory. Running inference on a single batch requires some more.
• The model supports up to a maximum sequence length of 2048 tokens.

Setup

Installation

Install PyTorch with your appropriate CUDA version, and then install from the requirements.txt (basically just tokenizers).

pip install -r requirements.txt

Download weights

Following the NeoX guide, download the model weights and tokenizer JSON file with the following command:

wget --cut-dirs=5 -nH -r --no-parent --reject "index.html*" https://mystic.the-eye.eu/public/AI/models/GPT-NeoX-20B/slim_weights/ -P 20B_checkpoints

You can also manually down them from here. Because of the size of the model, the model weights are broken into multiple files, based on the DeepSpeed save format.

Generate text

Here is some sample code to generate text. Note that since we are greedily decoding with no fancy tricks, there tends to be quite some repetitiion in generations.

import minimal20b
import torch
model = minimal20b.create_model(
    "/path/to/20B_checkpoints/global_step150000",
    use_cache=True,
    device="cuda:0",
)
tokenizer = minimal20b.create_tokenizer(
    "/path/to/20B_checkpoints/20B_tokenizer.json",
)
with torch.inference_mode():
    minimal20b.greedy_generate_text(
        model, tokenizer,
        "GPTNeoX20B is a 20B-parameter autoregressive Transformer model developed by EleutherAI.",
        max_seq_len=100,
    )

Evaluation

To run evaluation with the LM-eval-harness, you will need to install some additional dependencies (mostly just the eval harness library):

pip install -r scripts/eval/requirements.txt

Most datasets are automatically downloaded via Hugging Face datasets, but if you are evaluating on lambada, you will need to separately download the data.

mkdir -p data/lambada
wget http://eaidata.bmk.sh/data/lambada_test.jsonl -O data/lambada/lambada_test.jsonl

Then, you can run the following command.

python scripts/eval/eval_harness.py \
    --model_path /path/to/20B_checkpoints/global_step150000 \
    --tokenizer_path /path/to/20B_checkpoints/20B_tokenizer.json \
    --tasks lambada,anli_r1,anli_r2,anli_r3,wsc,winogrande,hellaswag,piqa

Task	Metric	NeoX Impl (2 GPU)	This Repo (1 GPU)
anli_r1	acc	0.3270	0.3300
	acc_stderr	0.0148	0.0149
anli_r2	acc	0.3410	0.3420
	acc_stderr	0.0150	0.0150
anli_r3	acc	0.3567	0.3617
	acc_stderr	0.0138	0.0139
hellaswag	acc	0.5351	0.5335
	acc_stderr	0.0050	0.0050
	acc_norm	0.7140	0.7126
	acc_norm_stderr	0.0045	0.0045
lambada	acc	0.7211	0.7223
	acc_stderr	0.0062	0.0062
	ppl	3.6760	3.6559
	ppl_stderr	0.0760	0.0757
piqa	acc	0.7748	0.7758
	acc_stderr	0.0097	0.0097
	acc_norm	0.7786	0.7856
	acc_norm_stderr	0.0097	0.0096
winogrande	acc	0.6598	0.6598
	acc_stderr	0.0133	0.0133
wsc	acc	0.5096	0.4808
	acc_stderr	0.0493	0.0492