樹莓派上運行 Stable Diffusion，260MB 的 RAM「hold」住 10 億參數(shù)大模型

Stable Diffusion 能在樹莓派上運行了！

11 個月前 Stable Diffusion 誕生，它能夠在消費級 GPU 上運行的消息讓不少研究者備受鼓舞。不僅如此，蘋果官方很快下場，將 Stable Diffusion「塞進」iPhone、iPad 和 Mac 中運行。這大大降低了 Stable Diffusion 對硬件設備的要求，讓其逐漸成為人人都能使用的「黑科技」。
現(xiàn)在，它甚至已經可以在 Raspberry Pi Zero 2 上運行了。
Raspberry Pi Zero 2 「Just as small. Five times as fast.」
這是怎樣一個概念？運行 Stable Diffusion 并不是一件容易的事，它包含一個 10 億參數(shù)的大型 Transformer 模型，建議使用的最低 RAM/VRAM 通常為 8GB。而 RPI Zero 2 只是內存為 512MB 的微型計算機。
這意味著在 RPI Zero 2 上運行 Stable Diffusion 是一個巨大的挑戰(zhàn)。而且，在運行過程中，作者沒有增加存儲空間，也沒有將中間結果卸載到磁盤上。
一般而言，主要的機器學習框架和庫都專注于最小化推理延遲和 / 或最大化吞吐量，但以上這些都以內存使用為代價。因此，作者決定寫一個超小的、可破解的推理庫，致力于將內存消耗最小化。
OnnxStream 做到了。
 項目地址 https://github.com/vitoplantamura/OnnxStream
OnnxStream 基于將推理引擎與負責提供模型權重的組件解耦的思路，后者是派生自 WeightsProvider 的一個類。一個 WeightsProvider 的專門化可以實現(xiàn)任何類型的模型參數(shù)加載、緩存和預取。例如，一個自定義的 WeightsProvider 可以決定直接從 HTTP 服務器下載數(shù)據，而不加載或寫入任何內容到磁盤（這也是 OnnxStream 命名中有 Stream 的原因）。有兩個默認的 WeightsProviders 可用：DiskNoCache 和 DiskPrefetch。
與微軟的推理框架 OnnxRuntime 相比，OnnxStream 只需要消耗 1/55 的內存就可以達到同樣的效果，但（在 CPU 上的）速度只比前者慢 0.5-2 倍。
接下來你將看到 Stable Diffusion 在 RPI Zero 2 上運行的效果，以及背后的方法。需要注意的是，雖然運行速度較慢，但是它是大模型在更小、更有限的設備上運行的嶄新嘗試。
 網友們認為這個項目很酷
將 Stable Diffusion 在 Raspberry Pi Zero 2 上運行
VAE ****是 Stable Diffusion 中唯一無法以單精度或半精度放入 RPI Zero 2 RAM 的模型。這是因為模型中存在殘差連接、非常大的張量和卷積。唯一的解決辦法就是靜態(tài)量化（8 bit）。
以下這些圖像是由作者 repo 中包含的 Stable Diffusion 示例實現(xiàn)在不同精度的 VAE ****下使用 OnnxStream 生成的。
第一張圖像是在作者的 PC 上生成的，使用了由 RPI Zero 2 生成的相同的 latent。
精度為 W16A16 的 VAE ****的生成效果
精度為 W8A32 的 VAE ****的生成效果
第三張圖由 RPI Zero 2 在大約 3 小時內生成。圖注：精度為 W8A8 的 VAE ****的生成效果
OnnxStream 的特點

• 推理引擎與 WeightsProvider 解耦
• WeightsProvider 可以是 DiskNoCache、DiskPrefetch 或自定義
• 注意力切片
• 動態(tài)量化（8 bit 無符號、非對稱、百分位數(shù)）
• 靜態(tài)量化（W8A8 無符號、非對稱、百分位數(shù)）
• 輕松校準量化模型
• 支持 FP16（使用或不使用 FP16 運算）
• 實現(xiàn)了 24 個 ONNX 算子（最常用的算子）
• 運算按順序執(zhí)行，但所有算子都是多線程的
• 單一實現(xiàn)文件 + header 文件
• XNNPACK 調用被封裝在 XnnPack 類中 (用于將來的替換)

