PyTorch Foundation 周三发布 PyTorch 2.12——自 2.11 以来 2,926 commits、457 贡献者。性能头条来自一个不起眼但非常重要的后端切换:CUDA 上批量的 `linalg.eigh` 加速最高达 100×,原因是放弃了遗留的 MAGMA 后端,改用 cuSolver,并把 dispatch 启发式更新为无条件使用 `syevj_batched`。以前要跑几分钟的工作负载——科学计算和任何需要批量矩阵特征分解的 ML 训练步骤里都很常见——现在几秒就跑完,终于补上了和 CuPy 之间的长期差距。另外,Adagrad 加入 Adam、AdamW、SGD 的行列支持 `fused=True`,把整步优化器塞进一个 CUDA kernel,而不是多个独立 launch。XPU 上基于 FMA 的 `addcdiv` lowering 让 `torch.compile` 与 eager 模式在 Intel GPU 上达到位精度一致——单看是小事,但如果你被编译后的优化器行为无法复现追着跑过,这就是承重项。
战略层面的故事是跨后端转向。PyTorch 2.12 引入 `torch.accelerator.Graph`,一个与设备无关的图捕获与回放 API,把 `torch.xpu.XPUGraph` 这种后端专属实现抽象统一起来。每个后端通过一个轻量 `GraphImplInterface` 注册;`c10::Stream` 和 `torch.Stream` 现在暴露 `is_capturing()`,取代了与设备绑定的 `is_current_stream_capturing`。初版支持 CUDA 与 XPU(Intel),并通过 `PrivateUse1` 向外树后端开放扩展。这正是本周早些时候 CUDA 护城河报道里讲的那个动态在框架层级上的对照物:PyTorch 在明确地让非 CUDA 加速器更容易参与到图捕获编程模型中,而图捕获就是 Inductor 与大部分性能工作所在的那一层。ROCm 用户在同一版本里获得可扩展内存段、rocSHMEM 对称内存集合通信、以及 FlexAttention 流水线——AMD 路径明显推进了一步。
还有两件值得 builder 注意的变化。第一,`torch.export.save` 和 `torch.export.load` 现在能正确地序列化 `float8_e8m0fnu` 这一 dtype——Microscaling(MX)量化格式(MXFP4、MXFP6、MXFP8)里作为块共享缩放指数使用的那个。过去,使用这种激进量化技术的模型无法通过 PyTorch 标准的导出路径走完;给成本受限或 edge 环境出货 LLM 的团队不得不自己发明自定义序列化。这个缺口被堵上了。第二,`torch.cond` 控制流现在可以借助 CUDA 12.4 的条件 IF 节点在 CUDA Graph 里被捕获和重放——以前数据依赖分支会强制回退到 CUDA graph trees,因为分支在 CPU 上计算。在 2.12,分支判断在同一次图捕获里就留在 GPU 上完成(目前支持 eager 与 cudagraphs 后端;Inductor 支持计划在以后)。对带数据依赖控制流的 agent 类或 RL 工作负载,这一改动消掉了一个实打实的性能悬崖。
对 builder:如果你已经在 2.x 路线上,这是一个低摩擦版本——大多数收益是透明的(eigh dispatch、fused Adagrad、addcdiv FMA),不需要改代码。如果你在出货量化模型,Microscaling 导出修复是立刻可用的——按文档里的示例 clone 一份,重测你的导出 pipeline。torch.accelerator.Graph API 主要对 Intel XPU 目标或自建加速器的人有意义——更长远的意义是:PyTorch 现在在明确地把图捕获跨后端抽象化,这是任何认真挑战 CUDA 的玩家必须建在上面的地基。分布式训练团队应该看一下 FlightRecorder 的 ncclx + gloo 支持,以及 NCCL 新的 `seq_num` 用于跨 rank 的集合通信关联——都是具体可感的 debugging 体验改进。如果你维护一条训练 pipeline,这份 release notes 值得从头到尾读一遍;上面这些头条只是采样,不是全部。