transformersでdeepspeedを使ってみる(config編)

はじめに

前回の記事で deepspeed を用いて文章分類をおこないました。

今回は前回使用した deepspeed の設定まわりをみていこうと思います。

deepspeed の config の各種パラメタについては、https://www.deepspeed.ai/docs/config-json/ に記載されています。 transformers で deepspeed を使用する際には、transformers の trainer に--deepspeedオプションで deepspeed 用の config ファイルを渡します。 その際に一部のパラメタは transformers 側のパラメタをもとに自動で設定することも可能です。

今回は、このあたりについて少し掘り下げてみます。

環境

transformers の deepspped まわりに開発は比較的活発なので、バージョンによって変わる可能性があります。 今回調査したのは以下の環境です。

• python: 3.8
• CuDA: 11.1
• pytorch: 1.8.1
• transformers: 4.6.1
• deepspeed: 0.3.16

サンプル ds 用 config

まず前回用いたサンプル用の config が以下になります。

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "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": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 6e7,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 6e7,
        "contiguous_gradients": true,
        "cpu_offload": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

optimizer や scheduler が未設定の場合に transformers 側のものが使用されますが、今回は deepspeed 側の optimizer,scheduler を利用する前提とします。

自動設定可能な項目

transformers では deepspeed 用の config を読み込んだあとに"auto"と設定された一部のパラメタを transformers のオプションをもとに設定して deepspeed に渡します。

これらは、DeepSpeedConfigHF クラスのconfig_process()とconfig_finalize()あたりで設定されていますが、まとめると以下のようになります。

(以下は ZeRO-3 のときのみ)

optimizer type

deepspeed のoptimizer.typeは transformers 側からは自動で設定されないため、あらかじめ ds 用 config で明示的に指定する必要があります。 transformers のデフォルトは AdamW ですが、deepspeed では、Adam, AdamW, OneBitAdam, Lamb, OneBitLamb をネイティブにサポートしているようです。

scheduler type

deepspeed のscheduler.typeは transformers 側からは自動で設定されないため、あらかじめ ds 用 config で明示的に指定する必要があります。

deepspeed のscheduler.typeと transformers の"--lr_scheduler_type"の対応は transformers のintegrations.py のコメント によると、以下のようになっています。

おわりに

今回は transformers+deepspeed を使用する際の deepspeed の設定まわりについて調べてみました。