本文是我们关于如何使用纯原生 PyTorch 加速生成式 AI 模型的多系列博客的第一篇文章的后续，重点关注延迟和弹性可扩展性。我们使用 torch.compile 和 torch.export 来创建 SAM2 的高度优化低延迟版本，该版本可以在新实例上快速扩展。
通过利用 AOTInductor (AOTI) 通过 torch.export 进行提前编译、降低精度、批量提示和 GPU 预处理，我们观察到与常规急切模式 PyTorch 相比，p90 执行延迟和排队时间最多提高了 13 倍。
我们计算了最终结果，并在 Modal 提供的自动伸缩云基础设施上展示了在实际部署中的改进。
| p50 执行延迟 （毫秒 / 改进） |
p90 执行延迟 （毫秒 / 改进） |
|||
| 急切 float32 | AOTI float16 | 急切 float32 | AOTI float16 | |
| AMG | 741 | 112 (6.6x) | 1140 | 176 (6.5x) |
| SPS | 98 | 20 (4.9x) | 130 | 28 (4.6x) |
| MPS | 269 | 38 (7.1x) | 714 | 52 (13.7x) |
| p50 排队时间（毫秒 / 改进） | p90 排队时间（毫秒 / 改进） | |||
| 急切 float32 | AOTI float16 | 急切 float32 | AOTI float16 | |
| AMG | 201 | 41 (4.9x) | 815 | 327 (2.6x) |
| SPS | 31 | 33 (0.9x) | 441 | 49 (9.0x) |
| MPS | 40 | 37 (1.1x) | 942 | 75 (12.6x) |
任务
第一篇文章重点关注每张图像处理少量不同提示（兴趣点）。这些点代表了真实掩码的中心点。对于本文，我们现在将关注更广泛的任务集。单提示分割 (SPS)、多提示分割 (MPS)、自动掩码生成 (AMG)，它在不给定提示集的情况下为输入图像生成完整的掩码集。第一篇文章只关注 MPS。
图像中的小星星代表用户提示。对于 AMG，没有提示，掩码从初始候选提示（猜测）的密集网格中启发式过滤。对于 SPS 和 MPS，用户提示来自 AMG 掩码的中心点。对于 SPS，我们选择面积最大的掩码。
请注意，SAM2 使用与 SAM1 不同的主干。特别是，对于本博客，我们只考虑最大且最准确的 sam2.1_hiera_large 主干。
我们将重现结果所需的脚本聚合到torchao 的示例文件夹中，并逐步将torchao 中 SAM2 模型的更稳定部分上游到主要的 SAM2 存储库。因此，如果您有兴趣查看尖端变体或希望贡献实验性功能，请随时联系 torchao 存储库和团队。有关更稳定和最新的模型版本，请直接访问 SAM2。
概述
我们将这里提出的更改分为两类。快速更改仅限于不影响模型准确性的技术。猛烈更改通过使用低精度数据类型等近似值来牺牲一些数值准确性以获得额外速度。
近似值可能会略微降低精度指标，以换取显著的性能改进，同时仍通过基于平均交并集 (mIoU) 的端到端检查。
为了衡量性能改进，我们处理了 1000 张图像，这些图像是从 SAM2 验证数据集中随机选择的。我们查看每张图像的 p50 和 p90 延迟。为了衡量准确性，我们考虑 mIoU。最值得注意的是，对于 AMG 任务，我们还定义了一个失败计数指标。如果掩码数量不同，我们认为比较失败。这被证明是一个相当不稳定的量，我们可以看到其他任务不如 AMG 对小的数值变化敏感。
设置
我们正在常规 H100 开发服务器上运行离线实验，这是一台相当强大且性能良好的机器。
然而，我们尝试在实际约束下查看这些任务。特别是，我们希望模拟服务器端推理环境。这意味着我们不使用 DataLoader 来隐藏图像预处理或解码例程的延迟。
对于延迟计算，我们包括解码、分割以及将掩码转换为运行长度编码掩码字典。换句话说，我们排除了将图像加载到内存主机字节数组中以及将生成的字典存储为磁盘上的 json 文件。这旨在模拟更真实的设置。
更具体地说，考虑下面我们测量中包含的例程的代码。对于任何任务，gen_masks 都会生成一个批处理的布尔张量位掩码，表示相应的对象掩码。然后，我们将此位掩码压缩为运行长度编码 (rle) 格式，该格式可以更有效地将结果从远程服务器传回。
image_tensors = decode_img_bytes(...)
masks = gen_masks(image_tensors, ...)
rle_dicts = [rle_dict_from_masks(m) for m in masks]
优化
ao：急切代码优化
这项工作最有效的工具是 PyTorch autograd 分析器与 record_function 的结合使用。为了构建此软件，我们反复使用分析器来观察程序并确认任何更改的有效性。同样重要的是要记住，分析器本身具有开销。您收集的数据越多，例如堆栈跟踪，引入的开销就越大，这可能会使收集到的跟踪失真。但它非常适合查找同步点、内核之间的空间以及耗时长的 GPU 内核。
GPU 跟踪可帮助您了解编译不一定容易解决的瓶颈。我们发现，特别是 AutomaticMaskGeneration，其主要瓶颈在于用于存储掩码的数据结构以及用于将掩码转换为运行长度编码压缩格式的例程。我们还发现，AMG 性能的很大一部分瓶颈在于作为单个批次创建的大量掩码。有时，在后处理阶段，通过重新排序操作，可以更早地将候选掩码过滤到更少的候选掩码中。这反过来又显著加快了后续操作的速度。
为了确认我们实现的准确性，我们首先在不更改设置和使用 float32 精度的情况下进行比较。我们看到，在使用完全相同的设置时，mIoU 保持不变，并且掩码完美匹配。这意味着这些急切模式的更改没有影响这些任务的准确性。
AMG
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU / 失败计数 | |
| 基准 | 864 | 1144 | 4350 | 参考 |
| AO | 693 | 786 | 4010 | 1 / 0 |
ao：批处理提示
我们能够应用的另一个无损性能优化是批量处理用户输入提示计算。当在 H100 等服务器级 GPU 上以批量大小 1 优化延迟时，我们通常会剩下大量空闲内存。我们可以通过一次处理更多兴趣点（也称为用户提示）来轻松地用内存换取更多性能。请记住，SAM2 分为两部分：首先是主干（图像编码器），其次是基于一组用户提示/兴趣点进行掩码预测和解码。我们可能会期望第二部分有更大或甚至变化的输入数量，我们对第二部分应用批量处理。
这导致内存大幅增加，但延迟也大大改善。基线在循环中为每个提示生成一个掩码。对于 AMG，基线一次处理 64 个提示，所需要做的就是将其更改为 1024，即生成的候选提示数量。对于 SPS，我们一次处理一个提示，但为了完整性，它仍包含在下面。
AMG
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU / 失败计数 | |
| 基准 | 864 | 1144 | 4350 | 参考 |
| AO + 批处理 | 613 | 706 | 33786 | 0.9999995 / 0 |
SPS
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | |
| 基准 | 116 | 181 | 1337 | 参考 |
| AO | 110 | 170 | 1339 | 1 |
MPS
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | |
| 基准 | 276 | 681 | 1337 | 参考 |
| AO + 批处理 | 126 | 225 | 8021 | 0.9999992 |
作为一个技术附注：最值得注意的是，为了实现 MPS 的批处理，并避免对代码库进行大量手动重写以同时支持多个提示，我们使用了我们称之为 MapTensor 的 Tensor 子类。MapTensor 允许我们传递 N 个提示的批处理，但它会声明批处理大小为 1。然后，任何操作都会自动广播到包装的 Tensor，并在模型的预测部分传播。这之所以有效，是因为单个提示预测彼此独立。这与 torch.vmap 非常相似。
center_points_torch = to_map_tensor(center_points_torch)
center_points_label_torch = to_map_tensor(center_points_label_torch)
masks, scores, _ = mask_generator.predictor.predict(
point_coords=center_points_torch,
point_labels=center_points_label_torch,
multimask_output=True,
return_logits=False,
return_type="torch",
)
# Unwrapping MapTensor
masks = masks.elems
scores = scores.elems
快速：全图编译
正如我们的第一篇文章一样，我们首先移除 GPU 同步和图中断，以便在适当的情况下使用带有最大自动调优内核的全图编译模型代码。经过一些重写后，我们能够编译图像编码器和掩码预测。
我们运行实验两次，以了解编译带来的开销。我们在 TORCHINDUCTOR_CACHE_DIR 为空的环境中运行一次，然后在摄取上次运行的工件时再次运行。特别是，自动调优可能需要很长时间，并且在原始环境中的第一次调用时发生。我们将第二次运行称为“热”。第一次迭代通常由于各种其他相关的初始化过程而预计会很慢，但即使使用现有缓存并再次馈送完全相同的形状，编译也会显著增加它。话虽如此，在热环境中，第一次调用的几秒钟开销通常仍然可以接受。
这些缺点大多可以缓解，并且编译会显著改善延迟并减少内存。
AMG
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU / 失败计数 |
第一次迭代 （毫秒） |
|
| AO + 批处理 | 613 | 706 | 33786 | 0.9999995 / 0 | 1125 |
| + 编译（冷） | 423 | 513 | 29349 | 跳过 | 404866 |
| + 编译（热） | 439 | 530 | 29349 | 0.994 / 190 | 8544 |
使用自动掩码分割时，每个掩码生成的掩码数量可能会略有不同。模型可能为每个对象生成的掩码数量存在模糊性。例如，一辆汽车可能被细分为车架、车窗和车门，或作为一个整体处理。当修改导致掩码数量发生变化时，我们认为比较失败，我们只计算精确匹配掩码的 mIoU。这不适用于其他任务。我们发现生成的掩码数量对微小的数值变化非常敏感。其他任务使用相同的代码，特别是 MPS 可以帮助我们进一步验证正确性。
SPS
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | 第一次迭代 （毫秒） |
|
| AO | 110 | 170 | 1339 | 1 | 562 |
| + 编译（冷） | 102 | 158 | 1343 | 跳过 | 319954 |
| + 编译（热） | 100 | 160 | 1302 | 0.9999 | 8947 |
MPS
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | 第一次迭代 （毫秒） |
|
| AO + 批处理 | 126 | 225 | 8021 | 0.9999992 | 504 |
| + 编译（冷） | 129 | 215 | 8021 | 跳过 | 333308 |
| + 编译（热） | 113 | 213 | 8021 | 0.998 | 8617 |
猛烈：TF32、float16 和 GPU 预处理
我们发现，对于模型的几个重要子组件来说，使用 float16 是正确的精度级别。特别是，图像编码器和掩码解码器权重可以完全转换为 float16。我们还可以将 TensorFloat32 精度用于剩余的 float32 矩阵操作。应该可以进一步降低精度，我们可能会在未来的文章中解决这个问题。我们还在猛烈模式下将图像预处理（如图像归一化）移至 GPU。我们不能使用 GPU 解码 (nvJPEG) 例程，因为差异太大，模型在 mIoU 中会遭受显著退化，因此图像解码仍然在 CPU 上进行。
AMG
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU / 失败计数 |
|
| AO + 批处理 + 编译（热） |
439 | 530 | 29349 | 0.994 / 190 |
| + 猛烈 | 165 | 240 | 28335 | 0.978 / 306 |
这导致 AMG 任务的 mIoU 显著下降，但未影响其他任务。经过深入调查，我们仍将其归因于数值不稳定性和操作的重新排序。需要更多工作来进一步调查这一点，并且以较低精度运行 AMG 任务可能没有意义。然而，其他任务在延迟方面受益匪浅，而 mIoU 变化极小。
SPS
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | |
| AO + 编译（热） |
100 | 160 | 1302 | 0.9999 |
| + 猛烈 | 32 | 63 | 861 | 0.9997 |
MPS
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | |
| AO + 批处理 + 编译（热） |
113 | 213 | 8021 | 0.998 |
| + 猛烈 | 36 | 64 | 4222 | 0.997 |
AOTInductor (AOTI) 通过 torch.export 进行提前编译
在弹性扩展时，通常无法适应长时间的启动时间。这意味着第一次迭代不能慢，但我们必须快速交付结果。这时，torch.compile 当前的编译开销可能会妨碍。为了解决这个问题，我们可以使用 AOTInductor (AOTI) 通过 torch.export 进行提前编译。AOTI 允许我们根据代表性输入编译模型，并将生成的代码存储在可以快速加载和运行的二进制文件中。
通过 torch.export 实现 AOTI 是一项新功能，我们目前无法导出所有可编译的内容。我们已经能够为所有任务导出图像编码器，但由于提示的不同，我们只能为 AMG 和 SPS 任务导出掩码预测。torch.export 也支持动态形状，但我们需要投入更多时间来为此准备代码。
AMG: AO + 批处理 + 猛烈
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU / 失败计数 |
第一次迭代 （毫秒） |
|
| + 编译（热） | 165 | 240 | 28335 | 0.978 / 306 | 10341 |
| + 加载导出 （冷） |
162 | 233 | 27927 | 0.974 / 308 | 906 |
SPS: AO + 猛烈
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | 第一次迭代 （毫秒） |
|
| + 编译（热） | 32 | 63 | 861 | 0.9997 | 7989 |
| + 加载导出 （冷） |
35 | 66 | 1686 | 0.9997 | 763 |
请注意，加载导出的模型会显著增加内存。它可能只会增加峰值内存利用率，因为在加载导出的模型之前确实需要延迟初始化，以避免内存中同时存在两倍的权重。这是我们可以解决的问题，但内存消耗远未达到限制。我们在其他任务中没有看到增加，因为 AMG 和 MPS 的峰值内存主要由处理批量掩码支配。减少这种情况的一种方法可能是更早地以 rle 格式（或其他稀疏格式）操作掩码，但目前，鉴于当前的内存消耗和对延迟的关注，没有理由这样做。
MPS: AO + 批处理 + 猛烈
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | 第一次迭代 （毫秒） |
|
| + 编译（热） | 36 | 64 | 4222 | 0.997 | 9626 |
| + 加载导出 （冷） |
43 | 72 | 3813 | 0.997 | 747 |
单独使用导出似乎不会受益于大量预热，并且可以在原始的新 inductor 缓存目录中运行。但同样，我们不会驱逐 CUDA 缓存或其他缓存。在 Modal 部分，我们正在原始环境中运行其中一些实验。
当在新进程中只处理 1000 张图像时，使用导出确实值得，可以节省编译和其他冷启动开销。
额外福利：更多 GPU 预处理
此时，延迟相当低。特别是，对于 SPS 和 MPS 任务，我们以大约 30 毫秒到 40 毫秒的速度处理。让我们再次回顾设置部分的伪代码。
image_tensors = decode_img_bytes(...)
masks = gen_masks(image_tensors, ...)
rle_dicts = [rle_dict_from_masks(m) for m in masks]
进一步的分析表明，此时 decode_img_bytes 大约需要 10 毫秒。特别是，它使用 torchvision 的 ToTensor 转换将 numpy Tensor 转换为缩放的 float32 torch.Tensor。传递给 ToTensor 的字节已经被解码并转换为 numpy ndarray。通过稍微重写 ToTensor，使用 torchvision 的 v2 API，并先将 uint8 解码的较小整数 Tensor 移动到 GPU，然后再进行缩放，我们可以在延迟方面再减少 10 毫秒。如果不在我们的分析中包含 decode_img_bytes，我们将错过这个对服务器端推理具有实际影响的机会。
image_tensor = torch.from_numpy(image_tensor)
image_tensor = image_tensor.permute((2, 0, 1))
image_tensor = image_tensor.cuda()
image_tensor = v2.ToDtype(torch.float32, scale=True)( image_tensor)
特别注意，使用固定内存进行异步数据传输不适用，因为将 Tensor 移动到固定内存所需的时间不值得此数据移动的异步性增益。对于未来的工作，我们可能希望通过使用更高级的直接内存传输技术来进一步探索这里的改进。
AMG: AO + 批处理 + 猛烈
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU / 失败计数 |
第一次迭代 （毫秒） |
|
| + 加载导出 （冷） |
162 | 233 | 27927 | 0.974 / 308 | 906 |
| + 加载导出（热） | 157 | 230 | 27927 | 0.974 / 308 | 799 |
| + 加载导出（热） + 预处理 |
136 | 208 | 27950 | 0.977 / 311 | 908 |
SPS: AO + 猛烈
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | 第一次迭代 （毫秒） |
|
| + 加载导出 （冷） |
35 | 66 | 1686 | 0.9997 | 763 |
| + 加载导出（热） | 31 | 63 | 1686 | 0.9997 | 683 |
| + 加载导出（热） + 预处理 |
19 | 25 | 1711 | 0.9997 | 658 |
MPS: AO + 批处理 + 猛烈
| p50 延迟 (ms) | p90 延迟 (ms) | 内存 (MiB) | mIoU | 第一次迭代 （毫秒） |
|
| + 加载导出 （冷） |
43 | 72 | 3813 | 0.997 | 747 |
| + 加载导出（热） | 53 | 81 | 3813 | 0.997 | 807 |
| + 加载导出（热） + 预处理 |
31 | 41 | 3837 | 0.997 | 671 |
这一微小的改变对 SPS 和 MPS 任务产生了显著影响。
在 Modal 上部署
最后，我们将优化的推理部署到无服务器基础设施提供商 Modal 上，以证明这些优化的好处可以在更真实的部署环境中实现。
特别是，通过 torch.export 进行编译和 AOTI 需要额外的工作。在简单部署中，这些工作可能会添加到每一次推理执行中，从而增加延迟，使任何来自更快模型的改进相形见绌。这对于弹性或自动缩放基础设施来说尤其具有挑战性，在这种基础设施中，我们的推理服务的副本需要定期自动创建和销毁。
我们在 torchao 仓库中分享了一个部署脚本 (cli_on_modal.py)，以演示弹性部署的一种模式。我们提前构建导出的模型，然后将它们上传到 分布式存储。相对于急切执行，当副本启动时，这会增加一些额外的工作，因为它们需要通过网络读取这些数据，但这比编译或导出成本要低得多。
我们使用大型批量推理工作负载对此次部署进行了基准测试：发送 1000 张图像进行并发处理。部署在高峰期可扩展到 10 个 GPU 上的 10 个副本，在非活动状态时可缩减到零个 GPU。
首先，让我们看看执行延迟。
| p50 执行延迟 （毫秒 / 改进） |
p90 执行延迟 （毫秒 / 改进） |
|||||
| 急切 float32 | AOTI float16 | 急切 float32 | AOTI float16 | |||
| Modal | 离线 | Modal | 离线 | |||
| AMG | 741 | 112 (6.6x) | 136 (5.4x) | 1140 | 176 (6.5x) | 208 (5.5x) |
| SPS | 98 | 20 (4.9x) | 19 (5.2x) | 130 | 28 (4.6x) | 25 (5.2x) |
| MPS | 269 | 38 (7.1x) | 31 (8.7x) | 714 | 52 (13.7x) | 41 (17.4x) |
我们注意到 Modal 和离线上的执行延迟相当接近，特别是相对于基线，这表明离线优化部署是直接优化部署的合理替代方法。
除了执行延迟，我们的批量工作负载还有排队时间，因为副本少于输入，所以一些输入必须排队等待。
| p50 排队时间 (ms) | p90 排队时间 (ms) | |||
| 急切 float32 | AOTI float16 | 急切 float32 | AOTI float16 | |
| AMG | 201 | 41 (4.9x) | 815 | 327 (2.6x) |
| SPS | 31 | 33 (0.9x) | 441 | 49 (9.0x) |
| MPS | 40 | 37 (1.1x) | 942 | 75 (12.6x) |
尽管基础设施提供的排队系统未作更改，但当我们使用优化后的模型时，排队延迟也随之降低——在 p90 情况下，降低了 2 到 12 倍。这是因为我们更快地完成了先前的输入（由于执行延迟降低），我们就能更快地提取下一个输入（从而缩短了它们的排队时间）。
如果您有兴趣进一步优化 SAM2 推理或部署，请随时通过 torchao 存储库 与我们联系！
结论
我们用纯 PyTorch 重写了 Meta 的原始 SAM2，在准确性几乎没有损失的情况下，高度关注延迟。我们将优化的推理部署到无服务器基础设施提供商 Modal 上，以证明这些优化的好处可以在更真实的部署环境中实现。
通过利用 AOTInductor (AOTI) 通过 torch.export 进行提前编译、降低精度、批量提示和 GPU 预处理，我们观察到与常规急切模式 PyTorch 相比，p90 执行延迟和排队时间最多提高了 13 倍。
在弹性或自动缩放基础设施中，推理服务的副本需要定期自动创建和销毁，而 torch.compile 的简单部署会增加推理执行的工作量，使任何来自更快模型的改进相形见绌。通过利用 AOTInductor (AOTI) 通过 torch.export 进行提前编译，我们能够提前上传导出的模型并通过网络读取这些数据，这使我们能够在不显著增加工作量的情况下获得编译的好处。
有关如何重现此博客文章中的数据的更多详细信息，请查看 torchao 的 experiments 文件夹。如果您遇到任何技术问题，请随时联系我们或提交问题。