PyTorch团队重写「分割一切」模型，比原始实现快八倍 _PyTorch

编辑：陈萍

我们该如何优化 Meta 的「分割一切」模型，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/ target=_blank class=infotextkey>Python DSL 编写 GPU 操作，并通过自定义操作符注册轻松将其集成到 PyTorch 的各种组件中。

【PyTorch团队重写「分割一切」模型，比原始实现快八倍】

文章插图
PyTorch 原生特性所带来的吞吐量增加以及减少的内存开销。
SAM 由 Meta 提出，关于这项研究的更多内容请参考「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、batch 大小为 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) 来建立基线。之后就是优化过程了。