![]()
機器之心發布
大模型開發者常面臨一個兩難選擇:要速度,還是省顯存?
通常情況下,想要跑得快,顯存會爆;想省點顯存,計算效率又會被頻繁的同步和流水線氣泡大幅拖垮。原生的 torch.compile 雖然好用,但在面對復雜的跨層優化和 FSDP 顯存管理時,依然力不從心。
為了徹底解決這一痛點,Sand.ai 今天正式開源MagiCompiler—— 一款基于 torch.compile 深度優化的即插即用、訓推一體編譯框架。
MagiCompiler 徹底突破了傳統局部編譯的界限,實現了推理期整圖捕獲與訓練期 FSDP-Aware 整層編譯。
更重要的是,研發團隊創新提出Compiler as Manager理念 —— 將編譯器從單純的 “算子優化器” 進階為全局管理器。它全面接管了計算調度與顯存的生命周期,以系統級的底層解法,破解算力與顯存墻難題。
![]()
- 代碼倉庫:https://github.com/SandAI-org/MagiCompiler
核心技術
打破邊界的全局調度
1. 打破編譯邊界:整圖與整層編譯
傳統編譯常因復雜的 Python 邏輯頻繁觸發 Graph Break。研發團隊徹底改變了這一點:
- 推理期:捕獲完整的計算圖,最大化 Transformer Block 內的算子融合空間。
- 訓練期:利用 FSDP 在前向 / 反向傳播中 “單層權重全駐留” 的特性,將 Transformer Layer 作為編譯單元。這使得編譯器可以執行激進的跨算子融合,大幅減少 Kernel Launch 開銷和 Global Memory 讀寫。
2. 內存魔術:啟發式重計算(Heuristic Recompute)
在訓練大模型時,開發者通常需要手動插入 torch.utils.checkpoint 來控制顯存,既繁瑣又難以最優。MagiCompiler 引入了智能感知圖分割器:
- 徹底告別手動打點:框架自動分析計算圖,識別并優先保留 MatMul、Attention 等計算密集型算子的輸出。
- 極致摳顯存:對于顯存密集型算子,自動在反向傳播時進行重計算,從根本上壓縮顯存峰值而不損失吞吐量。
3. 榨干帶寬:JIT 極致 Offload 調度
針對顯存瓶頸,研發團隊實現了一套極其優雅的權衡調度引擎:
- 性價比常駐:基于 Profiling 數據,將最劃算的權重貪心地常駐在有限的 GPU 顯存中。
- JIT 最晚預取:調度器逆向推導精確的預取時間表,卡在計算前的 “最后一刻” 完成權重拉取,確保 GPU 不囤積多余權重,徹底消除流水線氣泡。
![]()
MagiCompiler Overview
性能實測
真正免費的性能午餐
憑借底層的全局調度,MagiCompiler 交出了亮眼的答卷:
- 訓練端表現:在極短時間內,提供高吞吐的保底方案。無需耗時數周死磕 Kernel 或手工魔改底層邏輯,開箱即可解決 Baseline 的 CPU 調度與算子碎片化難題,直接帶來 44.7% 提速與 6.2% 顯存下降,且精度完全對齊。
![]()
MagiCompiler v.s. baseline
- 推理端表現:在多模態視頻生成場景下,MagiCompiler 展現了極其扎實的硬件泛化能力
- H100:比最好更好
- 在單機 NVIDIA H100 上,面對主流視頻生成模型,MagiCompiler 比目前的領跑方案(如 LightX2V)還要快9%~26%
![]()
H100 性能測評
- RTX 5090:顯存受限,近乎實時
- 即便在顯存有限的 5090 上,通過 JIT Offload 調度,MagiCompiler 也讓 daVinci-MagiHuman 這種超大模型跑出了近乎實時的速度。
![]()
5090 daVinci-MagiHuman 性能指標
極簡體驗
一行代碼,即插即用
強悍的底層性能并不意味著復雜的接入成本。秉持對開發者友好的設計理念,MagiCompiler 只需兩個裝飾器即可完成接入。
- 基礎編譯增強
- 無需修改模型源碼,magi_compile 一鍵裝飾 TransformerBlock:
![]()
- 自定義算子注冊
- 對于 FlashAttention 或 MoE 等定制化算子,輕松注冊并無縫融入重計算策略:
![]()
此外,我們內置了強大的自省工具鏈:開啟環境變量,所有隱式的編譯產物(反編譯字節碼、Kernel 代碼、Guard 條件等)均會被持久化為人類可讀的 Python 文件與圖表,讓編譯器 Debug 變得簡單直觀。
結語與未來展望
MagiCompiler 正在打破傳統編譯器的邊界。它不僅讓我們看到了 torch.compile 邁向全局調度的巨大潛力,更為大模型與多模態架構的規模化落地提供了基礎設施。
目前,MagiCompiler 已全面開源。Sand.ai 將持續降低大模型底層的開發門檻,為 AI 社區持續做出貢獻。
了解更多信息,歡迎訪問 Sand.ai 官網:https://sand.ai
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.