ChatGPTやStable Diffusionを生み出した「機械学習ソフトウェア」の10年間の流れを専門家が解説

近年は対話型AIのChatGPTや画像生成AIのStable Diffusionなど、さまざまなAIが登場して世間の注目を集めています。これらのAIを開発するために使用される機械学習ソフトウェアの過去10年の変遷について、半導体分野のアナリストであるディラン・パテル氏が解説しています。

How Nvidia’s CUDA Monopoly In Machine Learning Is Breaking - OpenAI Triton And PyTorch 2.0
https://www.semianalysis.com/p/nvidiaopenaitritonpytorch

◆TensorFlow対PyTorch
数年前の機械学習フレームワークエコシステムではGoogleのTensorFlowがフロントランナーであり、機械学習業界を支配する準備ができているように見えたとパテル氏は指摘。また、Googleは機械学習に特化したプロセッサであるTPUを開発・展開した点も、先行者としてのメリットになっていたとのこと。

しかし、Googleは先行者としての優位性を生かして機械学習業界を支配することができず、Pythonの機械学習ライブラリであるPyTorchが存在感を強めています。PyTorchはFacebook(現Meta)の人工知能研究チームによって開発され、記事作成時点では画像生成AIのStable Diffusionでも使用されています。PyTorchがTensorFlowに対して優位に立った理由についてパテル氏は、TensorFlowは「グラフモード」というスクリプト実行方法を採用しているのに対し、PyTorchは「Eagerモード」という方法を採用している点が大きかったと指摘しています。

Eagerモードは他のPythonコードと同様に行ごとに各操作を実行するため、中間操作の結果を確認してモデルの動作を確認しやすいという利点があります。これにより、コードの理解やデバッグがしやすいとのこと。一方でグラフモードは、実行操作を表す計算グラフを定義するフェーズと、最適化されたバージョンの計算を実行するフェーズの2つしかありません。これが原因で、計算が終了するまで何が起こっているか確認できないため、コードの理解やデバッグが難しいそうです。

パテル氏は2つのモードの違いについて、グラフモードはC＋＋のような事前に一括でコードを変換するコンパイル型言語に似ており、EagerモードはPythonのように中間表現を逐次解釈しながら実行するインタプリタ言語に似ていると述べています。なお、その後TensorFlowにはデフォルトでEagerモードが搭載されるようになりましたが、研究コミュニティや大手テクノロジー企業では主にPyTorchが定着しているとのこと。

Googleは独自のソフトウェアスタックとハードウェアを優先し、PyTorchとGPUを使用していないことから、機械学習コミュニティ内ではやや孤立しています。また、OpenAIなどの大規模言語モデルやその他のAIスタートアップにより検索や自然言語処理の優位性が脅かされているとも言われていますが、依然としてGoogleは機械学習モデルの最前線にいるとのことです。

by ranveer cool

◆機械学習トレーニングにおける帯域幅の問題
パテル氏は、機械学習トレーニングを簡単な形式に分けると、「コンピューティング(FLOPS)」と「メモリ(帯域幅)」の2つがあるとしています。このうち、FLOPSは各レイヤー内における計算実行速度を制限する要素であり、帯域幅はデータが計算を実行するコンピューティングリソースに到達するまでの速度を制限しています。かつて、機械学習トレーニングにかかる時間はコンピューティング能力によって制限されていましたが、NVIDIAが高性能のGPUを開発した結果この点は解決され、記事作成時点では帯域幅が問題となっているとのこと。

AI開発の進展によってモデルのサイズが急増し続けているため、数百GB～数十TBものメモリを必要とする大規模な言語モデルでは、帯域幅が大きな問題となります。最もデータの読み込みが速いのはオンチップのSRAMメモリですが、人工知能企業のCerebrasが開発した20cm×22cmの巨大チップ「Wafer Scale Engine 2」でさえ40GBのオンチップSRAMが限界であり、価格は数億円に上るそうです。

2.6兆個ものトランジスタを搭載した大型チップ「Wafer Scale Engine 2」をCerebrasが開発 - GIGAZINE

