SandAI-org/MagiCompiler

A plug-and-play compiler that delivers free-lunch optimizations for both inference and training.

What it solves

MagiCompiler は、大規模Transformer類似アーキテクチャの性能を制限するメモリ壁と演算子オーバーヘッドに対処します。単なるローカル演算子最適化に留まらず、トレーニングとマルチモーダル推論の両方に対するシステムレベルの最適化を提供します。

How it works

torch.compile の上に構築された MagiCompiler は、実行ライフサイクルのグローバルマネージャとして機能します。サブグラフのディスパッチ管理、データフロー（オフロードやプリフェッチを含む）のオーケストレーション、メモリ割り当ての制御を行います。主な技術メカニズムは以下の通りです。

  • Full-graph capture：推論時にTransformer境界を越えたカーネル融合範囲を最大化します。
  • FSDP-aware layer-wise compilation：トレーニング時に積極的なクロスオペレーション融合を可能にし、分散パラメータシャーディングを透過的に保ちます。
  • Selective offloading：ホスト→デバイス (H2D) 転送と計算を重ね合わせ、メモリ制約環境でのパイプラインバブルを排除します。
  • Heuristic activation recomputation：計算バウンドかメモリバウンドかに基づき、保存すべき操作と再計算すべき操作を自動判定し、ピークメモリを削減します。

Who it’s for

大規模Transformerモデルを扱う開発者や研究者、特にマルチモーダルフレームワークを実装したり、FSDP などの分散戦略で大規模モデルをトレーニングする方を対象としています。

Highlights

  • Plug-and-play integration：シンプルなデコレータ（@magi_compile と @magi_register_custom_op）を使用するだけで、複雑なモデルリファクタリングなしに速度向上が得られます。
  • Significant acceleration：NVIDIA H100 GPU 上で動画生成モデルに対し、SOTA ソリューションに比べ 9%〜26% の性能向上を実証しています。
  • Introspection toolkit：magi_depyf を含み、コンパイルタイムライン、バイトコードフロー、サブグラフ分割のデバッグが可能です。
  • Hardware-aware：JIT オフロードを通じて、ハイエンド GPU（H100）とコンシューマ GPU（RTX 5090）をサポートします。