博客

PyTorch 2.12 发布博客

作者: 2026年5月13日2026年5月19日暂无评论

特色项目

我们很高兴地宣布 PyTorch® 2.12 版本正式发布(发布说明)!

PyTorch 2.12 版本包含以下变更:

  • 得益于 cuSolver 后端选择的更新,CUDA 上批量化的 linalg.eigh 速度最高提升了 100 倍
  • 全新的 torch.accelerator.Graph API 统一了 CUDA、XPU 及外部后端中的图捕获和重放功能
  • torch.export.save 现已支持 Microscaling (MX) 量化格式,能够实现对激进压缩模型的完整导出
  • Adagrad 现支持 fused=True,与 Adam、AdamW 和 SGD 一样,采用单内核优化器实现
  • torch.cond 控制流现可在 CUDA Graph 中捕获并重放
  • ROCm 用户现可使用可扩展内存段、rocSHMEM 对称内存集合通信以及 FlexAttention 流水线功能

自 PyTorch 2.11 发布以来,本版本共包含来自 457 位贡献者的 2,926 次提交。我们衷心感谢社区成员的鼎力贡献。一如既往,我们鼓励大家尝试新功能并反馈遇到的任何问题,以帮助我们不断改进 2.12 版本。有关 PyTorch 2 系列入门的更多信息,请访问我们的 入门 页面。

有疑问吗?欢迎参加我们于太平洋标准时间 5 月 20 日周三上午 10 点举行的直播问答环节,嘉宾包括 Joe Spisak、Andrey Talman 和 Alban Desmaison,主持人为 Chris Gottbrath。我们将简要概述本次发布并实时回答您的问题。 立即注册。

在 2.x 系列版本中,PyTorch 已从最初的研究导向框架演进为统一的、硬件无关的平台,支持大规模生产环境下的训练与推理。PyTorch 2.10 通过跨后端性能原语和正式弃用 TorchScript 奠定了基础。PyTorch 2.11 则通过分布式训练的可微集合通信、下一代 GPU 上的 FlashAttention-4 以及更广泛的导出覆盖范围进一步扩展了这一基石。

PyTorch 2.12 延续了这一方向:全新的设备无关的 torch.accelerator.Graph API 统一了 CUDA、XPU 及外部后端中的图捕获和重放;批量特征值分解速度最高提升 100 倍;torch.export 现已支持 Microscaling 量化格式,用于部署高度压缩的模型。在这一系列版本中,PyTorch 正跨后端变得更快,并可在更广泛的平台上使用,持续推动 AI 创新。

性能特性

CUDA 上批量特征值分解 (linalg.eigh) 速度最高提升 100 倍

