使用原始的 Megatron-LM 訓練 GPT2

設(shè)置訓練數(shù)據(jù)

運行未修改的Megatron-LM GPT2模型

開啟DeepSpeed

DeepSpeed 使用 GPT-2 進行評估

Zero概述

訓練環(huán)境

開啟Zero優(yōu)化

訓練一個1.5B參數(shù)的GPT2模型

訓練一個10b的GPT-2模型

使用ZeRO-Infinity訓練萬億級別的模型

使用ZeRO-Infinity將計算轉(zhuǎn)移到CPU和NVMe

分配大規(guī)模Megatron-LM模型

以內(nèi)存為中心的分塊優(yōu)化

提取權(quán)重

ZeRO-Offload概述

訓練環(huán)境

在單個 V100 GPU 上訓練10B的GPT2模型

Megatron-LM GPT-2 的啟動腳本更改：

DeepSpeed 配置更改

0x0. 前言

這篇文章主要翻譯DeepSpeed的Megatron-LM GPT2 ，Zero零冗余優(yōu)化器技術(shù)，ZeRO-Offload技術(shù)。關(guān)于DeepSpeed 的Zero和ZeRO-Offload的技術(shù)原理大家也可以查看圖解大模型訓練之：數(shù)據(jù)并行下篇(ZeRO，零冗余優(yōu)化) 這篇文章，文章里面對內(nèi)存的計算和通信量的分析都很棒。

0x1. Megatron-LM GPT2

如果你還沒有閱讀過入門指南，我們建議你在開始本教程之前先閱讀該指南（https://www.deepspeed.ai/getting-started/ 這個指南的翻譯在 【DeepSpeed 教程翻譯】開始，安裝細節(jié)和CIFAR-10 Tutorial）。

在本教程中，我們將向 Megatron-LM GPT2 模型添加 DeepSpeed，Megatron-LM GPT2 是一個大而強的 transformer。Megatron-LM 支持模型并行和多節(jié)點訓練。有關(guān)更多詳細信息，請參閱相應的論文：Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（https://arxiv.org/abs/1909.08053）。

首先，我們討論數(shù)據(jù)和環(huán)境設(shè)置，以及如何使用原始的Megatron-LM訓練GPT-2模型。接下來，我們逐步介紹如何使用DeepSpeed使該模型運行。最后，我們演示使用DeepSpeed所獲得的性能提升和內(nèi)存占用減少。

使用原始的 Megatron-LM 訓練 GPT2

我們將原始的Megatron-LM（https://github.com/NVIDIA/Megatron-LM）模型代碼復制到DeepSpeed Megatron-LM（https://github.com/microsoft/Megatron-DeepSpeed）中，并將其作為子模塊提供。要下載，請執(zhí)行以下操作：

git submodule update --init --recursive

設(shè)置訓練數(shù)據(jù)

按照Megatron的說明（https://github.com/NVIDIA/Megatron-LM#collecting-gpt-webtext-data）下載webtext數(shù)據(jù)，并在DeepSpeedExamples/Megatron-LM/data(在最新版本的DeepSpeedExamples中可以放置在 /home/zhangxiaoyu/DeepSpeedExamples/training/megatron )下放置一個符號鏈接。

運行未修改的Megatron-LM GPT2模型

對于單塊GPU：

修改 scripts/pretrain_gpt2.sh，將其 --train-data 參數(shù)設(shè)置為 "webtext"。

運行 bash scripts/pretrain_gpt2.sh

對于多個節(jié)點和多個GPU：

修改 scripts/pretrain_gpt2_model_parallel.sh

將其 --train-data 參數(shù)設(shè)置為 "webtext"。

GPUS_PER_NODE 指示了測試中每個節(jié)點使用的GPU數(shù)量。

NNODES 指示了測試中涉及的節(jié)點數(shù)量。

執(zhí)行 bash scripts/pretrain_gpt2_model_parallel.sh

開啟DeepSpeed

為了用上DeepSpeed需要更新三個文件：

arguments.py : 參數(shù)配置文件

pretrain_gpt2.py : 訓練的主入口點

utils.py : 模型保存和加載工具

參數(shù)解析

第一步是在 arguments.py 中使用 deepspeed.add_config_arguments() 將 DeepSpeed 參數(shù)添加到 Megatron-LM GPT2 模型中。

初始化和訓練

我們將修改 pretrain.py 以啟用使用 DeepSpeed 進行訓練。

初始化

我們使用 deepspeed.initialize 創(chuàng)建 model_engine、optimizer 和 LR scheduler。下面是其定義：

definitialize(args,
model,
optimizer=None,
model_parameters=None,
training_data=None,
lr_scheduler=None,
mpu=None,
dist_init_required=True,
collate_fn=None):

對于 Megatron-LM GPT2 模型，我們在其 setup_model_and_optimizer() 函數(shù)中進行 DeepSpeed 初始化，函數(shù)傳的參數(shù)包含原始model、optimizer、args、lr_scheduler 和 mpu。

