跳转到主要内容
博客

PyTorch 2.5 发布博客

作者: 2024 年 10 月 17 日2024 年 11 月 21 日暂无评论

我们很高兴地宣布 PyTorch® 2.5 已发布(发行说明)!此版本为 SDPA 提供了新的 cuDNN 后端,默认情况下为 H100 或更新 GPU 上的 SDPA 用户实现加速。此外,torch.compile 的区域编译提供了一种减少 torch.compile 冷启动时间的方法,它允许用户编译重复的 nn.Module(例如 LLM 中的 Transformer 层)而无需重新编译。最后,TorchInductor CPP 后端通过众多增强功能,如 FP16 支持、CPP 封装器、AOT-Inductor 模式和最大自动调优模式,提供了坚实的性能加速。

此版本包含自 PyTorch 2.4 以来来自 504 位贡献者的 4095 次提交。我们衷心感谢我们敬业的社区所做的贡献。一如既往,我们鼓励您试用这些功能并报告任何问题,以便我们改进 2.5。有关如何开始使用 PyTorch 2 系列的更多信息,请访问我们的 入门 页面。

此外,请查看我们与 TorchRec 和 TorchFix 一同发布的新生态系统项目。

测试版原型
SDPA 的 cuDNN 后端FlexAttention
torch.compile 区域编译无需重新编译编译自动微分
TorchDynamo 添加了对异常处理和 MutableMapping 类型的支持飞行记录器
TorchInductor CPU 后端优化支持带 GEMM 模板的 CPU 最大自动调优
Windows 上的 TorchInductor
Eager 模式和 TorchInductor CPP 后端均支持 CPU 路径上的 FP16
自动加载设备扩展
增强的英特尔 GPU 支持

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

测试版功能

[测试版] SDPA 的 cuDNN 后端

cuDNN “融合 Flash Attention”后端已集成到 torch.nn.functional.scaled_dot_product_attention。在 NVIDIA H100 GPU 上,这比 FlashAttentionV2 可提供高达 75% 的加速。对于 H100 或更新 GPU 上的所有 SDPA 用户,此加速默认启用。

[测试版] torch.compile 区域编译无需重新编译

通过 torch._dynamo.config.inline_inbuilt_nn_modules 实现的区域编译无需重新编译,在 2.5+ 中默认设置为 True。此选项允许用户编译重复的 nn.Module (例如 LLM 中的 Transformer 层)而无需重新编译。与编译完整模型相比,此选项可减少编译延迟,性能损失为 1%-5%。

有关更多信息,请参阅 教程

[测试版] TorchInductor CPU 后端优化

此功能提升了 Inductor 的 CPU 后端优化,包括 CPP 后端代码生成和带自定义 CPU 内核的 FX 融合。Inductor CPU 后端支持常见数据类型和所有 Inductor IR 操作的向量化,以及静态和符号形状。它与 Linux 和 Windows 操作系统兼容,并支持默认 Python 封装器、CPP 封装器和 AOT-Inductor 模式。

此外,它扩展了 GEMM 模板的最大自动调优模式(在 2.5 中原型),提供了进一步的性能提升。后端支持各种 FX 融合,可下沉到自定义内核,例如用于 Linear/Conv 操作的 oneDNN 和 SDPA。Inductor CPU 后端在三个基准测试套件(TorchBench、Hugging Face 和 timms)中始终实现性能加速,在测试的 193 个模型中,97.5% 的模型优于 eager 模式。

原型功能

[原型] FlexAttention

我们引入了一个灵活的 API,只需几行惯用的 PyTorch 代码即可实现各种注意力机制,例如滑动窗口、因果掩码和 PrefixLM。此 API 利用 torch.compile 生成融合的 FlashAttention 内核,这消除了额外的内存分配并实现了与手写实现相当的性能。此外,我们使用 PyTorch 的自动微分机制自动生成反向传播。此外,我们的 API 可以利用注意力掩码中的稀疏性,从而显着优于标准注意力实现。

有关更多信息和示例,请参阅 官方博客文章 和 Attention Gym

[原型] 编译自动微分

编译自动微分是 PT2 堆栈的扩展,允许捕获整个反向传播。与 AOT dispatcher 追踪的反向图不同,编译自动微分追踪会延迟到反向执行时,这使得它不受前向传播图中断的影响,并允许它将反向钩子记录到图中。

有关更多信息,请参阅 教程

[原型] 飞行记录器

飞行记录器是一种新的调试工具,可帮助调试停滞作业。该工具通过持续捕获有关集合的信息来工作。检测到停滞作业后,这些信息可用于快速识别行为异常的排名/机器以及代码堆栈跟踪。

有关更多信息,请参阅以下 教程

[原型] 支持带 GEMM 模板的 CPU 最大自动调优

torch.compile 中 Inductor CPU 后端的最大自动调优模式会在编译时分析操作的多种实现并选择性能最佳的实现。这对于 GEMM 相关操作特别有利,它使用基于 C++ 模板的 GEMM 实现作为 ATen 基于 oneDNN 和 MKL 库的方法的替代方案。我们支持 x86 CPU 上的 FP32、BF16、FP16 和 INT8,并支持尾部融合。我们已在 dynamo 基准测试套件上看到高达 7% 的几何平均加速,在 LLM 推理的下一个令牌延迟方面提高了高达 20%。

有关更多信息,请参阅 教程

[原型] Windows 上的 TorchInductor CPU

torch.compile 中的 Inductor CPU 后端现在可以在 Windows 上运行。我们目前支持适用于 Windows Inductor 的 MSVC (cl)、clang (clang-cl) 和 Intel 编译器 (icx-cl)。

有关更多详细信息,请参阅 教程

[原型] Eager 模式和 TorchInductor CPP 后端均支持 CPU 路径上的 FP16

Float16 是一种常用的缩减浮点类型,用于提高神经网络推理/训练的性能。自此版本起,CPU 路径上的 eager 模式和 TorchInductor 都支持 float16。

[原型] 自动加载设备扩展

PyTorch 现在支持树外设备扩展的自动加载,通过消除手动导入的需要来简化集成。此功能通过 torch.backends 入口点启用,通过确保无缝扩展加载来简化使用,同时允许用户在需要时通过环境变量禁用它。

有关更多信息,请参阅 教程

[原型] 增强的英特尔 GPU 支持

英特尔 GPU 支持增强功能现已适用于英特尔® 数据中心 GPU Max 系列和英特尔® 客户端 GPU(内置英特尔® Arc™ 显卡的英特尔® 酷睿™ 超处理器和用于 dGPU 部件的英特尔® Arc™ 显卡),旨在让在 PyTorch 2.5 版本中更轻松地加速英特尔 GPU 上的机器学习工作流。我们在此版本中还启用了 PyTorch 在 Windows 上对英特尔® 客户端 GPU 的初步支持。

  • 扩展了 PyTorch 硬件后端支持矩阵,以包括英特尔数据中心和客户端 GPU。  
  • SYCL* 内核的实现旨在增强 Aten 运算符在英特尔 GPU 上的覆盖范围和执行,以提高 PyTorch eager 模式的性能。
  • 增强了 torch.compile 的英特尔 GPU 后端,以提高各种深度学习工作负载的推理和训练性能。

这些功能可通过 PyTorch 预览版和每晚二进制 PIP wheel 获得。有关英特尔 GPU 支持的更多信息,请参阅文档