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

作者 | 林偉 @阿里PAI

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

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

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

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

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

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

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

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

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

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

但是這種方式會(huì)使得一些習(xí)慣單機(jī)開(kāi)發(fā)的研究人員覺(jué)得比較晦澀。同時(shí)也是因?yàn)榉植际降脑?，我們做系統(tǒng)的很自然會(huì)把模型的開(kāi)發(fā)過(guò)程分成構(gòu)圖和執(zhí)行兩個(gè)階段。構(gòu)圖的時(shí)候只是生成一個(gè)邏輯執(zhí)行計(jì)劃，然后通過(guò)顯式方式的提交（或者execute）過(guò)程進(jìn)行執(zhí)行。

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

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

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

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

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

其實(shí)從設(shè)計(jì)理念來(lái)看，模型并行正是回歸到原來(lái)TensorFlow一開(kāi)始設(shè)計(jì)時(shí)候的設(shè)想，只是那個(gè)時(shí)候因?yàn)槟Ｐ筒⑿械男枨蟛粔?，沒(méi)有必要提供比較好高層自動(dòng)分布式的抽象，寫(xiě)模型的人還是可以自己精細(xì)化去構(gòu)造每個(gè)計(jì)算節(jié)點(diǎn)的子圖，整體上TensorFlow的框架只是負(fù)責(zé)把這些子圖能夠自動(dòng)通過(guò)Send-Recv進(jìn)行連接，并且在Runtime能夠合法的進(jìn)行計(jì)算。

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

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

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

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

本文參考資料