請注意，當啟用FP16時，Megatron-LM GPT2會在Adam優(yōu)化器上添加一個包裝器。DeepSpeed有自己的FP16優(yōu)化器，因此我們需要直接將Adam優(yōu)化器傳遞給DeepSpeed，而不需要任何包裝器。當啟用DeepSpeed時，我們從 get_optimizer() 返回未包裝的Adam優(yōu)化器。

使用訓練API

由 deepspeed.initialize 返回的模型是 DeepSpeed 模型引擎，我們將基于該引擎使用 forward、backward 和 step API 訓練模型。

前向傳播

前向傳播API與PyTorch兼容，不需要進行任何更改。

反向傳播

通過在模型引擎上直接調(diào)用 backward(loss) 來進行反向傳播。

defbackward_step(optimizer,model,lm_loss,args,timers):
"""Backwardstep."""

#Totalloss.
loss=lm_loss

#Backwardpass.
ifargs.deepspeed:
model.backward(loss)
else:
optimizer.zero_grad()
ifargs.fp16:
optimizer.backward(loss,update_master_grads=False)
else:
loss.backward()

DeepSpeed會在使用小批量更新權(quán)重后自動處理梯度清零。此外，DeepSpeed在內(nèi)部解決了分布式數(shù)據(jù)并行和FP16，簡化了多個地方的代碼。

(A) DeepSpeed 還在梯度累積邊界處自動執(zhí)行梯度平均，因此我們跳過allreduce通信。

ifargs.deepspeed:
#DeepSpeed反向傳播已經(jīng)處理了allreduce通信。重置計時器以避免破壞下面的計時器日志。
timers('allreduce').reset()
else:
torch.distributed.all_reduce(reduced_losses.data)
reduced_losses.data=reduced_losses.data/args.world_size
ifnotUSE_TORCH_DDP:
timers('allreduce').start()
model.allreduce_params(reduce_after=False,
fp32_allreduce=args.fp32_allreduce)
timers('allreduce').stop()

(B) 我們也跳過更新主節(jié)點梯度，因為DeepSpeed在內(nèi)部解決了這個問題。

#Updatemastergradients.
ifnotargs.deepspeed:
ifargs.fp16:
optimizer.update_master_grads()

#Clippinggradientshelpspreventtheexplodinggradient.
ifargs.clip_grad>0:
ifnotargs.fp16:
mpu.clip_grad_norm(model.parameters(),args.clip_grad)
else:
optimizer.clip_master_grads(args.clip_grad)

returnlm_loss_reduced

更新模型參數(shù)

DeepSpeed引擎中的 step() 函數(shù)更新模型參數(shù)以及學習率。

ifargs.deepspeed:
model.step()
else:
optimizer.step()

#Updatelearningrate.
ifnot(args.fp16andoptimizer.overflow):
lr_scheduler.step()
else:
skipped_iter=1

損失縮放

GPT2訓練腳本在訓練過程中記錄了損失縮放值。在DeepSpeed優(yōu)化器內(nèi)部，該值存儲為 cur_scale，而不是Megatron的優(yōu)化器中的 loss_scale。因此，我們在日志字符串中適當?shù)剡M行了替換。

ifargs.fp16:
log_string+='lossscale{:.1f}|'.format(
optimizer.cur_scaleifargs.deepspeedelseoptimizer.loss_scale)

檢查點保存和加載

DeepSpeed引擎具有靈活的API，用于保存和加載檢查點，以處理來自客戶端模型和其自身內(nèi)部的狀態(tài)。

defsave_checkpoint(self,save_dir,tag,client_state={})
defload_checkpoint(self,load_dir,tag)

要使用DeepSpeed，我們需要更新utils.py，它是Megatron-LM GPT2保存和加載檢查點的腳本。

創(chuàng)建一個新的函數(shù) save_ds_checkpoint()，如下所示。新函數(shù)收集客戶端模型狀態(tài)，并通過調(diào)用DeepSpeed的 save_checkpoint() 將其傳遞給DeepSpeed引擎。

defsave_ds_checkpoint(iteration,model,args):
"""Saveamodelcheckpoint."""

sd={}
sd['iteration']=iteration
#rngstates.
ifnotargs.no_save_rng:
sd['random_rng_state']=random.getstate()
sd['np_rng_state']=np.random.get_state()
sd['torch_rng_state']=torch.get_rng_state()
sd['cuda_rng_state']=get_accelerator().get_rng_state()
sd['rng_tracker_states']=mpu.get_cuda_rng_tracker().get_states()

model.save_checkpoint(args.save,iteration,client_state=sd)

在 Megatron-LM GPT2 的 save_checkpoint() 函數(shù)中，添加以下行以調(diào)用上述 DeepSpeed 函數(shù)。

defsave_checkpoint(iteration,model,optimizer,
lr_scheduler,args):
"""Saveamodelcheckpoint."""
ifargs.deepspeed:
save_ds_checkpoint(iteration,model,args)
else:
......

在 load_checkpoint() 函數(shù)中，使用以下 DeepSpeed 檢查點加載API，并返回客戶端模型的狀態(tài)。

defload_checkpoint(model,optimizer,lr_scheduler,args):
"""Loadamodelcheckpoint."""

