谷歌宣布開源神經(jīng)網(wǎng)絡(luò)模型高效訓(xùn)練庫 GPipe
谷歌4日在博客中宣布開源大規(guī)模神經(jīng)網(wǎng)絡(luò)模型高效訓(xùn)練庫 GPipe,這是一款分布式機器學(xué)習(xí)庫,可以讓研究員在不調(diào)整超參數(shù)的情況下,部署更多的加速器以對大規(guī)模模型進行訓(xùn)練,有效擴展了模型性能。對此進行編譯如下。
深度神經(jīng)網(wǎng)絡(luò)(DNNs)推進諸多機器學(xué)習(xí)任務(wù)的進步,其中包括語音識別、視覺識別和語言處理等。BigGan、Bert 、GPT2.0 等最新成果表明,DNN 的模型越大,任務(wù)處理的表現(xiàn)就越好,而該結(jié)論也在過去的視覺識別任務(wù)中得到了驗證,表明模型大小與分類準確性之間存在很強的關(guān)聯(lián)性。舉個例子,2014 年 ImageNet 視覺識別挑戰(zhàn)賽的冠軍 GoogleNet 通過對 400 萬參數(shù)進行調(diào)整,最終取得 74.8 % 的精確度成績;僅僅過了三年,2017 年 ImageNet 挑戰(zhàn)賽冠軍 Squeeze-and-Excitation Networks 調(diào)整的參數(shù)便高達 1.458 億(36 倍以上),最終取得了 82.7% 的精確度成績。與此對應(yīng)的是,市面的 GPU 內(nèi)存僅僅提高了 3 倍左右,目前最先進的圖像模型早已達到云 TPUv2 可用內(nèi)存的極限。因此,我們迫切需要一種高效、可擴展的基礎(chǔ)設(shè)施,以實現(xiàn)大規(guī)模的深度學(xué)習(xí)訓(xùn)練,并克服當前的加速器內(nèi)存受限問題。
最具有代表性的圖像分類模型體現(xiàn)出了 ImageNet 精確度結(jié)果與模型大小的強關(guān)聯(lián)性
在《GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism》一文中,我們展示了如何通過流水并行技術(shù)(pipeline parallelism)對 DNN 的訓(xùn)練進行擴展以克服這一限制。簡單來說,GPipe 是一款分布式機器學(xué)習(xí)庫,基于同步隨機梯度下降與流水并行技術(shù)進行模型訓(xùn)練,適用于由多個連續(xù)層組成的任意 DNN。最重要的是,GPipe 可以讓研究員在不調(diào)整超參數(shù)的情況下,部署更多的加速器以訓(xùn)練大規(guī)模模型,由此有效擴展了性能。為了證明 GPipe 的能力,我們在 Google Cloud TPUv2s 上訓(xùn)練了一個具有 5.57 億模型參數(shù)、 480 x 480 輸入圖像尺寸的 AmoebaNet-B。該模型在多個流行數(shù)據(jù)集上表現(xiàn)良好,取得的成就包括:single-crop ImageNet 的精確度提高至 84.3%、 CIFAR-10 的精確度提高至 99%、CIFAR-100 的精確度提高至 91.3%。
從小批次至微批次
目前存在兩種標準方法可以對中等規(guī)模的 DNN 模型進行加速。數(shù)據(jù)并行方法(The data parallelism)可以納入更多的機器,并將輸入的數(shù)據(jù)區(qū)分開來。另一種方法則是將模型置于加速器上(比如 GPU 或 TPU)——這些加速器的特殊硬件可加速模型的訓(xùn)練進程。然而加速器卻面臨著內(nèi)存與主機通信帶寬兩方面受限的問題。因此,通過將模型進行分區(qū),并根據(jù)分區(qū)配置相應(yīng)的加速器,模型并行技術(shù)可以讓我們在加速器上訓(xùn)練更大規(guī)模的 DNN 模型。由于 DNN 存在順序性,這種策略最后可能變成計算期間只有一個加速器處于活躍狀態(tài),未能將加速器的計算能力充分利用起來。此外,標準的數(shù)據(jù)并行技術(shù)只允許在多個加速器上同時訓(xùn)練具有不同輸入數(shù)據(jù)的相同模型,卻無法提升加速器所能支持的最大模型規(guī)模。
為了實現(xiàn)跨加速器的高效訓(xùn)練,GPipe 先按照加速器對模型進行劃分,然后自動將小批次的訓(xùn)練示例拆分為更小的微批次。通過在微批次中執(zhí)行流水管理,加速器得以并行運行。此外,梯度將在微批次中持續(xù)累積,以免分區(qū)的數(shù)量影響到模型的質(zhì)量。
上圖:由于網(wǎng)絡(luò)存在順序性,模型并行策略導(dǎo)致嚴重的利用不充分問題。 每次都只有一個加速器處于活動狀態(tài)。
下圖:GPipe 將輸入的小批次拆成更小的微批次,使不同的加速器可以同時在單獨的微批次上運作。
內(nèi)存和效率的最大化
GPipe 會對模型參數(shù)的內(nèi)存分配進行最大化處理。我們在每個 TPUv2 均配有 8 個加速器核心以及 64 GB 內(nèi)存(每個加速器 8 GB)的云 TPUv2 上做了相關(guān)實驗。如果沒有 GPipe,由于內(nèi)存的限制,單個加速器頂多只能訓(xùn)練 8200 萬個模型參數(shù)。而通過反向傳播以及批量分割技術(shù)進行重新計算的 GPipe ,成功將中間激活內(nèi)存從 6.26 GB 減至 3.46GB,由此實現(xiàn)單個加速器上訓(xùn)練 3.18 億個參數(shù)的成果。此外,我們還發(fā)現(xiàn)在流水并行技術(shù)的作用下,模型的最大規(guī)模與分區(qū)數(shù)量成正比,正如事前所預(yù)料的那樣??偟膩碚f,GPipe 使 AmoebaNet 能在云 TPUv2 的 8 個加速器上納入 18 億個參數(shù),比起之前高出了 25 倍。
為了測試模型的效率,我們研究了 GPipe 對 AmoebaNet-D 模型吞吐量的影響情況。由于訓(xùn)練過程需要至少兩個加速器以適應(yīng)模型尺寸,因此我們只能對沒有實施流水并行技術(shù)的兩個分區(qū)案例的加速情況進行觀察。我們發(fā)現(xiàn)訓(xùn)練過程存在近乎線性的加速效果。與兩個分區(qū)案例相比,將模型分布在四倍數(shù)量的加速器上能有效實現(xiàn) 3.5 倍的加速效果。我們的實驗均使用了云 TPUv2,但我們了解到最新的的云 TPUv3 由于每個 TPUv3 均配備了 16 個加速器核心以及 256 GB(每個加速器 16 GB),因此擁有更理想的表現(xiàn)性能。當我們在所有 16 個加速器上對模型進行分發(fā),GPipe 能讓基于 1024-token 句子的 80 億參數(shù) Transformer 語言模型的訓(xùn)練速度提高 11 倍。
通過 GPipe 對 AmoebaNet-D 進行加速。這種模型不適用于單個加速器。naive-2 基線是將模型拆分為兩個分區(qū),最終取得的性能結(jié)果。Pipeline-k 對應(yīng)的是 GPipe 將模型分成帶有 k 個加速器的 k 個分區(qū)的最終性能結(jié)果。
在無需更改超參數(shù)的情況下,GPipe 還能通過使用更多加速器來擴展訓(xùn)練結(jié)果。因此,它可以與數(shù)據(jù)并行技術(shù)相結(jié)合,通過互補的方式使用更多的加速器來擴展神經(jīng)網(wǎng)絡(luò)的訓(xùn)練規(guī)模。
精準度測試
我們試圖通過 GPipe 證明,對現(xiàn)有的神經(jīng)網(wǎng)絡(luò)進行擴展,可以實現(xiàn)更理想的模型質(zhì)量。為此,我們在 ImageNet ILSVRC-2012 數(shù)據(jù)集上訓(xùn)練一個模型參數(shù)為 5.57 億、輸入圖像尺寸為 480 x 480 的 AmoebaNet-B。該網(wǎng)絡(luò)被分為 4 個分區(qū),在模型與數(shù)據(jù)上執(zhí)行了并行訓(xùn)練程序。該巨型模型在沒有任何外部數(shù)據(jù)的情況下,最終達到了最先進的 84.3% top-1 / 97% top-5 的single-crop 驗證準確度結(jié)果。這說明大型的神經(jīng)網(wǎng)絡(luò)不僅適用于 ImageNet 等數(shù)據(jù)集,還能通過遷移學(xué)習(xí)的方式作用于其他數(shù)據(jù)集。事實證明,更好的 ImageNet 模型擁有更理想的傳輸效果。我們在 CIFAR10 和 CIFAR100 數(shù)據(jù)集上進行了遷移學(xué)習(xí)實驗。我們的巨型模型成功將 CIFAR-10 的精確度提高至到 99%、CIFAR-100 的精確度提高到 91.3%。
結(jié)論
當下許多機器學(xué)習(xí)應(yīng)用(如自動駕駛和醫(yī)學(xué)成像)得以持續(xù)發(fā)展并取得成功的原因,在于實現(xiàn)了盡可能高的模型精確度。然而這也意味著我們需要構(gòu)建一個更大、更復(fù)雜的模型,我們很高興能夠為研究社區(qū)提供 GPipe,我們希望未來它可以成為高效訓(xùn)練大規(guī)模 DNN 的基礎(chǔ)設(shè)施。