FSDPとDeepSpeed

Accelerate は、分散学習のための非常に強力な2つのツール、すなわち Pytorch FSDP と Microsoft DeepSpeed を統合することで、学習フレームワークの柔軟性を提供します。本チュートリアルの目的は、これら2つのフレームワークの類似点と潜在的な相違点を示し、ユーザーがこれら2つのフレームワークの間でシームレスに切り替えられるようにすることです。

フレームワークを切り替えるには、--config_file で正しい設定ファイルを指定して accelerate launch を実行するか、FSDP と DeepSpeed の場合はそれぞれの引数を直接渡すことをお勧めします。

DeepSpeedおよびFSDPのAccelerate設定の例は、“Launch Configurations”の例のコレクションで確認できます。

このチュートリアルは、単一ノード・マルチGPUのシナリオのみを対象としています。

機能の設定

モデルのテンソルは、モデルサイズを拡大する目的で異なるGPUに分割されます。これはFSDPではシャーディング、DeepSpeedではパーティショニングと呼ばれます。FSDPのシャーディングとDeepSpeed ZeRO（パーティショニング）のステージは、それぞれ--fsdp_sharding_strategyおよび--zero_stageによって設定されます。特に、FSDPのFULL_SHARDはDeepSpeed ZeROステージ3に対応します。FSDPのシャーディングとDeepSpeed ZeROの設定間の包括的な対応関係については、こちらを参照してください。以下の表は、類似した設定を要約してグループ化しています。

上記の詳細な説明については、Accelerateの起動ドキュメントを参照してください。

混合精度設定などの他の DeepSpeed 設定にアクセスするには、--deepspeed_config_file を渡す必要があります。ドキュメントを参照してください。

DeepSpeed は DeepSpeedPlugin 経由でも設定できます。例えば、DeepSpeedPlugin.zero_stage は --zero_stage と同等であり、DeepSpeedPlugin.hf_ds_config を使用して --deepeed_config_file. を渡すことができます。

FSDPは、FullyShardedDataParallelPluginでも設定できます。例えば、FullyShardedDataParallelPlugin.sharding_strategyは--fsdp_sharding_strategyと同等です。

チェックポイント

FSDPは--fsdp_state_dict_typeで設定することで、フルまたはシャーディングされたチェックポイントを保存できる点に注意してください。

DeepSpeed Zero3 の場合、--zero3_save_16bit_model true を渡すことで、モデルを単一のランクに集約して保存できます。これは fsdp_state_dict_type: FULL_STATE_DICT の FSDP における同等の機能です。

大規模モデルの場合、モデルを単一のランクに集約する処理は非常に遅くなる可能性があります。

より高速なチェックポイント保存のために、FSDP では fsdp_state_dict_type: SHARDED_STATE_DICT を使用し、DeepSpeed Zero3 では zero_to_fp32.py スクリプトを使用してシャード化されたチェックポイントを後処理で変換します。

オフローディング

FSDPはオールオアナッシングのオフロードのみを許可します（つまり、パラメータ、勾配、オプティマイザーをすべてオフロードするか、すべてGPUに保持するかのどちらかです）。一方、DeepSpeedはパラメータとオプティマイザーを異なる方法でオフロードできます。さらに、DeepSpeedはNVMEへのオフロードもサポートしています。

プリフェッチ

FSDPでは、--fsdp_forward_prefetchと--fsdp_backward_prefetchという2つのプリフェッチ設定により、追加のメモリをコストとして、通信と計算のオーバーラップを改善できます。FSDPのドキュメントを参照してください。 DeepSpeedの場合、必要に応じてプリフェッチが有効化され、stage3_param_persistence_thresholdやstage3_max_reuse_distanceなどの特定のハイパーパラメータに基づいて有効化されます。これらはZero3用に設定できます。accelerateは、deepspeed設定ファイルでこれらのハイパーパラメータを明示的に設定しない場合、自動的に設定する場合があります。

