Transformers 文档

DeepSpeed

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-bit Adam,这些功能在PyPI发行版中不可用。

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

这意味着您要么需要一个没有CPU卸载的80GB GPU,要么需要一个8GB GPU和一个约60GB的CPU来进行卸载(这些只是参数、优化器状态和梯度的内存需求,您还需要一些额外的内存用于CUDA内核和激活)。您还应该考虑成本和速度之间的权衡,因为租用或购买较小的GPU会更便宜,但训练模型的时间会更长。

如果你有足够的GPU内存,请确保禁用CPU/NVMe卸载以使一切运行得更快。

选择一个ZeRO阶段

在您安装了DeepSpeed并更好地了解了您的内存需求之后,下一步是选择一个ZeRO阶段来使用。按照最快和最高内存效率的顺序:

最快 内存效率最高
ZeRO-1 ZeRO-3 + 卸载
ZeRO-2 ZeRO-3
ZeRO-2 + offload ZeRO-2 + offload
ZeRO-3 ZeRO-2
ZeRO-3 + 卸载 ZeRO-1

为了找到最适合您的方法,首先从最快的方法开始,如果内存不足,尝试下一个阶段,虽然速度较慢但内存效率更高。您可以自由选择您偏好的方向(从内存效率最高或速度最快开始)来发现速度和内存使用之间的适当平衡。

你可以使用的一般流程是(从批量大小为1开始):

  1. 启用梯度检查点
  2. 尝试 ZeRO-2
  3. 尝试使用ZeRO-2并卸载优化器
  4. 尝试 ZeRO-3
  5. 尝试使用ZeRO-3并将参数卸载到CPU
  6. 尝试使用ZeRO-3并将参数和优化器卸载到CPU
  7. 如果你正在使用generate()方法,尝试降低各种默认值,比如更窄的搜索束
  8. 尝试混合半精度(在较旧的GPU架构上使用fp16,在Ampere上使用bf16)而不是全精度权重
  9. 如果可能,增加更多硬件或启用Infinity以将参数和优化器卸载到NVMe
  10. 一旦你不再遇到内存不足的问题,测量有效吞吐量,然后尝试尽可能增加批量大小,以最大化GPU效率
  11. 最后,尝试通过禁用一些卸载功能或使用更快的ZeRO阶段以及增加/减少批量大小来优化您的训练设置,以找到速度和内存使用之间的最佳平衡

DeepSpeed 配置文件

DeepSpeed 通过一个包含所有参数的配置文件与 Trainer 类一起工作,该配置文件用于配置您希望设置的训练运行。当您执行训练脚本时,DeepSpeed 会将从 Trainer 接收到的配置记录到控制台,以便您可以确切地看到使用了哪些配置。

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

path to file
nested dict
TrainingArguments(..., deepspeed="path/to/deepspeed_config.json")

DeepSpeed 和 Trainer 参数

有三种类型的配置参数:

  1. 一些配置参数由Trainer和DeepSpeed共享,当存在冲突定义时,可能难以识别错误。为了简化这一过程,这些共享的配置参数是从Trainer命令行参数配置的。

  2. 一些配置参数是从模型配置中自动派生的,因此您不需要手动调整这些值。Trainer 使用配置值 auto 来确定设置最正确或最有效的值。您可以显式设置自己的配置参数,但必须注意确保 Trainer 参数和 DeepSpeed 配置参数一致。不匹配可能会导致训练失败,且难以检测!

  3. 一些仅适用于DeepSpeed的配置参数,需要根据您的训练需求手动设置。

你也可以修改DeepSpeed配置并从中编辑TrainingArguments

  1. 创建或加载一个DeepSpeed配置以用作主配置
  2. 基于这些DeepSpeed配置值创建一个TrainingArguments对象

一些值,例如 scheduler.params.total_num_steps 是由 Trainer 在训练期间计算的。

ZeRO配置

有三种配置,每种配置对应不同的ZeRO阶段。阶段1在可扩展性方面不太有趣,本指南主要关注阶段2和阶段3。zero_optimization配置包含了所有要启用和如何配置它们的选项。有关每个参数的更详细解释,请查看DeepSpeed Configuration JSON参考。

DeepSpeed doesn’t validate parameter names and any typos fallback on the parameter's default setting. You can watch the DeepSpeed engine startup log messages to see what values it is going to use.

以下配置必须与DeepSpeed一起设置,因为Trainer没有提供等效的命令行参数。

ZeRO-1
ZeRO-2
ZeRO-3

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内存不足的情况,或者增加批量大小以提高性能。要启用此功能:

  1. 对于Hugging Face模型,在Trainer中设置model.gradient_checkpointing_enable()--gradient_checkpointing
  2. 对于非Hugging Face模型,使用DeepSpeed的Activation Checkpointing API。你也可以替换Transformers的建模代码,并用DeepSpeed API替换torch.utils.checkpoint。这种方法更加灵活,因为你可以将前向激活卸载到CPU内存中,而不是重新计算它们。

优化器和调度器

只要不启用offload_optimizer,DeepSpeed和Transformers的优化器和调度器可以混合使用。当启用offload_optimizer时,只要优化器具有CPU和GPU实现(除了LAMB),就可以使用非DeepSpeed优化器。

