极限优化：298MB RAM 运行 Stable Diffusion XL

📣 更新（2023 年 10 月）📣 已新增对 Stable Diffusion XL 1.0 Base 的支持！而且它可以在 RPI Zero 2 上运行！请参阅下文相关章节 👇

这项挑战的目标，是在 Raspberry Pi Zero 2 上运行 Stable Diffusion 1.5。Stable Diffusion 1.5 包含一个拥有近 10 亿参数的大型 Transformer 模型，而 Raspberry Pi Zero 2 是一台只有 512MB RAM 的微型计算机。整个过程不能增加额外的 swap 空间，也不能把中间结果卸载到磁盘。通常情况下，Stable Diffusion 建议至少配备 8GB RAM/VRAM。

主流机器学习框架和库通常专注于最大限度地降低推理延迟和/或提升吞吐量，而这些优化往往以增加 RAM 占用为代价。因此，我决定编写一个体积极小、易于修改，并且专门以降低内存消耗为目标的推理库：OnnxStream。

OnnxStream 的核心思路，是将推理引擎与负责提供模型权重的组件解耦，后者是一个派生自 WeightsProvider 的类。不同的 WeightsProvider 实现可以采用任意方式加载、缓存和预取模型参数。例如，自定义 WeightsProvider 可以直接从 HTTP 服务器下载数据，完全不需要从磁盘加载或向磁盘写入任何内容（这正是 OnnxStream 名称中「Stream」一词的由来）。目前提供了两个默认的 WeightsProvider：DiskNoCache 和 DiskPrefetch。

与 OnnxRuntime 相比，OnnxStream 的内存消耗最多可降低到其 1/55，而速度仅慢 0.5～2 倍（在 CPU 上运行，详见下文的性能部分）。

下面这些图片由本仓库内置的 Stable Diffusion 示例实现生成。该实现使用 OnnxStream，并分别采用不同精度的 VAE decoder。VAE decoder 是 Stable Diffusion 中唯一无法以单精度或半精度装入 Raspberry Pi Zero 2 RAM 的模型。这是因为模型中存在 residual connection，以及体积非常大的 tensor 和 convolution。唯一的解决方案是 static quantization（8 bit）。第三张图片由我的 RPI Zero 2 生成，大约耗时 3 小时、1.5 小时（编译时使用 MAX_SPEED 选项）。作为对比，第一张图片是在我的 PC 上使用 RPI Zero 2 生成的相同 latents 制作的：

采用 W16A16 精度的 VAE decoder：

采用 W8A32 精度的 VAE decoder：

采用 W8A8 精度的 VAE decoder，由我的 RPI Zero 2 生成，大约耗时 3 小时、1.5 小时（编译时使用 MAX_SPEED 选项）：

OnnxStream 的 Stable Diffusion 示例实现现在已经支持 SDXL 1.0（不包含 Refiner）。ONNX 文件导出自 Hugging Face Diffusers 库的 SDXL 1.0 实现（版本 0.19.3）。

SDXL 1.0 的计算开销远高于 SD 1.5。两者最显著的区别，是 SDXL 1.0 可以生成 1024×1024 的图片，而不是 512×512。作为参考，在我配备 32GB RAM 的 12 核 PC 上，使用 Hugging Face Diffusers 生成一张 10 steps 的图片需要 26 分钟。通常情况下，SDXL 建议至少配备 12GB VRAM。

OnnxStream 可以在不到 300MB RAM 的环境中运行 SDXL 1.0，因此它能够轻松地在 RPI Zero 2 上运行，无须增加额外的 swap 空间，推理期间也无须向磁盘写入任何内容。在我的 RPI Zero 2 上生成一张 10 steps 的图片大约需要 11 小时。

SDXL 1.0 使用了与 SD 1.5 相同的一系列优化，但存在以下不同之处。

对于 UNET 模型，为了让它能在 RPI Zero 2 上以不到 300MB RAM 运行，我们使用了 UINT8 dynamic quantization，但仅将其应用于一组特定的大型中间 tensor。

VAE decoder 的情况比 SD 1.5 更复杂。SDXL 1.0 的 VAE decoder 大小是 SD 1.5 的 4 倍，使用 OnnxStream 以 FP32 精度运行时会消耗 4.4GB RAM。

