PyTorch Foundationは最近、Intel GPUのサポートを含むPyTorchバージョン2.5をリリースした。このリリースには、FlexAttention API、TorchInductor CPUバックエンドの最適化、コンパイル時間を短縮するリージョナルコンパイル機能など、いくつかのパフォーマンス強化も含まれている。全体として、このリリースには PyTorch 2.4 以降の 4095 のコミットが含まれている。
Intel GPUサポートは最近のPyTorchカンファレンスでプレビューされた。IntelのエンジニアEikan Wang氏とMin Jean Cho氏は、ハードウェアをサポートするために行われたPyTorchの変更について説明した。これには、PyTorchのランタイムとデバイスレイヤーを一般化し、新しいハードウェアバックエンドを統合しやすくすることが含まれる。torch.compileとtorch.distributedにはIntel固有のバックエンドも実装された。IntelのAIフレームワーク・エンジニアリング担当副社長であるKismat Singh氏によると、以下の通りである。
我々はPyTorch 2.5でIntelクライアントGPUのサポートを追加した。これは基本的に、最新のIntelプロセッサーを使って作られたIntelのラップトップやデスクトップでPyTorchを実行できるようになることを意味する。今年中に4,000万台のラップトップとデスクトップがPyTorchユーザーのためにアンロックされ、来年末までにその数は1億台程度になると考えている。
このリリースには新しいFlexAttention APIが含まれており、PyTorchユーザーがモデルで異なるアテンションメカニズムを簡単に試すことができる。通常、新しい注意のバリエーションを試したい研究者は、PyTorchのオペレータから直接ハンドコーディングする必要がある。しかし、これは「遅い実行時間とCUDA OOM」を引き起こす可能性がある。新しいAPIは、代わりに「数行の慣用的なPyTorchコード」でこれらを書くことをサポートする。そしてコンパイラーは、これらを「余分なメモリを生成せず、手書きのものと遜色のないパフォーマンスを持つ」最適化されたカーネルに変換する。
いくつかの性能向上がベータ版としてリリースされた。新しいバックエンドのFused Flash Attentionは、NVIDIA H100 GPU向けに「FlashAttentionV2より最大75%の速度向上」している。torch.compileの地域コンパイル機能は、フルモデルコンパイルの必要性が軽減され、代わりにTransformer層などの繰り返しnn.Moduleがコンパイルされる。これにより、コンパイルの待ち時間が短縮される一方で、性能は数パーセントしか低下しない。また、TorchInductor CPUバックエンドにもいくつかの最適化が施されている。
スタックしたジョブの新しいデバッグツールであるFlight Recorderもリリースに含まれる。ジョブのスタックは分散トレーニング中に発生する可能性があり、データの枯渇、ネットワークの問題、ソフトウェアのバグなど、多くの根本原因が考えられる。Flight Recorderは、メモリ内の循環バッファを使用して診断情報をキャプチャする。スタックしたジョブを検出すると、診断情報をファイルにダンプし、そのデータをヒューリスティックのスクリプトを使って分析し、根本原因を特定ができる。
Redditでのリリースに関する議論では、多くのユーザーがIntel GPUのサポートを喜んでおり、「ゲームチェンジャー」と呼んでいる。別のユーザーはこう書いている。
torch.compileの改良、特にコンパイルを高速化するために繰り返されるモジュールを再利用できるようになったことに興奮している。似たようなコンポーネントがたくさんある大規模なモデルにとっては、画期的なことだ。FlexAttention APIも本当に期待できそうだ。ほんの数行のコードで様々なアテンションメカニズムを実装でき、手書きに近いパフォーマンスが得られるのは大きい。PyTorchチームとコントリビュータの皆様、またしても堅実なリリースをありがとう！
PyTorch 2.5のコードとリリースノートはGitHubで公開されている。