曉查 發自 凹非寺量子位 報導 | 公眾號 QbitAI
對土豪來說最痛苦的是什麼，就是有一大堆硬體卻不能實現1+1=2的效果。
AI訓練中的並行計算就是如此，縱使你有一千張GPU，也無法實現單機訓練一千倍的效果。
最近，不差錢的谷歌，開源了一種SEED RL框架，可以更容易地將AI訓練放在幾千臺機器上運行，效果比之前的方法最高提升近4倍。
如果你也不差錢，在雲端上進行大規模並行計算，那麼可以節約80%的訓練成本。考慮到現在一個大型AI模型動輒上百萬的訓練費用，真的是非常可觀了。
在訓練AI玩足球遊戲這件事情上，SEED RL的處理速度可以達到每秒240萬幀。如果以60fps計算，相當於每秒可處理11小時的遊戲畫面。
SEED RL體系架構
上一代的分布式強化學習智能體IMPALA，其體系架構中包含Actor和Learner兩部分。
Actor通常在CPU上運行，並且在環境中採取的步驟與對模型進行推斷之間進行迭代，以預測下一個動作。
Actor經常會更新推理模型的參數，並且在收集到足夠數量的觀測結果後，會將觀測結果和動作的軌跡發送給Learner，從而對Learner進行優化。
在這種架構中，Learner使用來自數百臺機器上的分布式推理輸入在GPU上訓練模型。
但IMPALA存在著許多缺點：
1、使用CPU進行神經網絡推理，效率低下。而且隨著模型變大、運算量變大，問題會越來越嚴重。
2、Actor和Learner之間模型參數的帶寬成為性能的瓶頸。
3、資源利用效率低，Actor在環境和推理兩個任務之間交替進行，而這兩個任務的計算要求不同，很難在同一臺機器上充分利用資源。
SEED RL體系架構解決了以上這些缺點。Actor可以在GPU、TPU這類AI硬體加速器上完成推理，通過確保將模型參數和狀態保持在本地來加快推理速度，並避免數據傳輸瓶頸。
與IMPALA體系結構相反，SEED RL中的Actor僅在環境中執行操作。Learner在硬體加速器上使用來自多個Actor的成批數據來集中執行推理。
SEED RL在每個環境步驟將觀測結果發送給Learner的同時，使用gPRC框架的網絡庫，將延遲保持在較低水平。這使SEED RL在一臺機器上每秒最多可以實現一百萬個查詢。
Learner可以擴展到幾千個核心上，Actor的數量可以擴展到幾千臺機器，從而實現每秒百萬幀的訓練速度。
SEED RL用到了兩種最先進的算法：V-trace和R2D2。
V-trace負責從採樣的動作中預測動作的分布，R2D2負責根據動作的預測未來值選擇一個動作。
V-trace是基於策略梯度的方法，最早被IMPALA採用。由於Actor和Learner是異步執行，而V-trace在異步體系架構中的效果很好
第二種算法是R2D2，這是一種Q學習方法，DeepMind曾用這種算法將強化學習智能體在Atari遊戲上的水平提高了4倍，並在52款遊戲上超過了人類水平。
這種方法允許Q學習算法大規模硬體上運行的同時仍然可以使用RNN。
實驗結果
谷歌在DeepMind最近研究的開源足球遊戲項目Google Research Football中進行了基準測試。
使用64個Cloud TPU核心實現了每秒240萬幀的數據傳輸速度，與之前的最新分布式IMPALA相比，提高了80倍。
IMPALA要想達到相同的速度，需要14000個CPU，而SEED RL只用了4160個CPU。對於相同的速度，IMPALA需要的CPU是SEED RL的3~4倍。
通過對並行計算的硬體加速器進行優化，我們就可以放心大膽地提高模型的大小。
在上面的足球遊戲任務中，通過增加模型大小和輸入解析度，可以解決從前未能解決的一些困難，讓訓練模型的效率大幅提高。
傳送門
論文地址：https://arxiv.org/abs/1910.06591
GitHub地址：https://github.com/google-research/seed_rl