PyTorch 2.4 发布博客

作者：PyTorch 团队

我们很高兴宣布 PyTorch® 2.4 版本发布 (发布说明)！PyTorch 2.4 增加了对最新版本 Python (3.12) 的支持，用于 torch.compile。AOTInductor freezing（冻结）通过允许序列化 MKLDNN 权重，为运行 AOTInductor 的开发者提供了更多基于性能的优化。此外，还引入了一个利用 libuv 的新默认 TCPStore 服务器后端，这将显著减少运行大规模作业的用户的初始化时间。最后，一个新的 Python 自定义算子 API 使将自定义内核集成到 PyTorch 中比以往更容易，特别是对于 torch.compile。

自 PyTorch 2.3 以来，此版本包含 3661 次提交和 475 位贡献者。我们衷心感谢我们专注的社区所做的贡献。一如既往，我们鼓励您尝试这些功能并报告任何问题，以便我们改进 2.4 版本。有关如何开始使用 PyTorch 2 系列的更多信息，请参见我们的开始页面。

*要查看公共特性提交的完整列表，请单击此处。

测试版功能 (Beta Features)

[测试版] *torch.compile* 支持 Python 3.12

torch.compile() 之前仅支持 Python **3.8-3.11**。用户现在可以使用 Python **3.12** 通过 torch.compile() 优化模型。

[测试版] CPU 的 AOTInductor Freezing

此功能允许用户在使用 CPU 上的 AOTInductor 时开启 freezing 标志。凭借此功能，AOTInductor 可以涵盖与 Inductor CPP 后端相同的算子场景集，并达到同等性能。在此支持之前，当模型包含 MKLDNN 算子（涉及计算密集型算子，例如 Convolution、Linear、ConvTranspose 等）并且 freezing 开启时，这些模型将无法运行，因为 AOTInductor 不支持序列化具有不透明格式的 MKLDNN 权重。

工作流程如 AOTInductor 教程中所述，此外用户现在可以添加 freezing 标志以获得更好的性能

export TORCHINDUCTOR_FREEZING=1

[测试版] 新的高级 Python 自定义算子 API

我们添加了一个新的高级 Python 自定义算子 API，与之前相比，使用行为类似 PyTorch 内置算子的自定义算子扩展 PyTorch 变得更加容易。使用新的高级 torch.library API 注册的算子保证与 torch.compile 和其他 PyTorch 子系统兼容；而使用之前的低级 torch.library API 在 Python 中编写自定义算子需要深入理解 PyTorch 内部机制，且存在许多陷阱。

更多信息请参见教程。

[测试版] 将 TCPStore 的默认服务器后端切换到 *libuv*

引入了一个用 libuv 构建的新默认 TCPStore 服务器后端，这将显著降低初始化时间和提高可伸缩性。这应该能为运行大规模作业的用户带来更短的启动时间，从而让他们受益匪浅。

有关此更改的动机 + 回退说明，请参阅此教程。

原型功能 (Prototype Features)

[原型] FSDP2：基于 DTensor 的按参数分片 FSDP

FSDP2 是一种新的完全分片数据并行实现，它使用 dim-0 按参数分片来解决 FSDP1 平坦参数分片所带来的基本可组合性挑战。

有关 FSDP2 的动机/设计，请参阅Github 上的 RFC。

[原型] *torch.distributed.pipelining*，简化的流水线并行

流水线并行是深度学习的一种基本并行技术。它允许对模型执行进行分区，以便多个微批次可以并发执行模型的不同部分代码。

torch.distributed.pipelining 提供了一个工具包，可以轻松地在通用模型上实现流水线并行，同时还提供了与 DDP、FSDP 或张量并行等其他常见 PyTorch 分布式功能的组合性。

有关此功能的更多信息，请参阅我们的文档和教程。

[原型] 可通过源码构建获取 Intel GPU 支持

PyTorch 在 Linux 系统上的 Intel GPU 为英特尔® 数据中心 GPU Max 系列提供了基本功能：eager 模式和 torch.compile。

对于 eager 模式，常用的 Aten 算子使用 SYCL 编程语言实现。使用 oneAPI 深度神经网络 (oneDNN) 对性能最关键的图和算子进行了高度优化。对于 torch.compile 模式，Intel GPU 后端已集成到基于 Triton 的 Inductor 中。

有关 Intel GPU 源码构建的更多信息，请参阅我们的博客文章和文档。

性能改进

针对 AWS Graviton (aarch64-linux) 处理器的 *torch.compile* 优化

AWS 优化了 PyTorch 的 torch.compile 功能，使其适用于 AWS Graviton3 处理器。与在基于 AWS Graviton3 的 Amazon EC2 实例上针对多个自然语言处理 (NLP)、计算机视觉 (CV) 和推荐模型的默认 eager 模式推理相比，此优化使 Hugging Face 模型推理的性能提升高达 2 倍（基于 33 个模型的性能提升几何平均值），TorchBench 模型推理的性能提升高达 1.35 倍（基于 45 个模型的性能提升几何平均值）。

有关具体技术细节的更多信息，请参阅博客文章。

TorchInductor 中的 BF16 符号形状优化

PyTorch 用户现在可以通过测试版 BF16 符号形状支持体验到改进的质量和性能提升。虽然静态形状与符号形状相比可能提供额外的优化机会，但对于具有不同批次大小和序列长度的推理服务，或具有数据依赖输出形状的检测模型等场景来说，静态形状是不够的。

使用 TorchBench、Huggingface 和 timms_model 进行验证显示，通过率与 BF16 静态形状场景相似，加速效果相当。结合符号形状的优势、英特尔 CPU 提供的 BF16 AMX 指令硬件加速以及适用于 PyTorch 2.4 中静态和符号形状的通用 Inductor CPU 后端优化，BF16 符号形状的性能与 PyTorch 2.3 相比显著提升。

使用此功能的 API

model = ….
model.eval()
with torch.autocast(device_type=”cpu”, dtype=torch.bfloat16), torch.no_grad():
   compiled_model = torch.compile(model, dynamic=True)

利用 CPU 设备的 GenAI 项目性能优化

通过为“Segment Anything Fast”和“Diffusion Fast”项目进行的优化，突显了 PyTorch 在 CPU 上的增强性能。然而，模型中只支持 CUDA 设备。我们已将 CPU 支持整合到项目中，使用户能够利用 CPU 增强的性能运行项目实验。同时，我们还为 SDPA 采用了块级注意力掩码，这可以显著降低峰值内存使用量并提高性能。我们还优化了 Inductor CPU 中的一系列布局传播规则以提高性能。

为了方便使用，我们更新了 README 文件。使用此功能的 API 如下所示，只需在命令行中提供 --device cpu

• 对于 Segment Anything Fast

  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
  

• 对于 Diffusion Fast

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

用户可以按照指南运行实验，亲身观察性能改进，并探索跨 FP32 和 BF16 数据类型的性能改进趋势。

此外，用户还可以使用 torch.compile 和 SDPA 获得良好的性能。通过观察这些不同因素的性能趋势，用户可以更深入地了解各种优化如何提升 PyTorch 在 CPU 上的性能。