編輯：LRS
【新智元導讀】到底多大的引數量能滿足研究人員的胃口？一百萬億夠不夠？華人團隊最近開源了一個史上最大的訓練系統Persia，最高支援百萬億級模型的訓練，但前提是得有卡！
不夠大！還不夠大？
在NLP領域，從BERT的億級引數，到OpenAI的1700億引數，再到Google Switch Transformer的1.6萬億引數，研究人員對引數量增長的渴望從來沒有停止過，而我等吃瓜群眾對引數量也早已經麻了。
那，100000000000000夠不夠？不用數了，後面14個零！一百萬億！
最近來自快手和蘇黎世聯邦理工學院的研究人員提出了一個新的推薦系統Persia，最高支援100萬億級引數的模型訓練，比目前最大的FB-ZionEX推薦系統還高出一個數量級。
論文地址：https://arxiv.org/pdf/2111.05897.pdf
開源地址：https://github.com/PersiaML/Persia
為了支援如此龐大規模模型的訓練，同時保證訓練效率和訓練精度，文中提出了一種全新的混合訓練演算法：透過混合非同步、同步機制對embedding層和dense層進行分別訓練。
用了這套機制的推薦系統Persia（parallel recommendation training system with hybrid acceleration，即混合加速的並行推薦系統）在一百萬億級引數量的模型上進行了理論和實驗驗證，證實了這種設計和實現的合理性。
作者表示，使用這套機制，任何人都可以很容易訓練一個百萬億級引數量的模型。（只有3060Ti可以嗎？）
值得一提的是，除了2016年Alibaba-XDL和2018年Baidu-AIBox以外，其他大規模推薦系統模型都沒有開源。
Persia的理論基礎
在推薦系統中，深度學習模型目前也成了主流，同樣遵循著「引數量大一級壓死人」原則。
是不是感覺購物app、短影片app越來越容易猜到了你的想法了？
實際上，現代推薦系統取得的進展很大程度上就是來源於不斷增大的模型規模，Google旗下的Youtube在2016年首次將推薦系統的模型規模推向十億引數，從此模型引數量開始一路狂奔，Facebook(Meta)最新的模型將規模拉到12萬億，比Youtube的引數量高出12000倍！
模型引數量的規模每次上一個新臺階，都能帶來效能上的明顯改進，所以劍指100萬億也是有實際意義的。
一般來說，推薦系統模型首先需要將不同的ID特徵（如使用者ID和session ID）對映到一個固定長度的低維向量，而系統中的使用者ID、交叉特徵數量都特別多，就需要更大規模的模型來捕獲特徵和對映。但更大規模的embedding layer也需要更大的記憶體來載入，不得不說大模型太費錢了！
有了embedding後，剩下的工作就簡單了，設計後續layer來適配不同的任務。通常只佔據整個模型的0.1%，無需大記憶體，主要是一些計算密集型的工作。
雖然效能好了，但隨著模型規模的不斷擴大，模型的訓練也是越來越難！
而想要訓練好一個推薦系統，主要看這5個關鍵步驟：
- 準備訓練樣例的embedding
- 神經網路的前向傳播（forward propagation）
- 神經網路的反向傳播（backward propagation）
- 神經網路引數需要同步更新（Synchronization）
- 根據相應的梯度對embedding進行更新
由於embedding層和推薦任務是相關聯的，所以上述五個任務必須順序執行，而這種序列機制也導致了模型訓練的硬體效率很低。
如果不嚴格順序執行這五步，而是採用非同步分散式訓練（Asynchronous distributed training）方式的話，就能顯著提升訓練速度。
目前業界對神經網路的非同步分散式訓練主要有兩個觀察結論：
- 非同步更新（Asynchronous update）對於稀疏訪問（sparse access）來說是有效的。例如SGD每次更新只會影響到模型引數的一小部分，所以就算並行SGD最佳化，引數覆蓋（overwrite）情況也是很少見的，幾乎對模型不會產生影響。
- 滯後性（staleness）限制了非同步SGD的可擴充套件性和收斂。對於稠密的（dense）的神經網路，SGD覆蓋更新的情況下可能會非常多，對模型的最終效果有影響。
在不考慮模型結構的時候，肯定全非同步的硬體效率最高，但在實際情況中，全非同步最佳化會降低模型的效能，對於商用推薦模型來說是不能接受的，所以「同步、非同步兩手抓」才是最終解決方案！
文中提出的混合訓練演算法（Hybrid training）對模型的embedding層和dense層分別使用不同的最佳化機制來提升硬體效率。
因為embedding層更大、更稀疏，採用非同步訓練的方式。
而dense層更稠密，就採用同步訓練的方式。
有了思路，下一步就是系統實現了！
Persia實現細節
Persia系統的設計上主要有兩個難點：
- 在異構叢集上部署訓練工作流（training workflow）
- 在混合基礎架構上部署對應的訓練過程（training procedure）
對於第一個難題，研究人員為了支援基於深度學習的推薦模型的分散式訓練，就直接使用了Tensforflow框架提供的PS正規化，同時將embedding和神經網路引數的儲存和更新放在一組PS節點（即CPU機器）中，將前向和後向傳播的計算放在一組工作節點（即GPU機器）中。
但只使用這種方式還遠遠不夠高效，甚至在實際情況中可能根本沒法完成部署。
例如推薦模型太大，GPU的RAM很容易就不夠用了，而且將embedding和神經網路模組進行統一檢視管理可能還會引入大量不必要的網路流量。
雖然之前的研究也提出過一些PS正規化最佳化、帶快取的colocated PS framework來減少通訊開銷等等，但這些方法根本無法支撐百萬億級模型的訓練。
Persia的研究人員分解出4大模組來解決這個問題：
- 使用一個data loader專門從分散式的儲存如Hadoop, Kafka中取訓練資料；
- 專門建立一個embedding引數伺服器(PS)來管理embedding層引數的儲存和更新；
- 使用一組embedding workers執行最佳化演算法從embedding PS中獲取引數更新；並把embedding向量聚合後放回到embedding PS中；
- 使用一組NN workers執行dense神經網路的前向和反向傳播。
還沒完！上面還提到了第二個難點就是分散式的訓練過程，研究人員又來了一個「七步走」解決了這個難題：
- data loader會對ID型別特徵進行排程分配給embedding worker，然後embedding worker會生成一個獨一無二的訓練樣例，將樣例的ID在本地快取起來後返回給data loader，最後data loader會將該樣例的Non-ID型別特徵和標籤與生成的ID繫結起來；
- data loader會把這個Non-ID型別特徵和標籤分配給一個NN worker;
- NN worker接收到不完整的訓練樣例後，會發起一個請求從embedding worker處pull回ID型別特徵的embedding，同時進行前向傳播；
- embedding worker會對原始的embedding向量進行一些potential aggregation，然後將聚合後的向量發起pull request傳送給NN worker；
- 上面四步完成後，NN網路就成功接收到一組完整的輸入了，然後建立mini-batch進行訓練。此時計算的引數結果都是放在記憶體中的，NN網路後續會採用AllReduce正規化對梯度進行同步；
- NN worker將計算好的梯度送回到embedding worker；
- embedding worker根據樣例ID找到快取的ID型別特徵，然後計算embedding引數的梯度，並將梯度送回到embedding PS中。至此完成全部引數的更新。
最後的訓練結果也可以看到，Persia的訓練速度相比其他大規模模型來說要快了很多，從而能夠支撐百萬億級模型的訓練。
參考資料：
https://arxiv.org/pdf/2111.05897.pdf
