作者：PyTorch 团队

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

自 PyTorch 2.3 以来，此版本包含 3661 个提交和 475 个贡献者。我们衷心感谢我们敬业的社区的贡献。与往常一样，我们鼓励您尝试这些功能，并在我们改进 2.4 版本的过程中报告任何问题。有关如何开始使用 PyTorch 2 系列的更多信息，请访问我们的 入门 页面。

测试版 原型 性能改进
Python 3.12 对 torch.compile 的支持 FSDP2：基于 DTensor 的按参数分片 FSDP torch.compile 对 AWS Graviton（aarch64-linux）处理器的优化
AOTInductor 对 CPU 的冻结 torch.distributed.pipelining，简化的管道并行 TorchInductor 中的 BF16 符号形状优化
新的更高层级的 Python 自定义操作符 API Intel GPU 可通过源代码构建获得 针对使用 CPU 设备的 GenAI 项目的性能优化
将 TCPStore 的默认服务器后端切换到 libuv

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

测试版功能

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

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

[测试版] AOTInductor 对 CPU 的冻结

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

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

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 新默认服务器后端，这应该会显着减少初始化时间并提高可扩展性。从理想情况来看，这将使用户在考虑大型作业时缩短启动时间。

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

原型功能

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

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

有关 FSDP2 的动机和设计的更多信息，请参阅 Github 上的 RFC。

[原型] torch.distributed.pipelining，简化的管道并行

管道并行是深度学习的基本并行技术之一。它允许将模型的执行进行分区，以便多个微批次可以并发地执行模型代码的不同部分。

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

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

[原型] Intel GPU 可通过源代码构建获得

Linux 系统上的 PyTorch 中的 Intel GPU 提供了 Intel® 数据中心 GPU Max 系列上的基本功能：急切模式和 torch.compile。

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

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

性能改进

torch.compile 对 AWS Graviton（aarch64-linux）处理器的优化

AWS 针对 AWS Graviton3 处理器优化了 PyTorch 的 torch.compile 功能。与默认的急切模式推理相比，这种优化可以使 Hugging Face 模型推理的性能提高高达 2 倍（基于 33 个模型的性能改进的几何平均值），并且可以使 TorchBench 模型推理的性能提高高达 1.35 倍（基于 45 个模型的性能改进的几何平均值），这些模型涵盖了自然语言处理 (NLP)、计算机视觉 (CV) 和推荐模型等多种模型，这些模型都运行在基于 AWS Graviton3 的 Amazon EC2 实例上。

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

TorchInductor 中的 BF16 符号形状优化

Pytorch 用户现在可以使用测试版的 BF16 符号形状支持体验改进的质量和性能提升。虽然静态形状可能比符号形状提供更多优化机会，但对于以下场景来说它是不够的：例如，具有不同批大小和序列长度的推理服务，或者具有数据相关输出形状的检测模型。

使用 TorchBench、Huggingface 和 timms_model 进行的验证表明，与 BF16 静态形状场景相比，其通过率相似，加速效果也相当。将符号形状的优势与 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 上的性能。