当前位置：首页 > 技术类 > 树莓派上运行 Stable Diffusion，260MB 的 RAM“hold”住 10 亿参数大模型

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

微信用户2年前 (2023-07-24)技术类1252

编辑：梓文、张倩

Stable Diffusion 能在树莓派上运行了！

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

树莓派上运行 Stable Diffusion，260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技提供

Raspberry Pi Zero 2 「Just as small. Five times as fast.」这是怎样一个概念？运行 Stable Diffusion 并不是一件容易的事，它包含一个 10 亿参数的大型 Transformer 模型，建议使用的最低 RAM/VRAM 通常为 8GB。而 RPI Zero 2 只是内存为 512MB 的微型计算机。这意味着在 RPI Zero 2 上运行 Stable Diffusion 是一个巨大的挑战。而且，在运行过程中，作者没有增加存储空间，也没有将中间结果卸载到磁盘上。一般而言，主要的机器学习框架和库都专注于最小化推理延迟和 / 或最大化吞吐量，但以上这些都以内存使用为代价。因此，作者决定写一个超小的、可破解的推理库，致力于将内存消耗最小化。OnnxStream 做到了。

树莓派上运行 Stable Diffusion，260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技提供

OnnxStream 基于将推理引擎与负责提供模型权重的组件解耦的思路，后者是派生自 WeightsProvider 的一个类。一个 WeightsProvider 的专门化可以实现任何类型的模型参数加载、缓存和预取。例如，一个自定义的 WeightsProvider 可以决定直接从 HTTP 服务器下载数据，而不加载或写入任何内容到磁盘（这也是 OnnxStream 命名中有 Stream 的原因）。有两个默认的 WeightsProviders 可用：DiskNoCache 和 DiskPrefetch。与微软的推理框架 OnnxStream 相比，OnnxStream 只需要消耗 1/55 的内存就可以达到同样的效果，但（在 CPU 上的）速度只比前者慢 0.5-2 倍。接下来你将看到 Stable Diffusion 在 RPI Zero 2 上运行的效果，以及背后的方法。需要注意的是，虽然运行速度较慢，但是它是大模型在更小、更有限的设备上运行的崭新尝试。

树莓派上运行 Stable Diffusion，260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技提供

网友们认为这个项目很酷将 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。

树莓派上运行 Stable Diffusion，260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技提供

精度为 W16A16 的 VAE 解码器的生成效果
树莓派上运行 Stable Diffusion，260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技提供

精度为 W8A32 的 VAE 解码器的生成效果第三张图由 RPI Zero 2 在大约 3 小时内生成。
树莓派上运行 Stable Diffusion，260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技提供

图注：精度为 W8A8 的 VAE 解码器的生成效果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 精度）。

树莓派上运行 Stable Diffusion，260MB 的 RAM「hold」住 10 亿参数大模型© 由 ZAKER科技提供

解决方案是垂直分割 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。

免责声明

    本网站在国家相关法律法规规定的范围内，只按现有状况提供文章发布第三方网络平台服务，本网站及其所有者非交易一方，也非交易任何一方之代理人或代表；同时，本网站及其所有者也未授权任何人代表或代理本网站及其所有者从事任何网络交易行为或做出任何承诺、保证或其他类似行为，除非有明确的书面授权。
    鉴于互联网的特殊性，本网站无法鉴别和判断相关交易各主体之民事权利和行为能力、资质、信用等状况，也无法鉴别和判断虚拟交易或正在交易或已交易之虚拟物品来源、权属、真伪、性能、规格、质量、数量等权利属性、自然属性及其他各种状况。因此，交易各方在交易前应加以仔细辨明，并慎重考虑和评估交易可能产生的各项风险。