用于图像生成的大型扩散模型 的激增导致模型大小和推理工作量大幅增加。在移动环境中进行设备上的 ML 推理需要细致的性能优化，并考虑资源限制带来的权衡。在设备上运行大型扩散模型 (LDM) 的推理，出于对成本效率和用户隐私的需求，由于这些模型的大量内存要求和计算需求，带来了更大的挑战。
我们在题为“速度就是一切：通过 GPU 感知优化实现大型扩散模型的设备加速”的论文中解决了这一挑战（将在CVPR 2023计算机视觉高效深度学习研讨会上发表），重点关注在移动 GPU 上优化执行基础 LDM 模型。在这篇博文中，我们总结了在现代智能手机上成功执行全分辨率（512x512 像素）和 20 次迭代等大型扩散模型所采用的核心技术，原始模型的高性能推理速度无需蒸馏即可在 12 秒内完成。正如我们之前的博文所讨论的那样，GPU 加速的 ML 推理通常受内存性能的限制，LDM 的执行也不例外。因此，我们优化的核心主题是高效的内存输入/输出 (I/O)，即使这意味着选择内存高效的算法而不是优先考虑算术逻辑单元效率的算法。最终，我们的主要目标是减少 ML 推理的总体延迟。
移动 GPU 上的 LDM 示例输出，提示文本为：“一张可爱的小狗周围有鲜花的逼真高分辨率图像”。
增强注意力模块，提高记忆效率
ML 推理引擎通常提供各种优化的 ML 操作。尽管如此，实现最佳性能仍然具有挑战性，因为在 GPU 上执行单个神经网络运算符会产生一定开销。为了减轻这种开销，ML 推理引擎采用了广泛的运算符融合规则，将多个运算符合并为一个运算符，从而减少跨张量元素的迭代次数，同时最大限度地提高每次迭代的计算能力。例如，TensorFlow Lite利用运算符融合将计算成本高昂的运算（如卷积）与后续激活函数（如整流线性单元）合并为一个。
一个明显的优化机会是LDM 中降噪器模型中采用的大量使用的注意力模块。注意力模块通过为重要区域分配更高的权重，使模型能够专注于输入的特定部分。有多种方法可以优化注意力模块，我们会根据哪种优化效果更好，选择性地采用下面介绍的两种优化方法之一。
第一个优化，我们称之为部分融合softmax，它消除了在注意力模块中的 softmax 和矩阵乘法之间进行大量内存写入和读取的需要。让注意力模块只是一个简单的矩阵乘法，形式为Y = softmax( X ) * W，其中X和W分别是形状为a × b和b × c的二维矩阵（如下图上半部分所示）。
为了实现数值稳定性，T = softmax( X ) 通常分三步计算：
确定列表中的最大值，即矩阵 X 中每一行的
将每个列表项的指数与最大值的差值相加（来自第 1 遍）
将项的指数减去最大值，除以第 2 轮的总和
如果简单地执行这些过程，将导致临时中间张量 T 需要大量内存写入，因为临时中间张量T保存着整个 softmax 函数的输出。如果我们只存储过程 1 和 2 的结果（分别标记为m和s ），我们就可以绕过这种大量内存写入，这两个结果都是小向量，每个向量都有a个元素，而T有a·b个元素。使用这种技术，我们能够将数十甚至数百兆字节的内存消耗减少多个数量级（如下图下半部分所示）。
注意力模块。顶部：一个简单的注意力模块，由一个 SOFTMAX（包含所有三个通道）和一个MATMUL组成，需要为大中间张量T写入大量内存。底部：我们在 MATMUL 中部分融合 softmax 的内存高效注意力模块只需要为m和 s 存储两个小的中间张量。
另一项优化涉及使用FlashAttention，这是一种 I/O 感知的精确注意算法。该算法减少了 GPU 高带宽内存访问的次数，非常适合内存带宽有限的用例。但是，我们发现这种技术仅适用于具有特定大小的SRAM，并且需要大量寄存器。因此，我们仅在选定的一组 GPU 上对具有特定大小的注意矩阵利用这种技术。
3×3 卷积层的 Winograd 快速卷积
常见 LDM 的主干严重依赖于 3×3 卷积层（过滤器大小为 3×3 的卷积），占解码器中 90% 以上的层。尽管内存消耗和数值误差增加，但我们发现Winograd 快速卷积能够有效加速卷积。与卷积中使用的过滤器大小3×3 不同，图块大小是指一次处理的输入张量的子区域的大小。增加图块大小可提高卷积在算术逻辑单元(ALU) 使用方面的效率。但是，这种改进是以增加内存消耗为代价的。我们的测试表明，4×4 的图块大小在计算效率和内存利用率之间实现了最佳平衡。
内存使用情况
瓷砖尺寸 FLOPS节省 中间张量 重量
2×2 2.25× 4.00× 1.77×
4×4 4.00× 2.25× 4.00×
6×6 5.06× 1.80× 7.12×
8×8 5.76× 1.56× 11.1×
具有不同图块尺寸的 Winograd 对 3×3 卷积的影响。
专门的运算符融合以提高内存效率
我们发现，在移动 GPU 上高效地推断 LDM 需要比当前现成的设备 GPU 加速 ML 推理引擎提供的更大的融合窗口来处理 LDM 中常用的层和单元。因此，我们开发了专门的实现，可以执行比典型融合规则允许的更大范围的神经运算符。具体来说，我们专注于两个专门化：高斯误差线性单元(GELU) 和组规范化层。
使用双曲正切函数 近似 GELU需要写入和读取七个辅助中间张量（如下图所示为浅橙色圆角矩形），从输入张量x读取三次，并在八个 GPU 程序中将输出张量y写入一次，每个程序都实现标记的操作（浅蓝色矩形）。在单个着色器中执行八个操作的自定义 GELU 实现（如下图底部所示）可以绕过中间张量的所有内存 I/O。
GELU 实现。顶部：具有内置操作的简单实现需要 8 次内存写入和 10 次读取。底部：我们的自定义 GELU 仅需要 1 次内存读取（针对x）和 1 次写入（针对y）。
结果
在应用所有这些优化后，我们在高端移动设备上对 Stable Diffusion 1.5（图像分辨率 512x512，20 次迭代）进行了测试。使用我们的 GPU 加速 ML 推理模型运行 Stable Diffusion 时，权重占用 2,093MB，中间张量占用 84MB。使用最新的高端智能手机，Stable Diffusion 可以在 12 秒内运行。
Stable Diffusion 在现代智能手机上运行时间不到 12 秒。请注意，每次迭代后运行解码器以显示此动画 GIF 中的中间输出会导致速度减慢约 2 倍。
结论
事实证明，在设备上执行大型模型的 ML 推理是一项巨大的挑战，包括模型文件大小的限制、大量的运行时内存要求以及长时间的推理延迟。通过认识到内存带宽使用是主要瓶颈，我们将努力优化内存带宽利用率并在 ALU 效率和内存效率之间取得微妙的平衡。结果，我们实现了大型扩散模型的最先进的推理延迟。您可以在论文中了解有关这项工作的更多信息。
致谢
我们要感谢 Yu-Hui Chen、Jiuqiang Tang、Frank Barchard、Yang Chao、Joe Zou、Khanh LeViet、Chuo-Ling Chang、Andrei Kulik、Lu Wang 和 Matthias Grundmann。
评论