10月17日、PyTorchは最新バージョン「PyTorch 2.5」をリリースした。本リリースには、AIモデルのパフォーマンス向上を実現する新機能が数多く搭載されている。
以下では、主な改良点を紹介する。
• CuDNNバックエンド(SDPA向け): NVIDIA H100 GPU上で75%の高速化が可能であり、SDPAを利用する全ユーザーに対しデフォルトで提供される。
• トーチコンパイルの「リージョナルコンパイル」: 特定のモジュール(例:LLMのトランスフォーマーレイヤー)を再コンパイルすることなくコンパイルでき、初期起動時間の短縮を実現する。
• TorchInductor CPPバックエンド: FP16サポートや自動チューニングモードの強化を含む最適化が導入され、パフォーマンス向上が図られた。
詳細な機能の紹介
CuDNNバックエンド for SDPA
NVIDIA H100 GPUで最大75%の高速化を可能にする「Fused Flash Attention」バックエンドを搭載。高いスピードが求められるアテンション演算に適している。
トーチコンパイルのリージョナルコンパイル
モジュールを部分的にコンパイルできるため、起動時間の短縮が図れる機能。全モデルをコンパイルするより1%-5%の性能劣化で済む。
TorchInductor CPUバックエンド最適化
独自のCPUカーネルへの最適化で、97.5%のモデルにおいてパフォーマンス向上を達成。これにより、LinuxおよびWindowsでの互換性が向上し、トレーニングや推論において効率化が実現した。
プロトタイプ機能
FlexAttention
スライディングウィンドウや因果マスクといった様々な注意機構を数行のコードで実装可能。性能を手書きの実装と同等にするための最適化が施されている。
Compiled Autograd
後方計算のフローをキャプチャする新機能。AOTディスパッチャとは異なり、フックも含めて後方グラフ全体を追跡できるため、安定性が高い。
Flight Recorder
新しいデバッグツール。ジョブが停止した際の情報を記録し、問題の迅速な特定を支援する機能を提供する。
詳細はPyTorch 2.5 Release Blogを参照していただきたい。