对于 SD 1.5，VAE decoder 使用了 static quantization（UINT8 精度），这足以将 RAM 消耗降低到 260MB。相比之下，SDXL 1.0 的 VAE decoder 使用 FP16 arithmetic 运行时会发生 overflow，而且其 activation 的数值范围太大，无法通过 UINT8 quantization 获得质量良好的图片。

因此，我们面对的是一个会消耗 4.4GB RAM、无法使用 FP16 精度运行，同时也无法量化为 UINT8 精度的模型。（注意：FP16 问题至少存在一种解决方案，但我没有继续深入研究，因为即使能以 FP16 精度运行 VAE decoder，总内存消耗也只是减半，模型最终仍会消耗 2.2GB，而不是 4.4GB；对于 RPI Zero 2 来说，这依然高得离谱。）

解决方案的灵感来自 Hugging Face Diffusers 库对 VAE decoder 的实现，也就是使用 tiled decoding。最终结果与完整 decoder 解码出的图片完全无法分辨，而这种方式能够把 RAM 消耗从 4.4GB 降低到 298MB！

思路很简单。diffusion 过程的结果是一个形状为 (1,4,128,128) 的 tensor。我们将这个 tensor 拆分成 5×5、也就是 25 个相互重叠且形状为 (1,4,32,32) 的 tensor，然后分别对这些 tensor 进行解码。每个 tensor 都会与左侧和上方的 tile 重叠 25%。解码结果是一个形状为 (1,3,256,256) 的 tensor，随后通过适当的 blending 合并到最终图片中。

例如，下面这张图片由 tiled decoder 生成，代码中的 blending 被手动关闭。你可以清楚地看到图片中的各个 tile：

下面则是启用 blending 后的同一张图片。这就是最终结果：

这是另一张由我的 RPI Zero 2 生成的图片，大约耗时 11 小时：（10 steps，Euler Ancestral）

• 推理引擎与 WeightsProvider 解耦
• WeightsProvider 可以是 DiskNoCache、DiskPrefetch 或自定义实现
• Dynamic quantization（8 bit unsigned、asymmetric、percentile）
• Static quantization（W8A8 unsigned、asymmetric、percentile）
• 可轻松校准 quantized model
• 支持 FP16（可以使用或不使用 FP16 arithmetic）
• 实现了 25 个 ONNX operator（最常用的那些）
• operation 按顺序执行，但所有 operator 都支持 multithreading
• 单个实现文件加一个 header 文件
• XNNPACK 调用封装在 XnnPack 类中（方便未来替换）

OnnxStream 的部分加速 primitive 依赖 XNNPACK，包括：MatMul、Convolution、逐元素 Add/Sub/Mul/Div、Sigmoid 和 Softmax。

Stable Diffusion 由三个模型组成：text encoder（672 个 operation、1.23 亿个参数）、UNET 模型（2050 个 operation、8.54 亿个参数）以及 VAE decoder（276 个 operation、4900 万个参数）。假设 batch size 等于 1，使用 Euler Ancestral scheduler 时，完整生成一张效果良好的 10 steps 图片，需要运行 text encoder 2 次、UNET 模型 20 次（即 2×10），以及 VAE decoder 1 次。

这张表展示了 Stable Diffusion 三个模型各自的推理时间和内存消耗，也就是 Windows 上的 Peak Working Set Size，或 Linux 上的 Maximum Resident Set Size。

以 UNET 模型为例，当它以 FP16 精度运行，并在 OnnxStream 中启用 FP16 arithmetic 时，OnnxStream 的内存消耗最多可降低到 OnnxRuntime 的 1/55，而速度仅慢 0.5～2 倍。

对于 OnnxRuntime，第一次运行属于 warm-up inference，因为它的 InferenceSession 会在首次运行前创建，并在之后的所有运行中复用。OnnxStream 不存在这种 warm-up，因为它在设计上完全采用 eager 模式。不过，后续运行仍可以受益于操作系统对权重文件的缓存。

目前 OnnxStream 不支持 batch size != 1 的输入，这一点与 OnnxRuntime 不同。OnnxRuntime 在运行 UNET 模型时可以使用 batch size = 2，从而显著加快整个 diffusion 过程。

