優化 Intel CPU 上的 Stable Diffusion 使用 NNCF 與 🤗 Optimum

總覽

Hugging Face 提供了一種方法，透過結合 Neural Network Compression Framework (NNCF) 量化感知訓練與 Token Merging，來優化 Intel CPU 上的 Stable Diffusion，使推理速度提升最高達 5.1 倍，且模型大小減少至原始 PyTorch 檢查點的 0.25 倍。

Stable Diffusion 優化

Stable Diffusion 管線的 UNet 元件是計算成本最高的部分，因此優化它可以帶來顯著的速度提升。傳統的訓練後 8‑bit 量化對此模型效果不佳，因為像素級預測任務對參數變化高度敏感，且模型在數億樣本的訓練中冗餘度低。為保持準確度，需要更複雜的量化方法，例如量化感知訓練（QAT）。

優化工作流程

我們從在 Pokemon 資料集（svjack/Stable-Diffusion-Pokemon-en）上微調的 Stable Diffusion 模型開始。利用 Diffusers text‑to‑image 微調範例，我們將基於 NNCF 的 QAT 整合到訓練腳本中，加入知識蒸餾損失（其中原始模型作為教師），並對模型參數（不包括量化器）應用指數移動平均（EMA）以提升訓練穩定性。梯度檢查點以及將 EMA 模型保留在 RAM 中，使得整個優化流程能在單塊擁有 24 GB VRAM 的 GPU 上，於不到一天的時間內完成 4096 次迭代。

超越量化感知訓練

僅使用量化即可減少模型佔用空間、載入時間、記憶體消耗與延遲。我們將 8‑bit 量化與 Token Merging (ToME) 方法疊加，該方法會在自注意力區塊之前合併冗餘的 token 以減少計算量。結合後的工作流程包含上述所述的知識蒸餾、EMA 與梯度檢查點。我們再次從已微調的 Pokemon 模型出發，在量化基礎上應用合併比例為 0.4 的 ToME。得到的模型適合在用戶端或邊緣 CPU 上進行推理。

結果

將 PyTorch 基線轉換為 OpenVINO FP32 可獲得 1.9 倍的加速。加入 8‑bit 量化後，相對於 PyTorch 的加速提升至 3.9 倍，且模型佔用空間減少至原始檢查點的 0.25 倍。在量化基礎上再疊加 Token Merging，可達到 5.1 倍的推理加速，同時模型佔用空間維持在相同的 0.25 倍水準。所有測試均在使用 OpenVINO 2022.3 的 Hugging Face Spaces CPU 升級實例上進行，該實例搭載第三代 Intel® Xeon® 可擴展處理器並具備 Intel® Deep Learning Boost 技術，使用預設的 50 步推理步驟。部落格指出，減少步驟數可提升速度但可能影響圖像品質，並建議嘗試不同的步驟數與排程器。

from optimum.intel import OVStableDiffusionPipeline

# Load and compile the pipeline for performance.
name = "OpenVINO/stable-diffusion-pokemons-tome-quantized-aggressive"
pipe = OVStableDiffusionPipeline.from_pretrained(name, compile=False)
pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
pipe.compile()

# Generate an image.
prompt = "a drawing of a green pokemon with red eyes"
output = pipe(prompt, num_inference_steps=50, output_type="pil\)).images[0]
output.save("image.png

訓練與量化程式碼可在 Optimum Intel 儲存庫中取得，並提供示範筆記本，優化後的模型可在 Hugging Face Hub 的 OpenVINO 組織下找到。即時示範在 Hugging Face Spaces 上運行。

關於通用目的的 Stable Diffusion 模型？

在 Pokemon 模型上展示的工作流程表明，即使訓練資源有限，也能達成顯著的優化。從零開始訓練通用目的的 Stable Diffusion 模型成本高昂，但在預算與硬體充足的情況下，仍可採用相同的方法。需要注意的是，Token Merging 會降低模型容量；因此，對於更複雜的資料集，在優化過程中應該使用較低的合併比例。

如需了解第四代 Intel Xeon CPU 的互補方法，請參閱相關的 Hugging Face 部落格文章，關於在 Intel 上進行 Stable Diffusion 推理。

Sources