Awesome

LLaMA Jax

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

Setup

In a conda env, run

pip install -r requirements.txt

Then in this repository

pip install -e .

Download

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.

Inference

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

Model Card

See MODEL_CARD.md

License

See the LICENSE file.

Testing

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: