PyTorch團(tuán)隊(duì)重寫「分割一切」模型，比原始實(shí)現(xiàn)快八倍
從年初到現(xiàn)在，生成式 AI 發(fā)展迅猛。但很多時(shí)候，我們又不得不面臨一個(gè)難題：如何加快生成式 AI 的訓(xùn)練、推理等，尤其是在使用 PyTorch 的情況下。
本文 PyTorch 團(tuán)隊(duì)的研究者為我們提供了一個(gè)解決方案。文章重點(diǎn)介紹了如何使用純?cè)?PyTorch 加速生成式 AI 模型，此外，文章還介紹了 PyTorch 新功能，以及如何組合這些功能的實(shí)際示例。
結(jié)果如何呢？PyTorch 團(tuán)隊(duì)表示，他們重寫了 Meta 的「分割一切」 (SAM) 模型，從而使代碼比原始實(shí)現(xiàn)快 8 倍，并且沒有損失準(zhǔn)確率，所有這些都是使用原生 PyTorch 進(jìn)行優(yōu)化的。
博客地址：https://pytorch.org/blog/accelerating-generative-ai/
看完本文，你將了解到：
- Torch.compile：PyTorch 模型編譯器， PyTorch 2.0 加入了一個(gè)新的函數(shù)，叫做 torch.compile ()，能夠通過一行代碼對(duì)已有的模型進(jìn)行加速；
- GPU 量化：通過降低運(yùn)算精度來加速模型；
- SDPA（Scaled Dot Product Attention ）：內(nèi)存高效的注意力實(shí)現(xiàn)方式；
- 半結(jié)構(gòu)化 (2:4) 稀疏性：一種針對(duì) GPU 優(yōu)化的稀疏內(nèi)存格式；
- Nested Tensor：Nested Tensor 把 {tensor, mask} 打包在一起，將非均勻大小的數(shù)據(jù)批處理到單個(gè)張量中，例如不同大小的圖像；
- Triton 自定義操作：使用 Triton Python DSL 編寫 GPU 操作，并通過自定義操作符注冊(cè)輕松將其集成到 PyTorch 的各種組件中。
PyTorch 原生特性所帶來的吞吐量增加以及減少的內(nèi)存開銷。
SAM 由 Meta 提出，關(guān)于這項(xiàng)研究的更多內(nèi)容請(qǐng)參考「CV 不存在了？Meta 發(fā)布「分割一切」AI 模型，CV 或迎來 GPT-3 時(shí)刻」。
接下來，文章介紹了 SAM 優(yōu)化過程，包括性能分析、瓶頸識(shí)別，以及如何將這些新功能整合進(jìn) PyTorch 以解決 SAM 面臨的這些問題。除此以外，本文還介紹了 PyTorch 的一些新特性：torch.compile、SDPA、Triton kernels、Nested Tensor 以及 semi-structured sparsity（半結(jié)構(gòu)化稀疏）。
本文內(nèi)容逐層深入，文章的最后會(huì)介紹快速版 SAM，感興趣的小伙伴可以去 GitHub 上下載，此外，本文還通過 Perfetto UI 對(duì)這些數(shù)據(jù)進(jìn)行了可視化，以此來闡釋 PyTorch 每項(xiàng)特性的應(yīng)用價(jià)值。
GitHub 地址：https://github.com/pytorch-labs/segment-anything-fast
對(duì)分割一切模型 SAM 的重寫
該研究表示，本文利用的 SAM 基線數(shù)據(jù)類型為 float32 dtype、batch 大小為 1，使用 PyTorch Profiler 查看內(nèi)核跟蹤的結(jié)果如下：
本文發(fā)現(xiàn) SAM 有兩個(gè)地方可以優(yōu)化：
第一個(gè)是對(duì) aten::index 的長調(diào)用，這是由張量索引操作（例如 []）產(chǎn)生的底層調(diào)用導(dǎo)致的。然而實(shí)際上 GPU 花費(fèi)在 aten::index 上的時(shí)間相對(duì)較低，原因在于 aten::index 在啟動(dòng)兩個(gè)內(nèi)核的過程中，兩者之間發(fā)生了阻塞 cudaStreamSynchronize。這意味著 CPU 會(huì)等待 GPU 完成處理，直到啟動(dòng)第二個(gè)內(nèi)核。因而為了優(yōu)化 SAM，本文認(rèn)為應(yīng)該致力于消除導(dǎo)致空閑時(shí)間的阻塞 GPU 同步。
第二個(gè)是 SAM 在矩陣乘法中花費(fèi)了大量的 GPU 時(shí)間（上圖中的深綠色），這在 Transformers 中很常見。如果能夠減少 SAM 模型在矩陣乘法上花費(fèi)的 GPU 時(shí)間，我們就可以顯著加快 SAM 的速度。
接下來本文用 SAM 的吞吐量 (img/s) 和內(nèi)存開銷 (GiB) 來建立基線。之后就是優(yōu)化過程了。
Bfloat16 半精度（加上 GPU 同步和批處理）
為了解決上述問題，即讓矩陣乘法花費(fèi)的時(shí)間更少，本文轉(zhuǎn)向 bfloat16。Bfloat16 是常用的半精度類型，通過降低每個(gè)參數(shù)和激活的精度，能夠節(jié)省大量的計(jì)算時(shí)間和內(nèi)存。
用 bfloat16 替換 padding 類型
此外，為了移除 GPU 同步，本文發(fā)現(xiàn)有兩個(gè)位置可以優(yōu)化。
具體來說（參考上圖更容易理解，出現(xiàn)的變量名都在代碼中），該研究發(fā)現(xiàn)在 SAM 的圖像編碼器中，有充當(dāng)坐標(biāo)縮放器（coordinate scalers）的變量 q_coords 和 k_coords，這些變量都是在 CPU 上分配和處理的。然而，一旦這些變量被用來在 rel_pos_resized 中建立索引，這些索引操作就會(huì)自動(dòng)的將這些變量移動(dòng)到 GPU 上，這種復(fù)制會(huì)導(dǎo)致 GPU 同步。為了解決上述問題，該研究注意到可以使用 torch.where 重寫這部分內(nèi)容來解決問題，如上所示。
內(nèi)核跟蹤
在應(yīng)用了這些更改之后，本文注意到單個(gè)內(nèi)核調(diào)用之間有著顯著的時(shí)間間隔，尤其在小批量（這里為 1）時(shí)更為突出。為了更深入的了解這一現(xiàn)象，本文開始對(duì)批大小為 8 的 SAM 推理進(jìn)行性能分析：
在查看每個(gè)內(nèi)核所花費(fèi)的時(shí)間時(shí)，本文觀察到 SAM 的大部分 GPU 時(shí)間都花費(fèi)在逐元素內(nèi)核（elementwise kernels）和 softmax 操作上。
現(xiàn)在可以看到矩陣乘法的相對(duì)開銷小了很多。
將 GPU 同步和 bfloat16 優(yōu)化結(jié)合在一起，SAM 性能提高了 3 倍。
Torch.compile（+graph breaks 和 CUDA graphs）
本文發(fā)現(xiàn)在深入研究 SAM 的過程中有很多小的操作，他們認(rèn)為使用編譯器來融合操作有很大的好處，因而 PyTorch 對(duì) torch.compile 做了以下優(yōu)化：
- 將 nn.LayerNorm 或 nn.GELU 等操作序列融合成一個(gè)單一的 GPU 內(nèi)核；
- 融合緊跟在矩陣乘法內(nèi)核之后的操作，以減少 GPU 內(nèi)核調(diào)用的數(shù)量。
通過這些優(yōu)化，該研究減少了 GPU 全局內(nèi)存往返次數(shù)（roundtrips），從而加快了推理速度。我們現(xiàn)在可以在 SAM 的圖像編碼器上嘗試 torch.compile。為了最大限度地提高性能，本文使用了一些高級(jí)編譯技術(shù)：
內(nèi)核跟蹤
結(jié)果顯示，torch.compile 工作得很好。
可以觀察到 softmax 占了很大一部分時(shí)間，然后是各種 GEMM 變體。以下測量的是批大小為 8 及以上的變化。
SDPA: scaled_dot_product_attention
接下來，本文又對(duì) SDPA（scaled_dot_product_attention）進(jìn)行了實(shí)驗(yàn)，研究的重點(diǎn)是注意力機(jī)制。一般來講，原生注意力機(jī)制在時(shí)間和內(nèi)存上隨序列長度呈二次方擴(kuò)展。PyTorch 的 SDPA 操作基于 Flash Attention、FlashAttentionV2 和 xFormer 的內(nèi)存高效注意力原理構(gòu)建，可以顯著加快 GPU 注意力。與 torch.compile 相結(jié)合，這個(gè)操作允許在 MultiheadAttention 的變體中表達(dá)和融合一個(gè)共同的模式。經(jīng)過一小部分更改后，現(xiàn)在模型可以使用 scaled_dot_product_attention。
內(nèi)核跟蹤
現(xiàn)在可以看到內(nèi)存高效的注意力內(nèi)核占用了 GPU 上大量的計(jì)算時(shí)間：
使用 PyTorch 的原生 scaled_dot_product_attention，可以顯著增加批處理大小。下圖為批大小為 32 及以上的變化。
之后，該研究又實(shí)驗(yàn)了 Triton，NestedTensor 、批處理 Predict_torch， int8 量化，半結(jié)構(gòu)化 (2:4) 稀疏性等操作。
例如本文使用自定義 positional Triton 內(nèi)核，觀察到批大小為 32 的測量結(jié)果。
使用 Nested Tensor，批大小為 32 及以上的變化。
添加量化后，批大小為 32 及以上變化的測量結(jié)果。
文章的最后是半結(jié)構(gòu)化稀疏性。該研究表示，矩陣乘法仍然是需要面對(duì)的一個(gè)瓶頸。解決的辦法是使用稀疏化來近似矩陣乘法。通過稀疏矩陣（即將值歸零）可以使用更少的位來存儲(chǔ)權(quán)重和激活張量。該研究將張量中哪些權(quán)重設(shè)置為零的過程稱為剪枝。剪枝掉較小的權(quán)重可以潛在地減小模型大小，而不會(huì)顯著損失準(zhǔn)確率。
剪枝的方法多種多樣，從完全非結(jié)構(gòu)化到高度結(jié)構(gòu)化。雖然非結(jié)構(gòu)化剪枝理論上對(duì)精度的影響最小，但 GPU 在進(jìn)行大型密集矩陣乘法方面盡管非常高效，然而在稀疏情況下可能還會(huì)遭受顯著的性能下降。PyTorch 最近支持的一種剪枝方法旨在尋求平衡，稱為半結(jié)構(gòu)化（或 2:4）稀疏性。這種稀疏存儲(chǔ)將原始張量減少了 50%，同時(shí)產(chǎn)生密集張量輸出。參見下圖的說明。
為了使用這種稀疏存儲(chǔ)格式和相關(guān)的快速內(nèi)核，接下來要做的是剪枝權(quán)重。本文在 2：4 的稀疏度下選擇最小的兩個(gè)權(quán)重進(jìn)行剪枝，將權(quán)重從默認(rèn)的 PyTorch（“strided”）布局更改為這種新的半結(jié)構(gòu)化稀疏布局很容易。要實(shí)現(xiàn) apply_sparse (model)，只需要 32 行 Python 代碼：
在 2:4 的稀疏度下，本文觀察到 vit_b 和批大小為 32 時(shí)的 SAM 峰值性能：
最后，一句話總結(jié)這篇文章：本文介紹了迄今為止在 PyTorch 上最快的 Segment Anything 實(shí)現(xiàn)方式，借助官方發(fā)布的一系列新功能，本文在純 PyTorch 中重寫了原始 SAM，并且沒有損失準(zhǔn)確率。
感興趣的讀者可以查看原博客了解更多內(nèi)容。
