nanoAlphaZero: A Single-File JAX Implementation That Trains to 2700 Elo in 24 Hours

One Person's AlphaZero Challenge

In 2017, DeepMind's AlphaZero burst onto the scene, achieving superhuman performance in chess, shogi, and Go purely through self-play—a milestone in reinforcement learning history. AlphaZero's core breakthrough was completely abandoning human game data, learning game strategies from scratch solely through self-play and reinforcement learning. Before this, DeepMind's AlphaGo series still relied on large amounts of human expert games as a training starting point. AlphaZero unified the learning framework across three different board games, proving the viability of general-purpose algorithms—its training process combined deep neural networks (for evaluating positions and predicting move probabilities) with Monte Carlo Tree Search (for lookahead planning during games), the two reinforcing each other in a positive feedback loop.

However, reproducing AlphaZero has always been an extremely high-barrier engineering challenge—it involves distributed self-play, Monte Carlo Tree Search (MCTS), large-scale neural network training, and complex engineering orchestration. The original paper used 5,000 TPUs for self-play and 64 TPUs for training, hardware requirements that made reproduction virtually impossible for academics and individual developers.

Recently, a developer (GitHub user wtedw) shared his open-source project nanoAlphaZero on Reddit: a complete, single-file AlphaZero implementation written in JAX. His original goal sounded quite ambitious—"speedrun a GM-level chess AI from scratch in one month."

As he candidly admitted: "It didn't go smoothly." The project ultimately took over two years, going through multiple complete rewrites before finally stabilizing.

Training to Elo 2700+ in 24 Hours

Despite the winding path, the final results are quite impressive. According to the author, on a single TPU v4-32 pod, nanoAlphaZero can train a chess model to Elo 2700+ within 24 hours.

TPU (Tensor Processing Unit) is a custom chip designed by Google specifically for machine learning workloads. A single TPU v4 chip has peak performance of approximately 275 TFLOPS (BF16 precision). A TPU v4-32 pod refers to a compute unit consisting of 32 TPU v4 chips connected via Google's proprietary ICI high-speed interconnect, with total compute power of approximately 8,800 TFLOPS. While this is still a significant resource investment for individuals, compared to the original AlphaZero's use of thousands of TPUs, it represents a two-order-of-magnitude reduction and can be accessed on-demand through Google Cloud.

For reference, Elo 2700 is solidly above the level of human chess Grandmasters (GM). The Elo rating system was designed by Hungarian-American physicist Arpad Elo, and its core idea is that the rating difference between two players can be converted into expected win probability—a 400-point difference means the higher-rated player has approximately a 91% win rate. In chess, Elo 2500+ qualifies as Grandmaster level, 2700+ is super-Grandmaster level, with only about 50 people worldwide able to maintain ratings at this level. The current top players in the world have ratings around 2800. This means the system—starting from zero, requiring no human game data, and learning purely through self-play—indeed achieved the project's original "GM-level" goal.

What's even more interesting is that its core logic is game-agnostic. Besides chess, it currently supports:

• Tic-Tac-Toe
• Connect Four
• Hex
• Small-size Go boards

For smaller, fully solvable games, nanoAlphaZero can learn perfect play. This also serves as indirect verification of the underlying algorithm's correctness. Support for larger Go boards is still under development.

The Entire RL Pipeline Is a Single JAX Function

The most engineering-elegant aspect of nanoAlphaZero lies in its extreme simplification of a complex system. The author's core design philosophy is: the entire reinforcement learning pipeline is one large, JIT-compiled JAX function.

JAX is a high-performance numerical computing library developed by Google, which can be understood as "differentiable, compilable, vectorizable NumPy." Its three core primitives make it particularly well-suited for large-scale reinforcement learning: jit (just-in-time compilation) compiles Python functions into XLA (Accelerated Linear Algebra) computation graphs that execute efficiently directly on TPU/GPU; vmap (vectorized map) automatically batches single-sample logic into parallel computation without manually writing batch processing code; pmap (parallel map) implements multi-device data parallelism. This means developers can write clean single-game logic and then automatically scale it to thousands of parallel games through vmap, achieving both code readability and computational efficiency. Additionally, JAX's functional programming paradigm means the entire computation graph consists of pure functions, naturally supporting compilation optimization and reproducibility.

