我们很高兴地宣布发布 PyTorch® 2.4 版本 (发布说明)！PyTorch 2.4 为 torch.compile 增加了对最新版 Python (3.12) 的支持。AOTInductor 冻结功能通过允许序列化 MKLDNN 权重，为运行 AOTInductor 的开发者提供了更多基于性能的优化。此外，还引入了一个利用 libuv 的新默认 TCPStore 服务器后端，这将显著减少运行大规模作业的用户的初始化时间。最后，一个新的 Python 自定义操作符 API 使集成自定义内核到 PyTorch 中比以前更容易，特别是对于 torch.compile。
自 PyTorch 2.3 以来，此版本由 3661 次提交和 475 位贡献者组成。我们衷心感谢我们敬业的社区所做的贡献。一如既往，我们鼓励您尝试这些功能，并在我们改进 2.4 版本时报告任何问题。有关如何开始使用 PyTorch 2 系列的更多信息，请访问我们的入门页面。
| Beta | 原型 | 性能改进 |
| torch.compile 支持 Python 3.12 | FSDP2: 基于 DTensor 的逐参数分片 FSDP | 针对 AWS Graviton (aarch64-linux) 处理器的 torch.compile 优化 |
| AOTInductor CPU 冻结 | torch.distributed.pipelining，简化流水线并行 | TorchInductor 中的 BF16 符号形状优化 |
| 新的高级 Python 自定义操作符 API | 可通过源码构建获得 Intel GPU 支持 | 利用 CPU 设备的 GenAI 项目的性能优化 |
| 将 TCPStore 的默认服务器后端切换到 libuv |
*要查看完整的公开功能提交列表，请点击此处。
Beta 版功能
[Beta] torch.compile 支持 Python 3.12
torch.compile() 以前只支持 Python 3.8-3.11。现在用户可以使用 Python 3.12 优化带有 torch.compile() 的模型。
[Beta] AOTInductor CPU 冻结
此功能允许用户在使用 CPU 上的 AOTInductor 时开启冻结标志。有了此功能，AOTInductor 可以覆盖相同的一组操作场景，并达到与 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 支持
PyTorch 在 Linux 系统上的 Intel GPU 提供 Intel® 数据中心 GPU Max 系列的基本功能：即时模式 (eager mode) 和 torch.compile。
对于即时模式，常用的 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 个模型的性能提升的几何平均值），TorchBench 模型推理性能提高高达 1.35 倍（基于 45 个模型的性能提升的几何平均值），与 AWS Graviton3-based Amazon EC2 实例上多个自然语言处理 (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 项目的性能优化
通过对 “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 上的性能。