【深度學(xué)習(xí)框架】Tensorflow是系統(tǒng)派，Pytorch是算法派

作者 | 林偉 @阿里PAI

https://zhuanlan.zhihu.com/p/375634204

大家好，我是老胡，周末愉快！今天分享一篇阿里林偉大佬的文章，關(guān)于深度學(xué)習(xí)框架的獨到見解，希望對你有幫助 

上次寫博客還是10多年前，然后還險些出現(xiàn)了事故，所以就一直沒有寫了。一晃就10幾年過去了，雖然有時候看著一些知乎大咖（前同事等）在消費我們原來做的一些事情，很想回應(yīng)幾句，但是最后也是不了了之。

最近團隊一位知乎大咖去創(chuàng)業(yè)了，所以為了團隊的建設(shè)，所以決定開始寫寫自己對于系統(tǒng)領(lǐng)域的一些感悟，也希望通過這個能夠讓大家更加理解我們團隊，以及分享我本人對于系統(tǒng)的一些理解。純屬自己的觀點，大家有不同觀點，歡迎指正和討論。

我一直工作在分布式系統(tǒng)的領(lǐng)域，從大數(shù)據(jù)到AI工程，其實有不少做系統(tǒng)的我們這行的人，很多都有這個路徑，所以第一篇就先聊聊自己一個感悟，說說最近做的深度學(xué)習(xí)框架，算是開一個頭。

我覺得做深度學(xué)習(xí)框架其實有兩個派別的人，一派是從分布式系統(tǒng)的人來做的，另外一派是從做算法的人來做的。不同的人的背景不同，所以做這個事情的角度也會不同，從而產(chǎn)生不同門派。tensorflow屬于系統(tǒng)派，而pytorch屬于算法派。

像我們這種做系統(tǒng)，特別是做過超大規(guī)模分布式系統(tǒng)的人，往往最擔(dān)心的就是要對一個已經(jīng)部署在成千上萬臺的計算集群上的平臺軟件需要做重大重構(gòu)，這個中間困難沒有做過這個事情的人可能不會太有體感，這么大一個平臺，公司不可能財力讓你能夠去鏡像一個集群去做任務(wù)的遷移，并且越大公司的平臺上用戶數(shù)眾多，業(yè)務(wù)都會耦合在一起去完成公司的使命，基本上你不可能有時間點可以讓全公司的業(yè)務(wù)團隊放下他們自己手頭的優(yōu)先級來配合你做這種遷移，哪怕你工程能力非常強，這種遷移中間不會出現(xiàn)任何意外。

何況很復(fù)雜系統(tǒng)要做到這一點基本上是很難的。所以我們做系統(tǒng)的，往往會把系統(tǒng)設(shè)計得更加具有可擴展性，從而盡最大可能去避免這種大的重構(gòu)和推倒重來。當(dāng)我們在面對需要構(gòu)建一個深度學(xué)習(xí)框架的時候，我們第一時間就在設(shè)想這個框架需要能夠從規(guī)模上很好支持分布式，能夠很好的擴展到任意大的深度模型的框架，我們希望構(gòu)建一個系統(tǒng)，能夠像人腦一樣能夠把視覺，語音，語言等多種模型能夠一同訓(xùn)練。

其實這個就是tensorflow這樣系統(tǒng)構(gòu)造的時候的原始想法，把整個計算構(gòu)成一個Tensor的Flow圖。因為分布式本身就很復(fù)雜，需要處理各種節(jié)點相互的數(shù)據(jù)和執(zhí)行中的各種依賴關(guān)系。這些事情由人來寫代碼，太繁瑣且容易出錯，所以自然地，我們就會設(shè)想由系統(tǒng)來負(fù)責(zé)這種依賴關(guān)系。這也就是為什么我們希望整個分布式執(zhí)行的計劃是一個靜態(tài)圖，然后系統(tǒng)再根據(jù)用戶指定的或者系統(tǒng)智能的決定的placement進行分圖，并在這些分圖中添加合適的Send-Recv的Op從而構(gòu)成一個分布式的執(zhí)行計劃。

