跳至主要内容

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