在 iPhone 上運(yùn)行 Stable Diffusion 到底難不難？今天我們要介紹的這篇文章，作者給出了答案：不難，而且 iPhone 還剩余 50% 的性能。

眾所周知，每年蘋果都會(huì)推出一款聲稱在各方面都更快、更好的新款 iPhone，這主要得益于新的視覺模型和圖像傳感器的快速發(fā)展。就拿拍照來說，如果回到 10 年前，你能用 iPhone 拍出高質(zhì)量的圖片嗎，答案是不能，因?yàn)榧夹g(shù)的發(fā)展是漸進(jìn)式的，10 年時(shí)間，足夠提高手機(jī)拍照技術(shù)。

由于技術(shù)的這種發(fā)展模式（漸進(jìn)式），在一段時(shí)間里有些程序即使運(yùn)行在最好的計(jì)算設(shè)備上，也幾乎無法使用。但是這些帶有新啟用場(chǎng)景的新程序吸引了一些用戶的注意力，人們?cè)敢庋芯克?/span>

本文的作者就是被吸引的其中之一，在過去 3 周里，作者開發(fā)了一個(gè)應(yīng)用程序，可以通過 Stable Diffusion 來生成（summon）圖像，然后按你喜歡的方式編輯它。該應(yīng)用在最新的 iPhone 14 Pro 上生成圖片僅需一分鐘，使用大約 2GiB 的應(yīng)用內(nèi)存，另外還需要下載大約 2GiB 的初始數(shù)據(jù)才能開始使用。

應(yīng)用商店鏈接：https://apps.apple.com/us/app/draw-things-ai-generation/id6444050820

這一結(jié)果引來眾多網(wǎng)友討論，有人開始擔(dān)心手機(jī)耗電問題，并開玩笑的說：這很酷，不過這看起來是一個(gè)消耗手機(jī)電池的好方法。

「我從來沒有像現(xiàn)在這樣開心地感受 iPhone 的熱度?！?/span>

「這個(gè)寒冬，可以把手機(jī)當(dāng)暖手器用了?！?/span>

不過在大家調(diào)侃手機(jī)發(fā)熱問題的同時(shí)，他們也給與這項(xiàng)工作極高的評(píng)價(jià)。

「這簡(jiǎn)直不可思議。在我的 iPhone SE3 上生成一張完整的圖像大約需要 45 秒——這幾乎和我的 M1 Pro macbook 用原始版本生成的速度一樣快！」

內(nèi)存、硬件同時(shí)優(yōu)化

這是如何做到的呢？接下來我們看看作者的實(shí)現(xiàn)過程：

想要完成在 iPhone 上運(yùn)行 Stable Diffusion，還能結(jié)余 50% 的性能，面臨的一大挑戰(zhàn)是需要在 6GiB RAM 的 iPhone 設(shè)備上將程序運(yùn)行起來。6GiB 聽起來很多，但如果你在 6GiB 設(shè)備上使用超過 2.8GiB，或在 4GiB 設(shè)備上使用超過 2GiB，iOS 就會(huì)殺死你的應(yīng)用程序。

那么 Stable Diffusion 模型究竟需要多少內(nèi)存來進(jìn)行推理？

這還要從模型的結(jié)構(gòu)說起。通常 Stable Diffusion 模型包含 4 個(gè)部分：1. 文本編碼器，它生成文本特征向量以指導(dǎo)圖像生成；2. 可選的圖像編碼器，將圖像編碼到潛在空間 (用于圖像到圖像生成)；3. 降噪器模型，它從噪聲中緩慢地去噪圖像的潛在表示；4. 圖像解碼器，從潛在表示中解碼圖像。

第 1、第 2 和第 4 個(gè)模塊在推理過程中運(yùn)行一次，最大需要約 1GiB。而降噪器模型占用了大約 3.2GiB（全浮點(diǎn)數(shù)），并且還需要執(zhí)行多次，因此作者想讓該模塊在 RAM 中保存得更久。

最初的 Stable Diffusion 模型需要接近 10GiB 才能執(zhí)行單個(gè)圖像推理。在單個(gè)輸入（2x4x64x64）與輸出（2x4x64x64）之間，其中夾雜著許多輸出層。并不是所有層的輸出都可以被立即復(fù)用，它們中一部分必須保留一些參數(shù)以供后續(xù)使用（殘差網(wǎng)絡(luò)）。

一段時(shí)間以來，研究者圍繞 PyTorch Stable Diffusion 進(jìn)行了一番優(yōu)化，對(duì) PyTorch 用到的 NVIDIA CUDNN 和 CUBLAS 庫，他們保留了暫存空間，這些優(yōu)化都是為了降低內(nèi)存使用量，因此 Stable Diffusion 模型可以用低至 4GiB 的卡運(yùn)行。

但這仍然超出了作者的預(yù)期。因此作者開始專注于蘋果硬及優(yōu)化。

