現(xiàn)在的模型動輒數(shù)百、數(shù)千億參數(shù)，普通人訓不動怎么辦？

前不久，谷歌發(fā)布了參數(shù)量為 1.6 萬億的語言模型Swith Transformer，將 GPT-3 創(chuàng)下的參數(shù)量記錄（1750 億）推至新高。這些大模型的出現(xiàn)讓普通研究者越發(fā)絕望：沒有「鈔能力」、沒有一大堆 GPU 就做不了 AI 研究了嗎？

在此背景下，部分研究者開始思考：如何讓這些大模型的訓練變得更加接地氣？也就是說，怎么用更少的卡訓練更大的模型？

為了解決這個問題，來自微軟、加州大學默塞德分校的研究者提出了一種名為「 ZeRO-Offload 」的異構(gòu)深度學習訓練技術(shù)，可以在單個 GPU 上訓練擁有 130 億參數(shù)的深度學習模型，讓普通研究者也能著手大模型的訓練。與 Pytorch 等流行框架相比，ZeRO-Offload 將可訓練的模型規(guī)模提升了 10 倍，而且不需要數(shù)據(jù)科學家對模型做出任何改變，也不會犧牲計算效率。

論文鏈接：https://arxiv.org/pdf/2101.06840.pdf

ZeRO-Offload 通過將數(shù)據(jù)和計算卸載（offload）至 CPU 來實現(xiàn)大規(guī)模模型訓練。為了不降低計算效率，它被設(shè)計為最小化與 GPU 之間的數(shù)據(jù)往來，并在盡可能節(jié)省 GPU 內(nèi)存的同時降低 CPU 的計算時間。因此，對于一個參數(shù)量為 100 億的模型，ZeRO-Offload 可以在單個 NVIDIA V100 GPU 上實現(xiàn) 40 TFlops/GPU。相比之下，使用 PyTorch 訓練一個參數(shù)量為 14 億的模型僅能達到 30TFlops，這是在不耗盡內(nèi)存的情況下所能訓練的最大模型。ZeRO-Offload 還可以擴展至多 GPU 設(shè)置并實現(xiàn)線性加速，最多可在 128 個 GPU 上實現(xiàn)近似線性加速。

此外，ZeRO-Offload 還可以和模型并行一起使用，在一個 DGX-2 box AI 服務(wù)器上訓練參數(shù)量超 700 億的模型。與單獨使用模型并行相比，這一參數(shù)量實現(xiàn)了 4.5 倍的規(guī)模提升。

在下文中，我們將結(jié)合 Medium 博主 LORENZ KUHN 的一篇博客來詳細了解這篇論文。

ZeRO-Offload 是什么？

ZeRO-Offload 是一種通過將數(shù)據(jù)和計算從 GPU 卸載到 CPU，以此減少神經(jīng)網(wǎng)絡(luò)訓練期間 GPU 內(nèi)存占用的方法，該方法提供了更高的訓練吞吐量，并避免了移動數(shù)據(jù)和在 CPU 上執(zhí)行計算導(dǎo)致的減速問題。

借助 ZeRO-offload，使用相同的硬件能訓練以往 10 倍大的模型，即使在單個 GPU 上也是如此。比如在一個 32GB RAM 的 V100 GPU 上訓練百億參數(shù)的 GPT-2。

此外，ZeRO-offload 還能實現(xiàn)在多 GPU 設(shè)置中的近似線性擴展。

對于研究者來說，ZeRO-offload 適用的情況包括：

想訓練更大的模型，或者想更快地訓練現(xiàn)在的模型，因為 ZeRO-offload 允許訓練更大的 batch size；

你正在使用 PyTorch，并且愿意 / 能夠使用微軟的 DeepSpeed 庫（ZeRO-offload 的其他實現(xiàn)形式暫未推出），你也可以嘗試根據(jù)官方實現(xiàn)自行調(diào)整；

愿意接受一些建模時的限制，比如當前版本的 ZeRO-Offload 需要搭配使用 Adam 的混合精度訓練。

如何使用？

ZeRO-Offload 在微軟的 DeepSpeed 庫中實現(xiàn)，官方實現(xiàn)地址：https://github.com/microsoft/DeepSpeed/blob/6e65c2cc084ecfc393c67a2f64639e8d08d325f6/deepspeed/runtime/zero/stage2.py。

在 DeepSpeed 中設(shè)置完畢后，使用 ZeRO-Offload 就不需要太多額外的工作了，只需要修改一些標志和配置文件。

