Llama 3.1, from PyTorch to JAX
We rewrote Llama 3.1 in pure JAX, so one codebase runs on TPU, Trainium, AMD and NVIDIA. There are no per-chip code paths: XLA compiles the same model for each, from a single 8-core TPU VM up to a 6,000-core pod.
Felafax builds RL environments that teach frontier models to write fast JAX and kernels for XLA chips.
Environments target TPU, Trainium and AMD. Models that speed up their own training and inference are a first step toward recursive self‑improvement.
PRIOR WORK
Making JAX fast on hardware that isn't NVIDIA, written by hand before coding agents could help. Now we're teaching models to do the same work.
Llama 3.1 405B on eight AMD MI300X
We fine-tuned the 405B model in JAX on a 2D mesh of eight GPUs: model parallel for the weights, data parallel for the batch. A few sharding annotations replace hand-written communication, and scaling was near-linear.
FastPitch + HiFi-GAN on TPU
We converted NVIDIA's text-to-speech models to run on Google TPUs with PyTorch/XLA. The speedup came from the XLA layer: static shapes so each graph compiles once, bfloat16, and fewer host–device copies. About 160× real time on a v5p.
Whisper JAX, served on TPU
Open-source speech transcription with OpenAI's Whisper in JAX. The same code runs on CPU, GPU and TPU: audio is decoded, chunked and transcribed, and an HTTP server with a Docker image makes it a one-command deploy on a TPU VM.
We build environments with a few frontier labs that run JAX and XLA at scale.
founders@felafax.ai