注意
跳转到最后 下载完整的示例代码。
（原型）使用 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
据我们进行的测试（从 186.16ms 到 25.43ms），仅此快速更改就将运行时间提高了约 7 倍。
接下来，让我们使用 torch.compile 和我们的模型，看看性能会有多少提升。
# bf16 compiled runtime of the block is 19.95ms and peak memory 2.24GB
首次运行时，您应该会看到一系列 AUTOTUNE 输出，这是当 Inductor 在各种内核参数之间比较内核性能时出现的。这只发生一次（除非您删除缓存），因此如果您再次运行该单元格，只会得到基准测试输出。
torch.compile 带来了大约 27% 的额外改进。这使得模型达到了一个合理的基线，我们现在需要付出更多努力才能获得改进。
接下来，让我们应用量化。在 torchao 中，GPU 量化主要有三种形式，它只是原生的 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.28ms 到 729.65ms）。我们可以通过增加批量大小和优化模型的其他部分来进一步提高这一点。例如，可以使用某种形式的闪存注意力（flash attention）来实现。
更多信息请访问 torchao，并在您自己的模型上进行尝试。
# %%%%%%RUNNABLE_CODE_REMOVED%%%%%%
脚本总运行时间： (0 分钟 0.003 秒)
