現在的模型動輒數百、數千億參數，普通人訓不動怎么辦？

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

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

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

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

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

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

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

ZeRO-Offload 是什么？

ZeRO-Offload 是一種通過將數據和計算從 GPU 卸載到 CPU，以此減少神經網絡訓練期間 GPU 內存占用的方法，該方法提供了更高的訓練吞吐量，并避免了移動數據和在 CPU 上執行計算導致的減速問題。

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

此外，ZeRO-offload 還能實現在多 GPU 設置中的近似線性擴展。

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

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

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

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

如何使用？

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

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

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

Facebook 研究院的 fairscale 有 ZeRO 的部分實現，ZeRO-Offload 正是基于 ZeRO 這一多 GPU 內存優化方法構建的。目前還不支持 CPU 卸載。

ZeRO-Offload 的工作原理

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

與標準數據并行基準相比，ZeRO 在三個階段中節省的內存和通信用量。

讓我們來回顧一下 ZeRO：

為了解決數據并行和模型并行存在的問題，ZeRO 提供了三階段的優化方法，分別為優化器狀態分割、梯度分割、參數分割，三個階段按順序實施。

在優化器分割狀態：ZeRO 降低了 3/4 的內存，通信量和數據并行相同；

加入梯度分割：降低了 7/8 的內存，通信量和數據并行相同；

加入參數分割：內存減少與數據并行度呈線性關系。例如，在 64 個 GPU 上進行分割的時候，可以將內存降至 1/64。在通信量上有 50% 的提升。

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

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

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

ZeRO-Offload 概覽。

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

實驗結果

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

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

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

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