# DeepSpeed ZeRO

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

```text
                  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 をインストールしてください。

```shell
pip install deepspeed
pip install transformers[deepspeed]
```

CUDA関連のインストールエラーが発生した場合は、[DeepSpeed CUDA](./debugging#deepspeed-cuda)のドキュメントを確認してください。[ソースからのインストール](https://www.deepspeed.ai/tutorials/advanced-install/#install-deepspeed-from-source)の方が、お使いのハードウェアに正確に適合し、PyPIリリースではまだ利用できない機能も含まれるため、より確実な選択肢です。

## 設定

[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer) は [deepspeed](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments.deepspeed) 引数を通じて DeepSpeed を統合しており、この引数には JSON 設定ファイルを指定できます。代わりに、[TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) の代わりに [Accelerate 設定ファイル](./accelerate#accelerate-config-file) を使用することもできます。

DeepSpeed に [TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) から値を埋めてほしい項目には、設定で `"auto"` を使用します。値を明示的に指定する場合は、DeepSpeed の引数と [TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) の両方で *同じ* 値を使用してください。

> [!NOTE]
> DeepSpeed の設定オプションの完全なリストについては、[DeepSpeed Configuration JSON](https://www.deepspeed.ai/docs/config-json/) リファレンスを参照してください。

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

| DeepSpeed 設定 | [TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.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`引数に渡します。

```py
from transformers import TrainingArguments

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

```cli
# DeepSpeed launcher
deepspeed --num_gpus 4 train.py

# torchrun
torchrun --nproc_per_node 4 train.py

# Accelerate
accelerate launch --num_processes 4 train.py
```

> [!NOTE]
> Accelerate は [TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) の `deepspeed` 引数を無視します。

[accelerate config](https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-config) コマンドを実行し、ハードウェアや学習環境に関する質問に答えることで、キャッシュに `default_config.yaml` ファイルを作成します。

```yaml
distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 4
```

[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer)ベースのスクリプトで[accelerate launch](https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-launch)を実行します。

```shell
accelerate launch --config_file deepspeed_config.yaml train.py
```

## ZeROステージ

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

```json
{
    "bf16": { "enabled": "auto" },
    "zero_optimization": { "stage": 1 },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}
```

```json
{
    "bf16": { "enabled": "auto" },
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}
```

> [!WARNING]
> ZeRO-3 は初期化時にパラメータをシャードします。モデルを読み込む前に [TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) をインスタンス化する必要があります。DeepSpeed が設定される前にモデルがすでに各 GPU に配置されている場合、メモリは節約されません。

```json
{
    "bf16": { "enabled": "auto" },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_gather_16bit_weights_on_model_save": true,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}
```

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

- `zero_optimization` は ZeRO のステージを設定します。

    ```json
    { "zero_optimization": { "stage": 3 } }
    ```

- バッチサイズと勾配蓄積の引数を`"auto"`に設定してください。これらを[TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments)と不一致する値に手動で設定した場合、誤った値のままトレーニングが静かに続行されます。

    ```json
    {
        "train_micro_batch_size_per_gpu": "auto",
        "train_batch_size": "auto",
        "gradient_accumulation_steps": "auto",
        "gradient_clipping": "auto"
    }
    ```

- `bf16` は学習の精度を設定します。`"auto"` に設定すると、[TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) の `bf16` フラグに追従します。

    ```json
    { "bf16": { "enabled": "auto" } }
    ```

- `stage3_gather_16bit_weights_on_model_save` は、保存前にすべての GPU 間で all-gather を実行し、シャードから完全なテンソルを再構築します。これは ZeRO-3 の引数です。

    ```json
    {
        "zero_optimization": {
            "stage": 3,
            "stage3_gather_16bit_weights_on_model_save": true
        }
    }
    ```

- `overlap_comm` を `true` に設定すると、バックワードパスの処理中に all-reduce のレイテンシを隠すことができます。`allgather_bucket_size` と `reduce_bucket_size` は、通信速度と GPU メモリ使用量のトレードオフを調整します。値を下げるとメモリ使用量は減りますが、通信速度は低下します。

    ```json
    {
        "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 がロックされます。

    ```json
    {
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": { "device": "cpu", "pin_memory": true },
            "offload_param":     { "device": "cpu", "pin_memory": true }
        }
    }
    ```

- `optimizer` と `scheduler` のデフォルト値は、[TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) で設定されたオプティマイザーとスケジューラーです。LAMB などの DeepSpeed ネイティブのオプティマイザーが必要でない限り、[TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments) から値を読み込むように `"auto"` に設定してください。

    ```json
    {
        "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` に設定してください。

    ```json
    {
        "zero_force_ds_cpu_optimizer": false
    }
    ```

## チェックポイント

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

```py
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 はエラーを発生させます。

> [!TIP]
> 異なる並列構成間でレジュームを行う方法については、DeepSpeedの[Universal Checkpointing](https://www.deepspeed.ai/tutorials/universal-checkpointing)ガイドを参照してください。

## 次のステップ

- ZeRO の仕組みについて詳しく知るには、The Ultra-Scale Playbook の [Zero Redundancy Optimizer](https://nanotron-ultrascale-playbook.static.hf.space/index.html#zero_redundancy_optimizer_(zero)) 章をお読みください。
- ZeROの論文をお読みください:[Memory Optimizations Toward Training Trillion Parameter Models](https://hf.co/papers/1910.02054)、[Democratizing Billion-Scale Model Training](https://hf.co/papers/2101.06840)、[Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://hf.co/papers/2104.07857)。