为什么值得关注

它借助 JAX 的 GSPMD 编译器原生支持数据并行、全分片数据并行和张量并行，代码量远少于等效的 PyTorch 方案；内建光谱 μTransfer 用于超参数缩放；可在笔记本电脑到大型 TPU Pod 上运行。

适合谁使用

训练大语言模型的研究人员和工程师
希望获得分布式训练参考实现的 JAX 用户
从 PyTorch 转向 JAX 进行大语言模型开发的开发者
对 μTransfer（最大更新参数化）感兴趣的实践者

典型使用场景

从零开始训练具有精确 RoPE、SwiGLU 和分组查询注意力的 LLaMA 风格模型
从单 GPU 扩展到多主机 TPU 集群，自动分片训练
使用 μTransfer 在小代理模型上找到最佳超参数并迁移到大模型
利用 jax.lax.scan 实现内存高效的梯度累积以支持大批量训练

项目优势

原生支持 3D 并行（数据并行、FSDP、张量并行），无需手动编写集合通信——JAX 的 XLA/GSPMD 自动处理 all-gather 和 reduce-scatter
忠实复现 Meta 的 LLaMA 架构，包括 RoPE、SwiGLU 和分组查询注意力
内建光谱 μTransfer：根据模型宽度/深度自动缩放学习率和初始化范围
通过 jax.lax.scan 实现内存高效的微批次处理，降低梯度累积期间的 HBM 用量

使用前须知

需要 JAX 及兼容硬件（GPU 或 TPU），不支持大型模型的 CPU 训练
对于不熟悉 JAX 函数式编程模型和 Equinox 的用户学习曲线陡峭
紧密依赖特定库（Equinox、Optax、Orbax），可能难以与其他 JAX 生态集成

README 快速开始

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 AllGather and ReduceScatter collectives 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.scan loop cleanly calculates gradient accumulations without blowing out accelerator HBM memory.
Modern Infrastructure: Uses equinox for clean PyTree OOP modules, optax for optimization, orbax for async multi-host checkpointing, and wandb for 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

项目描述

A minimal, 3D-parallel LLM trainer written in pure JAX and Equinox.

jaxotron