分布式系统与 TPU
AlphaGo 的成功不仅是算法的胜利，也是工程的胜利。要在合理时间内训练出超越人类的围棋 AI，需要精心设计的分布式系统和专用硬件的支援。
本文将深入解析 AlphaGo 背后的系统架构，包括训练流程、推理架构、并行 MCTS，以及 TPU 的关键角色。
训练架构总览
原版 AlphaGo 的训练架构
原版 AlphaGo（击败李世石的版本）的训练分为多个阶段，每个阶段使用不同的资源配置：
AlphaGo Zero 的训练架构
AlphaGo Zero 大幅简化了训练流程，使用单一的端到端训练循环：
这个架构的优势：
- 持续学习：Self-play 和 Training 同时进行，不需要等待
- 资源效率：所有资源都在做有用的工作
- 快速迭代：网络更新后立即用于产生新数据
自我对弈工作站（Self-play Workers）
任务分配
Self-play Workers 负责用当前最强的网络进行自我对弈，产生训练数据。
| 配置 | AlphaGo Zero |
|---|---|
| Worker 数量 | 数十个 |
| 每个 Worker | 1-4 TPU |
| 每局 MCTS | 1600 次模拟 |
| 每天产生 | ~100,000 局 |
工作流程
每个 Self-play Worker 的工作流程：
while True:
# 1. 下载最新的网络权重
network = download_latest_checkpoint()
# 2. 进行多局自我对弈
for game in range(batch_size):
positions = []
board = EmptyBoard()
while not board.is_terminal():
# 执行 MCTS
mcts = MCTS(network, board)
policy = mcts.search(num_simulations=1600)
# 选择落子
action = sample(policy)
# 记录
positions.append((board.state, policy))
# 落子
board = board.play(action)
# 3. 获取胜负结果
result = board.get_result()
# 4. 上传数据
upload_to_replay_buffer(positions, result)
负载平衡
多个 Worker 需要负载平衡：
- 网络同步：所有 Worker 使用相同版本的网络
- 数据平衡：确保不同 Worker 的数据都被使用
- 容错处理：单个 Worker 失败不影响整体训练
训练工作站（Training Workers）
任务分配
Training Workers 负责从 Replay Buffer 取样数据，训练神经网络。
| 配置 | AlphaGo Zero |
|---|---|
| Worker 数量 | 1-4 |
| 每个 Worker | 4 TPU |
| Batch Size | 2048（每个 TPU 512） |
| 训练步数 | 每天数万步 |
分布式训练
大规模训练使用数据并行（Data Parallelism）：
每个 TPU 处理不同的 mini-batch，计算出本地梯度，然后聚合更新全局参数。
同步 vs. 异步更新
| 更新方式 | 优点 | 缺点 |
|---|---|---|
| 同步 | 稳定、可重现 | Worker 需要等待最慢的 |
| 异步 | 吞吐量高 | 梯度可能过时 |
AlphaGo Zero 使用同步更新，确保训练的稳定性。
TPU 的角色
什么是 TPU？
TPU（Tensor Processing Unit） 是 Google 专门为深度学习设计的加速器：
| 特性 | TPU | GPU | CPU |
|---|---|---|---|
| 设计目标 | 矩阵运算 | 通用并行 | 通用计算 |
| 精度 | FP16/BF16 优化 | FP32/FP16 | FP64/FP32 |
| 功耗 | 相对低 | 较高 | 最高 |
| 延迟 | 低 | 中等 | 高 |
TPU 的架构
TPU 的核心是 MXU（Matrix Multiply Unit）：
MXU 每个周期可以执行 16K 次乘加运算，这对于神经网络的矩阵乘法至关重要。
为什么 AlphaGo 需要 TPU？
围棋 AI 的计算瓶颈在于神经网络推理：
| 操作 | 占比 |
|---|---|
| 神经网络前向传播 | ~95% |
| MCTS 树操作 | ~4% |
| 其他 | ~1% |
每一步 MCTS 需要执行 1600 次神经网络推理。TPU 的高吞吐量让这成为可能。
TPU 使用的演进
| 版本 | 训练 TPU | 推理 TPU |
|---|---|---|
| AlphaGo Lee | 50 GPU | 48 TPU（v1） |
| AlphaGo Master | 4 TPU（v2） | 4 TPU（v2） |
| AlphaGo Zero | 4 TPU（v2） | 4 TPU（v2）（可扩展） |
AlphaGo Zero 使用的 TPU 数量大幅减少，这要归功于更高效的架构和更新的 TPU 版本。
并行 MCTS 与虚拟损失
并行化的挑战
MCTS 的标准实现是串行的：
for i in range(num_simulations):
1. Selection：从根向下选择
2. Expansion：扩展叶节点
3. Evaluation：神经网络评估
4. Backup：回传更新
但神经网络评估是 GPU/TPU 友好的批次操作。如何让多个模拟同时进行？
叶节点并行（Leaf Parallelization）
最简单的并行方式：同时执行多个完整的模拟，最后合并结果。
问题：每个模拟都从根开始，会重复探索相同的路径。
虚拟损失（Virtual Loss）
DeepMind 采用了虚拟损失技术来实现树并行（Tree Parallelization）。
基本概念
当一个线程正在探索某个节点时，临时降低该节点的价值，让其他线程选择其他路径。
正常的 UCB：Q(s,a) + c * P(s,a) * sqrt(N(s)) / (1 + N(s,a))
加入虚拟损失后：
(Q(s,a) * N(s,a) - v * n_virtual) / (N(s,a) + n_virtual) + c * P(s,a) * sqrt(N(s)) / (1 + N(s,a) + n_virtual)
其中：
n_virtual是正在探索该节点的线程数v是虚拟损失的值（通常为 1 或胜率对应值）
运作流程
时间 T1：
Thread 1 选择路径 A → B → C
节点 C 获得虚拟损失 -1
时间 T2：
Thread 2 选择路径 A → B → D（因为 C 被「惩罚」了）
节点 D 获得虚拟损失 -1
时间 T3：
Thread 1 完成评估，更新 C 的实际值，移除虚拟损失
Thread 3 现在可能选择 C（如果实际值够好）
虚拟损失的效果
| 方面 | 效果 |
|---|---|
| 探索多样性 | 强制探索不同路径 |
| 批次效率 | 可以同时评估多个叶节点 |
| 收敛性 | 虚拟损失最终被真实值覆盖，不影响收敛 |
批次神经网络评估
通过虚拟损失，可以收集多个待评估的叶节点，进行批次推理：
TPU 的批次推理效率远高于逐个推理，这让并行 MCTS 成为可能。
推理架构
比赛时的配置
AlphaGo 在正式比赛中的推理架构：
| 版本 | 硬件配置 |
|---|---|
| AlphaGo Fan | 176 GPU |
| AlphaGo Lee | 48 TPU + 多台服务器 |
| AlphaGo Master | 4 TPU |
| AlphaGo Zero | 4 TPU（可扩展） |
分布式推理流程
比赛时的推理流程（以 AlphaGo Lee 为例）：
思考时间管理
AlphaGo 的时间管理策略：
| 局面 | 思考时间 | MCTS 次数 |
|---|---|---|
| 开局（有定式） | 较短 | ~10,000 |
| 中盘（复杂） | 较长 | ~100,000 |
| 简明局面 | 较短 | ~5,000 |
| 读秒 | 固定 | ~1,600 |
更多的 MCTS 模拟通常意味着更好的落子品质。
通讯与同步
数据格式
训练数据的传输格式：
message TrainingExample {
// 棋盘状态（17 × 19 × 19）
repeated float board_planes = 1;
// MCTS 搜索结果（362）
repeated float mcts_policy = 2;
// 胜负结果（1 = 当前方胜，-1 = 当前方负）
float game_result = 3;
}
网络带宽需求
| 数据流 | 大小 | 频率 |
|---|---|---|
| 训练样本 | ~10 KB/样本 | 每秒数千样本 |
| 网络权重 | ~200 MB | 每小时数次 |
| 控制消息 | < 1 KB | 持续 |
总带宽需求：~100 Mbps（内部网络足够）
故障处理
分布式系统的故障处理：
| 故障类型 | 处理方式 |
|---|---|
| Worker 挂掉 | 重启，继续使用最新 checkpoint |
| 网络断线 | 缓冲数据，重连后续传 |
| TPU 故障 | 自动切换到备用 TPU |
| 数据损坏 | 校验后丢弃，重新生成 |
成本分析
硬件成本估算
以 Google Cloud 的 TPU 定价估算 AlphaGo Zero 的训练成本：
| 资源 | 数量 | 单价/小时 | 总价/天 |
|---|---|---|---|
| TPU v2 Pod | 4 | ~$32 | ~$3,000 |
| 高内存 VM | 数台 | ~$5 | ~$500 |
| 存储空间 | 10 TB | ~$0.02/GB | ~$200 |
| 网络 | - | 包含 | - |
每天约 150,000。
注意：这是 2017 年的估算，DeepMind 作为 Google 子公司可能有内部折扣。
与人类训练的对比
| 方面 | AlphaGo Zero | 人类职业棋手 |
|---|---|---|
| 达到职业水平 | 2 天 | 10-15 年 |
| 训练成本 | ~$7,500 | 数百万（学费、生活费、机会成本） |
| 持续成本 | 电费 | 生活费 |
| 可复制性 | 完美复制 | 不可复制 |
当然，这个对比不完全公平——人类在学棋过程中学到的不只是围棋。
推理成本
正式比赛的推理成本：
| 配置 | 每局成本 |
|---|---|
| 48 TPU（AlphaGo Lee） | ~$500 |
| 4 TPU（AlphaGo Zero） | ~$50 |
| 单 GPU（KataGo） | ~$1 |
推理成本随着技术进步大幅下降。
技术演进
从 AlphaGo 到 AlphaZero
| 方面 | AlphaGo Lee | AlphaGo Zero | AlphaZero |
|---|---|---|---|
| 训练 TPU | 50+ GPU → TPU | 4 TPU | 4 TPU |
| 推理 TPU | 48 TPU | 4 TPU | 4 TPU |
| MCTS/步 | ~100,000 | ~1,600 | ~800 |
| 训练时间 | 数月 | 40 天 | 数小时-数天 |
效率提升约 100 倍。
对开源社区的影响
AlphaGo 的架构启发了多个开源项目：
| 项目 | 特点 |
|---|---|
| Leela Zero | 社区分布式训练，复现 AlphaGo Zero |
| KataGo | 单 GPU 高效训练，超越 AlphaGo Zero |
| ELF OpenGo | Facebook 开源，使用 PyTorch |
| Minigo | Google 开源，使用 TensorFlow |
这些项目让普通研究者也能训练强大的围棋 AI。
动画对应
本文涉及的核心概念与动画编号：
| 编号 | 概念 | 物理/数学对应 |
|---|---|---|
| 🎬 C9 | 并行 MCTS | 多体问题 |
| 🎬 E9 | 分布式训练 | 分布式计算 |
| 🎬 C5 | 虚拟损失 | 排斥势 |
| 🎬 D15 | 批次推理 | 向量化计算 |
延伸阅读
- 上一篇：从零训练的过程 — 训练曲线的详细分析
- 下一篇：AlphaGo 的遗产 — AlphaGo 对 AI 领域的深远影响
- 相关文章：MCTS 与神经网络的结合 — MCTS 的基础知识
参考资料
- Silver, D., et al. (2017). "Mastering the game of Go without human knowledge." Nature, 550, 354-359.
- Jouppi, N., et al. (2017). "In-Datacenter Performance Analysis of a Tensor Processing Unit." ISCA 2017.
- Dean, J., et al. (2012). "Large Scale Distributed Deep Networks." NeurIPS 2012.
- Chaslot, G., et al. (2008). "Parallel Monte-Carlo Tree Search." CIG 2008.
- Segal, R. (2010). "On the Scalability of Parallel UCT." CIG 2010.