DeepSpeed

DeepSpeed

DeepSpeedとは

DeepSpeed(ディープスピード)とは、Microsoftが開発した大規模深層学習の最適化ライブラリです。ZeRO(Zero Redundancy Optimizer)テクノロジーを中核に、メモリ効率の大幅な改善と学習の高速化を実現します。数兆パラメータ規模のモデル学習を可能にする技術として広く採用されています。

表:DeepSpeedの要点まとめ
ひとことで言うとMicrosoft製の、大規模モデルを効率よく学習するためのライブラリ。
主な機能ZeRO(メモリ分割)、混合精度、オフロード、パイプライン並列。
注意点設定項目が多く習熟が必要。小規模学習では効果が出にくい。

ZeROの仕組み

ZeROは分散学習におけるメモリの冗長性を排除する技術です。3つのステージがあり、ZeRO Stage 1はオプティマイザ状態の分割、Stage 2は勾配の分割、Stage 3はモデルパラメータ自体の分割を行います。Stage 3ではモデル全体を各GPUに保持する必要がなくなり、GPUメモリの使用量を劇的に削減できます。

図:DeepSpeed ZeROの段階
1ZeRO-1最適化状態を分割
▶
2ZeRO-2勾配も分割
▶
3ZeRO-3重みも分割
▶
4オフロードCPUメモリも活用
▶
5巨大モデルが学習可少ないGPUでも動く

DeepSpeedの主要機能

ZeROに加えて、DeepSpeedは混合精度学習、勾配チェックポインティング、パイプライン並列、3D並列、DeepSpeed-Inference(推論最適化)、DeepSpeed-Chat(RLHFの効率化)など多くの機能を提供します。設定ファイル(JSONベース)で機能の有効化・パラメータ調整が可能で、コード変更を最小限に抑えられます。

DeepSpeedの活用

多くの大規模言語モデルの学習にDeepSpeedが活用されており、Hugging Face TransformersのTrainerクラスとも統合されています。Azureとの親和性が高く、Azure Machine Learning上でのスケールアウトが容易です。