オンチップSRAMの次に読み込み速度が速いのはDRAMであり、オンチップSRAMと比較すると桁違いにレイテンシが増すものの、コスト面では100分の1程度にまで下がるとのこと。それでも2012年以降はDRAMのコスト改善が進んでおらず、2022年の時点ではサーバーに占める総コストのうち50％をDRAMが占めているとパテル氏は指摘しています。

NVIDIAが開発するA100 GPUと次世代のH100 GPUでは、FLOPSは6倍以上に増加するものの、メモリ帯域幅は1.65倍にしか増加しません。大規模な言語モデルのトレーニングにおいては、FLOPSの使用率が60％でも「非常に使用率が高い」と見なされるそうで、H100 GPUでは使用率のさらなる低下が懸念されています。

◆PyTorchの演算子の増加と対策
GoogleのTensorFlowに対して優位に立っているPyTorchですが、Eagerモードは各操作においてメモリからデータを読み取って計算し、次の計算が処理される前に結果をメモリに送信しています。これにより、大規模な最適化が行われない場合はメモリ帯域幅が大幅に増加してしまうとのこと。

そのため、Eagerモードで実行されるモデルではメモリの読み取り／書き込みを最小限に抑えるため、複数の関数を1つのパスにまとめる演算子融合が行われています。これによってメモリ帯域幅やメモリサイズのコストが削減できますが、PyTorchの演算子は過去数年で2000以上に膨れ上がったため、今度は開発者が適切な演算子を選択することが困難になっているそうです。

そんな中、PyTorchはMetaから独立してLinux Foundationに移行し、2022年12月には「PyTorch 2.0」を公開しました。

画像生成AIのStable Diffusionにも使われるAIフレームワーク「PyTorch」がMetaから独立してLinux Foundationに移行 - GIGAZINE

PyTorch 2.0には多くの変更が加えられていますが、大きな違いがグラフ実行モデルをサポートするコンパイル済みソリューションが搭載されている点であり、これによりハードウェアリソースの最適化がはるかに簡単になるとのこと。また、演算子の数を250以下に減らして実装を簡単にした「PrimTorch」や、メモリに書き込む必要がある中間操作と融合できる中間操作を決定する「TorchDynamo」、コンパイラ作業の負荷を軽減する「TorchInductor」などの新機能も搭載されており、メモリ帯域幅と容量の要件が大幅に削減されているそうです。PyTorch 2.0は2023年3月に正式版の公開が予定されています。

◆OpenAIのAI向け言語「Triton」
AI研究団体のOpenAIは2021年7月、オープンソースのニューラルネットワーク向けプログラミング言語「Triton」を公開しました。Tritonは、NVIDIAが開発・提供するGPU向けの汎用並列コンピューティングプラットフォームである「CUDA」を超える生産性と高速コード記述を可能にし、GPUプログラミングに慣れていない機械学習エンジニアや研究者も利用しやすいとされています。

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

パテル氏は、CUDAはアクセラレーテッドコンピューティングを専門とする人々に使われているものの、機械学習研究者やデータサイエンティストにはあまり使われていないと指摘。その理由は、CUDAを効率的に使用するのは難しく、ハードウェアアーキテクチャを深く理解する必要があるため、開発プロセスが遅れてしまう可能性があるからだとのこと。

Tritonはそのギャップを埋める存在となり得るため、近い将来にTritonがNVIDA以外のGPUを公式にサポートするようになれば、NVIDIAの牙城を崩すことになりかねないとパテル氏は見ています。また、NVIDIAが抱える問題については「NVIDIAの巨大なソフトウェア組織は、機械学習ハードウェアとソフトウェアにおける彼らの巨大な優位性を活用し、機械学習のためのデフォルトコンパイラになるための先見性を欠いていました。NVIDIAがユーザービリティを重視していなかったため、OpenAIやMetaといった外部の人間が、他のハードウェアに移植可能なソフトウェアスタックを作ることができたのです」と述べました。