コンパイル

概要

PyTorch 2.0 では、PyTorch コードを最適化されたカーネルに JIT コンパイルすることで実行速度を向上させる強力な機能 torch.compile が導入されました。torch.compile の主な機能は以下の通りです。

• パフォーマンスの向上: 計算グラフを最適化することで、モデルの実行を大幅に高速化します。
• 使いやすさ: 実装に必要なコード変更が最小限で、非常に導入しやすいです。
• 互換性: 既存のPyTorchコードやモデルとシームレスに動作します。

Accelerate と組み合わせて使用する場合、torch.compile は分散学習ワークフローにスムーズに統合され、分散実行とコンパイル最適化の両方の恩恵を同時に得ることができます。

コンパイル済みコードの最初の実行はコンパイル時間を含むため通常より長くなりますが、その後の実行は大幅に高速化されます。異なるシナリオでの最適なパフォーマンスのために、torch.compileは"default"、"reduce-overhead"（CUDA グラフを使用してオーバーヘッドをさらに削減します）、"max-autotune"（モデルに最適なカーネルを見つけるために広範なオートチューニングを実行します）などのさまざまなモードを提供しています。

Accelerateでのtorch.compileの使用

Accelerate は、torch.compile を学習スクリプトに簡単にシームレスに統合するための TorchDynamoPlugin を提供しています。

from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    backend="inductor",  # Options: "inductor", "aot_eager", "aot_nvfuser", etc.
    mode="default",      # Options: "default", "reduce-overhead", "max-autotune"
    fullgraph=True,
    dynamic=False
)

# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply torch.compile to your model
model = accelerator.prepare(model)

Accelerateの他のすべての機能やプラグイン、例えば混合精度、分散学習（DDP、FSDP、Deepspeed）などとも互換性があります。

地域別コンパイル

通常、モデル全体をコンパイルしようとすると、最適化の対象となる問題空間が大きくなるため、問題が生じることがあります。そこで、リージョナルコンパイルでは、同じクラスの繰り返しブロックを対象とし、それらを順次コンパイルすることでコンパイラのキャッシュにヒットさせます。例えば、GPT2LMHeadModelでは、繰り返しブロック/クラスはGPT2Blockであり、model.transformer.h[0]としてアクセスできます。モデルの残りの部分（例: model.lm_head）は、別途コンパイルされます。

これにより、LLMやTransformersなどのモデルのコンパイルオーバーヘッドやコールドスタートを高速化できます。 詳細はhttps://pytorch.org/tutorials/recipes/regional_compilation.htmlを参照してください。

リージョナルコンパイルの使用方法

TorchDynamoPluginの設定でuse_regional_compilation=Trueを設定することで、これを有効にできます。

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    use_regional_compilation=True,
    ... # other parameters
)
# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply compile_regions to your model
model = accelerator.prepare(model)

accelerate.utils.compile_regions ユーティリティを、torch.compile と同じ方法で直接使用することもできます。

地域コンパイルの利点

PyTorchのtorch.compile機能を用いて、フルコンパイルとリージョナルコンパイルを比較する大規模なベンチマークを実施しました。完全な結果はaccelerateリポジトリで確認できます。ベンチマークから得られた主な知見は以下の通りです。

1. 同等のパフォーマンス: リージョンコンパイルは、フルコンパイルと類似したパフォーマンスの高速化を実現します。特に大規模なモデルにおいて顕著です。
2. より高速なコンパイル: リージョン別コンパイルにより、モデルのコンパイルにかかる時間が大幅に短縮され、デプロイメントにおいてより効率的な選択肢となります。
3. バッチサイズの影響: バッチサイズが大きくなると、コンパイル戦略間のパフォーマンス差が縮小します。これは、コンパイルのオーバーヘッドがこれらのシナリオでは影響が小さいことを示しています。
4. モデルサイズに関する考慮事項: 地域コンパイルの利点は、コンパイル時間の削減効果が顕著な大規模モデルにおいてより明確に現れます。
5. 実践的な応用: 実際のアプリケーションでは、特に大規模モデルを扱う場合、トレーニングのコールドスタート時間を最適化するための現実的な選択肢として、リージョンコンパイルが挙げられます。

結論

フルコンパイルとリージョナルコンパイルの両方が、モデルを大幅に高速化できます。リージョナルコンパイルは、コンパイル時間と実行時パフォーマンスの間に実用的なバランスを提供し、特に大きなバッチサイズで大きなモデルをトレーニングする場合に有効です。