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

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

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

為了徹底解決這一痛點(diǎn)，Sand.ai 今天正式開(kāi)源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)階為全局管理器。它全面接管了計(jì)算調(diào)度與顯存的生命周期，以系統(tǒng)級(jí)的底層解法，破解算力與顯存墻難題。

• 代碼倉(cāng)庫(kù)：https://github.com/SandAI-org/MagiCompiler

核心技術(shù)

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

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

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

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

2. 內(nèi)存魔術(shù)：?jiǎn)l(fā)式重計(jì)算（Heuristic Recompute）

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

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

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

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

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

MagiCompiler Overview

性能實(shí)測(cè)

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

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

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

MagiCompiler v.s. baseline

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

H100 性能測(cè)評(píng)

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

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

極簡(jiǎn)體驗(yàn)

一行代碼，即插即用

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

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

• 自定義算子注冊(cè)
• 對(duì)于 FlashAttention 或 MoE 等定制化算子，輕松注冊(cè)并無(wú)縫融入重計(jì)算策略：

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

結(jié)語(yǔ)與未來(lái)展望

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

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

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