如何評價Google即將發表在ICML17上的Placement Optimization的論文？

論文鏈接如下：[1706.04972] Device Placement Optimization with Reinforcement Learning

接要： 過去幾年，我們見證了神經網路的訓練和推理推動著計算能力需求的增長。目前解決這種需求的常用辦法是搭建異構分散式計算環境，其中包含 CPU 和 GPU 這樣的硬體設備。重要的是，神經網路模型在設備上的計算資源分配通常是由人類專家依靠簡單的探索和直覺做出的。在本文中，我們提出了一種為 TensorFlow 計算圖學習並優化設備任務分配的方法。其中最關鍵的思路是使用序列到序列模型來預測 TensorFlow 圖中哪些子集應該運行在哪些可用設備上，然後將預測分配的執行時間用作獎勵信號以優化序列到序列模型的參數。我們的結果在用於 ImageNet 分類的 Inception-V3 任務、RNN LSTM 語言建模和神經機器翻譯任務上的表現優於普通手動設備配置（device placement）方式。

因為個人興趣和工作關係，一直在關注Deep Learning的auto-placement相關的技術進展。5月份了解到這個工作Google一直在內部進行，之前還只能看到摘要，這兩天也終於在arXiv上看到paper正文了。

總體的感覺，論文的作法還是蠻fancy且principal的，通過DRL的方法來系統化地去為給定DL model探索placement的優化方案。不過僅僅從實際效果來看，當前談不上significant，直接附上論文里的實驗結果：

可以看到，在由多層sequential layer構成的NMT model上，相較於最優的hand-crafted placement strategy，最好的計算加速是23.5%（這裡的計算加速等價於收斂加速，因為文章里所propose的placement作法並沒有像weak-scaling或是Asynchronous演算法那樣引入對訓練收斂趨勢的影響），在典型的wide model Inception-V3上，相較於單卡最好的加速比也是將將20%的樣子。

數字上的提升看起來並不算顯著，不過考慮到這樣一個前提：文章里提出的RL-based的placement方法並沒有引入任何data-parallelism層面的並行加速（這篇文章里提出的是純粹的model parallelism），從這個角度上來說，能夠通過機器自動找到一種model-parallelism placement方法，可以在某些模型上獲得比單卡運行上更快的性能還是不錯的結果。不過，再考慮到實驗中是使用了2張卡甚至四張卡才達到這樣的加速效果，收益又打了一個折扣。因為加速效果比較margin的話，在單卡可以hold住完整訓練過程的時候，每張卡各啟動一個任務對用戶其實是建模效率更高的實驗策略。如果模型大到單卡hold不住（比如多層的NMT模型，以及GAN這種由兩個甚至更多模型組成的composite模型），那麼我會覺得文章中的placement方案會更具價值一些。

不過我還是會對這篇文章給出比較高的評價，原因是，auto-placement是一個本質難的問題，從演算法理論的角度來看，實際上可以說這是一個NP-hard的問題，通常對於這類問題，一個比較practical的作法是基於經驗或啟發式的方法去尋找近優解（這其實也是在EDA綜合工具領域經常使用的手段），而Google的同學直接選擇了一個新的思路，使用data-driven的方法來解決這個問題，在RL的框架上基於實際運行的反饋信號來制導學習，省去了標註樣本的環節，這個思路的價值其實高於這個思路目前取得的實際結果（Google IO 2017上發布的AutoML以及ICLR 17上發表的《Neural Combinatorial Optimization with Reinforcement Learning》也是這類工作的代表）。基於這個思路，其實我們可以在auto-placement里嘗試引入更多的啟發式規則，或是在model parallelism之外，再加入data parallelism的考慮，在增加求解探索空間的同時，也增加了所找到的優化placement方案的更顯著加速收益的可能性，實際上，這個工作提供了一個系列優化工作的框架性的思路，這是我個人最看重的。

