CODA: GEMM-Epilogue FusionによるTransformerトレーニングの最適化
Transformerのトレーニングシステムは、本質的に密な線形代数、特にGeneral Matrix Multiplication (GEMM) に基づいて構築されています。しかし、トレーニングスタックが高度に最適化されるにつれて、エンドツーエンドの実行時間の大部分が、正規化、活性化、残差更新、およびリダクションといった「周辺」オペレータに費やされるようになっています。これらのメモリバウンドな操作は、演算をほとんど行わずに膨大な中間テンソルをグローバルメモリ経由で繰り返し移動させるため、ハードウェアの全体的な効率を制限するボトルネックとなります。
これに対処するため、研究者たちは、これらの計算を「GEMM-plus-epilogue」プログラムとして表現するように設計されたGPUカーネル抽象化であるCODAを導入しました。Transformerブロックの実行方法を再考することで、CODAは不要なデータ移動を排除し、オンチップ利用率を最大化することを目指しています。
問題点: メモリバウンドのボトルネック
標準的なTransformerブロックでは、主要な処理はGEMMオペレータによって行われます。しかし、これらは通常、一連の要素ごとの操作やリダクション操作に続きます。従来のフレームワークレベルの実行では、各オペレータは個別のカーネルとして起動されます。これは、GEMMの出力がグローバルメモリに書き込まれ、それが直ちに次のカーネル（例：LayerNormやReLU活性化）によって読み戻されることを意味します。
このグローバルメモリへの書き込みと読み出しのサイクルはコストがかかります。計算能力が増力するにつれて、演算スループットとメモリ帯域幅の差が広がり、これらの「メモリバウンド」なオペレータがトレーニングの総時間の無視できない割合を占めるようになります。
CODAのアプローチ: GEMM-Epilogue Fusion
CODAは、多くのTransformerオペレータが、GEMMの出力タイルがグローバルメモリに書き込まれる前に、チップ上（共有メモリまたはレジスタ）に留まっている間に実行されるように代数的に再パラメータ化できるという観察に基づいています。
仕組み
CODAは、行列計算の高度に最適化されたコアであるGEMMメインループを固定し、合成可能な一連の小さなエピローグ・プリミティブを公開します。これらのプリミティブにより、開発者は以下のような操作を定義できます：
- Scaling: 出力を定数またはベクトルで乗算する。
- Reductions: 特定の次元にわたって合計または平均を算出する。
- Pairwise Transformations: 活性化関数やその他の要素ごとの関数を適用する。
- Accumulation: 結果を残差接続に追加する。
インターフェースをこれらのプリミティブに制限することで、CODAは専門家が記述したGEMMの性能特性を維持しつつ、標準的なTransformerブロックの順伝播および逆伝播の両方における、ほぼすべての非アテンション計算をカバーするのに十分な表現力を提供します。
LLMによるコード生成への影響
この研究における最も重要な発見の一つは、人間が作成したCODAカーネルとLLMが作成したCODAカーネルの両方が高い性能を達成していることです。これは、低レベルハードウェア最適化へのアプローチが変化することを示唆しています。
LLMは、レジスタ圧迫の管理や複雑なメモリタイリングといった低レベルハードウェア最適化の複雑な詳細には苦戦することが多いですが、高レベルの合成には優れています。制限された合成可能なAPIを提供することで、CODAはハードウェアの複雑さを効果的に「サンドボックス化」します。LLMはGEMMメインループを記述する必要はなく、専門家が記述したエピローグ・ブロックを結合するだけで済みます。
Hacker Newsのコミュニティ観察者によって次のように指摘されています：
"設計されたコンパイラ抽象化が、制限された、合成可能なAPIを持つことで、LLMが専門家が記述したブロックを簡単に結合できるようにすることは、賢明な動きです。エージェント的開発へと移行するにつれて、これが最終的にコード生成の標準になるのではないかと推測しています。"
結論
CODAは、フレームワークレベルの生産性とハードウェアレベルの効率を組み合わせるための実用的な道筋を示しています。 TransformerのメモリバウンドなオペレータをGEMMのエピローグとして扱うことで、CODAはグローバルメモリへのラウンドトリップを大幅に削減し、次世代の大規模モデルのトレーニングを最適化する方法を変革します。