分散式系統與 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.
本文重點:
- AlphaGo Zero 的訓練架構:Self-play Workers 持續產生資料、Training Workers 同步訓練、Replay Buffer 儲存最近 50 萬局
- TPU 的關鍵作用:神經網路推理佔計算量 95%,TPU 的 MXU 單元每週期執行 16K 次乘加運算,是高效搜索的基礎
- 虛擬損失技術:讓多執行緒同時探索不同路徑,透過批次神經網路評估大幅提升並行 MCTS 效率
常見問題
什麼是虛擬損失(Virtual Loss),它如何讓 MCTS 並行化?
虛擬損失是一種技術,當一個執行緒正在探索某個節點時,臨時降低該節點的價值,讓其他執行緒選擇不同路徑。這樣可以同時收集多個待評估的葉節點,進行批次神經網路推理,大幅提升效率。探索完成後虛擬損失會被真實值覆蓋,不影響最終結果。
為什麼 AlphaGo 需要 TPU 而不是 GPU?
圍棋 AI 的計算瓶頸是神經網路推理(佔 95%),每步 MCTS 需要執行 1600 次推理。TPU 是 Google 專為矩陣運算設計的加速器,其 MXU 單元特別擅長這類計算,功耗更低、延遲更少。從 AlphaGo Lee 的 48 TPU 到 AlphaGo Zero 的 4 TPU,效率提升超過 10 倍。
AlphaGo Zero 的完整訓練成本大約是多少?
以 Google Cloud 定價估算,每天約 $3,700(4 TPU v2 Pod + 高記憶體 VM + 儲存空間),完整 40 天訓練約 $150,000。但這只是雲端租用成本,DeepMind 作為 Google 子公司可能有內部折扣。對比之下,培養一位職業棋手需要 10-15 年和數百萬的總成本。