7月25日、PyTorchチームはPyTorch 2.4をリリースした。今回のリリースでは、最新バージョンのPython（3.12）への対応を始め、AOTInductorフリーズ機能の強化、新しいデフォルトTCPStoreサーバーバックエンドの導入、そして新しいPythonカスタムオペレーターAPIの提供が行われた。

主要な新機能と改善点

• Python 3.12サポート: torch.compileがPython 3.12で動作するようになり、最新のPythonバージョンを活用したモデルの最適化が可能となった。
• AOTInductorフリーズ機能: CPU上でのAOTInductor使用時にフリーズフラグを設定することで、MKLDNNオペレーターのシリアル化を可能にし、パフォーマンスの最適化が図られた。以下の環境変数を設定することでこの機能を有効にできる:

  export TORCHINDUCTOR_FREEZING=1
  

• 新しいPythonカスタムオペレーターAPI: 新たに導入された高レベルPythonカスタムオペレーターAPIにより、従来よりも簡単にカスタムオペレーターをPyTorchに統合することが可能となった。このAPIはtorch.compileと互換性があり、PyTorchの他のサブシステムとも連携する。
• TCPStoreサーバーバックエンドのlibuvへの変更: TCPStoreのデフォルトサーバーバックエンドがlibuvに変更され、大規模ジョブの初期化時間が大幅に短縮された。

プロトタイプ機能

• FSDP2: DTensorベースのパラメーターシャーディング: FSDP1の課題を解決する新しい完全シャーディングデータ並列実装。
• torch.distributed.pipelining: 簡易化されたパイプライン並列化ツールキット。

パフォーマンス向上

• AWS Graviton向けtorch.compile最適化: AWS Graviton3プロセッサ向けに最適化されたtorch.compileにより、最大2倍の推論性能向上。
• BF16シンボリック形状最適化: インファレンスサービスの変動するバッチサイズやシーケンス長に対応した性能向上。

詳細はPyTorch 2.4 Release Blogを参照していただきたい。