PyTorch團隊重新實現「分割一切」模型，速度比原始實現提升八倍

從年初到現在，生成式 AI 發展迅速。但很多時候，我們又必須面對一個難題：如何加速生成式 AI 的訓練、推理等，尤其是在使用 PyTorch 的情況下。

本文 PyTorch 團隊的研究者為我們提供了一個解決方案。文章重點介紹如何使用純原生 PyTorch 加速生成式 AI 模型，此外，文章還介紹了 PyTorch 新功能，以及如何組合這些功能的實際範例。

結果如何呢？ PyTorch 團隊表示，他們重寫了 Meta 的「分割一切」 (SAM) 模型，從而使程式碼比原始實現快 8 倍，並且沒有損失準確率，所有這些都是使用原生 PyTorch 進行優化的。

部落格網址：https://pytorch.org/blog/accelerating-generative-ai/

在閱讀本文後，你將會獲得以下的了解：

• Torch.compile：PyTorch 模型編譯器， PyTorch 2.0 加入了一個新的函數，叫做torch .compile ()，能夠透過一行程式碼對現有的模型進行加速；
• GPU 量化：透過降低運算精度來加速模型；
• ##SDPA（Scaled Dot Product Attention ）：記憶體高效的注意力實作方式；
• 半結構化(2:4) 稀疏性：一種針對GPU 最佳化的稀疏記憶體格式;
• Nested Tensor：Nested Tensor 把{tensor, mask} 打包在一起，將非均勻大小的資料批次處理到單張量中，例如不同大小的圖片；
• Triton 自訂操作：使用Triton Python DSL 編寫GPU 操作，並透過自訂操作符註冊輕鬆將其整合到PyTorch 的各種元件中。

PyTorch 原生特性所帶來的吞吐量增加以及減少的記憶體開銷。

有關此研究的更多信息，請參考Meta提出的SAM。詳細文章可在「CV不存在了？Meta發布「分割一切」AI模型，CV或迎來GPT-3時刻」找到

# #接下來，我們將介紹SAM的最佳化過程，包括效能分析、瓶頸識別，以及如何將這些新功能整合進PyTorch以解決SAM所面臨的問題。此外，我們也會介紹PyTorch的一些新特性，包括torch.compile、SDPA、Triton kernels、Nested Tensor以及semi-structured sparsity（半結構化稀疏）

內容的逐層深入，本文最後將介紹快速版SAM。對於有興趣的讀者，可以前往 GitHub 下載。此外，透過使用Perfetto UI 對這些數據進行了視覺化，以展示PyTorch 各項特性的應用價值

GitHub 位址：https://github.com/pytorch-labs/segment -anything-fast 可以找到這個專案的原始碼

對分割一切模型SAM 的重寫

該研究指出，本文所使用的SAM基準資料類型為float32 dtype，批次大小為1，並使用PyTorch Profiler來查看核心追蹤的結果如下：

############################################# ###本文發現SAM 有兩個地方可以優化：############第一個是對aten::index 的長調用，這是由張量索引操作（例如[]）產生的底層調用導致的。然而實際上 GPU 花在 aten::index 上的時間相對較低，原因在於 aten::index 在啟動兩個核心的過程中，兩者之間發生了阻塞 cudaStreamSynchronize。這意味著 CPU 會等待 GPU 完成處理，直到啟動第二個核心。因而為了優化 SAM，本文認為應該致力於消除導致空閒時間的阻塞 GPU 同步。 ######

第二个问题是在矩阵乘法中，SAM花费了大量的GPU时间（如图所示的深绿色部分），这在Transformers模型中非常普遍。如果我们能够减少SAM模型在矩阵乘法上的GPU时间，那么我们就能够显著提高SAM的速度

接下来，我们将以SAM的吞吐量（img/s）和内存开销（GiB）来建立基准。然后就是优化过程

需要进行改写的句子是：Bfloat16 半精度（加上 GPU 同步和批处理）

为了解决上述问题，即减少矩阵乘法所需的时间，本文转向bfloat16。bfloat16是常用的半精度类型，通过降低每个参数和激活的精度，能够节省大量的计算时间和内存

将填充类型替换为 bfloat16

此外，本文发现有两个位置可以进行优化，以移除 GPU 同步

具体来说，根据上图更容易理解，该研究发现在SAM的图像编码器中，有两个变量q_coords和k_coords充当坐标缩放器，这些变量都在CPU上进行分配和处理。然而，一旦这些变量用于在rel_pos_resized中建立索引，索引操作会自动将这些变量移动到GPU上，从而导致GPU同步的问题。为了解决这个问题，该研究指出可以使用torch.where函数重写这部分内容来解决问题，具体如上所示

