众所周知,每年苹果都会推出一款声称在各方面都更快、更好的新款 iPhone,这主要得益于新的视觉模型和图像传感器的快速发展。就拿拍照来说,如果回到 10 年前,你能用 iPhone 拍出高质量的图片吗,答案是不能,因为技术的发展是渐进式的,10 年时间,足够提高手机拍照技术。
由于技术的这种发展模式(渐进式),在一段时间里有些程序即使运行在最好的计算设备上,也几乎无法使用。但是这些带有新启用场景的新程序吸引了一些用户的注意力,人们愿意研究它。
本文的作者就是被吸引的其中之一,在过去 3 周里,作者开发了一个应用程序,可以通过 Stable Diffusion 来生成(summon)图像,然后按你喜欢的方式编辑它。该应用在最新的 iPhone 14 Pro 上生成图片仅需一分钟,使用大约 2GiB 的应用内存,另外还需要下载大约 2GiB 的初始数据才能开始使用。
应用商店链接:https://apps.apple.com/us/app/draw-things-ai-generation/id6444050820
这一结果引来众多网友讨论,有人开始担心手机耗电问题,并开玩笑的说:这很酷,不过这看起来是一个消耗手机电池的好方法。
「我从来没有像现在这样开心地感受 iPhone 的热度。」
「这个寒冬,可以把手机当暖手器用了。」
不过在大家调侃手机发热问题的同时,他们也给与这项工作极高的评价。
「这简直不可思议。在我的 iPhone SE3 上生成一张完整的图像大约需要 45 秒——这几乎和我的 M1 Pro macbook 用原始版本生成的速度一样快!」
内存、硬件同时优化
这是如何做到的呢?接下来我们看看作者的实现过程:
想要完成在 iPhone 上运行 Stable Diffusion,还能结余 50% 的性能,面临的一大挑战是需要在 6GiB RAM 的 iPhone 设备上将程序运行起来。6GiB 听起来很多,但如果你在 6GiB 设备上使用超过 2.8GiB,或在 4GiB 设备上使用超过 2GiB,iOS 就会杀死你的应用程序。
那么 Stable Diffusion 模型究竟需要多少内存来进行推理?
这还要从模型的结构说起。通常 Stable Diffusion 模型包含 4 个部分:1. 文本编码器,它生成文本特征向量以指导图像生成;2. 可选的图像编码器,将图像编码到潜在空间 (用于图像到图像生成);3. 降噪器模型,它从噪声中缓慢地去噪图像的潜在表示;4. 图像解码器,从潜在表示中解码图像。
第 1、第 2 和第 4 个模块在推理过程中运行一次,最大需要约 1GiB。而降噪器模型占用了大约 3.2GiB(全浮点数),并且还需要执行多次,因此作者想让该模块在 RAM 中保存得更久。
最初的 Stable Diffusion 模型需要接近 10GiB 才能执行单个图像推理。在单个输入(2x4x64x64)与输出(2x4x64x64)之间,其中夹杂着许多输出层。并不是所有层的输出都可以被立即复用,它们中一部分必须保留一些参数以供后续使用(残差网络)。
一段时间以来,研究者围绕 PyTorch Stable Diffusion 进行了一番优化,对 PyTorch 用到的 NVIDIA CUDNN 和 CUBLAS 库,他们保留了暂存空间,这些优化都是为了降低内存使用量,因此 Stable Diffusion 模型可以用低至 4GiB 的卡运行。
但这仍然超出了作者的预期。因此作者开始专注于苹果硬及优化。
起初作者考虑的是 3.2GiB 或 1.6GiB 半浮点数,如果不想触发苹果的 OOM(Out of Memory,指的是 App 占用的内存达到了 iOS 系统对单个 App 占用内存上限后,而被系统强杀掉的现象),作者大约有 500MiB 的空间可以使用。
第一个问题,每个中间输出的大小到底是多少?
事实证明,它们中的大多数都相对较小,每个都低于 6MiB (2x320x64x64)。作者使用的框架 (s4nnc) 可以合理地将它们打包到小于 50MiB,以备复用。
值得一提的是,降噪器有一个自注意机制,它以自己的图像潜在表示作为输入。在自注意力计算期间,有一个大小为 16x4096x4096 的批处理矩阵,对该矩阵应用 softmax 后,大约是 FP16 中的 500MiB,并且可以「inplace」完成,这意味着它可以安全地重写其输入而不会损坏。幸运的是,Apple 和 NVIDIA 低级库都提供了 inplace softmax 实现,然而 PyTorch 等更高级的库中没有。
那么是否真的使用 550MiB + 1.6GiB 左右的内存就能完成?
在 Apple 硬件上,实现神经网络后端的一个常用选择是使用 MPSGraph 框架。于是作者首先尝试使用 MPSGraph 实现了所有的神经网络操作。在 FP16 精度下峰值内存使用量大约是 6GiB,显然比预期的内存使用量多太多,这是怎么回事?
作者详细分析了原因,首先他没有按照常见的 TensorFlow 方式使用 MPSGraph。MPSGraph 需要对整个计算图进行编码,然后使用输入 / 输出张量,进而处理内部分配,并让用户提交整个图以供执行。
而作者使用 MPSGraph 的方式很像 PyTorch 的做法——当作一个操作执行引擎。为了执行推理任务,许多已编译的 MPSGraphExecutable 在 Metal 命令队列上执行,它们中的每一个都可能持有一些中间分配内存。如果一次性提交,那么所有这些命令都持有分配内存,直到完成执行。
一种解决这个问题的简单方法是调整提交速度,没有必要一次性提交所有命令。实际上,Metal 的每个队列有 64 个并发提交的限制。作者尝试改成一次提交 8 个操作,峰值内存就降低到了 4GiB。
然而,这仍然比 iPhone 能承受的多 2 GiB。
为了使用 CUDA 计算自注意力,原始 Stable Diffusion 代码实现中有一个常见技巧:使用置换而不是转置。这个技巧很有效,因为 CUBLAS 可以直接处理置换的跨步(strided)张量,避免使用专用内存来转置张量。
但是 MPSGraph 没有跨步张量支持,一个置换的张量无论如何都会在内部被转置,这需要中间分配内存。通过显式转置,分配将由更高级别的层处理,避免了 MPSGraph 内部效率低下。利用这个技巧,内存使用量将接近 3GiB。
事实证明,从 iOS 16.0 开始,MPSGraph 不能再为 softmax 做出最优分配决策。即使输入和输出张量都指向相同的数据,MPSGraph 也会分配一个额外的输出张量,然后将结果复制到指向的位置。
作者发现使用 Metal Performance Shaders 替代方案完全符合要求,并将内存使用量降至 2.5GiB,而不会出现任何性能下降。
另一方面,MPSGraph 的 GEMM 内核需要内部转置。显式转置在此也无济于事,因为这些转置不是更高级别层的「inplace」操作,对于特定的 500MiB 大小的张量,这种额外的分配是不可避免的。通过切换到 Metal Performance Shaders,项目作者又回收了 500MiB,性能损失约为 1%,最终将内存使用量减到了理想的 2GiB。