Accelerate documentation
低精度トレーニング手法
低精度トレーニング手法
Accelerate は、TransformersEngine および torchao パッケージを介して、指定された対応ハードウェア上で低精度手法を用いた学習を行うための統合機能を提供します。このドキュメントでは、対応するハードウェア、Accelerator を低精度手法を活用するように設定する方法、および学習時に期待できる内容について、ご案内します。
FP8でのトレーニングとは
PyTorchとAccelerateによるFP8でのトレーニングの詳細な仕組みをさらに調べるには、なぜこれが難しいのかを説明したconcept_guideを参照してください。基本的には、BF16でのトレーニングではなく、モデルのトレーニングのいくつか(またはすべて)の側面を16ビットではなく8ビットで実行できます。課題は、最終的なパフォーマンスを劣化させずにこれを達成することです。
これは特定のNVIDIAハードウェアでのみ有効化されます。具体的には、以下のハードウェアです:
- 3000シリーズのコンシューマー向けグラフィックスカード(4090など)以降の製品
- HopperベースのGPUアーキテクチャ(
H100やH200など)
これにより、メモリ使用量が削減され(トレーニングの一部で必要なメモリを半分に削減したため)、特定のレイヤーをFP8対応のレイヤーに置き換え可能な大規模モデルでは、スループットの向上も期待できます。
Acceleratorの設定
現在、FP8用にアクティブにメンテナンスされている2つのバックエンド(TransformersEngineとtorchao)がサポートされており、それぞれ異なる機能と設定を持っています。
どちらを使用する場合も、同じコアAPIが使用されます。accelerate configの実行中に混合精度について質問された際に、またはconfig.yamlファイルのmixed_precisionキーとして、Acceleratorにmixed_precision="fp8"を渡すだけです:
from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp8")バックエンドを指定し(FP8混合精度設定の他の部分もカスタマイズするには)、utils.AORecipeKwargs や utils.TERecipeKwargs などの RecipeKwargs データクラスを利用できます。また、設定の yaml 内や accelerate launch 実行時に明示することも可能です。新規プロジェクトでは、TransformersEngine または torchao の使用を推奨します。
from accelerate import Accelerator
from accelerate.utils import TERecipeKwargs, AORecipeKwargs
# Use TransformersEngine
kwargs = [TERecipeKwargs()]
# Or to use torchao
# kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)mixed_precision: fp8
fp8_config:
amax_compute_algo: max
amax_history_len: 1024
backend: TE
fp8_format: HYBRID
interval: 1
margin: 0
override_linear_precision: (false, false, false)
use_autocast_during_eval: falseTransformersEngineの構成
TransformersEngineには、FP8計算の実行方法や対象をカスタマイズするための多くのオプションがあります。サポートされている引数の完全なリストとその意味はNVIDIAのドキュメントで確認できますが、ご参考までにFP8KwargsHandlerのdocstringにも記載されています。
Accelerateは適切なデフォルト値を設定しようとしますが、さまざまなパラメータを自分で探索・調整することで、パフォーマンスが向上する可能性があります。
使用するには、backend="te"を指定し、kwargハンドラーの一部として、変更したい引数を修正します。
from accelerate import Accelerator
from accelerate.utils import TERecipeKwargs
kwargs = [TERecipeKwargs(...)]
accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=kwargs)または、--fp8_backend=te ... を介して accelerate launch の実行時です。関連する引数を確認するには、accelerate launch --fp8_backend=te -h を使用してください。
同様に、config.yamlで設定することもできます:
mixed_precision: fp8
fp8_config:
amax_compute_algo: max
amax_history_len: 1024
backend: TE
fp8_format: HYBRID
interval: 1
margin: 0
override_linear_precision: (false, false, false)
use_autocast_during_eval: falsetorchaoの構成
torchao は PyTorch によるハック可能な FP8 バックエンドであり、従来の 2 つのエンジンよりも扱いやすくすることを目標としています。ao が従来の 2 つと根本的に異なる点の一つは、数値の安定性の観点から、モデルの最初 と 最後の層を通常の精度(FP32 または BF16)のままにし、それ以外の層を FP8 に量子化するのが一般的により良いとされていることです。そのため、ao の設定は少し異なっています:
注: このAPIは実験的であり、変更される可能性があります。
from accelerate import Accelerator
from accelerate.utils import AORecipeKwargs, TorchDynamoPlugin, FullyShardedDataParallelPlugin
from torchao.float8 import Float8LinearConfig
fsdp2_plugin = FullyShardedDataParallelPlugin(
fsdp_version=2,
cpu_ram_efficient_loading=False, # CPU RAM efficient loading CANNOT work with fp8 torchao
fsdp_auto_wrap_policy="TRANSFORMER_BASED_WRAP",
)
dynamo_plugin = TorchDynamoPlugin(
backend="inductor",
use_regional_compilation=True,
)
fp8_config = Float8LinearConfig(
enable_fsdp_float8_all_gather=True, # Use FP8 all_gather in FSDP2
pad_inner_dim=True,
)
kwargs = [AORecipeKwargs(
config=fp8_config
)]
accelerator = Accelerator(
mixed_precision="fp8",
fsdp_plugin=fsdp2_plugin,
dynamo_plugin=dynamo_plugin,
kwarg_handlers=kwargs,
)または、--fp8_backend=ao ... を介して accelerate launch 実行時です。関連する引数を確認するには、accelerate launch --fp8_backend=ao -h を使用してください。
同様に、config.yamlで設定することもできます:
mixed_precision: fp8
fsdp_config:
fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
fsdp_cpu_ram_efficient_loading: false
fsdp_version: 2
fp8_config:
backend: AO
pad_inner_dim: true
enable_fsdp_float8_all_gather: true
dynamo_config:
dynamo_backend: INDUCTOR
dynamo_use_regional_compilation: true使用する具体的なパラメータの詳細については、公式のtorchaoリポジトリを参照してください。
例のズー
accelerate リポジトリには、accelerate およびその基盤となる実装を用いた FP8 での学習を示すサンプルがあります。 現在、以下のことを示すスクリプトをサポートしています:
- シングルGPU
- 分散データ並列性(マルチGPU)
- 完全シャード化データ並列
- DeepSpeed ZeRO 1から3
詳細はこちらをご覧ください。
さらに詳しく
FP8での学習について詳しく知るには、以下のリソースをご覧ください:
- TransformersEngineとtorchaoについて詳しく説明したコンセプトガイド
transformers-engineのドキュメントtorchaoのドキュメント