iteration,release=get_checkpoint_iteration(args)

ifargs.deepspeed:
checkpoint_name,sd=model.load_checkpoint(args.load,iteration)

ifcheckpoint_nameisNone:
ifmpu.get_data_parallel_rank()==0:
print("Unabletoloadcheckpoint.")
returniteration
else:
......

DeepSpeed Activation Checkpoints（可選）

DeepSpeed可以通過在模型并行GPU之間劃分激活檢查點或者將其轉(zhuǎn)移到CPU來減少模型并行訓練過程中激活的內(nèi)存消耗。這些優(yōu)化措施是可選的,除非激活內(nèi)存成為瓶頸,否則可以跳過。要啟用Activation checkpoint,我們使用deepspeed.checkpointing API來替換Megatron的Activation checkpoint和隨機狀態(tài)跟蹤器API。這個替換應該發(fā)生在首次調(diào)用這些API之前。

a) 替換 pretrain_gpt.py 中的：

#OptionalDeepSpeedActivationCheckpointingFeatures
#
ifargs.deepspeedandargs.deepspeed_activation_checkpointing:
set_deepspeed_activation_checkpointing(args)

defset_deepspeed_activation_checkpointing(args):

deepspeed.checkpointing.configure(mpu,
deepspeed_config=args.deepspeed_config,
partition_activation=True)

mpu.checkpoint=deepspeed.checkpointing.checkpoint
mpu.get_cuda_rng_tracker=deepspeed.checkpointing.get_cuda_rng_tracker
mpu.model_parallel_cuda_manual_seed=
deepspeed.checkpointing.model_parallel_cuda_manual_seed

替換 mpu/transformer.py 中的：

ifdeepspeed.checkpointing.is_configured():
globalget_cuda_rng_tracker,checkpoint
get_cuda_rng_tracker=deepspeed.checkpoint.get_cuda_rng_tracker
checkpoint=deepspeed.checkpointing.checkpoint

通過這些替換，可以使用 deepspeed.checkpointing.configure 或 deepspeed_config 文件指定各種 DeepSpeed Activation checkpoint優(yōu)化，例如activation partitioning, contiguous checkpointing 和 CPU checkpointing。

關(guān)于DeepSpeed Activation CheckPoint的更多信息我們可以參考 https://deepspeed.readthedocs.io/en/latest/activation-checkpointing.html#configuring-activation-checkpointing ，我翻譯一下主要的 configure 和 is_configured接口。

deepspeed.checkpointing.configure(mpu_, deepspeed_config=None, partition_activations=None, contiguous_checkpointing=None, num_checkpoints=None, checkpoint_in_cpu=None, synchronize=None, profile=None)

配置 DeepSpeed Activation Checkpointing.

參數(shù)：

    mpu – 可選：一個實現(xiàn)以下方法的對象：get_model_parallel_rank/group/world_size 和 get_data_parallel_rank/group/world_size。
    deepspeed_config – 可選：當提供DeepSpeed配置JSON文件時，將用于配置DeepSpeed激活檢查點。
    partition_activations – 可選：啟用后在模型并行GPU之間Partitions activation checkpoint。默認為False。如果提供，將覆蓋deepspeed_config。
    contiguous_checkpointing – 可選：將activation checkpoint復制到一個連續(xù)的內(nèi)存緩沖區(qū)中。僅在啟用Partitions activation checkpoint時與同構(gòu)檢查點一起使用。必須提供num_checkpoints。默認為False。如果提供，將覆蓋deepspeed_config。
    num_checkpoints – 可選：在模型的前向傳播期間存儲的activation checkpoint數(shù)。用于計算連續(xù)checkpoint緩沖區(qū)的大小。如果提供，將覆蓋deepspeed_config。
    checkpoint_in_cpu – 可選：將activation checkpoint移動到CPU。僅在Partitions activation checkpoint時工作。默認值為false。如果提供，將覆蓋deepspeed_config。
    synchronize – 可選：在每次調(diào)用deepspeed.checkpointing.checkpoint的前向和反向傳遞的開始和結(jié)束處執(zhí)行g(shù)et_accelerator().synchronize()。默認為false。如果提供，將覆蓋deepspeed_config。
    profile – 可選：記錄每個deepspeed.checkpointing.checkpoint調(diào)用的前向和反向傳播時間。如果提供，將覆蓋deepspeed_config。

deepspeed.checkpointing.is_configured()

如果已配置deepspeed activation checkpoint，則為True

否則返回false，需要通過調(diào)用deepspeed.checkpointing.configure來進行配置。

訓練腳本

我們假設(shè)在先前的步驟中準備好了 webtext 數(shù)據(jù)。要開始使用 DeepSpeed 訓練 Megatron-LM GPT2 模型，請執(zhí)行以下命令開始訓練。

單GPU運行：bash scripts/ds_pretrain_gpt2.sh

多GPU/節(jié)點運行：bash scripts/ds_zero2_pretrain_gpt2_model_parallel.sh

DeepSpeed 使用 GPT-2 進行評估

