树莓派上运行 Stable Diffusion，260MB 的 RAM「hold」住 10 亿参数大模型

11 个月前 Stable Diffusion 诞生，它能够在消费级 GPU 上运行的消息让不少研究者备受鼓舞。不仅如此，苹果官方很快下场，将 Stable Diffusion「塞进」iPhone、iPad 和 Mac 中运行。这大大降低了 Stable Diffusion 对硬件设备的要求，让其逐渐成为人人都能使用的「黑科技」。

现在，它甚至已经可以在 Raspberry Pi Zero 2 上运行了。  

这是怎样一个概念？运行 Stable Diffusion 并不是一件容易的事，它包含一个 10 亿参数的大型 Transformer 模型，建议使用的最低 RAM/VRAM 通常为 8GB。而 RPI Zero 2 只是内存为 512MB 的微型计算机。

这意味着在 RPI Zero 2 上运行 Stable Diffusion 是一个巨大的挑战。而且，在运行过程中，作者没有增加存储空间，也没有将中间结果卸载到磁盘上。

一般而言，主要的机器学习框架和库都专注于最小化推理延迟和 / 或最大化吞吐量，但以上这些都以内存使用为代价。因此，作者决定写一个超小的、可破解的推理库，致力于将内存消耗最小化。

OnnxStream 做到了。  

  OnnxStream 基于将推理引擎与负责提供模型权重的组件解耦的思路，后者是派生自 WeightsProvider 的一个类。一个 WeightsProvider 的专门化可以实现任何类型的模型参数加载、缓存和预取。例如，一个自定义的 WeightsProvider 可以决定直接从 HTTP 服务器下载数据，而不加载或写入任何内容到磁盘（这也是 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 的模型。这是因为模型中存在残差连接、非常大的张量和卷积。唯一的解决办法就是静态量化（8 bit）。

以下这些图像是由作者 repo 中包含的 Stable Diffusion 示例实现在不同精度的 VAE 解码器下使用 OnnxStream 生成的。

第一张图像是在作者的 PC 上生成的，使用了由 RPI Zero 2 生成的相同的 latent。  

第三张图由 RPI Zero 2 在大约 3 小时内生成。

OnnxStream 的特点

• 推理引擎与 WeightsProvider 解耦
• WeightsProvider 可以是 DiskNoCache、DiskPrefetch 或自定义
• 注意力切片
• 动态量化（8 bit 无符号、非对称、百分位数）
• 静态量化（W8A8 无符号、非对称、百分位数）
• 轻松校准量化模型
• 支持 FP16（使用或不使用 FP16 运算）
• 实现了 24 个 ONNX 算子（最常用的算子）
• 运算按顺序执行，但所有算子都是多线程的
• 单一实现文件 + header 文件
• XNNPACK 调用被封装在 XnnPack 类中 (用于将来的替换)

并且需要注意的是，OnnxStream 依赖 XNNPACK 来加速某些原语：MatMul、Convolution、element-wise Add/Sub/Mul/Div、Sigmoid 和 Softmax。

性能对比

Stable Diffusion 由三个模型组成：文本编码器（672 次运算和 1.23 亿个参数）、UNET 模型（2050 次运算和 8.54 亿个参数）和 VAE 解码器（276 次运算和 4900 万个参数。
假设批大小等于 1，生成完整图像则需要 10 步，这需要运行 2 次文本编码器、运行 20 次（即 2*10）UNET 模型和运行 1 次 VAE 解码器，才能获得良好效果（使用 Euler Ancestral 调度器）。
该表显示了 Stable Diffusion 的三个模型不同的推理时间，以及内存消耗（即 Windows 中的 Peak Working Set Size 或 Linux 中的 Maximum Resident Set Size）。
可以发现，在 UNET 模型中（以 FP16 精度运行时，OnnxStream 中启用了 FP16 算术），OnnxStream 的内存消耗量仅为 OnnxRuntime 的 1/55，但速度只慢 0.5-2 倍。
这次测试需要注明的几点是：

• OnnxRuntime 的第一次运行是预热推理，因为它的 InferenceSession 是在第一次运行前创建的，并在随后的所有运行中重复使用。而 OnnxStream 没有预热推理，因为它的设计是纯粹「eager」的（不过，后续运行可以受益于操作系统对权重文件的缓存）。
• 目前 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 上运行的水平至关重要。

虽然互联网上有很多关于量化神经网络的信息，但关于「注意力切片」的却很少。

这里的想法很简单：目标是在计算 UNET 模型中各种多头注意力的缩放点积注意力时，避免生成完整的 Q @ K^T 矩阵。在 UNET 模型中，注意力头数为 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