跳转到主要内容
博客

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 函数模式的支持,允许用户覆盖任何 *torch.** 操作以实现自定义用户定义行为。
  • Mega Cache,允许用户对 torch 进行端到端可移植缓存;
  • FlexAttention 的新功能——LLM 首个 token 处理、LLM 吞吐量模式优化和推理用 Flex Attention。

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

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

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

测试版功能

[测试版] Torch.Compile 对 Torch 函数模式的支持

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

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

[测试版] 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 Toolkit 发布说明

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

更多上下文信息也可在此处找到。

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

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

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

请参阅此处的教程。

[原型] 增强 Intel GPU 加速

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

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

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

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

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

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

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

通过采用新的 C++ micro-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 融合支持

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