AI 四巨頭 Google、DeepMind、Microsoft、Uber 深度學習框架大比拼

編者按：Google、Uber、DeepMind和Microsoft這四大科技公司是當前將深度學習研究廣泛應用於自身業務的典型代表，躋身全球深度學習研究水平最高的科技公司之列。GPipe、Horovod、TF Replicator和DeepSpeed分別是這四家公司開發應用的深度學習框架，它們結合了深度學習研究及其基礎設施的前沿技術，以提高深度學習模型的訓練效率。這四個深度學習框架各有什麼特點和優勢呢？本文將對此做一個基本介紹。

以下為譯文：

      

大規模深度學習訓練是在現實世界中構建深度學習解決方案最具挑戰性的方面之一。 正如俗語所說，你最大的優點可以成為你最大的缺點，這一點自然也適用於深度學習模型。整個深度學習空間在一定程度上取決於深度神經網絡（DNN）跨越GPU拓撲的能力。然而，同樣的擴展能力導致了計算密集型程序的創建，這給大多數組織帶來了操作上的挑戰。從訓練到優化，深度學習項目的生命周期需要健壯的基礎設施構建塊，以便能夠並行化和擴展計算工作負載。

儘管深度學習框架正在快速發展，但相應的基礎架構模型仍處於初期階段。在過去的幾年裡，技術巨頭谷歌、微軟、優步（Uber）、DeepMind和其他公司定期發布了各自的研究成果，以便在大型GPU等基礎設施上實現深度學習模型的並行化。

分布式和並行化計算的原理與深度學習程序生命周期的幾乎所有階段都息息相關。訓練一個深度學習模型可能是一個非常昂貴的任務，運行也是如此。顯而易見的答案是，可以利用大型GPU網絡來分配深度學習程序的工作負載，但這絕非易事。眾所周知，並發和並行編程是非常複雜的，尤其在應用於大型神經網絡時更是如此。大型科技公司每天都在面臨這些挑戰，因為它們必須為關鍵業務應用運行極其複雜的深層神經網絡。

今天，我想回顧一下谷歌、微軟、DeepMind、和優步（Uber）等科技巨頭用於並行化大規模深度學習模型訓練的一些頂級框架。具體來說，我們將討論以下四個框架：

谷歌的 GPipe

優步的Horovod

DeepMind的TF-Replicator

微軟的 DeepSpeed

谷歌的GPipe

GPipe專注於為深度學習項目提高訓練工作量。從基礎設施的角度來看，訓練過程的複雜性是深度學習模型中經常被忽視的一個方面。訓練數據集越來越大，也越來越複雜。例如，在醫療保健領域，遇到需要使用數百萬高解析度圖像進行訓練的模型已不罕見。因此，訓練過程通常需要很長的時間才能完成，並且由於內存和CPU消耗而導致的成本也非常高。

將深度學習模型的並行性分為數據並行性和模型並行性，是一種研究深度學習模型並行性的有效方法。數據並行方法使用大量的機器集群來拆分輸入數據。模型並行性嘗試將模型移至加速器上，如GPU或TPU，它們具有加速模型訓練的特殊硬體。在較高的層次上，幾乎所有的訓練數據集都可以按照一定的邏輯進行並行化，但在模型上卻不是這樣。例如，一些深度學習模型是由可以獨立訓練的並行分支組成的。在這種情況下，經典的策略是將計算劃分為多個分區，並將不同的分區分配給不同的分支。然而，這種策略在連續堆疊層的深度學習模型中存在不足，這給高效並行計算帶來了挑戰。

GPipe採用一種稱為管道的技術，將數據和模型並行性結合起來。從概念上講，GPipe是一個分布式機器學習庫，它使用同步隨機梯度下降和管道並行性進行訓練，適用於任何由多個連續層組成的深度神經網絡（DNN）。GPipe將一個模型劃分到不同的加速器上，並自動將一個小批量的訓練實例分割成更小的微批量。該模型允許GPipe的加速器並行運行，以最大限度地提高了訓練過程的可伸縮性。

下圖說明了具有多個連續層的神經網絡的GPipe模型在四個加速器之間的劃分。Fk是第k個劃分的複合正向計算函數。Bk是對應的反向傳播函數。Bk既依賴於上層的Bk+1，也依賴於Fk的中間激活。在圖片的上部中，我們可以看到網絡的順序特性如何導致資源利用不足的。圖片的下部顯示了GPipe方法，其中輸入的小批量（mini-batch）被劃分為更小的宏批量（macro-batch），這些宏批量（macro-batch）可以由加速器同時處理。

             

谷歌開源了GPipe的一種實現，作為TensorFlow項目的一部分。

開源地址：https://github.com/tensorflow/lingvo/blob/master/lingvo/core/gpipe.py

優步的 Horovod

Horovod是Uber的機器學習（ML）堆棧之一，它已經在社區中非常流行，並且已經被DeepMind和OpenAI等人工智慧巨頭的研究團隊採用。從概念上講，Horovod是一個用於大規模運行分布式深度學習訓練工作的框架。

