page contents

PyTorch 2.15 发布:foreach,scaled_addmm,优化器参数传递方式调整,clamp 饱和处理

PyTorch 2.15 来了。在这个版本中,Python 3.10 wheel 停止发布,Python 3.15 支持更完整;torch.foreach 从内部接口走向公开 beta API;大词表训练和低精度矩阵乘法继续补齐;优化器的参数传递也有变化;新增 scaled_addmm 和它的本地版本;一批基础算子和自动求导修复也值得注意。

attachments-2026-11-lzMNmrhj6ae7e7d0687b1.png

PyTorch 2.15 来了。在这个版本中,Python 3.10 wheel 停止发布,Python 3.15 支持更完整;torch.foreach 从内部接口走向公开 beta API;大词表训练和低精度矩阵乘法继续补齐;优化器的参数传递也有变化;新增 scaled_addmm 和它的本地版本;一批基础算子和自动求导修复也值得注意。

目录

Python 版本相关变化

新特性

错误修复

限制情况

升级建议

Python 和 CUDA 依赖

PyTorch 2.15 的二进制包不再支持 Python 3.10。还停留在 Python 3.10 的项目,需要先升级到 Python 3.11 或更新版本,或者继续留在 PyTorch 2.14。

Python 3.15 这边则更完整了。PyTorch 2.14 已经提供 Python 3.15 wheel,但 torch.compile 还不能在 Python 3.15 上使用;到 2.15,torch.compile 已经支持 Python 3.15。如果你想尝试新 Python 版本,同时项目里依赖编译模式,2.15 才是更合适的起点。

CUDA wheel 也有变化:PyPI 默认 CUDA wheel 从 CUDA 13.0 升到 CUDA 13.2,同时提供 CUDA 13.4 wheel;CUDA 12.6 和 CUDA 13.0 wheel 不再发布。CUDA 13 wheel 需要较新的 NVIDIA 驱动和 Turing 或更新架构的 GPU。这个变化对模型代码本身影响不大,但会影响开发机、CI、Docker 镜像和线上推理镜像,升级前要先核对环境。

新特性

torch.foreach 成为公开 beta API

PyTorch 很多优化器内部会同时更新一组张量。例如参数更新可以理解成:

for p in params:
    p.add_(p.grad, alpha=-lr)

这种 Python 循环容易理解,但当参数很多时,每个小张量操作都要经过一次 Python 调度,开销会比较明显。PyTorch 早就有一组内部 _foreach_* 接口,用来把“一组张量上的同类操作”批量提交给底层实现。2.15 把它们整理成公开的 torch.foreach 模块,目前是 beta。

过去如果想用这类接口,可能会写内部 API:

torch._foreach_mul_(params, 0.99)

现在可以改用公开模块:

import torch

torch.foreach.mul_(params, 0.99)

这对普通模型训练脚本不是必需的,因为 torch.optim 已经会在合适场景下使用 foreach 路径。但对写自定义优化器、EMA、梯度裁剪、参数批量后处理的用户来说,公开 API 的意义很大:不用再依赖带下划线的内部接口。

一个简单的 EMA 更新可以写成这样:

@torch.no_grad()
def update_ema(ema_params, model_params, decay=0.999):
    torch.foreach.mul_(ema_params, decay)
    torch.foreach.add_(ema_params, model_params, alpha=1 - decay)

大词表训练:linear_cross_entropy 继续补齐

PyTorch 2.14 里,linear_cross_entropy 已经是一个值得关注的新方向:把最后一层线性投影和交叉熵合在一起,避免完整物化 [batch, seq, vocab] 这样巨大的 logits。对大词表语言模型来说,这能降低峰值显存。

传统写法通常是:

import torch.nn.functional as F

logits = F.linear(hidden, weight)
loss = F.cross_entropy(
    logits.flatten(0, 1),
    target.flatten(0, 1),
)

融合写法是:

import torch.nn.functional as F

loss = F.linear_cross_entropy(
    hidden,
    weight,
    target,
    reduction="mean",
)

