業界 | Google開源大規模神經網路模型高效訓練庫 GPipe

• 2019 年 11 月 1 日
• 筆記

AI 科技評論按：Google昨日在部落格中宣布開源大規模神經網路模型高效訓練庫 GPipe，這是一款分散式機器學習庫，可以讓研究員在不調整超參數的情況下，部署更多的加速器以對大規模模型進行訓練，有效擴展了模型性能。雷鋒網 AI 科技評論對此進行編譯如下。

深度神經網路（DNNs）推進諸多機器學習任務的進步，其中包括語音識別、視覺識別和語言處理等。BigGan、Bert 、GPT2.0 等最新成果表明，DNN 的模型越大，任務處理的表現就越好，而該結論也在過去的視覺識別任務中得到了驗證，表明模型大小與分類準確性之間存在很強的關聯性。舉個例子，2014 年 ImageNet 視覺識別挑戰賽的冠軍 GoogleNet 通過對 400 萬參數進行調整，最終取得 74.8 ％ 的精確度成績；僅僅過了三年，2017 年 ImageNet 挑戰賽冠軍 Squeeze-and-Excitation Networks 調整的參數便高達 1.458 億（36 倍以上），最終取得了 82.7％ 的精確度成績。與此對應的是，市面的 GPU 記憶體僅僅提高了 3 倍左右，目前最先進的影像模型早已達到雲 TPUv2 可用記憶體的極限。因此，我們迫切需要一種高效、可擴展的基礎設施，以實現大規模的深度學習訓練，並克服當前的加速器記憶體受限問題。

最具有代表性的影像分類模型體現出了 ImageNet 精確度結果與模型大小的強關聯性

在《GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism》（https://arxiv.org/pdf/1811.06965.pdf）一文中，我們展示了如何通過流水並行技術（pipeline parallelism）對 DNN 的訓練進行擴展以克服這一限制。簡單來說，GPipe 是一款分散式機器學習庫，基於同步隨機梯度下降與流水並行技術進行模型訓練，適用於由多個連續層組成的任意 DNN。最重要的是，GPipe 可以讓研究員在不調整超參數的情況下，部署更多的加速器以訓練大規模模型，由此有效擴展了性能。為了證明 GPipe 的能力，我們在 Google Cloud TPUv2s 上訓練了一個具有 5.57 億模型參數、 480 x 480 輸入影像尺寸的 AmoebaNet-B。該模型在多個流行數據集上表現良好，取得的成就包括：single-crop ImageNet 的精確度提高至 84.3％、 CIFAR-10 的精確度提高至 99％、CIFAR-100 的精確度提高至 91.3％。

核心 GPipe 庫已在 Lingvo 框架下進行開源： https://github.com/tensorflow/lingvo/blob/master/lingvo/core/gpipe.py

從小批次至微批次

目前存在兩種標準方法可以對中等規模的 DNN 模型進行加速。數據並行方法（The data parallelism）可以納入更多的機器，並將輸入的數據區分開來。另一種方法則是將模型置於加速器上（比如 GPU 或 TPU）——這些加速器的特殊硬體可加速模型的訓練進程。然而加速器卻面臨著記憶體與主機通訊頻寬兩方面受限的問題。因此，通過將模型進行分區，並根據分區配置相應的加速器，模型並行技術可以讓我們在加速器上訓練更大規模的 DNN 模型。由於 DNN 存在順序性，這種策略最後可能變成計算期間只有一個加速器處於活躍狀態，未能將加速器的計算能力充分利用起來。此外，標準的數據並行技術只允許在多個加速器上同時訓練具有不同輸入數據的相同模型，卻無法提升加速器所能支援的最大模型規模。

為了實現跨加速器的高效訓練，GPipe 先按照加速器對模型進行劃分，然後自動將小批次的訓練示例拆分為更小的微批次。通過在微批次中執行流水管理，加速器得以並行運行。此外，梯度將在微批次中持續累積，以免分區的數量影響到模型的品質。

