DeepSpeed
DeepSpeed 是一个 PyTorch 优化库，它使分布式训练内存高效且快速。它的核心是 零冗余优化器 (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 位 Adam，这些功能在 PyPI 分发版中不可用。
pip install deepspeed
内存需求
在开始之前，最好检查一下您是否有足够的 GPU 和 CPU 内存来容纳您的模型。DeepSpeed 提供了一个工具来估计所需的 CPU/GPU 内存。例如，要估计在单个 GPU 上使用 bigscience/T0_3B 模型的内存需求
$ 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)'
[...]
Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 1 GPU per node.
SW: Model with 2783M total params, 65M largest layer params.
per CPU | per GPU | Options
70.00GB | 0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
70.00GB | 0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
62.23GB | 5.43GB | offload_param=none, offload_optimizer=cpu , zero_init=1
62.23GB | 5.43GB | offload_param=none, offload_optimizer=cpu , zero_init=0
0.37GB | 46.91GB | offload_param=none, offload_optimizer=none, zero_init=1
15.56GB | 46.91GB | offload_param=none, offload_optimizer=none, zero_init=0这意味着您需要一个 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 上使用 bf16）而不是全精度权重
- 如果可能，添加更多硬件或启用 Infinity 将参数和优化器卸载到 NVMe
- 在不再遇到内存不足的情况后，测量有效吞吐量，然后尝试将批次大小尽可能地增大以最大限度地提高 GPU 效率
- 最后，尝试通过禁用某些卸载功能来优化您的训练设置，或者使用更快的 ZeRO 阶段并增加/减少批次大小以找到速度和内存使用之间的最佳权衡
DeepSpeed 配置文件
DeepSpeed 通过配置文件与 Trainer 类协同工作，配置文件包含所有参数，用于配置您希望如何设置训练运行。当您执行训练脚本时，DeepSpeed 会将从 Trainer 收到的配置记录到控制台中，这样您就可以准确地看到使用了什么配置。
在 DeepSpeed 配置 JSON 参考中查找 DeepSpeed 配置选项的完整列表。您还可以在 DeepSpeedExamples 存储库或主要的 DeepSpeed 存储库中找到更多实际的 DeepSpeed 配置示例。要快速找到特定示例，您可以
git clone https://github.com/microsoft/DeepSpeedExamples
cd DeepSpeedExamples
find . -name '*json'
# find examples with the Lamb optimizer
grep -i Lamb $(find . -name '*json')DeepSpeed 配置文件作为 JSON 文件的路径传递，如果您是从命令行界面进行训练，或者作为嵌套的 dict 对象传递，如果您在笔记本设置中使用 Trainer。
TrainingArguments(..., deepspeed="path/to/deepspeed_config.json")DeepSpeed 和 Trainer 参数
有三种类型的配置参数
一些配置参数由 Trainer 和 DeepSpeed 共享，当存在冲突定义时，很难识别错误。为了简化操作，这些共享的配置参数是从 Trainer 命令行参数进行配置的。
一些配置参数会自动从模型配置中推断出来，因此您不需要手动调整这些值。 Trainer 使用配置值
auto来确定设置最正确或最有效的配置值。您可以显式设置自己的配置参数，但您必须注意确保 Trainer 参数和 DeepSpeed 配置参数一致。不匹配可能会导致训练以非常难以检测的方式失败！一些仅针对 DeepSpeed 的特定配置参数，需要根据您的训练需求手动设置。
您还可以修改 DeepSpeed 配置并从其中编辑 TrainingArguments
- 创建或加载 DeepSpeed 配置以用作主配置
- 基于这些 DeepSpeed 配置值创建一个 TrainingArguments 对象
在训练过程中，一些值（例如 scheduler.params.total_num_steps）由 Trainer 计算。
ZeRO 配置
有三种配置，分别对应于不同的 ZeRO 阶段。阶段 1 对于可扩展性不那么有趣，本指南侧重于阶段 2 和 3。 zero_optimization 配置包含所有选项，用于启用这些选项以及如何配置它们。有关每个参数的更详细说明，请查看 DeepSpeed 配置 JSON 参考。
以下配置必须在 DeepSpeed 中设置，因为 Trainer 没有提供等效的命令行参数。
ZeRO-1 将优化器状态跨 GPU 分片，你可以预期会有一点加速。ZeRO-1 配置可以这样设置
{
"zero_optimization": {
"stage": 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，但你也可以手动添加这些值。
{
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16,
"hysteresis": 2,
"min_loss_scale": 1
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "nvme",
"nvme_path": "/local_nvme",
"pin_memory": true,
"buffer_count": 4,
"fast_init": false
},
"offload_param": {
"device": "nvme",
"nvme_path": "/local_nvme",
"pin_memory": true,
"buffer_count": 5,
"buffer_size": 1e8,
"max_in_cpu": 1e9
},
"aio": {
"block_size": 262144,
"queue_depth": 32,
"thread_count": 1,
"single_submit": false,
"overlap_events": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}DeepSpeed 特性
本节简要介绍了 DeepSpeed 配置文件中需要指定的一些重要参数。
激活/梯度检查点
激活和梯度检查点会以更多 GPU 内存换取速度，这使你能够克服 GPU 内存不足的情况，或者增加你的批次大小以获得更好的性能。要启用此功能
- 对于 Hugging Face 模型，在 Trainer 中设置
model.gradient_checkpointing_enable()或--gradient_checkpointing。 - 对于非 Hugging Face 模型，请使用 DeepSpeed 激活检查点 API。你也可以替换 Transformers 模型代码，并将
torch.utils.checkpoint替换为 DeepSpeed API。这种方法更加灵活，因为你可以将正向激活卸载到 CPU 内存，而不是重新计算它们。
优化器和调度器
只要你没有启用 offload_optimizer，DeepSpeed 和 Transformers 的优化器和调度器就可以混合和匹配。当启用 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" 或手动输入你自己的所需值。
{
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
}
}你也可以通过在顶级配置中添加以下内容来使用不支持的优化器。
{
"zero_allow_untested_optimizer": true
}从 DeepSpeed==0.8.3 开始，如果你想使用卸载，你还需要将以下内容添加到顶级配置中，因为卸载与 DeepSpeed 的 CPU Adam 优化器配合得最好。
{
"zero_force_ds_cpu_optimizer": false
}精度
Deepspeed 支持 fp32、fp16 和 bf16 混合精度。
如果你的模型不能很好地使用混合精度，例如，如果它没有用混合精度进行预训练，你可能会遇到溢出或下溢问题，这会导致 NaN 损失。对于这些情况，你应该通过显式禁用默认的 fp16 模式来使用完整的 fp32 精度。
{
"fp16": {
"enabled": false
}
}对于 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。
{
"train_micro_batch_size_per_gpu": "auto",
"train_batch_size": "auto"
}梯度累积
梯度累积可以自动配置或显式设置。如果你选择使用 "auto" 选项，Trainer 会将它设置为 args.gradient_accumulation_steps 的值。
{
"gradient_accumulation_steps": "auto"
}
梯度裁剪
梯度裁剪可以自动配置或显式设置。如果你选择使用 "auto" 选项，Trainer 会将它设置为 args.max_grad_norm 的值。
{
"gradient_clipping": "auto"
}通信数据类型
对于像缩减、收集和散射操作这样的通信集体，使用单独的数据类型。
所有收集和散射操作都在数据所在的相同数据类型中执行。例如，如果您使用 bf16 进行训练，数据也会以 bf16 形式收集，因为收集是非损耗操作。
缩减操作是有损耗的，例如当梯度在多个 GPU 上进行平均时。当通信在 fp16 或 bf16 中完成时，它更有可能是有损耗的，因为在低精度下添加多个数字并不精确。对于 bf16 来说尤其如此，因为它的精度低于 fp16。为此，fp16 是缩减操作的默认值，因为在对梯度进行平均时，损耗最小。
您可以通过在配置文件中设置 communication_data_type 参数来选择通信数据类型。例如，选择 fp32 会增加少量开销，但可以确保缩减操作以 fp32 形式累加，并且在准备好后，它将被下转换为您正在训练的任何半精度数据类型。
{
"communication_data_type": "fp32"
}部署
DeepSpeed 可以通过不同的启动器进行部署，例如 torchrun、deepspeed 启动器或 Accelerate。要部署，请将 --deepspeed ds_config.json 添加到 Trainer 命令行。建议使用 DeepSpeed 的 add_config_arguments 实用程序为您的代码添加任何必要的命令行参数。
本指南将向您展示如何使用 deepspeed 启动器针对不同的训练设置部署 DeepSpeed。您可以查看这个 帖子，了解更实用的使用示例。
要在多个 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 以允许在没有共享文件系统的情况下加载
{
"checkpoint": {
"use_node_local_storage": true
}
}您也可以使用 Trainer 的 --save_on_each_node 参数自动将上面的 checkpoint 添加到您的配置中。
对于 torchrun，您必须 ssh 到每个节点并在两个节点上运行以下命令。启动器会等待这两个节点同步，然后启动训练。
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 # name
#SBATCH --nodes=2 # nodes
#SBATCH --ntasks-per-node=1 # crucial - only 1 task per dist per node!
#SBATCH --cpus-per-task=10 # number of cores per tasks
#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。如果您想使用超过 1 个 GPU，则必须使用多进程环境才能使 DeepSpeed 正常工作。这意味着您必须使用 deepspeed 启动器，而无法像这里显示的那样进行模拟。
# DeepSpeed requires a distributed environment even when only one process is used.
# This emulates a launcher in the notebook
import os
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "9994" # modify if RuntimeError: Address already in use
os.environ["RANK"] = "0"
os.environ["LOCAL_RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
# Now proceed as normal, plus pass the DeepSpeed config file
training_args = TrainingArguments(..., deepspeed="ds_config_zero3.json")
trainer = Trainer(...)
trainer.train()如果您想在笔记本的当前目录中动态创建配置文件，您可以使用一个专门的单元格。
%%bash
cat <<'EOT' > ds_config_zero3.json
{
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16,
"hysteresis": 2,
"min_loss_scale": 1
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
EOT如果训练脚本位于文件中而不是笔记本单元格中，您可以从笔记本单元格的 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 权重存储在自定义检查点优化器文件中（glob 模式看起来像 global_step*/*optim_states.pt），并保存在正常检查点下。
使用 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_optimization": {
"stage3_gather_16bit_weights_on_model_save": true
}
}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
非训练器 DeepSpeed 集成
DeepSpeed 也可以与 Transformers 一起使用，而无需 Trainer 类。这由 HfDeepSpeedConfig 处理，它只负责收集 ZeRO-3 参数，并在调用 from_pretrained() 时将模型拆分到多个 GPU 上。
如果您希望所有事情都自动为您处理，请尝试将 DeepSpeed 与 Trainer 一起使用！您需要遵循 DeepSpeed 文档，并手动在配置文件中配置参数值（您不能使用 "auto" 值）。
为了有效地部署 ZeRO-3，您必须在模型之前实例化 HfDeepSpeedConfig 对象，并保持该对象存活。
from transformers.integrations import HfDeepSpeedConfig
from transformers import AutoModel
import deepspeed
ds_config = {...} # deepspeed config object or path to the file
# must run before instantiating the model to detect zero 3
dschf = HfDeepSpeedConfig(ds_config) # keep this object alive
model = AutoModel.from_pretrained("openai-community/gpt2")
engine = deepspeed.initialize(model=model, config_params=ds_config, ...)非训练器 ZeRO 推理
要在无法将模型容纳到单个 GPU 的情况下，在没有 Trainer 的情况下运行 ZeRO 推理，请尝试使用额外的 GPU 或/和将模型卸载到 CPU 内存中。这里需要理解的一个重要细微差别是，ZeRO 的设计方式允许您在不同的 GPU 上并行处理不同的输入。
请确保
- 如果 GPU 内存足够，则禁用 CPU 卸载（因为它会降低速度）。
- 如果您的 GPU 是 Ampere 或更新的 GPU，则启用 bf16 以提高速度。如果您没有这些 GPU 之一，您可以启用 fp16，但前提是您不使用以 bf16（T5 模型）预训练的模型，因为这会导致溢出错误。
查看以下脚本，更好地了解如何在无法容纳到单个 GPU 上的模型上，在没有 Trainer 的情况下运行 ZeRO 推理。
#!/usr/bin/env python
# This script demonstrates how to use Deepspeed ZeRO in an inference mode when one can't fit a model
# into a single GPU
#
# 1. Use 1 GPU with CPU offload
# 2. Or use multiple GPUs instead
#
# First you need to install deepspeed: pip install deepspeed
#
# Here we use a 3B "bigscience/T0_3B" model which needs about 15GB GPU RAM - so 1 largish or 2
# small GPUs can handle it. or 1 small GPU and a lot of CPU memory.
#
# To use a larger model like "bigscience/T0" which needs about 50GB, unless you have an 80GB GPU -
# you will need 2-4 gpus. And then you can adapt the script to handle more gpus if you want to
# process multiple inputs at once.
#
# The provided deepspeed config also activates CPU memory offloading, so chances are that if you
# have a lot of available CPU memory and you don't mind a slowdown you should be able to load a
# model that doesn't normally fit into a single GPU. If you have enough GPU memory the program will
# run faster if you don't want offload to CPU - so disable that section then.
#
# To deploy on 1 gpu:
#
# deepspeed --num_gpus 1 t0.py
# or:
# python -m torch.distributed.run --nproc_per_node=1 t0.py
#
# To deploy on 2 gpus:
#
# deepspeed --num_gpus 2 t0.py
# or:
# 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" # To avoid warnings about parallelism in tokenizers
# distributed setup
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
# batch size has to be divisible by world_size, but can be bigger than world_size
train_batch_size = 1 * world_size
# ds_config notes
#
# - enable bf16 if you use Ampere or higher GPU - this will run in mixed precision and will be
# faster.
#
# - for older GPUs you can enable fp16, but it'll only work for non-bf16 pretrained models - e.g.
# all official t5 models are bf16-pretrained
#
# - set offload_param.device to "none" or completely remove the `offload_param` section if you don't
# - want CPU offload
#
# - if using `offload_param` you can manually finetune stage3_param_persistence_threshold to control
# - which params should remain on gpus - the larger the value the smaller the offload size
#
# For in-depth info on Deepspeed config see
# https://huggingface.co/docs/transformers/main/main_classes/deepspeed
# keeping the same format as json for consistency, except it uses lower case for 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
# next line instructs transformers to partition the model directly over multiple gpus using
# deepspeed.zero.Init when model's `from_pretrained` method is called.
#
# **it has to be run before loading the model AutoModelForSeq2SeqLM.from_pretrained(model_name)**
#
# otherwise the model will first be loaded normally and only partitioned at forward time which is
# less efficient and when there is little CPU RAM may fail
dschf = HfDeepSpeedConfig(ds_config) # keep this object alive
# now a model can be loaded.
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# initialise Deepspeed ZeRO and store only the engine object
ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
ds_engine.module.eval() # inference
# Deepspeed ZeRO can process unrelated inputs on each GPU. So for 2 gpus you process 2 inputs at once.
# If you use more GPUs adjust for more.
# And of course if you have just one input to process you then need to pass the same string to both gpus
# If you use only one GPU, then you will have only 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这是一个非常基本的示例，您需要根据自己的用例进行调整。
生成
使用 ZeRO-3 和多个 GPU 进行生成需要通过在 generate() 方法中设置 synced_gpus=True 来同步 GPU。否则，如果一个 GPU 在另一个 GPU 完成生成之前就完成了生成，整个系统就会挂起，因为剩余的 GPU 还没有收到从第一个完成的 GPU 发送的权重分片。
对于 Transformers>=4.28，如果在生成过程中检测到多个 GPU，则 synced_gpus 将自动设置为 True。
疑难解答
当您遇到问题时，您应该考虑 DeepSpeed 是否是问题的原因，因为通常它不是（除非它是超级明显的，并且您可以在异常中看到 DeepSpeed 模块）！第一步应该是尝试在没有 DeepSpeed 的情况下重新运行您的设置，如果问题仍然存在，那么您可以报告问题。如果问题是 DeepSpeed 的核心问题，与 Transformers 集成无关，请在 DeepSpeed 存储库 上打开一个 Issue。
对于与 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 损失
当模型以 bf16 方式预训练，然后您尝试使用 fp16（尤其适用于 TPU 训练的模型）时，经常会出现 NaN 损失。为了解决这个问题，如果您硬件支持（TPU、Ampere GPU 或更新的 GPU），请使用 fp32 或 bf16。
另一个问题可能与使用 fp16 相关。例如，如果您的 fp16 配置如下
{
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16,
"hysteresis": 2,
"min_loss_scale": 1
}
}您可能会在日志中看到以下 OVERFLOW! 消息
0%| | 0/189 [00:00<?, ?it/s] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 262144, reducing to 262144 1%|▌ | 1/189 [00:00<01:26, 2.17it/s] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 262144, reducing to 131072.0 1%|█▏ [...] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 1, reducing to 1 14%|████████████████▌ | 27/189 [00:14<01:13, 2.21it/s] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 1, reducing to 1 15%|█████████████████▏ | 28/189 [00:14<01:13, 2.18it/s] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 1, reducing to 1 15%|█████████████████▊ | 29/189 [00:15<01:13, 2.18it/s] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 1, reducing to 1 [...]
这意味着 DeepSpeed 损失缩放器无法找到缩放系数来克服损失溢出。要解决此问题，请尝试使用更高的 initial_scale_power 值（32 通常有效）。
资源
DeepSpeed ZeRO 是一项强大的技术，可以用于训练和加载非常大的模型，以便在 GPU 资源有限的情况下进行推理，使其更容易被所有人使用。要了解有关 DeepSpeed 的更多信息，请随时阅读 博客文章、文档 和 GitHub 存储库。
以下论文也是学习有关 ZeRO 的更多信息的宝贵资源
< > 更新 在 GitHub 上