nanoAlphaZero：单文件JAX实现24小时训练2700分棋力

一个人的AlphaZero挑战

2017年，DeepMind的AlphaZero横空出世，仅凭自我对弈就在国际象棋、将棋和围棋上达到超越人类顶尖水平的能力，成为强化学习历史上的里程碑。AlphaZero的核心突破在于完全抛弃了人类棋谱数据，仅通过自我对弈（self-play）和强化学习从零开始学习博弈策略。在此之前，DeepMind的AlphaGo系列仍依赖大量人类专家棋局作为训练起点——2016年击败李世石的AlphaGo Lee使用了约16万局人类职业棋手对弈数据进行监督学习预训练，随后才通过强化学习自我提升。2017年的AlphaGo Zero首次证明了纯自我对弈的可行性（但仅限围棋），而AlphaZero则将这一方法统一推广到三种棋类，证明了通用算法的可行性——其训练过程结合了深度神经网络（用于评估局面和预测走法概率）与蒙特卡洛树搜索（用于在对弈中进行前瞻规划），两者相互增强形成正循环。具体而言，神经网络为MCTS提供先验概率和价值估计来引导搜索方向，而MCTS的搜索结果反过来作为更准确的训练目标来改进神经网络，这种「策略迭代」式的循环是整个系统不断自我提升的核心动力。

然而，复现AlphaZero一直是一件门槛极高的工程——它涉及分布式自我对弈、蒙特卡洛树搜索（MCTS）、大规模神经网络训练，以及复杂的工程编排。蒙特卡洛树搜索是一种基于随机采样的决策规划算法，其核心循环包含四个步骤：选择（沿树向下选择最有潜力的节点）、扩展（在叶节点处展开新节点）、模拟/评估（使用神经网络评估该局面的价值）、回溯（将评估结果沿路径向上更新）。在AlphaZero中，每次走棋前会执行数百次这样的循环来构建搜索树，最终根据各走法被访问的次数分布来选择实际着法。原始论文中AlphaZero使用了5000个TPU进行自我对弈、64个TPU进行训练，这一硬件需求使得学术界和个人开发者几乎无法复现。尽管此后出现了Leela Chess Zero（一个分布式志愿者计算项目）和KataGo等社区努力，但它们要么需要数千名志愿者贡献算力、积累数月训练，要么在代码复杂度上仍然相当高。

近日，一位开发者（GitHub 用户 wtedw）在 Reddit 上分享了他的开源项目 nanoAlphaZero：一个用 JAX 编写的、完整且单文件的 AlphaZero 实现。他最初的目标听起来颇具野心——「在一个月内，从零开始 speedrun（极速通关）出 GM（特级大师）级别的国际象棋 AI」。

结果如他自己坦言：「进展并不顺利。」这个项目最终耗费了超过两年时间，经历数次彻底重写，才终于稳定下来。

24小时训练出Elo 2700+棋力

尽管过程曲折，最终成果相当亮眼。据作者介绍，在一个 TPU v4-32 pod 上，nanoAlphaZero 可以在 24 小时内训练出 Elo 2700+ 的国际象棋模型。

TPU（Tensor Processing Unit）是Google专为机器学习工作负载设计的定制芯片，其架构与通用GPU有本质区别——TPU采用脉动阵列（Systolic Array）设计，专门优化大规模矩阵乘法运算，这正是神经网络推理和训练的核心计算模式。TPU v4单芯片峰值算力约275 TFLOPS（BF16精度），其中BF16（Brain Floating Point 16）是Google专门为深度学习设计的16位浮点格式，保留了与FP32相同的指数位宽（8位）以维持数值范围，同时将尾数缩减到7位以换取双倍吞吐量。TPU v4-32 pod表示由32个TPU v4芯片通过Google自研的ICI（Inter-Chip Interconnect）高速互联组成的计算单元，ICI提供每芯片数百GB/s的双向带宽，使得多芯片间的张量通信开销极低，总算力约8800 TFLOPS。虽然这对个人而言仍是不小的资源投入（按Google Cloud当前定价估算，24小时训练的成本大约在数百至数千美元量级），但相比原始AlphaZero使用数千块TPU的规模，已经缩小了两个数量级，且可通过Google Cloud按需租用获取。