DeepSpeed 通過先進的 ZeRO 優(yōu)化器有效地訓練非常大的模型。在2020年2月，我們在 DeepSpeed 中發(fā)布了 ZeRO 的一部分優(yōu)化，該優(yōu)化執(zhí)行優(yōu)化器狀態(tài)切分。我們將其稱為 ZeRO-1。在2020年5月，我們在 DeepSpeed 中擴展了 ZeRO-1，包括來自 ZeRO 的其它優(yōu)化，包括梯度和激活切分，以及連續(xù)內(nèi)存優(yōu)化。我們將此版本稱為 ZeRO-2。

ZeRO-2顯著降低了訓練大型模型的內(nèi)存占用，這意味著可以使用更少的模型并行度和更大的批量大小來訓練大型模型。較小的模型并行度通過增加計算的粒度（例如矩陣乘法）來提高訓練效率，其中性能與矩陣的大小直接相關(guān)。此外，較小的模型并行度還導致模型并行GPU之間的通信更少，進一步提高了性能。較大的批量大小具有類似的效果，增加了計算粒度，并減少了通信，也獲得了更好的性能。因此，通過DeepSpeed和ZeRO-2集成到Megatron中，與僅使用Megatron相比，我們將模型規(guī)模和速度提升到了一個全新的水平。

更具體地說，DeepSpeed和ZeRO-2在以下四個方面表現(xiàn)出色（如圖2所示），支持比現(xiàn)有模型大一個數(shù)量級的模型，速度快了多達10倍，具有超線性的可擴展性，并提高了可用性以實現(xiàn)大型模型訓練的平民化。以下詳細介紹了這四個方面。

模型大小：目前最先進的大型模型，例如OpenAI GPT-2、NVIDIA Megatron-LM、Google T5和Microsoft Turing-NLG，分別具有1.5B、8.3B、11B和17B個參數(shù)。ZeRO-2提供了系統(tǒng)支持，可以高效地運行1700億個參數(shù)的模型，比這些最大的模型大一個數(shù)量級（圖2，左上角）。速度: 改進的內(nèi)存效率提高了吞吐量和訓練速度。圖2（左下角）顯示了ZeRO-2和ZeRO-1（兩者都將由ZeRO驅(qū)動的數(shù)據(jù)并行與NVIDIA Megatron-LM模型并行結(jié)合在一起）的系統(tǒng)吞吐量，以及僅使用最先進的模型并行方法Megatron-LM（圖2，左下方的基準線）。ZeRO-2可以在一個由400個NVIDIA V100 GPU組成的集群上運行1000億參數(shù)的模型，每個GPU的性能超過38 teraflops，聚合性能超過15 petaflops。對于相同大小的模型，與僅使用Megatron-LM相比，ZeRO-2的訓練速度快10倍，與ZeRO-1相比快5倍。擴展性: 我們觀察到超線性加速（圖2，右上角），即當GPU數(shù)量翻倍時，性能增長超過兩倍。ZeRO-2降低了模型狀態(tài)的內(nèi)存占用，使我們可以在每個GPU上容納更大的批量大小，從而提高了性能。隨著數(shù)據(jù)并行度的增加，ZeRO-2減少了模型狀態(tài)的內(nèi)存占用，這也導致了超線性加速的效果。平民化大模型訓練: ZeRO-2使模型科學家能夠高效地訓練高達130億個參數(shù)的模型，而無需進行通常需要模型重構(gòu)的模型并行（圖2，右下角）。130億個參數(shù)比大多數(shù)最先進的模型（如具有110億個參數(shù)的Google T5）都要大。因此，模型科學家可以自由地嘗試大型模型，而不必擔心模型并行。相比之下，經(jīng)典數(shù)據(jù)并行方法的實現(xiàn)（如PyTorch分布式數(shù)據(jù)并行）在1.4億個參數(shù)的模型上會耗盡內(nèi)存，而ZeRO-1則支持最多6億個參數(shù)。

此外，在沒有模型并行的情況下，這些模型可以在帶寬較低的集群上進行訓練，同時仍然比使用模型并行獲得顯著更高的吞吐量。例如，使用40 Gbps Infiniband互連連接的四個節(jié)點集群（每個節(jié)點具有四個連接到PCI-E的NVIDIA 16GB V100 GPU），使用ZeRO提供的數(shù)據(jù)并行比使用模型并行快近4倍，可將GPT-2模型訓練得更快。因此，通過這種性能提升，大型模型訓練不再僅限于具有超快速互連的GPU集群，而且也可以在帶寬有限的中等集群上進行。

目前，Megatron的倉庫也集成了DeepSpeed提供的大量Feature，所以Megatron-DeepSpeed這個倉庫用的人很少，一般還是獨立使用Megatron或者DeepSpeed來煉丹。不過這里的教程為我們指出了要將DeepSpeed用在Megatron倉庫需要做的修改，我們也可以看到DeepSpeed的擴展性是還不錯的。如果你對DeepSpeed和Megatron的聯(lián)合使用感興趣，可以參考下面DeepSpeedExamples中的例子：https://github.com/microsoft/DeepSpeedExamples/tree/bdf8e59aede8c8e0577e8d4d557298ca8515268f/Megatron-LM

