OpenAIは最新の言語Tritonをリリースした。このオープンソースプログラミング言語によって、研究者はAI処理用に非常に効率的なGPUコードを記述できる。これは、Python互換で、ユーザはわずか25行のコードでエキスパート品質の結果を得ることができる。
TritonはPythonをベースとして使っている。開発者は、Tritonのライブラリを使ってPythonでコードを記述する。その後、GPUで実行するためにJITコンパイルされる。これにより、現在機械学習ソリューションを開発するための最大の目的となっているPythonエコシステムとの統合が実現される。
NumPyを連想させるTritonのライブラリは、たとえば、さまざまな行列演算や、何らかの基準に従った配列の縮小を実行する関数を提供する。ユーザは、@triton.jitデコレータを追加することで、これらのプリミティブを自身のコードに結合し、GPUで実行するようにコンパイルできる。この意味で、TritonはNumbaにも似ている。Numbaは、数値処理がメインとなるPythonコードをマシンネイティブアセンブリにJITコンパイルして、高速化を実現するプロジェクトだ。
稼働中のTritonの簡単な例として、ベクトル加算カーネルと融合ソフトマックス演算がある。後者の例では、完全にGPUメモリで実行できる操作について、ネイティブのPyTorch融合ソフトマックスよりも何倍も高速に実行できる。
OpenAIによると、Tritonにより、汎用ライブラリよりもはるかに高速処理できる特殊なカーネルの開発がシンプルになる。そのコンパイラによってコードがシンプルになり、自動的に最適化、並列化が行われる。最近のNvidia GPUで実行できるコードに変換される。
最新のGPUのアーキテクチャは、DRAM、SRAM、ALUの3つの主要コンポーネントに分類できる。CUDAコードを最適化するときは、それぞれを考慮する必要がある。GPUプログラミングに伴う課題があることを見逃すことはできない。例えば、現在のメモリインターフェイスでより大きなバス幅を活用するために、DRAMからのメモリ転送がコアレッシングすることを確認しなければならない。データは、取得時に他の共有メモリバンクと競合しないように、再利用する前に手動でSRAMSに格納する必要がある。
OpenAIは、これにより多くの労力をかけずに最高のハードウェアパフォーマンスに到達できるようになり、これまで以上に複雑なワークフローの作成が簡単になると言っている。キャッシュ形式のデータがプロセスに対して大きすぎる場合には常に、Tritonは必要なメモリのほぼ2倍の量を転送し、他のデータよりも高速化する。Tritonでのコーディングは、PyTorchのCUDAカーネルに比べて簡単で信頼性が高く、プロセスを理解して維持するのも簡単となる。
Tritonはコミュニティ主導のプロジェクトになることを目指している。GitHubでリポジトリをフォークできる。