作为参照，Elo 2700 已经稳稳处于人类国际象棋特级大师（GM）的水准线之上。Elo等级分是由匈牙利裔美国物理学家Arpad Elo在1960年代设计的竞技评分系统，最初用于美国国际象棋联盟，后被国际棋联（FIDE）采纳并推广到几乎所有竞技领域。其核心数学模型基于logistic分布：两位选手的等级分差可以转化为预期胜率——具体公式为 E_A = 1/(1+10^((R_B-R_A)/400))，因此分差400分意味着高分方约有91%的胜率，分差200分对应约76%的胜率。在国际象棋中，Elo 2500以上即为特级大师水平（获得GM称号还需满足特定锦标赛表现要求），2700以上为超级特级大师，全球仅约50人能维持在此水平。当前世界顶尖棋手如Magnus Carlsen的等级分大约在2830左右。值得注意的是，计算机国际象棋的Elo评分通常基于CCRL（Computer Chess Rating Lists）等专门的计算机对弈测试平台，其评分池与人类FIDE等级分不完全可比，但2700+的评分无论在哪个评分体系中都代表着极强的棋力。这意味着这套从零开始、无需任何人类棋谱、仅靠自我对弈的系统，确实达到了项目最初设定的「GM 级别」目标。

更有意思的是，其核心逻辑是**游戏无关（game-agnostic）**的。这意味着算法本身不包含任何特定游戏的领域知识——它只需要一个游戏环境提供合法动作列表、执行动作后的新状态、以及游戏是否结束和胜负结果，就能从零学习任何完全信息、双人零和博弈游戏的最优策略。除了国际象棋，它目前还支持：

• 井字棋（Tic-Tac-Toe）
• 四子棋（Connect Four）
• 六贯棋（Hex）——一种在六边形棋盘上双方尝试连接对边的策略游戏，因其分支因子适中且具有丰富的策略深度，常被用作AI算法的测试平台
• 小尺寸围棋棋盘

对于那些规模较小、可完全求解的游戏（如井字棋和小尺寸棋盘），nanoAlphaZero 能够学到完美策略（perfect play）——即在任何局面都能选择理论最优着法，在先手必胜的游戏中必定获胜，在理论和棋的游戏中永不失误。这也从侧面验证了其底层算法的正确性，因为完美策略是可以通过穷举验证的。更大棋盘的围棋支持仍在开发中。

整个强化学习管线就是一个JAX函数

nanoAlphaZero 最具工程美感的地方，在于它对复杂系统的极致简化。作者的核心设计理念是：整个强化学习管线，就是一个巨大的、经过 JIT 编译的 JAX 函数。

JAX是Google Brain团队（现Google DeepMind）开发的高性能数值计算库，可以理解为「可微分、可编译、可向量化的NumPy」。它建立在XLA（Accelerated Linear Algebra）编译器之上，后者最初是为TensorFlow开发的优化编译器，能够将高层数学运算融合、优化后生成针对特定硬件的高效机器代码。JAX的三大核心函数变换（function transformation）原语使其特别适合大规模强化学习：jit（即时编译）将Python函数编译为XLA计算图，消除Python解释器的逐行执行开销，直接在TPU/GPU上以编译后的原生代码高效执行，首次调用时有编译延迟但后续调用几乎无开销；vmap（向量化映射）是一种自动批量化变换，可以将针对单个样本编写的函数自动提升为处理整个批次的并行版本，编译器会将循环转化为硬件原生的向量运算，无需开发者手动编写批处理维度管理代码；pmap（并行映射）则将计算自动分发到多个设备上，实现数据并行，并在需要时自动插入跨设备的集合通信操作（如all-reduce）。这意味着开发者可以编写单局对弈的简洁逻辑，然后通过vmap自动扩展为数千局并行对弈，代码可读性和计算效率兼得。此外JAX的函数式编程范式要求所有函数都是纯函数（无副作用），状态通过显式参数传递——这一约束虽然需要适应，但使得整个计算图天然支持编译优化、自动微分和可复现性（相同输入必然产生相同输出）。

