“零”程式碼改動，靜態編譯讓太乙Stable Diffusion推理速度翻倍

 

作者｜梁德澎
 

AI 作圖領域的工具一直不盡如人意，直到去年 8 月 Stable Diffusion 開源，成為AI 影象生成領域無可爭辯的劃時代模型。

 

為了提升其推理效率，OneFlow 首度將 Stable Diffusion 模型加速至“一秒出圖”時代，極大提升了文生圖的速度，在AIGC領域引發巨大反響，並得到了 Stability.ai 官方的支援。至今，OneFlow 還在不斷重新整理 SOTA 紀錄。

 

不過，由於目前大部分團隊主要是基於翻譯 API + 英文 Stable Diffusion 模型進行開發，所以在使用中文獨特的敘事和表達時，英文版模型就很難給出正確匹配的圖片內容，這對部分國內使用者來說不太方便。

為了解決這一問題，國內的IDEA 研究院認知計算與自然語言研究中心（IDEA CCNL）也開源了第一箇中文版本的“太乙 Stable Diffusion”，基於0.2億篩選過的中文圖文對訓練。上個月，太乙 Stable Diffusion 在 HuggingFace 上有近 15 萬下載量，是下載量最大的中文 Stable Diffusion。

 

近期，OneFlow 團隊為太乙 Stable Diffusion 適配了 OneFlow 後端，大大提升了推理效能，也可以做到一秒出圖。不少開發者好奇OneFlow使用了哪些優化“祕笈”，後文將進行簡要解讀。

 

歡迎Star、執行 OneFlow 版太乙 Stable Diffusion：
 

https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion#without-docker

 

1
對比 PyTorch，OneFlow 將“太乙 Stable Diffusion”推理速度提升1倍以上
 

下面的圖表分別展示了在 A100 (PCIe 40GB / SXM 80GB)，V100 ( SXM2 32GB )， RTX 2080，RTX 3080 Ti，RTX 3090， 和 T4 不同型別的 GPU 硬體上分別使用 PyTorch, 和 OneFlow對 太乙 Stable Diffusion 進行推理的效能表現。

 

 

 

可以看到，對於 A100 顯示卡，無論是 PCIe 40GB 的配置還是 SXM 80GB 的配置，OneFlow 的效能相比 PyTorch 能提升 1 倍以上，推理速度達到了 50it/s 以上，生成一張圖片所需要的時間在 1 秒以內。

 

其他硬體資料：

 

注：3090上的AIT資料由 IDEA 研究院提供

 

 

 

 

 

 

綜上，在各種硬體的對比中，對比 PyTorch， OneFlow 能將太乙 Stable Diffusion 的推理效能提升 1 倍多。

 

2

生成圖片展示

 

滔滔江水, 連綿不絕, 唯美, 插畫

 

 

長城, 清晨, 朦朧, 唯美, 插畫

 

 

夢迴江南，中國古代小鎮，唯美，插畫

 

 

中國的未來城市, 科幻插畫

 

 

古代建築, 白雪紛飛

 

 

螺螄粉

 

注：上述圖片均基於 OneFlow 版太乙 Stable Diffusion 生成
 

3

無縫相容 PyTorch 生態

 

想體驗 OneFlow 版的太乙 Stable Diffusion？只需要修改兩行程式碼：

 

 

之所以能這麼輕鬆遷移模型，是因為 OneFlow Stable Diffusion 有兩個出色的特性：

 

1. OneFlowStableDiffusionPipeline.from_pretrained 能夠直接使用 PyTorch 權重。

2. OneFlow 本身的 API 和 PyTorch 對齊，因此 import oneflow as torch 之後，torch.autocast、torch.float16 等表示式完全不需要修改。

 

上述特性使得 OneFlow 相容了 PyTorch 的生態，這不僅在 OneFlow 對 太乙 Stable Diffusion 的遷移中發揮了作用，也大大加速了 OneFlow 使用者遷移其它許多模型，比如在和 torchvision 對標的 flowvision 中，許多模型只需通過在 torchvision 模型檔案中加入 import oneflow as torch 即可得到。

 

此外，OneFlow 還提供全域性 “mock torch” 功能，在命令列執行 eval $(oneflow-mock-torch) 就可以讓接下來執行的所有 Python 腳本里的 import torch 都自動指向 oneflow。
 

4

動靜一體的程式設計體驗

 

深度學習演算法原型開發階段需要快速修改和除錯，動態圖執行（Eager mode, define by run）最優。但在部署階段，模型已經固定下來，計算效率變得更重要，靜態圖執行（Lazy mode，define and run）可以藉助編譯器做靜態優化來獲得更好的效能。因此，推理階段主要使用靜態圖模式。

 

最近，PyTorch 升級到2.0引入了compile()這個API，可以把一個模型或一個Module從動態圖執行變成靜態圖執行。OneFlow裡也有一個類似的機制，不過介面名是nn.Graph()，它可以把傳入Module轉成靜態圖執行模式。

 

不僅如此，OneFlow的nn.Graph模式基於MLIR實現了一系列計算圖的圖層優化 ，譬如記憶體佈局、運算元融合等。
 

這不僅使得計算圖表示的深度學習模型可以在各種硬體上達到最高效能，更重要的是，使得深度學習框架匯入的計算圖更方便地在不同硬體之間實現遷移，有助於克服國產硬體軟體生態薄弱的問題。未來，我們將釋出更多內容來揭示OneFlow深度學習編譯器的設計和實現。

 

歡迎Star、執行 OneFlow 版太乙 Stable Diffusion：
 

https://github.com/Oneflow-Inc/diffusers/wiki/How-to-Run-OneFlow-Stable-Diffusion#without-docker
 

OneFlow 地址：https://github.com/Oneflow-Inc/oneflow/
 

其他人都在看

• 35張圖，直觀理解Stable Diffusion

• ChatGPT的一小步，NLP正規化轉變的一大步

• 2023年AI十大展望：GPT-4領銜大模型變革

• 李白：你的模型權重很不錯，可惜被我沒收了

• OpenAI掌門Sam Altman：AI下一個發展階段

• 比快更快，開源Stable Diffusion重新整理作圖速度

• OneEmbedding:單卡 訓練TB級推薦模型不是夢

歡迎Star、試用OneFlow最新版本：https://github.com/Oneflow-Inc/oneflow/

---