起初作者考慮的是 3.2GiB 或 1.6GiB 半浮點(diǎn)數(shù)，如果不想觸發(fā)蘋果的 OOM（Out of Memory，指的是 App 占用的內(nèi)存達(dá)到了 iOS 系統(tǒng)對(duì)單個(gè) App 占用內(nèi)存上限后，而被系統(tǒng)強(qiáng)殺掉的現(xiàn)象），作者大約有 500MiB 的空間可以使用。

第一個(gè)問題，每個(gè)中間輸出的大小到底是多少？

事實(shí)證明，它們中的大多數(shù)都相對(duì)較小，每個(gè)都低于 6MiB (2x320x64x64)。作者使用的框架 (s4nnc) 可以合理地將它們打包到小于 50MiB，以備復(fù)用。

值得一提的是，降噪器有一個(gè)自注意機(jī)制，它以自己的圖像潛在表示作為輸入。在自注意力計(jì)算期間，有一個(gè)大小為 16x4096x4096 的批處理矩陣，對(duì)該矩陣應(yīng)用 softmax 后，大約是 FP16 中的 500MiB，并且可以「inplace」完成，這意味著它可以安全地重寫其輸入而不會(huì)損壞。幸運(yùn)的是，Apple 和 NVIDIA 低級(jí)庫都提供了 inplace softmax 實(shí)現(xiàn)，然而 PyTorch 等更高級(jí)的庫中沒有。

那么是否真的使用 550MiB + 1.6GiB 左右的內(nèi)存就能完成？

在 Apple 硬件上，實(shí)現(xiàn)神經(jīng)網(wǎng)絡(luò)后端的一個(gè)常用選擇是使用 MPSGraph 框架。于是作者首先嘗試使用 MPSGraph 實(shí)現(xiàn)了所有的神經(jīng)網(wǎng)絡(luò)操作。在 FP16 精度下峰值內(nèi)存使用量大約是 6GiB，顯然比預(yù)期的內(nèi)存使用量多太多，這是怎么回事？

作者詳細(xì)分析了原因，首先他沒有按照常見的 TensorFlow 方式使用 MPSGraph。MPSGraph 需要對(duì)整個(gè)計(jì)算圖進(jìn)行編碼，然后使用輸入 / 輸出張量，進(jìn)而處理內(nèi)部分配，并讓用戶提交整個(gè)圖以供執(zhí)行。

而作者使用 MPSGraph 的方式很像 PyTorch 的做法——當(dāng)作一個(gè)操作執(zhí)行引擎。為了執(zhí)行推理任務(wù)，許多已編譯的 MPSGraphExecutable 在 Metal 命令隊(duì)列上執(zhí)行，它們中的每一個(gè)都可能持有一些中間分配內(nèi)存。如果一次性提交，那么所有這些命令都持有分配內(nèi)存，直到完成執(zhí)行。

一種解決這個(gè)問題的簡(jiǎn)單方法是調(diào)整提交速度，沒有必要一次性提交所有命令。實(shí)際上，Metal 的每個(gè)隊(duì)列有 64 個(gè)并發(fā)提交的限制。作者嘗試改成一次提交 8 個(gè)操作，峰值內(nèi)存就降低到了 4GiB。

然而，這仍然比 iPhone 能承受的多 2 GiB。

為了使用 CUDA 計(jì)算自注意力，原始 Stable Diffusion 代碼實(shí)現(xiàn)中有一個(gè)常見技巧：使用置換而不是轉(zhuǎn)置。這個(gè)技巧很有效，因?yàn)?CUBLAS 可以直接處理置換的跨步（strided）張量，避免使用專用內(nèi)存來轉(zhuǎn)置張量。

但是 MPSGraph 沒有跨步張量支持，一個(gè)置換的張量無論如何都會(huì)在內(nèi)部被轉(zhuǎn)置，這需要中間分配內(nèi)存。通過顯式轉(zhuǎn)置，分配將由更高級(jí)別的層處理，避免了 MPSGraph 內(nèi)部效率低下。利用這個(gè)技巧，內(nèi)存使用量將接近 3GiB。

事實(shí)證明，從 iOS 16.0 開始，MPSGraph 不能再為 softmax 做出最優(yōu)分配決策。即使輸入和輸出張量都指向相同的數(shù)據(jù)，MPSGraph 也會(huì)分配一個(gè)額外的輸出張量，然后將結(jié)果復(fù)制到指向的位置。

作者發(fā)現(xiàn)使用 Metal Performance Shaders 替代方案完全符合要求，并將內(nèi)存使用量降至 2.5GiB，而不會(huì)出現(xiàn)任何性能下降。

另一方面，MPSGraph 的 GEMM 內(nèi)核需要內(nèi)部轉(zhuǎn)置。顯式轉(zhuǎn)置在此也無濟(jì)于事，因?yàn)檫@些轉(zhuǎn)置不是更高級(jí)別層的「inplace」操作，對(duì)于特定的 500MiB 大小的張量，這種額外的分配是不可避免的。通過切換到 Metal Performance Shaders，項(xiàng)目作者又回收了 500MiB，性能損失約為 1%，最終將內(nèi)存使用量減到了理想的 2GiB。