AIとMLの日本語ニュース
フォロー
PyTorch/XLA 2.4の発表:改善されたPallasと開発者体験、プラス「即時モード」
PyTorch/XLA 2.4は、TPUとGPU上での深層学習向けに大幅に向上を実現します。このリリースでは、カスタムカーネル言語PallasがTPUとGPUの両方をサポートし、Pythonコードで性能を向上させる改善が行われました。新しいAPIコール、例えばtorch_xla.sync()は、既存のPyTorchワークフローとの統合を簡単化します。実験的な即時実行モードでは、ターゲットハードウェア上で即座に操作が実行されますが、TPUでは"mark_step"コールが必要です。Pallasの改善には、Flash AttentionとPaged Attentionのサポート、およびグループ行列乗算のためのビルトインMegablocksブロックスパースカーネルが含まれます。さらに、新しいTPUコマンドラインインターフェースtpu-infoは、Nvidiaのnvidia-smiツールと同様に、デバッグを容易にするために、使用率とデバイス情報を表示します。既存のコードはこれらの変更に対応しているため、開発者にとっては無痛でアップグレードが可能です。これらの更新は、PyTorch/XLAの使いやすさと性能を向上させることを目指しています。