DeepSpeedについてのメモ

機械学習

目次

DeepSpeedとは

基本的に公式の日本語解説がとてもわかりやすいのでそちらを参照。
ここでは概要だけまとめておく。

ZeRO（基本的にDeepSpeedを使用する場合はこれがメインとなる）

  • ZeRO
    データ並列のメモリ利用を効率化してくれる機能
    通常のデータ並列とは違って、それぞれのGPUのVRAMにはモデル全体を置かずに分散して置いておき、計算時にパラメータを転送して利用することで、メモリ使用量を削減してくれる。
  • ZeRO-Offload
    学習パラメータをCPUメモリにオフロードすることで、大きいモデルも学習できる
  • ZeRO-Infinity
    CPUメモリだけでなくNVMeストレージにもオフロードするらしい（流石にそれなりに計算速度が低下するのでは？）
    実験では、V100-32GB×16 + CPUメモリ1.5TB + NVMeストレージ20TBで1兆パラメータのモデルが学習できたらしい…

基本的に、上記のような手法を適用するとデータの通信オーバヘッドが大きくなり、計算速度が大幅に低下してしまうが、この辺りをいい感じにやりくりしながらメモリ使用量を抑えてくれるのが、DeepSpeedの恩恵となっている。

基本的な用途としては、VRAMだけだとギリギリでしか動かない（バッチサイズ2とかで学習させないといけない）モデルに適用して大きいバッチサイズで学習させたり、複数台のGPUを並列で使用して学習させたり、巨大なモデルを計算効率度外視で動かしたいという際に使用する。

また、メモリの削減によってバッチサイズを大きくできれば、学習に要する時間も短縮できるので、そういう使い方も良さそう。

3D Parallelism

モデル並列（巨大なモデルを複数GPUに展開して学習）を行う機能。
ZeROのように簡単な設定だけで…とはいかなさそうだが、ちょっとコードを変更すれば使えそう…
今後挑戦してみたい

その他

DeepSpeed InferenceやDeepSpeed-MII、DeepSpeed Compressionなどの推論時に計算速度や量子化によるメモリ削減やモデルの圧縮をしてくれる機能や、DeepSpeed-MoEなどのMixture of Experts（MOE）を適用してモデルの一部だけを計算するようにしたり、Communication Compressionで独自のoptimizerを使用することでオーバヘッドを小さくする機能、DeepSpeed Data Efficiencyなどの学習データを効率的に使用して学習させる（学習時に使用するデータを簡単なデータ→難しいデータのようにスケジューリングする）機能、DeepSpeed-ChatというRLHFを高速・低コストで実施できる機能などがある。

DeepSpeed（ZeRO）とTrainerを使用して学習を行う場合

Hugging FaceのTrainer（SFTTrainerのような派生系も可）を使用している場合は、DeepSpeedの設定ファイルをJSONで書いて以下のように実行すればOK。
JSONも簡単な設定だけであり、配布されているサンプルをそのまま使ってもちゃんと動く。
なお、--num_gpusは使用するGPUの数を選択でき、指定しない場合は使用できるGPUをすべて使用して計算する。

pip install deepspeed
deepspeed --num_gpus=2 your_program.py <normal cl args> --deepspeed ds_config.json

もしくは、Trainerに設定ファイルを指定して実行することもできる。

training_args = TrainingArguments(
    output_dir='./output',  # 出力パス
    num_train_epochs=1,  # エポック数
    per_device_train_batch_size=1,  # 学習バッチサイズ
    per_device_eval_batch_size=1,  # 評価バッチサイズ
    gradient_accumulation_steps=1,  # 勾配の累積ステップ数
    gradient_checkpointing=True,  # 勾配チェックポイント
    optim='adafactor',  # オプティマイザの種類
    deepspeed='./zero_train.json',  # deepspeed設定ファイル
    logging_steps=100,  # 途中経過を表示する間隔
)

# 学習の実行
trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    max_seq_length=512,
    formatting_func=formatting_prompts_func,
    args=training_args,
)
trainer.train()

DeepSpeedで学習する際の注意点

DeepSpeedとTrainerで設定項目が重複している部分があり、両方で違う設定を使用しているとトラブルの原因になる。例えば、optimizerや学習率などの設定項目などが共通しているため、同じ設定にする必要がある。

以下設定の優先順位をまとめる。
ただし、DeepSpeedとTrainerのドキュメントで確認した情報ではないので誤りがある可能性はそこそこ高い。（結局公式のサンプルコードやライブラリのコードを読まないとよくわからないかも…）

  • DeepSpeedとTrainerで設定が競合する場合はTrainerの設定が学習で使用されるらしい（ただし、DeepSpeedの設定はそのままな可能性がある）
  • Trainerにのみ設定があり、DeepSpeedで明示的に設定されていない場合は、Trainerは設定値を使用し、DeepSpeedはデフォルトの値を使用するらしい
  • DeepSpeedにのみ設定があり、Trainerで明示的に設定されていない場合は、TrainerはDeepSpeedの設定を使用するらしい

学習ではTrainerの設定が優先されるようだが、DeepSpeedの設定が勝手に上書きされたりはしなさそうなので、設定が競合すると問題が起こりそう…
考えられる解決策は以下の３つ。

1.手動でDeepSpeedとTrainerの設定項目を合わせる（面倒だけど確実）

2.DeepSpeedでautoの設定を積極的に使用する（公式だとこれが推奨されている）
optimizerの種類としてAdamWを使用するなどの設定はしないといけないが、それ以外はautoにしてDeepSpeedに任せることができる

"optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

3.DeepSpeedの設定ファイルを参照して、TrainerのTrainingArgumentsを設定する
設定についてのルールを見る限り、Trainerで設定を行わなければデフフォルトでDeepSpeedの設定が使用されるようなので、以下のようにわざわざ指定しなくても良い気もする…

import json
from transformers import Trainer, TrainingArguments

# DeepSpeedの設定ファイルを読み込む
with open("./zero_train.json", "r") as json_file:
    ds_config = json.load(json_file)

# DeepSpeedの設定に基づいてTrainingArgumentsを作成
training_args = TrainingArguments(
    output_dir='./output',
    overwrite_output_dir=True,
    per_device_train_batch_size=ds_config["train_batch_size"],  # バッチサイズをds_configから取得
    learning_rate=ds_config["optimizer"]["params"]["lr"],  # 学習率をds_configから取得
    # 他のトレーニング関連の設定を追加できます
)

DeepSpeed（ZeRO）を使用してTrainerを使わずに学習する場合

メチャクチャ面倒くさそう…
余裕ができたら更新する