跳至主要内容

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