并且需要注意的是，OnnxStream 依賴 XNNPACK 來加速某些原語：MatMul、Convolution、element-wise Add/Sub/Mul/Div、Sigmoid 和 Softmax。
性能對比
Stable Diffusion 由三個模型組成：文本編碼器（672 次運算和 1.23 億個參數(shù)）、UNET 模型（2050 次運算和 8.54 億個參數(shù)）和 VAE ****（276 次運算和 4900 萬個參數(shù)。
假設批大小等于 1，生成完整圖像則需要 10 步，這需要運行 2 次文本編碼器、運行 20 次（即 2*10）UNET 模型和運行 1 次 VAE ****，才能獲得良好效果（使用 Euler Ancestral 調度器）。
該表顯示了 Stable Diffusion 的三個模型不同的推理時間，以及內存消耗（即 Windows 中的 Peak Working Set Size 或 Linux 中的 Maximum Resident Set Size）。
可以發(fā)現(xiàn)，在 UNET 模型中（以 FP16 精度運行時，OnnxStream 中啟用了 FP16 算術），OnnxStream 的內存消耗量僅為 OnnxRuntime 的 1/55，但速度只慢 0.5-2 倍。
這次測試需要注明的幾點是：

• OnnxRuntime 的第一次運行是預熱推理，因為它的 InferenceSession 是在第一次運行前創(chuàng)建的，并在隨后的所有運行中重復使用。而 OnnxStream 沒有預熱推理，因為它的設計是純粹「eager」的（不過，后續(xù)運行可以受益于操作系統(tǒng)對權重文件的緩存）。
• 目前 OnnxStream 不支持 batch size ! = 1 的輸入，這與 OnnxRuntime 不同，后者在運行 UNET 模型時使用 batch size = 2 可以大大加快整個擴散過程。
• 在測試中，改變 OnnxRuntime 的 SessionOptions（如 EnableCpuMemArena 和 ExecutionMode）對結果沒有產生明顯影響。
• 在內存消耗和推理時間方面，OnnxRuntime 的性能與 NCNN（另一個框架）非常相似。
• 測試的運行條件：Windows Server 2019、16GB 內存、8750H CPU (AVX2)、970 EVO Plus SSD, VMWare 上的 8 個虛擬內核。

注意力切片與量化
在運行 UNET 模型時，采用「注意力切片」技術，并對 VAE ****使用 W8A8 量化，這對于將模型內存消耗降低到適合在 RPI Zero 2 上運行的水平至關重要。
雖然互聯(lián)網上有很多關于量化神經網絡的信息，但關于「注意力切片」的卻很少。
這里的想法很簡單：目標是在計算 UNET 模型中各種多頭注意力的縮放點積注意力時，避免生成完整的 Q @ K^T 矩陣。在 UNET 模型中，注意力頭數(shù)為 8 時，Q 的形狀為 (8,4096,40)，同時 K^T 為 (8,40,4096)。因此，第一個 MatMul 的最終形狀為 (8,4096,4096)，這是一個 512MB 的張量（FP32 精度）。

解決方案是垂直分割 Q，然后在每個 Q 塊上正常進行注意力操作。Q_sliced 形狀為 (1,x,40)，其中 x 為 4096（在本例中），除以 onnxstream::Model::m_attention_fused_ops_parts（默認值為 2，但可以自定義。
這個簡單的技巧可以將 UNET 模型以 FP32 精度運行時的整體內存消耗從 1.1GB 降低到 300MB。一個更高效的替代方案是使用 FlashAttention，但是 FlashAttention 需要為每個支持的架構（AVX, NEON）等編寫自定義內核，在作者給出的例子中繞過 XnnPack。
更多信息參見該項目的 GitHub 界面。
參考鏈接：https://www.reddit.com/r/MachineLearning/comments/152ago3/p_onnxstream_running_stable_diffusion_in_260mb_of/https://github.com/vitoplantamura/OnnxStream