DeepSpeed¶
DeepSpeed 是一个 PyTorch 优化库，它可以让你在分布式环境中高效地训练大型模型。DeepSpeed 的核心功能是 ZeRO Redundancy Optimizer (ZeRO)，它能够实现在大规模上训练大型模型。ZeRO 分为几个阶段：
- ZeRO-1：将优化器状态分布在多个 GPU 上
- ZeRO-2：将梯度分布在多个 GPU 上
- ZeRO-3：将参数分布在多个 GPU 上
在 GPU 资源有限的环境下，ZeRO 还可以将优化器的内存和计算从 GPU 卸载到 CPU，以便在单个 GPU 上训练非常大的模型。DeepSpeed 已经与 Transformers 的 Trainer 类集成，支持所有 ZeRO 阶段和卸载功能。你只需要提供一个配置文件，或者使用提供的模板。对于推理，Transformers 支持 ZeRO-3 和卸载，因为这允许加载巨大的模型。
本指南将引导你如何部署 DeepSpeed 训练，介绍你可以启用的功能，如何为不同的 ZeRO 阶段和卸载设置配置文件，以及如何在不使用 Trainer 的情况下使用 DeepSpeed。
安装¶
DeepSpeed 可以通过 PyPI 或 Transformers 安装。更多详细的安装选项，请参阅 DeepSpeed 的 安装详情 或 GitHub 的 README。
如果你在安装 DeepSpeed 时遇到困难，可以查看 DeepSpeed CUDA 安装 指南。尽管 DeepSpeed 有一个可通过 pip 安装的 PyPI 包，但强烈建议 从源码安装，以更好地匹配你的硬件并支持某些特性（如 1-bit Adam），这些特性在 PyPI 分发中不可用。
通过 PyPI 安装¶
pip install deepspeed
通过 Transformers 安装¶
pip install transformers[deepspeed]
内存需求¶
在开始之前，最好检查一下你的 GPU 和 CPU 是否有足够的内存来容纳你的模型。DeepSpeed 提供了一个工具来估算所需的 CPU/GPU 内存。例如，要估计 bigscience/T0_3B 模型在一个 GPU 上的内存需求：
python -c 'from transformers import AutoModel; \
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
model = AutoModel.from_pretrained("bigscience/T0_3B"); \
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=1, num_nodes=1)'
输出可能如下所示：
这意味着你需要一个 80GB 的 GPU（不使用 CPU 卸载）或一个 8GB 的 GPU 和大约 60GB 的 CPU 内存来进行卸载（这些只是参数、优化器状态和梯度的内存需求，你还需要一些额外的内存用于 CUDA 内核和激活）。你还应该考虑成本和速度之间的权衡，因为租用或购买较小的 GPU 会更便宜，但训练模型的时间会更长。
如果 GPU 内存足够，确保禁用 CPU/NVMe 卸载以加快速度。
选择 ZeRO 阶段¶
安装了 DeepSpeed 并了解了内存需求后，下一步是选择一个 ZeRO 阶段。按速度和内存效率排序：
| 最快 | 内存最省 |
|---|---|
| ZeRO-1 | ZeRO-3 + 卸载 |
| ZeRO-2 | ZeRO-3 |
| ZeRO-2 + 卸载 | ZeRO-2 + 卸载 |
| ZeRO-3 | ZeRO-2 |
| ZeRO-3 + 卸载 | ZeRO-1 |
为了找到最适合你的方案，可以从最快的方式开始尝试，如果内存不足，再尝试下一级别，虽然速度会稍慢但内存效率更高。你可以根据自己的偏好从最节省内存的方案开始尝试，直到找到速度和内存使用的最佳平衡点。
一个通用的过程是（从批量大小为 1 开始）：
- 启用梯度检查点
- 尝试 ZeRO-2
- 尝试 ZeRO-2 并卸载优化器
- 尝试 ZeRO-3
- 尝试 ZeRO-3 并将参数卸载到 CPU
- 尝试 ZeRO-3 并将参数和优化器都卸载到 CPU
- 尝试降低默认值，如在使用 generate() 方法时使用更窄的搜索束
- 尝试混合精度（在较旧的 GPU 架构上使用 fp16，在 Ampere GPU 上使用 bf16）
- 如果可能，添加更多硬件或启用 Infinity 以将参数和优化器卸载到 NVMe
- 一旦不再出现内存不足，测量有效吞吐量，然后尝试尽可能增加批量大小以最大化 GPU 效率
- 最后，尝试优化训练配置，禁用某些卸载功能或使用更快的 ZeRO 阶段，并调整批量大小以找到速度和内存使用之间的最佳权衡
DeepSpeed 配置文件¶
DeepSpeed 通过包含所有配置参数的配置文件与 Trainer 类配合使用。当你执行训练脚本时，DeepSpeed 会将从 Trainer 接收到的配置日志记录到控制台，以便你可以看到确切的配置。
完整的 DeepSpeed 配置选项列表可以在 DeepSpeed 配置 JSON 参考中找到。你还可以在 DeepSpeedExamples 存储库或主 DeepSpeed 存储库中找到更多实用的 DeepSpeed 配置示例。要快速找到特定示例，可以：
git clone https://github.com/microsoft/DeepSpeedExamples
cd DeepSpeedExamples
find . -name '*json'
# 查找使用 Lamb 优化器的示例
grep -i Lamb $(find . -name '*json')
DeepSpeed 配置文件可以通过命令行接口传递为 JSON 文件的路径，也可以在笔记本环境中作为嵌套的 dict 对象传递。
配置文件路径¶
嵌套字典¶
DeepSpeed 和 Trainer 参数¶
有三种类型的配置参数：
- 共享参数：有些配置参数由 Trainer 和 DeepSpeed 共享，当定义冲突时识别错误可能会很困难。为了简化问题，这些共享配置参数应从 Trainer 命令行参数中配置。
- 自动派生参数：有些配置参数是从模型配置中自动派生的，因此你不需要手动调整这些值。Trainer 使用配置值
auto来确定最正确或高效的值。你可以显式设置自己的配置参数，但必须确保 Trainer 参数和 DeepSpeed 配置参数一致。不一致可能导致训练失败，而且很难检测到！ - 仅限 DeepSpeed 的参数：有些配置参数仅适用于 DeepSpeed，需要根据你的训练需求手动设置。
你还可以修改 DeepSpeed 配置并编辑 TrainingArguments：
- 创建或加载一个 DeepSpeed 配置作为主要配置
- 基于这些 DeepSpeed 配置值创建一个 TrainingArguments 对象
一些值，如 scheduler.params.total_num_steps，是在训练过程中由 Trainer 计算的。
ZeRO 配置¶
有三个配置，每个对应一个不同的 ZeRO 阶段。阶段 1 对于可扩展性来说并不是很有吸引力，本指南关注的是阶段 2 和 3。zero_optimization 配置包含了所有启用和配置的选项。有关每个参数的详细说明，请参阅 DeepSpeed 配置 JSON 参考。
DeepSpeed 不会验证参数名称，任何拼写错误都会回退到参数的默认设置。你可以观看 DeepSpeed 引擎启动日志消息，以了解它将使用哪些值。
以下配置必须通过 DeepSpeed 设置，因为 Trainer 没有提供相应的命令行参数。
ZeRO-1¶
ZeRO-1 将优化器状态分片分布在多个 GPU 上，你可以期望有轻微的速度提升。ZeRO-1 配置可以这样设置：
NVMe 配置¶
ZeRO-Infinity 允许将模型状态卸载到 CPU 和/或 NVMe 以节省更多内存。智能分区和切片算法允许每个 GPU 在卸载期间发送和接收非常少量的数据，使得现代 NVMe 可以提供比你可用的训练进程更大的总内存池。ZeRO-Infinity 需要 ZeRO-3。
根据可用的 CPU 和/或 NVMe 内存，你可以卸载 优化器状态 和 参数，只卸载其中一个，或者都不卸载。你还需要确保 nvme_path 指向 NVMe 设备，因为尽管它仍然可以使用普通硬盘或固态硬盘，但速度会显著变慢。使用现代 NVMe，你可以期望读取峰值速度约为 3.5GB/s，写入峰值速度约为 3GB/s。最后，运行基准测试 以确定你的训练设置的最佳 aio 配置。
以下示例 ZeRO-3/Infinity 配置文件将大多数参数值设置为 auto，但你也可以手动添加这些值。
DeepSpeed 功能¶
在 DeepSpeed 配置文件中指定许多重要的参数，本节简要介绍了这些参数。
激活/梯度检查点¶
激活和梯度检查点交换了速度以换取更多的 GPU 内存，这可以帮助你在 GPU 内存不足或增加批量大小以提高性能的情况下解决问题。要启用此功能：
- 对于 Hugging Face 模型，设置
model.gradient_checkpointing_enable()或在 Trainer 中使用--gradient_checkpointing。 - 对于非 Hugging Face 模型，使用 DeepSpeed 的 激活检查点 API。你还可以替换 Transformers 的建模代码，将
torch.utils.checkpoint替换为 DeepSpeed API。这种方法更灵活，因为你可以将前向激活卸载到 CPU 内存而不是重新计算它们。
优化器和调度器¶
DeepSpeed 和 Transformers 的优化器和调度器可以混合使用，只要你没有启用 offload_optimizer。当 offload_optimizer 被启用时，你可以使用非 DeepSpeed 优化器（除了 LAMB），但前提是它同时具有 CPU 和 GPU 实现。
配置文件中的优化器和调度器参数可以从命令行设置，以避免难以发现的错误。例如，如果学习率在其他地方设置了不同的值，你可以从命令行覆盖它。除了优化器和调度器参数外，你还需要确保 Trainer 命令行参数与 DeepSpeed 配置一致。
优化器¶
DeepSpeed 提供了几种 优化器（Adam, AdamW, OneBitAdam, 和 LAMB），但你也可以从 PyTorch 中导入其他优化器。如果你没有在配置中配置优化器，Trainer 会自动选择 AdamW，并使用命令行提供的值或默认值设置以下参数：lr，adam_beta1，adam_beta2，adam_epsilon，weight_decay。
你可以将参数设置为 "auto" 或手动输入你希望的值。
你还可以使用不受支持的优化器，只需在顶级配置中添加以下内容：
从 DeepSpeed==0.8.3 开始，如果你想使用卸载，还需要在顶级配置中添加以下内容，因为卸载与 DeepSpeed 的 CPU Adam 优化器配合效果最佳。
精度¶
Deepspeed 支持 fp32、fp16 和 bf16 混合精度。
- fp32：全精度浮点数
- fp16：半精度浮点数
- bf16：脑浮点数（BFloat16）
如果模型不支持混合精度，例如没有在混合精度下预训练，你可能会遇到溢出或下溢问题，导致损失为 NaN。在这种情况下，你应该显式禁用默认的 fp16 模式，使用全 fp32 精度。
对于 Ampere GPU 和 PyTorch > 1.7，它会自动将某些操作切换为更高效的 tf32 格式，但结果仍然是 fp32。你可以在 Trainer 中通过设置 --tf32 来启用它，通过 --tf32 0 或 --no_tf32 来禁用它。
批量大小¶
批量大小可以自动配置或显式设置。如果你选择使用 "auto" 选项，Trainer 会将 train_micro_batch_size_per_gpu 设置为 args.per_device_train_batch_size 的值，将 train_batch_size 设置为 args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps。
通信数据类型¶
对于诸如归约、聚集和散射等通信集体操作，使用不同的数据类型。
所有聚集和散射操作都在相同的数据类型下进行。例如，如果你使用 bf16 训练，数据也会以 bf16 形式聚集，因为聚集是非损耗操作。
归约操作是损耗的，例如当梯度在多个 GPU 上平均时。当通信以 fp16 或 bf16 进行时，更容易发生损耗，因为在低精度下加法运算不是精确的。特别是 bf16，其精度低于 fp16。因此，默认情况下，归约操作使用 fp16，因为在平均梯度时损失最小。
你可以在配置文件中通过设置 communication_data_type 参数来选择通信数据类型。例如，选择 fp32 会增加少量开销，但确保归约操作在 fp32 下累积，当准备就绪时，再降精度到你正在训练的半精度数据类型。
部署¶
DeepSpeed 可以通过不同的启动器部署，如 torchrun、deepspeed 启动器或 Accelerate。要部署，可以在 Trainer 命令行中添加 --deepspeed ds_config.json。推荐使用 DeepSpeed 的 add_config_arguments 实用程序来添加任何必要的命令行参数到你的代码中。
本指南将展示如何使用 deepspeed 启动器在不同的训练设置中部署 DeepSpeed。你可以查看这个 帖子 获取更多实际使用示例。
多 GPU 部署¶
要在多个 GPU 上部署 DeepSpeed，添加 --num_gpus 参数。如果你想要使用所有可用的 GPU，不需要添加 --num_gpus。以下示例使用 2 个 GPU。
deepspeed --num_gpus=2 examples/pytorch/translation/run_translation.py \
--deepspeed tests/deepspeed/ds_config_zero3.json \
--model_name_or_path google-t5/t5-small --per_device_train_batch_size 1 \
--output_dir output_dir --overwrite_output_dir --fp16 \
--do_train --max_train_samples 500 --num_train_epochs 1 \
--dataset_name wmt16 --dataset_config "ro-en" \
--source_lang en --target_lang ro
多节点部署¶
一个节点是用于运行工作负载的一个或多个 GPU。更强大的设置是多节点设置，可以使用 deepspeed 启动器启动。假设有两个节点，每个节点有 8 个 GPU。第一个节点可以通过 ssh hostname1 访问，第二个节点可以通过 ssh hostname2 访问。两个节点必须能够在本地通过 ssh 无密码通信。
默认情况下，DeepSpeed 期望你的多节点环境使用共享存储。如果这不是情况，每个节点只能看到本地文件系统，你需要调整配置文件，包括一个 checkpoint，以允许在没有共享文件系统访问的情况下加载：
torchrun --nproc_per_node=8 --nnode=2 --node_rank=0 --master_addr=hostname1 \
--master_port=9901 your_program.py <normal cl args> --deepspeed ds_config.json
使用 SLURM¶
在 SLURM 环境中，你需要根据你的具体 SLURM 环境调整 SLURM 脚本。一个示例 SLURM 脚本可能如下所示：
#SBATCH --job-name=test-nodes # job name
#SBATCH --nodes=2 # number of nodes
#SBATCH --ntasks-per-node=1 # crucial - only 1 task per dist per node!
#SBATCH --cpus-per-task=10 # number of cores per task
#SBATCH --gres=gpu:8 # number of GPUs
#SBATCH --time 20:00:00 # maximum execution time (HH:MM:SS)
#SBATCH --output=%x-%j.out # output file name
export GPUS_PER_NODE=8
export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MASTER_PORT=9901
srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run \
--nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
--master_addr $MASTER_ADDR --master_port $MASTER_PORT \
your_program.py <normal cl args> --deepspeed ds_config.json'
然后你可以使用以下命令安排多节点部署，该命令在同一时间在所有节点上启动训练。
sbatch launch.slurm
笔记本¶
deepspeed 启动器不支持从笔记本部署，因此你需要模拟分布式环境。然而，这仅适用于 1 个 GPU。如果你想要使用多个 GPU，必须使用多进程环境，这意味着你必须使用 deepspeed 启动器，而不能像这里所示那样模拟。
# DeepSpeed 需要在只有一个进程的情况下也需要分布式环境
# 这在笔记本中模拟启动器
import os
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "9994" # 如果出现 RuntimeError: Address already in use，修改此端口
os.environ["RANK"] = "0"
os.environ["LOCAL_RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
# 现在正常进行，加上 DeepSpeed 配置文件
training_args = TrainingArguments(..., deepspeed="ds_config_zero3.json")
trainer = Trainer(...)
trainer.train()
如果你想要在当前目录中的笔记本中动态创建配置文件，可以在一个专用的单元格中这样做。
如果训练脚本在一个文件中而不是笔记本单元格中，你可以从笔记本单元格中的 shell 正常启动 deepspeed。例如，启动 run_translation.py：
!git clone https://github.com/huggingface/transformers
!cd transformers; deepspeed examples/pytorch/translation/run_translation.py ...
你也可以使用 %%bash 魔法命令编写多行代码来运行 shell 程序，但你将无法在训练完成前查看日志。使用 %%bash 魔法命令时，你不需要模拟分布式环境。
%%bash
git clone https://github.com/huggingface/transformers
cd transformers
deepspeed examples/pytorch/translation/run_translation.py ...
保存模型权重¶
DeepSpeed 将主要的全精度 fp32 权重存储在自定义检查点优化器文件中（路径模式类似于 global_step*/*optim_states.pt），并保存在常规检查点下。
- fp16：使用 ZeRO-2 训练的模型将
pytorch_model.bin权重保存为 fp16。要在使用 ZeRO-3 训练的模型中以 fp16 保存模型权重，你需要设置"stage3_gather_16bit_weights_on_model_save": true，因为模型权重被分片分布在多个 GPU 上。否则，Trainer 不会以 fp16 保存权重，也不会创建pytorch_model.bin文件。这是因为 DeepSpeed 的state_dict包含占位符而不是实际权重，你将无法加载它们。
ZeRO 推理¶
ZeRO 推理 将模型权重放置在 CPU 或 NVMe 内存中，以减轻 GPU 的负担，从而在 GPU 上运行巨大模型的推理成为可能。推理不需要大量的额外内存来存储优化器状态和梯度，因此你可以在相同的硬件上处理更大的批量或序列长度。
ZeRO 推理使用与 ZeRO-3 相同的配置文件，ZeRO-2 和 ZeRO-1 的配置文件不会起作用，因为它们对推理没有好处。
要运行 ZeRO 推理，将常用的训练参数传递给 TrainingArguments 类，并添加 --do_eval 参数。
deepspeed --num_gpus=2 your_program.py <normal cl args> --do_eval --deepspeed ds_config.json
无 Trainer 的 DeepSpeed 集成¶
DeepSpeed 也可以在不使用 Trainer 类的情况下与 Transformers 一起使用。这由 HfDeepSpeedConfig 处理，它负责在调用 from_pretrained() 时收集 ZeRO-3 参数并将模型拆分到多个 GPU 上。
如果你希望一切都能自动处理，尝试使用 DeepSpeed 与 Trainer！你需要遵循 DeepSpeed 文档，并在配置文件中手动配置参数值（不能使用 "auto" 值）。
要高效地部署 ZeRO-3，必须在实例化模型之前实例化 HfDeepSpeedConfig 对象，并保持该对象存活：
预训练模型¶
from transformers.integrations import HfDeepSpeedConfig
from transformers import AutoModel
import deepspeed
ds_config = {...} # deepspeed 配置对象或文件路径
# 必须在实例化模型之前运行以检测零 3
dschf = HfDeepSpeedConfig(ds_config) # 保持此对象存活
model = AutoModel.from_pretrained("openai-community/gpt2")
engine = deepspeed.initialize(model=model, config_params=ds_config, ...)
非预训练模型¶
from transformers.integrations import HfDeepSpeedConfig
from transformers import AutoModel
import deepspeed
ds_config = {...} # deepspeed 配置对象或文件路径
# 必须在实例化模型之前运行以检测零 3
dschf = HfDeepSpeedConfig(ds_config) # 保持此对象存活
model = AutoModel.from_config(config=your_custom_config)
engine = deepspeed.initialize(model=model, config_params=ds_config, ...)
无 Trainer 的 ZeRO 推理¶
要在无法将模型放入单个 GPU 的情况下运行 ZeRO 推理，可以尝试使用额外的 GPU 或/和卸载到 CPU 内存。理解的重要细节是 ZeRO 的设计方式，你可以在不同的 GPU 上并行处理不同的输入。
确保：
- 如果你有足够的 GPU 内存，禁用 CPU 卸载（因为它会减慢速度）。
- 如果你有 Ampere 或更新的 GPU，启用 bf16 以加快速度。如果没有这些 GPU，你可以启用 fp16，只要你不使用在 bf16 上预训练的模型（如 T5 模型），因为这可能导致溢出错误。
查看以下脚本，以更好地了解如何在无法将模型放入单个 GPU 的情况下运行 ZeRO 推理。
#!/usr/bin/env python
# 本脚本演示了如何在无法将模型放入单个 GPU 的情况下使用 Deepspeed ZeRO 进行推理。
#
# 1. 使用 1 个 GPU 和 CPU 卸载
# 2. 或者使用多个 GPU
#
# 首先需要安装 deepspeed: pip install deepspeed
#
# 这里我们使用 3B 的 "bigscience/T0_3B" 模型，需要大约 15GB GPU 内存 - 因此 1 个大 GPU 或 2
# 个小 GPU 可以处理它。或者 1 个小 GPU 和大量 CPU 内存。
#
# 要使用更大的模型如 "bigscience/T0"，需要大约 50GB，除非你有一个 80GB GPU -
# 你需要 2-4 个 GPU。然后你可以调整脚本以处理更多 GPU，如果你想一次处理多个输入。
#
# 提供的 deepspeed 配置还启用了 CPU 内存卸载，所以如果有大量可用的 CPU 内存并且不介意速度变慢，你应该能够加载一个通常无法放入单个 GPU 的模型。如果你有足够的 GPU 内存，程序会在不卸载到 CPU 的情况下运行得更快 - 因此禁用该部分。
#
# 在 1 个 GPU 上部署：
#
# deepspeed --num_gpus 1 t0.py
# 或者：
# python -m torch.distributed.run --nproc_per_node=1 t0.py
#
# 在 2 个 GPU 上部署：
#
# deepspeed --num_gpus 2 t0.py
# 或者：
# python -m torch.distributed.run --nproc_per_node=2 t0.py
from transformers import AutoTokenizer, AutoConfig, AutoModelForSeq2SeqLM
from transformers.integrations import HfDeepSpeedConfig
import deepspeed
import os
import torch
os.environ["TOKENIZERS_PARALLELISM"] = "false" # 避免 tokenizers 并行警告
# 分布式设置
local_rank = int(os.getenv("LOCAL_RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
torch.cuda.set_device(local_rank)
deepspeed.init_distributed()
model_name = "bigscience/T0_3B"
config = AutoConfig.from_pretrained(model_name)
model_hidden_size = config.d_model
# 批量大小必须能被 world_size 整除，但可以大于 world_size
train_batch_size = 1 * world_size
# ds_config 注释
#
# - 如果你使用 Ampere 或更高版本的 GPU，启用 bf16 - 这将以混合精度运行并更快。
#
# - 对于较旧的 GPU，你可以启用 fp16，但它只适用于非 bf16 预训练的模型 - 例如，所有官方的 t5 模型都是 bf16 预训练的
#
# - 如果不想使用 CPU 卸载，将 offload_param.device 设置为 "none" 或完全删除 offload_param 部分
#
# - 如果使用 offload_param，你可以手动微调 stage3_param_persistence_threshold 以控制哪些参数保留在 GPU 上 - 值越大，卸载的大小越小
#
# 更多关于 Deepspeed 配置的详细信息请参阅
# https://huggingface.co/docs/transformers/main/main_classes/deepspeed
# 保持与 json 一致的格式，只是将 true/false 用小写表示
# fmt: off
ds_config = {
"fp16": {
"enabled": False
},
"bf16": {
"enabled": False
},
"zero_optimization": {
"stage": 3,
"offload_param": {
"device": "cpu",
"pin_memory": True
},
"overlap_comm": True,
"contiguous_gradients": True,
"reduce_bucket_size": model_hidden_size * model_hidden_size,
"stage3_prefetch_bucket_size": 0.9 * model_hidden_size * model_hidden_size,
"stage3_param_persistence_threshold": 10 * model_hidden_size
},
"steps_per_print": 2000,
"train_batch_size": train_batch_size,
"train_micro_batch_size_per_gpu": 1,
"wall_clock_breakdown": False
}
# fmt: on
# 下一行指示 transformers 在调用模型的 from_pretrained 方法时直接将模型拆分到多个 GPU 上。
#
# **必须在加载模型 AutoModelForSeq2SeqLM.from_pretrained(model_name) 之前运行**
#
# 否则模型将首先正常加载，仅在前向传播时拆分，这效率较低，且当 CPU 内存较少时可能会失败
dschf = HfDeepSpeedConfig(ds_config) # 保持此对象存活
# 现在可以加载模型。
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# 初始化 Deepspeed ZeRO 并仅保留引擎对象
ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
ds_engine.module.eval() # 推理
# Deepspeed ZeRO 可以在每个 GPU 上处理不相关的输入。所以对于 2 个 GPU，你可以一次处理 2 个输入。
# 如果使用更多 GPU，相应地调整输入数量。
# 当然，如果只有一个输入要处理，你需要将相同的字符串传递给所有 GPU
# 如果只使用一个 GPU，则只有 rank 0
rank = torch.distributed.get_rank()
if rank == 0:
text_in = "Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy"
elif rank == 1:
text_in = "Is this review positive or negative? Review: this is the worst restaurant ever"
tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer.encode(text_in, return_tensors="pt").to(device=local_rank)
with torch.no_grad():
outputs = ds_engine.module.generate(inputs, synced_gpus=True)
text_out = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"rank{rank}:\n in={text_in}\n out={text_out}")
将脚本保存为 t0.py 并启动：
$ deepspeed --num_gpus 2 t0.py
rank0:
in=Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy
out=Positive
rank1:
in=Is this review positive or negative? Review: this is the worst restaurant ever
out=negative
这是一个非常基本的示例，你需要根据自己的用途进行调整。
生成¶
使用多个 GPU 和 ZeRO-3 进行生成时，需要通过在 generate() 方法中设置 synced_gpus=True 来同步 GPU。否则，如果一个 GPU 比另一个 GPU 先完成生成，整个系统会挂起，因为剩余的 GPU 尚未从第一个完成的 GPU 收到权重分片。
对于 Transformers >= 4.28，如果在生成过程中检测到多个 GPU，synced_gpus 会自动设置为 True。
故障排除¶
当你遇到问题时，应该考虑是否是 DeepSpeed 导致的问题，因为往往并非如此（除非你在异常中看到了 DeepSpeed 模块）！第一步应该是不使用 DeepSpeed 重试你的设置，如果问题仍然存在，那么你可以报告问题。如果问题是 DeepSpeed 核心问题且与 Transformers 集成无关，请在 DeepSpeed 仓库 中打开一个问题。
对于与 Transformers 集成相关的问题，请提供以下信息：
- 完整的 DeepSpeed 配置文件
- Trainer 的命令行参数，或如果你自己编写 Trainer 设置的 TrainingArguments 参数（不要转储 TrainingArguments，因为它包含数十个无关条目）
- 以下命令的输出：
python -c 'import torch; print(f"torch: {torch.__version__}")'
python -c 'import transformers; print(f"transformers: {transformers.__version__}")'
python -c 'import deepspeed; print(f"deepspeed: {deepspeed.__version__}")'
- 一个指向 Google Colab 笔记本的链接以重现问题
- 如果不可能，提供一个标准的非自定义数据集，我们也尝试使用现有的示例来重现问题
以下部分提供了解决最常见的两个问题的指南。
DeepSpeed 进程在启动时被杀死¶
当 DeepSpeed 进程在启动时被杀死且没有堆栈跟踪时，通常意味着程序尝试分配的 CPU 内存超过了系统的可用内存，或者你的进程尝试分配的 CPU 内存超过了允许的内存，导致操作系统内核终止进程。这种情况下，检查你的配置文件是否有 offload_optimizer、offload_param 或两者都设置为卸载到 CPU。
如果你有 NVMe 和 ZeRO-3 设置，尝试将数据卸载到 NVMe (估算 你的模型的内存需求)。
NaN 损失¶
NaN 损失通常发生在模型在 bf16 上预训练，然后你尝试使用 fp16 时（特别是在 TPU 上训练的模型）。要解决这个问题，使用 fp32 或 bf16（如果硬件支持，如 TPU、Ampere GPU 或更新版本）。
另一个可能是与使用 fp16 相关。例如，如果你的 fp16 配置如下：
你可能会在日志中看到以下 OVERFLOW! 消息：