谷歌宣布開源神經(jīng)網(wǎng)絡(luò)模型高效訓(xùn)練庫(kù) GPipe

谷歌4日在博客中宣布開源大規(guī)模神經(jīng)網(wǎng)絡(luò)模型高效訓(xùn)練庫(kù) GPipe，這是一款分布式機(jī)器學(xué)習(xí)庫(kù)，可以讓研究員在不調(diào)整超參數(shù)的情況下，部署更多的加速器以對(duì)大規(guī)模模型進(jìn)行訓(xùn)練，有效擴(kuò)展了模型性能。對(duì)此進(jìn)行編譯如下。

深度神經(jīng)網(wǎng)絡(luò)（DNNs）推進(jìn)諸多機(jī)器學(xué)習(xí)任務(wù)的進(jìn)步，其中包括語(yǔ)音識(shí)別、視覺識(shí)別和語(yǔ)言處理等。BigGan、Bert 、GPT2.0 等最新成果表明，DNN 的模型越大，任務(wù)處理的表現(xiàn)就越好，而該結(jié)論也在過去的視覺識(shí)別任務(wù)中得到了驗(yàn)證，表明模型大小與分類準(zhǔn)確性之間存在很強(qiáng)的關(guān)聯(lián)性。舉個(gè)例子，2014 年 ImageNet 視覺識(shí)別挑戰(zhàn)賽的冠軍 GoogleNet 通過對(duì) 400 萬(wàn)參數(shù)進(jìn)行調(diào)整，最終取得 74.8 ％ 的精確度成績(jī)；僅僅過了三年，2017 年 ImageNet 挑戰(zhàn)賽冠軍 Squeeze-and-Excitation Networks 調(diào)整的參數(shù)便高達(dá) 1.458 億（36 倍以上），最終取得了 82.7％ 的精確度成績(jī)。與此對(duì)應(yīng)的是，市面的 GPU 內(nèi)存僅僅提高了 3 倍左右，目前最先進(jìn)的圖像模型早已達(dá)到云 TPUv2 可用內(nèi)存的極限。因此，我們迫切需要一種高效、可擴(kuò)展的基礎(chǔ)設(shè)施，以實(shí)現(xiàn)大規(guī)模的深度學(xué)習(xí)訓(xùn)練，并克服當(dāng)前的加速器內(nèi)存受限問題。

最具有代表性的圖像分類模型體現(xiàn)出了 ImageNet 精確度結(jié)果與模型大小的強(qiáng)關(guān)聯(lián)性

在《GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism》一文中，我們展示了如何通過流水并行技術(shù)（pipeline parallelism）對(duì) DNN 的訓(xùn)練進(jìn)行擴(kuò)展以克服這一限制。簡(jiǎn)單來(lái)說，GPipe 是一款分布式機(jī)器學(xué)習(xí)庫(kù)，基于同步隨機(jī)梯度下降與流水并行技術(shù)進(jìn)行模型訓(xùn)練，適用于由多個(gè)連續(xù)層組成的任意 DNN。最重要的是，GPipe 可以讓研究員在不調(diào)整超參數(shù)的情況下，部署更多的加速器以訓(xùn)練大規(guī)模模型，由此有效擴(kuò)展了性能。為了證明 GPipe 的能力，我們?cè)?Google Cloud TPUv2s 上訓(xùn)練了一個(gè)具有 5.57 億模型參數(shù)、 480 x 480 輸入圖像尺寸的 AmoebaNet-B。該模型在多個(gè)流行數(shù)據(jù)集上表現(xiàn)良好，取得的成就包括：single-crop ImageNet 的精確度提高至 84.3％、 CIFAR-10 的精確度提高至 99％、CIFAR-100 的精確度提高至 91.3％。

從小批次至微批次

目前存在兩種標(biāo)準(zhǔn)方法可以對(duì)中等規(guī)模的 DNN 模型進(jìn)行加速。數(shù)據(jù)并行方法（The data parallelism）可以納入更多的機(jī)器，并將輸入的數(shù)據(jù)區(qū)分開來(lái)。另一種方法則是將模型置于加速器上（比如 GPU 或 TPU）——這些加速器的特殊硬件可加速模型的訓(xùn)練進(jìn)程。然而加速器卻面臨著內(nèi)存與主機(jī)通信帶寬兩方面受限的問題。因此，通過將模型進(jìn)行分區(qū)，并根據(jù)分區(qū)配置相應(yīng)的加速器，模型并行技術(shù)可以讓我們?cè)诩铀倨魃嫌?xùn)練更大規(guī)模的 DNN 模型。由于 DNN 存在順序性，這種策略最后可能變成計(jì)算期間只有一個(gè)加速器處于活躍狀態(tài)，未能將加速器的計(jì)算能力充分利用起來(lái)。此外，標(biāo)準(zhǔn)的數(shù)據(jù)并行技術(shù)只允許在多個(gè)加速器上同時(shí)訓(xùn)練具有不同輸入數(shù)據(jù)的相同模型，卻無(wú)法提升加速器所能支持的最大模型規(guī)模。

為了實(shí)現(xiàn)跨加速器的高效訓(xùn)練，GPipe 先按照加速器對(duì)模型進(jìn)行劃分，然后自動(dòng)將小批次的訓(xùn)練示例拆分為更小的微批次。通過在微批次中執(zhí)行流水管理，加速器得以并行運(yùn)行。此外，梯度將在微批次中持續(xù)累積，以免分區(qū)的數(shù)量影響到模型的質(zhì)量。

上圖：由于網(wǎng)絡(luò)存在順序性，模型并行策略導(dǎo)致嚴(yán)重的利用不充分問題。 每次都只有一個(gè)加速器處于活動(dòng)狀態(tài)。