编辑注:这里的“融合”不只是少启动几个 CUDA kernel。torch.compile 确实可以把很多普通算子合成更紧凑的执行图,减少中间张量读写;但 linear_cross_entropy 的核心价值,是把线性层和交叉熵看成一个整体来重新组织计算,尽量不生成完整的 logits。普通 F.linear(...); F.cross_entropy(...) 在语义上已经产生了一个真实的 [batch, seq, vocab] 张量,编译器不能轻易假设它一定可以被丢掉,尤其还要考虑 autograd、别名、hook、动态 shape、不同 target 形式和二阶梯度等边界。linear_cross_entropy 相当于显式告诉 PyTorch:用户要的不是中间 logits,而是最终 loss;实现就可以按分块方式边算 logits、边更新 logsumexp 和目标类别项,用更专门的前向/反向路径换取显存下降。这类优化更接近 scaled_dot_product_attention 或 binary_cross_entropy_with_logits 的意义:不是用户不能手写旧组合,而是把高层语义暴露出来后,框架才能做普通算子图里很难稳妥推断的算法级融合。

2.15 继续补齐这个路径:当目标是 soft label,并且使用分块实现时,linear_cross_entropy 现在支持 reduction="none"。soft label 常见于标签平滑、知识蒸馏、多教师模型融合等场景;reduction="none" 则表示保留每个样本或每个位置的损失,后面自己加权。

import torch
import torch.nn.functional as F

hidden = torch.randn(2, 4, 16, device="cuda")
weight = torch.randn(100, 16, device="cuda")

# soft label:最后一维是词表概率分布
target = torch.randn(2, 4, 100, device="cuda").softmax(dim=-1)

loss = F.linear_cross_entropy(
    hidden,
    weight,
    target,
    reduction="none",
)
print(loss.shape)

这类接口不是每个分类任务都需要;小词表分类直接 F.cross_entropy 就很好。但如果你在训练语言模型、推荐模型或其他超大类别数模型,2.15 的改进说明这条融合路径正在继续成熟。

低精度矩阵乘法:scaled_addmm

这一版还新增了 torch.nn.functional.scaled_addmm 和原地版本 scaled_addmm_。它可以在 CUDA 上把 beta * input + alpha * scaled_mm(mat1, mat2) 合成一次 cuBLASLt GEMM,并且支持编译模式。对普通用户来说,可以先把它理解成面向低精度矩阵乘法和缩放矩阵乘法的更融合接口,主要服务 FP8、量化、推理或大模型训练中的高性能路径。

import torch
import torch.nn.functional as F

# scale_recipe_a / scale_recipe_b 与 scaled_mm 中的缩放方案一致。
# 具体取值取决于你的 FP8 或块缩放 recipe。
out = F.scaled_addmm(
    input,
    mat1,
    mat2,
    scale_a,
    scale_recipe_a,
    scale_b,
    scale_recipe_b,
    beta=1.0,
    alpha=1.0,
)

如果你的项目没有直接使用 scaled_mm 或 FP8 路径,可以先不关心它;如果你维护低精度训练或推理代码,2.15 提供了一个更标准的入口。

clamp 和无符号整数支持更自然

torch.clamp、torch.maximum、torch.minimum 在 CPU 和 CUDA 上新增了 uint16、uint32、uint64 支持。与此同时,整数张量遇到越界标量边界时,clamp 和 hardtanh 会采用饱和语义,而不是包装或直接报错。

import torch

x = torch.tensor([0, 10, 255], dtype=torch.uint8)

# min=-255 对 uint8 来说低于可表示范围;
# 2.15 中会按饱和边界处理,相当于不改变这些值
print(torch.clamp(x, min=-255))

这对图像、量化、低精度整数张量处理更友好。以 uint8 图像为例,用户写 min=-255 时,大概率是想表达“下界比所有合法值都小”,而不是想让 -255 在无符号整数里发生包装。

错误修复

PyTorch 2.15 的修复列表很长。这里不逐条翻译 release notes,只挑几类 Python 用户容易受益或容易踩到的正确性修复。

