DeepSpeedを使用して大規模モデルトレーニングを高速化する
Using DeepSpeed to accelerate large-scale model training.
この投稿では、Accelerate ライブラリを活用して、ユーザーが DeeSpeed の ZeRO 機能を利用して大規模なモデルをトレーニングする方法について説明します。
大規模なモデルをトレーニングしようとする際にメモリ不足 (OOM) エラーに悩まされていますか？私たちがサポートします。大規模なモデルは非常に高性能ですが、利用可能なハードウェアでトレーニングするのは困難です。大規模なモデルのトレーニングに利用可能なハードウェアの最大限の性能を引き出すために、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 に統合されており、DeepSpeed および FairScale を使用した ZeRO を活用するための優れたブログ Fit More and Train Faster With ZeRO via DeepSpeed and FairScale [4] および PyTorch Fully Sharded Data Parallel を使用した大規模モデルトレーニングの高速化 [5] があります。裏側で何が行われているかの詳しい説明はこれらのブログに譲り、主に DeepSpeed ZeRO を活用する方法に焦点を当てます。
ハードウェアのセットアップ : 2X24GB NVIDIA Titan RTX GPUs。60GB の RAM。
テキスト分類のためのエンコーダのみのモデルのファインチューニングタスクを見てみましょう。ファインチューニングには、MRPC GLUE データセット上で事前学習済みの microsoft/deberta-v2-xlarge-mnli (900M パラメータ) を使用します。
コードはこちらの run_cls_no_trainer.py で利用できます。こちらは公式のテキスト分類の例と似ていますが、トレーニングと評価時間を測定するためのロジックが追加されています。分散データ並列 (DDP) と DeepSpeed ZeRO ステージ 2 のパフォーマンスをマルチGPU環境で比較しましょう。
コードの変更なしで DeepSpeed ZeRO ステージ 2 を有効にするには、accelerate config を実行し、Accelerate DeepSpeed プラグインを活用してください。
ZeRO ステージ 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 ステージ 2 は OOM エラーなしでバッチサイズ 40 を実現します。したがって、DeepSpeed は DDP に比べて1つのGPUあたりデータを 5倍 も多くフィットさせることができます。以下は、wandb の実行からのプロットのスナップショットと、DDP と DeepSpeed のベンチマークテーブルの比較です。
表1：DeBERTa-XL（900M）モデルでのDeepSpeed ZeRO Stage-2のベンチマーキング
このより大きなバッチサイズでは、コードを変更することなく、全体のトレーニング時間を約3.5倍短縮し、パフォーマンスメトリクスに一切の低下がなくなりました。やったね！ 🤗。
さらにオプションを微調整するには、DeepSpeedの設定ファイルと最小限のコード変更が必要です。それでは、これをどのように行うか見てみましょう。
まず、独自のチャットボットをトレーニングするためのシーケンス-シーケンスモデルのファインチューニングのタスクを見てみましょう。具体的には、smangrul/MuDoConv（Multi-Domain Conversation）データセット上でfacebook/blenderbot-400M-distillをファインチューニングします。このデータセットには、パーソナリティ、特定の感情的な文脈での接地、目標指向型（例：レストラン予約）および一般的なWikipediaトピック（例：クリケット）をカバーする10の異なるデータソースからの会話が含まれています。
この例では、EngagingnessとHumannessのChatbotの効果的な測定手法は、高価な人間による評価です[6]。したがって、この例では、追跡されるメトリクスはBLEUスコアです（理想的ではありませんが、このようなタスクに対する従来のメトリクスです）。もしGPUがbfloat16精度をサポートしている場合は、コードを大きなT5モデルのトレーニングに適応させることができます。さもないと、NaNの損失値になります。DeepSpeed vs DDPについてのクイックベンチマークを10000のトレーニングサンプルと1000の評価サンプルで実行します。
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を実行し、設定ファイルのパスを入力してください。詳細については、🤗のDeepSpeed Config Fileの公式ドキュメントを参照してください。
ZeRO Stage-2 DeepSpeed Configファイルの例
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では、OOMエラーが発生せずにバッチサイズ200が可能です。したがって、DeepSpeedはDDPと比較して、GPUごとにデータを2倍多く格納できるようになります。トレーニングでは約1.44倍の高速化、評価では約1.23倍の高速化が観察されます。これはVoAGIサイズのモデルですので、非常に興奮するわけではありませんが、モデルが大きくなると改善されるでしょう。あなたは🤗 Space smangrul/Chat-Eで、全データを使用してトレーニングされたChatbotとチャットすることができます。チャットボットにパーソナリティを与え、特定の感情に対するグラウンド会話、目標指向のタスク、あるいは自由な流れの方法で使用することができます。以下は、チャットボットとの楽しい会話の例です 💬。ここで、さまざまなコンテキストを使用した他の会話のスナップショットを見つけることができます。
GPUメモリに収まらない巨大なモデルのトレーニングを可能にするためのCPU/ディスクオフロード
24GBのNVIDIA Titan RTX GPU1枚では、バッチサイズ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"\Table 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
We will continue to update VoAGI; if you have any questions or suggestions, please contact us!
Was this article helpful?
93 out of 132 found this helpful
Related articles