配置文件中优化器和调度器的参数可以从命令行设置,以避免难以发现的错误。例如,如果在其他地方设置了不同的学习率,你可以从命令行覆盖它。除了优化器和调度器参数外,你还需要确保你的Trainer命令行参数与DeepSpeed配置匹配。

optimizer
scheduler

DeepSpeed 提供了几种优化器(Adam、AdamW、OneBitAdam 和 LAMB),但你也可以从 PyTorch 导入其他优化器。如果你没有在配置中配置优化器,Trainer 会自动选择 AdamW,并使用命令行中提供的值或以下参数的默认值:lradam_beta1adam_beta2adam_epsilonweight_decay

你可以将参数设置为"auto"或手动输入你所需的值。

{
   "optimizer": {
       "type": "AdamW",
       "params": {
         "lr": "auto",
         "betas": "auto",
         "eps": "auto",
         "weight_decay": "auto"
       }
   }
}

你也可以通过将以下内容添加到顶层配置中来使用不受支持的优化器。

{
   "zero_allow_untested_optimizer": true
}

从DeepSpeed==0.8.3开始,如果你想使用offload,你还需要在顶层配置中添加以下内容,因为offload与DeepSpeed的CPU Adam优化器配合使用效果最佳。

{
   "zero_force_ds_cpu_optimizer": false
}

精度

Deepspeed 支持 fp32、fp16 和 bf16 混合精度。

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 可以通过不同的启动器进行部署,例如 torchrundeepspeed 启动器或 Accelerate。要部署,请将 --deepspeed ds_config.json 添加到 Trainer 命令行中。建议使用 DeepSpeed 的 add_config_arguments 工具将任何必要的命令行参数添加到您的代码中。

本指南将向您展示如何为不同的训练设置使用deepspeed启动器部署DeepSpeed。您可以查看此帖子以获取更多实际使用示例。

multi-GPU
single-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,以允许在没有访问共享文件系统的情况下进行加载:

{
  "checkpoint": {
    "use_node_local_storage": true
  }
}

你也可以使用Trainer--save_on_each_node参数来自动将上述checkpoint添加到你的配置中。

torchrun
deepspeed

对于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 权重存储在自定义的检查点优化器文件中(全局模式看起来像 global_step*/*optim_states.pt),并保存在常规检查点下。

fp16
fp32

使用ZeRO-2训练的模型以fp16格式保存pytorch_model.bin权重。对于使用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 Inference 将模型权重放置在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 也可以在没有 Trainer 类的情况下与 Transformers 一起使用。这是通过 HfDeepSpeedConfig 处理的,它只负责在调用 from_pretrained() 时收集 ZeRO-3 参数并将模型分布在多个 GPU 上。

如果你希望一切都自动为你处理,尝试使用DeepSpeed与Trainer!你需要遵循DeepSpeed文档,并手动在配置文件中配置参数值(你不能使用"auto"值)。

要高效部署 ZeRO-3,您必须在模型之前实例化 HfDeepSpeedConfig 对象,并保持该对象存活:

pretrained model
non-pretrained model
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, ...)

非训练器零推理

在无法将模型放入单个GPU的情况下,尝试使用额外的GPU或/和卸载到CPU内存来运行ZeRO推理,而无需使用Trainer。这里需要理解的重要细微差别是,ZeRO的设计方式允许你在不同的GPU上并行处理不同的输入。

确保:

  • 如果你有足够的GPU内存,请禁用CPU卸载(因为它会减慢速度)。
  • 如果您有Ampere或更新的GPU,可以启用bf16以加快速度。如果您没有这些GPU,只要您不使用在bf16中预训练的模型(如T5模型),可以启用fp16,因为这可能会导致溢出错误。

看看下面的脚本,以更好地了解如何在不使用Trainer的情况下,在无法适应单个GPU的模型上运行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生成时,需要通过设置synced_gpus=True来同步GPU,该设置在generate()方法中。否则,如果一个GPU在另一个GPU之前完成生成,整个系统将挂起,因为剩余的GPU尚未接收到首先完成的GPU的权重分片。

对于Transformers>=4.28,如果在生成过程中检测到多个GPU,synced_gpus会自动设置为True

故障排除

当你遇到问题时,你应该考虑是否是DeepSpeed导致的问题,因为通常不是(除非问题非常明显并且你可以在异常中看到DeepSpeed模块)!第一步应该是尝试在没有DeepSpeed的情况下重新设置,如果问题仍然存在,那么你可以报告这个问题。如果问题是DeepSpeed的核心问题并且与Transformers集成无关,请在DeepSpeed仓库上提交问题。

对于与Transformers集成相关的问题,请提供以下信息:

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_optimizeroffload_param或两者都配置为卸载到CPU。

如果您有NVMe和ZeRO-3设置,可以尝试将数据卸载到NVMe(估计您的模型的内存需求)。

NaN 损失

NaN损失通常发生在模型以bf16预训练后,然后尝试使用fp16时(特别是与TPU训练的模型相关)。要解决此问题,如果您的硬件支持(TPU、Ampere 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的绝佳资源:

< > Update on GitHub