鵝廠內部萬億大模型訓練加速框架上線公有云！

自研 AngelPTM 大模型訓練框架上線公有云，助力降本增效

最近 ChatGPT 憑藉其強大的語言理解能力、文本生成能力、對話能力等在多個領域均取得了巨大成功，掀起了新一輪的人工智能浪潮。ChatGPT、GPT3、以及 GPT3.5 都是基於 Transformer 架構堆疊而成，研究發現隨著訓練數據量和模型容量的增加可以持續提高模型的泛化能力和表達能力，研究大模型成為了近兩年的趨勢。國內外頭部科技公司均有佈局，發佈了若干千億規模以上的大模型，如下圖所示：

近幾年 NLP 預訓練模型規模的發展，模型已經從億級發展到了萬億級參數規模。具體來說，2018 年 BERT 模型最大參數量為3.4億，2019 年 GPT-2 為十億級參數的模型。2020 年發佈的百億級規模有 T5 和 T-NLG，以及千億參數規模的 GPT-3。在 2021 年末，Google 發佈了 SwitchTransformer，首次將模型規模提升至萬億。然而 GPU 硬件發展的速度難以滿足 Transformer 模型規模發展的需求。近四年中，模型參數量增長了十萬倍，但 GPU 的顯存僅增長了4倍。

舉例來說，萬億模型的模型訓練僅參數和優化器狀態便需要 1.7TB 以上的存儲空間，至少需要 425 張 A100-40G，這還不包括訓練過程中產生的激活值所需的存儲。在這樣的背景下，大模型訓練不僅受限於海量的算力, 更受限於巨大的存儲需求。

業內大模型預訓練的解決方案主要包括微軟的 DeepSpeed 和英偉達的 Megatron-LM。DeepSpeed 引入 ZeRO（Zero Redundancy Optimizer ）優化器，將模型參數、梯度、優化器狀態按需分配到不同的訓練卡上，滿足大模型對存儲的極致要求；Megatron-LM 基於3D並行（張量並行、流水並行、數據並行）將模型參數進行切分，滿足大模型在有限顯存資源下的訓練訴求。

騰訊內部也有非常多的大模型預訓練業務，為了以最小的成本和最快的性能訓練大模型，太極機器學習平臺對 DeepSpeed 和 Megatron-LM 進行了深度定製優化，推出了 AngelPTM 訓練框架。2022 年 4 月騰訊發佈的混元 AI 大模型便是基於 AngelPTM 框架訓練而來。鑑於最近大模型的火熱趨勢，我們決定將內部成熟落地的 AngelPTM 框架推廣給廣大公有云用戶，幫助業務降本增效。

AngelPTM 技術原理簡介

ZeRO-Cache 優化策略

ZeRO-Cache 是一款超大規模模型訓練的利器，如下圖所示，其通過統一視角去管理內存和顯存，在去除模型狀態冗餘的同時擴大單個機器的可用存儲空間上限，通過 Contiguous Memory 顯存管理器管理模型參數的顯存分配/釋放進而減少顯存碎片，通過多流均衡各個硬件資源的負載，通過引入 SSD 進一步擴展單機模型容量。

統一視角存儲管理

大模型訓練時模型狀態都位於 CPU 內存中，在訓練時會拷貝到 GPU 顯存，這就導致模型狀態的冗餘存儲（CPU 和 GPU 同時存在一份），此外大模型的訓練中會用到大量的 pin memory，pin memory 的使用會提升性能，但同時也會導致物理內存的大量浪費，如何科學合理的使用 pin memory 是 ZeRO-Cache 著重要解決的問題。

本著極致化去冗餘的理念，AngelPTM 引入了 chunk 對內存和顯存進行管理，保證所有模型狀態只存儲一份，通常模型會存儲在內存 or 顯存上，ZeRO-Cache 引入異構統一存儲，採用內存和顯存共同作為存儲空間，擊破了異構存儲的壁壘，極大擴充了模型存儲可用空間，如下圖左圖所示：

在 CPU 時原生 Tensor 的底層存儲機制對於實際佔用的內存空間利用率極不穩定，對此問題 AngelPTM 實現了 Tensor 底層分片存儲的機制，在極大的擴展了單機可用存儲空間的同時，避免了不必要的 pin memory 存儲浪費，使得單機可負載的模型上限獲得了極大提升，如上圖右圖所示。

ZeRO-Cache 顯存管理器

PyTorch 自帶的顯存管理器可以 Cache 幾乎所有顯存進行二次快速分配，在顯存壓力不大的情況下這種顯存分配方式可以達到性能最優，但是對於超大規模參數的模型，導致顯存壓力劇增，且由於參數梯度頻繁的顯存分配導致顯存碎片明顯增多，PyTorch Allocator 嘗試分配顯存失敗的次數增加，導致訓練性能急劇下降。

為此，ZeRO-Cache 引入了一個 Contiguous Memory 顯存管理器，如下圖所示，其在 PyTorch Allocator 之上進行二次顯存分配管理，模型訓練過程中參數需要的顯存的分配和釋放都由 Contiguous Memory 統一管理，在實際的大模型的訓練中其相比不使用 Contiguous Memory 顯存分配效率以及碎片有顯著提升，模型訓練速度有質的飛越。

PipelineOptimizer

ZeRO-Infinity 利用 GPU 或者 CPU 更新模型的參數，特別是針對大模型，只可以通過 CPU 來更新參數，由於 CPU 更新參數的速度相比 GPU 更新參數有數十倍的差距，且參數更新幾乎可以佔到整個模型訓練時間的一半，在 CPU 更新參數的時候 GPU 空閒且顯存閒置，造成資源的極大浪費。

如下圖所示，ZeRO-Cache 會在模型參數算出梯度之後開始 Cache 模型的優化器狀態到 GPU 顯存，並在參數更新的時候異步 Host 和 Device 之間的模型狀態數據傳輸，同時支持 CPU 和 GPU 同時更新參數。ZeRO-Cache pipeline 了模型狀態 H2D、參數更新、模型狀態 D2H，最大化地利用硬件資源，避免硬件資源閒置。

此外，AngelPTM 自研了異構 Adafactor 優化器，如下圖所示，支持 CPU 和 GPU 同時進行參數的更新，其可以減少 33% 的模型狀態存儲空間，同時可提高模型訓練精度。

多流異步化

大模型訓練過程中有大量的計算和通信，包括 GPU 計算、H2D 和 D2H 單機通信、NCCL 多機通信等，涉及的硬件有 GPU、CPU、PCIE 等。如下圖所示，ZeRO-Cache 為了最大化地利用硬件，多流異步化 GPU 計算、H2D 和 D2H 單機通信、NCCL 多機通信，參數預取採用用時同步機制，梯度後處理採用多 buffer 機制，優化器狀態拷貝採用多流機制。

ZeRO-Cache SSD 框架

為了更加低成本地擴展模型參數，如下圖所示，ZeRO-Cache 進一步引入了 SSD 作為三級存儲，針對 GPU 高計算吞吐、高通信帶寬和 SSD 低 PCIE 帶寬之間的 GAP，ZeRO-Cache 放置所有 fp16 參數和梯度到內存中，讓 foward 和 backward 的計算不受 SSD 低帶寬影響，同時通過對優化器狀態做半精度壓縮來緩解 SSD 讀寫對性能的影響。