linspace、logspace 和 arange 的数值更可靠

torch.linspace、torch.logspace、torch.arange 常被用来生成时间步、采样网格、位置编码和索引。它们通常不在模型主体里,但一旦结果错了,后面可能表现为形状不匹配、采样点偏移,或者某个边界 batch 才出现异常。

PyTorch 2.15 修复了 float16 / bfloat16 的 linspace 在 CPU 上因为步长先被舍入到输出 dtype 而丢精度的问题,也修复了 CUDA 上 linspace / logspace 以及大范围序列生成中的若干索引溢出问题。

import torch

# 用半精度生成采样点时,步长计算更可靠
t = torch.linspace(0, 10, 38, dtype=torch.float16)
print(t[3])

# 大范围索引或时间步生成也更适合作为后续张量操作的输入
idx = torch.arange(0, 4096, device="cuda")

这类修复的意义不是让某个模型突然变快,而是减少“预处理和索引生成看起来很普通,却在边界形状下出错”的风险。

torch.load 在 meta 和 FakeTensorMode 下保持 storage sharing

很多大模型工具会先把 checkpoint 加载到 meta 设备或 FakeTensor 环境里,只读取形状、dtype 和共享关系,而不真正分配完整权重。这样可以在转换、切分、检查模型结构时节省内存。

PyTorch 2.15 修复了 torch.load 配合 map_location="meta" 或 FakeTensorMode 时丢失 storage sharing 的问题。storage sharing 指多个张量共享同一块底层存储,例如一个权重和它的 view。这个关系如果在加载时丢失,后续工具可能会误判内存占用,也可能在重建权重别名关系时出错。

import torch

base = torch.arange(8)
state = {
    "weight": base,
    "weight_view": base.view(2, 4),
}

torch.save(state, "shared.pt")
meta_state = torch.load("shared.pt", map_location="meta")

print(meta_state["weight"].shape)
print(meta_state["weight_view"].shape)

对普通训练用户来说,这个修复可能不显眼;对模型转换、checkpoint 检查、懒加载和导出工具来说,它很实用。

torch.asarray 更尊重 buffer protocol 的 dtype

torch.asarray 可以从已有 Python 对象创建张量,包括支持 buffer protocol 的对象。buffer protocol 会携带底层数据的格式信息,例如元素是 16 位整数还是 32 位浮点数。

PyTorch 2.15 修复了 torch.asarray 对 buffer-protocol 对象取 dtype 的问题。升级后,这类对象转换成张量时会更尊重 buffer 自身声明的格式。

import array
import torch

values = array.array("h", [1, 2, 3])  # signed short
x = torch.asarray(values)

print(x)
print(x.dtype)

这对数据读取、二进制协议解析、和 NumPy / Python 标准库对象互操作的代码更有意义。以前如果 dtype 被推错,后面的模型输入可能已经不是你以为的精度或整数范围。

线性代数和二阶自动求导更稳

PyTorch 2.15 修复了 torch.linalg.det、torch.linalg.slogdet、torch.linalg.solve 的二阶 forward AD 问题,也修复了 torch.linalg.matrix_exp 在部分输入上的错误结果。

这些函数常见于科学计算、物理建模、可微优化、几何深度学习和概率模型。普通反向传播只需要一阶梯度,但一些算法会继续对梯度求导,或者需要更高阶的雅可比 / 海森信息。线性代数算子的高阶导数如果错了,模型不一定立刻报错,却可能让优化结果偏掉。

import torch

A = torch.randn(4, 4, dtype=torch.double, requires_grad=True)
b = torch.randn(4, dtype=torch.double)

def objective(mat):
    x = torch.linalg.solve(mat @ mat.mT + torch.eye(4), b)
    return x.square().sum()

grad, = torch.autograd.grad(objective(A), A, create_graph=True)
score = grad.square().sum()
score.backward()

如果你的项目里有线性方程求解、行列式、矩阵指数或二阶梯度,升级后建议重新跑数值一致性测试。

complex dropout 可以正常使用

