Re Reference AI

インフラ

DeepSpeedとは

ZeRO

ZeROによる分散学習を中心に、大規模モデルの学習・推論を効率化するMicrosoftのオープンソースライブラリ

フレームワーク並列計算

ひとことで言うと

1枚のGPUには載らない巨大なAIモデルを、たくさんのGPUに分担させて学習するためのMicrosoft製ツール。

概要

DeepSpeedとは、Microsoftが開発する、大規模モデルの学習・推論を効率化するオープンソースの深層学習最適化ライブラリを指す。 オプティマイザ状態・勾配・パラメータをGPU間で分割して保持するZeROにより、1枚のGPUに載らない規模のモデルの分散学習を可能にする。 PyTorch上で動作し、パイプライン並列・テンソル並列・CPUへのオフロードなどを組み合わせて、数百億から数千億パラメータ級のモデル学習を支える標準的な基盤の1つとして使われてきた。

歴史

2020年2月にMicrosoftがZeROの論文とともに公開した。

コード例

ZeRO Stage 2を有効にする設定ファイルの例

{
  "train_batch_size": 16,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 0.00015 }
  },
  "zero_optimization": {
    "stage": 2
  }
}

利点

  • ZeROによりオプティマイザ状態・勾配・パラメータをGPU間で分割し、単一GPUのメモリを超える規模のモデルを学習できる
  • PyTorchとの統合が進んでおり、既存の学習コードへ比較的小さな変更で組み込める
  • パイプライン並列・テンソル並列・CPU/NVMeへのオフロードなど複数の並列化手法を組み合わせられる

欠点

  • 設定項目が多く、モデル規模やハードウェア構成に応じたチューニングは専門知識を要する
  • 分散学習特有のデバッグの難しさ(通信のボトルネックや同期の問題など)が伴う

関連用語

PyTorchGPURay

よくある質問

DeepSpeedはどのモデルでも使えるか。

PyTorchベースのモデルであれば適用しやすく、Hugging FaceのTransformersなど主要なライブラリとの統合も進んでいるが、モデルやタスクによっては設定の調整が必要になる。

ZeROのステージとは何か。

ZeROが持つメモリ最適化の段階を示す区分で、段階が進むほどオプティマイザ状態・勾配・パラメータの分割範囲が広がり、メモリ効率と通信コストのトレードオフが変化する。

参考文献

関連Zenn記事