分散式系統與 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.