0x2. Zero Redundancy Optimizer （零冗余優(yōu)化器）

在閱讀這個 Tutorial 之前可以先瀏覽一下0x1節(jié)，在本教程中，我們將把ZeRO優(yōu)化器應用于Megatron-LM GPT-2模型。ZeRO是一組強大的內(nèi)存優(yōu)化技術(shù)，可以有效地訓練具有數(shù)萬億參數(shù)的大型模型，如GPT-2和Turing-NLG 17B。與其它用于訓練大型模型的模型并行方法相比，ZeRO的一個關(guān)鍵優(yōu)勢是不需要對模型代碼進行修改。正如本教程將演示的那樣，在DeepSpeed模型中使用ZeRO非?？旖莺秃唵?，因為你只需要在DeepSpeed配置JSON中更改一些配置即可。不需要進行代碼更改。

Zero概述

ZeRO利用數(shù)據(jù)并行的計算和內(nèi)存資源來降低模型訓練所需的每個設(shè)備（GPU）的內(nèi)存和計算要求。ZeRO通過在分布式訓練硬件中的可用設(shè)備（GPU和CPU）之間分區(qū)各種模型訓練狀態(tài)（權(quán)重、梯度和優(yōu)化器狀態(tài)）來降低每個GPU的內(nèi)存消耗。具體而言，ZeRO被實現(xiàn)為逐步優(yōu)化的階段，其中早期階段的優(yōu)化在后期階段可用。如果您想深入了解ZeRO，請參見Zero的論文（https://arxiv.org/abs/1910.02054v3）。

Stage 1 。優(yōu)化器狀態(tài)（例如Adam優(yōu)化器的32位權(quán)重和第一、第二階矩估計）在進程間被切分，以便每個進程僅更新它持有的部分。

Stage 2。用于更新模型權(quán)重的32位梯度也被切分，以便每個進程僅保留與其優(yōu)化器狀態(tài)部分對應的梯度。

Stage 3。16位模型參數(shù)被在進程間被切分。ZeRO-3將在前向和后向傳遞期間自動收集和切分它們。

此外，ZeRO-3還包括無限卸載引擎以形成ZeRO-Infinity（https://arxiv.org/abs/2104.07857），可以卸載到CPU和NVMe內(nèi)存以實現(xiàn)巨大的內(nèi)存節(jié)省。

訓練環(huán)境

我們使用DeepSpeed Megatron-LM GPT-2代碼作為例子。你可以按照Megatron-LM教程逐步操作，熟悉代碼。我們將在配備32GB RAM的NVIDIA Tesla V100-SXM3 Tensor Core GPU（https://www.nvidia.com/en-us/data-center/v100/）上訓練本教程中的模型。

開啟Zero優(yōu)化

要為DeepSpeed模型啟用ZeRO優(yōu)化，我們只需要將zero_optimization鍵添加到DeepSpeed JSON配置中。有關(guān)zero_optimization鍵的配置的完整描述，請參見此處(https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training)。

訓練一個1.5B參數(shù)的GPT2模型

我們通過展示ZeROStage 1的優(yōu)點來演示它使得在八個V100 GPU上進行1.5億參數(shù)的GPT-2模型的數(shù)據(jù)并行訓練成為可能。我們將訓練配置為每個設(shè)備使用1個批次，以確保內(nèi)存消耗主要由模型參數(shù)和優(yōu)化器狀態(tài)引起。我們通過對deepspeed啟動腳本應用以下修改來創(chuàng)建這個訓練場景：

--model-parallel-size 1 
       --num-layers 48 
       --hidden-size 1600 
       --num-attention-heads 16 
       --batch-size 1 
       --deepspeed_config ds_zero_stage_1.config 

在沒有ZeRO的情況下訓練這個模型會失敗，并顯示出內(nèi)存不足（OOM）錯誤，如下所示：

這個模型不能適應GPU內(nèi)存的一個重要原因是Adam優(yōu)化器狀態(tài)消耗了18GB的內(nèi)存，這是32GB RAM的一個相當大的部分。通過使用ZeRO Stage1將優(yōu)化器狀態(tài)在八個數(shù)據(jù)并行 rank 之間進行切分，每個設(shè)備的內(nèi)存消耗可以降低到2.25GB，從而使得模型可訓練。為了啟用ZeRO Stage1，我們只需要更新DeepSpeed JSON配置文件如下：

{
    "zero_optimization": {
        "stage": 1,
        "reduce_bucket_size": 5e8
    }
}

如上所示，我們在zero_optimization鍵中設(shè)置了兩個字段。具體來說，我們將stage字段設(shè)置為1，并將可選的reduce_bucket_size設(shè)置為500M。啟用ZeRO Stage1后，模型現(xiàn)在可以在8個GPU上平穩(wěn)地訓練，而不會耗盡內(nèi)存。以下是模型訓練的一些屏幕截圖：

在這里插入圖片描述

