大規模訓練是在現實世界中構建深度學習解決方案最具挑戰性的方面之一。。從訓練到優化，深度學習程序的生命週期需要健壯的基礎設施構建塊，以便能夠並行化和擴展計算工作負載。儘管深度學習框架正在快速發展，但相應的基礎架構模型仍處於初期階段。在過去的幾年中，技術巨頭Google，Microsoft，Uber，DeepMind和其他公司定期發佈了研究成果，以實現跨大型GPU基礎架構的深度學習模型的並行化。
分佈式和並行化計算的原理與深度學習程序生命週期的幾乎所有階段都息息相關。訓練一個深度學習模型是一項非常昂貴的練習，它的執行也是如此。顯而易見的答案是，可以利用大型GPU網絡來分配深度學習程序的工作負載，但這絕非易事。眾所周知，併發和並行編程是非常複雜的，在應用於大型神經網絡時更是如此。今天，我想回顧一下在谷歌、DeepMind、Microsoft和Uber中用於並行化大規模深度學習模型訓練的一些頂級架構。
谷歌的GPipe
GPipe專注於為深度學習項目擴展訓練工作量。從基礎架構的角度來看，訓練過程的複雜性是深度學習模型中經常被忽視的一個方面。訓練數據集越來越大，越來越複雜。例如，在衛生保健領域，遇到需要使用數百萬高分辨率圖像進行訓練的模型並不罕見。結果，訓練過程通常要花很長時間才能完成，由於內存和CPU消耗而導致的成本也非常高。
將深度學習模型的並行性分為數據並行性和模型並行性是研究深度學習模型並行性的有效方法。數據並行方法使用大量的機器集群來拆分輸入數據。模型並行性嘗試將模型移至具有專用硬件以加速模型訓練的加速器，例如GPU或TPU。幾乎所有的訓練數據集都可以按照一定的邏輯並行化，但模型卻不是這樣。例如，一些深度學習模型是由可以獨立訓練的並行分支組成的。在這種情況下，典型的策略是將計算劃分為分區，並將不同的分區分配給不同的分支。然而，這種策略在順序堆疊層的深度學習模型中存在不足，這給高效地並行化計算帶來了挑戰。
GPipe利用一種稱為管道的技術，將數據和模型並行性結合起來。從概念上講，GPipe是一個分佈式機器學習庫，它使用同步隨機梯度下降和管道並行性進行訓練，適用於任何由多個連續層組成的DNN。GPipe將一個模型劃分到不同的加速器上，並自動將一個小批量的訓練實例分割成更小的微批量。該模型允許GPipe的加速器並行運行，以最大限度地提高訓練過程的可伸縮性。
下圖說明了神經網絡的GPipe模型在四個加速器之間的劃分。Fk是第k個劃分的複合正向計算函數。Bk是對應的反向傳播函數。Bk既依賴於上層的Bk+1，也依賴於Fk的中間激活。在模型上部中，我們可以看到網絡的順序特性是如何導致資源利用不足的。圖片下部顯示了GPipe方法，其中輸入的mini-batch被分成更小的macro-batches，這些macro-batches可以由加速器同時處理。
Uber Horovod
Horovod是在社區中變得極為流行的Uber ML堆棧之一，並已被DeepMind或OpenAI等AI巨頭的研究團隊採用。從概念上講，Horovod是用於大規模運行分佈式深度學習訓練工作的框架。
Horovod利用諸如OpenMPI之類的消息傳遞接口棧來使訓練作業能夠在高度並行且分佈式的基礎架構上運行，而無需進行任何修改。通過以下四個簡單的步驟即可在Horovod中運行分佈式TensorFlow訓練工作：
- hvd.init（）初始化Horovod。
- config.gpu_options.visible_device_list = str（hvd.local_rank（））為每個TensorFlow進程分配一個GPU。
- opt = hvd.DistributedOptimizer（opt）使用Horovod優化器包裝任何常規的TensorFlow優化器，該優化器使用ring-allreduce來平均梯度。
- hvd.BroadcastGlobalVariablesHook（0）將變量從第一個進程廣播到所有其他進程，以確保一致的初始化。
這是基本TensorFlow訓練作業的模板：
DeepMind的TF-Replicator
TF-Replicator專注於與TensorFlow程序如何利用Tensor處理單元（TPU）相關的可伸縮性的不同方面。TPU被認為是最先進的AI芯片之一，可為機器學習工作負載提供本地可擴展性。但是，在TensorFlow程序中使用TPU 需要專用的API，這會給不熟悉底層硬件模型的數據科學家帶來可移植性問題和採用障礙。DeepMind的TF複製器通過提供一個更簡單，友好的開發程序模型來利用TensorFlow程序中的TPU，從而解決了這一難題。
TF-Replicator關注的是可伸縮性的另一個方面，即TensorFlow程序如何利TPU。TPU被認為是最先進的人工智能芯片之一，它為機器學習工作負載提供了本機可伸縮性。然而，在TensorFlow程序中使用TPU需要專門的api，這會導致可移植性問題。DeepMind的TF-Replicator提供了一個更簡單、對開發人員友好的編程模型，用於在TensorFlow程序中利用TPU。
TF-Replicator的魔力在於“in-graph replication”模式，其中每個設備的計算都複製到同一TensorFlow圖中。設備之間的通信是通過連接設備相應子圖中的節點來實現的。為了達到該級別的並行化，TF-Replicator利用TensorFlow的圖重寫模型在圖中的設備之間插入本地通信。當提供一個TensorFlow圖時，TF-Replicator首先獨立地構建每個設備的計算，並在用戶指定跨設備計算的地方留下佔位符。一旦構建了所有設備的子圖，TF-Replicator就會用實際的跨設備計算代替佔位符來連接它們。
從編程模型的角度來看，使用TF-Replicator編寫的代碼看起來類似於為單個設備編寫的本機TensorFlow代碼。用戶只需要定義（1）公開數據集的輸入函數，（2）定義其模型邏輯的階躍函數（例如，梯度下降的單個階躍）。以下Python代碼片段顯示了一個簡單的TF-Replicator程序：
為了優化不同設備之間的通信，TF-Replicator利用了最先進的MPI接口。在一些實驗中，DeepMind能夠在TPUv3 pod的512個核上，以2048的batches size訓練著名的BigGAN模型。目前，TF-Replicator是DeepMind的主要TPU編程接口。
微軟DeepSpeed
微軟的DeepSpeed是一個新的開源框架，專注於優化大型深度學習模型的訓練。當前版本包括ZeRO的第一個實現以及其他優化方法。從編程的角度來看，DeepSpeed是建立在PyTorch之上的，它提供了一個簡單的API，讓工程師們只需幾行代碼就可以利用並行化技術進行訓練。DeepSpeed抽象了大規模訓練的所有困難方面，例如並行化、混合精度、梯度累積和檢查點，使開發人員可以專注於模型的構建。
DeepScale在四個關鍵方面都表現出色：
· 規模： DeepSpeed提供系統支持以運行多達1000億個參數的模型，與其他訓練優化框架相比，這代表了10倍的改進。
· 速度：在最初的測試中，DeepSpeed顯示出的吞吐量是其他庫的4到5倍。
· 成本：使用DeepSpeed訓練模型的成本比其他方法低三倍。
· 可用性： DeepSpeed不需要重構PyTorch模型，僅需幾行代碼即可使用。
將深度學習模型的訓練並行化是一項非常複雜的練習，超出了大多數機器學習團隊的專業能力。利用谷歌、微軟(Microsoft)、優步(Uber)或DeepMind等科技公司創建的框架和架構，肯定可以簡化這些工作。
閱讀更多 不靠譜的貓 的文章
