機器之心發(fā)布

大模型開發(fā)者常面臨一個兩難選擇：要速度，還是省顯存？

通常情況下，想要跑得快，顯存會爆；想省點顯存，計算效率又會被頻繁的同步和流水線氣泡大幅拖垮。原生的 torch.compile 雖然好用，但在面對復雜的跨層優(yōu)化和 FSDP 顯存管理時，依然力不從心。

為了徹底解決這一痛點，Sand.ai 今天正式開源MagiCompiler—— 一款基于 torch.compile 深度優(yōu)化的即插即用、訓推一體編譯框架。

MagiCompiler 徹底突破了傳統(tǒng)局部編譯的界限，實現(xiàn)了推理期整圖捕獲與訓練期 FSDP-Aware 整層編譯。

更重要的是，研發(fā)團隊創(chuàng)新提出Compiler as Manager理念 —— 將編譯器從單純的 “算子優(yōu)化器” 進階為全局管理器。它全面接管了計算調度與顯存的生命周期，以系統(tǒng)級的底層解法，破解算力與顯存墻難題。

• 代碼倉庫：https://github.com/SandAI-org/MagiCompiler

核心技術

打破邊界的全局調度

1. 打破編譯邊界：整圖與整層編譯

傳統(tǒng)編譯常因復雜的 Python 邏輯頻繁觸發(fā) Graph Break。研發(fā)團隊徹底改變了這一點：

• 推理期：捕獲完整的計算圖，最大化 Transformer Block 內的算子融合空間。
• 訓練期：利用 FSDP 在前向 / 反向傳播中 “單層權重全駐留” 的特性，將 Transformer Layer 作為編譯單元。這使得編譯器可以執(zhí)行激進的跨算子融合，大幅減少 Kernel Launch 開銷和 Global Memory 讀寫。

2. 內存魔術：啟發(fā)式重計算（Heuristic Recompute）

在訓練大模型時，開發(fā)者通常需要手動插入 torch.utils.checkpoint 來控制顯存，既繁瑣又難以最優(yōu)。MagiCompiler 引入了智能感知圖分割器：

• 徹底告別手動打點：框架自動分析計算圖，識別并優(yōu)先保留 MatMul、Attention 等計算密集型算子的輸出。
• 極致?lián)革@存：對于顯存密集型算子，自動在反向傳播時進行重計算，從根本上壓縮顯存峰值而不損失吞吐量。

3. 榨干帶寬：JIT 極致 Offload 調度

針對顯存瓶頸，研發(fā)團隊實現(xiàn)了一套極其優(yōu)雅的權衡調度引擎：

• 性價比常駐：基于 Profiling 數(shù)據(jù)，將最劃算的權重貪心地常駐在有限的 GPU 顯存中。
• JIT 最晚預取：調度器逆向推導精確的預取時間表，卡在計算前的 “最后一刻” 完成權重拉取，確保 GPU 不囤積多余權重，徹底消除流水線氣泡。

MagiCompiler Overview

性能實測

真正免費的性能午餐

憑借底層的全局調度，MagiCompiler 交出了亮眼的答卷：

• 訓練端表現(xiàn)：在極短時間內，提供高吞吐的保底方案。無需耗時數(shù)周死磕 Kernel 或手工魔改底層邏輯，開箱即可解決 Baseline 的 CPU 調度與算子碎片化難題，直接帶來 44.7% 提速與 6.2% 顯存下降，且精度完全對齊。

MagiCompiler v.s. baseline

• 推理端表現(xiàn)：在多模態(tài)視頻生成場景下，MagiCompiler 展現(xiàn)了極其扎實的硬件泛化能力
• H100：比最好更好
• 在單機 NVIDIA H100 上，面對主流視頻生成模型，MagiCompiler 比目前的領跑方案（如 LightX2V）還要快9%~26%

H100 性能測評

• RTX 5090：顯存受限，近乎實時
• 即便在顯存有限的 5090 上，通過 JIT Offload 調度，MagiCompiler 也讓 daVinci-MagiHuman 這種超大模型跑出了近乎實時的速度。

5090 daVinci-MagiHuman 性能指標

極簡體驗

一行代碼，即插即用

強悍的底層性能并不意味著復雜的接入成本。秉持對開發(fā)者友好的設計理念，MagiCompiler 只需兩個裝飾器即可完成接入。

• 基礎編譯增強
• 無需修改模型源碼，magi_compile 一鍵裝飾 TransformerBlock：

• 自定義算子注冊
• 對于 FlashAttention 或 MoE 等定制化算子，輕松注冊并無縫融入重計算策略：

此外，我們內置了強大的自省工具鏈：開啟環(huán)境變量，所有隱式的編譯產物（反編譯字節(jié)碼、Kernel 代碼、Guard 條件等）均會被持久化為人類可讀的 Python 文件與圖表，讓編譯器 Debug 變得簡單直觀。

結語與未來展望

MagiCompiler 正在打破傳統(tǒng)編譯器的邊界。它不僅讓我們看到了 torch.compile 邁向全局調度的巨大潛力，更為大模型與多模態(tài)架構的規(guī)?；涞靥峁┝嘶A設施。

目前，MagiCompiler 已全面開源。Sand.ai 將持續(xù)降低大模型底層的開發(fā)門檻，為 AI 社區(qū)持續(xù)做出貢獻。

了解更多信息，歡迎訪問 Sand.ai 官網(wǎng)：https://sand.ai