在我的测试中，修改 OnnxRuntime 的 SessionOptions，例如 EnableCpuMemArena 和 ExecutionMode，并不会让结果出现显著差异。

OnnxRuntime 的性能与 NCNN（我评估的另一个框架）非常接近，无论是内存消耗还是推理时间都是如此。如果有需要，我会在未来加入 NCNN 的 benchmark。

测试在我的开发机器上完成：Windows Server 2019、16GB RAM、8750H CPU（AVX2）、970 EVO Plus SSD，以及 VMWare 中的 8 个虚拟核心。

在运行 UNET 模型时使用「attention slicing」，以及对 VAE decoder 使用 W8A8 quantization，是将内存消耗降低到足以在 RPI Zero 2 上运行的关键。

虽然网上有大量关于 neural network quantization 的资料，但关于「attention slicing」的信息却很少。其思路很简单：在计算 UNET 模型中各个 multi-head attention 的 scaled dot-product attention 时，避免实体化完整的 Q @ K^T 矩阵。UNET 模型的 attention head 数量为 8，Q 的形状是 (8,4096,40)，而 K^T 的形状是 (8,40,4096)。因此，第一个 MatMul 的结果最终形状为 (8,4096,4096)，这是一个大小为 512MB 的 tensor（FP32 精度）：

解决方案是纵向拆分 Q，然后像往常一样，分别对 Q 的每个 chunk 执行 attention operation。Q_sliced 的形状为 (1,x,40)，其中 x 等于 4096（在本例中）除以 onnxstream::Model::m_attention_fused_ops_parts。后者的默认值为 2，但可以自定义。这个简单的技巧，可以把 UNET 模型的整体内存消耗从 1.1GB 降低到 300MB（模型以 FP32 精度运行时）。另一个可能的替代方案是使用 FlashAttention，它无疑会更高效，但 FlashAttention 需要针对每一种受支持的架构（AVX、NEON 等）编写自定义 kernel，在我们的场景中还需要绕过 XnnPack。

下面的代码可以运行定义在 path_to_model_folder/model.txt 中的模型：（所有模型 operation 都定义在 model.txt 文本文件中；OnnxStream 还要求在同一目录中找到所有权重文件，这些文件以一系列 .bin 文件的形式存在。）

#include "onnxstream.h"

using namespace onnxstream;

int main()
{
    Model model;

    //
    // Optional parameters that can be set on the Model object:
    //
    // model.set_weights_provider( ... ); // specifies a different weights provider (default is DiskPrefetchWeightsProvider)
    // model.read_range_data( ... ); // reads a range data file (which contains the clipping ranges of the activations for a quantized model)
    // model.write_range_data( ... ); // writes a range data file (useful after calibration)
    // model.m_range_data_calibrate = true; // calibrates the model
    // model.m_use_fp16_arithmetic = true; // uses FP16 arithmetic during inference (useful if weights are in FP16 precision)
    // model.m_use_uint8_arithmetic = true; // uses UINT8 arithmetic during inference
    // model.m_use_uint8_qdq = true; // uses UINT8 dynamic quantization (can reduce memory consumption of some models)
    // model.m_fuse_ops_in_attention = true; // enables attention slicing
    // model.m_attention_fused_ops_parts = ... ; // see the "Attention Slicing" section above
    //

    model.read_file("path_to_model_folder/model.txt");

    tensor_vector<float> data;
    
    ... // fill the tensor_vector with the tensor data. "tensor_vector" is just an alias to a std::vector with a custom allocator.

    Tensor t;
    t.m_name = "input";
    t.m_shape = { 1, 4, 64, 64 };
    t.set_vector(std::move(data));
    model.push_tensor(std::move(t));

    model.run();
    
    auto& result = model.m_data[0].get_vector<float>();
    
    ... // process the result: "result" is a reference to the first result of the inference (a tensor_vector<float> as well).

    return 0;
}

model.txt 文件以 ASCII 格式包含模型的所有 operation，它们均由原始 ONNX 文件导出。每一行对应一个 operation。例如，下面这一行表示 quantized model 中的一个 convolution：

Conv_4:Conv*input:input_2E_1(1,4,64,64);post_5F_quant_5F_conv_2E_weight_nchw.bin(uint8[0.0035054587850383684,134]:4,4,1,1);post_5F_quant_5F_conv_2E_bias.bin(float32:4)*output:input(1,4,64,64)*dilations:1,1;group:1;kernel_shape:1,1;pads:0,0,0,0;strides:1,1

