raijinspecial / gpt-neox-depricated

An implementation of model parallel GPT-3-like models on GPUs, based on the DeepSpeed library. Designed to be able to train models in the hundreds of billions of parameters or larger.

GPT-NeoX

An implementation of model parallel GPT-3-like models on GPUs, based on the DeepSpeed library. Designed to be able to train models in the hundreds of billions of parameters or larger.

Requirements

$ pip install -r requirements.txt

Test deepspeed locally

$ deepspeed train_enwik8.py \
	--deepspeed \
	--deepspeed_config ./configs/base_deepspeed.json

Sparse Attention

To use sparse attention in your GPTNeoX model, you first need to make sure Deepspeed is installed with sparse attention enabled. You can use the following script to install all the dependencies as well as reinstall Deepspeed.

$ ./install_deepspeed.sh

Then

model = GPTNeoX(
    num_tokens = 20000,
    dim = 512,
    seq_len = SEQ_LEN,
    depth = 12,
    heads = 8,
    sparse_attn = True,
)

Or if you want it for specific layers

model = GPTNeoX(
    num_tokens = 20000,
    dim = 512,
    seq_len = SEQ_LEN,
    depth = 12,
    heads = 8,
    sparse_attn = (True, False) * 6, # interleaved
)

Resources

If you have trouble getting the model to run, consider consulting this guide to installing in a GCE virtual machine.