传统的 AlphaZero 实现往往需要管理线程、服务器和分布式工作节点，工程复杂度极高。例如在典型的分布式AlphaZero架构中，需要数百个自我对弈Worker进程不断生成对局数据，一个或多个Learner进程从Replay Buffer中采样训练更新网络权重，还需要参数服务器来同步最新模型到所有Worker，以及复杂的数据管道和心跳监控来处理节点故障——整体代码量动辄数万行。而 nanoAlphaZero 把所有这些都消化进了一个反复执行「自我对弈 + 模型更新」的 run_alphazero 函数中：

state = make_alphazero()

def run_alphazero(state):
    state, games = selfplay(state)  # 使用 Gumbel MuZero
    # 将进行中的对局移入 selfplay buffer
    # 将结束的对局移入 replay buffer
    state = train(state, replay_buffer.sample())
    return state

while True:
    state = run_alphazero(state)

作者强调：「没有线程，没有服务器，也没有需要管理的分布式工作节点。整个 RL 管线就是一个大的 JAX 函数。」

这种设计之所以可行，关键在于JAX的编译器能够将整个自我对弈+训练循环视为一个统一的计算图进行全局优化。当这个函数被jit编译后，XLA编译器可以跨越自我对弈和训练两个阶段进行算子融合、内存布局优化和通信调度，这是传统分布式架构中各组件独立运行所无法实现的。同时，由于TPU pod内部的ICI互联带宽极高，数据在设备间的搬移成本远低于通过网络在独立机器间传输，因此将整个管线放在单个pod上运行的效率可以超过分布式方案。

值得一提的是，自我对弈部分采用了 Gumbel MuZero 的搜索策略。Gumbel MuZero是DeepMind在2022年发表的论文「Policy improvement by planning with Gumbel」中提出的改进型树搜索算法，解决了传统MCTS在有限模拟次数下策略改进不稳定的问题。传统AlphaZero的MCTS使用PUCT（Predictor + Upper Confidence bounds applied to Trees）公式来平衡探索与利用，依赖大量模拟（通常800次以上）来获得可靠的走法分布——模拟次数不足时，访问次数分布的噪声很大，作为训练目标的质量就会下降。而Gumbel MuZero从一个全新的理论视角重新推导了树搜索：它将走法选择建模为从带Gumbel噪声扰动的对数概率中取最大值的过程（即Gumbel-Max技巧），这保证了在任何模拟次数下采样的走法分布都是对真实最优策略的无偏估计。具体实现上，它首先根据Gumbel噪声对所有走法进行初始排序，然后使用Sequential Halving（序贯减半）策略逐轮淘汰一半候选走法，将有限的模拟预算集中分配给最有希望的走法。这一特性对大规模并行训练至关重要——当需要同时运行数千局自我对弈时（这正是nanoAlphaZero使用vmap实现的），每局分配的搜索预算必须很小（例如仅16或32次模拟），此时Gumbel MuZero的效率优势就非常显著，它能在极低预算下仍然提供高质量的策略改进信号。

项目定位：让大规模AlphaZero实验更平易近人

作者对项目的定位相当清醒。他明确表示，nanoAlphaZero 的首要目标是让大规模 AlphaZero 实验变得更易上手，代码在追求速度与内存效率的同时，保持了紧凑（compact）和可魔改（hackable）的特性。

至于训练出强力模型，反而是次要目标——主要用来当作「健全性检查」（sanity check），验证底层逻辑没有问题。这种「工具优先、性能其次」的思路，与 Andrej Karpathy 的 nanoGPT 一脉相承。nanoGPT是前Tesla AI总监、OpenAI联合创始人Karpathy在2023年发布的教学项目，用约600行PyTorch代码实现了完整的GPT-2训练流程。它开创了AI领域「nano」项目的风格——用最少、最清晰的代码，把一个复杂算法的本质讲透，让研究者和爱好者能够快速理解、修改和实验。此后社区涌现了nanoGPT-speedrun（竞速优化训练时间）、nanoDiffusion等一系列类似项目，形成了一种以极简代码探索前沿算法的开源文化。

对于教学与研究而言，这类「nano」项目的价值往往超过工业级框架——它们剥离了工程噪音，直击算法核心。研究者可以在此基础上快速验证新想法，比如修改搜索策略（将Gumbel MuZero替换为其他变体）、调整网络架构（尝试Transformer替代ResNet）、尝试不同的训练调度方案（如修改学习率策略或replay buffer的采样权重），甚至探索将AlphaZero扩展到新的游戏类型（如不完全信息博弈或多人游戏），而不需要花费数周时间理解复杂的分布式系统代码。