为了从 ONNX 文件中导出 model.txt 及其权重（以一系列 .bin 文件的形式），供 OnnxStream 使用，项目提供了一个只有单个 cell 的 notebook：onnx2txt.ipynb。

将 Pytorch nn.Module（本项目的情况）导出为 ONNX 并供 OnnxStream 使用时，需要注意以下事项：

调用 torch.onnx.export 时，dynamic_axes 应当留空，因为 OnnxStream 不支持 dynamic shape 的输入。

强烈建议先使用优秀的 ONNX Simplifier 处理导出的 ONNX 文件，再将其转换为 model.txt 文件。

仅限 Windows：启动以下命令提示符：Visual Studio Tools > x64 Native Tools Command Prompt。

仅限 Mac：请确保已安装 cmake：brew install cmake。

首先需要构建 XNNPACK。

由于 XnnPack 的函数原型随时可能发生变化，因此我加入了一条 git checkout 命令，用于确保 OnnxStream 能够与本文编写时兼容的 XnnPack 版本正确编译：

git clone https://github.com/google/XNNPACK.git
cd XNNPACK
git rev-list -n 1 --before="2023-06-27 00:00" master
git checkout <COMMIT_ID_FROM_THE_PREVIOUS_COMMAND>
mkdir build
cd build
cmake -DXNNPACK_BUILD_TESTS=OFF -DXNNPACK_BUILD_BENCHMARKS=OFF ..
cmake --build . --config Release

接下来，就可以构建 Stable Diffusion 示例了。

<DIRECTORY_WHERE_XNNPACK_WAS_CLONED> 例如可以是 /home/vito/Desktop/XNNPACK，或者 Windows 上的 C:\Projects\SD\XNNPACK：

git clone https://github.com/vitoplantamura/OnnxStream.git
cd OnnxStream
cd src
mkdir build
cd build
cmake -DMAX_SPEED=ON -DXNNPACK_DIR=<DIRECTORY_WHERE_XNNPACK_WAS_CLONED> ..
cmake --build . --config Release

重要：MAX_SPEED 选项在 Windows 上可以将性能提升约 10%，在 Raspberry Pi 上则能提升 50% 以上。该选项会在构建期间消耗更多内存，而且生成的 executable 可能无法运行（我的测试中，Termux 就出现了这种情况）。因此，如果遇到问题，首先应该尝试将 MAX_SPEED 设置为 OFF。

现在，你可以运行 Stable Diffusion 示例了。

对于 Stable Diffusion 1.5，示例所需的权重可以从本仓库的 Releases 下载（约 2GB）。对于 Stable Diffusion XL 1.0 Base，权重可以从 Hugging Face 下载（约 8GB）：

git lfs install
git clone --depth=1 https://huggingface.co/vitoplantamura/stable-diffusion-xl-base-1.0-onnxstream

下面是 Stable Diffusion 示例的命令行选项：

--xl                Runs Stable Diffusion XL 1.0 instead of Stable Diffusion 1.5.
--models-path       Sets the folder containing the Stable Diffusion models.
--ops-printf        During inference, writes the current operation to stdout.
--output            Sets the output PNG file.
--decode-latents    Skips the diffusion, and decodes the specified latents file.
--prompt            Sets the positive prompt.
--neg-prompt        Sets the negative prompt.
--steps             Sets the number of diffusion steps.
--save-latents      After the diffusion, saves the latents in the specified file.
--decoder-calibrate (ONLY SD 1.5) Calibrates the quantized version of the VAE decoder.
--decoder-fp16      (ONLY SD 1.5) During inference, uses the FP16 version of the VAE decoder.
--not-tiled         (ONLY SDXL 1.0) Don't use the tiled VAE decoder.
--rpi               Configures the models to run on a Raspberry Pi.
--rpi-lowmem        (ONLY SDXL 1.0) Configures the models to run on a Raspberry Pi Zero 2.

sd.cpp 中的 Stable Diffusion 实现基于另一个项目，而那个项目又基于 @EdVince 的项目。原始代码经过了修改，以便使用 OnnxStream 代替 NCNN。