使用原始的 Megatron-LM 訓練 GPT2

設置訓練數據

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

開啟DeepSpeed

DeepSpeed 使用 GPT-2 進行評估

Zero概述

訓練環境

開啟Zero優化

訓練一個1.5B參數的GPT2模型

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

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

使用ZeRO-Infinity將計算轉移到CPU和NVMe

分配大規模Megatron-LM模型

以內存為中心的分塊優化

提取權重

ZeRO-Offload概述

訓練環境

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

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

DeepSpeed 配置更改

0x0. 前言

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

0x1. Megatron-LM GPT2

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

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

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

使用原始的 Megatron-LM 訓練 GPT2

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

git submodule update --init --recursive

設置訓練數據

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

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

對于單塊GPU：

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

運行 bash scripts/pretrain_gpt2.sh

對于多個節點和多個GPU：

修改 scripts/pretrain_gpt2_model_parallel.sh

將其 --train-data 參數設置為 "webtext"。

GPUS_PER_NODE 指示了測試中每個節點使用的GPU數量。

NNODES 指示了測試中涉及的節點數量。

執行 bash scripts/pretrain_gpt2_model_parallel.sh

開啟DeepSpeed

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

arguments.py : 參數配置文件

pretrain_gpt2.py : 訓練的主入口點

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

參數解析

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

初始化和訓練

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

初始化

我們使用 deepspeed.initialize 創建 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() 函數中進行 DeepSpeed 初始化，函數傳的參數包含原始model、optimizer、args、lr_scheduler 和 mpu。

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

使用訓練API

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

前向傳播

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

反向傳播

通過在模型引擎上直接調用 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會在使用小批量更新權重后自動處理梯度清零。此外，DeepSpeed在內部解決了分布式數據并行和FP16，簡化了多個地方的代碼。

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

ifargs.deepspeed:
#DeepSpeed反向傳播已經處理了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) 我們也跳過更新主節點梯度，因為DeepSpeed在內部解決了這個問題。

#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

更新模型參數

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

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

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

損失縮放

GPT2訓練腳本在訓練過程中記錄了損失縮放值。在DeepSpeed優化器內部，該值存儲為 cur_scale，而不是Megatron的優化器中的 loss_scale。因此，我們在日志字符串中適當地進行了替換。

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

檢查點保存和加載

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

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

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

創建一個新的函數 save_ds_checkpoint()，如下所示。新函數收集客戶端模型狀態，并通過調用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() 函數中，添加以下行以調用上述 DeepSpeed 函數。

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

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

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之間劃分激活檢查點或者將其轉移到CPU來減少模型并行訓練過程中激活的內存消耗。這些優化措施是可選的,除非激活內存成為瓶頸,否則可以跳過。要啟用Activation checkpoint,我們使用deepspeed.checkpointing API來替換Megatron的Activation checkpoint和隨機狀態跟蹤器API。這個替換應該發生在首次調用這些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優化，例如activation partitioning, contiguous checkpointing 和 CPU checkpointing。

關于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.

參數：

    mpu – 可選：一個實現以下方法的對象：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復制到一個連續的內存緩沖區中。僅在啟用Partitions activation checkpoint時與同構檢查點一起使用。必須提供num_checkpoints。默認為False。如果提供，將覆蓋deepspeed_config。
    num_checkpoints – 可選：在模型的前向傳播期間存儲的activation checkpoint數。用于計算連續checkpoint緩沖區的大小。如果提供，將覆蓋deepspeed_config。
    checkpoint_in_cpu – 可選：將activation checkpoint移動到CPU。僅在Partitions activation checkpoint時工作。默認值為false。如果提供，將覆蓋deepspeed_config。
    synchronize – 可選：在每次調用deepspeed.checkpointing.checkpoint的前向和反向傳遞的開始和結束處執行get_accelerator().synchronize()。默認為false。如果提供，將覆蓋deepspeed_config。
    profile – 可選：記錄每個deepspeed.checkpointing.checkpoint調用的前向和反向傳播時間。如果提供，將覆蓋deepspeed_config。

deepspeed.checkpointing.is_configured()

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

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

訓練腳本

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

單GPU運行：bash scripts/ds_pretrain_gpt2.sh

