knightXun/jax-llm-examples

JAX LLM examples

A collection (in progress) of example high-performance large language model implementations, written with JAX.

Current contents include:

• DeepSeek R1
• Llama 4
• Llama 3
• Qwen 3
• Kimi K2
• OpenAI GPT OSS

---

For multi-host cluster setup and distributed training, see multi_host_README.md and the tpu_toolkit.sh script.