DeepSpeed里面和Zero相關(guān)技術(shù)教程

使用原始的 Megatron-LM 訓(xùn)練 GPT2

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

運(yùn)行未修改的Megatron-LM GPT2模型

開啟DeepSpeed

DeepSpeed 使用 GPT-2 進(jìn)行評估

Zero概述

訓(xùn)練環(huán)境

開啟Zero優(yōu)化

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

訓(xùn)練一個(gè)10b的GPT-2模型

使用ZeRO-Infinity訓(xùn)練萬億級別的模型

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

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

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

提取權(quán)重

ZeRO-Offload概述

訓(xùn)練環(huán)境

在單個(gè) V100 GPU 上訓(xùn)練10B的GPT2模型

Megatron-LM GPT-2 的啟動(dòng)腳本更改：

DeepSpeed 配置更改

0x0. 前言

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

0x1. Megatron-LM GPT2

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

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

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

使用原始的 Megatron-LM 訓(xùn)練 GPT2

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

git submodule update --init --recursive

設(shè)置訓(xùn)練數(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 )下放置一個(gè)符號(hào)鏈接。

運(yùn)行未修改的Megatron-LM GPT2模型

對于單塊GPU：

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

運(yùn)行 bash scripts/pretrain_gpt2.sh

對于多個(gè)節(jié)點(diǎn)和多個(gè)GPU：

修改 scripts/pretrain_gpt2_model_parallel.sh

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

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

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

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

開啟DeepSpeed

為了用上DeepSpeed需要更新三個(gè)文件：

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

pretrain_gpt2.py : 訓(xùn)練的主入口點(diǎn)

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

參數(shù)解析

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

初始化和訓(xùn)練

我們將修改 pretrain.py 以啟用使用 DeepSpeed 進(jìn)行訓(xùn)練。

初始化

我們使用 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ù)中進(jìn)行 DeepSpeed 初始化，函數(shù)傳的參數(shù)包含原始model、optimizer、args、lr_scheduler 和 mpu。

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

使用訓(xùn)練API

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

前向傳播

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

反向傳播

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

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

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

ifargs.deepspeed:
#DeepSpeed反向傳播已經(jīng)處理了allreduce通信。重置計(jì)時(shí)器以避免破壞下面的計(jì)時(shí)器日志。
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é)點(diǎn)梯度，因?yàn)镈eepSpeed在內(nèi)部解決了這個(gè)問題。

#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ù)以及學(xué)習(xí)率。

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

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

損失縮放

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

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

檢查點(diǎn)保存和加載

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

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

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