Traditional AlphaZero implementations often require managing threads, servers, and distributed worker nodes, making engineering complexity extremely high. nanoAlphaZero absorbs all of this into a single run_alphazero function that repeatedly executes "self-play + model update":

state = make_alphazero()

def run_alphazero(state):
    state, games = selfplay(state)  # Uses Gumbel MuZero
    # Move ongoing games into selfplay buffer
    # Move finished games into replay buffer
    state = train(state, replay_buffer.sample())
    return state

while True:
    state = run_alphazero(state)

The author emphasizes: "No threads, no servers, no distributed workers to manage. The entire RL pipeline is one big JAX function."

Behind this design is JAX's core advantage—through jit compilation and primitives like vmap, it maps large-scale parallel computation onto TPU hardware, achieving extremely high throughput efficiency while keeping code concise. It's worth noting that the self-play component uses the Gumbel MuZero search strategy. Gumbel MuZero is an improved tree search algorithm proposed by DeepMind in 2022 that solves the instability of policy improvement in traditional MCTS with limited simulations. Traditional AlphaZero's MCTS relies on many simulations (typically 800+) to obtain reliable move distributions, while Gumbel MuZero introduces Gumbel noise sampling and Sequential Halving strategies to guarantee policy improvement even with very few simulations (e.g., 16 or 32). Its mathematical foundation comes from the Gumbel-Top-k trick, which guarantees optimality of sampling. This property is crucial for large-scale parallel training—when thousands of self-play games need to run simultaneously, the search budget allocated per game must be very small, and Gumbel MuZero's efficiency advantage becomes very significant in this regime.

Project Positioning: Making Large-Scale AlphaZero Experiments More Accessible

The author has a clear-eyed positioning for the project. He explicitly states that nanoAlphaZero's primary goal is to make large-scale AlphaZero experiments more accessible, with code that pursues speed and memory efficiency while remaining compact and hackable.

As for training strong models, that's actually a secondary goal—mainly serving as a "sanity check" to verify the underlying logic is correct. This "tools-first, performance-second" approach is in the same spirit as Andrej Karpathy's nanoGPT: using the least and clearest code possible to explain the essence of a complex algorithm, enabling researchers and enthusiasts to quickly understand, modify, and experiment.

For teaching and research, these "nano" projects often provide more value than industrial-grade frameworks—they strip away engineering noise and go straight to the algorithmic core. Researchers can quickly validate new ideas on this foundation, such as modifying search strategies, adjusting network architectures, or trying different training schedules, without spending weeks understanding complex distributed systems code.

v2: A Faster Game Engine Is on the Way

The author has also previewed an upcoming v2 major refactor, which includes several substantial improvements:

• Switching to the general-purpose KataGo neural network architecture: KataGo is an open-source Go AI developed by David Wu, widely considered one of the strongest open-source Go programs available. Its neural network architecture has been iteratively optimized over many years, incorporating Global Pooling layers that allow the network to perceive global information, nested bottleneck residual blocks that increase network depth while controlling parameter count, and auxiliary training objectives designed for board game characteristics (such as territory prediction, score prediction, etc.). Adopting this architecture means users can directly benefit from network design experience validated at scale by the community;
• 1000x speedup for the chess environment on TPU: This will dramatically accelerate self-play data generation. Traditionally, game environment logic runs on CPU and becomes the training pipeline bottleneck; compiling it into XLA computation graphs for direct execution on TPU eliminates CPU-TPU data transfer latency;
• CPU+TPU hybrid rewrite of MCTX: For evaluating positions, supporting search budgets up to 10,000, with approximately 5x speed improvement. This hybrid architecture executes tree search branching logic on CPU while neural network inference runs on TPU, leveraging the strengths of each.

If these improvements materialize, they will further lower the barrier to reproducing and experimenting with AlphaZero while improving the reliability of model quality evaluation.

Conclusion

nanoAlphaZero's story is, in a sense, a microcosm of open-source spirit and individual geek perseverance—an optimistic "one-month speedrun" goal that ultimately evolved into over two years of continuous refinement. It not only provides a runnable AlphaZero implementation with a browser-based local demo, but more importantly, it condenses what was once a complex system belonging to large research labs into a single file of code that anyone can read, modify, and run.

For any developer interested in reinforcement learning, MCTS, JAX, or TPU-accelerated training, this is a textbook-quality open-source case study worth studying.

Project: github.com/wtedw/nanoAlphaZero Online demo (runs locally in browser): nanoalphazero.wtedw.com