PyTorch Foundationは100%後方互換のアップデートであるPyTorch version 2.0をリリースした。今回のリリースの目玉、ディープラーニングモデルのコンパイル機能ではトレーニングのスピードアップを実現しており、163個のオープンソースAIプロジェクトにおけるベンチマークではトレーニング中のモデルの動作が平均43%高速化されたことが示された。
2.0リリースの計画は、2022年12月のPyTorch Conferenceで発表された。このリリースには、新しいコンパイル機能の他にも、scaled dot product attention (SDPA) 関数の導入による、LLMやdiffusionなどのTransformerベースのモデルの性能向上も含まれている。Appleシリコン上でのトレーニングは、改良されたMetal Performance Shaders (MPS)によって加速され、現在では300個の演算がMPSで実装されている。コアリリースに加えTorchAudio、TorchVision、TorchTextなどの個別ライブラリが新しいベータ機能で更新された。本リリースには1.13.1リリース以降から数えると、428人の開発者から4,500以上のコミットによる貢献を経ている。
PyTorch Foundationブログは次のように発信している。
2022/12/2のPyTorch Conferenceで紹介したPyTorch®2.0のリリースを発表できることを嬉しく思います！PyTorch 2.0は、Eager-Modeによる開発とユーザー体験はそのままに、PyTorchのコンパイラレベルでの動作を根本的に変え、高速化し、Dynamic ShapesとDistributedをサポートするようになりました。
PyTorchの共同開発者であるSoumith Chintala氏は、PyTorch Conference 2022の基調講演で、GPU計算能力の向上により、既存のPyTorchワークロードの多くがメモリ帯域幅やPyTorchフレームのオーバーヘッドによって制約を受けていると指摘した。PyTorchチームはこれまで、コアコンポーネントの一部をC++で記述することでパフォーマンスの問題に対処してきた。Chintala氏はPyTorchを「基本的にC++コードベース」と表現し、C++コンポーネントへの貢献は「嫌い」であると述べている。
新しいコンパイル機能は、Python製の4つのコンポーネントに基づいている。
- TorchDynamo : ディープラーニングモデルを表すPythonコードを計算グラフのブロックに書き換えることでグラフ取得を実行。
- AOTAutograd : バックプロップステップのための差分先読みを自動実行。
- PrimTorch : 2000個以上のPyTorchの演算子を約250個のプリミティブな演算子に正規化。
- TorchInductor : アクセラレータのハードウェアに特化したバックエンドコードを高速に生成。
コンパイル機能の性能向上と使いやすさを実証するために、PyTorchチームはベンチマークとして163個のオープンソースのディープラーニングプロジェクトを定めた。これらのプロジェクトは、コンピュータビジョン、自然言語処理、強化学習など、さまざまな用途を目的としている。チームはコンパイル関数を呼び出す1行だけのコード変更を加えた。そしてこのたった1行の変更が93％のプロジェクトで機能し、コンパイルされたモデルは、NVIDIA A100 GPUで43％速く動作したのである。
Hacker Newsのリリースに関する議論で、とあるユーザーがこう述べている。
PyTorchと他のフレームワークの比較から学んだ大きな教訓は、生産性は性能の向上より優先されるということだ。CaffeもMXNetも高速であることを売りにしていたが、どうやらここが何％速くても、それほど重要ではないことがわかった。一方、一度システムを成功させ、それを普及させれば、競合他社が期待するよりも早く、コミュニティが性能差を縮めてくれるであろうということ。もう1つの教訓は、古くさいかもしれないが、何度も言うに値する。つまりオープンソースプロジェクトには、投資と専門的なブラッシュアップが［重要］ってことだ。
PyTorchのコードとバージョン2.0リリースノートは、GitHubに公開されている。