創(chuàng)建一個(gè)新的函數(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 檢查點(diǎn)加載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之間劃分激活檢查點(diǎn)或者將其轉(zhuǎn)移到CPU來減少模型并行訓(xùn)練過程中激活的內(nèi)存消耗。這些優(yōu)化措施是可選的,除非激活內(nèi)存成為瓶頸,否則可以跳過。要啟用Activation checkpoint,我們使用deepspeed.checkpointing API來替換Megatron的Activation checkpoint和隨機(jī)狀態(tài)跟蹤器API。這個(gè)替換應(yīng)該發(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 – 可選：一個(gè)實(shí)現(xiàn)以下方法的對象：get_model_parallel_rank/group/world_size 和 get_data_parallel_rank/group/world_size。
    deepspeed_config – 可選：當(dāng)提供DeepSpeed配置JSON文件時(shí)，將用于配置DeepSpeed激活檢查點(diǎn)。
    partition_activations – 可選：啟用后在模型并行GPU之間Partitions activation checkpoint。默認(rèn)為False。如果提供，將覆蓋deepspeed_config。
    contiguous_checkpointing – 可選：將activation checkpoint復(fù)制到一個(gè)連續(xù)的內(nèi)存緩沖區(qū)中。僅在啟用Partitions activation checkpoint時(shí)與同構(gòu)檢查點(diǎn)一起使用。必須提供num_checkpoints。默認(rèn)為False。如果提供，將覆蓋deepspeed_config。
    num_checkpoints – 可選：在模型的前向傳播期間存儲(chǔ)的activation checkpoint數(shù)。用于計(jì)算連續(xù)checkpoint緩沖區(qū)的大小。如果提供，將覆蓋deepspeed_config。
    checkpoint_in_cpu – 可選：將activation checkpoint移動(dòng)到CPU。僅在Partitions activation checkpoint時(shí)工作。默認(rèn)值為false。如果提供，將覆蓋deepspeed_config。
    synchronize – 可選：在每次調(diào)用deepspeed.checkpointing.checkpoint的前向和反向傳遞的開始和結(jié)束處執(zhí)行g(shù)et_accelerator().synchronize()。默認(rèn)為false。如果提供，將覆蓋deepspeed_config。
    profile – 可選：記錄每個(gè)deepspeed.checkpointing.checkpoint調(diào)用的前向和反向傳播時(shí)間。如果提供，將覆蓋deepspeed_config。

deepspeed.checkpointing.is_configured()

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

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

訓(xùn)練腳本

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

單GPU運(yùn)行：bash scripts/ds_pretrain_gpt2.sh

多GPU/節(jié)點(diǎn)運(yùn)行：bash scripts/ds_zero2_pretrain_gpt2_model_parallel.sh

DeepSpeed 使用 GPT-2 進(jìn)行評估

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

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

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

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

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

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

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

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

Zero概述

ZeRO利用數(shù)據(jù)并行的計(jì)算和內(nèi)存資源來降低模型訓(xùn)練所需的每個(gè)設(shè)備（GPU）的內(nèi)存和計(jì)算要求。ZeRO通過在分布式訓(xùn)練硬件中的可用設(shè)備（GPU和CPU）之間分區(qū)各種模型訓(xùn)練狀態(tài)（權(quán)重、梯度和優(yōu)化器狀態(tài)）來降低每個(gè)GPU的內(nèi)存消耗。具體而言，ZeRO被實(shí)現(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)重和第一、第二階矩估計(jì)）在進(jìn)程間被切分，以便每個(gè)進(jìn)程僅更新它持有的部分。

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

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

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

訓(xùn)練環(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/）上訓(xùn)練本教程中的模型。

開啟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)。

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

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

--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的情況下訓(xùn)練這個(gè)模型會(huì)失敗，并顯示出內(nèi)存不足（OOM）錯(cuò)誤，如下所示：

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

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

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

在這里插入圖片描述

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

訓(xùn)練一個(gè)10b的GPT-2模型

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

首先，我們需要配置一個(gè)啟用了Activation Checkpoint的10B參數(shù)模型。這可以通過對DeepSpeed啟動(dòng)腳本應(yīng)用以下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)化選項(xiàng)。例如，我們啟用了contiguous_gradients，以減少反向傳播期間的內(nèi)存碎片。這些優(yōu)化設(shè)置的完整描述可在（https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training）找到。有了這些更改，我們現(xiàn)在可以啟動(dòng)訓(xùn)練。

以下是訓(xùn)練日志的截圖：

以下是訓(xùn)練期間nvidia-smi顯示的GPU活動(dòng)的截圖：

在這里插入圖片描述

使用ZeRO-Infinity訓(xùn)練萬億級別的模型

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

使用ZeRO-Infinity將計(jì)算轉(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實(shí)現(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ù)，并具有更有效的帶寬利用和計(jì)算與通信的重疊。

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

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

以可擴(kuò)展內(nèi)存的方式分配模型。模型參數(shù)將被分配并立即切分到數(shù)據(jù)并行g(shù)roup中。如果remote_device是“cpu”或“nvme”，模型也將被分配到CPU / NVMe內(nèi)存中而不是GPU內(nèi)存中。有關(guān)更多詳細(xì)信息，請參閱完整的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)重以進(jìn)行初始化。DeepSpeed 在 module 的構(gòu)造函數(shù)和前向/反向傳遞期間會(huì)自動(dòng)收集 module 的參數(shù)。但是，如果需要額外的訪問，則必須與DeepSpeed進(jìn)行協(xié)調(diào)以確保參數(shù)數(shù)據(jù)被收集并隨后被切分。如果修改了張量，則還應(yīng)使用modifier_rank參數(shù)，以確保所有進(jìn)程對數(shù)據(jù)有一致的視角。有關(guān)更多詳細(xì)信息，請參閱完整的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)化就類似于編譯器中的公共子表達(dá)式消除，由于embeding層在模型中只在init中聲明，然后每一層的Transformer塊都要訪問embedding對象的weight，所以可以采用這種聲明為公共變量的特殊優(yōu)化來減少顯存占用。

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

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

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

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