目前，Hugging Face 的 transformers 庫與 DeepSpeed 進行了實驗性集成，使用方法和基準測試結(jié)果參見：https://huggingface.co/blog/zero-deepspeed-fairscale。

Facebook 研究院的 fairscale 有 ZeRO 的部分實現(xiàn)，ZeRO-Offload 正是基于 ZeRO 這一多 GPU 內(nèi)存優(yōu)化方法構(gòu)建的。目前還不支持 CPU 卸載。

ZeRO-Offload 的工作原理

ZeRO-Offload 是基于 Zero Redundancy Optimizer (ZeRO) 構(gòu)建的。ZeRO 是微軟在 2020 年 2 月提出的一種萬億級模型參數(shù)訓練方法，用于數(shù)據(jù)并行和模型并行訓練中的內(nèi)存優(yōu)化，其中梯度、參數(shù)和優(yōu)化器狀態(tài)分布在多 GPU 內(nèi)存中，沒有任何冗余。這使得 GPU 之間的通信開銷保持在比較低的狀態(tài)。

與標準數(shù)據(jù)并行基準相比，ZeRO 在三個階段中節(jié)省的內(nèi)存和通信用量。

讓我們來回顧一下 ZeRO：

為了解決數(shù)據(jù)并行和模型并行存在的問題，ZeRO 提供了三階段的優(yōu)化方法，分別為優(yōu)化器狀態(tài)分割、梯度分割、參數(shù)分割，三個階段按順序?qū)嵤?/p>

在優(yōu)化器分割狀態(tài)：ZeRO 降低了 3/4 的內(nèi)存，通信量和數(shù)據(jù)并行相同；

加入梯度分割：降低了 7/8 的內(nèi)存，通信量和數(shù)據(jù)并行相同；

加入?yún)?shù)分割：內(nèi)存減少與數(shù)據(jù)并行度呈線性關(guān)系。例如，在 64 個 GPU 上進行分割的時候，可以將內(nèi)存降至 1/64。在通信量上有 50% 的提升。

在去年 9 月份的博客中，微軟這么介紹 ZeRO-Offload：

ZeRO-Offload 繼承了 ZeRO-2 的優(yōu)化器狀態(tài)和梯度分割。但與 ZeRO-2 不同的是，ZeRO-Offload 不在每塊 GPU 上保持優(yōu)化器狀態(tài)和梯度的分割，而是將二者卸載至主機 CPU 內(nèi)存。在整個訓練階段，優(yōu)化器狀態(tài)都保存在 CPU 內(nèi)存中；而梯度則在反向傳播過程中在 GPU 上利用 reduce-scatter 進行計算和求均值，然后每個數(shù)據(jù)并行線程將屬于其分割的梯度平均值卸載到 CPU 內(nèi)存中（參見下圖 g offload），將其余的拋棄。一旦梯度到達 CPU，則每個數(shù)據(jù)并行線程直接在 CPU 上并行更新優(yōu)化器狀態(tài)分割（參見下圖 p update）。

之后，將參數(shù)分割移回 GPU，再在 GPU 上執(zhí)行 all-gather 操作，收集所有更新后的參數(shù)（參見下圖 g swap）。ZeRO-Offload 還利用單獨的 CUDA 流來窮盡通信與計算中的重疊，從而最大化訓練效率。

ZeRO-Offload 概覽。

值得注意的是，ZeRO-Offload 專為使用 Adam 的混合精度訓練而設(shè)計。也就是說，當前版本的 ZeRO-Offload 使用 Adam 的優(yōu)化版本 DeepCPUAdam。其主要原因是避免 CPU 計算成為整個過程中的瓶頸。DeepCPUAdam 的速度是 Adam PyTorch 實現(xiàn)的 6 倍。

實驗結(jié)果

最后來看一下 ZeRO-Offload 論文中提供的一些實驗結(jié)果。

下圖 7 展示了利用 ZeRO-Offload 技術(shù)在 1 個、4 個或 16 個 GPU（一個 DGX-2）上可以訓練的最大模型情況。

下圖 11 展示了每個 GPU 的吞吐量隨 GPU 數(shù)量增加而呈現(xiàn)的變化情況。可以看出，在 GPU 數(shù)量逐漸增加至 128 個的過程中，ZeRO-Offload 幾乎可以實現(xiàn)吞吐量的線性加速。

下圖 8 展示了使用 PyTorch、L2L 和 ZeRO-Offload 實現(xiàn)的每個 GPU 吞吐量差異。從中可以看出，利用 ZeRO-Offload 實現(xiàn)的每個 GPU 吞吐量比 L2L 平均高出 14%（最多高出 22%）。