在 iPhone、iPad 和 Mac 上通过 Core ML 实现更快的 Stable Diffusion

TL;DR

Apple 和 Hugging Face 推出了新的 Core ML 优化，特别是 6-bit palettization（6 位调色板量化）和更新的 attention layer 实现，旨在让 Stable Diffusion 在 iPhone、iPad 和 Mac 上运行得更快，并降低内存开销。这些改进可以在极小精度损失的情况下实现模型权重的显著压缩，从而实现更高效的设备端生成式 AI。

Core ML 优化与 6-Bit Palettization

Core ML 利用 Apple 设备的 CPU、GPU 和 Neural Engine 在设备端运行机器学习模型。为了提高像 Stable Diffusion 这样的大型模型的效率，Apple 通过 coremltools.optimize 子模块引入了一种称为 6-bit palettization 的新压缩技术。

Palettization 的工作原理

Palettization 是量化的一种形式，它将模型权重从 16 位浮点表示压缩到每个参数 6 位。其工作原理类似于计算机图形学中的调色板：定义固定数量的颜色（调色板），然后用该调色板中最接近的可用颜色的索引来替换图像颜色。

内存与执行效率

在之前的 Core ML 版本中，压缩后的权重在从磁盘加载时会被解压缩，这意味着内存占用仍与未压缩的模型大小相同。更新后的框架现在会在推理过程中从一层到另一层时即时转换调色板权重。这种方法更快，因为它减少了内存传输量（执行过程中的主要瓶颈），而不是增加了解压缩的计算负载。

更新后的 Stable Diffusion 实现

基于 diffusers 库的 Apple ml-stable-diffusion 仓库已更新，以纳入这些新优化：

  • 量化支持： 用户现在可以使用 --quantize-nbits 标志将模型量化为 8、6、4 或 2 位。建议使用 6 位量化，因为它是精度和性能之间的最佳平衡。
  • Attention Layer 优化： 一种名为 SPLIT_EINSUM_V2 的新方法将查询序列分成 512 大小的块，以避免创建大型中间张量。这种优化可以在 Neural Engine 上提高 10% 到 30% 的性能。

可用的优化模型

以下官方 Stable Diffusion 模型的 6 位调色板版本现已在 Hugging Face Hub 上提供：

模型 未压缩 6-bit Palettized
Stable Diffusion 1.4 Core ML float16 Core ML 6-bit palettized
Stable Diffusion 1.5 Core ML float16 Core ML 6-bit palettized
Stable Diffusion 2 base Core ML float16 Core ML 6-bit palettized
Stable Diffusion 2.1 base Core ML float16 Core ML 6-bit palettized

部署与自定义模型转换

系统要求

要使用 6 位模型，设备必须运行 iOS/iPadOS 17 或 macOS 14 (Sonoma) 的开发版本，因为这些版本包含了必要的更新版 Core ML 框架。

转换自定义模型

开发人员可以使用 ml-stable-diffusion 仓库转换微调过的或 Dreambooth 模型。过程包括：

  1. 使用 coremltools 7.0 beta 和 Xcode 15.0 beta。
  2. 运行带有 --quantize-nbits 6 标志的 torch2coreml 转换脚本。
  3. 在 ORIGINAL attention 实现（通常对 macOS 更好）和 SPLIT_EINSUM_V2（通常对 iOS 更快）之间进行选择。

高级压缩：训练时量化

虽然 6-bit palettization 是训练后压缩的一个例子，但 coremltools 还引入了训练时压缩。这允许开发人员在进行权重压缩的同时微调模型。通过使用用于权重聚类的可微算法，可以在训练期间优化量化表以最大限度地减少损失，从而可能实现比训练后方法质量下降更低的 4 位或 2 位压缩。

Sources