オープンソースのニューラルネットワーク向けプログラミング言語「Triton」をOpenAIが公開

NVIDIAが開発・提供するGPU向けの汎用並列コンピューティングプラットフォームである「CUDA」を超える生産性と高速コード記述が可能になるようなオープンソースのニューラルネットワーク向けプログラミング言語「Triton」が公開されました。非常に効率的なカスタムディープラーニングプリミティブを作成するための言語コンパイラとなっており、GitHub上で開発リポジトリが公開されています。

Introducing Triton: Open-Source GPU Programming for Neural Networks
https://www.openai.com/blog/triton/

OpenAI debuts Python-based Triton for GPU-powered machine learning | InfoWorld
https://www.infoworld.com/article/3627243/openai-debuts-python-based-triton-for-gpu-powered-machine-learning.html

OpenAI proposes open-source Triton language as an alternative to Nvidia’s CUDA | ZDNet
https://www.zdnet.com/article/openai-proposes-triton-language-as-an-alternative-to-nvidias-cuda/

人工知能(AI)について研究する非営利団体のOpenAIが、Pythonをベースとしたオープンソースのニューラルネットワーク向けプログラミング言語「Triton」のバージョン1.0をリリースしました。

OpenAIがCUDAに代わるプログラミング言語としてTritonを開発した理由はとても単純で、「CUDAを含むNVIDIAのグラフィックスプロセッシングユニットはプログラミングが難しすぎるため」です。具体的には、GPU用のネイティブカーネルや関数を作成する場合、マルチコアGPUのメモリ階層全体で割り当てデータと命令を移動しなければならず、プログラミングが非常に複雑になります。

OpenAIによると、Tritonを使えばCUDAを扱ったことがない研究者でも、GPUコーダーが作成したものと同等レベルの非常に効率的なGPUコードを記述可能となります。例えば、25行未満のコードでcuBLASのパフォーマンスに匹敵するFP16乗算カーネルを作成することが可能であり、「これはほとんどのGPUプログラマーが実現できないレベルのパフォーマンス」とOpenAIは説明しました。

テクノロジーメディアのZDNetがOpenAIに問い合わせたところ、Philippe Tillet氏から「我々の目標はディープラーニングのためのCUDAの実行可能な代替手段を作り出すことです」というコメントが返ってきたそうです。加えて、Tillet氏はTritonのターゲットは「優れたソフトウェアエンジニアリングスキルを持っているにもかかわらず、GPUプログラミングには慣れていないという機械学習エンジニアや研究者向け」と語っています。

OpenAIの研究者たちはすでにTritonを利用しており、PyTorchでのプログラミングと比べて最大2倍も効率的なカーネルの作成に成功しているとアピールしています。

Tritonを使用する場合、開発者は専用ライブラリを使用してPythonでコードを記述し、GPUで実行するためにJITコンパイルします。これにより、機械学習ソリューションを開発するためのPythonエコシステムの残りの部分との統合が可能になるとのこと。

テクノロジーメディアのInfoWorldは、「TritonのライブラリはNumPyを彷彿とさせる一連のプリミティブを提供しています」と記し、NumPyとの類似性を指摘。具体的には、行列演算や何らかの基準に従って配列の縮小を実行する関数を提供する点などを挙げています。また、プリミティブを独自のコードと統合し、GPUで実行するようにコンパイルするという点では、「Numbaにも似ている」と記しました。

Tritonは2019年に発表された論文をベースにしたものですが、まだまだスタートしたばかりのプロジェクトです。記事作成時点ではLinuxでのみ利用可能となっており、用意されているドキュメントも最小限であるため、InfoWorldは「早期に採用する開発者はソースと例を綿密に調べる必要があるかもしれません」と警告しています。

Tritonの開発リポジトリには以下からアクセス可能です。

GitHub – openai/triton: Development repository for the Triton language and compiler
https://github.com/openai/triton