PyTorch 2.15 修复了 torch.nn.functional.dropout 对 complex 输入直接报 NotImplementedError 的问题。复数张量在信号处理、频域神经网络、量子模拟和部分科学计算任务里会出现,dropout 虽然不是所有复数模型都需要,但它作为基础神经网络操作,应该能自然参与训练图。

import torch
import torch.nn.functional as F

x = torch.randn(8, 16, dtype=torch.cfloat, requires_grad=True)
y = F.dropout(x, p=0.1, training=True)

loss = y.abs().square().mean()
loss.backward()

这一修复和后面“complex 排序 / 比较异常类型变化”并不矛盾。前者是一个本应支持的神经网络操作补齐,后者是对复数没有大小顺序的操作给出更合适的异常类型。

fmin 和 fmax 遇到 NaN 时梯度更正确

torch.fmin 和 torch.fmax 会在两个输入之间逐元素选择较小值或较大值,并按浮点语义处理 NaN。它们常出现在带边界的数值计算、损失函数裁剪和数据清洗逻辑里。

PyTorch 2.15 修复了其中一个输入是 NaN 时的梯度问题。对普通前向计算来说,NaN 已经是需要关注的信号;但如果梯度传播路径也不对,排查会更困难。

import torch

x = torch.tensor([1.0, float("nan"), 3.0], requires_grad=True)
y = torch.tensor([2.0, 2.0, float("nan")], requires_grad=True)

out = torch.fmin(x, y).nan_to_num(0.0).sum()
out.backward()

print(x.grad)
print(y.grad)

如果你的训练里会用 fmin / fmax 处理带 NaN 的中间值,升级后相关梯度路径会更可靠。但更好的工程习惯仍然是尽早定位 NaN 来源,而不是只在末端吞掉它。

限制情况

这一节只列升级时需要注意的限制和行为变化。

优化器参数不要传 set

PyTorch 2.15 开始,当优化器参数直接传入裸 set 时,会发出 FutureWarning,未来版本会变成错误。原因是 set 的迭代顺序不稳定,而优化器状态通常按参数顺序保存和恢复。顺序变了,AdamW 这类优化器里的动量、二阶矩估计等状态就可能对到错误参数上。

# 不推荐:set 的顺序不稳定
optimizer = torch.optim.AdamW({p for p in model.parameters()}, lr=1e-3)

# 推荐:直接使用 model.parameters()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

model.parameters() 返回的是按模块注册顺序遍历的迭代器,直接传给优化器没有问题。如果需要筛选部分参数,建议用列表推导式保留 named_parameters() 的顺序。

params = [
    p for name, p in model.named_parameters()
    if not name.startswith("backbone.")
]
optimizer = torch.optim.AdamW(params, lr=1e-3)

torch.load 的默认 allowlist 更严格

torch.load(..., weights_only=True) 的默认 allowlist 更严格。checkpoint 里包含 nested tensor 时,需要把 nested tensor 的重建函数和类型显式加入安全列表。

import torch
from torch.nested._internal.nested_tensor import NestedTensor, _rebuild_njt

with torch.serialization.safe_globals([_rebuild_njt, NestedTensor]):
    loaded = torch.load("nested_tensor.pt", weights_only=True)

如果是完全可信的旧 checkpoint,也可以退回 weights_only=False,但这会重新接受普通 pickle 反序列化风险。

torch 不再支持 reload

torch 不再支持在同一个进程里重新初始化。普通 import torch 不受影响,但 importlib.reload(torch)、删除 sys.modules["torch"] 后再导入这类热重载写法会报错。

import importlib
import torch

importlib.reload(torch)  # 2.15: ImportError

编辑注:这里的关键不是“重复 import”本身,而是重新执行 torch 的初始化代码。importlib.reload() 并不会把一个模块干净卸载后再导入,它会在原模块对象上重新执行模块级代码;del sys.modules["torch"] 也只是删除 Python 导入缓存里的名字,不会卸载已经加载的 torch._C 动态库,更不会清空 PyTorch 在 C++ 侧维护的 dispatcher、算子注册表、设备后端、内存分配器、线程池等进程级全局状态。因此第二次初始化可能撞上第一次留下的注册记录,例如重复注册同一个 TORCH_LIBRARY namespace;如果第一次导入中途失败,同一进程内重试还可能遇到半初始化状态。2.15 的变化可以理解为:与其让用户继续得到误导性的重复注册错误甚至崩溃,不如直接说明 torch 每个进程只能初始化一次。

