背景：
目前，大模型的發(fā)展已經(jīng)非?；馃幔P于大模型的訓練、微調(diào)也是各個公司重點關注方向。但是大模型訓練的痛點是模型參數(shù)過大，動輒上百億，如果單靠單個GPU來完成訓練基本不可能。所以需要多卡或者分布式訓練來完成這項工作。
一、分布式訓練
1.1 目前主流的大模型分布式訓練主要包括兩種：
- 數(shù)據(jù)并行訓練
- 模型并行訓練
二、DeepSpeed
DeepSpeed是由Microsoft提供的分布式訓練工具，旨在支持更大規(guī)模的模型和提供更多的優(yōu)化策略和工具。對于更大模型的訓練來說，DeepSpeed提供了更多策略，例如：Zero、Offload等。
2.1 基礎組件
分布式訓練需要掌握分布式環(huán)境中的基礎配置，包括節(jié)點變化、全局進程編號、局部進程編號、全局總進程數(shù)、主節(jié)點等。這些組件都跟分布式訓練緊密相關，同時組件之間也有非常大的聯(lián)系，例如通信聯(lián)系等。
2.2 通信策略
既然是分布式訓練，那機器之間必須要保持通信，這樣才可以傳輸模型參數(shù)，梯度參數(shù)等信息。
DeepSpeed提供了mpi、gioo、nccl等通信策略
| 通信策略 | 通信作用 |
|---|---|
| mpi | 它是一種跨界點的通信庫，經(jīng)常用于CPU集群的分布式訓練 |
| gloo | 它是一種高性能的分布式訓練框架，可以支持CPU或者GPU的分布式訓練 |
| nccl | 它是nvidia提供的GPU專用通信庫，廣泛用于GPU上的分布式訓練 |
我們在使用DeepSpeed進行分布式訓練的時候，可以根據(jù)自身的情況選擇合適的通信庫，通常情況下，如果是GPU進行分布式訓練，可以選擇nccl。
2.3 Zero（零冗余優(yōu)化器）
Microsoft開發(fā)的Zero可以解決分布式訓練過程中數(shù)據(jù)并行和模型并行的限制。比如： Zero通過在數(shù)據(jù)并行過程中劃分模型狀態(tài)（優(yōu)化器、梯度、參數(shù)），來解決數(shù)據(jù)并行成可能出現(xiàn)內(nèi)存冗余的情況（正常數(shù)據(jù)并行訓練，模型全部參數(shù)是復制在各個機器上的）；同時可以在訓練期間使用動態(tài)通信計劃，在分布式設備之間共享重要的狀態(tài)變量，這樣保持計算粒度和數(shù)據(jù)并行的通信量。
Zero是用于大規(guī)模模型訓練優(yōu)化的技術，它的主要目的是減少模型的內(nèi)存占用，讓模型可以在顯卡上訓練，內(nèi)存占用主要分為Model States和Activation兩個部分，Zero主要解決的是Model States的內(nèi)存占用問題。
Zero將模型參數(shù)分成三個部分：
| 狀態(tài) | 作用 |
|---|---|
| Optimizer States | 優(yōu)化器在進行梯度更新的時候需要用到的數(shù)據(jù) |
| Gradient | 在反向轉(zhuǎn)播過程中產(chǎn)生的數(shù)據(jù)，其決定參數(shù)的更新方向 |
| Model Parameter | 模型參數(shù)，在模型訓練過程中通過數(shù)據(jù)“學習”的信息 |
Zero的級別如下：
| 級別 | 作用 |
|---|---|
| Zero-0 | 不使用所有類型的分片，僅使用DeepSpeed作為DDP |
| Zero-1 | 分割Optimizer States， 減少4倍內(nèi)存，通信容量和數(shù)據(jù)并行性相同 |
| Zero-2 | 分割Optimizer States和Gradients，減少8倍內(nèi)存，通信容量和數(shù)據(jù)并行性相同 |
| Zero-3 | 分割Optimizer States、gradients、Parametes，內(nèi)存減少與數(shù)據(jù)并行度呈線性關系。例如，在64個GPU（Nd=64）之間進行拆分將產(chǎn)生64倍的內(nèi)存縮減。通信量有50%的適度增長 |
| Zero-Infinity | Zero-Infinity是Zero-3的擴展，它允許通過使用 NVMe 固態(tài)硬盤擴展 GPU 和 CPU 內(nèi)存來訓練大型模型 |
2.4 Zero-Offload：
相比GPU，CPU就相對比較廉價，所以Zero-Offload思想是將訓練階段的某些模型狀態(tài)放（offload）到內(nèi)存以及CPU計算。
Zero-Offload不希望為了最小化顯存占用而讓系統(tǒng)計算效率下降，但如果使用CPU也需要考慮通信和計算的問題（通信：GPU和CPU的通信；計算：CPU占用過多計算就會導致效率降低）。
Zero-Offload想做的是把計算節(jié)點和數(shù)據(jù)節(jié)點分布在GPU和CPU上，計算節(jié)點落到哪個設備上，哪個設備就執(zhí)行計算，數(shù)據(jù)節(jié)點落到哪個設備上，哪個設備就負責存儲。
Zero-Offload切分思路：
下圖中有四個計算類節(jié)點：FWD、BWD、Param update和float2half，前兩個計算復雜度大致是 O(MB)， B是batch size，后兩個計算復雜度是 O(M)。為了不降低計算效率，將前兩個節(jié)點放在GPU，后兩個節(jié)點不但計算量小還需要和Adam狀態(tài)打交道，所以放在CPU上，Adam狀態(tài)自然也放在內(nèi)存中，為了簡化數(shù)據(jù)圖，將前兩個節(jié)點融合成一個節(jié)點FWD-BWD Super Node，將后兩個節(jié)點融合成一個節(jié)點Update Super Node。如下圖右邊所示，沿著gradient 16和parameter 16兩條邊切分。
Zero-Offload計算思路：
在GPU上面進行前向和后向計算，將梯度傳給CPU，進行參數(shù)更新，再將更新后的參數(shù)傳給GPU。為了提高效率，可以將計算和通信并行起來，GPU在反向傳播階段，可以待梯度值填滿bucket后，一遍計算新的梯度一遍將bucket傳輸給CPU，當反向傳播結束，CPU基本上已經(jīng)有最新的梯度值了，同樣的，CPU在參數(shù)更新時也同步將已經(jīng)計算好的參數(shù)傳給GPU，如下圖所示。
2.5 混合精度：
混合精度訓練是指在訓練過程中同時使用FP16（半精度浮點數(shù)）和FP32（單精度浮點數(shù)）兩種精度的技術。使用FP16可以大大減少內(nèi)存占用，從而可以訓練更大規(guī)模的模型。但是，由于FP16的精度較低，訓練過程中可能會出現(xiàn)梯度消失和模型坍塌等問題。
DeepSpeed支持混合精度的訓練，可以在config.json配置文件中設置來啟動混合精度（"fp16.enabled":true）。在訓練的過程中，DeepSpeed會自動將一部分操作轉(zhuǎn)化為FP16格式，并根據(jù)需要動態(tài)調(diào)整精度縮放因子，來保證訓練的穩(wěn)定性和精度。
在使用混合精度訓練時，需要注意一些問題，例如梯度裁剪（Gradient Clipping）和學習率調(diào)整（Learning Rate Schedule）等。梯度裁剪可以防止梯度爆炸，學習率調(diào)整可以幫助模型更好地收斂。
三、總結
DeepSpeed方便了我們在機器有限的情況下來訓練、微調(diào)大模型，同時它也有很多優(yōu)秀的性能來使用，后期可以繼續(xù)挖掘。
目前主流的達模型訓練方式： GPU + PyTorch + Megatron-LM + DeepSpeed
優(yōu)勢
- 存儲效率：DeepSpeed提供了一種Zero的新型解決方案來減少訓練顯存的占用，它與傳統(tǒng)的數(shù)據(jù)并行不同，它將模型狀態(tài)和梯度進行分區(qū)來節(jié)省大量的顯存；
- 可擴展性：DeepSpeed支持高效的數(shù)據(jù)并行、模型并行、pipeline并行以及它們的組合，這里也稱3D并行；
- 易用性： 在訓練階段，只需要修改幾行代碼就可以使pytorch模型使用DeepSpeed和Zero。
參考：
1. http://wed.xjx100.cn/news/204072.html?action=onClick
2. https://zhuanlan.zhihu.com/p/513571706
作者：京東物流 鄭少強文章來源：http://www.zghlxwxcb.cn/news/detail-711374.html
來源：京東云開發(fā)者社區(qū) 轉(zhuǎn)載請注明來源文章來源地址http://www.zghlxwxcb.cn/news/detail-711374.html
到了這里，關于DeepSpeed： 大模型訓練框架的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關文章，希望大家以后多多支持TOY模板網(wǎng)！