上圖：由於網路存在順序性，模型並行策略導致嚴重的利用不充分問題。 每次都只有一個加速器處於活動狀態。

下圖：GPipe 將輸入的小批次拆成更小的微批次，使不同的加速器可以同時在單獨的微批次上運作。

記憶體和效率的最大化

GPipe 會對模型參數的記憶體分配進行最大化處理。我們在每個 TPUv2 均配有 8 個加速器核心以及 64 GB 記憶體（每個加速器 8 GB）的雲 TPUv2 上做了相關實驗。如果沒有 GPipe，由於記憶體的限制，單個加速器頂多只能訓練 8200 萬個模型參數。而通過反向傳播以及批量分割技術進行重新計算的 GPipe ，成功將中間激活記憶體從 6.26 GB 減至 3.46GB，由此實現單個加速器上訓練 3.18 億個參數的成果。此外，我們還發現在流水並行技術的作用下，模型的最大規模與分區數量成正比，正如事前所預料的那樣。總的來說，GPipe 使 AmoebaNet 能在雲 TPUv2 的 8 個加速器上納入 18 億個參數，比起之前高出了 25 倍。

為了測試模型的效率，我們研究了 GPipe 對 AmoebaNet-D 模型吞吐量的影響情況。由於訓練過程需要至少兩個加速器以適應模型尺寸，因此我們只能對沒有實施流水並行技術的兩個分區案例的加速情況進行觀察。我們發現訓練過程存在近乎線性的加速效果。與兩個分區案例相比，將模型分布在四倍數量的加速器上能有效實現 3.5 倍的加速效果。我們的實驗均使用了雲 TPUv2，但我們了解到最新的的雲 TPUv3 由於每個 TPUv3 均配備了 16 個加速器核心以及 256 GB（每個加速器 16 GB），因此擁有更理想的表現性能。當我們在所有 16 個加速器上對模型進行分發，GPipe 能讓基於 1024-token 句子的 80 億參數 Transformer 語言模型的訓練速度提高 11 倍。

通過 GPipe 對 AmoebaNet-D 進行加速。這種模型不適用於單個加速器。naive-2 基準線是將模型拆分為兩個分區，最終取得的性能結果。Pipeline-k 對應的是 GPipe 將模型分成帶有 k 個加速器的 k 個分區的最終性能結果。

在無需更改超參數的情況下，GPipe 還能通過使用更多加速器來擴展訓練結果。因此，它可以與數據並行技術相結合，通過互補的方式使用更多的加速器來擴展神經網路的訓練規模。

精準度測試

我們試圖通過 GPipe 證明，對現有的神經網路進行擴展，可以實現更理想的模型品質。為此，我們在 ImageNet ILSVRC-2012 數據集上訓練一個模型參數為 5.57 億、輸入影像尺寸為 480 x 480 的 AmoebaNet-B。該網路被分為 4 個分區，在模型與數據上執行了並行訓練程式。該巨型模型在沒有任何外部數據的情況下，最終達到了最先進的 84.3％ top-1 / 97％ top-5 的single-crop 驗證準確度結果。這說明大型的神經網路不僅適用於 ImageNet 等數據集，還能通過遷移學習的方式作用於其他數據集。事實證明，更好的 ImageNet 模型擁有更理想的傳輸效果。我們在 CIFAR10 和 CIFAR100 數據集上進行了遷移學習實驗。我們的巨型模型成功將 CIFAR-10 的精確度提高至到 99％、CIFAR-100 的精確度提高到 91.3％。

結論

當下許多機器學習應用（如自動駕駛和醫學成像）得以持續發展並取得成功的原因，在於實現了儘可能高的模型精確度。然而這也意味著我們需要構建一個更大、更複雜的模型，我們很高興能夠為研究社區提供 GPipe，我們希望未來它可以成為高效訓練大規模 DNN 的基礎設施。

via https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html