PyTorch/XLA 2.4でさらに使いやすく：開発者体験の向上とeager modeの導入

PyTorch/XLAチームは、開発者の課題に対処するためのいくつかの顕著な改善を含む、PyTorch/XLA 2.4のリリースを発表しました。

今回のリリースで特に興味深いのは、カスタムカーネル言語であるPallasの強化と、実験的なeager modeの導入です。

Pallasは、TPUとGPUの両方をサポートしており、C++のような低レベルで複雑な言語を使用せずに、ハードウェアに近いパフォーマンスの高いコードをPythonで記述することを可能にします。今回のアップデートでは、Flash Attentionのサポートや、推論のためのPaged Attentionの組み込みサポートなどが追加され、より使いやすくなりました。

eager modeは、PyTorch/XLAのモデルが「遅延実行」されていたのに対し、操作をコンパイルしてすぐにターゲットハードウェア上で実行するモードです。TPU自体は真のeager modeを持っていないため、PyTorch/XLA 2.4では、各PyTorch操作の後に「mark_step」呼び出しを追加することで、コンパイルと実行を強制することで実現しています。

まだ実験的な機能ではありますが、このeager modeの導入により、開発者は大規模なデバイスフリートにデプロイすることなく、自分のマシン上でローカルにモデルをデバッグすることが容易になります。

さらに、Cloud TPU用のコマンドラインユーティリティ「tpu-info」も導入されました。これは、Nvidia GPUの「nvidia-smi」ツールと同様に、TPUのワークロードのデバッグ、使用されているコアの特定、ワークロードが消費しているメモリ量の確認などを簡単に行うことができます。

今回のPyTorch/XLA 2.4は、開発者にとってより使いやすく、より強力なツールとなるための重要なアップデートと言えるでしょう。

参照元サイト：Announcing PyTorch/XLA 2.4: A better Pallas and developer experience, plus “eager mode”

NewsPilot編集部

世界の様々なニュースを、日本語でわかりやすく、いち早くお届けします！