auto-placement是一個很難的問題，也是一個蠻耗費計算力的問題，即便是這篇文章中這個相對克制的優化場景，也需要調用分散式的計算資源（ 個controller，每個掛接 到 個worker，一個parameter server，按照一個worker一張GPU卡來配比的話，起碼也得需要 到 張卡才可能跑得起這個優化訓練的過程，即便是對於一些深度應用Deep Learning的公司，這也幾乎已經是他們的全部GPU計算資源預算了，更不用說為了把這個訓練過程高效跑起來所需要的網路基礎設施和框架層面的優化支持），耗費 個小時左右才能找到so far看到的marginal的加速效果，按主流GPU卡的租賃成本合算下來大約是 (綜合了國內和國外的雲計算平台提供的商用GPU的租賃報價，6 RMB/小時算是一個合理的估價，甚至有些under-estimate了)，這至少是1W RMB的成本。花1W塊錢，找到一個20%左右的加速比placement方案，嗯，這買賣看起來還是土豪才付得起的感覺。如果我們再把問題泛化到多機（文章中沒有明確提到，不過從規模上來推敲，加上現在的主流GPU機型已經基本上是單機八卡，所以我suppose目前文章中的工作還只是考慮到了單機內的placement），再泛化到不僅僅做model parallelism，還可以做data-parallelism，那麼求解空間的複雜性會再上一個量級，求解過程自然會更為time consuming了。所以我個人認為，這篇文章里提出的作法，如果不做額外優化和剪枝約束的話，還處於research prototype的階段，還不具備productionize的可能。

這，也許就是接下來，工業界和學術界可以繼續開展的工作了（考慮到這個算力資源，以及對真實workload的考慮，也許工業界更適合主導這個工作的演進趨勢）。當然，也許G家一如既往，內部已經在開展這部分的工作甚至取得了新的進展，只是還沒有public出來罷了：）。

另外，如果僅僅是考慮到DL模型的訓練加速的話，除了文章里提到的RL-based的方法，其實還有很多更具ROI的技巧和方法，比如前不久FAIR發布的《Accurate, Large Minibatch SGD:Training ImageNet in 1 Hour》(如何評價Facebook Training ImageNet in 1 Hour這篇論文?)，以及語音訓練領域裡最早使用的Model Average的技術及其變種（我們團隊前段時間跟阿里翻譯團隊合作的項目里也基於MA的變種取得了還不錯的加速效果，可以參考這裡： 上帝的巴別塔在崩塌？阿里翻譯一年2500億次調用，節省25億美元 - 51CTO.COM）以及DL領域從高性能計算領域借鑒過來的優化技巧（比如HPC領域慣用的Ring Allreduce的通用優化手段）等等。從公司實際產出的角度來說，要腳踏實地，快速且有效地找到這些對公司來說更具ROI的策略來解決實際問題，從技術積累的角度來說 ，也需要舉望星空，對於業界、學術界新的方法不斷保持關注並keep curiosity，只有這樣才可能對技術的tipping point保持足夠的關注度和敏銳度。

這篇論文我看了差不多一天，一些粗淺的理解在這裡和大家分享一下，也作為一種記錄方式記錄下來。水平有限，多多指教~

使用了增強學習的方法，以及sequence to sequence的模型也就是LSTM。整體框架如下：

placement作為RL框架的輸入，Environment就是TensorFlow按照這種placement在實際系統中運行一下，其實際的運行時間runtime作為RL框架的輸出。該輸出和某個baseline進行計算得到Loss，根據得到的Loss優化改進placement再次輸入給RL系統。

其中，Loss是由該次的runtime和歷史所有平均runtime得到的，並且該歷史平均runtime在計算時忽略掉第一次的runtime，因為第一次runtime的值一般比較不靠譜，可以當做是outlier。

而placement是由一個sequence to sequence模型得到的，也即是一個LSTM模型。該模型的框架如下：

模型輸入：op的一些屬性。圖示中表明，該計算圖一共100個op, 因此會有100個op相關的屬性作為特徵輸入。op的屬性包括以下三類：

