现在,它甚至已经可以在 Raspberry Pi Zero 2 上运行了。
图片
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 做到了。
项目地址 https://github.com/vitoplantamura/OnnxStream
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 在 Raspberry Pi Zero 2 上运行
VAE 解码器是 Stable Diffusion 中唯一无法以单精度或半精度放入 RPI Zero 2 RAM 的模型。这是因为模型中存在残差连接、非常大的张量和卷积。唯一的解决办法就是静态量化(8 bit)。
以下这些图像是由作者 repo 中包含的 Stable Diffusion 示例实现在不同精度的 VAE 解码器下使用 OnnxStream 生成的。
第一张图像是在作者的 PC 上生成的,使用了由 RPI Zero 2 生成的相同的 latent。
精度为 W16A16 的 VAE 解码器的生成效果
精度为 W8A32 的 VAE 解码器的生成效果
第三张图由 RPI Zero 2 在大约 3 小时内生成。
图注:精度为 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 精度)。
图片
解决方案是垂直分割 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 界面。