跳转到主要内容
博客

PyTorch 2.7 发布

作者: 2025 年 4 月 23 日2025 年 5 月 15 日暂无评论

我们很高兴地宣布 PyTorch® 2.7 版本发布(版本说明)!此版本主要包括:

  • 支持 NVIDIA Blackwell GPU 架构,并提供适用于 Linux x86 和 arm64 架构的 CUDA 12.8 预构建 wheel 包。
  • torch.compile 支持 Torch Function Modes,这使用户能够覆盖任何 *torch.** 操作,以实现自定义的用户定义行为。
  • Mega Cache,允许用户对 torch 进行端到端的可移植缓存;
  • FlexAttention 的新功能——LLM 首个 token 处理、LLM 吞吐量模式优化以及用于推理的 Flex Attention。

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

Beta 原型
Torch.Compile 支持 Torch Function Modes NVIDIA Blackwell 架构支持
Mega Cache PyTorch 原生上下文并行
增强英特尔 GPU 加速
FlexAttention LLM x86 CPU 上的首个 token 处理
FlexAttention LLM x86 CPU 上的吞吐量模式优化
Foreach Map
用于推理的 Flex Attention
Inductor 中的 Prologue Fusion 支持

*要查看所有已提交的公共功能列表,请点击此处

BETA 功能

[Beta] Torch.Compile 支持 Torch Function Modes

此功能使用户能够覆盖任何 *torch.** 操作,以实现自定义的用户定义行为。例如,可以重写操作以适应特定的后端。这在 FlexAttention 中用于重写索引操作。

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

[Beta] Mega Cache

Mega Cache 允许用户对 torch 进行端到端的可移植缓存。预期的用例是在编译和执行模型后,用户调用 torch.compiler.save_cache_artifacts(),它将以可移植形式返回编译器工件。之后,可能在不同的机器上,用户可以使用这些工件调用 torch.compiler.load_cache_artifacts(),以预填充 torch.compile 缓存,从而启动他们的缓存。

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

原型功能

[原型] NVIDIA Blackwell 架构支持

PyTorch 2.7 引入了对 NVIDIA 新的 Blackwell GPU 架构的支持,并提供了 CUDA 12.8 的预构建 wheel 包。有关 CUDA 12.8 的更多详细信息,请参阅CUDA 工具包发布说明

  • cuDNN、NCCL 和 CUTLASS 等核心组件和库已升级,以确保与 Blackwell 平台兼容。
  • PyTorch 2.7 包含 Triton 3.3,它增加了对 Blackwell 架构与 torch.compile 兼容性的支持。
  • 要利用这些新功能,请使用以下命令安装 PyTorch 和 CUDA 12.8:pip install torch==2.7.0 –index-url https://download.pytorch.org/whl/cu128

更多上下文信息也可在此处找到:https://github.com/pytorch/pytorch/issues/145949

[原型] PyTorch 原生上下文并行

PyTorch 上下文并行 API 允许用户创建一个 Python 上下文,以便其中每个 *torch.nn.functional.scaled_dot_product_attention() *调用都将以上下文并行方式运行。目前,PyTorch 上下文并行支持 3 种注意力后端:1. Flash attention;2. Efficient attention;3. cuDNN attention。

例如,这在 TorchTitan 中用作 LLM 训练的上下文并行解决方案

请参阅此处的教程。

[原型] 增强英特尔 GPU 加速

此最新版本引入了针对英特尔 GPU 架构的增强性能优化。这些改进通过以下关键增强功能加速了各种英特尔 GPU 上的工作负载:

  • 在 Windows 11 上为英特尔 GPU 启用 torch.compile,提供与 Linux 上相同的优于 eager 模式的性能优势。
  • 优化英特尔 GPU 上 PyTorch 2 Export 训练后量化 (PT2E) 的性能,提供具有增强计算效率的完整图模式量化管道。
  • 改进 bfloat16 和 float16 的 Scaled Dot-Product Attention (SDPA) 推理性能,以加速英特尔 GPU 上的基于注意力的模型。
  • 在 Linux 上启用 AOTInuctor 和 torch.export 以简化部署工作流。
  • 实现更多 Aten 算子,以增强英特尔 GPU 上算子执行的连续性,并提高英特尔 GPU 在 eager 模式下的性能。
  • 在 Windows 和 Linux 上启用 profiler,以方便模型性能分析。
  • 将英特尔 GPU 支持扩展到 Intel® Core™ Ultra 系列 2 与 Intel® Arc™ 显卡,以及 Windows 和 Linux 上的 Intel® Arc™ B 系列显卡

有关英特尔 GPU 支持的更多信息,请参阅入门指南

另请参阅此处此处的教程。

[原型] FlexAttention LLM x86 CPU 上的首个 token 处理

FlexAttention x86 CPU 支持首次在 PyTorch 2.6 中引入,通过 TorchInductor C++ 后端提供了优化的实现(例如 PageAttention,这对于 LLM 推理至关重要)。在 PyTorch 2.7 中,支持更多用于 LLM 首个 token 处理的注意力变体。通过此功能,用户可以在 x86 CPU 上更流畅地运行 FlexAttention,将特定的 scaled_dot_product_attention 算子替换为统一的 FlexAttention API,并在使用 torch.compile 时受益于通用支持和良好性能。

[原型] FlexAttention LLM 吞吐量模式优化

通过采用新的 C++ 微型 GEMM 模板功能,FlexAttention 在 x86 CPU 上用于 LLM 推理吞吐量场景的性能得到了进一步提升。这解决了 PyTorch 2.6 中存在的大批量场景的性能瓶颈。通过此增强功能,用户在使用 FlexAttention API 和 torch.compile 进行 x86 CPU 上的 LLM 吞吐量服务时,可以透明地受益于更好的性能和更流畅的体验。

[原型] Foreach Map

此功能使用 torch.compile 允许用户将任何逐点或用户定义函数(例如 torch.add)应用于张量列表,类似于现有的 *torch.foreach** 算子。与现有的 *torch.foreach** 算子相比,主要优点是可以将标量或张量列表的任意组合作为参数提供,甚至可以将用户定义的 python 函数提升以应用于张量列表。Torch.compile 将自动生成水平融合的内核以实现最佳性能。

请参阅此处的教程。

[原型] 用于推理的 Flex Attention

在 2.5.0 版本中,FlexAttention* torch.nn.attention.flex_attention* 被引入,旨在为希望自定义注意力内核而无需编写内核代码的机器学习研究人员提供便利。此更新引入了一个针对推理进行优化的解码后端,支持 GQA 和 PagedAttention,以及功能更新,包括嵌套锯齿张量支持、性能调优指南和可训练偏差支持。

[原型] Inductor 中的 Prologue Fusion 支持

Prologe 融合通过将矩阵乘法 (matmul) 之前的操作融合到 matmul 内核本身中来优化矩阵乘法操作,通过减少全局内存带宽来提高性能。