本博客是系列文章中的第五篇，重点介绍如何使用纯原生 PyTorch 加速生成式 AI 模型。我们将展示在英特尔®至强®处理器上对 GPTFast、Segment Anything Fast 和 Diffusion Fast 进行 GenAI 加速。
首先，我们回顾了 GPTFast，这是一项了不起的工作，它在不到 1000 行原生 PyTorch 代码中加速了文本生成。最初，GPTFast 只支持 CUDA 后端。我们将向您展示如何在 CPU 上运行 GPTFast，并通过仅权重量化 (WOQ) 实现额外的性能加速。
在 Segment Anything Fast 中，我们集成了对 CPU 后端的支持，并将通过利用 BFloat16、torch.compile 和带有块级注意力掩码的 scaled_dot_product_attention (SDPA) 增加 CPU 的能力来展示性能加速。与 FP32 相比，vit_b 的加速比可以达到 2.91 倍，vit_h 的加速比可以达到 3.95 倍。
最后，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.compiletorch.compile 是 PyTorch 2.0 引入的一个 PyTorch 函数，旨在解决 PyTorch 中准确图捕获的问题，最终使软件工程师能够更快地运行他们的 PyTorch 程序。
- 仅权重量子化仅权重量子化 (WOQ) 是性能和精度之间的权衡，因为文本生成中自回归解码阶段的瓶颈是加载权重的内存带宽，并且通常与传统的量化方法（例如 W8A8）相比，WOQ 可以带来更好的精度。GPTFast 支持两种类型的 WOQ：W8A16 和 W4A16。具体来说，激活以 BFloat16 存储，模型权重可以量化为 int8 和 int4，如图 1 所示。
图 1. 仅权重量子化模式。来源：马鸣飞，英特尔
- 权重预打包和微内核设计。为了最大化吞吐量，GPTFast 允许使用内部 PyTorch ATen API 将模型权重预打包为 int4 上的硬件特定布局。受 Llama.cpp 的启发，我们将模型权重从 [N, K] 预打包到 [N/kNTileSize, K, kNTileSize/2]，其中 kNTileSize 在 avx512 上设置为 64。首先，模型权重沿 N 维度阻塞，然后转置两个最内部的维度。为了最小化内核计算中的反量化开销，我们以交错模式打乱同一行中的 64 个数据元素，将 Lane2 和 Lane0 打包在一起，Lane3 和 Lane1 打包在一起，如图 2 所示。
图 2. Int4 上的权重预打包。来源：马鸣飞，英特尔
在生成阶段，torch.nn.Linear 模块将降级到 PyTorch ATen 内部的高性能内核进行计算，其中量化权重将首先进行反量化，然后与寄存器级别的融合乘加 (FMA) 累积，如图 3 所示。
图 3. 微内核设计。来源：马鸣飞，英特尔
Segment Anything Fast
Segment Anything Fast 为 Segment Anything Model (SAM) 提供了简单高效的 PyTorch 原生加速，SAM 是一种用于生成可提示图像掩码的零样本视觉模型。以下是用于提升 PyTorch CPU 性能的优化策略
- BFloat16Bfloat16 是一种常用的半精度类型。通过每个参数和激活的精度降低，我们可以显着节省计算时间和内存。
- Torch.compiletorch.compile 是 PyTorch 2.0 引入的一个 PyTorch 函数，旨在解决 PyTorch 中准确图捕获的问题，最终使开发人员能够更快地运行他们的 PyTorch 程序。
- 缩放点积注意力 (SDPA)缩放点积注意力 (SDPA) 是 Transformer 模型中的一个关键机制。PyTorch 提供了一种融合实现，其性能显著优于朴素方法。对于 Segment Anything Fast，我们以块级方式将注意力掩码从 bfloat16 转换为 float32。这种方法不仅降低了峰值内存使用，使其非常适合内存资源有限的系统，而且还提高了性能。
Diffusion Fast
Diffusion Fast 为文本到图像扩散模型提供了简单高效的 PyTorch 原生加速。以下是用于提升 PyTorch CPU 性能的优化策略
- BFloat16Bfloat16 是一种常用的半精度类型。通过每个参数和激活的精度降低，我们可以显着节省计算时间和内存。
- Torch.compiletorch.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
Diffusion Fast
python run_benchmark.py --compile_unet --compile_vae --device=cpu
性能评估
GPTFast
我们基于 test branch 和上述硬件配置在 PyTorch 上运行了 llama-2-7b-chat 模型。在应用以下步骤后，我们发现与 eager 模式下的基线相比，性能提升了 3.8 倍
- 使用
torch.compile自动融合元素级运算符。 - 使用 WOQ-int8 减少内存占用。
- 使用 WOQ-int4 进一步减少内存占用。
- 使用 AVX512，可以在微内核中实现更快的反量化。
图 4. Llama2-7b-chat 中 GPTFast 的性能加速
Segment Anything Fast
我们在上述硬件配置下在 PyTorch 上运行了 Segment Anything Fast，并实现了 BFloat16 与 torch.compile 和 SDPA 相比 FP32 的性能加速，如图 5 所示。与 FP32 相比，vit_b 的加速比可以达到 2.91 倍，vit_h 的加速比可以达到 3.95 倍。
图 5. vit_b/vit_h 中 Segment Anything Fast 的性能加速
Diffusion Fast
我们在上述硬件配置下在 PyTorch 上运行了 Diffusion Fast，并实现了 BFloat16 与 torch.compile 和 SDPA 相比 FP32 的性能加速，如图 6 所示。与 FP32 相比，Stable Diffusion XL (SDXL) 的加速比可以达到 3.91 倍。
图 6. Stable Diffusion XL 中 Diffusion Fast 的性能加速
结论和未来工作
在本博客中，我们介绍了仅权重量子化、torch.compile 和 SDPA 的软件优化，展示了如何使用原生 PyTorch 在 CPU 上加速文本生成。通过支持 AMX-BF16 指令集以及使用 CPU 上的 torchao 优化动态 int8 量化，预计会进一步改进。我们将继续将软件优化工作扩展到更广泛的范围。
鸣谢
本博客中呈现的结果是 Meta 和英特尔 PyTorch 团队的共同努力。特别感谢 Meta 的 Michael Gschwind 花费宝贵时间提供了大量帮助。我们共同在改进 PyTorch CPU 生态系统的道路上又迈进了一步。
相关博客
第 1 部分：如何使用 Segment Anything Fast 将 Segment Anything 加速 8 倍以上。
第 2 部分：如何借助 GPTFast 将 Llama-7B 加速近 10 倍。
第 3 部分：如何使用 Diffusion Fast 将 文本到图像扩散模型加速高达 3 倍。
第 4 部分：如何将 FAIR 的 Seamless M4T-v2 模型加速 2.7 倍。
产品和性能信息
图 4：英特尔至强可扩展处理器：使用第四代英特尔至强可扩展处理器进行测量：2 个英特尔(R) 至强(R) 白金 8480+，56 核，超线程开启，睿频开启，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 运行。英特尔于 2024 年 10 月 15 日测试。
图 5：英特尔至强可扩展处理器：使用第四代英特尔至强可扩展处理器进行测量：2 个英特尔(R) 至强(R) 白金 8480+，56 核，超线程开启，睿频开启，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 运行。英特尔于 2024 年 10 月 15 日测试。
图 6：英特尔至强可扩展处理器：使用第四代英特尔至强可扩展处理器进行测量：2 个英特尔(R) 至强(R) 白金 8480+，56 核，超线程开启，睿频开启，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 运行。英特尔于 2024 年 10 月 15 日测试。
注意事项和免责声明
性能因使用、配置和其他因素而异。在性能指数网站上了解更多信息。性能结果基于所示日期在配置中的测试，可能无法反映所有公开可用的更新。有关配置详情，请参见备份。任何产品或组件都不能绝对安全。您的成本和结果可能会有所不同。英特尔技术可能需要启用硬件、软件或服务激活。
英特尔公司。英特尔、英特尔徽标和其他英特尔标志是英特尔公司或其子公司的商标。其他名称和品牌可能属于他人所有。
AI 免责声明
AI 功能可能需要购买软件、订阅或由软件或平台提供商启用，或者可能有特定的配置或兼容性要求。详情请访问 www.intel.com/AIPC。结果可能会有所不同。