DeepSpeed ZeRO

DeepSpeed の ZeRO (Zero Redundancy Optimizer) は、オプティマイザの状態、勾配、パラメータを GPU 間でシャード化することで、分散学習におけるメモリ冗長性を排除します。ZeRO は3つのステージがあり、各ステージでより多くの状態をシャード化します。DeepSpeed は、さらなるメモリ削減のために CPU または NVMe メモリへのオフロードもサポートしています。ステージやオフロードレベルを1つ増やすごとにピークメモリは減少しますが、その代わり GPU 間の通信量が増加します。

                  params        grads       opt states
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-1          │██████████│ │██████████│ │███░░░░░░░│  GPU 0
                │██████████│ │██████████│ │░░░███░░░░│  GPU 1
                │██████████│ │██████████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-2          │██████████│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │██████████│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │██████████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-3          │███░░░░░░░│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │░░░███░░░░│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │░░░░░░████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
  █ resident    ░ held on another GPU

ZeRO-2 は、ZeRO-3 よりも通信オーバーヘッドが低い方法で勾配とオプティマイザの状態をシャードします。ZeRO-2 ではモデルが GPU にまたがって収まらない場合のみ、ZeRO-3 を使用してください。

インストール

DeepSpeed は PyPI からインストールするか、deepspeed エクストラ付きで Transformers をインストールしてください。

pip install deepspeed
pip install transformers[deepspeed]

CUDA関連のインストールエラーが発生した場合は、DeepSpeed CUDAのドキュメントを確認してください。ソースからのインストールの方が、お使いのハードウェアに正確に適合し、PyPIリリースではまだ利用できない機能も含まれるため、より確実な選択肢です。

設定

Trainer は deepspeed 引数を通じて DeepSpeed を統合しており、この引数には JSON 設定ファイルを指定できます。代わりに、TrainingArguments の代わりに Accelerate 設定ファイルを使用することもできます。

TrainingArguments

Accelerate設定ファイル

DeepSpeed に TrainingArguments から値を埋めてほしい項目には、設定で "auto" を使用します。値を明示的に指定する場合は、DeepSpeed の引数と TrainingArguments の両方で同じ値を使用してください。

DeepSpeed の設定オプションの完全なリストについては、DeepSpeed Configuration JSON リファレンスを参照してください。

{
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto" }
    },
    "fp16": { "enabled": "auto" }
}

DeepSpeed 設定	TrainingArguments
`train_micro_batch_size_per_gpu`	`per_device_train_batch_size`
`gradient_accumulation_steps`	`gradient_accumulation_steps`
`optimizer.params.lr`	`learning_rate`
`fp16.enabled`	`fp16`

設定をdeepspeed引数に渡します。

from transformers import TrainingArguments

args = TrainingArguments(
    deepspeed="path/to/deepspeed_config.json",
    ...
)

# DeepSpeed launcher
deepspeed --num_gpus 4 train.py

# torchrun
torchrun --nproc_per_node 4 train.py

# Accelerate
accelerate launch --num_processes 4 train.py

ZeROステージ

開始点として使用する ZeRO ステージ設定を選択します。

ZeRO-1

ZeRO-2

ZeRO-3

以下のフィールドは、トレーニングのカスタマイズにおいて重要です。

zero_optimization は ZeRO のステージを設定します。
```
{ "zero_optimization": { "stage": 3 } }
```
バッチサイズと勾配蓄積の引数を"auto"に設定してください。これらをTrainingArgumentsと不一致する値に手動で設定した場合、誤った値のままトレーニングが静かに続行されます。
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16 は学習の精度を設定します。"auto" に設定すると、TrainingArguments の bf16 フラグに追従します。
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save は、保存前にすべての GPU 間で all-gather を実行し、シャードから完全なテンソルを再構築します。これは ZeRO-3 の引数です。
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
overlap_comm を true に設定すると、バックワードパスの処理中に all-reduce のレイテンシを隠すことができます。allgather_bucket_size と reduce_bucket_size は、通信速度と GPU メモリ使用量のトレードオフを調整します。値を下げるとメモリ使用量は減りますが、通信速度は低下します。
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer は、オプティマイザーを CPU メモリにオフロードします。さらにメモリを節約するには、offload_param を使用してモデルのパラメータもオフロードします（ZeRO-3 のみ）。CPU と GPU 間の転送を高速化するには pin_memory を true に設定しますが、これにより他のプロセスで使用できない RAM がロックされます。
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```
optimizer と scheduler のデフォルト値は、TrainingArguments で設定されたオプティマイザーとスケジューラーです。LAMB などの DeepSpeed ネイティブのオプティマイザーが必要でない限り、TrainingArguments から値を読み込むように "auto" に設定してください。
```
{
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": { "total_num_steps": "auto", "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" }
    }
}
```
オプティマイザーをオフロードする場合、DeepSpeed の CPU Adam オプティマイザーを使用するには、zero_force_ds_cpu_optimizer を false に設定してください。
```
{
    "zero_force_ds_cpu_optimizer": false
}
```

チェックポイント

DeepSpeed は、from_pretrained() で直接読み込めないシャーディング形式でチェックポイントを保存します。学習の終了時に最良のチェックポイントをトラッキングして再読み込みさせるには、load_best_model_at_end() を True に設定してください。

from transformers import TrainingArguments, Trainer

args = TrainingArguments(
    deepspeed="ds_config_zero3.json",
    load_best_model_at_end=True,
    ...
)
# after training, save a normal transformers checkpoint
trainer.save_model("./best-model")

save_only_model=True を設定すると、オプティマイザーの完全な状態の保存がスキップされるため、トレーニング終了時に最良のモデルを再読み込みできません。また、シャードから完全な重みを再構築するには、stage3_gather_16bit_weights_on_model_save: true も設定してください。これは、ZeRO-3 で統合された 16 ビットのモデル成果物や 16 ビットの state dict を保存するために必要です。save_only_model=True と load_best_model_at_end=True を組み合わせた場合、Transformers はエラーを発生させます。

異なる並列構成間で再開を行う場合は、DeepSpeedのUniversal Checkpointingガイドを参照してください。

次のステップ

ZeRO の仕組みについて詳しく知るには、The Ultra-Scale Playbook の Zero Redundancy Optimizer 章をお読みください。
ZeROの論文をお読みください:Memory Optimizations Toward Training Trillion Parameter Models、Democratizing Billion-Scale Model Training、Breaking the GPU Memory Wall for Extreme Scale Deep Learning。

Update on GitHub