CUDA 上 linalg.eigh 的后端选择机制经过了全面重构。弃用了传统的 MAGMA 后端,转而采用 cuSolver(由 Grayson Derossi 贡献,PR #174619),并将 cuSolver 调度启发式算法更新为无条件使用 syevj_batched(由 Johannes Z 贡献,PR #175403)。对于批量对称/埃尔米特特征值问题,这比上一版本带来了高达 100 倍的加速,解决了长期以来与 CuPy 之间的性能差距。

此前由于 PyTorch 对每个矩阵求解器单独调度的低效,往往需要数分钟才能完成的工作负载,现在利用 cuSolver 的 syevj_batched 内核(旨在将多个中小型矩阵作为单个 GPU 操作处理),仅需几秒钟即可运行。这些提升对于依赖批量矩阵特征值分解的科学计算和机器学习工作负载尤为重要。(文档中的用法示例

融合版 (Fused) Adagrad 优化器

Adagrad 优化器现已支持 fused=True,它在一个 CUDA 内核中执行整个优化器步骤,而不是为每个操作启动单独的内核。这减少了内核启动开销和内存流量。Adagrad 现已加入 Adam、AdamW 和 SGD 的行列,提供融合版本。底层的 CUDA 内核由 @MeetThePatel 在 2.11 周期内贡献(PR #159008),Jane Xu 在 2.12 版本中完成了 Python 前端对其的封装(PR #177672)。

跨硬件编译与导出

torch.accelerator.Graph:设备无关的加速器图捕获与流 API

`torch.accelerator.Graph` 是一种全新的、与设备无关的图捕获与重放 API,它为后端特定的实现(如 `torch.xpu.XPUGraph`)提供了一种统一的抽象。每个后端都可以通过轻量级的 GraphImplInterface 注册自己的实现,在保持后端自主性的同时实现一致的用户 API。

与此同时,`c10::Stream` 和 `torch.Stream` 现已公开 `is_capturing()` 方法,取代了特定设备的 `is_current_stream_capturing`,提供了后端无关的替代方案。Stream 上下文管理器的重入问题也已修复。这些变更共同实现了流和图管理在跨后端的一致性,最初支持 XPU 后端,并通过 `PrivateUse1` 支持扩展到外部后端。
由 Intel 的 Guangye Yu 贡献,共涉及六个 PR,以 C++ 接口 (PR #171269) 和 Python 前端 (PR #171285) 为核心。(文档注释中的用法示例

torch.export 现支持 Microscaling (MX) 量化格式

随着模型从研究走向生产,torch.export 已成为序列化 PyTorch 模型进行部署的标准途径。然而,之前使用 Microscaling (MX) 量化(一种日益流行的减小模型大小和推理成本的技术)的模型无法导出,因为 torch.export.save 无法处理 MX 格式(MXFP4, MXFP6, MXFP8)中用作共享块缩放指数的 float8_e8m0fnu 数据类型。

在 PyTorch 2.12 中,torch.export.savetorch.export.load 现在可以正确序列化和反序列化具有此数据类型的张量,从而打通了利用 Microscaling 量化模型的完整导出到部署工作流程。这对将大语言模型部署到成本受限或边缘环境(在这些环境中激进的量化至关重要)的团队尤其重要。由 ARM 的 Chizkiyahu Raful 贡献 (PR #176270)。

在 CUDA Graph 中使用 torch.cond 捕获控制流

现在,使用 torch.cond 的控制流区域可以作为 CUDA Graph 的一部分进行捕获和重放。此前,由于分支需要在 CPU 上求值,依赖数据的控制流被迫回退到 CUDA 图树(graph trees)。通过利用 CUDA 12.4 的条件 IF 节点,torch.cond 分支现在可以在单个图捕获内完全在 GPU 上求值。

该功能由 NVIDIA 的 Daniel Galvez 和 Ting-Yang Kuei 贡献 (PR #168912),Inductor 排序支持由 Meta 的 Paul Zhang 添加 (PR #179457)。目前该功能适用于 eager 和 cudagraphs 后端;Inductor 支持计划在未来的版本中提供。

基于 FMA 的 XPU addcdiv 指令降低

Inductor 现在对 addcdiv 操作使用融合乘加 (FMA) 指令,在保持 Triton 内核融合优势的同时,实现了与 eager CUDA 执行的位级数值对等。

addcdiv 是一种融合算术运算(result = input + value × (tensor1 / tensor2)),处于许多优化器更新规则的核心,包括 Adam、AdamW 和 RMSprop。此前,Inductor 的降低过程使用独立的乘法和除法指令,与 eager 模式相比引入了微小的浮点舍入差异。这些差异在数千次训练步骤后累积,使得验证编译模型是否产生数值一致的结果变得困难。

此功能最初由 Meta 的 Michael Lazos 为 CUDA 实现 (PR #174912),随后由 Intel 的 Guangye Yu 扩展至 XPU (PR #176163),修复了 Intel GPU 上的多个数值正确性问题。现在,任何在训练循环中使用 torch.compile 并伴有大量优化器计算的用户,都可以在不牺牲数值可复现性的前提下获得编译带来的性能提升——无论是在 NVIDIA 还是 Intel 硬件上。

分布式训练

自定义算子中的 ProcessGroup 支持

自定义算子现在可以直接接受 ProcessGroup 对象作为参数,而无需调用方将其转换为字符串组名称并在全局注册表中查找。所有 c10d 函数式集合通信算子(all_reduce, reduce_scatter 等)已更新,现同时接受 ProcessGroup 对象和字符串名称。由 Meta 的 Aaron Orenstein 贡献 (PR #172795)。

多 GPU/多节点性能分析优化

PyTorch Profiler Events API 现公开流 ID、流类型、活动类型、未完成事件和 Python 函数事件——使 events() 与 Chrome trace JSON 输出对齐,并支持更丰富的程序化事后分析。此外,现在可以使用新的 seq_num 字段关联跨 rank 的 NCCL 集合通信跟踪——参与同一集合通信的所有 rank 在进程组内共享相同的序列号。这些变更共同显著改进了调试跨多个 GPU 和节点分布式训练性能的工具。API 增强由 Meta 的 Ryan Zhang 贡献 (PR #177888),NCCL seq_num 由 Meta 的 Marvin Dsouza 添加 (PR #177148)。

FlightRecorder:支持 ncclx + gloo 后端

FlightRecorder 的跟踪分析器现在除了现有的 nccl 和 xccl 后端外,还支持 ncclx 和 gloo 后端,从而能够在更广泛的集合通信后端实现分布式通信跟踪。此外,FlightRecorder 现在可以识别此前未追踪的 torchcomms 操作(例如 all_gather_single, reduce_scatter_v, barrier)。本次周期内还修复了一个当多个进程组同时访问 FlightRecorder 单例时可能导致无限循环的竞态条件。后端允许列表由 Lily Janjigian (Meta) 添加 (PR #180268),torchcomms 操作支持由 Tushar Jain 贡献 (PR #178359)。

平台相关更新

CUDA

CUDA Graph 内核标注

torch.cuda.graph 现接受 enable_annotations 关键字参数,将标注元数据(如集合通信算子名称、进程组、消息大小)注入到已捕获 CUDA 图中的各个内核。使用配套的处理脚本(python -m torch.cuda._annotate_cuda_graph_trace)对跟踪进行后处理后,标注会合并到跟踪文件中。这些标注将出现在 Perfetto/Chrome 分析器跟踪中,从而更容易理解重放图中每个内核的操作。由 Meta 的 Shangdi Yu 贡献 (PR #179768)。

CUDA Green Context 工作队列限制

CUDA Green Context 现支持指定工作队列限制,从而实现对 GPU 资源划分的更细粒度控制。此实验性功能允许用户限制 green context 内并发工作提交的数量,从而在并发工作负载中实现更可预测的资源共享。由 NVIDIA 的 Matthias Jouanneaux 贡献 (PR #177242)。

ROCm

ROCm:可扩展段 (Expandable segments)

AMD GPU(ROCm >= 7.02)现支持 PyTorch 缓存分配器中的可扩展内存段,与 CUDA 的该特性保持一致,通过虚拟内存 API 动态增加分配,减少内存碎片。由 AMD 的 Prachi Gupta 添加 (PR #173330)

ROCm:rocSHMEM 支持

rocSHMEM 支持在 AMD GPU 上启用对称内存集合通信操作 (torch.ops.symm_mem.*),将基于 NVSHMEM 的 GPU 上通信原语(包括点对点、广播、全对全以及面向 MoE 的 2D AllToAllv)移植到了 ROCm。rocSHMEM 实现使用专门的编译单元来处理 NVSHMEM 和 rocSHMEM 之间的 API 和 warp 大小差异。由 Prachi Gupta 贡献 (PR #173518)。

ROCm:hipSPARSELt 和 FP8 半结构化稀疏性

在 ROCm >= 7.12 的 PyTorch 构建中,hipSPARSELt 现已默认启用,为 AMD GPU 带来半结构化 (2:4) 稀疏性支持。现在通过 hipSPARSELt 在 MI350X (gfx950) 上也支持 FP8 (float8_e4m3fn) 输入,并支持 FP32 输出。这启用了与此前仅在 CUDA 上可用的相同的 torch._cslt_sparse_mm 稀疏加速路径。hipSPARSELt 由 rraminen (AMD) 启用 (PR #170852),FP8 半结构化稀疏性由 Benji Beck (Meta) 添加 (PR #179310)。

ROCm:Inductor FlexAttention 流水线

AMD GPU 上的 FlexAttention 现已在 Triton 后端使用两级流水线,在 MI350X 上针对多种注意力模式(causal, alibi, sliding window)和形状带来了 5-26% 的速度提升。这是一个配置更改(num_stages 从 1 改为 2),能够释放更高效的内存-计算重叠。由 nithinsubbiah 贡献 (PR #176676)。

Apple MPS

MPS:Metal-4 离线着色器编译

Apple Silicon 二进制 wheel 文件现随附提前编译的 Metal-4 着色器,基于 macOS 26 和 metal-4 标准构建。这消除了首次运行时的着色器编译开销,减少了 MPS 工作负载的启动延迟。由 Isalia20 (Irakli Salia) 贡献 (PR #179378)。

弃用与破坏性变更

分布式:torchcomms 的计划内破坏性变更

我们一直致力于将 torchcomms 直接集成到 PyTorch Distributed 中,以便每个人都能开箱即用。在即将到来的版本 (2.13+) 中,我们计划默认使用 torchcomms,这包括一些关于 ProcessGroup 如何运行的破坏性变更。我们的目标是使这些变更对大多数模型自动生效,并修复生态系统中的任何不兼容问题,但尽管如此,仍会有部分模型受到影响。

我们仍在完善 torchcomms,但您现在就可以使用它并获得新的 API、故障恢复、窗口、可扩展性和可调试性特性。要开始使用,请运行 pip install torchcomms 并设置 TORCH_DISTRIBUTED_USE_TORCHCOMMS=1

详情请参见 https://github.com/meta-pytorch/torchcomms

关键变更

  • Eager 初始化:我们将要求所有 ProcessGroup/通信器必须在 dist.init_process_group 期间进行 eager 初始化,且仅支持单一后端设备。这意味着在初始化期间必须指定设备。
  • P2P 操作:我们的目标是使每个 ProcessGroup/通信器与底层通信器 1:1 匹配。这意味着在同一组/流上发出的 P2P 操作将无法保证并发运行。并发 P2P 操作将需要使用批处理 API 或单独的组/通信器。
  • torchcomms 依赖:我们计划将 torchcomms 作为 PyTorch Distributed 的必需包,并弃用现有的 c10d::Backends,转而采用单一、更现代的通信定义。

torchcomms 集成由 PyTorch Distributed 团队领导,2.12 版本中已奠定基础,包括 Yifan Mao 的后端包装重构 (PR #177157) 和 Tushar Jain 的 FlightRecorder 集成 (PR #175270)。

Torchscript 现已弃用

Torchscript 已在 2.10 中弃用,应使用 torch.export 替换 jit trace 和 script API,并应使用 Executorch 替换嵌入式运行时。更多详细信息,请参阅 PTC 的 此讲座

CUDA 12.8 Wheel 的弃用

从 PyTorch 2.12 开始,CUDA 12.8 二进制 wheel 已弃用,将不再作为标准发布矩阵的一部分进行发布。默认的 wheel 依然是 CUDA 13.0(通过 PyPI 的 pip install torch 安装),并添加了 CUDA 13.2 作为实验性构建。

运行在较旧架构(如 Pascal, Volta)上的用户应切换到 CUDA 12.6 wheel,该版本在本版本中依然得到支持。运行在较新 GPU(如 Blackwell)上的用户应使用 CUDA 13.0+ 的 wheel;请注意,这需要将 NVIDIA 驱动程序升级至 580.65.06 (Linux) 或 580.88 (Windows)。

更新 (2026-05-19): 移除了以下句子,因为 pytorch/pytorch#177276 未在 2.12 版本中合并:此外还添加了一个配套 API (torch._C._mps_loadMetallib),用于直接加载预编译的 .metallib 二进制数据,支持 Triton Apple MPS 后端的编译时 metallib 工作流。