# 關於DeepSpeed的嘗試
知乎部落格地址： https://zhuanlan.zhihu.com/p/673359684
## 參考
Repo：https://github.com/microsoft/DeepSpeedExamples
https://github.com/microsoft/DeepSpeedExamples/blob/master/training/HelloDeepSpeed/train_bert_ds.py，程式碼複製到了本專案的：https://github.com/Coobiw/MiniGPT4Qwen/blob/master/deepspeed_tutorials/train_bert_official.py
## EVA-ViT-G（1B）上的實驗
| Setting（bf16，不開gradient-checkpointing） | max_allocated_memory/GB（若無說明則bs=32） | time per epoch/s（bs=32 8卡 共400條資料） |
| ------------------------------------------- | ------------------------------------------ | ----------------------------------------- |
| ZERO-0 (DDP) | bs=32 OOM bs=16 18.36 | 11.57 |
| ZERO-1 | 21.46 | 9.68 |
| ZERO-1 (offload optimizer) | 20.40 | 13.45 |
| ZERO-2 | 22.26 | 8.51 |
| ZERO-2 (offload optimizer) | 20.80 | 13.34 |
| ZERO-3 | 22.37 | 7.94 |
| ZERO-3 (offload optimizer) | 20.39 | 12.67 |
| ZERO-3 (offload optimizer + params) | 20.39 | 12.08 |
- DeepSpeed x MiniGPT4Qwen
- 一、官方教程
- 1. 瞭解DeepSpeed
- 2. DeepSpeed入門
- 2.1 DeepSpeed引擎 & 分散式環境
- 2.2 API for Training
- 2.3 模型檢查點
- 2.4 DeepSpeed配置
- 2.5 啟動 DeepSpeed 訓練
- 資源配置（多節點）
- 多節點環境變數
- 資源配置（單節點）
- 3. 自動張量並行
- 示例指令碼
- 4. Megatron-LM GPT2
- 5. Zero Redundancy Optimizer
- 5.1 ZeRO概述
- 5.2 訓練環境
- 5.3 開啟ZeRO最佳化
- 5.4 訓練1.5B引數的GPT-2模型
- 5.5 訓練10B引數的GPT-2模型
- 5.6 使用ZeRO-Infinity訓練萬億規模的模型
- 使用ZeRO-Infinity解除安裝到CPU和NVMe
- 分配巨大的Megatron-LM模型
- 以記憶體為中心的平鋪? tiling ???
- 註冊外部引數
- 5.7 提取權重
- 6. 自動微調？Autotuning
- 引數解析
- 二、結合MiniGPT4Qwen
- 1. 初始化分散式訓練模式
- 2.
- 五、DeepSpeed引數配置
- 一、官方教程
原文：https://zhuanlan.zhihu.com/p/673359684
DeepSpeed x MiniGPT4Qwen
作者對MiniGPT4Qwen支援了更加簡單易用，Plug-in的DeepSpeed，給出了ZERO-0(等價於DDP), ZERO-1, ZERO-2的配置。
lavis框架的分散式使用的是最基本的pytorch的DDP，雖然簡單易用，但如今還是算是有些out-of-date了。
FSDP、DeepSpeed、Megatron等各種分散式並行訓練框架更加受到青睞。
一、官方教程
1. 瞭解DeepSpeed
DeepSpeed API 是 PyTorch 的輕量級包裝器。這意味著你可以使用你在 PyTorch 中喜歡的一切，而無需學習一個新平臺。
速度
DeepSpeed 使用 1024 個 V100 GPU（64 臺 DGX-2 機箱）在 44 分鐘內訓練 BERT-large 達到一致性，並在使用 256 個 GPU（16 臺 DGX-2 機箱）時在 2.4 小時內完成。
DeepSpeed 訓練 GPT2（15 億引數）比最先進的 NVIDIA Megatron 在 Azure GPU 上快 3.75 倍。
記憶體效率
DeepSpeed提供記憶體高效的資料並行，並支援在沒有模型並行性的情況下訓練模型。
例如，DeepSpeed 可以在單個 GPU 上訓練多達 130 億引數的模型。
相比之下，現有框架（例如，PyTorch 的 Distributed Data Parallel）在 14 億引數模型時記憶體不足。
DeepSpeed 透過一種稱為零冗餘最佳化器（ZeRO）的新穎解決方案減少了訓練記憶體佔用。與基本資料並行性不同，後者在資料並行程序中複製記憶體狀態，ZeRO 將模型狀態和梯度分割槽以節省大量記憶體。此外，它還減少了啟用記憶體和碎片記憶體。當前實現（ZeRO-2）相對於最先進的技術減少了多達 8 倍的記憶體。你可以在我們的論文中閱讀更多關於 ZeRO 的內容，以及在我們的部落格文章與 中瞭解到ZeRO-1 和 ZeRO-2。
對於 GPU 資源有限的模型科學家，ZeRO-Offload 利用 CPU 和 GPU 記憶體進行大型模型訓練。使用一臺帶有一個 GPU 的機器，我們的使用者可以執行多達 130 億引數的模型而不耗盡記憶體，比現有方法大 10 倍，同時獲得有競爭力的吞吐量。這一功能使多億引數模型訓練民主化，併為許多深度學習從業者探索更大更好的模型開啟了視窗。
可擴充套件性
DeepSpeed支援高效的資料並行、模型並行、流水線並行及其組合，我們稱之為3D並行。
DeepSpeed的3D並行性為執行具有數萬億個引數的模型提供了系統支援，請在我們的新聞稿和教程中閱讀更多資訊。
DeepSpeed可以更高效地執行大型模型，對於跨越1.5B到千億的各種大小的模型，速度最高可達10倍。更具體地說，由ZeRO提供支援的資料並行是互補的，可以與不同型別的模型並行性相結合。它允許DeepSpeed使用較低程度的模型並行性和更高的批次大小來擬合模型，與單獨使用模型並行性相比，提供顯著的效能提升。
通訊效率
DeepSpeed 的流水線並行性在分散式訓練期間減少了通訊量，這使得使用者能夠在網路受限的叢集上以 2-7 倍的速度訓練多億引數模型。
1-bit Adam、0/1 Adam 和 1-bit LAMB 透過最多減少 26 倍的通訊量，同時實現與 Adam 相似的收斂效率，從而允許擴充套件到不同型別的 GPU 叢集和網路。
1-bit Adam blog post, 1-bit Adam tutorial, 0/1 Adam tutorial, 1-bit LAMB tutorial
資料效率
DeepSpeed Data Efficiency Library 透過課程學習提供高效的資料取樣，並透過隨機分層令牌丟棄提供高效的資料路由。組合解決方案可在 GPT-3/BERT 預訓練和 GPT/ViT 微調期間節省多達 2 倍的資料和 2 倍的時間，或在相同的資料/時間下進一步提高模型質量。在教程中檢視更多資訊。
支援長序列
DeepSpeed 提供了稀疏注意力核（sparse attention kernels），這是一種關鍵技術，用於支援模型輸入的長序列，無論是文字、影像還是聲音。與傳統的密集型 Transformer 模型相比，它能夠處理數量級更長的輸入序列，並且執行速度可以快達 6 倍，同時保持相當的準確性。此外，與現有的最先進的稀疏實現相比，DeepSpeed 的稀疏核執行速度還快了 1.5 到 3 倍。更重要的是，DeepSpeed 的稀疏核支援靈活的稀疏格式的高效執行，並使使用者能夠在自定義的稀疏結構上進行創新。
快速收斂的有效性
DeepSpeed 支援高階超引數調整和大批次大小的最佳化器，例如 LAMB。
這些技術提高了模型訓練的有效性，並減少了達到期望精度所需的樣本數量。
易用性
只需幾行程式碼更改即可使PyTorch模型能夠使用DeepSpeed和ZeRO。
與當前的模型並行庫相比，DeepSpeed不需要程式碼重新設計或模型重構。它也不會對模型維度（例如注意力頭的數量、隱藏大小等）、批次大小或任何其他訓練引數進行限制。對於多達130億引數的模型，您可以方便地使用ZeRO驅動的資料並行，而無需模型並行，而相比之下，標準資料並行會對引數超過14億的模型執行記憶體溢位。此外，DeepSpeed方便地支援ZeRO驅動的並行資料與自定義模型並行的靈活組合，例如NVIDIA的Megatron-LM的張量切片。
2. DeepSpeed入門
2.1 DeepSpeed引擎 & 分散式環境
-
DeepSpeed模型訓練是使用DeepSpeed引擎完成的。該引擎可以包裝任何torch. nn.module型別的任意模型，並具有用於訓練和檢查點模型的最小API集。
-
初始化DeepSpeed engine
model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args, model=model, model_parameters=params)deepspeed.initialize 確保在幕後正確執行所需的分散式資料並行或混合精度訓練設定。
DeepSpeed 還可以根據傳遞給 deepspeed.initialize 和 DeepSpeed 配置檔案的引數構建和管理訓練最佳化器、資料載入器和學習速率排程器。
請注意，DeepSpeed 會在每個 training step 自動執行學習速率排程。
-
如果您已經有分散式環境設定，則需要替換
torch.distributed.init_process_group(...)為deepspeed.init_distributed()-
DeepSpeed的
init_distributed函式來初始化分散式環境 -
PyTorch的
init_process_group函式來初始化分散式環境
-
-
預設情況下使用 NCCL 後端
-
如果在
deepspeed.initialize()之後才需要設定分散式環境，可以不使用deepspeed.init_distributed()函式，因為 DeepSpeed 將在其初始化期間自動初始化分散式環境。 不過，如果已經處於就緒狀態，則需要刪除torch.distributed.init_process_group。
2.2 API for Training
一旦DeepSpeed引擎初始化完成，就可以使用三個簡單的API來訓練模型，用於前向傳播、反向傳播和權重更新（step）.
for step, batch in enumerate(data_loader):
#forward() method
loss = model_engine(batch)
#runs backpropagation
model_engine.backward(loss)
#weight update
model_engine.step()
在幕後，DeepSpeed 自動執行必要的操作，以混合精度進行分散式資料並行訓練，並具有預定義的學習率排程器：
-
梯度平均化：在分散式資料並行訓練中，後向傳播確保在訓練一個
train_batch_size後對梯度進行跨資料並行程序平均。 -
損失縮放：在 FP16/混合精度訓練中，DeepSpeed 引擎會自動處理損失的縮放，以避免梯度精度丟失。
-
學習率排程器：當使用DeepSpeed的學習率排程器（在
ds_config.json檔案中指定），DeepSpeed在每個訓練步驟（執行model_engine.step()時）呼叫排程器的step()方法。當不使用DeepSpeed的學習率排程器時：
- 如果排程器在每個 training step 時執行，則使用者可以在初始化 DeepSpeed 引擎時將排程程式傳遞給
deepspeed.initialize，讓 DeepSpeed 管理它以進行更新或儲存/恢復。 - 如果排程器在其他間隔執行排程（例如，training epochs ），則使用者在初始化期間不應將排程程式傳遞給 DeepSpeed，並必須明確管理它。
- 如果排程器在每個 training step 時執行，則使用者可以在初始化 DeepSpeed 引擎時將排程程式傳遞給
2.3 模型檢查點
透過 DeepSpeed 中的 save_checkpoint 和 load_checkpoint API 處理訓練狀態的儲存和載入，這兩個 API 需要傳入兩個引數來唯一標識一個檢查點: ckpt_dir 和 ckpt_id(在目錄中唯一標識檢查點的識別符號。在下面的程式碼片段中，我們使用損失值作為檢查點識別符號。)
#load checkpoint
_, client_sd = model_engine.load_checkpoint(args.load_dir, args.ckpt_id)
step = client_sd['step']
#advance data loader to ckpt step
dataloader_to_step(data_loader, step + 1)
for step, batch in enumerate(data_loader):
#forward() method
loss = model_engine(batch)
#runs backpropagation
model_engine.backward(loss)
#weight update
model_engine.step()
#save checkpoint
if step % args.save_interval:
client_sd['step'] = step
ckpt_id = loss.item()
model_engine.save_checkpoint(args.save_dir, ckpt_id, client_sd = client_sd)
DeepSpeed 可以在對使用者隱藏細節的情況下自動儲存和恢復模型、最佳化器和學習率排程器狀態。
然而，使用者可能希望儲存與給定模型訓練相關的其他資料。為了支援這些專案，save_checkpoint 接受一個客戶端狀態字典 client_sd 進行儲存。這些專案可以從 load_checkpoint 中作為返回引數進行檢索。在上面的示例中，step 值被儲存為 client_sd 的一部分。
重要提示：所有程序都必須呼叫此方法，而不僅僅是rank為0的程序。這是因為每個程序都需要儲存其主權重和排程器+最佳化器狀態。如果僅為rank為0的程序呼叫此方法，它將等待與其他程序進行同步。
2.4 DeepSpeed配置
DeepSpeed功能可以啟用，禁用或被指定為args.deepspeed_config 的JSON配置設定。
以下是一個簡單的JSON配置
{
"train_batch_size": 8,
"gradient_accumulation_steps": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.00015
}
},
"fp16": {
"enabled": true
},
"zero_optimization": true
}
2.5 啟動 DeepSpeed 訓練
資源配置（多節點）
DeepSpeed 使用與 OpenMPI 和 Horovod 相容的 hostfiles 配置多節點計算資源。
一個 hostfile 是主機名（或 SSH 別名）列表，這些主機可以透過無密碼 SSH 訪問，並且包含 slot counts，指定系統中可用的 GPU 數量。例如，
worker-1 slots=4
worker-2 slots=4
...
另外，DeepSpeed 還允許你將模型的分散式訓練限制在可用節點和 GPU 的子集上。透過兩個命令列引數 --num_nodes 和 --num_gpus 啟用此功能。例如，可以透過以下命令將分散式訓練限制在只使用兩個節點：
deepspeed --num_nodes=2 \
<client_entry.py> <client args> \
--deepspeed --deepspeed_config ds_config.json
您可以使用--include和--exclude標誌來包含或排除特定資源。例如，要使用除了節點worker-2上的GPU 0和worker-3上的GPU 0和1之外的所有可用資源：
deepspeed --exclude="worker-2:0@worker-3:0,1" \
<client_entry.py> <client args> \
--deepspeed --deepspeed_config ds_config.json
同樣，您可以僅在worker-2上使用GPU 0和1。
deepspeed --include="worker-2:0,1" \
<client_entry.py> <client args> \
--deepspeed --deepspeed_config ds_config.json
多節點環境變數
...
資源配置（單節點）
如果我們只在單節點上執行（具有一個或多個GPU），那麼 DeepSpeed 不需要像上面描述的主機檔案。
如果未檢測到或傳遞主機檔案，則 DeepSpeed 將查詢本地計算機上的 GPU 數量，以發現可用的插槽數。-- include 和 --exclude 引數的工作方式與往常相同，但使用者應指定“localhost”作為主機名。
另外請注意，CUDA_VISIBLE_DEVICES 不能與 DeepSpeed 一起使用來控制應該使用哪些裝置。例如，要僅使用當前節點的 gpu1，執行以下操作：
deepspeed --include localhost:1 ...
3. 自動張量並行
展示用於推理的新自動張量並行特性。
以前，使用者需要向 DeepSpeed 提供一個注入策略以啟用張量並行性。DeepSpeed現在預設支援HuggingFace模型的自動張量並行性，只要未啟用核心注入並且未提供注入策略。這允許我們的使用者在不提供注入策略的情況下提高當前未透過核心注入支援的模型的效能。下面是新方法的示例
# ---------------------------------------
# New automatic tensor parallelism method
# ---------------------------------------
import os
import torch
import transformers
import deepspeed
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
# 建立pipeline
pipe = transformers.pipeline(task = "text2text-generation", model="google/t5-v1_1-small", device = local_rank)
# 初始化DeepSpeed推理引擎
pipe.model = deepspeed.init_inference(
pipe.model,
mp_size = world_size,
dtype = torch.float
)
output = pipe('Input String')
以前，要為不支援核心注入的模型執行僅具有張量並行性的推理，您可以傳遞一個注入策略，該策略顯示 Transformer Encoder/Decoder 的兩個特定線性層：1）attention output GeMM 和2）layer output GeMM。
我們需要該層的這些部分來新增GPU之間所需的 all-reduce communication，以合併模型並行 ranks 中的部分結果。下面，我們展示了前面方法的一個示例：
# ----------------------------------
# Previous tensor parallelism method
# ----------------------------------
import os
import torch
import transformers
import deepspeed
from transformers.models.t5.modeling_t5 import T5Block
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
# create the model pipeline
pipe = transformers.pipeline(task="text2text-generation", model="google/t5-v1_1-small", device=local_rank)
# Initialize the DeepSpeed-Inference engine
pipe.model = deepspeed.init_inference(
pipe.model,
mp_size=world_size,
dtype=torch.float,
injection_policy={T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')}
)
output = pipe('Input String')
使用自動張量並行性，我們不需要為支援的模型提供注入策略。注入策略將在執行時確定並自動應用。
示例指令碼
我們可以使用 inference test suite 觀察到自動張量並行性的效能改進。
此指令碼用於測試文字生成模型，包括每個token延遲、頻寬、吞吐量和記憶體檢查以進行比較。
需要提前下載該專案 https://github.com/microsoft/DeepSpeedExamples.git
-
使用以下命令在沒有DeepSpeed和張量並行性的情況下執行。設定test_performance標誌以收集效能資料：
deepspeed --num_gpus 1 DeepSpeedExamples/inference/huggingface/text-generation/inference-test.py --model /root/autodl-tmp/cache/ckpt/Qwen7B-chat --batch_size 1 --test_performance -
要啟用張量並行性，您需要為相容模型使用標誌ds_inference
deepspeed --num_gpus 1 DeepSpeedExamples/inference/huggingface/text-generation/inference-test.py --name /root/autodl-tmp/cache/ckpt/Qwen7B-chat --batch_size 1 --test_performance --ds_inference
4. Megatron-LM GPT2
我們將向Megatron-LM GPT2模型新增DeepSpeed，這是一個大型、強大的transformer。
Megatron-LM支援模型並行和多節點訓練。
首先，我們討論資料和環境設定，以及如何使用 original Megatron-LM 訓練GPT-2模型。接下來，我們逐步使該模型能夠與DeepSpeed一起執行。最後，我們展示了使用DeepSpeed的效能提升和記憶體佔用減少。
https://www.deepspeed.ai/tutorials/megatron/
5. Zero Redundancy Optimizer
在本節，我們將把ZeRO最佳化器應用於 Megatron-LM GPT-2模型。
ZeRO是一套強大的記憶體最佳化技術，能夠有效訓練具有數萬億個引數的大型模型，例如GPT-2和Turing-NLG 17B。
與模型並行方法相比，ZeRO的一個關鍵吸引力在於不需要修改模型程式碼。
與模型並行方法相比，ZeRO的一個關鍵吸引力在於不需要修改模型，所需要的只是更改DeepSpeed配置JSON中的一些配置。
5.1 ZeRO概述
ZeRO利用資料並行的聚合計算和記憶體資源來減少用於模型訓練的每個裝置（GPU）的記憶體和計算需求。
ZeRO透過在分散式訓練硬體中的可用裝置（GPU和CPU）之間劃分各種模型訓練狀態（權重、梯度和最佳化器狀態）來減少每個GPU的記憶體消耗。具體來說，ZeRO被實現為最佳化的增量階段，早期階段的最佳化在後期階段可用。要深入瞭解ZeRO，請參閱論文。
- 最佳化器狀態（例如，對於Adam最佳化器，32位權重以及第一和第二矩估計(動量估計? moment estimates )）在程序之間進行分割槽，以便每個程序僅更新其分割槽。
- 用於更新模型權重的32位梯度也被分割槽，以便每個程序僅保留與其最佳化器狀態部分相對應的梯度。
- 16位模型引數在程序之間進行分割槽。ZeRO-3將在 forward and backward 過程中自動收集和分割槽它們。
此外，ZeRO-3包括 infinity offload engine 以形成 ZeRO-Infinity ，它可以解除安裝到CPU和NVMe記憶體以節省大量記憶體。
5.2 訓練環境
在本練習中使用DeepSpeed Megatron-LM GPT-2程式碼，on NVIDIA Tesla V100-SXM3 Tensor Core GPUs with 32GB RAM.
5.3 開啟ZeRO最佳化
要為DeepSpeed模型啟用ZeRO最佳化，我們只需將zero_optimization鍵新增到DeepSpeed JSON配置中。
"zero_optimization": {
"stage": [0|1|2|3],
"allgather_partitions": [true|false],
"allgather_bucket_size": 5e8,
"overlap_comm": false,
"reduce_scatter": [true|false],
"reduce_bucket_size": 5e8,
"contiguous_gradients" : [true|false],
"offload_param": {
...
},
"offload_optimizer": {
...
},
"stage3_max_live_parameters" : 1e9,
"stage3_max_reuse_distance" : 1e9,
"stage3_prefetch_bucket_size" : 5e8,
"stage3_param_persistence_threshold" : 1e6,
"sub_group_size" : 1e12,
"elastic_checkpoint" : [true|false],
"stage3_gather_16bit_weights_on_model_save": [true|false],
"ignore_unused_parameters": [true|false]
"round_robin_gradients": [true|false]
"zero_hpz_partition_size": 1
"zero_quantized_weights": [true|false]
"zero_quantized_gradients": [true|false]
}
5.4 訓練1.5B引數的GPT-2模型
我們展示了ZeRO第1階段的優勢，展示了它可以在8個V100 GPU上對15億引數GPT-2模型進行資料並行訓練。
我們將訓練配置為每個裝置使用1個批處理大小，以確保記憶體消耗主要歸因於模型引數和最佳化器狀態。
--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階段1將最佳化器狀態在八個資料並行的ranks中進行分割槽，每個裝置上的記憶體消耗可以減少到2.25GB，從而使模型可以進行訓練。
-
為了啟用ZeRO階段1，我們只需簡單地更新DeepSpeed的JSON配置檔案。
{ "zero_optimization": { "stage": 1, "reduce_bucket_size": 5e8 # } }
從上面的NVIDIA-smi截圖中我們可以看到，只有編號為6和7的GPU被用於訓練模型。這表明在進行深度學習模型訓練時，只有部分顯示卡被利用。
使用ZeRO第一階段，我們可以透過增加資料並行度來進一步減少每個裝置上的記憶體消耗。這些節省下來的記憶體可以用來增加模型的大小和/或增加批次大小。這意味著透過減少記憶體消耗，可以有更多的資源來處理更大的模型或者更大的資料批次。
5.5 訓練10B引數的GPT-2模型
ZeRO第2階段最佳化進一步增加了可以使用資料並行訓練的模型的大小。
使用32個V100 GPU訓練具有10B引數的模型來進行展示。
-
首先，我們需要配置一個具有10B引數的模型，並且啟用啟用檢查點。
透過對GPT-2模型配置的修改以實現10B引數。
--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配置來啟用ZeRO階段2最佳化。
{ "zero_optimization": { "stage": 2, "contiguous_gradients": true, # 減少反向傳遞期間的記憶體碎片 "overlap_comm": true, # 將梯度的減少與反向傳播計算過程重疊進行 "reduce_scatter": true, # 梯度平均的過程中，使用了“reduce”或“reduce scatter”操作，而不是“allreduce”操作 "reduce_bucket_size": 5e8, # 控制reduce操作中處理的元素數量 "allgather_bucket_size": 5e8 # } }
5.6 使用ZeRO-Infinity訓練萬億規模的模型
ZeRO-3是ZeRO的第三階段，它對完整模型狀態（即權重、梯度和最佳化器狀態）進行分割槽，隨著資料並行度的增加，記憶體節省的效果也會線性增加。
使用ZeRO-Infinity解除安裝到CPU和NVMe
ZeRO-Infinity使用 DeepSpeed's infinity offload engine，專門用於將模型的狀態（即模型的引數和中間資料）從GPU解除安裝到CPU或NVMe儲存器。這意味著透過解除安裝模型狀態，可以支援訓練更大的模型，因為不需要在GPU上儲存所有模型資料。
{
"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 vs ZeRO-Offload
DeepSpeed首先包括ZeRO-Offload的解除安裝功能，這是一個將最佳化器和梯度狀態解除安裝到ZeRO-2中CPU記憶體的系統。
ZeRO-Infinity是ZeRO-3可訪問的下一代解除安裝功能。ZeRO-Infinity能夠解除安裝比ZeRO-Offload更多的資料，並且具有更有效的頻寬利用和計算和通訊的重疊。
分配巨大的Megatron-LM模型
我們對模型初始化進行了兩項進一步的更改，以支援超過本地系統記憶體但不超過總系統記憶體的模型。
-
以記憶體可擴充套件的方式分配模型：模型引數將被分配並立即在資料並行組中進行分割槽。如果remote_device是“cpu”或“nvme”，模型也將分配在CPU/NVMe記憶體中，而不是GPU記憶體中。
with deepspeed.zero.Init(data_parallel_group=mpu.get_data_parallel_group(), # 資料並行組，用於在多個裝置上分配模型引數 remote_device=get_args().remote_device, # 指定模型引數應分配到的裝置型別（如CPU、NVMe或GPU） enabled=get_args().zero_stage==3): # 當zero_stage==3時，啟用這種初始化方式 model = GPT2Model(num_tokentypes=0, parallel_output=True) -
收集嵌入權重以進行初始化：DeepSpeed將在其建構函式期間以及前向和後向傳播期間自動收集模組的引數。但是，額外的訪問必須與DeepSpeed協調，以確保收集引數資料並隨後進行分割槽。如果張量被修改，還應使用modifier_rank引數來確保所有 ranks 對資料具有一致的檢視。
self.position_embeddings = torch.nn.Embedding(...) # 上下文管理器，用於確保在初始化引數時，所有節點對引數有一致的檢視 with deepspeed.zero.GatheredParameters(self.position_embeddings.weight, modifier_rank=0): # Initialize the position embeddings. self.init_method(self.position_embeddings.weight) ... self.tokentype_embeddings = torch.nn.Embedding(...) # 上下文管理器 with deepspeed.zero.GatheredParameters(self.tokentype_embeddings.weight, modifier_rank=0): # Initialize the token-type embeddings. self.init_method(self.tokentype_embeddings.weight)
以記憶體為中心的平鋪? tiling ???
-
ZeRO-Infinity包含了對線性層的替代方案，這種方案能夠進一步減少記憶體的使用。我們可以選擇性的將每一個transformer層中的模型並行線性層進行 tile 操作 ???。模型並行和分割 (tiling) 可以透過在構建層時指定相應的基類來結合使用。
deepspeed.zero.TiledLinear模組利用了 data fetch and release pattern of ZeRO-3 來減少工作記憶體需求透過將一個大的操作分解為可以順序執行的小的tiles。模型並行是一種將模型的不同部分分配到不同的裝置（如GPU）上進行訓練的方法。
分割（tiling）是一種技術，將一個大的模型操作（如線性層）分解成多個小的操作，這些小操作可以並行執行，從而減少記憶體的使用。
-
我們包括來自 Megatron-LM的ParallelMLP的一個示例的更改。transformer.py中還有三個模型並行層類似地進行。
-
Megatron-LM的模型並行層有一種特殊的形式，其中層的加性偏置
bias被延遲，而是從 forward 返回並與後來的運算子融合。DeepSpeed的TiledLinear的deepspeed.zero.TiledLinearReturnBias子類也簡單地轉發返回的偏置bias引數而不累積。
@@ -1,6 +1,9 @@
-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)
請注意，我們按 input_size 和 output_size 成比例的縮放 in_splits 和 out_splits。
這會導致固定大小的 tiles [hidden/tile_factor, hidden/tile_factor]。
註冊外部引數
已棄用：DeepSpeed版本0.3.15引入了自動外部引數註冊，不再需要此步驟。
5.7 提取權重
如果您需要從DeepSpeed中取出預訓練的權重，您可以透過以下方式獲得fp16權重：
-
在ZeRO-2中，
state_dict包含 fp16 的模型權重，可以使用torch.save正常儲存。 -
在ZeRO-3中，
state_dict只包含佔位符，因為模型權重是跨多個GPU分割槽的。如果想要獲取權重，需要開啟：
"zero_optimization": { "stage3_gather_16bit_weights_on_model_save": true },然後透過如下程式碼進行儲存：
if self.deepspeed: self.deepspeed.save_16bit_model(output_dir, output_file)因為它需要將權重合併到一個GPU上，所以它可能會很慢並且需要記憶體，因此僅在需要時使用此功能。
如果
stage3_gather_16bit_weights_on_model_save = False，沒有權重會被儲存。也可以使用這種方法來節省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)
儲存檢查點時會自動建立 zero_to_fp32.py 指令碼。
注意：目前此指令碼使用最終檢查點大小的2倍記憶體（通用RAM）。
或者，如果您有大量空閒CPU記憶體，並且您希望模型更新為其fp32權重而不是獲取檔案，您可以在訓練結束時執行以下操作
from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
fp32_model = load_state_dict_from_zero_checkpoint(deepspeed.module, checkpoint_dir)
請注意，該模型將有利於儲存，但不再適合繼續訓練，並且將需要重新使用 deepspeed.initialize()。
If you just want the state_dict, you can do:
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)
6. 自動微調？Autotuning
https://www.deepspeed.ai/tutorials/autotuning/
引數解析
DeepSpeed使用argparse庫為DeepSpeed執行時提供命令列配置。
使用 deepspeed.add_config_arguments() 將DeepSpeed的內建引數新增到應用程式的解析器中。
import argparse
import deepspeed
parser = argparse.ArgumentParser(description="train")
parser.add_argument("--local_rank", type=int, default=-1, help="local rank passed from distributed launcher")
# 包含deepspeed配置引數
parser = deepspeed.add_config_arguments(parser)
cmd_args = parser.parse_args()
print(cmd_args)
'''
Namespace(deepscale=False, deepscale_config=None, deepspeed=False, deepspeed_config=None, local_rank=-1)
二、結合MiniGPT4Qwen
1. 初始化分散式訓練模式
- DeepSpeed的
init_distributed函式來初始化分散式環境 - PyTorch的
init_process_group函式來初始化分散式環境
def init_deepspeed_distributed_mode(args):
import deepspeed
# 檢查環境變數中是否存在RANK和WORLD_SIZE
# 這些變數通常在分散式訓練環境中設定，用於指示當前程序的全域性索引和總程序數
if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
args.rank = int(os.environ["RANK"])
args.world_size = int(os.environ["WORLD_SIZE"])
# 當前程序應使用的GPU索引
args.gpu = int(os.environ["LOCAL_RANK"])
# 檢查是否存在SLURM_PROCID環境變數，通常在SLURM作業排程系統中使用
elif "SLURM_PROCID" in os.environ:
args.rank = int(os.environ["SLURM_PROCID"])
args.gpu = args.rank % torch.cuda.device_count()
else:
print("Not using distributed mode")
args.distributed = False
return
args.distributed = True
# 當前程序的GPU設定為args.gpu指定的裝置
torch.cuda.set_device(args.gpu)
# 分散式後端為NCCL, GPU通訊庫
args.dist_backend = "nccl"
print(
"| distributed init (rank {}, world {}): {}".format(
args.rank, args.world_size, args.dist_url
),
flush=True,
)
# 初始化分散式環境
deepspeed.init_distributed(
dist_backend=args.dist_backend,
init_method=args.dist_url, # "env://"
world_size=args.world_size,
rank=args.rank,
timeout=datetime.timedelta(
days=365
), # allow auto-downloading and de-compressing
)
# 在所有程序中同步，確保所有程序都到達此點後再繼續執行。
torch.distributed.barrier()
# 僅在主程序（rank為0）上執行一些設定
setup_for_distributed(args.rank == 0)
def init_distributed_mode(args):
if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
args.rank = int(os.environ["RANK"])
args.world_size = int(os.environ["WORLD_SIZE"])
args.gpu = int(os.environ["LOCAL_RANK"])
elif "SLURM_PROCID" in os.environ:
args.rank = int(os.environ["SLURM_PROCID"])
args.gpu = args.rank % torch.cuda.device_count()
else:
print("Not using distributed mode")
args.distributed = False
return
args.distributed = True
torch.cuda.set_device(args.gpu)
args.dist_backend = "nccl"
print(
"| distributed init (rank {}, world {}): {}".format(
args.rank, args.world_size, args.dist_url
),
flush=True,
)
# PyTorch的init_process_group函式來初始化分散式環境
torch.distributed.init_process_group(
backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank,
timeout=datetime.timedelta(
days=365
), # allow auto-downloading and de-compressing
)
torch.distributed.barrier()
setup_for_distributed(args.rank == 0)