注意
跳到末尾 下载完整示例代码。
（原型）使用 TorchAO 进行 GPU 量化#
作者：HDCharles
在本教程中，我们将引导您完成流行Segment Anything 模型的量化和优化。这些步骤将模仿开发 segment-anything-fast 存储库时所采取的一些步骤。本分步指南演示了如何应用这些技术来加速您自己的模型，尤其是那些使用 transformer 的模型。为此，我们将重点关注广泛适用的技术，例如使用 torch.compile 进行性能优化和量化，并衡量其影响。
设置您的环境#
首先，让我们配置您的环境。本指南是为 CUDA 12.1 编写的。我们已在 A100-PG509-200 上运行此教程，功率限制为 330.00 W。如果您使用不同的硬件，您可能会看到不同的性能数据。
> conda create -n myenv python=3.10
> pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
> pip install git+https://github.com/facebookresearch/segment-anything.git
> pip install git+https://github.com/pytorch/ao.git
Segment Anything 模型检查点设置
转到 segment-anything 存储库检查点并下载
vit_h检查点。或者，您可以使用wget（例如，wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth --directory-prefix=<path>）。通过编辑下面的代码，传递该目录：
{sam_checkpoint_base_path}=<path>
在本教程中，我们专注于量化 image_encoder，因为其输入是静态大小的，而提示编码器和掩码解码器具有可变大小，这使得它们更难量化。
为了简化分析，我们首先只关注一个块。
让我们从测量基线运行时间开始。
通过将模型转换为 bfloat16，我们可以立即获得性能提升。我们选择 bfloat16 而不是 fp16 的原因在于其动态范围，它与 fp32 相当。bfloat16 和 fp32 都具有 8 个指数位，而 fp16 只有 4 个。这种更大的动态范围有助于保护我们免受溢出错误以及由于量化而缩放和重缩放张量时可能出现的问题。
# bf16 runtime of the block is 25.43ms and peak memory 3.17GB
仅此一项快速更改，在我们的测试中运行时提高了约 7 倍（186.16 毫秒到 25.43 毫秒）。
接下来，让我们对模型使用 torch.compile，看看性能提高了多少。
# bf16 compiled runtime of the block is 19.95ms and peak memory 2.24GB
首次运行时，您应该会看到一系列 AUTOTUNE 输出，这发生在 inductor 比较各种内核参数的性能时。这只发生一次（除非您删除缓存），因此如果您再次运行该单元格，您应该只获得基准输出。
torch.compile 再次带来了约 27% 的提升。这使模型达到了合理的基线，我们现在必须更加努力地工作才能获得改进。
接下来，让我们应用量化。GPU 的量化在 torchao 中主要有三种形式，它们只是原生 pytorch+python 代码。这包括
int8 动态量化
int8 仅权重量化
int4 仅权重量化
不同的模型，或者有时是模型中的不同层，可能需要不同的技术。对于严重受计算限制的模型，动态量化往往效果最好，因为它将通常昂贵的浮点矩阵乘法运算替换为整数版本。仅权重量化在内存受限的情况下效果更好，其优势来自于加载更少的权重数据，而不是执行更少的计算。torchao API
Int8DynamicActivationInt8WeightConfig()、Int8WeightOnlyConfig() 或 Int4WeightOnlyConfig()
可用于轻松应用所需的量化技术，然后一旦模型使用 torch.compile 和 max-autotune 编译，量化就完成了，我们可以看到我们的加速。
注意
在旧版 PyTorch 上，您可能会遇到这些问题。如果您遇到问题，可以使用 apply_dynamic_quant 和 apply_weight_only_int8_quant 作为上述两个的直接替代（int4 没有替代）。
两个 API 的区别在于，Int8DynamicActivationInt8WeightConfig API 改变了线性模块的权重张量，因此它不是执行正常的线性运算，而是执行量化运算。当您有执行多于一个功能的非标准线性运算时，这会很有帮助。apply API 直接将线性模块替换为量化模块，这在旧版本上有效，但对非标准线性模块无效。
在这种情况下，Segment Anything 是计算受限的，因此我们将使用动态量化。
# bf16 compiled runtime of the quantized block is 19.04ms and peak memory 3.58GB
通过量化，我们进一步提高了性能，但内存使用量显著增加。
这有两个原因：
量化会增加模型的开销，因为我们需要对输入和输出进行量化和反量化。对于小批量大小，这种开销实际上会使模型变慢。
尽管我们正在进行量化矩阵乘法，例如
int8 x int8，乘法的结果会存储在 int32 张量中，其大小是未量化模型结果的两倍。如果我们能避免创建这个 int32 张量，我们的内存使用量将大大改善。
我们可以通过将整数矩阵乘法与随后的重新缩放操作融合来解决 #2，因为最终输出将是 bf16，如果我们立即将 int32 张量转换为 bf16 并存储它，我们将在运行时和内存方面获得更好的性能。
这样做的方法是启用 inductor 配置中的选项 force_fuse_int_mm_with_mul。
# bf16 compiled runtime of the fused quantized block is 18.78ms and peak memory 2.37GB
融合将性能又提高了少许（总共比基线提高了约 6%），并消除了几乎所有内存增加，剩余的量（量化后 2.37GB vs 未量化后 2.24GB）是由于量化开销，这是无法避免的。
然而，我们还没有完成，我们可以应用一些通用优化来获得最终的最佳性能。
我们有时可以通过禁用尾声融合来提高性能，因为自动调优过程可能会被融合混淆并选择糟糕的内核参数。
我们可以在所有方向上应用坐标下降调优，以扩大内核参数的搜索区域。
# bf16 compiled runtime of the final quantized block is 18.16ms and peak memory 2.39GB
正如您所看到的，我们从模型中又挤出了一点点改进，使我们的总改进比原始模型提高了 10 倍以上。为了最终估计量化的影响，让我们对完整模型进行公平的比较，因为实际的改进将因所涉及的形状而异。
结论#
在本教程中，我们以 Segment Anything 模型为例，学习了量化和优化技术。
最终，我们实现了完整模型在批量大小为 16 时约 7.7% 的量化加速（677.28 毫秒到 729.65 毫秒）。我们可以通过增加批量大小和优化模型的其他部分来进一步提升性能。例如，这可以通过某种形式的 Flash Attention 来实现。
欲了解更多信息，请访问 torchao，并在您自己的模型上尝试一下。
# %%%%%%RUNNABLE_CODE_REMOVED%%%%%%
脚本总运行时间：（0 分 0.002 秒）