從上面的nvidia-smi截圖中，我們可以看到只有第6-7個GPU被用于訓練模型。通過使用ZeRO Stage1，我們可以進一步減少每個設(shè)備的內(nèi)存消耗，通過增加數(shù)據(jù)并行度來實現(xiàn)這些內(nèi)存節(jié)省，這些內(nèi)存節(jié)省可以用于增加模型大小和/或 batch 大小。相比之下，僅使用數(shù)據(jù)并行無法實現(xiàn)這樣的好處。

訓練一個10b的GPT-2模型

ZeRO Stage2 優(yōu)化進一步增加了可以使用數(shù)據(jù)并行訓練的模型大小。我們通過使用32個V100 GPU訓練一個具有10B參數(shù)的模型來展示這一點。

首先，我們需要配置一個啟用了Activation Checkpoint的10B參數(shù)模型。這可以通過對DeepSpeed啟動腳本應用以下GPT-2模型配置更改來完成。

       --model-parallel-size 1 
       --num-layers 50 
       --hidden-size 4096 
       --num-attention-heads 32 
       --batch-size 1 
       --deepspeed_config ds_zero_stage_2.config 
       --checkpoint-activations

接下來，我們需要更新DeepSpeed JSON配置，如下所示，以啟用ZeRO Stage2優(yōu)化：

{
    "zero_optimization": {
        "stage": 2,
        "contiguous_gradients": true,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e8,
        "allgather_bucket_size": 5e8
    }
}

在上面的更改中，我們將stage字段設(shè)置為2，并配置了在ZeRO Stage2 中可用的其他優(yōu)化選項。例如，我們啟用了contiguous_gradients，以減少反向傳播期間的內(nèi)存碎片。這些優(yōu)化設(shè)置的完整描述可在（https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training）找到。有了這些更改，我們現(xiàn)在可以啟動訓練。

以下是訓練日志的截圖：

以下是訓練期間nvidia-smi顯示的GPU活動的截圖：

在這里插入圖片描述

使用ZeRO-Infinity訓練萬億級別的模型

ZeRO-3是ZeRO的第三個階段，它可以將完整的模型狀態(tài)（即權(quán)重、梯度和優(yōu)化器狀態(tài)）進行切分，以線性地擴展內(nèi)存節(jié)省量和數(shù)據(jù)并行度。可以在JSON配置中啟用ZeRO-3。這里（https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training）提供了這些配置的完整描述。

使用ZeRO-Infinity將計算轉(zhuǎn)移到CPU和NVMe

ZeRO-Infinity使用DeepSpeed的無限卸載引擎將完整的模型狀態(tài)轉(zhuǎn)移到CPU或NVMe內(nèi)存中，從而使更大的模型尺寸成為可能。卸載可以在DeepSpeed配置中啟用：

{
    "zero_optimization": {
        "stage": 3,
        "contiguous_gradients": true,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_prefetch_bucket_size": 1e7,
        "stage3_param_persistence_threshold": 1e5,
        "reduce_bucket_size": 1e7,
        "sub_group_size": 1e9,
        "offload_optimizer": {
            "device": "cpu"
         },
        "offload_param": {
            "device": "cpu"
       }
   }
}

ZeRO-Infinity與ZeRO-Offload的區(qū)別：DeepSpeed最初通過ZeRO-Offload實現(xiàn)了Offload功能，這是一種將優(yōu)化器和梯度狀態(tài)轉(zhuǎn)移到ZeRO-2中的CPU內(nèi)存的系統(tǒng)。ZeRO-Infinity是下一代基于ZeRO-3的Offload功能。ZeRO-Infinity能夠比ZeRO-Offload更多地卸載數(shù)據(jù)，并具有更有效的帶寬利用和計算與通信的重疊。

分配大規(guī)模Megatron-LM模型

我們進行了兩項進一步的模型初始化更改，以支持超出本地系統(tǒng)內(nèi)存但未超出總系統(tǒng)內(nèi)存的模型。

