跳至主要内容

分散式系統與 TPU

AlphaGo 的成功不僅是演算法的勝利,也是工程的勝利。要在合理時間內訓練出超越人類的圍棋 AI,需要精心設計的分散式系統和專用硬體的支援。

本文將深入解析 AlphaGo 背後的系統架構,包括訓練流程、推理架構、並行 MCTS,以及 TPU 的關鍵角色。


訓練架構總覽

原版 AlphaGo 的訓練架構

原版 AlphaGo(擊敗李世乭的版本)的訓練分為多個階段,每個階段使用不同的資源配置:

AlphaGo Zero 的訓練架構

AlphaGo Zero 大幅簡化了訓練流程,使用單一的端到端訓練循環:

這個架構的優勢:

  1. 持續學習:Self-play 和 Training 同時進行,不需要等待
  2. 資源效率:所有資源都在做有用的工作
  3. 快速迭代:網路更新後立即用於產生新資料

自我對弈工作站(Self-play Workers)

任務分配

Self-play Workers 負責用當前最強的網路進行自我對弈,產生訓練資料。

配置AlphaGo Zero
Worker 數量數十個
每個 Worker1-4 TPU
每局 MCTS1600 次模擬
每天產生~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
每個 Worker4 TPU
Batch Size2048(每個 TPU 512)
訓練步數每天數萬步

分散式訓練

大規模訓練使用資料並行(Data Parallelism)

每個 TPU 處理不同的 mini-batch,計算出本地梯度,然後聚合更新全局參數。

同步 vs. 非同步更新

更新方式優點缺點
同步穩定、可重現Worker 需要等待最慢的
非同步吞吐量高梯度可能過時

AlphaGo Zero 使用同步更新,確保訓練的穩定性。


TPU 的角色

什麼是 TPU?

TPU(Tensor Processing Unit) 是 Google 專門為深度學習設計的加速器:

特性TPUGPUCPU
設計目標矩陣運算通用並行通用計算
精度FP16/BF16 優化FP32/FP16FP64/FP32
功耗相對低較高最高
延遲中等

TPU 的架構

TPU 的核心是 MXU(Matrix Multiply Unit)

MXU 每個週期可以執行 16K 次乘加運算,這對於神經網路的矩陣乘法至關重要。

為什麼 AlphaGo 需要 TPU?

圍棋 AI 的計算瓶頸在於神經網路推理

操作佔比
神經網路前向傳播~95%
MCTS 樹操作~4%
其他~1%

每一步 MCTS 需要執行 1600 次神經網路推理。TPU 的高吞吐量讓這成為可能。

TPU 使用的演進

版本訓練 TPU推理 TPU
AlphaGo Lee50 GPU48 TPU(v1)
AlphaGo Master4 TPU(v2)4 TPU(v2)
AlphaGo Zero4 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 Fan176 GPU
AlphaGo Lee48 TPU + 多台伺服器
AlphaGo Master4 TPU
AlphaGo Zero4 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 Pod4~$32~$3,000
高記憶體 VM數台~$5~$500
儲存空間10 TB~$0.02/GB~$200
網路-包含-

每天約 3,700,完整訓練(40天)約3,700**,完整訓練(40 天)約 **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 LeeAlphaGo ZeroAlphaZero
訓練 TPU50+ GPU → TPU4 TPU4 TPU
推理 TPU48 TPU4 TPU4 TPU
MCTS/步~100,000~1,600~800
訓練時間數月40 天數小時-數天

效率提升約 100 倍。

對開源社群的影響

AlphaGo 的架構啟發了多個開源項目:

項目特點
Leela Zero社群分散式訓練,複現 AlphaGo Zero
KataGo單 GPU 高效訓練,超越 AlphaGo Zero
ELF OpenGoFacebook 開源,使用 PyTorch
MinigoGoogle 開源,使用 TensorFlow

這些項目讓普通研究者也能訓練強大的圍棋 AI。


動畫對應

本文涉及的核心概念與動畫編號:

編號概念物理/數學對應
🎬 C9並行 MCTS多體問題
🎬 E9分散式訓練分散式計算
🎬 C5虛擬損失排斥勢
🎬 D15批次推理向量化計算

延伸閱讀


參考資料

  1. Silver, D., et al. (2017). "Mastering the game of Go without human knowledge." Nature, 550, 354-359.
  2. Jouppi, N., et al. (2017). "In-Datacenter Performance Analysis of a Tensor Processing Unit." ISCA 2017.
  3. Dean, J., et al. (2012). "Large Scale Distributed Deep Networks." NeurIPS 2012.
  4. Chaslot, G., et al. (2008). "Parallel Monte-Carlo Tree Search." CIG 2008.
  5. 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 年和數百萬的總成本。