-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。這會(huì)導(dǎo)致固定大小的小塊[hidden/tile_factor，hidden/tile_factor]。

提取權(quán)重

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

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

在ZeRO-3下，state_dict僅包含占位符，因?yàn)槟Ｐ蜋?quán)重被切分到多個(gè)GPU上。如果你想要獲取這些權(quán)重，請啟用：

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

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

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

因?yàn)樗枰谝粋€(gè)GPU上合并權(quán)重，所以這可能會(huì)很慢并且需要大量內(nèi)存，因此只在需要時(shí)使用此功能。

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

如果你想獲取fp32權(quán)重，我們提供了一種特殊的腳本，可以進(jì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)

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

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

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

請注意，該模型適合保存，但不再適合繼續(xù)訓(xùn)練，并且需要重新執(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)存和計(jì)算從GPU轉(zhuǎn)移到主機(jī)CPU。 ZeRO-Offload 可以在單個(gè)GPU上高效地訓(xùn)練具有多達(dá)130億個(gè)參數(shù)的大模型。在本教程中，我們將使用 ZeRO-Offload 在 DeepSpeed 中訓(xùn)練一個(gè)具有100億個(gè)參數(shù)的 GPT-2 模型。此外，在 DeepSpeed 模型中使用 ZeRO-Offload 很快很容易，因?yàn)槟阒恍枰?DeepSpeed 配置json中更改一些配置即可，無需進(jìn)行代碼更改。

ZeRO-Offload概述

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

在這里插入圖片描述

訓(xùn)練環(huán)境

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

在單個(gè) V100 GPU 上訓(xùn)練10B的GPT2模型

我們需要對 Megatron-LM 的啟動(dòng)腳本和 DeepSpeed 配置 json 進(jìn)行更改。

Megatron-LM GPT-2 的啟動(dòng)腳本更改：

我們需要對 DeepSpeed Megatron-LM GPT-2 模型的啟動(dòng)腳本進(jìn)行兩個(gè)更改。第一個(gè)更改是配置一個(gè)啟用activation checkpointing的 10B 參數(shù) GPT-2 模型，可以通過以下一組更改來實(shí)現(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ù)標(biāo)志應(yīng)該是熟悉的。

其次，我們需要應(yīng)用以下更改，以確保只使用一個(gè)GPU進(jìn)行訓(xùn)練。

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

DeepSpeed 配置更改

ZeRO-Offload 利用了一些 ZeRO Stage 1和 Stage 2 機(jī)制，因此啟用 ZeRO-Offload 需要的配置更改是啟用 ZeRO Stage 1和 Stage 2 所需的擴(kuò)展。下面是啟用 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)化標(biāo)志，例如overlap_comm來調(diào)整ZeRO-Offload性能。通過這些更改，我們現(xiàn)在可以運(yùn)行模型。我們在下面分享了一些訓(xùn)練的截圖。

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

在這里插入圖片描述

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

在這里插入圖片描述

0x4. 總結(jié)

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