以可擴展內(nèi)存的方式分配模型。模型參數(shù)將被分配并立即切分到數(shù)據(jù)并行g(shù)roup中。如果remote_device是“cpu”或“nvme”，模型也將被分配到CPU / NVMe內(nèi)存中而不是GPU內(nèi)存中。有關(guān)更多詳細信息，請參閱完整的ZeRO-3初始化文檔(https://deepspeed.readthedocs.io/en/latest/zero3.html#deepspeed.zero.Init)。

withdeepspeed.zero.Init(data_parallel_group=mpu.get_data_parallel_group(),
remote_device=get_args().remote_device,
enabled=get_args().zero_stage==3):
model=GPT2Model(num_tokentypes=0,parallel_output=True)

收集額外的嵌入權(quán)重以進行初始化。DeepSpeed 在 module 的構(gòu)造函數(shù)和前向/反向傳遞期間會自動收集 module 的參數(shù)。但是，如果需要額外的訪問，則必須與DeepSpeed進行協(xié)調(diào)以確保參數(shù)數(shù)據(jù)被收集并隨后被切分。如果修改了張量，則還應使用modifier_rank參數(shù)，以確保所有進程對數(shù)據(jù)有一致的視角。有關(guān)更多詳細信息，請參閱完整的GatheredParameters文檔(https://deepspeed.readthedocs.io/en/latest/zero3.html#deepspeed.zero.GatheredParameters)。

self.position_embeddings=torch.nn.Embedding(...)
withdeepspeed.zero.GatheredParameters(self.position_embeddings.weight,
modifier_rank=0):
#Initializethepositionembeddings.
self.init_method(self.position_embeddings.weight)

...

self.tokentype_embeddings=torch.nn.Embedding(...)
withdeepspeed.zero.GatheredParameters(self.tokentype_embeddings.weight,
modifier_rank=0):
#Initializethetoken-typeembeddings.
self.init_method(self.tokentype_embeddings.weight)

閱讀了一下文檔，這里的優(yōu)化就類似于編譯器中的公共子表達式消除，由于embeding層在模型中只在init中聲明，然后每一層的Transformer塊都要訪問embedding對象的weight，所以可以采用這種聲明為公共變量的特殊優(yōu)化來減少顯存占用。

以內(nèi)存為中心的分塊優(yōu)化

ZeRO-Infinity包括一個用于進一步降低內(nèi)存使用的Linear層替代方案。我們可以選擇性地對每個Transformer層中的模型并行線性層進行分塊。請注意，可以通過在構(gòu)建層時指定相應的基類來將模型并行性和分塊相結(jié)合。deepspeed.zero.TiledLinear模塊利用ZeRO-3的數(shù)據(jù)獲取和釋放模式，將一個大的運算符分解成較小的塊，可以順序執(zhí)行，從而降低工作內(nèi)存需求。

我們在代碼中包含了一個來自Megatron-LM的ParallelMLP（https://github.com/microsoft/DeepSpeedExamples/blob/bdf8e59aede8c8e0577e8d4d557298ca8515268f/Megatron-LM-v1.1.5-ZeRO3/megatron/model/transformer.py#L82）的示例更改。transformer.py中的另外三個模型并行層的處理方式類似。

Megatron-LM的模型并行層具有特殊的形式，其中層的加性bias被延遲，并在forward()中返回，以便與后續(xù)運算符融合。DeepSpeed的deepspeed.zero.TiledLinearReturnBias是TiledLinear的子類，它只是將返回的偏置參數(shù)轉(zhuǎn)發(fā)而不進行累加。

-self.dense_h_to_4h=mpu.ColumnParallelLinear(
+self.dense_h_to_4h=deepspeed.zero.TiledLinearReturnBias(
args.hidden_size,
4*args.hidden_size,
+in_splits=args.tile_factor,
+out_splits=4*args.tile_factor,
+linear_cls=mpu.ColumnParallelLinear,
gather_output=False,
init_method=init_method,
skip_bias_add=True)

注意，我們按比例縮放in_splits和out_splits與input_size和output_size。這會導致固定大小的小塊[hidden/tile_factor，hidden/tile_factor]。

提取權(quán)重

如果您需要從Deepspeed中獲取預訓練權(quán)重，則可以按以下步驟獲取fp16權(quán)重：

在ZeRO-2下，state_dict包含fp16模型權(quán)重，可以使用torch.save正常保存這些權(quán)重。

在ZeRO-3下，state_dict僅包含占位符，因為模型權(quán)重被切分到多個GPU上。如果你想要獲取這些權(quán)重，請啟用：

"zero_optimization": {
        "stage3_gather_16bit_weights_on_model_save": true
    },

然后使用如下的代碼保存模型：

ifself.deepspeed:
self.deepspeed.save_16bit_model(output_dir,output_file)

因為它需要在一個GPU上合并權(quán)重，所以這可能會很慢并且需要大量內(nèi)存，因此只在需要時使用此功能。

請注意，如果stage3_gather_16bit_weights_on_model_save為False，則不會保存任何權(quán)重（因為state_dict中沒有這些權(quán)重）。你也可以使用此方法保存ZeRO-2權(quán)重。

如果你想獲取fp32權(quán)重，我們提供了一種特殊的腳本，可以進行離線合并。它不需要配置文件或GPU。以下是其使用示例：

$ cd /path/to/checkpoint_dir
$ ./zero_to_fp32.py . pytorch_model.bin
Processing zero checkpoint at global_step1
Detected checkpoint of type zero stage 3, world_size: 2
Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)

當你保存checkpoint時，zero_to_fp32.py腳本會自動生成。注意：目前該腳本使用的內(nèi)存（通用RAM）是最終checkpoint大小的兩倍。

或者，如果你有足夠的CPU內(nèi)存，并且想要將模型更新為其fp32權(quán)重，您可以在訓練結(jié)束時執(zhí)行以下操作：

fromdeepspeed.utils.zero_to_fp32importload_state_dict_from_zero_checkpoint
fp32_model=load_state_dict_from_zero_checkpoint(deepspeed.module,checkpoint_dir)

請注意，該模型適合保存，但不再適合繼續(xù)訓練，并且需要重新執(zhí)行deepspeed.initialize()。如果你只想要state_dict，可以執(zhí)行以下操作：

fromdeepspeed.utils.zero_to_fp32importget_fp32_state_dict_from_zero_checkpoint
state_dict=get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)

0x3. Zero-Offload

Zero-Offload有一篇介紹的博客（https://www.microsoft.com/en-us/research/blog/zero-infinity-and-deepspeed-unlocking-unprecedented-model-scale-for-deep-learning-training/），本來想翻譯下發(fā)現(xiàn)智源的一篇博客基本算是翻譯版本了，所以大家可以看這篇中文版的Zero-Offload博客 https://hub.baai.ac.cn/view/7905。這里做一下Zero-Offload的教程翻譯。

ZeRO-Offload 是一種 ZeRO 優(yōu)化，它將優(yōu)化器內(nèi)存和計算從GPU轉(zhuǎn)移到主機CPU。 ZeRO-Offload 可以在單個GPU上高效地訓練具有多達130億個參數(shù)的大模型。在本教程中，我們將使用 ZeRO-Offload 在 DeepSpeed 中訓練一個具有100億個參數(shù)的 GPT-2 模型。此外，在 DeepSpeed 模型中使用 ZeRO-Offload 很快很容易，因為你只需要在 DeepSpeed 配置json中更改一些配置即可，無需進行代碼更改。

ZeRO-Offload概述

對于大型模型訓練，如Adam等優(yōu)化器可能會消耗大量GPU計算和內(nèi)存資源。ZeRO-Offload通過利用主機CPU的計算和內(nèi)存資源執(zhí)行優(yōu)化器來減少此類模型的GPU計算和內(nèi)存要求。此外，為了防止優(yōu)化器成為瓶頸，ZeRO-Offload使用DeepSpeed高度優(yōu)化的CPU Adam實現(xiàn)，稱為DeepSpeedCPUAdam（https://github.com/microsoft/DeepSpeed/tree/master/deepspeed/ops/adam）。 DeepSpeedCPUAdam比標準的PyTorch實現(xiàn)快5倍到7倍。要深入了解ZeRO-Offload的設(shè)計和性能，請參閱我們的博客文章（就是上面提到的）。截圖：

在這里插入圖片描述

訓練環(huán)境

在本教程中，我們將使用 DeepSpeed Megatron-LM GPT-2 代碼配置一個具有 100 億參數(shù)的 GPT-2 模型。如果你之前沒有使用過 Megatron-LM，請建議你先完成 Megatron-LM 教程（也就是本文中的0x1節(jié)）。我們將使用一張 NVIDIA Tesla V100-SXM3 Tensor Core GPU，其具有 32GB 的內(nèi)存。

在單個 V100 GPU 上訓練10B的GPT2模型

我們需要對 Megatron-LM 的啟動腳本和 DeepSpeed 配置 json 進行更改。

Megatron-LM GPT-2 的啟動腳本更改：

我們需要對 DeepSpeed Megatron-LM GPT-2 模型的啟動腳本進行兩個更改。第一個更改是配置一個啟用activation checkpointing的 10B 參數(shù) GPT-2 模型，可以通過以下一組更改來實現(xiàn)：

--model-parallel-size 1 
       --num-layers 50 
       --hidden-size 4096 
       --num-attention-heads 32 
       --batch-size 10 
       --deepspeed_config ds_zero_offload.config 
       --checkpoint-activations

如果你已經(jīng)完成了 Megatron-LM 教程，上述更改中的大多數(shù)標志應該是熟悉的。

其次，我們需要應用以下更改，以確保只使用一個GPU進行訓練。

deepspeed --num_nodes 1 --num_gpus 1 ...

DeepSpeed 配置更改

ZeRO-Offload 利用了一些 ZeRO Stage 1和 Stage 2 機制，因此啟用 ZeRO-Offload 需要的配置更改是啟用 ZeRO Stage 1和 Stage 2 所需的擴展。下面是啟用 ZeRO-Offload 的 zero_optimization 配置：

{
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
        }
        "contiguous_gradients": true,
        "overlap_comm": true
    }
}

如上所述，除了將stage字段設(shè)置為2（啟用ZeRO Stage 2，但Stage 1也可以），我們還需要將offload_optimizer設(shè)備設(shè)置為cpu以啟用ZeRO-Offload優(yōu)化。此外，我們可以設(shè)置其他ZeRO Stage 2優(yōu)化標志，例如overlap_comm來調(diào)整ZeRO-Offload性能。通過這些更改，我們現(xiàn)在可以運行模型。我們在下面分享了一些訓練的截圖。

以下是 nvidia-smi 的截圖，顯示僅在訓練期間激活了 GPU 0

在這里插入圖片描述

最后，以下是 htop 的截圖，顯示了在優(yōu)化器計算期間主機CPU和內(nèi)存的活動情況：

在這里插入圖片描述

0x4. 總結(jié)

本篇文章主要翻譯了DeepSpeed里面和Zero相關(guān)的技術(shù)教程，對DeepSpeed感興趣的讀者可以對照官方文檔學習一下。
責任編輯：彭菁