v2版本：更快的棋类引擎正在路上

作者还预告了即将到来的 v2 大重构，其中包含几项颇具分量的改进：

• 切换到通用的 KataGo 神经网络架构：KataGo是由David Wu（前Jane Street量化研究员）独立开发的开源围棋AI，自2019年发布以来经过持续改进，被广泛认为是当前最强的开源围棋程序之一，其棋力已接近甚至在某些配置下超越早期版本的AlphaGo。其神经网络架构经过多年迭代优化，包含多项创新设计：全局池化（Global Pooling）层通过对整个棋盘特征图取均值和最大值，使网络能在任何层级感知全局信息（如整体形势判断），弥补了纯卷积网络局部感受野的限制；嵌套瓶颈残差块（Nested Bottleneck Residual Blocks）在标准残差连接内部再嵌套一层瓶颈结构，通过1x1卷积先降维再升维，在控制参数量和计算量的同时显著增加网络深度和表达能力；此外还有针对棋类游戏特性设计的辅助训练目标（如领地归属预测、最终分数分布预测、走法产权预测等），这些辅助任务通过多任务学习为网络提供更丰富的梯度信号，加速收敛并提升主任务性能。采用这一架构意味着用户可以直接受益于围棋社区数百万GPU/TPU小时验证过的网络设计经验；

• 国际象棋环境在 TPU 上提速 1000 倍：这将极大加速自我对弈数据的生成。传统上棋类环境的游戏逻辑（规则判断、合法着法生成、棋盘状态更新等）运行在CPU上，而神经网络推理在加速器上执行，两者之间需要频繁数据搬移，CPU成为整个训练管线的吞吐瓶颈。将棋类环境逻辑用JAX的纯函数重写后，整个环境可以被XLA编译为在TPU上直接执行的计算图，消除CPU-TPU之间的数据传输延迟（PCIe带宽通常仅为TPU HBM带宽的几十分之一），并且配合vmap可以在TPU上同时并行运行数万局游戏。这种「环境上加速器」的思路也被用于Brax（物理模拟）、Pgx（棋类游戏）等项目中；

• MCTX 的 CPU+TPU 混合重写：MCTX是DeepMind开源的JAX树搜索库，原始版本将整个MCTS过程编译在加速器上执行。但当搜索预算增大到数千甚至10,000次模拟时，树结构的动态扩展和复杂分支逻辑在TPU的固定计算图模式下效率下降。混合架构的核心思想是让树搜索的逻辑控制流（节点选择、树的扩展和回溯）在CPU上执行——CPU擅长处理条件分支和动态数据结构——而将神经网络前向推理（局面评估和策略预测）批量化后发送到TPU执行，充分发挥各自的架构优势。这种设计支持高达 10,000 的大搜索预算，运行速度约提升 5 倍，使得在评估阶段可以运行更深入的搜索来获得更准确的棋力测量结果。

这些改进如果落地，将进一步降低复现和实验 AlphaZero 的门槛，同时提升模型质量评估的可靠性。特别是环境加速和搜索预算的提升，将使得研究者能够在更短时间内完成从假设到验证的实验循环。

结语

nanoAlphaZero 的故事，某种程度上是开源精神与个人极客毅力的缩影——一个「一个月 speedrun」的乐观目标，最终演变成两年多的持续打磨。它不仅提供了一份可运行、可在浏览器本地演示的 AlphaZero 实现（浏览器端运行依赖WebAssembly和ONNX Runtime Web等技术，使得训练好的模型可以在用户本地设备上进行推理，无需服务器支持），更重要的是，它把原本属于大型实验室的复杂系统，浓缩成了一个人人都能读懂、修改和运行的单文件代码。

对于任何对强化学习、MCTS、JAX 或 TPU 加速训练感兴趣的开发者来说，这都是一份值得研读的教科书级开源案例。它展示了现代编译器技术（XLA）、硬件加速（TPU）和优雅的软件抽象（JAX的函数变换）如何共同作用，将曾经需要数百人团队和数千台机器的系统，转变为单人可以驾驭的项目。

项目地址：github.com/wtedw/nanoAlphaZero 在线演示（浏览器本地运行）：nanoalphazero.wtedw.com