使用 DeepSpeed 加速大型模型训练
使用DeepSpeed加速大型模型训练
在本文中，我们将介绍如何利用Accelerate库来训练大型模型，使用户能够利用DeeSpeed的ZeRO功能。
在尝试训练大型模型时，厌倦了内存不足（OOM）错误吗？我们为您提供解决方案。大型模型具有很高的性能[1]，但使用可用的硬件进行训练却很困难。为了充分利用可用的硬件来训练大型模型，可以利用使用ZeRO（Zero Redundancy Optimizer）的数据并行性[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中，并伴有很棒的博客文章Fit More and Train Faster With ZeRO via DeepSpeed和FairScale [4]以及使用PyTorch Fully Sharded Data Parallel加速大型模型训练[5]。我们将把幕后的解释推迟到那些博客文章中，主要关注利用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允许使用批量大小为40而不会出现OOM错误。因此，与DDP相比，DeepSpeed能够每个GPU容纳更多的数据，达到5倍。下面是wandb运行的图表快照，以及DDP与DeepSpeed的性能对比表。
表格1：在DeBERTa-XL（900M）模型上对DeepSpeed ZeRO Stage-2进行基准测试
通过使用更大的批量大小，我们观察到总训练时间加快了约3.5倍，而且性能指标没有降低，而且所有这些都不需要更改任何代码。太棒了！🤗
要能够调整更多选项，您需要使用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进行训练。有关各种配置功能的详细信息，请参阅DeepSpeed文档。
{
"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
+ )表2：在BlenderBot（400M）模型上基准测试DeepSpeed ZeRO Stage-2
在我们的单节点多GPU设置中，DDP支持的最大批量大小（无OOM错误）为100。相比之下，DeepSpeed Zero-Stage 2使批量大小增大到200，而不会遇到OOM错误。因此，与DDP相比，DeepSpeed可以使每个GPU处理的数据量增加2倍。我们观察到训练速度提升了约1.44倍，评估速度提升了约1.23倍，因为我们能够在相同的可用硬件上处理更多的数据。由于这个模型是VoAGI尺寸的，所以这个加速并不那么令人兴奋，但对于更大的模型来说会有所改善。您可以与使用全部数据训练的Chatbot进行交流，网址为🤗 Space smangrul/Chat-E。您可以为机器人设定角色、将对话引导到特定的情感、在目标导向任务中使用它，或者以自由流动的方式使用它。下面是与聊天机器人的一次有趣对话💬。您可以在这里找到使用不同上下文的更多对话的快照。
使用CPU/磁盘卸载来训练无法适应GPU内存的庞大模型
在单个24GB的NVIDIA Titan RTX GPU上，即使批量大小为1，也无法训练GPT-XL模型（15亿个参数）。我们将看看如何使用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"\表3：在GPT-XL（1.5B）模型上对DeepSpeed ZeRO Stage-3 CPU卸载进行基准测试
使用DDP即使批大小为1也会导致OOM错误。另一方面，使用DeepSpeed ZeRO Stage-3 CPU卸载，我们可以使用批大小为16进行训练。
最后，请记住，🤗 Accelerate 只集成了DeepSpeed，因此如果您对DeepSpeed的使用有任何问题或疑问，请在DeepSpeed GitHub上提出问题。
[1] Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers
[2] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
[3] DeepSpeed: Extreme-scale model training for everyone – Microsoft Research
[4] Fit More and Train Faster With ZeRO via DeepSpeed and FairScale
[5] Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel
[6] Recipes for building an open-domain chatbot
