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 邊界的 kernel 融合範圍。
  • 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 上對影片生成模型展示 9%~26% 的效能提升，超越現有最先進解決方案。
  • Introspection toolkit：內含 magi_depyf 用於除錯編譯時間線、位元碼流程與子圖切分。
  • Hardware-aware：支援高階 GPU（H100）與消費級 GPU（RTX 5090），透過 JIT 卸載實現硬體感知。