Transformers documentation
DeepSpeed
DeepSpeed
DeepSpeedは、Zero Redundancy Optimizer (ZeRO) を活用した、非常に大規模なモデルを GPU に載せて学習・フィッティングするための最適化ライブラリです。複数の ZeRO ステージが用意されており、各ステージでオプティマイザ状態、勾配、パラメータを分割し、CPU や NVMe へのオフロードを有効化することで、段階的に GPU メモリをより多く節約します。DeepSpeed は Trainer クラスと統合されており、セットアップの大部分は自動的に処理されます。
ただし、Trainer を使わずに DeepSpeed を使用したい場合は、Transformers が HfDeepSpeedConfig クラスを提供しています。
HfDeepSpeedConfig
class transformers.integrations.HfDeepSpeedConfig
< source >( config_file_or_dict )
このオブジェクトには DeepSpeed の設定辞書が含まれており、zero stage などの情報を簡単に照会できます。
このオブジェクトのweakrefがモジュールのグローバル変数に保存されており、Trainer オブジェクトなどが利用できない領域（from_pretrainedや_get_resized_embeddingsなど）から設定にアクセスできるようにするためです。そのため、プログラムが実行されている間、このオブジェクトが生存し続けることが重要です。
Trainerは、代わりにHfTrainerDeepSpeedConfigサブクラスを使用します。このサブクラスには、"auto"という特別なプレースホルダー値を置き換えることで、TrainingArgumentsの値と設定を同期するロジックが含まれています。この特別なロジックがない場合、DeepSpeedの設定は変更されません。