核心追踪

在对这些更改进行应用之后，我们注意到单个内核调用之间存在明显的时间间隔，特别是在小批量（这里为1）的情况下更为明显。为了更深入地了解这一现象，我们开始对批大小为8的SAM推理进行性能分析

在分析每个内核所花费的时间时，我们注意到 SAM 的大部分 GPU 时间都用于逐元素内核和 softmax 操作

现在可以看到矩阵乘法的相对开销小了很多。

将 GPU 同步和 bfloat16 优化结合在一起，SAM 性能提高了 3 倍。

Torch.compile（ graph breaks 和 CUDA graphs）

在研究SAM的过程中发现了许多细小的操作。研究人员认为使用编译器来整合这些操作非常有益，因此PyTorch对torch.compile进行了以下优化

• 将 nn.LayerNorm 或 nn.GELU 等操作序列融合成一个单一的 GPU 内核；
• 融合紧跟在矩阵乘法内核之后的操作，以减少 GPU 内核调用的数量。

通过这些优化，该研究减少了 GPU 全局内存往返次数（roundtrips），从而加快了推理速度。我们现在可以在 SAM 的图像编码器上尝试 torch.compile。为了最大限度地提高性能，本文使用了一些高级编译技术：

核心追蹤

#根據結果顯示，torch.compile 的表現非常出色

#可以觀察到softmax 佔了很大一部分時間，然後是各種GEMM 變體。以下測量的是批次大小為 8 及以上的變化。

SDPA: scaled_dot_product_attention

##接下來，本文又對SDPA（ scaled_dot_product_attention）進行了實驗，研究的重點是注意力機制。一般來講，原生注意力機制在時間和記憶體上隨序列長度呈二次方擴展。 PyTorch 的 SDPA 操作基於 Flash Attention、FlashAttentionV2 和 xFormer 的記憶體高效注意力原理構建，可以顯著加快 GPU 注意力。與 torch.compile 結合，這個操作允許在 MultiheadAttention 的變體中表達和融合一個共同的模式。經過一小部分變更後，現在模型可以使用 scaled_dot_product_attention。

#核心追蹤

#現在可以看到記憶體高效的注意力核心佔用了GPU 上大量的運算時間：

使用PyTorch 的原生scaled_dot_product_attention，可以顯著增加批次大小。下圖為批次大小為 32 以上的變化。

接下來，研究進行了對Triton、NestedTensor、批次Predict_torch、int8 量化、半結構化(2:4) 稀疏性等操作的實驗

例如本文使用自訂positional Triton 內核，觀察到批次大小為32 的測量結果。

採用Nested Tensor 技術，並調整批次大小為32 以上

##添加量化後，批大小為32 以上變化的測量結果。

文章的最後是半結構化稀疏性。研究表示，矩陣乘法仍然是需要面對的瓶頸。解決的辦法是使用稀疏化來近似矩陣乘法。透過稀疏矩陣（即將值歸零）可以使用更少的位元來儲存權重和激活張量。該研究將張量中哪些權重設為零的過程稱為剪枝。剪枝掉較小的權重可以潛在地減少模型大小，而不會顯著損失準確率。

剪枝的方法有很多种，从完全非结构化到高度结构化都有。虽然理论上来说非结构化剪枝对精度的影响最小，但是在稀疏情况下，GPU可能会遇到显著的性能下降，尽管在进行大型密集矩阵乘法时非常高效。最近PyTorch支持的一种剪枝方法是半结构化（或2:4）稀疏性，旨在寻求平衡。这种稀疏存储方式将原始张量减少了50%，同时产生密集张量的输出。请参考下图进行说明

为了使用这种稀疏存储格式和相关的快速内核，接下来要做的是剪枝权重。本文在 2：4 的稀疏度下选择最小的两个权重进行剪枝，将权重从默认的 PyTorch（“strided”）布局更改为这种新的半结构化稀疏布局很容易。要实现 apply_sparse (model)，只需要 32 行 Python 代码：

在稀疏度为2:4的情况下，我们观察到vit_b和批大小为32时的SAM峰值性能

最终，对这篇文章的概括如下：本文介绍了截至目前在PyTorch上实现Segment Anything的最快方法，借助官方发布的一系列新功能，本文在纯PyTorch中重新编写了原始的SAM，并且没有损失准确度

对于感兴趣的读者，可以查看原博客以获取更多信息

以上是PyTorch團隊重新實現「分割一切」模型，速度比原始實現提升八倍的詳細內容。更多資訊請關注PHP中文網其他相關文章！