作者：Intel PyTorch 团队

本博客是系列文章中的第五篇，重点介绍如何使用纯原生 PyTorch 加速生成式 AI 模型。我们将展示 GPTFast、Segment Anything Fast 和 Diffusion Fast 在 Intel® Xeon® 处理器上的 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 中优化了卷积 (convolution)、cat 和 permute 的布局传播规则以提高性能。相对于 FP32，其加速比在 Stable Diffusion XL (SDXL) 中可达 3.91 倍。

提升 PyTorch CPU 性能的优化策略

GPTFast

过去一年，生成式 AI 在各种语言任务中取得了巨大成功，并变得越来越受欢迎。然而，由于自回归解码过程中的内存带宽瓶颈，生成模型的推理成本很高。为了解决这些问题，PyTorch 团队发布了 GPTFast，旨在仅使用纯原生 PyTorch 加速文本生成。该项目从头开始开发了一个 LLM，其速度比基线快了近 10 倍，代码行数不足 1000 行。最初，GPTFast 只支持 CUDA 后端，并在大约四个月内获得了约 5,000 个 star。受 Llama.cpp 的启发，Intel 团队从 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 所示。

flow diagram

图 1. 仅权重量化模式。来源：Mingfei Ma，Intel

  • 权重预打包与微内核设计。

    为了最大化吞吐量，GPTFast 允许使用内部 PyTorch ATen API 将模型权重预打包成 int4 上的硬件特定布局。受 Llama.cpp 启发，我们将模型权重从 [N, K] 预打包为 [N/kNTileSize, K, kNTileSize/2]，其中 avx512 上的 kNTileSize 设置为 64。首先，沿 N 维度对模型权重进行分块，然后对最里面的两个维度进行转置。为了最小化内核计算中的反量化开销，我们以交错模式打乱同一行的 64 个数据元素，将 Lane2 和 Lane0 打包在一起，将 Lane3 和 Lane1 打包在一起，如图 2 所示。

flow diagram

图 2. Int4 上的权重预打包。来源：Mingfei Ma，Intel

在生成阶段，torch.nn.Linear 模块将被下沉到 PyTorch ATen 内部的高性能内核进行计算，其中量化后的权重将首先进行反量化，然后在寄存器级别与融合乘加 (FMA) 进行累积，如图 3 所示。

flow diagram

图 3. 微内核设计。来源：Mingfei Ma，Intel

Segment Anything Fast

Segment Anything Fast 为 Segment Anything Model (SAM) 提供了一种简单高效的 PyTorch 原生加速，SAM 是一种用于生成可提示图像掩码的零样本视觉模型。以下是用于提升 PyTorch CPU 性能的优化策略：

  • BFloat16

    Bfloat16 是一种常用的半精度类型。通过降低每个参数和激活的精度，我们可以显著节省计算时间和内存。

  • Torch.compile

    torch.compile 是自 PyTorch 2.0 以来引入的一个 PyTorch 函数，旨在解决 PyTorch 中准确捕获图的问题，并最终使开发者能够更快地运行他们的 PyTorch 程序。

  • 缩放点积注意力 (SDPA)

    缩放点积注意力 (SDPA) 是 Transformer 模型中的关键机制。PyTorch 提供了一种融合实现，其性能显著优于朴素方法。对于 Segment Anything Fast，我们以块状方式将注意力掩码从 bfloat16 转换为 float32。这种方法不仅减少了峰值内存使用量，使其非常适合内存资源有限的系统，而且还提升了性能。

Diffusion Fast

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

Diffusion Fast

python run_benchmark.py --compile_unet --compile_vae --device=cpu

性能评估

GPTFast

我们基于 测试分支 和上述硬件配置在 PyTorch 上运行了 llama-2-7b-chat 模型。应用以下步骤后，我们看到了相对于 eager 模式基线的 3.8 倍提升。

  • 使用 torch.compile 自动融合逐元素算子。
  • 使用 WOQ-int8 减少内存占用。
  • 使用 WOQ-int4 进一步减少内存占用。
  • 使用 AVX512，可在微内核中实现更快的反量化。

