使用 DeepSpeed 加速大型模型训练

发布于 2022 年 6 月 28 日

在 GitHub 上更新

在这篇文章中，我们将探讨如何利用 Accelerate 库来训练大型模型，该库使用户能够利用 DeeSpeed 的 ZeRO 功能。

动机 🤗

在尝试训练大型模型时，是否厌倦了内存不足（OOM）错误？我们已为您解决。大型模型性能非常出色 [1]，但难以用现有硬件进行训练。为了最大限度地利用现有硬件训练大型模型，可以使用 ZeRO - 零冗余优化器 [2] 进行数据并行。.

下面是使用 ZeRO 进行数据并行的简要描述，并附有来自这篇博客文章的图表

（来源：链接）

a. 阶段 1：在数据并行工作器/GPU 之间分片优化器状态

b. 阶段 2：在数据并行工作器/GPU 之间分片优化器状态 + 梯度

c. 阶段 3：在数据并行工作器/GPU 之间分片优化器状态 + 梯度 + 模型参数

d. 优化器卸载：在 ZERO 阶段 2 的基础上，将梯度 + 优化器状态卸载到 CPU/磁盘

e. 参数卸载：在 ZERO 阶段 3 的基础上，将模型参数卸载到 CPU/磁盘

在这篇博客文章中，我们将探讨如何使用 Accelerate 来利用 ZeRO 进行数据并行。DeepSpeed、FairScale 和 PyTorch FullyShardedDataParallel (FSDP) 已经实现了 ZERO 论文的核心思想。这些功能已经集成到 🤗 transformers Trainer 和 🤗 accelerate 中，并附有精彩的博客文章通过 DeepSpeed 和 FairScale 使用 ZeRO 适应更多并训练更快 [4] 和使用 PyTorch Fully Sharded Data Parallel 加速大型模型训练 [5]。我们将其背后的解释推迟到这些博客文章，主要关注使用 Accelerate 利用 DeepSpeed ZeRO。

Accelerate 🚀：无需任何代码更改即可利用 DeepSpeed ZeRO

硬件设置：2X24GB NVIDIA Titan RTX GPU。60GB RAM。

我们将研究用于文本分类的仅编码器模型微调任务。我们将使用预训练的 microsoft/deberta-v2-xlarge-mnli（900M 参数）在 MRPC GLUE 数据集上进行微调。

代码可在此处找到 run_cls_no_trainer.py。它与此处的官方文本分类示例类似，但增加了测量训练和评估时间的逻辑。让我们比较多 GPU 设置中分布式数据并行（DDP）和 DeepSpeed ZeRO Stage-2 的性能。

要在不更改任何代码的情况下启用 DeepSpeed ZeRO Stage-2，请运行 accelerate config 并利用 Accelerate DeepSpeed 插件。

ZeRO Stage-2 DeepSpeed 插件示例

compute_environment: LOCAL_MACHINE
deepspeed_config:
 gradient_accumulation_steps: 1
 gradient_clipping: 1.0
 offload_optimizer_device: none
 offload_param_device: none
 zero3_init_flag: false
 zero_stage: 2
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

现在，运行以下命令进行训练

accelerate launch run_cls_no_trainer.py \
  --model_name_or_path "microsoft/deberta-v2-xlarge-mnli" \
  --task_name "mrpc" \
  --ignore_mismatched_sizes \
  --max_length 128 \
  --per_device_train_batch_size 40 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir "/tmp/mrpc/deepspeed_stage2/" \
  --with_tracking \
  --report_to "wandb" \

在我们的单节点多 GPU 设置中，DDP 在不出现 OOM 错误的情况下支持的最大批量大小为 8。相比之下，DeepSpeed Zero-Stage 2 可以在不出现 OOM 错误的情况下支持批量大小为 40。因此，与 DDP 相比，DeepSpeed 使得每个 GPU 可以容纳 5 倍 的数据。以下是 wandb 运行的图表快照以及比较 DDP 与 DeepSpeed 的基准测试表。

方法	最大批量大小	每个 epoch 的训练时间（秒）	每个 epoch 的评估时间（秒）	F1 分数	准确率
DDP（分布式数据并行）	8	103.57	2.04	0.931	0.904
DeepSpeed ZeRO Stage 2	40	28.98	1.79	0.936	0.912

表 1：DeBERTa-XL (900M) 模型上的 DeepSpeed ZeRO Stage-2 基准测试

在不改变任何代码的情况下，使用更大的批量大小，我们观察到总训练时间加速了约 3.5 倍，而性能指标没有任何下降。太棒了！🤗。

为了调整更多选项，您将需要使用 DeepSpeed 配置文件和最少的代码更改。让我们看看如何做到这一点。

Accelerate 🚀：利用 DeepSpeed 配置文件调整更多选项

首先，我们将研究微调序列到序列模型以训练我们自己的聊天机器人的任务。具体来说，我们将在 smangrul/MuDoConv（多领域对话）数据集上微调 facebook/blenderbot-400M-distill。该数据集包含来自 10 个不同数据源的对话，涵盖了人物、特定情感背景下的基础、目标导向（例如，餐厅预订）和一般维基百科主题（例如，板球）。