下圖：GPipe 將輸入的小批次拆成更小的微批次，使不同的加速器可以同時(shí)在單獨(dú)的微批次上運(yùn)作。

內(nèi)存和效率的最大化

GPipe 會(huì)對(duì)模型參數(shù)的內(nèi)存分配進(jìn)行最大化處理。我們?cè)诿總€(gè) TPUv2 均配有 8 個(gè)加速器核心以及 64 GB 內(nèi)存（每個(gè)加速器 8 GB）的云 TPUv2 上做了相關(guān)實(shí)驗(yàn)。如果沒有 GPipe，由于內(nèi)存的限制，單個(gè)加速器頂多只能訓(xùn)練 8200 萬(wàn)個(gè)模型參數(shù)。而通過反向傳播以及批量分割技術(shù)進(jìn)行重新計(jì)算的 GPipe ，成功將中間激活內(nèi)存從 6.26 GB 減至 3.46GB，由此實(shí)現(xiàn)單個(gè)加速器上訓(xùn)練 3.18 億個(gè)參數(shù)的成果。此外，我們還發(fā)現(xiàn)在流水并行技術(shù)的作用下，模型的最大規(guī)模與分區(qū)數(shù)量成正比，正如事前所預(yù)料的那樣?？偟膩?lái)說，GPipe 使 AmoebaNet 能在云 TPUv2 的 8 個(gè)加速器上納入 18 億個(gè)參數(shù)，比起之前高出了 25 倍。

為了測(cè)試模型的效率，我們研究了 GPipe 對(duì) AmoebaNet-D 模型吞吐量的影響情況。由于訓(xùn)練過程需要至少兩個(gè)加速器以適應(yīng)模型尺寸，因此我們只能對(duì)沒有實(shí)施流水并行技術(shù)的兩個(gè)分區(qū)案例的加速情況進(jìn)行觀察。我們發(fā)現(xiàn)訓(xùn)練過程存在近乎線性的加速效果。與兩個(gè)分區(qū)案例相比，將模型分布在四倍數(shù)量的加速器上能有效實(shí)現(xiàn) 3.5 倍的加速效果。我們的實(shí)驗(yàn)均使用了云 TPUv2，但我們了解到最新的的云 TPUv3 由于每個(gè) TPUv3 均配備了 16 個(gè)加速器核心以及 256 GB（每個(gè)加速器 16 GB），因此擁有更理想的表現(xiàn)性能。當(dāng)我們?cè)谒?16 個(gè)加速器上對(duì)模型進(jìn)行分發(fā)，GPipe 能讓基于 1024-token 句子的 80 億參數(shù) Transformer 語(yǔ)言模型的訓(xùn)練速度提高 11 倍。

通過 GPipe 對(duì) AmoebaNet-D 進(jìn)行加速。這種模型不適用于單個(gè)加速器。naive-2 基線是將模型拆分為兩個(gè)分區(qū)，最終取得的性能結(jié)果。Pipeline-k 對(duì)應(yīng)的是 GPipe 將模型分成帶有 k 個(gè)加速器的 k 個(gè)分區(qū)的最終性能結(jié)果。

在無(wú)需更改超參數(shù)的情況下，GPipe 還能通過使用更多加速器來(lái)擴(kuò)展訓(xùn)練結(jié)果。因此，它可以與數(shù)據(jù)并行技術(shù)相結(jié)合，通過互補(bǔ)的方式使用更多的加速器來(lái)擴(kuò)展神經(jīng)網(wǎng)絡(luò)的訓(xùn)練規(guī)模。

精準(zhǔn)度測(cè)試

我們?cè)噲D通過 GPipe 證明，對(duì)現(xiàn)有的神經(jīng)網(wǎng)絡(luò)進(jìn)行擴(kuò)展，可以實(shí)現(xiàn)更理想的模型質(zhì)量。為此，我們?cè)?ImageNet ILSVRC-2012 數(shù)據(jù)集上訓(xùn)練一個(gè)模型參數(shù)為 5.57 億、輸入圖像尺寸為 480 x 480 的 AmoebaNet-B。該網(wǎng)絡(luò)被分為 4 個(gè)分區(qū)，在模型與數(shù)據(jù)上執(zhí)行了并行訓(xùn)練程序。該巨型模型在沒有任何外部數(shù)據(jù)的情況下，最終達(dá)到了最先進(jìn)的 84.3％ top-1 / 97％ top-5 的single-crop 驗(yàn)證準(zhǔn)確度結(jié)果。這說明大型的神經(jīng)網(wǎng)絡(luò)不僅適用于 ImageNet 等數(shù)據(jù)集，還能通過遷移學(xué)習(xí)的方式作用于其他數(shù)據(jù)集。事實(shí)證明，更好的 ImageNet 模型擁有更理想的傳輸效果。我們?cè)?CIFAR10 和 CIFAR100 數(shù)據(jù)集上進(jìn)行了遷移學(xué)習(xí)實(shí)驗(yàn)。我們的巨型模型成功將 CIFAR-10 的精確度提高至到 99％、CIFAR-100 的精確度提高到 91.3％。

結(jié)論

當(dāng)下許多機(jī)器學(xué)習(xí)應(yīng)用（如自動(dòng)駕駛和醫(yī)學(xué)成像）得以持續(xù)發(fā)展并取得成功的原因，在于實(shí)現(xiàn)了盡可能高的模型精確度。然而這也意味著我們需要構(gòu)建一個(gè)更大、更復(fù)雜的模型，我們很高興能夠?yàn)檠芯可鐓^(qū)提供 GPipe，我們希望未來(lái)它可以成為高效訓(xùn)練大規(guī)模 DNN 的基礎(chǔ)設(shè)施。