PyTorch 2.4 发布博客

PyTorch 团队

我们很高兴地宣布 PyTorch® 2.4 的发布（发布说明）！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 系列的更多信息，请访问我们的入门页面。

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

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 算子（当涉及计算密集型算子时，例如 Convolution、Linear、ConvTranspose 等）并且 freezing 开启时，这些模型将无法运行，因为 AOTInductor 不支持序列化具有不透明格式的 MKLDNN 权重。

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

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 系统上 PyTorch 中的 Intel GPU 在 Intel® Data Center GPU Max 系列上提供基本功能：eager 模式和 torch.compile。

对于 eager 模式，常用的 Aten 算子通过使用 SYCL 编程语言实现。性能最关键的图和算子通过使用 oneAPI Deep Neural Network (oneDNN) 进行高度优化。对于 torch.compile 模式，Intel GPU 后端集成到 Triton 之上的 Inductor 中。

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

性能改进

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

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

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

TorchInductor 中的 BF16 符号形状优化

Pytorch 用户现在可以通过 beta BF16 符号形状支持体验到改进的质量和性能提升。虽然与符号形状相比，静态形状可以提供额外的优化机会，但它不足以满足诸如具有可变批大小和序列长度的推理服务或具有数据相关输出形状的检测模型等场景。

使用 TorchBench、Huggingface 和 timms_model 进行的验证表明，BF16 静态形状场景具有相似的通过率和相当的加速。结合符号形状的优势与 Intel CPU 提供的 BF16 AMX 指令硬件加速以及 PyTorch 2.4 中适用于静态和符号形状的通用 Inductor CPU 后端优化，与 PyTorch 2.3 相比，BF16 符号形状的性能得到了显着提升。

使用此功能的 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 上的性能。