但是這樣的設(shè)計理念也會帶來一些困惱，我們在模型訓(xùn)練時候有時候有些類似控制圖的部分，在這種設(shè)計理念下，我們必須要把這些控制流圖的代碼也op化，然后把這些op也整體串聯(lián)在Tensor的Flow執(zhí)行圖中，大家有興趣了解細(xì)節(jié)的話也可以看看我的老朋友也是前同事Yuan Yu寫的文章：Dynamic Control Flow in Large-Scale Machine Learning, Eurosys2018.^[1]

但是這種方式會使得一些習(xí)慣單機開發(fā)的研究人員覺得比較晦澀。同時也是因為分布式的原因，我們做系統(tǒng)的很自然會把模型的開發(fā)過程分成構(gòu)圖和執(zhí)行兩個階段。構(gòu)圖的時候只是生成一個邏輯執(zhí)行計劃，然后通過顯式方式的提交（或者execute）過程進行執(zhí)行。

這種方式讓研究人員覺得不能一邊寫代碼一邊就能夠馬上看到代碼片段的結(jié)果，所以這也造成很多人詬病TensorFlow的模式不太容易debug自己的模型程序，其實這也是因為分布式帶來負(fù)擔(dān)。但是也是因為TensorFlow是靜態(tài)圖的方式，其可以做到訓(xùn)推一體，在訓(xùn)練出來的模型能夠?qū)С瞿Ｐ蛨D，并且在這個圖上進行系統(tǒng)化的推理優(yōu)化從而能夠非常方便部署到線上。這個系統(tǒng)性化的方法帶來另外一個優(yōu)勢。

框架的另外一派是算法派，特別是感知類模型（圖像，語音，語言類）訓(xùn)練，因為這類訓(xùn)練一般都是同步訓(xùn)練，然后“分布式訓(xùn)練”也不像前者那樣設(shè)想是任意異構(gòu)的分布式執(zhí)行圖（即每個分布式節(jié)點的執(zhí)行邏輯可以不同），因為是數(shù)據(jù)并行，這樣我們就可以利用MPI的AllReduce的通訊源語來進行梯度的匯集計算。

算法同學(xué)需要是一種豐富的可擴展的在GPU上能夠很好運行的，并且能夠很好進行自動梯度的算子庫，并且因為面向是數(shù)據(jù)并行的場景，這樣話在神經(jīng)網(wǎng)絡(luò)部分其實都是單機程序，從而可以利用任何python的語法糖去構(gòu)建任何的動態(tài)的訓(xùn)練控制邏輯（大家也把這種稱作動態(tài)圖），對于算法研究人員來講，這種方式寫代碼比較隨性也方便debug，所以在研究界pytorch得到大量的關(guān)注和使用。

剛才說過TensorFlow從設(shè)計之初就在考慮可以超大的模型分布式訓(xùn)練的場景，但是沒有預(yù)想到硬件的發(fā)展也非常迅速，顯存越來越大以及訓(xùn)練技術(shù)的發(fā)展，還有非常精細(xì)化優(yōu)化顯存的工作，比如DeepSpeed等把optimizer所需要的顯存sharding化掉，使得除了超大規(guī)模稀疏模型訓(xùn)練外，感知類的SOTA模型一直可以利用數(shù)據(jù)并行的方式來進行訓(xùn)練。從而使得TensorFlow這種設(shè)計理念看上去有overdesign的嫌疑。并且就算超大規(guī)模稀疏模型訓(xùn)練，因為TensorFlow整體化的設(shè)計理念，不把Parameter Server作為游離在Flow圖之外，使得他在超大規(guī)模場景下的scalability上出現(xiàn)了問題，從而催生一堆自建PS+深度學(xué)習(xí)框架的（稀疏）模型訓(xùn)練框架。這是另外一個話題，我會在日后寫一寫在這個領(lǐng)域上我們一些工作。

好在隨著transformer的出現(xiàn)，我們終于有方法能夠回歸到最初那個夢想，使得我們可以把多種數(shù)據(jù)（視覺的，文字的）合在一起訓(xùn)練多模態(tài)的模型，因為問題規(guī)模的增大，必然需要更多參數(shù)的模型來支持，所以我們迅速將模型大小從幾十億增加到萬億規(guī)模，這個時候就必然需要能夠支持很好模型并行框架，這也是為什么最近這個領(lǐng)域重新變得火熱，比如類似OneFlow，MindSpore，PaddlePaddle，Mesh Tensorflow，GShard，以及我們阿里的Whale框架。