多GPU/節點運行：bash scripts/ds_zero2_pretrain_gpt2_model_parallel.sh

DeepSpeed 使用 GPT-2 進行評估

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

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

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

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

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

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

0x2. Zero Redundancy Optimizer （零冗余優化器）

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

Zero概述

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

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

Stage 2。用于更新模型權重的32位梯度也被切分，以便每個進程僅保留與其優化器狀態部分對應的梯度。

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

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

訓練環境

我們使用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優化

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

訓練一個1.5B參數的GPT2模型

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

--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的情況下訓練這個模型會失敗，并顯示出內存不足（OOM）錯誤，如下所示：

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

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

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

在這里插入圖片描述

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

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

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

首先，我們需要配置一個啟用了Activation Checkpoint的10B參數模型。這可以通過對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優化：

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

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

以下是訓練日志的截圖：

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

在這里插入圖片描述

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

ZeRO-3是ZeRO的第三個階段，它可以將完整的模型狀態（即權重、梯度和優化器狀態）進行切分，以線性地擴展內存節省量和數據并行度?？梢栽贘SON配置中啟用ZeRO-3。這里（https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training）提供了這些配置的完整描述。

使用ZeRO-Infinity將計算轉移到CPU和NVMe

ZeRO-Infinity使用DeepSpeed的無限卸載引擎將完整的模型狀態轉移到CPU或NVMe內存中，從而使更大的模型尺寸成為可能。卸載可以在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的區別：DeepSpeed最初通過ZeRO-Offload實現了Offload功能，這是一種將優化器和梯度狀態轉移到ZeRO-2中的CPU內存的系統。ZeRO-Infinity是下一代基于ZeRO-3的Offload功能。ZeRO-Infinity能夠比ZeRO-Offload更多地卸載數據，并具有更有效的帶寬利用和計算與通信的重疊。

分配大規模Megatron-LM模型

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

以可擴展內存的方式分配模型。模型參數將被分配并立即切分到數據并行group中。如果remote_device是“cpu”或“nvme”，模型也將被分配到CPU / NVMe內存中而不是GPU內存中。有關更多詳細信息，請參閱完整的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)

收集額外的嵌入權重以進行初始化。DeepSpeed 在 module 的構造函數和前向/反向傳遞期間會自動收集 module 的參數。但是，如果需要額外的訪問，則必須與DeepSpeed進行協調以確保參數數據被收集并隨后被切分。如果修改了張量，則還應使用modifier_rank參數，以確保所有進程對數據有一致的視角。有關更多詳細信息，請參閱完整的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)

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

以內存為中心的分塊優化

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

我們在代碼中包含了一個來自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()中返回，以便與后續運算符融合。DeepSpeed的deepspeed.zero.TiledLinearReturnBias是TiledLinear的子類，它只是將返回的偏置參數轉發而不進行累加。

-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]。

提取權重

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

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

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

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

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

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

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

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

如果你想獲取fp32權重，我們提供了一種特殊的腳本，可以進行離線合并。它不需要配置文件或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腳本會自動生成。注意：目前該腳本使用的內存（通用RAM）是最終checkpoint大小的兩倍。

或者，如果你有足夠的CPU內存，并且想要將模型更新為其fp32權重，您可以在訓練結束時執行以下操作：

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

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

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/），本來想翻譯下發現智源的一篇博客基本算是翻譯版本了，所以大家可以看這篇中文版的Zero-Offload博客 https://hub.baai.ac.cn/view/7905。這里做一下Zero-Offload的教程翻譯。

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

ZeRO-Offload概述

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

在這里插入圖片描述

訓練環境

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

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

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

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

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

--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

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

其次，我們需要應用以下更改，以確保只使用一個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字段設置為2（啟用ZeRO Stage 2，但Stage 1也可以），我們還需要將offload_optimizer設備設置為cpu以啟用ZeRO-Offload優化。此外，我們可以設置其他ZeRO Stage 2優化標志，例如overlap_comm來調整ZeRO-Offload性能。通過這些更改，我們現在可以運行模型。我們在下面分享了一些訓練的截圖。

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

在這里插入圖片描述

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

在這里插入圖片描述

0x4. 總結

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