Axe：機械学習コンパイラのためのシンプルで統一されたレイアウト抽象化

TL;DR（結論）

現代の深層学習ワークロードのスケールアップに伴い、デバイスメッシュやメモリ階層、異種アクセラレータ間でのデータと計算の調整が不可欠となっていますが、本論文は論理的なテンソル座標を「名前付き軸」を介して多軸物理空間にマッピングする、ハードウェアを意識した抽象化「Axe Layout」を提案しています。 Axeは、デバイス間の分散処理からデバイス内のスレッドレベルのレイアウトまで、タイリング、シャーディング、レプリケーション、オフセットを統一的に表現することを可能にし、単一のカーネル内でスレッドローカルな制御と集合演算子を組み合わせた多粒度かつ分散対応のDSLおよびコンパイラを実現しています。 検証の結果、Axeを用いた手法は最新のGPU（NVIDIA B200等）やAIアクセラレータ（AWS Trainium-1等）において、手動で調整されたカーネルに近い性能を達成し、MoEレイヤーや分散GEMM、マルチヘッドアテンションなどの演算で既存のライブラリを上回る最大1.44倍の高速化を記録しました。

なぜこの問題か

大規模言語モデル（LLM）に代表される現代の深層学習モデルは、検索エンジンやコーディング・アシスタント、クリエイティブ・ツールなどの基盤となっており、モデルの規模やデータ量、推論サーバーのフリートが拡大するにつれて、実行時の最適化がシステム上の中心的な課題となっています。深層学習システムを最適化するためには、ソフトウェアとハードウェアのスタックにおける複数のレイヤーで発生する課題に対処しなければなりません。まず、分散実行の観点では、巨大なモデルを複数のデバイスやマシンにわたって実行する必要があり、フレームワークやコンパイラはデバイスメッシュ上でのデータのシャーディングやレプリケーションの選択を明示的に行い、通信と計算のオーバーラップを最適化する必要があります。 次に、デバイスレベルのメモリとスレッドの階層構造が挙げられます。GPUやAIアクセラレータは、グリッド、ブロック、ワープ、レーンといった入れ子状の並列構造と複雑なメモリ階層を持っており、カーネルライブラリはこれらのメモリ範囲にわたってデータがどのようにタイリングされるかを慎重に調整しなければなりません。…

核心：何を提案したのか

本論文が提案する「Axe Layout」は、分散デバイス階層と異種ハードウェア設定の両方にわたって、計算とデータのマッピングを統一するシンプルかつ効果的な抽象化です。Axeの最大の特徴は、スレッド軸、アクセラレータのメモリバンク、分散ワーカーなどのハードウェア構成要素を明示的に表現するために「名前付き軸（named axes）」を導入した点にあります。この名前付き軸に基づき、Axeはタイルの論理座標を、GPUデバイス、スレッド、メモリにわたる多軸物理空間へどのようにマッピングするかを定義します。…

続きはログイン/プランで閲覧できます。