其實從設(shè)計理念來看，模型并行正是回歸到原來TensorFlow一開始設(shè)計時候的設(shè)想，只是那個時候因為模型并行的需求不夠，沒有必要提供比較好高層自動分布式的抽象，寫模型的人還是可以自己精細(xì)化去構(gòu)造每個計算節(jié)點的子圖，整體上TensorFlow的框架只是負(fù)責(zé)把這些子圖能夠自動通過Send-Recv進行連接，并且在Runtime能夠合法的進行計算。

而現(xiàn)在，因為需求增多，算法迭代需求的增多，迫切需要一種高層次的自動分布式框架，從而使得算法同學(xué)能夠去快速簡單構(gòu)造一個邏輯圖的方式去構(gòu)造自己神經(jīng)網(wǎng)絡(luò)，而由系統(tǒng)層來幫助他來進行復(fù)雜模型并行的構(gòu)成。所以其實可以看到TensorFlow的設(shè)計理念正好就是為這個考慮的，利用靜態(tài)圖，我們可以邏輯性去描述一個網(wǎng)絡(luò)訓(xùn)練，然后在執(zhí)行時候在進行系統(tǒng)化的分圖和分布式訓(xùn)練。所以說自動分布式的需求并沒有超越原來設(shè)計的基本范疇，也是因為這樣，我們采取和谷歌Gshard類似技術(shù)路線去提供自動分布式的能力。正是站在原有框架基礎(chǔ)上去做增量。

不同于GShard更加關(guān)注于谷歌TPU集群，我們關(guān)注于異構(gòu)的GPU集群，這里所說異構(gòu)是因為我們不如谷歌這么有錢，構(gòu)建非常大的同構(gòu)化TPU集群，我們集群中有不同年代的GPU和CPU，這些GPU各自算力和顯存都大小不一。也正是因為這樣，其實給我們系統(tǒng)提出更大挑戰(zhàn)，我們在進行自動分布式時候需要在cost model上考慮好這些差異點。這樣才能做到比較優(yōu)化的分布式訓(xùn)練。這也是我們自動分布式框架Whale一種差異性和核心能力之一。

其實系統(tǒng)派的框架和算法派的框架也在進行一定的融合，TensorFlow提出了Eager模式，通過TF.function在eager模式下可能單步執(zhí)行計算，得到Tensor來提高可調(diào)式性；而Pytorch通過Trace或者Parse的方式轉(zhuǎn)化為TorchScript的圖描述，從而能夠更好支持訓(xùn)練到推理的工程化。但是這種動靜結(jié)合其實只是在一定層次的，比如如果考慮分布式，Trace的方式去得到TorchScript就不足夠。需要進一步去限制構(gòu)圖能夠使用的API，這也是像NV的megatron以及微軟DeepSpeed在Pytorch上去支持分布式所帶來的一些約束，感興趣的可以讀讀OneFlow的Blog：《GPT-3難以復(fù)現(xiàn)，為什么說PyTorch走上了一條“大彎路”？”》

總結(jié)下我的觀點：我覺得現(xiàn)在深度學(xué)習(xí)框架主要流行的兩個TensorFlow和Pytorch是有其設(shè)計理念的原因的。我們做Whale正是在這種理解的基礎(chǔ)上進行路線選擇，并且認(rèn)為應(yīng)該站在已有的工作基礎(chǔ)上去做增量的東西。而不是再去造一個別人做過的輪子。接下來我還陸續(xù)展開我們分布式框架Whale，大規(guī)模稀疏模型訓(xùn)練工作，編譯系統(tǒng)Ansor和DISC，以及我們?nèi)绾伟逊植际剑幾g和調(diào)度有機結(jié)合方面一系列系統(tǒng)工作，敬請關(guān)注。如果大家對于我們PAI團隊的工作有興趣，非常歡迎和我們聯(lián)系，我的郵箱是[email protected]

本文參考資料