Published on

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

Author
株式会社イエローバックの機械学習エンジニアです

はじめに

前回の記事で 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 が以下になります。

~/ldn/ds_config_zero2.json
{
    "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()あたりで設定されていますが、まとめると以下のようになります。

| deepspeed config | transformers args | | ------------------------------ | ------------------------------------------------------------------------------------------------------------------------------- | | train_micro_batch_size_per_gpu | args.per_device_train_batch_size | | gradient_accumulation_steps | args.gradient_accumulation_steps | | train_batch_size | args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps | | gradient_clipping | args.max_grad_norm | | optimizer.lr | args.learning_rate | | optimizer.betas | [args.adam_beta1, args.adam_beta2] | | optimizer.eps | args.adam_epsilon | | optimizer.weight_decay | args.weight_decay | | scheduler.warmup_min_lr | 0 | | scheduler.warmup_max_lr | args.learning_rate | | scheduler.warmup_num_steps | args.warmup_steps | | scheduler.total_num_steps | 読み込んだ学習データから計算された値 | | fp16.enabled | fp16_backend == "amp"または"apex"のときに"enabled"。 ただし"apex"のときには ZeRO 関連機能が有効にならないため apex は避けるべき |

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

| deepspeed config | transformers args | | ---------------------------------------------------- | --------------------------------- | | zero_optimization.reduce_bucket_size | hidden_size * hidden_size | | zero_optimization.stage3_prefetch_bucket_size | 0.9 * hidden_size * hidden_size | | zero_optimization.stage3_param_persistence_threshold | 10 * hidden_size |

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 のコメント によると、以下のようになっています。

| DS name | --lr_scheduler_type | HF func | Notes | | ------------- | -------------------- | --------------------------------- | ------------------ | | LRRangeTest | na | na | LRRT | | OneCycle | na | na | 1CLR | | WarmupLR | constant_with_warmup | get_constant_schedule_with_warmup | w/ warmup_min_lr=0 | | WarmupDecayLR | linear | get_linear_schedule_with_warmup | |

おわりに

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

Discuss on Twitter