Horovod利用諸如OpenMPI之類的消息傳遞接口棧，來使訓練作業能夠在高度並行和分布式的基礎設施上運行，而無需進行任何修改。通過以下四個簡單的步驟即可在Horovod中運行分布式TensorFlow訓練工作：

config.gpu_options.visible_device_list = str(hvd.local_rank())為每個TensorFlow進程分配一個GPU。opt=hvd.DistributedOptimizer(opt)使用Horovod優化器包裝任何常規的TensorFlow優化器，該優化器使用ring-allreduce來處理平均梯度。hvd.BroadcastGlobalVariablesHook(0)將變量從第一個進程廣播到所有其他進程，以確保一致的初始化。下面這個代碼示例是一個基本的TensorFlow訓練作業的模板，你可以從中看到上面的四個步驟：

1import tensorflow as tf
2import horovod.tensorflow as hvd# Initialize Horovod
3hvd.init()# Pin GPU to be used to process local rank (one GPU per process)
4config = tf.ConfigProto()
5config.gpu_options.visible_device_list = str(hvd.local_rank())# Build model…
6loss = …
7opt = tf.train.AdagradOptimizer(0.01)# Add Horovod Distributed Optimizer
8opt = hvd.DistributedOptimizer(opt)# Add hook to broadcast variables from rank 0 to all other processes during
9# initialization.
10hooks = [hvd.BroadcastGlobalVariablesHook(0)]# Make training operation
11train_op = opt.minimize(loss)# The MonitoredTrainingSession takes care of session initialization,
12# restoring from a checkpoint, saving to a checkpoint, and closing when done
13# or an error occurs.
14with tf.train.MonitoredTrainingSession(checkpoint_dir=「/tmp/train_logs」,
15                                      config=config,
16                                      hooks=hooks) as mon_sess:
17 while not mon_sess.should_stop():
18   # Perform synchronous training.
19   mon_sess.run(train_op)

TF-Replicator專注於TensorFlow程序如何利用Tensor處理單元（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）一個定義其模型邏輯的階躍函數（例如，梯度下降的單個階躍）。下面的代碼片段展示了一個簡單的TF-Replicator程序：

1# Deploying a model with TpuReplicator.
2repl = tf_replicator.TpuReplicator(
3num_workers=1, num_tpu_cores_per_worker=8
4)
5with repl.context():
6model = resnet_model()
7base_optimizer = tf.train.AdamOptimizer()
8optimizer = repl.wrap_optimizer(base_optimizer)# ... code to define replica input_fn and step_fn.per_replica_loss = repl.run(step_fn, input_fn)
9train_op = tf.reduce_mean(per_replica_loss)with tf.train.MonitoredSession() as session:
10repl.init(session)
11for i in xrange(num_train_steps):
12session.run(train_op)
13repl.shutdown(session)

為了優化不同設備之間的通信，TF-Replicator利用了最先進的MPI接口。在一些實驗中，DeepMind能夠在一個TPUv3 pod的512個核心上，以2048的batch size批量訓練著名的BigGAN模型。目前，TF-Replicator是DeepMind公司的TPU的主要編程接口。

微軟的DeepSpeed是一個新的開源框架，專注於優化大型深度學習模型的訓練。當前版本包含了ZeRO的第一個實現以及其他優化方法。從編程的角度來看，DeepSpeed是在PyTorch之上構建的，它提供了一個簡單的API，允許工程師只需幾行代碼就可以利用並行化技術來進行訓練。DeepSpeed抽象了大規模訓練的所有困難方面，例如並行化、混合精度、梯度累積和檢查點，使得開發人員可以專注於模型的構建。從功能角度來看，DeepSpeed在以下四個關鍵方面表現出色：規模：DeepSpeed可以為運行高達1,000億個參數的模型提供系統支持，這和其他訓練優化框架相比，提高了10倍。速度：在最初的測試中，DeepSpeed表現出的吞吐量比其他庫高出4-5倍。成本：可以使用DeepSpeed進行訓練的模型，其成本比其他替代方案少三倍。可用性：DeepSpeed不需要重構PyTorch模型，僅需幾行代碼即可使用。   將深度學習模型的訓練並行化是一項非常複雜的工作，超出了大多數機器學習團隊的專業知識。利用谷歌、微軟、優步或DeepMind等科技公司創建的框架和架構，肯定可以簡化這些工作。在不久的將來，我們希望看到這些框架的一些版本包含在主流深度學習堆棧中，以加速核心深度學習社區準入的民主化。

https://towardsdatascience.com/the-frameworks-that-google-deepmind-microsoft-and-uber-use-to-train-deep-learning-models-at-scale-30be6295725

【END】

今日福利：評論區留言入選，可獲得價值299元的「2020 AI開發者萬人大會」在線直播門票一張。  快來動動手指，寫下你想說的話吧。

點擊閱讀原文，精彩繼續！