by Team PyTorch

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

自 PyTorch 2.3 以来，这个版本包含了 3661 个提交和 475 个贡献者的工作。我们衷心感谢所有敬业的社区成员作出的贡献。 一如既往，我们鼓励您尝试这些新功能并报告遇到的任何问题，以便我们改进 2.4 版本。关于如何开始使用 PyTorch 2 系列的更多信息可以在我们的 Getting Started 页面找到。

Beta 原型 性能改进
torch.compile 支持 Python 3.12 FSDP2：基于 DTensor 的参数分片 FSDP torch.compile 针对 AWS Graviton（aarch64-linux）处理器的优化
CPU 的 AOTInductor Freezing torch.distributed.pipelining，简化pipeline并行 TorchInductor 中的 BF16 符号形状优化
新的高级 Python 自定义算子 API Intel GPU 可通过源码构建使用。 针对使用CPU设备的生成式AI项目的性能优化
将 TCPStore 的默认服务器后端切换到 libuv

*查看完整的公开功能提交列表，请点击这里。

Beta 功能

[Beta] torch.compile 支持 Python 3.12

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

[Beta] CPU 的 AOTInductor Freezing

该功能使用户在 CPU 上使用 AOTInductor 时能够启用 Freezing 标志。通过该功能，AOTInductor 可以涵盖与 Inductor CPP 后端相同的一组操作场景，并达到相当的性能。在此支持之前， 当模型包含 MKLDNN 操作（如卷积、线性运算、反卷积等计算密集型操作）并且开启冻结时，这些模型将无法运行，因为 AOTInductor 不支持序列化具有不透明格式的MKLDNN权重。

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

export TORCHINDUCTOR_FREEZING=1

[Beta] 新的 Higher-level Python 自定义算子 API

我们添加了一个新的高级 Python 自定义算子 API，使得用自定义算子扩展 PyTorch 变得比以前更容易，这些自定义算子的行为类似于 PyTorch 的内置算子。 使用新的 high-level torch.library API 注册的算子保证与 torch.compile 和其他 PyTorch 子系统兼容；使用之前的low-level torch.library API API 编写 Python 自定义操作需要对 PyTorch 内部机制有深入理解，并且存在许多潜在风险。

更多信息请参见教程。

[Beta] 将 TCPStore 的默认服务器后端切换到 libuv

引入了一个新的默认 TCPStore 服务器后端，使用 libuv 构建，这应该带来显著降低的初始化时间和更好的可扩展性。理想情况下，应使用户在处理大规模作业时受益于更短的启动时间。

有关动机和回退说明的更多信息，请参阅此教程。

原型功能

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

FSDP2 是一种新的全分片数据并行实现，使用按参数维度 0 的分片方法，解决了 FSDP1 的扁平参数分片在组合性方面的根本性挑战。

有关 FSDP2 的初衷/设计的更多信息，请参阅 Github 上的 RFC。

[原型] torch.distributed.pipelining，简化 pipeline 并行

Pipeline 并行是深度学习的基本并行技术之一。它允许将模型的执行分区，使得多个微批次可以同时执行模型代码的不同部分。

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

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

[原型] Intel GPU 可通过源码构建使用

PyTorch 中的 Intel GPU 在 Linux 系统上为 Intel® 数据中心 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 针对 AWS Graviton3 处理器优化了 PyTorch 的 torch.compile 功能。这项优化使 Hugging Face 模型推理性能提高了最多 2 倍（基于 33 个模型性能改进的几何平均值），并且在 AWS Graviton3 基础的 Amazon EC2 实例上，TorchBench 模型推理性能相比默认的即时模式推理提高了最多 1.35 倍（45 个模型性能改进的几何平均值），涵盖了多个自然语言处理（NLP）、计算机视觉（CV）和推荐模型。

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

TorchInductor 中的 BF16 symbolic shape 优化

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

使用 TorchBench、Huggingface 和 timms_model 进行验证，显示通过率与 BF16 symbolic shape 场景相似，并且速度提升相当。 结合符号形状的优势与 Intel 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 项目性能优化

突出了PyTorch在CPU上的性能提升，通过对 “Segment Anything Fast” 和“Diffusion Fast” 项目的优化得到了证明。 然而，模型中仅支持 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 上的性能。