本博客是该系列第五篇，专注于使用纯原生 PyTorch 加速生成式 AI 模型。我们展示了 GPTFast、Segment Anything Fast 和 Diffusion Fast 在英特尔® Xeon®处理器上的 AI 加速。
首先，我们回顾了 GPTFast，这是一项令人瞩目的工作，它通过不到 1000 行原生 PyTorch 代码加快了文本生成速度。最初，GPTFast 仅支持 CUDA 后端。我们将向您展示如何在 CPU 上运行 GPTFast，并通过仅权重量化（WOQ）实现额外的性能提升。
在 Segment Anything Fast 中，我们增加了对 CPU 后端的支持，并将通过利用 CPU 的增强能力以及 BFloat16、torch.compile 和块状注意力掩码的 scaled_dot_product_attention（SDPA）来展示性能加速。与 FP32 相比，速度提升比可达 2.91 倍在 vit_b 和 3.95 倍在 vit_h。
终于，Diffusion Fast 现在支持 CPU 后端，并利用 BFloat16、torch.compile 和 SDPA 提升了 CPU 的性能。我们还优化了 Inductor CPU 中的卷积、cat 和 permute 的布局传播规则，以提升性能。与 FP32 相比，在 Stable Diffusion XL (SDXL) 中速度提升比可达 3.91 倍。
提升 PyTorch CPU 性能的优化策略
GPTFast
过去一年，生成式 AI 在各个语言任务中取得了巨大成功，并越来越受欢迎。然而，由于自回归解码过程中的内存带宽瓶颈，生成模型面临着高昂的推理成本。为了解决这些问题，PyTorch 团队发布了 GPTFast，该工具旨在通过纯、原生的 PyTorch 加速文本生成。该项目从头开始开发了一个LLM，其速度比基线快近 10 倍，且在 1000 行原生 PyTorch 代码内完成。最初，GPTFast 仅支持 CUDA 后端，在约四个月内获得了大约 5,000 个星标。受到 Llama.cpp 的启发，英特尔团队从 PyTorch 2.4 版本开始提供了 CPU 后端支持，进一步增强了项目在无 GPU 环境中的可用性。以下是在 PyTorch CPU 上提升性能的优化策略：
-
Torch.compile
torch.compile 是自 PyTorch 2.0 以来引入的 PyTorch 函数，旨在解决 PyTorch 中精确图捕获的问题，最终使软件工程师能够更快地运行他们的 PyTorch 程序。
-
权重仅量化
仅权重量化（WOQ）是性能和精度之间的权衡，因为文本生成中自回归解码阶段的瓶颈是加载权重的内存带宽，通常 WOQ 相比于传统的 W8A8 量化方法能带来更好的精度。GPTFast 支持两种类型的 WOQ：W8A16 和 W4A16。具体来说，激活值存储在 BFloat16 中，模型权重可以量化为 int8 和 int4，如图 1 所示。
图 1. 仅权重量化模式。来源：Mingfei Ma，英特尔
-
权重预打包与微内核设计。
为了最大化吞吐量，GPTFast 允许使用内部 PyTorch ATen API 将模型权重预打包到针对 int4 的硬件特定布局中。受 Llama.cpp 的启发，我们将模型权重从 [N, K] 预打包到 [N/kNTileSize, K, kNTileSize/2]，在 avx512 上 kNTileSize 设置为 64。首先，模型权重在 N 维上分块，然后两个最内层维度进行转置。为了最小化内核计算中的去量化开销，我们将同一行的 64 个数据元素以交错模式进行洗牌，将 Lane2 与 Lane0 一起打包，将 Lane3 与 Lane1 一起打包，如图 2 所示。
图 2. Int4 上的权重预打包。来源：Mingfei Ma，Intel
在生成阶段，torch.nn.Linear 模块将被降低为在 PyTorch ATen 中的高性能内核中进行计算，其中量化权重首先进行去量化，然后在寄存器级别使用融合乘加 (FMA) 进行累加，如图 3 所示。
图 3. 微内核设计。来源：Mingfei Ma，Intel
Segment Anything Fast
Segment Anything Fast 提供了 Segment Anything 模型（SAM）的简单高效的 PyTorch 原生加速，SAM 是一种零样本视觉模型，用于生成可提示的图像掩码。以下是在 PyTorch CPU 上提升性能的优化策略：
-
BFloat16
Bfloat16 是一种常用的半精度类型。通过每个参数和激活的精度降低，我们可以在计算中节省大量时间和内存。
-
Torch.compile
torch.compile 是 PyTorch 2.0 以来引入的一个函数，旨在解决 PyTorch 中精确图捕获的问题，并最终使开发者能够更快地运行他们的 PyTorch 程序。
-
缩放点积注意力（SDPA）
缩放点积注意力（SDPA）是 Transformer 模型中的一个关键机制。PyTorch 提供了一个融合的实现，其性能显著优于原始方法。对于 Segment Anything Fast，我们将注意力掩码从 bfloat16 转换为 float32，采用分块方式。这种方法不仅降低了峰值内存使用，使其非常适合内存资源有限的系统，而且也提高了性能。
扩散快速
Diffusion Fast 为文本到图像扩散模型提供了一种简单高效的 PyTorch 原生加速。以下是在 PyTorch CPU 上提升性能的优化策略：
-
BFloat16
Bfloat16 是一种常用的半精度类型。通过每个参数和激活的精度降低，我们可以在计算中节省大量时间和内存。
-
Torch.compile
torch.compile 是自 PyTorch 2.0 以来引入的 PyTorch 函数，旨在解决 PyTorch 中精确图捕获的问题，最终使软件工程师能够更快地运行他们的 PyTorch 程序。
-
扩放点积注意力（SDPA）
SDPA 是 Transformer 模型中使用的核心机制，PyTorch 提供了融合实现，以展示相较于原始实现的大幅性能提升。
在原生 PyTorch CPU 上的模型使用
GPTFast
在 GPTFast 中启动 WOQ 之前，首先对模型权重进行量化。例如，使用 int4 和组大小为 32 进行量化：
python quantize.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --mode int4 –group size 32
然后通过传递 int4 检查点给 generate.py 运行生成
python generate.py --checkpoint_path checkpoints/$MODEL_REPO/model_int4.g32.pth --compile --device $DEVICE
要在 GPTFast 中使用 CPU 后端，只需将 DEVICE 变量从 cuda 更改为 CPU。
Segment Anything Fast
cd experiments
export SEGMENT_ANYTHING_FAST_USE_FLASH_4=0
python run_experiments.py 16 vit_b <pytorch_github> <segment-anything_github> <path_to_experiments_data> --run-experiments --num-workers 32 --device cpu
python run_experiments.py 16 vit_h <pytorch_github> <segment-anything_github> <path_to_experiments_data> --run-experiments --num-workers 32 --device cpu
扩散快速
python run_benchmark.py --compile_unet --compile_vae --device=cpu
性能评估
GPTFast
我们在测试分支和上述硬件配置上运行了基于 llama-2-7b-chat 模型的 PyTorch。在执行以下步骤后，与基线相比，在 eager 模式下我们看到了 3.8 倍的提升：
- 使用
torch.compile自动融合逐元素运算符。 - 使用 WOQ-int8 减少内存占用。
- 进一步使用 WOQ-int4 减少内存占用。
- 使用 AVX512，它可以在微内核中实现更快的量化和反量化。
图 4. GPTFast 在 Llama2-7b-chat 中的性能加速
Segment Anything Fast
我们在上述硬件配置上使用 PyTorch 运行 Segment Anything Fast，并通过 torch.compile 和 SDPA 实现了 BFloat16 的性能加速，与 FP32 相比，如图 5 所示，在 vit_b 中速度提升可达 2.91 倍，在 vit_h 中可达 3.95 倍。
图 5. Segment Anything Fast 在 vit_b/vit_h 中的性能加速
扩散快速
我们在上述硬件配置上使用 PyTorch 运行了扩散快速，并通过 torch.compile 和 SDPA 实现了 BFloat16 的性能加速，与 FP32 相比，如图 6 所示。与 FP32 相比，在 Stable Diffusion XL（SDXL）中，加速比可达 3.91 倍。
图 6. Stable Diffusion XL 中扩散快速性能加速
结论与未来工作
在本文中，我们介绍了针对权重量化、torch.compile 和 SDPA 的软件优化，展示了如何使用原生 PyTorch 在 CPU 上加速文本生成。随着 AMX-BF16 指令集的支持和 torchao 在 CPU 上对动态 int8 量化的优化，我们预计将进一步改进。我们将继续扩大我们的软件优化工作范围。
致谢
本博客中展示的结果是 Meta 和 Intel PyTorch 团队共同努力的成果。特别感谢 Meta 的 Michael Gschwind，他花费宝贵时间提供了大量帮助。我们一起在改善 PyTorch CPU 生态系统的道路上迈出了新的一步。
相关博客
第 1 部分：如何使用 Segment Anything Fast 将 Segment Anything 加速 8 倍。
第二部分：如何在 GPTFast 的帮助下将 Llama-7B 加速近 10 倍。
第三部分：如何将文本到图像扩散模型加速至 3 倍（使用 Diffusion Fast）。
第四部分：如何将 FAIR 的 Seamless M4T-v2 模型加速 2.7 倍。
产品和性能信息
图 4：英特尔至强可扩展处理器：使用 2x 英特尔(R)至强(R)铂金 8480+、56 核、HT 开启、Turbo 开启、NUMA 2、集成加速器可用[使用]：DLB 2 [0]、DSA 2 [0]、IAA 2 [0]、QAT 2 [0]、总内存 512GB（16x32GB DDR5 4800 MT/s [4800 MT/s]）、BIOS 3B07.TEL2P1、微代码 0x2b000590、三星 SSD 970 EVO Plus 2TB、CentOS Stream 9、5.14.0-437.el9.x86_64、单插槽运行（总共 1 个实例，每个实例 56 核，每个实例批大小 1），使用 PyTorch 2.5 wheel 运行的模型。英特尔于 10/15/24 进行测试。
图 5：英特尔至强可扩展处理器：使用 2x 英特尔(R)至强(R)铂金 8480+、56 核、HT 开启、Turbo 开启、NUMA 2、集成加速器可用[使用]：DLB 2 [0]、DSA 2 [0]、IAA 2 [0]、QAT 2 [0]、总内存 512GB（16x32GB DDR5 4800 MT/s [4800 MT/s]）、BIOS 3B07.TEL2P1、微代码 0x2b000590、三星 SSD 970 EVO Plus 2TB、CentOS Stream 9、5.14.0-437.el9.x86_64、单插槽运行（总共 1 个实例，每个实例 56 核，每个实例批大小 16），使用 PyTorch 2.5 wheel 运行的模型。英特尔于 10/15/24 进行测试。
图 6：英特尔至强可扩展处理器：使用 2x 英特尔(R)至强(R)铂金 8480+、56 核、HT 开启、Turbo 开启、NUMA 2、集成加速器可用[使用中]：DLB 2 [0]、DSA 2 [0]、IAA 2 [0]、QAT 2 [0]、总内存 512GB（16x32GB DDR5 4800 MT/s [4800 MT/s]）、BIOS 3B07.TEL2P1、微代码 0x2b000590、三星 SSD 970 EVO Plus 2TB、CentOS Stream 9、5.14.0-437.el9.x86_64、运行单插槽（总共 1 个实例，每个实例 56 核，每个实例批量大小 1），使用 PyTorch 2.5 wheel 运行的模型。由英特尔于 24 年 10 月 15 日测试。
通知和免责声明
性能因使用、配置和其他因素而异。请在性能指数网站上了解更多信息。性能结果基于所示配置的测试日期，可能无法反映所有公开可用的更新。请参阅备份以获取配置详细信息。没有产品或组件可以绝对安全。您的成本和结果可能会有所不同。英特尔技术可能需要启用硬件、软件或服务激活。
英特尔公司。英特尔、英特尔标志以及其他英特尔商标均为英特尔公司或其子公司的商标。其他名称和品牌可能属于他人所有。
AI 免责声明：
AI 功能可能需要购买软件、订阅或由软件或平台提供商启用，或可能具有特定的配置或兼容性要求。详细信息请访问 www.intel.com/AIPC。结果可能有所不同。
