This repository is intended as a minimal, hackable and readable example to load LLaMA (arXiv) models and run inference. In order to download the checkpoints and tokenizer, fill this google form
In a conda env with pytorch / cuda available, run
pip install -r requirements.txt
Then in this repository
pip install -e .
The provided example.py
can be run on a single or multi-gpu node with torchrun
and will output completions for two pre-defined prompts. Using TARGET_FOLDER
as defined in download.sh
:
torchrun --nproc_per_node MP example.py --ckpt_dir $TARGET_FOLDER/model_size --tokenizer_path $TARGET_FOLDER/tokenizer.model
Different models require different MP values:
Model | MP |
---|---|
7B | 1 |
13B | 2 |
33B | 4 |
65B | 8 |
Use the gradio demo to play around with inference in an interactive way. Make sure to change the model paths in app.py.
torchrun --nproc_per_node 1 app.py
This demo requires ~16GB GPU memory, but could also run on CPU if you adapt a few lines. See this example for the app in use:
Recording.2023-03-02.225512.mp4
See MODEL_CARD.md
See the LICENSE file.