PyTorch 2.12 版本亮点（7分钟阅读）

# PyTorch 2.12 发布博客 – PyTorch 来源：https://pytorch.org/blog/pytorch-2-12-release-blog/ ### 特色项目 - PyTorch 标志 (https://pytorch.org/projects/pytorch/) 我们激动地宣布 PyTorch® 2.12 版本正式发布！(发布说明: https://github.com/pytorch/pytorch/releases/tag/v2.12.0) PyTorch 2.12 版本包含以下变更： - 由于更新了 cuSolver 后端选择，CUDA 上的批处理 `linalg.eigh` 速度提升高达 100 倍 - 全新的 `torch.accelerator.Graph` API 统一了 CUDA、XPU 以及外部后端上的图捕获与回放 - `torch.export.save` 现在支持微缩 (MX) 量化格式，使得能够完整导出经过大幅压缩的模型 - Adagrad 现在支持 `fused=True`，加入了 Adam、AdamW 和 SGD 的单内核优化器实现行列 - `torch.cond` 控制流现在可以在 CUDA Graph 内部被捕获并回放 - ROCm 用户获得了可扩展内存段、rocSHMEM 对称内存集合以及 FlexAttention 流水线 自 PyTorch 2.11 以来，本次版本包含来自 457 位贡献者的 2,926 次提交。我们衷心感谢我们热心的社区所做出的贡献。一如既往，我们鼓励您尝试这些新功能，并在我们改进 2.12 的过程中报告任何问题。有关如何开始使用 PyTorch 2 系列的更多信息，请访问我们的入门指南页面 (https://pytorch.org/get-started/pytorch-2-x/)。 有问题吗？请于太平洋夏令时 5 月 20 日星期三上午 10 点参加我们的直播问答，小组成员包括 Joe Spisak、Andrey Talman 和 Alban Desmaison，主持人 Chris Gottbrath。我们将简要介绍本次版本的内容，并现场回答您的问题。立即注册 (https://pytorch.org/event/pytorch-2-12-release-live-qa/)。 在整个 2.x 系列中，PyTorch 正从一个以研究为主的框架演变为一个统一的、跨硬件平台的规模化生产训练与推理平台。PyTorch 2.10 (https://pytorch.org/blog/pytorch-2-10-release-blog/) 通过跨后端性能基元和正式弃用 TorchScript 奠定了这一基础。PyTorch 2.11 (https://pytorch.org/blog/pytorch-2-11-release-blog/) 通过可微分的分布式训练集合、在下一代 GPU 上的 FlashAttention-4 以及更广泛的导出覆盖扩展了这一基础。 PyTorch 2.12 延续了这一方向：新的设备无关的 torch.accelerator.Graph API 统一了 CUDA、XPU 和外部后端上的图捕获与回放；批处理特征值分解速度提升高达 100 倍；torch.export 现在支持微缩量化格式，用于部署经过大幅压缩的模型。在这些版本中，PyTorch 在各后端上变得更快，在更广泛的平台上可用，持续推动 AI 创新。 ### **性能特性** **CUDA 上批处理特征分解速度提升高达 100 倍 (`linalg.eigh`)** linalg.eigh 在 CUDA 上的后端选择已经过彻底改革。传统的 MAGMA 后端已被弃用，转而使用 cuSolver（Grayson Derossi 的 PR #174619），并且 cuSolver 分派启发式策略已更新为无条件使用 syevj_batched（Johannes Z 的 PR #175403）。对于批处理对称/厄米特征值问题，这相比之前版本实现了高达 100 倍的加速，解决了与 CuPy 之间长期存在的性能差距。以前需要数分钟的工作负载（因为 PyTorch 低效地逐个处理每个矩阵求解）现在通过使用专为将多个小/中型矩阵作为单个 GPU 操作处理的 cuSolver syevj_batched 内核，只需数秒即可完成。这些提升对依赖批处理矩阵特征分解的科学计算和机器学习工作负载尤其重要。（示例用法见文档 (https://docs.pytorch.org/docs/2.11/generated/torch.linalg.eigh.html)） **融合 Adagrad 优化器** Adagrad 优化器现在支持 fused=True，将整个优化器步骤在单个 CUDA 内核中执行，而不是为每个操作启动独立的内核。这减少了内核启动开销和内存流量。Adagrad 加入了 Adam、AdamW 和 SGD 的行列，提供了融合变体。底层的 CUDA 内核由 @MeetThePatel 在 2.11 周期贡献（PR #159008），Python 前端在 2.12 中由 Jane Xu 最终完成，向用户开放访问（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`。流上下文管理器的重入问题也已修复。这些变更共同为流和图管理带来了跨后端的一致性，初始支持 XPU 后端，并通过 `PrivateUse1` 实现对外部后端的可扩展性。由 Intel 的 Guangye Yu 在六个 PR 中贡献，基于 C++ 接口（PR #171269）和 Python 前端（PR #171285）。（示例用法见文档字符串 (https://github.com/pytorch/pytorch/blob/1d803512199040e98738e95d0dc074acbde9fb5c/torch/accelerator/graphs.py#L11-L48)） **`torch.export` 现在支持微缩 (MX) 量化格式** 随着模型从研究走向生产，`torch.export` 是序列化 PyTorch 模型以进行部署的标准路径。然而，使用微缩 (MX) 量化（一种日益流行的减小模型尺寸和推理成本的技术）的模型之前无法导出，因为 `torch.export.save` 无法处理用作 MX 格式 (MXFP4, MXFP6, MXFP8) 共享块缩放指数的 `float8_e8m0fnu` 数据类型。 在 PyTorch 2.12 中，`torch.export.save` 和 `torch.export.load` 现在能正确序列化和反序列化带有该数据类型的张量，从而解除了利用微缩量化的模型从导出到部署的完整工作流阻塞。这对于将大型语言模型部署到成本受限或边缘环境、需要激进量化的团队尤其重要。由 ARM 的 Chizkiyahu Raful 贡献（PR #176270）。 **使用 torch.cond 在 CUDA Graph 内捕获控制流** 使用 torch.cond 的控制流区域现在可以作为 CUDA Graph 的一部分被捕获和回放。以前，数据相关的控制流因分支在 CPU 上评估而被迫回退到 CUDA 图树。通过利用 CUDA 12.4 的条件 IF 节点，torch.cond 分支现在完全在 GPU 上单次图捕获中评估。该功能由 NVIDIA 的 Daniel Galvez 和 Ting-Yang Kuei 贡献（PR #168912），并由 Meta 的 Paul Zhang 添加了 Inductor 排序支持（PR #179457）。目前该功能与 eager 和 cudagraphs 后端兼容；Inductor 支持计划在未来版本中实现。 **基于 FMA 的 XPU 上 addcdiv 降低** Inductor 现在使用融合乘加 (FMA) 指令进行 addcdiv 操作，实现了与 eager 模式 CUDA 执行的位级数值一致性，同时保留了 Triton 内核融合的优势。 `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 性能分析器事件 API 现在公开了流 ID、流类型、活动类型、未完成事件和 Python 函数事件——使 events() 与 Chrome 跟踪 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 的跟踪分析器现在支持 ncclx 和 gloo 后端，与现有的 nccl 和 xccl 后端并列，实现了跨更广泛集合后端的分布式通信跟踪。此外，FlightRecorder 现在能够识别以前未被跟踪的 torchcomms 操作（如 all_gather_single, reduce_scatter_v, barrier）。在此周期中还修复了一个竞态条件，该条件可能导致多个进程组同时访问 FlightRecorder 单例时出现无限循环。后端白名单由 Meta 的 Lily Janjigian 添加（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 Contexts 现在支持指定工作队列限制，从而对 GPU 资源分区实现更细粒度的控制。这个实验性功能允许用户在绿色上下文中限制并发工作提交的数量，使得跨并发工作负载的资源共享更具可预测性。由 NVIDIA 的 Matthias Jouanneaux 贡献（PR #177242）。 ### **ROCm** **ROCm: 可扩展段** 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 半结构化稀疏性** hipSPARSELt 现在在 ROCm >= 7.12 的 PyTorch 构建中默认启用，为 AMD GPU 带来了半结构化 (2:4) 稀疏性支持。FP8 (float8_e4m3fn) 输入现在也通过 hipSPARSELt 在 MI350X (gfx950) 上得到支持，并输出 FP32。这使得之前仅限 CUDA 的 torch._cslt_sparse_mm 稀疏加速路径变得可用。hipSPARSELt 由 AMD 的 rraminen 启用（PR #170852），FP8 半结构化稀疏性由 Meta 的 Benji Beck 添加（PR #179310）。 **ROCm: Inductor FlexAttention 流水线** AMD GPU 上的 FlexAttention 现在在 Triton 后端使用两级流水线，在 MI350X 上的多种注意力模式（因果、alibi、滑动窗口）和形状上实现了 5-26% 的加速。这是一行配置更改（num_stages=1 变为 2），解锁了更高效的内存-计算重叠。由 nithinsubbiah 贡献（PR #176676）。 ### **Apple MPS** **MPS: Metal-4 离线着色器编译** Apple Silicon 二进制轮子现在附带预编译的 Metal-4 着色器，基于 macOS 26 并使用 metal-4 标准构建。这消除了首次运行时的运行时着色器编译开销，减少了 MPS 工作负载的启动延迟。还添加了一个配套 API（`torch._C._mps_loadMetallib`），用于直接加载预编译的 .metallib 二进制文件，支持 Triton Apple MPS 后端的编译时 metallib 工作流。由 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。 主要变更： - 即时初始化：我们将要求所有 ProcessGroup/通信器在 dist.init_process_group 期间进行即时初始化，并且仅支持单个后端设备。这意味着在初始化期间必须指定设备。 - 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 (https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/export.html) 来替换 jit trace 和 script API，使用 Executorch (https://docs.pytorch.org/executorch/stable/index.html) 来替换嵌入式运行时。