bar chart

图 4. Llama2-7b-chat 中 GPTFast 的性能加速

Segment Anything Fast

我们在上述硬件配置下在 PyTorch 上运行 Segment Anything Fast，与 FP32 相比，使用 BFloat16、torch.compile 和 SDPA 实现了性能加速，如图 5 所示。相对于 FP32，其加速比在 vit_b 中可达 2.91 倍，在 vit_h 中可达 3.95 倍。

bar chart

图 5. vit_b/vit_h 中 Segment Anything Fast 的性能加速

Diffusion Fast

我们在上述硬件配置下在 PyTorch 上运行 Diffusion Fast，与 FP32 相比，使用 BFloat16、torch.compile 和 SDPA 实现了性能加速，如图 6 所示。相对于 FP32，其加速比在 Stable Diffusion XL (SDXL) 中可达 3.91 倍。

bar chart

图 6. Stable Diffusion XL 中 Diffusion Fast 的性能加速

结论与未来工作

在本博客中，我们介绍了仅权重量化、torch.compile 和 SDPA 的软件优化，演示了如何在 CPU 上使用原生 PyTorch 加速文本生成。随着对 AMX-BF16 指令集的支持以及在 CPU 上使用 torchao 优化动态 int8 量化，预计将会有进一步的改进。我们将继续扩大我们的软件优化范围。

致谢

本博客中展示的结果是 Meta 和 Intel 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：Intel Xeon 可扩展处理器：在第 4 代 Intel Xeon 可扩展处理器上测得，配置为：2x Intel(R) Xeon(R) Platinum 8480+, 56 核, HT On, Turbo On, 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, Samsung SSD 970 EVO Plus 2TB, CentOS Stream 9, 5.14.0-437.el9.x86_64, 运行单路 (总共 1 个实例，每个实例 56 核，每个实例 Batch Size 1), 模型运行使用 PyTorch 2.5 wheel。由 Intel 于 24 年 10 月 15 日测试。

图 5：Intel Xeon 可扩展处理器：在第 4 代 Intel Xeon 可扩展处理器上测得，配置为：2x Intel(R) Xeon(R) Platinum 8480+, 56 核, HT On, Turbo On, 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, Samsung SSD 970 EVO Plus 2TB, CentOS Stream 9, 5.14.0-437.el9.x86_64, 运行单路 (总共 1 个实例，每个实例 56 核，每个实例 Batch Size 16), 模型运行使用 PyTorch 2.5 wheel。由 Intel 于 24 年 10 月 15 日测试。

图 6：Intel Xeon 可扩展处理器：在第 4 代 Intel Xeon 可扩展处理器上测得，配置为：2x Intel(R) Xeon(R) Platinum 8480+, 56 核, HT On, Turbo On, 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, Samsung SSD 970 EVO Plus 2TB, CentOS Stream 9, 5.14.0-437.el9.x86_64, 运行单路 (总共 1 个实例，每个实例 56 核，每个实例 Batch Size 1), 模型运行使用 PyTorch 2.5 wheel。由 Intel 于 24 年 10 月 15 日测试。

注意事项与免责声明

性能因使用情况、配置及其他因素而异。请访问性能指标网站了解更多信息。性能结果基于所示配置在所示日期的测试，可能无法反映所有公开发布的更新。请参阅备份以获取配置详情。任何产品或组件都无法做到绝对安全。您的成本和结果可能会有所不同。Intel 技术可能需要启用硬件、软件或服务激活。

Intel Corporation。Intel、Intel 徽标和其他 Intel 标记是 Intel Corporation 或其子公司的商标。其他名称和品牌可能是他人的财产。

AI 免责声明

AI 功能可能需要购买软件、订阅或由软件或平台提供商启用，或者可能具有特定的配置或兼容性要求。详情请见 www.intel.com/AIPC。结果可能有所不同。