从我开发过的Tensorflow、飞桨、无量框架看深度学习这几年
文 | Peter潘欣@知乎
和深度學習框架打交道已有多年時間。從Google的TensorFlow, 到百度的PaddlePaddle，再到現在騰訊的無量。很慶幸在AI技術爆發的這些年橫跨中美幾家公司，站在一個比較好的視角看著世界發生巨大的變化。在這些經歷中，視角在不斷切換，從最早的算法研究，到后來的框架開發，到機器學習平臺和更多基礎架構，每一段都有不同的感受和更深的領悟。
清明節這幾天有些時間寫了這篇文章，從我的視角，用幾個深度學習框架串起來這些年歷史上的一些有趣的插曲，和技術背后的一些故事，免得寶貴的記憶隨著時間在腦中淡去。
TensorFlow
入門
故事開始在2015年底，我結束了在Google Core Storage和Knowledge Engine的工作，加入了Google Brain，在Samy Bengio下擔任一名Research Software Engineer，簡稱RSWE。RSWE角色的產生主要是因為Google Brain和DeepMind發現Research Scientist很難在研究中解決復雜的工程問題，并最終技術落地。因此需要卷入一些工程能力比較強的Engineer和Scientist一起工作。而我比較“幸運”的成為Google Brain第一個RSWE。
加入新組的前一個周末，非常興奮的提前探訪了Google Brain的辦公地點。想到能近距離在Jeff Dean旁邊工作還是有些小激動，畢竟是讀著MapReduce, BigTable, Spanner那些論文一路成長起來的。辦公場所沒有特別，Jeff和大家一樣坐在一起，比較意外的是發現我工位旁幾米的辦公室門牌上寫著谷歌創始人Larry Page&Sergey Brin，辦公室被許多獎杯，證書，太空服之類的雜物包圍著。看來公司對于AI技術的重視程度真實非常的高。
言歸正傳，早期的TensorFlow比較缺模型示例，相關API文檔還不太規范，于是先開始給TensorFlow搭建模型庫。我花了一年時間把Speech Recognition, Language Model, Text Summarization, Image Classification, Object Detection, Segmentation, Differential Privacy, Frame Prediction等模型寫了一遍，后來成為TensorFlow github上model zoo的雛形。那年還是個到處都是低垂果實的時候，沒有GPT3這種極其燒錢的大模型，只要對模型做一些小的調整，擴大模型的規模，就能刷新State-Of-The-Art。Bengio大佬經常在世界各地云游，偶爾回來后的1v1還是能給我不少的指引。印象深刻的第一次面聊，這位寫過幾百篇論文的一字眉大神給剛剛入門的我在白板上手推了gradient descent的一些公式。另外一次1v1，他發給我一本Ian Goodfellow寫的書（當時還是草稿pdf），然后我每天晚上就躺在茶水間的沙發上一邊做實驗一邊讀書。
那年還發生了件有意思的插曲，AlphaGo大戰人類棋手。DeepMind和Brain有非常緊密的合作關系，組里組織了一輪paper reading，仔細研讀了相關的paper，然后大家帶上啤酒和零食組織了觀戰活動，感覺就有點像是在看球賽。那次學會了兩件事，強化學習算法，還有圍棋的英文是Go。
16年是TensorFlow高速發展的一年。Jeff的演講里經常有TensorFlow代碼被引用次數指數級暴漲的圖。但是16年也是TensorFlow被噴的比較慘的一年。TF的Operator粒度是非常細的，據說這是從內部上一代框架DistBelif上吸取的教訓。細粒度的Operator可以通過組合形成各種高層的Layer，具有更好的靈活性和擴展性。然而，對于性能和易用性來說卻是比較嚴重的問題，一個模型隨便就有幾千個甚至跟多的Operator。舉幾個例子：
當時我要實現第一個基于TensorFlow的ResNet，光為了寫一個BatchNormalization（查了好幾個內部版本竟然都有些問題），需要通過5～10個細粒度的算子通過加減乘除的方式組裝起，1001層的ResNet有非常多的BatchNorm，整個ResNet有成千上萬個Operator，可想而知，性能也不怎么樣。不久后我朋友Yao搞了個Fused BatchNormalization，據說能讓整個ResNet提速好幾成。
BatchNormalization只是初級難度，做Speech Recognition時為了在Python層用TensorFlow完成BeamSearch也花了不少功夫。當時寫了個End-to-End的模型，用的是Seq2Seq with Attention，能夠一個模型直接把聲音轉成文字。為了把搜索生產線上語音識別的數據訓到最后收斂，用128個GPU整整花了2個月的時間。每天早上上班第一件事就是打開TensorBoard，放大后才能看到Loss又下降了那么一點點。
16年時TensorFlow訓練模式主要是基于Jeff等幾位的Paper，基于參數服務器的異步訓練是主流。訓練速度線性擴展性不錯，但是今天基于ring的同步訓練在NLP，CV這些領域的聲音更響一些。記得第一次和Jeff單獨交流是關于Speech Recognition分布式訓練的實驗情況，加到128個GPU做異步訓練基本能保證線性擴展，但是基于SyncOptimizer的同步訓練速度會慢很多。當時Jeff問了下收斂效果有沒有收到影響，我懵了一下，說沒有仔細分析過，趕緊回去查一下。順便八卦一下，Jeff真是非常瘦，握手的時候感覺幾乎就剩皮包骨了。
開發過一些模型后發現算法研究員其實還有不少痛點。1. 不知道怎么Profile模型。2. 不知道怎么優化性能。為了解決這兩個問題，我抽空寫了個tf.profiler。tf.profiler的原理比較簡單，就是把Graph, RunMeta和一些其他的產物做一些分析，然后用戶可以通過CLI，UI或者python API快速的去分析模型的結構，Parameter, FLOPs, Device Placement, Runtime等屬性。另外還做了個內部數據的抓取任務，去抓算法研究員的訓練任務的metrics，如果發現GPU利用率異常，網絡通行量過大，數據IO慢時會自動發郵件提醒，并給出一些修改的建議。
讓一個專心搞算法研究的人寫一個白板的數學公式不難，但是讓他去搞明白復雜的任務配置，分布式系統里的性能、資源、帶寬問題確是件很困難的事。無論多么牛的研究員都會問為什么任務沒能跑起來，是資源不夠還是配置不對。記得有天傍晚，人不多，Geoffrey Hinton大神突然走過來問到Can you do me a favor？My job cannot start...（正當我準備答應時，Quoc Lee已經搶先接單了，真是個精神的越南小哥。。。）
Moonshots
Google Brain每年會組織一次Moonshots提案，許多后來比較成功的項目都是這樣孵化出來的，比如AutoML，Neural Machine Translation等等。團隊成員會提出一些當時技術比較難達到的項目，大家組成類似興趣小組的形式投入到這些項目中。
現在火的一塌糊涂的AutoML有點因為商業化或者其他原因，感覺已經對原始的定義做了極大的拓展。當時Brain孵化這個項目的時候有兩隊人在做LearningToLearn的項目，一個小隊希望通過遺傳算法來搜索更優的模型結構，另一個小隊則決定使用強化學習算法搜索。遺傳算法小隊在使用資源時比較謹慎，通常只使用幾百個GPU。而另一個小隊則使用了幾千個GPU。最后強化學習小隊更早的做出了成果，也就是Neural Architecuture Search。而另一個小隊雖然后來使用更多的GPU也達到了類似的效果，但是要晚了不少。
一個比較有趣的插曲是Brain雖然很早就有幾萬張GPU，但是每當論文截稿的前一段時間總是不夠用，其中搞NAS的同學常常在郵件中被暗示。為了解決資源的分配問題，領導們被卷入了一個非常長的email，后來大概解決方案是每個人會被分配少量的高優先級GPU和適量的競爭級GPU資源。而NAS的同學因為已經完成了資本的原始積累成為了一個很火的項目，得到了特批的獨立資源池。為了支持這個策略，我又開發了個小工具，現在回頭想想還挺吃力不討好的。
Pytorch
動態圖
快速成長的時間總是過得很快，Megan加入Brain后，我被安排向她匯報，當時的RSWE團隊已經有十幾人，而Google Brain也從幾十人變成了幾百人。
2017年初，經Megan介紹，TensorFlow團隊一位資深專家Yuan Yu找到我，問有沒有關注Pytorch，約我調研后一塊聊聊。于是我就去網上搜集了一下Pytorch的資料，又試用了一下。作為一個TensorFlow的深度用戶，我的第一反應就是Pytorch解決了TensorFlow很大的痛點，用起來非常的“自然”。
和Yuan聊完后，我們快速的決定在TensorFlow上也嘗試支持類似Pytorch的imperative programming用法。Demo的開發過程還算比較順利，我大概花了一個多月的時間。記得當時我把項目命名為iTensorFlow, short for imperative TensorFlow。（后來被改名成eager，感覺好奇怪）。
Demo的設計思路其實也不復雜：1. TensorFlow graph可以被切分成任意粒度的Subgraph，可以通過函數調用的語法直接執行，2. TensorFlow對用戶透明的記錄執行過程以用于反向梯度計算。給用戶的感覺就就類似Python native的運行。
進而產生幾個推導：1. 當Subgraph的粒度是operator時，基本等價于Pytorch。2. 當Subgraph粒度由多個operator組成時，保留了graph-level optimization的能力，可以編譯優化。
最后再埋個伏筆：1. tf.Estimator可以自動的去融合Subgraph，形成更大的Subgraph。用戶在開發階段基于imperative operator-level Subgraph可以簡單的調試。用戶在部署階段，可以自動融合大的Subgraph，形成更大的optimization space。
做完之后，我非常興奮的和Yuan演示成果。Yuan也說要幫我在TensorFlow里面推這個方案。當時Pytorch的成長速度非常的快，TensorFlow的Director也召集了多名專家級的工程師同時進行方案的探索。當時我還沒能進入TensorFlow的決策層，最終得到的結論是1. 讓我們成立一個虛擬組專門做這個項目。2. 之前的Demo全部推倒重新做，TensorFlow 2.0作為最重要Feature 發布，默認使用Imperative Mode (后改名叫Eager Mode，中文常常叫動態圖)。我則作為團隊的一員在項目中貢獻來一些代碼。
后來Brain來了位新的大神，Chris Lattner，在編程語言和編譯領域研究的同學估計很多認識他。他提出來希望用Swift來實現Deep Learning Model的Progamming，也就是后來的Swift for TensorFlow。理由大概是Python是個動態的語言，很難靜態編譯優化。后來我和他深入討論來幾次，從技術上非常贊同他的觀點，但是也明確的表示Swift for TensorFlow是一條很難走的路。用Python并不是因為Python語言多么好，而是因為很多人用Python。和Chris的一些交流中我對編譯過程中的IR和Pass有了更深的了解，對后來在PaddlePaddle中的一些工作產生了不少的影響。
一個插曲是某位TensorFlow團隊的資深專家有次悄悄和我說：Python is such a bad language。這句話我品味了好久，不過和他一樣沒有勇氣大聲說出來。。。
當時動態圖的項目還延展出兩個比較有趣的項目。有兩個其他團隊的哥們想對Python做語法分析，進而編譯control flow。我很委婉的表示這個方案做成通用解決方案的可能性不太大，但是這個項目依然被很執著的做了很長一段時間，并且進行了開源，但是這個項目也就慢慢壽終正寢了。另一個很酷的項目是完全用numpy來構造一個deep learning model。通過隱式的tape來完成自動的求導，后來項目好像逐漸演化成來JAX項目。
API
后面我逐漸轉到了TensorFlow做開發。記得2017年還發生了一件印象深刻的事情，當TensorFlow收獲海量用戶時，網上一篇“TensorFlow Sucks"火了。雖然那篇文章很多觀點我不能茍同，許多想法比較膚淺。但是，有一點不能否認，TensorFlow API是比較讓人蛋疼的。1. 同一個功能往往幾套重復的API支持。2. API經常變動，而且經常發生不向后兼容的問題。3. API的易用性不高。
為什么會發生這個問題呢？可能要從Google這個公司的工程師文化說起。Google是非常鼓勵自由創新和跨團隊貢獻的。經常會有人給另一個團隊貢獻代碼，并以此作為有影響力的論據參與晉升。所以在早期TensorFlow還不是特別完善的時候，經常有外部的團隊給TensorFlow貢獻代碼，其中就包含了API。另外，在Google內部的統一代碼倉庫下，放出去的API是可以很容易的升級修改的，很多時候只需要grep和replace一下就行。但是github上放出去的API完全不一樣，Google的員工不能去修改百度，阿里，騰訊內部的TensorFlow使用代碼。對此TensorFlow團隊早期的確沒有非常有效的方案，后來才出現了API Committee對public API做統一的把關和規劃。
在我做視覺的時候，和Google內部一個視覺團隊有過很多合作，其中一個是slim API。這個視覺團隊非常的強，當年還拿了CoCo的冠軍。隨著他們模型的推廣流行，他們的tf.slim API也被廣為流傳。slim API的arg_scope使用了python context manager的特性。熟悉早期TensorFlow的人知道還有tf.variable_scope, tf.name_scope, tf.op_name_scope等等。with xxx_scope一層套一層，復雜的時候代碼幾乎沒有什么可讀性。另外就是各種global collection，什么global variable, trainable variable, local variable。這在傳統的編程語言課里，全局變量這種東西可能是拿來當反面教材的。然而，算法人員的視角是不一樣的，with xxx_scope和global collection能減少他們的代碼量。雖然我們知道合理的程序設計方法也可以做到，但是算法專家估計需要把時間用來讀paper，不太愿意研究這些程序設計的問題。
記得在早期內部還有兩個流派的爭論：面向對象和面向過程的API設計。
基于我教育歷史的洗腦，感覺這個是不需要爭論的問題。Keras的Layer class和Pytorch的Module class這些面向對象的接口設計無疑是非常優雅的。然而，其實當時的確發生了非常激烈的爭論。一些functional API的作者認為functional的調用非常節省代碼量：一個函數就可以解決的問題，為什么需要先構造一個對象，然后再call一下？
在TensorFlow動態圖能力開發的早期，我們也反復討論了2.0里面接口的設計方案。作為炮灰的我又接下了寫Demo的工作。
閉關兩周后，我給出了一個方案：1. 復用Keras的Layer接口。2. 但是不復用Keras的Network，Topology等其他更高層的復雜接口。
原因主要又兩點：1. Layer是非常簡潔優雅的，Layer可以套Layer，整個網絡就是一個大Layer。Layer抽象成construction和execution兩階段也非常自然。2. Keras有很多歷史上為了極簡設計的高層接口。我個人經驗覺得很難滿足用戶靈活的需求，并不需要官方提供。而且這樣可能會導致TensorFlow API層過度復雜。
后來方案被采納了一半，大佬們希望能夠更多的復用Keras接口。其實沒有完美的API，只有最適合某類人群的API。有個小插曲，當時Keras的作者Fran?ois也在Google Brain。為了在TensorFlow 2.0的動態圖和靜態圖同時使用Keras的接口，不得不在Keras API內做很多改造。通常Fran?ois在Review代碼時都非常的不情愿，但是最后又往往妥協。很多時候，特別是技術方面，真相可能在少數不被大多數人理解的人手上，需要時間來發現。
TPU
感覺互聯網公司那幾年，真正把AI芯片做得成熟且廣泛可用的，只有Google一家。TPU一直都是Google Brain和TensorFlow團隊關注的重點。原因可能是Jeff老是提起這件事，甚至一度在TensorFlow搞GPU優化是件很沒前途的事情。
TPU有個比較特別的地方，在于bfloat16的類型。如今bfloat16，還有英偉達最新GPU上的TF32都已經被廣為了解了。在當時還是個不小的創新。
bfloat16的原理非常簡單，就是把float32的后16bit全部截掉。和IEEE的float16相比，bfloat16的mantissa bits會少一些，但是exponential bits會多一些。保留更多的exponential bits有利于gradients很接近0時不會消失，保證bfloat16訓練時能夠更好的保留模型的效果。而傳統基于float16訓練時，往往需要做loss scaling等調試才能達到類似的效果。因此bfloat16能讓AI芯片更快的運算，同時又確保收斂效果通常不會有損失。
為了在TensorFlow上全面的支持bfloat16，我當時花了不少的功夫。雖然之前有基于bfloat16通信的方案，但是要在所有地方都無縫打通bfloat16，還有非常多的工作要做。比如eigen和numpy都不支持bfloat16這種特殊的東西。幸好他們都可以擴展數據類型（就是文檔太少了）。然后還要修復成百上千個fail掉的unit tests來證明bfloat16可以在python層完備的使用。
TPU是一個非常高難度，跨團隊，跨技術棧的復雜工程。據說Google有位非常優秀的工程師，為了在TPU上支持depthwise convolution一個TPU kernel，花掉了半年的時間。
其實這一點也不夸張，除了底層的硬件設計，單是將tensorflow graph編譯成硬件binary的XLA項目早期就至少卷入幾十人。從HLO到底層的target-specific code generation，幾乎又重寫了一遍TensorFlow C++層，遠比之前的解釋型執行器復雜。
TPU的訓練在底層跑通后，我基于底層接口的基礎上完成Python層的支撐API，然后去實現幾個模型。當時碰到了好幾個難題，有些在幾周時間內解決了，有些持續到我不再團隊后好些年。這里舉幾個例子。
當時一個TPU Pod（好像是512個chips）算得太快了，即使是很復雜模型的計算也會卡在數據的IO和預處理上。后來搞了個分布式的data processing，通過多個CPU機器來同時去處理數據，才能喂飽TPU。
早器的TPU API易用性比較弱。通常一個model需要在TPU上train幾百步然后再返回python層，否則TPU的性能會飛快的退化。這對于算法人員是很不友好的，這意味著debug能力的缺失，以及大量復雜模型無法實現。記得當年OKR被迫降低為支持常見的CV模型。
TPU如何支持動態圖。記得我當時迫于TPU的約束，做了個所謂的JIT的能力。就是Estimator先在CPU或者GPU上迭代N步，完成模型的初步調試，然后再自動的deploy到TPU上。從算法人員角度，既滿足單步調試的能力，又能在主要training過程用上TPU。
團隊
Google Brain是個很神奇的團隊，比較不客氣的說，在2015年后的幾年間包攬了全世界在深度學習領域一半以上的關鍵技術突破，比如TPU，TensorFlow, Transformer, BERT, Neural Machine Translation, Inception, Neural Architecture Search, GAN，Adverserial Training, Bidrectional RNN等等。這里不只有深度學習領域的圖靈獎獲得者，還有編程語言、編譯器、計算機體系結構、分布式系統的頂級專家，甚至還有生物，物理學專家。Jeff將這些人放在一起后，發生了神奇的化學反應，加快了技術改變世界的步伐。
PaddlePaddle 飛槳
Paddle其實誕生時間比較早，據說是大約13～14年的時候徐偉老師的作品。后來據說Andrew Ng覺得Paddle叫一次不過癮，就改名成了PaddlePaddle。Paddle和那個年代的框架Caffe有類似的問題，靈活性不夠。很多地方用C++寫成比較粗粒度的Layer，無法通過Python等簡單的編程語言完成模型的快速構造。
后來17年下半年，團隊開始完全從新寫一個框架，但是繼承了Paddle的名字。2017年底的時候，Paddle國內的團隊找到了我，邀請我擔任Paddle國內研發團隊的負責人。抱著打造國產第一框架的理想，我接受了邀請，一個月后就在北京入職了。
早期設計
加入團隊的時候，新的Paddle還是一個比較早期的原型系統，里面有一些設計已經被開發了出來。我發現其中有些設計理念和TensorFlow有明顯的差異，但是實現的時候卻又模仿了TensorFlow。
仿編程語言
設計者希望設計一種編程語言來完成深度學習模型的構建（有點類似Julia等把深度學習模型的特性嵌入到了編程語言中）。然而在實現上，我發現其實和TensorFlow比較類似。都是通過Python去聲明一個靜態模型結構，然后把模型結構交給執行器進行解釋執行。并沒有發明一種新的深度學習編程語言。
這塊我基本沒有對設計進行調整。本質上和TensorFlow早期靜態圖的沒有區別。但是在細節上，TF基于Graph的模型可以通過feed/fetch選擇性的執行任意一部分子圖，更加靈活。Paddle中與Graph對應的是Program。Program就像正常程序一樣，只能從頭到尾完整的執行，無法選擇性的執行。因此Paddle在這塊相對簡化了一些，但是可以通過在Python層構造多個Program的方式補全這部分靈活性的缺失，總體來說表達能力是足夠的。
Transpiler
Transpiler是對Program進行直接改寫，進而可以讓模型能夠被分布式運行，或者進行優化。初衷是比較好的，可以降低算法人員的使用難度。然而在實現上，最開始是在Python層直接對Program結構進行改寫。后來我從新設計了IR+Pass的Compiler體系，通過一種更系統性的方式做了實現。
LoDTensor
可能是因為團隊的NLP和搜索背景比較強，對于變長序列的重視程度很高。Paddle的底層數據是LoDTensor，而不是類似其他框架Tensor。LoDTensor相當于把變長序列信息耦合進了Tensor里面。這可能導致比較多的問題，比如很多Operator是完全序列無關的，根本無法處理序列信息在輸入Tensor和輸出Tensor的關系，進而比較隨機的處理，給框架的健壯性埋下隱患。雖然我一直想推動序列信息和Tensor的解耦合，但是因為種種原因，沒有徹底的完成這個重構的目標，希望后面能改掉。
性能
18年初的時候，Paddle還是個原型系統。由于OKR目標，團隊已經開始初步接入一些業務場景。其實一個比較大的痛點就是性能太差。單機單卡速度非常慢，單機4卡加速比只有1.x。但是性能問題的定位卻非常困難。我花了些時間寫了些profile的工具，比如timeline。一些明顯的性能問題可以被快速的定位出來并修復。
但是單機多卡的速度還是非常慢，timeline分析后發現其中有個ParallelOp，存在大量的Barrier。最后改寫成了ParallelExecutor，把Program復制了N份部署在多張卡上，在其中插入AllReduce通信算子，然后這N倍的算子基于圖依賴關系，不斷把ready的算子扔進線程池執行。即使這樣，我們也發現在多卡的性能上，不同模型需要使用不同的線程調度策略來達到最優。很難有一種完美的one-fits-all的方案。后面我們再聊如何通過IR+Pass的方法插件化的支持不同的算子調度策略。
分布式的訓練也碰到不少的問題。一開始使用grpc，花了挺大的功夫做并行請求，然后又切成了brpc，在RDMA等方面做了不少的優化。分布式訓練的性能逐步得到了提升。另外為了做到自動化分布式部署，前面提到的Transpiler隨著場景的增加，Python代碼也變得越來越復雜。
模型推理在公司內碰到來非常強勁的對手。Anakin的GPU推理速度的確很快，讓我吃驚的是他們竟然是用SASS匯編完成大量基礎算子的開發，針對Pascal架構做了異常極致的優化，甚至在某些場景遠超TensorRT。我一直主張訓練和推理要盡量用一樣的框架，并不需要一個單獨的推理框架來解決性能問題。使用不同的框架做推理會造成很多意外的精度問題和人工開銷。
因為推理性能的問題，我們和兄弟團隊發生來曠日持久競賽，作為狗頭軍事，我充分發揮來團隊在CPU這塊的技術積累、以及和Intel外援的良好關系，在CPU推理場景常常略勝一籌。在GPU方面苦于對手無底線使用匯編，和我方戰線太多、人員不夠，只能戰略性放棄了部分頭部模型，通過支持子圖擴展TensorRT引擎的方式，利用Nvidia的技術優勢在許多個通用場景下展開進攻。現在回想起來真實一段有趣的經歷。
Imtermediate Representation&Pass
Imtermediate Representation+Pass的模式主要是從LLVM的架構上借鑒來的。在編譯器上主要是用來解決把M個編程語言中任意一個編譯到N個硬件設備中任意一個執行的問題。簡單的解決方案是為每個編程語言和硬件單獨寫一個編譯器。這需要M*N個編譯器。顯然這對于復雜的編譯器開發來說，是非常高成本的。
Intermediate Representation是架構設計中抽象能力的典型體現。不同編程語言的層次不一樣，或者僅僅是單純的支持的功能有些差異。但是，這些編程語言終歸需要在某種硬件指令集上執行。所以在編譯的過程中，他們會在某個抽象層次上形成共性的表達。而IR+Pass的方法很好的利用了這一點。其基本思想是通過多層Pass (編譯改寫過程），逐漸的把不同語言的表達方式在某個層次上改寫成統一的IR的表達方式。在這個過程中，表達方式逐漸接近底層的硬件。而IR和Pass可以很好的被復用，極大的降低了研發的成本。
深度學習框架也有著非常類似的需求。
用戶希望通過高層語言描述模型的執行邏輯，甚至是僅僅聲明模型的結構，而不去關心模型如何在硬件上完成訓練或者推理。
深度學習框架需要解決模型在多種硬件上高效執行的問題，其中包括協同多個CPU、GPU、甚至大規模分布式集群進行工作的問題。也包括優化內存、顯存開銷、提高執行速度的問題。
更具體的。前文說到需要能夠自動的將用戶聲明的模型Program自動的在多張顯卡上并行計算、需要將Program拆分到多個機器上進行分布式計算、還需要修改執行圖來進行算子融合和顯存優化。
Paddle在一開始零散的開展了上面描述的工作，在分布式、多卡并行、推理加速、甚至是模型的壓縮量化上各自進行模型的改寫。這個過程非常容易產生重復性的工作，也很難統一設計模式，讓團隊不同的研發快速理解這些代碼。
意思到這些問題后，我寫了一個Single Static Assignment（SSA）的Graph，然后把Program通過第一個基礎Pass改寫成了SSA Graph。然后又寫了第二個Pass把SSA Graph改寫成了可以多卡并行的SSA Graph。
后面的事情就應該可以以此類推了。比如推理加速可以在這個基礎上實現OpFusionPass, InferenceMemoryOptimizationPass, PruningPass等等，進而達到執行時推理加速的目的。分布式訓練時則可以有DistributedTransPass。量化壓縮則可以有ConvertToInt8Pass等等。這一套東西基本解決了上層Program聲明到底層執行器的Compiler問題。
這個過程中的確碰到了不少的阻力。比如分布式早期通過Python完成了這個邏輯，需要遷移到C++層。壓縮量化的研發更喜歡寫Python，而IR&Pass是基于C++的。不同Pass間順序依賴和Debug等。
全套深度學習框架工具
TensorFlow Everywhere原本是TensorFlow團隊時的一個口號，意思是TensorFlow需要支持深度學習模型在任意的場景下運行，進而達到AI Everywhere的目標。可以說深度學習框架希望成為AI的“操作系統”，就像魚離不開水、App離不開iOS/Android一樣。
Paddle作為全面對標TensorFlow的國產深度學習框架，自然也希望提供全套的解決方案。在早期的時候，Paddle和公司其他團隊合作了PaddleMobile，提供了移動端的推理能力。后來又開展了Paddle.js，支持在H5、Web等場景的推理能力。為了在toB，在Linux的基礎上又新增了Windows的支持。為了支持無人車等設備、又支持了在更多不同設備上運行。
舉個PaddleMobile的例子。深度學習框架想再移動設備上部署面臨這比較多的挑戰。手機的空間和算力都比服務器小很多，而模型最開始在服務器訓練好后體積相對較大，需要從很多角度下手。1. 使用較小的模型結構。2. 通過量化，壓縮等手段削減模型體積。
另外移動段深度學習框架是通常基于ARM CPU，GPU則有Mali GPU, adreno GPU等等。為了最求比較極致性能，常常需要使用匯編語言。有個同學寫到后面幾乎懷疑人生，感覺自己大學學的東西不太對。為了不顯著增加APP的體積，框架編譯后的體積需要在KB～幾MB的級別，因此需要基于部署的模型結構本身用到的算子進行選擇性編譯。極端的時候甚至需要是通過C++ Code Gen的方法直接生成前向計算必須的代碼，而不是通過一個通用的解釋器。
回顧
隨著項目的復雜化，很多棘手的問題逐漸從深度學習的領域技術問題轉變成了軟件工程開發和團隊管理分工的問題。隨著團隊的不斷變化，自己有時候是作為一個leader的角色在處理問題，有的時候又是以一個independent contributor的角色在參與討論。很慶幸自己經歷過這么一段，有些問題在親身經歷后才能想得明白，想得開。時代有時候會把你推向風口浪尖，讓你帶船隊揚帆起航，在更多的時候是在不斷的妥協與摸索中尋找前進的方向。
無量
無量是騰訊PCG建設的一個深度學習框架，主要希望解決大規模推薦場景下的訓練和推理問題。深度學習在推薦場景的應用和CV、NLP、語音有些不一樣。
業務會持續的產生用戶的行為數據。當用戶規模達到數千萬或者上億時就會產生海量的訓練數據，比如用戶的畫像，用戶的點擊，點贊，轉發行為，還有Context等等。
這些數據是高度稀疏的，通常會編碼成ID類的特征進而通過Embedding的方式進入模型訓練。隨著業務規模的提升和特征工程日漸復雜，比如累計用戶數，商品，內容增加，以及特征交叉的使用，Embedding參數的體積可以達到GB，甚至TB級。
推薦場景是實時動態變化的，新用戶，內容，熱點不斷產生。用戶的興趣，意圖逐漸變化，因此模型需要持續不斷的適應這些變化，時刻保持最好的狀態。
調整
19年中這個項目時大概有2～3人。團隊希望開發一個新的版本，基于TensorFlow進行擴展加強，使得無量可以復用TensorFlow已有的能力，并且能夠支持推薦場景下的特殊需求。無量一開始采用的是基于參數服務器的架構。TensorFlow被復用來提供Python API以及完成基礎算子的執行。而參數服務器，分布式通信等方面則是自己開發，沒有復用TensorFlow。
這個選擇在團隊當時的情況下是比較合理的。如果選擇另一種方向，基于TensorFlow底層進行改造，研發難度會比較大，而且很可能與社區版TensorFlow走向不同的方向，進而導致TensorFlow版本難以升級。而把TensorFlow作為一個本地執行的lib則可以在外圍開發，不需要了解TensorFlow內部的復雜邏輯，也可以復用一些其他開源組件，比如pslib。
早期在軟件開發的流程上相對比較欠缺。為了保障工程的推進，我先幫忙做了些基礎工作，比如加上了第一個自動化測試和持續集成，對一些過度封裝和奇怪的代碼做了重構和簡化。
另外，在接口層也做了一些調整。原來框架開始執行后就進入C++執行器，無法從python層提供或者返回任何執行結果，也無法在python層執行邏輯進行插件化的擴展。為了滿足預期用戶將來需要進行調試的需求，我模擬tf.Session和tf.Estimator對執行層的接口做了重構。這樣用戶可以通過feed/fetch的方式單步調試執行的過程。也可以通過Hook的方式在執行前后擴展任意的邏輯，提高框架的適用場景。
另外一個問題是python層基本完全是全局變量，很難進行多模型的封裝。像TensorFlow有Graph實例或者Paddle有Program實例。因為python層需要重構的量比較大，我暫時先加入了Context的封裝，勉強將各種狀態和配置封裝在了Context下。考慮到短期可能不會有更復雜的需求，暫時沒有把這件事做完。
reader那塊也做了一些重構。最開始那塊的線程模型異常復雜，一部分是因為分布式文件系統等基礎設施無法提供比較好的SDK，導致許多邏輯不得不在深度學習框架里面，比如文件的本地緩存。考慮到特征加工的邏輯比較復雜，以及一些老的TensorFlow用戶可能習慣于tf.Example和tf.feature_column等基礎算子庫，我在reader層引入了基于TensorFlow的tf.dataset。不過后來發現用戶似乎更關心性能問題，喜歡自定義C++ lib的方式來解決特征處理的問題。
API設計是個老大難的問題。TensorFlow，Paddle，無量都沒能幸免。在一個多人協同的團隊里，每個研發更多還是關注每個獨立功能是否完成開發，而功能的接口往往需要考慮到整體的API設計風格，易用性，兼容性等許多因素，常常在高速迭代的過程中被忽略掉。不幸的是API常常不能像內部實現一樣后期優化。當API被放給用戶使用后，后續的修改往往會破壞用戶代碼的正確性。很多時候只能自己評審一下。
升級
無量經過一年基礎能力的打磨，逐漸的成為來整個事業群統一的大規模推薦模型訓練和推理框架，支撐數十個業務場景，每天都能生產數千個增量和全量的模型。簡單的完成功能已經不能滿足業務和團隊發展的需求，需要在技術上更加前沿。
數據處理
數據格式上要從原來的明文轉到更高效的二進制。另外基于CSR編碼的稀疏數據可以進一步的減少數據處理時的拷貝等額外開銷。
流水線
盡量挖掘訓練中可以并行的地方，通過流水線的方式提高并發度，進而提高訓練的速度。比如在數據讀取的過程中，就可以提前按照參數服務器的規模對數據進行預切分，并告知參數服務器需要提前準備哪些參數。這樣當pull/push的時候能夠更快的完成計算，進而提高每個minibatch的速度。
同樣，當使用GPU訓練時，也可以在數據IO的并行過程中，預計算未來需要用的的Embedding參數。這樣GPU訓練下一輪的數據時，需要用到的Embedding已經提前被計算好，可以直接開始訓練，減少來等待的時間。
定制化參數服務器
由于無量解決的一個關鍵問題是推薦模型的海量參數問題，因此參數服務器必須是高度優化過的。并且應該合理的將推薦模型的領域知識引入到設計中，通過特殊的策略進一步產生差異化的優勢。
定制化的線程模型，內存管理和HashMap。由于參數是被切分歸屬到不同線程上，所以可以通過無鎖化的把每次pull/push的參數處理好。另外由于海量參數消耗較大硬件成本，內存空間都需要通過定制化的內存池來管理。否則很可能有大量的空間碎片在默認內存庫中無法及時歸還給操作系統。另外也有無法精細化控制內存清理機制，導致內存OOM或者浪費。定制化的內存管理可以解決這些問題，甚至通過特殊的內存淘汰策略，在不損害模型效果的基礎上進一步降低內存的開銷。高性能HashMap則是需要解決Embedding快速的增刪改查的問題。
Embedding向量的管理也是有非常多可以改進的地方。1. 動態的改變Embeding向量的長度來支持模型的壓縮，提高模型效果。2. 擴展Embedding的元數據來記錄熱度，點擊展現等統計值，有助于提高訓練推理時高級分布式架構的Cache命中率，已經模型的訓練效果。3. 模型的恢復和導出機制在大規模Embedding場景對于Serving時能夠實時加載模型更新重要。另外還需要考慮到任務失敗重啟后資源伸縮等問題。
GPU訓練
傳統PS架構的訓練模式下，由于單臺機器的計算能力有限，需要幾十甚至上百個實例進行分布式訓練。但是這樣會導致大量的計算被用在來無效的開銷上。比如稀疏特征在網絡通信兩邊的處理。這種額外開銷甚至經常超過有效計算。
GPU和相應的高速網絡鏈接可以解決這一問題。單臺8卡機器通過NVLink連接起來，速度甚至可以超過幾十臺物理機，有更高的性價比。但是由于幾百GB，甚至TB級的參數問題，還有Embedding的GPU計算問題，導致GPU一直都沒有被廣泛的用起來。
然而實驗發現其實稀疏特征存在顯著的Power-law分布，少部分Hot特征使用遠多于其他大量不Hot的特征。因此，通過在數據處理時統計特征，然后批量將將來新需要的Embedding換入GPU，就可以讓GPU長時間的進行連續訓練，而不需要頻繁的和CPU內存交換參數。
GPU預測
隨著推薦模型復雜度提高，引入傳統CV，NLP的一些結構需要消耗更多的計算。CPU往往很難在有效的時間延遲下（幾十毫秒）完成大量候（幾百上千）選集在復雜推薦模型的推理。而GPU則成為了一個潛在的解決方案。
同樣，GPU推理也需要解決顯存遠小于Embedding參數的問題。通過在訓練時預先計算Hot Embedding，然后加載如推理GPU，可以一定程度的緩解這個問題。在推理時僅有少部分的Embedding沒有在GPU顯存中緩存，需要通過CPU內存拷貝進入GPU。
而通過模型的量化和壓縮能進一步減少Embedding參數的規模。實驗表明當大部分Embedding參數的值控制為0時，模型依然能夠表現出原來的效果，甚至略優。
總結
深度學習算法的發展和深度學習框架的發展是相輔相成，互相促進的。從2002年時Torch論文發表后，框架的技術發展相對緩慢，性能無法顯著提升導致無法探索更加復雜的算法模型，或者利用更加大規模的數據集。
在2010年后逐漸出現了Caffe, Theano等框架，通過將更高性能的GPU引入，可以訓練更加復雜的CNN和RNN模型，深度學習算法的發展出現來顯著的加速。
到了2014~2017年幾年間，TensorFlow的出現讓用戶可以通過簡單的Python語言將細粒度的算子組裝各種模型結構。并且模型可以簡單的被分布式訓練，然后自動化部署在服務器，手機，攝像頭等各種設備上。而Pytorch的動態圖用法滿足了研究人員對易用性和靈活性更高的要求，進一步推進算法研究。
國內的深度學習框架技術在這股浪潮中也緊跟這世界的步伐。Paddle在14年左右產生，在國內積累了一定的用戶，在當時基本能比肩其他的框架。雖然在TensorFlow和Pytorch等更先進的框架出現后，國內錯過了寶貴的幾年技術升級的窗口以及社區生態培育時機，但是我們看到從18年到20年間，新版的PaddlePaddle，OneFlow, MindSpore等深度學習框架陸續開源，技術上逐漸趕了上來。
推薦場景在電商，視頻，資訊等眾多頭部互聯網公司的火爆導致推薦系統對AI硬件的消耗在互聯網公司超過了傳統NLP，CV，語音等應用的總和。許多公司開始針對推薦場景（以及廣告，搜索）的特殊需求對深度學習框架進行定制優化。百度的abacus是比較早期的框架，和其他早期框架一樣，在易用性和靈活性上較弱。無量，XDL等框架則進行了改進，兼顧了社區兼容性，算法易用性和系統的性能等緯度。
深度學習的框架的觸角其實遠不止我們常見到的。隨著AI技術的推廣，Web、H5、嵌入式設備、手機等場景下都有許多優秀的深度學習框架產生，如PaddleMobile, TFLite，tensorflow.js等等。
深度學習框架的技術也逐漸從更多緯度開始拓展。ONNX被提出來作為統一的模型格式，雖然離目標有很長的距離和問題需要解決。但是從它的流行我們能看到社區對于框架間互通的渴望。隨著摩爾定律難以維持，框架開始更多的從新的硬件和異構計算領域尋求突破。為了支持海量的算子在CPU、FPGA、GPU、TPU、NPU、Cerebras等眾多AI芯片上運行，TVM、XLA等借鑒編譯技術幾十年來的積累，在更加艱巨的道路上進行來持續的探索，經常能傳來新進展的好消息。深度學習框架也不再僅應用于深度學習，還在科學計算，物理化學等領域發光發熱。
后臺回復關鍵詞【入群】
加入賣萌屋NLP/IR/Rec與求職討論群
后臺回復關鍵詞【頂會】
獲取ACL、CIKM等各大頂會論文集！
創作挑戰賽新人創作獎勵來咯，堅持創作打卡瓜分現金大獎總結
以上是生活随笔為你收集整理的从我开发过的Tensorflow、飞桨、无量框架看深度学习这几年的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： NLP十大研究方向Highlights!
- 下一篇： Longformer：超越RoBERTa