slvDev
Featured
esp32-ai
A 28.9 million parameter language model runs on an $8 ESP32-S3 microcontroller entirely on-device, generating simple stories at about 9.5 tokens per second.
AI repository intelligence
FR-AI / ANALYSISWhy it is worth attention
It demonstrates state-of-the-art 3D parallelism (DP, FSDP, TP) entirely through JAX's GSPMD compiler, requires far less code than equivalent PyTorch solutions, includes spectral muTransfer for hyperparameter scaling, and is designed to run from a laptop to large TPU pods.
Who it is for
- Researchers and engineers training large language models
- JAX users seeking a reference implementation for distributed training
- Developers transitioning from PyTorch to JAX for LLM work
- Practitioners interested in muTransfer (maximal update parameterization)
Use cases
- Training LLaMA-style models from scratch with exact RoPE, SwiGLU, and GQA
- Scaling training from a single GPU to multi-host TPU clusters with automatic sharding
- Using muTransfer to find optimal hyperparameters on a small proxy model and transfer them to a larger model
- Memory-efficient gradient accumulation using jax.lax.scan for large batch sizes
Strengths
- Native 3D parallelism (DP, FSDP, TP) without manual collective calls—JAX's XLA/GSPMD handles all-gather and reduce-scatter automatically
- Faithful replication of Meta's LLaMA architecture including RoPE, SwiGLU, and Grouped-Query Attention
- Spectral muTransfer (muP) built in: automatically scales learning rates and initialization bounds based on model width/depth
- Memory-efficient micro-batching via jax.lax.scan, reducing HBM usage during gradient accumulation
Considerations
- Requires JAX and compatible hardware (GPU or TPU); no CPU training for large models
- Steep learning curve for those unfamiliar with JAX's functional programming model and Equinox
- Tightly tied to specific libraries (Equinox, Optax, Orbax); may not integrate seamlessly with other JAX ecosystems
README quick start
Jaxotron 🚀
Jaxotron is a highly-scalable, educational, and production-ready JAX adaptation of Hugging Face's Nanotron, heavily inspired by the Google DeepMind "How to Scale Your Model" JAX book.
Jaxotron proves that you can build state-of-the-art 3D parallelized LLM training pipelines in a fraction of the code compared to PyTorch, all by leveraging JAX's GSPMD compiler magic and functional programming.
✨ Key Features
- Native 3D Parallelism: Seamlessly mix Data Parallelism (DP), Fully Sharded Data Parallelism (FSDP/ZeRO-3), and Tensor Parallelism (TP). JAX's XLA compiler automatically inserts
AllGatherandReduceScattercollectives based on semantic mesh rules. - Exact LLaMA Architecture: Faithfully replicates Meta's original Rotary Positional Embeddings (RoPE), SwiGLU gated MLPs, and Grouped-Query Attention (GQA). Perfect for training models like SmolLM.
- Spectral muTransfer (muP): Automatically scales learning rates and initialization bounds depending on width/depth ratios using Greg Yang's muTransfer rules. Train a proxy model and perfectly transfer hyperparameters to scale.
- Memory-Efficient Micro-batching: A
jax.lax.scanloop cleanly calculates gradient accumulations without blowing out accelerator HBM memory. - Modern Infrastructure: Uses
equinoxfor clean PyTree OOP modules,optaxfor optimization,orbaxfor async multi-host checkpointing, andwandbfor logging.
📁 Project Structure
jaxotron/
├── app.py # Entry point for training
├── config.py # Strongly typed dataclasses for topology/model/training
├── data/
│ ├── collator.py # Asynchronous JAX data prefetcher
│ └── nanoset.py # Fast memory-mapped binary dataset reader
├── models/
│ ├── attention.py # GQA Attention & SwiGLU MLP
│ ├── layers.py # RMSNorm & Exact RoPE math
│ └── llama.py # Main LlamaForCausalLM model
├── optim/
│ ├── build.py # Optax AdamW + LR Schedulers
│ └── mup.py # Spectral muP initialization & LR multipliers
├── parallel/
│ ├── mesh.py # JAX device mesh creation
│ ├── partitioning.py # Megatron/FSDP logical axis rules
│ └── sharding_utils.py # PyTree node sharding applier
└── serialize/
└── orbax_utils.py # Async P
Description
A minimal, 3D-parallel LLM trainer written in pure JAX and Equinox.
Related repositories
Similar projects matched by category, topics, and programming language.
gavamedia
Featured
deltafin
Deltafin is a research project that runs the 2.8-trillion-parameter Mixture-of-Experts model Kimi K3 on a single Apple Silicon Mac (e.g., M1 Max with 64 GB) at about 16 seconds per token, using exact, reproducible inference with local or streaming expert loading.
jamesob
Featured
local-llm
A comprehensive guide for building and configuring a high-end local machine to run state-of-the-art LLMs, with detailed hardware choices, BIOS tuning, and Docker-based model serving.