UPDATE: now with LLaMA-3 support!
This repository is a Huggingface compatible port of LLaMA (arXiv) models to Jax. In order to download the checkpoints and tokenizer for LLaMA, fill this google form
In a conda env, run
pip install -r requirements.txt
Then in this repository
pip install -e .
Once your request is approved, you will receive links to download the tokenizer and model files.
Edit the download.sh script with the signed url provided in the email to download the model weights and tokenizer.
The provided jax_example.py will output completions for two pre-defined prompts. Using TARGET_FOLDER as defined in download.sh:
python jax_example.py --ckpt_dir $TARGET_FOLDER/model_size --tokenizer_path $TARGET_FOLDER/tokenizer.model
See MODEL_CARD.md
See the LICENSE file.
The provided jax_test.py script runs a comparison between this jax model and the pytorch version provided by Meta (to test LLaMA 3, use the Meta LLaMA 3 repo instead). To run the tests, install Meta's code in the same environment and run the script with:
torchrun --nproc_per_node MP jax_test.py --ckpt_dir $TARGET_FOLDER/model_size --tokenizer_path $TARGET_FOLDER/tokenizer.model --is_llama3
(Note: some of the tests only run when MP=1)
if testing with LLaMA 1 or 2, drop the --is_llama3 flag (and make sure you install the correct repo for testing against).
Different models require different MP values:
| Model | MP |
|---|---|
| 7B | 1 |
| 8B | 1 |
| 13B | 2 |
| 33B | 4 |
| 65B | 8 |
| 70B | 8 |