代码可在此处找到 run_seq2seq_no_trainer.py。目前有效衡量聊天机器人“参与度”和“人性化”的方法是通过人工评估，但这成本高昂 [6]。因此，对于此示例，跟踪的指标是 BLEU 分数（尽管不理想，但它是此类任务的常规指标）。如果您可以使用支持 bfloat16 精度的 GPU，则可以调整代码以训练更大的 T5 模型，否则您将遇到 NaN 损失值。我们将对 10000 个训练样本和 1000 个评估样本进行快速基准测试，因为我们关注的是 DeepSpeed 与 DDP 的比较。

我们将利用 DeepSpeed Zero Stage-2 配置 zero2_config_accelerate.json（如下所示）进行训练。有关各种配置功能的详细信息，请参阅 DeeSpeed 文档。

{
    "fp16": {
        "enabled": "true",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 15,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

要使用上述配置启用 DeepSpeed ZeRO Stage-2，请运行 accelerate config 并在要求时提供配置文件路径。有关更多详细信息，请参阅 🤗 accelerate 官方文档中的 DeepSpeed 配置文件。

ZeRO Stage-2 DeepSpeed 配置文件示例

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /path/to/zero2_config_accelerate.json
 zero3_init_flag: false
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

现在，运行以下命令进行训练

accelerate launch run_seq2seq_no_trainer.py \
    --dataset_name "smangrul/MuDoConv" \
    --max_source_length 128 \
    --source_prefix "chatbot: " \
    --max_target_length 64 \
    --val_max_target_length 64 \
    --val_min_target_length 20 \
    --n_val_batch_generations 5 \
    --n_train 10000 \
    --n_val 1000 \
    --pad_to_max_length \
    --num_beams 10 \
    --model_name_or_path "facebook/blenderbot-400M-distill" \
    --per_device_train_batch_size 200 \
    --per_device_eval_batch_size 100 \
    --learning_rate 1e-6 \
    --weight_decay 0.0 \
    --num_train_epochs 1 \
    --gradient_accumulation_steps 1 \
    --num_warmup_steps 100 \
    --output_dir "/tmp/deepspeed_zero_stage2_accelerate_test" \
    --seed 25 \
    --logging_steps 100 \
    --with_tracking \
    --report_to "wandb" \
    --report_name "blenderbot_400M_finetuning"

当使用 DeepSpeed 配置时，如果用户在配置中指定了 optimizer 和 scheduler，则用户必须使用 accelerate.utils.DummyOptim 和 accelerate.utils.DummyScheduler。这些是用户必须进行的唯一微小更改。下面我们展示了使用 DeepSpeed 配置时所需的最少更改示例

- optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=args.learning_rate)
+ optimizer = accelerate.utils.DummyOptim(optimizer_grouped_parameters, lr=args.learning_rate)

- lr_scheduler = get_scheduler(
-     name=args.lr_scheduler_type,
-     optimizer=optimizer,
-     num_warmup_steps=args.num_warmup_steps,
-     num_training_steps=args.max_train_steps,
- )

+ lr_scheduler = accelerate.utils.DummyScheduler(
+     optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
+ )

方法	最大批量大小	最大评估大小	每个 epoch 的训练时间（秒）	每个 epoch 的评估时间（秒）
DDP（分布式数据并行）	100	50	27.36	48.41
DeepSpeed ZeRO Stage 2	200	100	19.06	39.27

表 2：在 BlenderBot (400M) 模型上对 DeepSpeed ZeRO Stage-2 进行基准测试

在我们的单节点多 GPU 设置中，DDP 在不出现 OOM 错误的情况下支持的最大批量大小为 100。相比之下，DeepSpeed Zero-Stage 2 可以在不出现 OOM 错误的情况下支持批量大小为 200。因此，与 DDP 相比，DeepSpeed 使得每个 GPU 可以容纳 2 倍 的数据。我们观察到训练加速约 1.44 倍，评估加速约 1.23 倍，因为我们能够在相同的可用硬件上容纳更多数据。由于该模型属于中等大小，因此加速效果并非特别令人兴奋，但对于更大的模型，这将有所改善。您可以通过 🤗 Space smangrul/Chat-E 与使用整个数据训练的聊天机器人进行交流。您可以为机器人设定一个角色，将对话置于特定的情感背景下，用于目标导向的任务或以自由流式方式进行。下面是与聊天机器人进行的一次有趣的对话 💬。您可以在此处找到使用不同上下文进行更多对话的快照。

CPU/磁盘卸载以启用无法适应 GPU 内存的巨型模型训练

在单个 24GB NVIDIA Titan RTX GPU 上，即使批量大小为 1，也无法训练 GPT-XL 模型（1.5B 参数）。我们将探讨如何使用 DeepSpeed ZeRO Stage-3 和优化器状态、梯度和参数的 CPU 卸载来训练 GPT-XL 模型。

我们将利用 DeepSpeed Zero Stage-3 CPU 卸载配置 zero3_offload_config_accelerate.json（如下所示）进行训练。使用 🤗 accelerate 配置的其余过程与上述实验类似。

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_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,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

ZeRO Stage-3 CPU 卸载 DeepSpeed 配置文件示例

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /path/to/zero3_offload_config_accelerate.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

现在，运行以下命令进行训练

accelerate launch run_clm_no_trainer.py \
--config_name "gpt2-xl" \
--tokenizer_name "gpt2-xl" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "/tmp/clm_deepspeed_stage3_offload__accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 16 \
--per_device_eval_batch_size 1 \
--num_train_epochs 1 \
--with_tracking \
--report_to "wandb"\