complex 排序和比较异常类型变化

对 complex 张量没有定义的排序和比较类操作,现在更倾向于抛出 TypeError,而不是 RuntimeError。如果测试或容错逻辑只捕获 RuntimeError,需要同步调整。

try:
    torch.sort(torch.randn(3, dtype=torch.cfloat))
except TypeError:
    pass

CUDA histc 的 out dtype 必须一致

CUDA 上的 torch.histc(..., out=...) 要求 out 和输入 dtype 一致。需要整数结果时,先按输入 dtype 计算,再显式转换。

x = torch.randn(1000, device="cuda")
out = torch.empty(10, dtype=torch.int64, device="cuda")
out.copy_(torch.histc(x, bins=10).to(out.dtype))

selective activation checkpointing 对副作用更保守

selective activation checkpointing 对带有有序副作用的算子更保守。CheckpointPolicy.PREFER_RECOMPUTE 可能改为保存输出,避免反向时重放副作用;自定义 checkpoint policy 需要检查。

policy = lambda ctx, op, *args, **kwargs: CheckpointPolicy.PREFER_RECOMPUTE

torch.ops 的 warning 进入 warnings 模块

通过 torch.ops.* 调用的自定义算子 warning 会进入 Python 的 warnings 模块。测试环境如果使用 -W error,过去只是打印的底层 warning 现在可能让测试失败。

import warnings

with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    torch.ops.my_lib.my_op(x)

升级建议

如果准备从 PyTorch 2.14 升级到 2.15,可以优先检查下面几类内容:

基础环境:确认 Python 版本、CUDA wheel、驱动和 GPU 架构是否匹配 2.15。

新 API 试用:如果维护自定义优化器、EMA 或梯度批量处理代码,可以评估是否从 _foreach_* 迁到 torch.foreach。

大词表训练:如果已经在试用 linear_cross_entropy,关注 soft label 和 reduction="none" 路径;如果仍手动物化巨大 logits,可以评估融合接口是否适合。

错误修复回归测试:重新跑 linspace / arange、checkpoint 结构检查、buffer 输入转换、线性代数和复数训练相关测试。

限制情况:搜索 torch.load(、importlib.reload(torch)、优化器参数里的 set、complex 排序比较、torch.histc(..., out=...) 和自定义 torch.ops warning。

总体来看,PyTorch 2.15 更像一次基础能力和行为边界的整理:新 Python 支持更完整,一些内部多张量接口走向公开,错误修复覆盖了序列生成、checkpoint、dtype 推断、线性代数和复数训练,默认行为也比过去更明确。普通用户升级时可以重点关注环境和限制情况;工具链和大模型训练用户,则更值得测试 torch.foreach、linear_cross_entropy、scaled_addmm 以及相关正确性修复。

更多相关技术内容咨询欢迎前往并持续关注好学星城论坛了解详情。

想高效系统的学习Python编程语言,推荐大家关注一个微信公众号:Python编程学习圈。每天分享行业资讯、技术干货供大家阅读,关注即可免费领取整套Python入门到进阶的学习资料以及教程,感兴趣的小伙伴赶紧行动起来吧。

attachments-2022-05-rLS4AIF8628ee5f3b7e12.jpg

 

你可能感兴趣的文章

相关问题

0 条评论

请先 登录 后评论
Pack
Pack

2557 篇文章

作家榜 »

  1. Pack 2557 文章
  2. 轩辕小不懂 2403 文章
  3. 小柒 2228 文章
  4. Nen 576 文章
  5. 王昭君 216 文章
  6. 文双 71 文章
  7. 小威 64 文章
  8. Cara 36 文章