メモリが許容する場合は、スループットを向上させるために FSDP で fsdp_backward_prefetch: BACKWARD_PRE を設定してください。

モデルの読み込み

FSDPでは効率的なモデルの読み込みを有効にするために明示的な--fsdp_cpu_ram_efficient_loading trueが必要ですが、transformersではDeepSpeed Zero3が使用される場合、同様の機能が自動的に有効になります。

FSDPでは、--fsdp_cpu_ram_efficient_loading trueを設定すると、accelerateは自動的にsync_module_statesをtrueに設定します。 RAM効率の読み込みでは、重みが単一のランクでのみ読み込まれるため、他のランクに重みをブロードキャストするためにsync_module_statesが必要です。

モデル

FSDPでは、all-gatherおよびreduce-scatter操作のスケジューリング方法をアルゴリズムが決定するために、--fsdp_auto_wrap_policyを明示的に指定する必要があります。一方、DeepSpeedでは、この処理はユーザーに対して透過的です。

FSDPでは、fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAPを設定するだけです。最新のtransformersバージョンでは、HF transformersモデルに適したfsdp_transformer_layer_cls_to_wrapを特定するように努めています。ただし、これに関するエラーが発生した場合は、この値を明示的に指定してください。

パラメータの召喚

torch.compile を使用する場合、FSDP には明示的な --fsdp_use_orig_params フラグが必要です。PyTorch のドキュメントを参照してください。DeepSpeed の場合は、ユーザーには透過的です。

FSDPを使用する際は、torch.compileを使う場合、fsdp_use_orig_params: Trueを設定してください。

トレーニング

DeepSpeedでは、--gradient_accumulation_stepsおよび--gradient_clippingフラグを明示的に指定する必要があります。FSDPの場合、これはユーザーにとって透過的です。

DeepSpeed を使用する場合、gradient_accumulation_steps: "auto" と gradient_clipping: "auto" を設定することで、Accelerator または TrainingArguments（transformers を使用している場合）で設定された値が自動的に取得されます。

データ精度処理の違いについて

FSDP と DeepSpeed におけるデータ精度の扱いについて議論するために、まずこれらのフレームワークでモデルパラメータがどのように扱われるかの概要を理解することが有益です。モデル / オプティマイザのパラメータが GPU 間で分散される前に、パラメータの準備段階として、それらを 1 次元の torch.Tensor に「フラット化」する処理が行われます。FSDP / DeepSpeed の実装は、これらの「フラット化」されたパラメータを保存する dtype に関して異なり、torch.Optimizer が dtype を割り当てる方法にも影響が及ぶ点に注意が必要です。下表は、両方のフレームワークのプロセスを概観したものです。「Local」列は GPU 単位で発生するプロセスを示しており、したがって、アップキャストによるメモリオーバーヘッドは、使用される GPU 数で分散されるものと理解すべきです。

経験則として、自動混合精度による安定した学習のためには、すべての学習可能なパラメータがtorch.float32である必要があります。

したがって、DeepSpeed を使用して少数の GPU を利用する場合は、準備段階での型変換（アップキャスト）により、メモリオーバーヘッドが大幅に発生する可能性がある点に注意してください。

FSDPでは、混合精度を使用しない場合、torch.Optimizerを低精度のtorch_dtypeで動作させることが可能です。これは、GPUの台数が少ない場合に有用です。

混合精度を使用する場合、FSDP と DeepSpeed はモデルの準備ステップでアップキャストを行います（上記の表を参照）。ただし、FSDP はアップキャストされた精度でチェックポイントを保存する点に注意してください。一方、DeepSpeed は --zero3_save_16bit_model が指定されている場合でも、低精度のチェックポイントを保存する可能性があります。

上記の表を明確にするため、以下に具体的な例を示します。簡潔にするため、オプティマイザーの事前ステップと実際のステップを統合しています。FSDPでは以下に示す2つのモードで動作させることが可能ですが、DeepSpeedでは1つのモードでのみ動作します。