（1）一個embedding向量，根據op類型比如MatMul(矩陣乘)，conv2d(卷積)等，對這種特徵進行embedding得到一個向量，可以通過參數調整改變embedding之後向量長度，因此是一個可變長的（tunable）向量。論文中並沒有說明如何做embedding.

(2)一個定長的由output的shape拼接成的向量。拼接shape的時候按照output中shape維度最高的分配空間，例如，有兩個output，其中一個shape是4維[20，40，20，30]，另一個shape, 其中間2維[20, 30], 拼接在一起之後就是[20, 40, 20, 30, 0，20， 30，0]，也就是論文中所說的fixed-size zero-padded.

(3) 一個用來表示op的輸入輸入的one hot encoding 向量。這也很明顯是個定長的向量，統計計算圖中一共有多少個op，並分別編號。然後對某op的輸入輸出進行one hot encoding編碼，這個比較簡單，不用解釋。

模型輸出：計算圖中有多少個op，就有多少個輸出。圖示中就會有100個輸出。

RL系統得到Loss之後如何優化現有模型，公式如下：

實際上，論文中使用了均值的方法代替公式中的期望值E：

公式的意思就是：根據誤差改變參數theta, 再根據新的theta修改LSTM模型得到新的placement。其中R(Pi)表示某個placement的runtime，B表示歷史平均Runtime。對於一些特殊情況，例如某個策略使得大部分op放置在同一個設備上導致內存不足，那麼該次的runtime就設置為一個很大的常數，該常數一定是帶來巨大的誤差的。

以上就是整個論文中實驗的框架，如果要進行該實驗，需要準備n台機器，每台機器上有些設備，這些設備就是計算圖op要進行放置的候選。但是這麼做可能比較慢，因此論文中設置了分散式訓練該RL模型的方案，即分散式訓練。其架構如下：

圖中每個的worker有多個device，這些device負責執行controller發送來的placement policy。訓練過程就是：首先，每個Controller根據上面截圖的公式(3)生成K個placement，這裡的K是該controller控制的worker的個數，然後將這K個placement分別發給worker，每個worker一個placement, 每個worker就根據收到的placement設置op應該放在哪個設備上執行。其次，worker執行完之後，Controller收集對應的每個worker的runtime，每個worker的runtime也就是公式（3）的R(Pi)。這裡有個小trick，在計算歷史平均runtime時，第一個step的時間不計，並且是每10個step的runtime才計算一次R(Pi)。也就是RL系統每10個step才更新一次模型參數theta。

每個Controller的參數theta都存放在parameter server中，每次Controller根據公式（3）得到新梯度之後，都會和PS通信更新其參數，然後Controller再根據最新得到的參數做K次蒙特卡洛採樣得到新的K個placement policy發給K個worker。

實驗的結果，個人感覺並不明顯，只有20%左右的提升。

以上就是論文的細節。

另外還有一點我覺得對實驗影響較大，就是3.3節op的Co-locating那部分。由於有些模型的op個數比較多，並且由模型輸入部分的組成可以看出，真實的輸入是op個數乘以一個比較大的常數。這可能會導致模型訓練時梯度vanishing難以收斂或者過擬合，因此論文提出由用戶根據先驗將某些op劃分到一個group中，同一個group中的op放置在同一個設備上去。論文中給出了一個Co-locating前後的對比：

我之所以覺得這是個缺點，因為如果一個小白用戶做了一些不合理的cl-locating，這將直接導致RL訓練得不到好結果。主要是因為，用戶一旦設置了group，該group內的op不能被拆解，只能merger新的op進入該op。

最後是關於效果的好壞，我覺得這篇論文一個比較大的意義就是，提供了一個用機器學習方法解決placement問題的思路，效果雖然沒有想像的好但是也很有價值。還有就是關於並行策略的問題，我認為也是可以用這個方法學到對某個op做模型並行還是數據並行的，只需要在輸入時加入並行策略的特徵，輸出時輸出可能的並行策略就可以。