PyTorch 2.4 版本发布博客

由 PyTorch 团队

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

本次发布由自 PyTorch 2.3 以来的 3661 个提交和 475 位贡献者组成。我们衷心感谢我们忠诚的社区成员的贡献。一如既往，我们鼓励大家尝试这些新功能，并在我们改进 2.4 版本时报告任何问题。有关如何开始使用 PyTorch 2 系列的信息，请参阅我们的入门页面。

点击此处查看所有公共功能提交的完整列表。

测试功能

[Beta] 支持 Python 3.12 的 torch.compile

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

[Beta] CPU AOTInductor 冻结

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

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

export TORCHINDUCTOR_FREEZING=1

[Beta] 新增高级 Python 自定义运算符 API

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

请参阅教程以获取更多信息。

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

引入了一个新的默认服务器后端，用于 TCPStore，该后端是用 libuv 构建的，这应该会显著降低初始化时间并提高可扩展性。在考虑大规模作业时，这应该理想地使用户的启动时间大大缩短。

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

原型功能

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

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

关于 FSDP2 的动机/设计，请参阅 GitHub 上的 RFC。

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

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

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

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

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

在 Linux 系统上，Intel GPU 在 PyTorch 中提供了 Intel®数据中心 GPU Max 系列的基本功能：eager 模式以及 torch.compile。

对于 eager 模式，常用的 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 模型推理（基于 33 个模型的性能改进几何平均值）的性能提升高达 2 倍，与 AWS Graviton3 基于的 Amazon EC2 实例上默认的 eager 模式推理相比，TorchBench 模型推理（基于 45 个模型的性能改进几何平均值）的性能提升高达 1.35 倍，涉及多个自然语言处理（NLP）、计算机视觉（CV）和推荐模型。

如需了解具体技术细节，请参阅博客文章。

TorchInductor 中的 BF16 符号形状优化

PyTorch 用户现在可以体验 beta 版本的 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 上的性能。