Sand.ai開源MagiCompiler：突破局部編譯界限，定義訓(xùn)推性能上限

機(jī)器之心發(fā)布

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

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

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

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

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

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

核心技術(shù)

打破邊界的全局調(diào)度

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

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

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

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

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

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

3. 榨干帶寬：JIT 極致 Offload 調(diào)度

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

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

MagiCompiler Overview

性能實(shí)測

真正免費(fèi)的性能午餐

憑借底層的全局調(diào)度，MagiCompiler 交出了亮眼的答卷：

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

MagiCompiler v.s. baseline

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

H100 性能測評

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

5090 daVinci-MagiHuman 性能指標(biāo)

極簡體驗(yàn)

一行代碼，即插即用

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

• 基礎(chǔ)編譯增強(qiáng)
• 無需修改模型源碼，magi_compile 一鍵裝飾 TransformerBlock：

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

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

結(jié)語與未來展望

MagiCompiler 正在打破傳統(tǒng)編譯器的邊界。它不僅讓我們看到了 torch.compile 邁向全局調(diào)度的巨大潛力，更為大模型與多模態(tài)架構(gòu)的規(guī)模化落地提供了基礎(chǔ)設(shè)施。

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

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