llamax

A basic Flax library which contains reference implementations of various LLMs. Currently, just Llama3.

The code is very much a work in-progress; use at your own risk.

TODO

• Verify that Flax code matches reference implementation at 6 significant digits in float64 on CPU.
• Support passing in kv_mask and pairwise_mask to support padded sequences. Test that this works properly.
• Implement sharding, at least model-parallel + data-parallel.
• Change the main transformer loop to use a scan instead.
• Add throughput benchmarks.
• Test against the actual llama3 weights.
• Add code that can run tests using Modal Labs.
• (Stretch goal) Implement some of the fancier sharding methods listed in Meta's Movie Gen paper.
• Add tests on GPU.
• Pin dependency versions.

Setup

Prerequisites

Install uv for Python package management:

curl -LsSf https://astral.sh/uv/install.sh | sh

Installation

1. Clone the repository and install dependencies:

uv sync --extra dev

2. Install pre-commit hooks:

uv run pre-commit install

Usage

Running Tests

uv run pytest

Linting and Formatting

Run the linter and formatter:

uv run ruff check .
uv run ruff format .

Or use the Makefile (if available):

make style && make quality

Pre-commit hooks will automatically run ruff on staged files before each commit.

Docker

Run tests with docker build -t llamax . && docker run -e JAX_ENABLE_X64=True -e HF_TOKEN=$HF_TOKEN llamax pytest.

Get your Huggingface token by following their instructions. If you want to run the integration test, download weights as per instructions below, and then pass them in via -v $WEIGHTS_DIR:/data in the Docker command above.

Weights

I don't include any of the weights needed to test the implementations. You'll need to get them yourself. Here's where I got them:

1. Llama 3.2 1B: Copy the "original" folder from Huggingface.