[Bug已解决] linear 加 cross_entropy 在 CUDA float16 下重分发失效 DISABLED test 绕过方案解决方案

[Bug已解决] nn.functional.linear cross_entropy 在 CUDA float16 下 torch_function 重分发失效DISABLED test绕过方案解决方案一、现象长什么样你在 CUDA 上用float16半精度跑分类组合了F.linear和F.cross_entropyimport torch import torch.nn.functional as F logits F.linear(x, W).half().cuda() # float16 loss F.cross_entropy(logits, labels) # 反向/重分发出问题如果你在使用__torch_function__重载机制比如自定义 Tensor 子类、或某些库拦截torch_function遇到一个被官方DISABLED的测试DISABLED test_redispatch_nn_functional_linear_cross_entropy_cuda_float16 (__main__.TestTorchFunctionRedispatchOpsCUDA)即 pytorch/pytorch#186203。含义nn.functional.linear与nn.functional.cross_entropy的组合在 CUDA float16 下当走「torch_function 重分发redispatch」路径时行为不正确 / 崩溃测试被禁用。 本文聚焦torch_function 重分发是什么、为什么 float16 下的 linearcross_entropy 组合会坏、作为用户怎么绕开。二、背景torch_function与重分发PyTorch 的__torch_function__协议让自定义 Tensor如torch.Tensor子类、或FakeTensor、MetaTensor、量化张量能拦截对torch.xxx/F.xxx的调用先做自己的处理再**重分发redispatch**回 PyTorch 默认实现。 典型重分发流程class MyTensor(torch.Tensor): classmethod def __torch_function__(cls, func, types, args(), kwargsNone): # 先拦截 print(调用了:, func) # 再重分发回默认实现 return super().__torch_function__(func, types, args, kwargs)当你的自定义 Tensor 调F.linear时会进入__torch_function__再 redispatch 到F.linear的真实实现。问题出在某些算子组合在「redispatch」路径上对 dtype尤其 float16/ 设备CUDA的处理有缺口。三、为什么 float16 linear cross_entropy 重分发会坏F.linear与F.cross_entropy在 CUDA float16 下内部会做dtype 提升 / 特殊内核选择F.linear在 float16 下可能调用 Tensor Core 内核或对权重做某种转换F.cross_entropy在 float16 下为了数值稳定内部常把 logits提升promote到 float32算 softmax / log再降回 float16避免半精度下 softmax 溢出。 当这两者被包在__torch_function__重分发里时缺陷在于redispatch 时 dtype 信息丢失 / 被错误转换自定义 Tensor 的 redispatch 可能把 float16 当成 float32 传进去或反过来导致内核不匹配cross_entropy 的「内部 promote 到 float32」在重分发路径下没生效于是 float16 下直接算 softmax数值溢出 / NaN或触发不支持的内核组合顺序的重分发边界F.linear的输出作为F.cross_entropy的输入两次 redispatch 之间状态没对齐CUDA float16 内核拿到错误布局。 本质torch_function重分发对「多算子组合 float16 CUDA」的 dtype/设备传播支持不完整测试被禁用。四、最小可运行复现带守卫下面演示「linear cross_entropy 在 CUDA float16」的正常用法以及「自定义 Tensor 重分发」的陷阱结构import torch import torch.nn.functional as F def demo_fp16_linear_ce(): if not torch.cuda.is_available(): print([skip] 无 GPU仅说明 fp16 用法) return x torch.randn(4, 16, devicecuda, dtypetorch.float16) W torch.randn(8, 16, devicecuda, dtypetorch.float16) labels torch.randint(0, 8, (4,), devicecuda) logits F.linear(x, W) loss F.cross_entropy(logits, labels) # 内部会 promote 到 fp32 算 print(fp16 linearCE loss:, loss.dtype, loss.item()) def demo_custom_tensor_redispatch(): # 自定义 Tensor 走 __torch_function__ 重分发float16 组合易踩坑 class MyTensor(torch.Tensor): classmethod def __torch_function__(cls, func, types, args(), kwargsNone): kwargs kwargs or {} # 注意redispatch 时 dtype 信息要原样传递否则 float16 组合出错 return super().__torch_function__(func, types, args, kwargs) # 演示用实际触发依赖内部 API print(自定义 Tensor 重分发需谨慎处理 dtype/设备传播)要点裸的 float16 linearCE 是稳定的PyTorch 内部会正确 promotebug 只在「被 torch_function 重分发包裹」时。五、解决方案一避免在 float16 关键路径用 torch_function 拦截最直接绕过别在F.linearF.cross_entropy这种「float16 敏感组合」上挂__torch_function__拦截或让自定义 Tensor 在 redispatch 时原样传递 dtype/设备class MyTensor(torch.Tensor): classmethod def __torch_function__(cls, func, types, args(), kwargsNone): kwargs kwargs or {} # 关键不要修改 args/kwargs 的 dtype 与设备原样 redispatch return super().__torch_function__(func, types, args, kwargs)即redispatch 时保持 float16 与 CUDA 设备不变让 PyTorch 内部自己做 promote不要你在拦截层做 dtype 转换那正是触发缺口的地方。六、解决方案二在调用前显式 promote 到 float32如果你控制调用代码可主动规避 float16 下的内部 promote 缺口在cross_entropy前把 logits 提到 float32loss 算完再按需降回import torch import torch.nn.functional as F def safe_fp16_linear_ce(x, W, labels): logits F.linear(x, W).float() # 主动 promote 到 float32 loss F.cross_entropy(logits, labels) # float32 下稳定 return loss # 需要半精度结果再 .half() if torch.cuda.is_available(): x torch.randn(4, 16, devicecuda, dtypetorch.float16) W torch.randn(8, 16, devicecuda, dtypetorch.float16) labels torch.randint(0, 8, (4,), devicecuda) loss safe_fp16_linear_ce(x, W, labels) print(主动 promote 后 loss:, loss.item())原理绕开「float16 下 cross_entropy 内部 promote 在重分发路径失效」的缺口——你在外面就 promote 了内部不再需要做这个可能坏的转换。七、解决方案三用 AMP 自动混合精度标准做法更标准的做法是用torch.amp自动混合精度它正确处理 dtype 传播不依赖手动 redispatchimport torch import torch.nn as nn import torch.nn.functional as F model nn.Linear(16, 8).cuda() x torch.randn(4, 16, devicecuda) labels torch.randint(0, 8, (4,), devicecuda) scaler torch.amp.GradScaler() with torch.amp.autocast(cuda, dtypetorch.float16): logits model(x) # autocast 自动管理 dtype loss F.cross_entropy(logits, labels) scaler.scale(loss).backward() scaler.step(torch.optim.SGD(model.parameters(), lr1e-3)) scaler.update()autocast会按算子黑白名单自动决定「linear 用 float16、cross_entropy 内部用 float32」且这套 dtype 管理在官方路径上经过验证不踩 torch_function 重分发的未修复缺口。八、解决方案四升级并关注该测试启用#186203 是内部测试被禁用官方会修 redispatch 的 dtype 传播。升级pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu124判断修复该 DISABLED 测试被启用且通过即代表 float16 下 linearCE 的 torch_function 重分发修好。修复前用「原样 redispatch / 主动 promote / autocast」绕过。九、排查清单自定义 TensorFakeTensor/量化/MetaTensor float16 下 linearCE 数值错/崩 → 确认 #186203。裸 float16 linearCE 是稳的崩在「被 torch_function 重分发包裹」→ 基本可定位。绕开redispatch 时原样保持 float16 CUDA 设备不在拦截层做 dtype 转换。主动 promotecross_entropy 前.float()绕开内部 promote 缺口。用 AMPtorch.amp.autocast标准混合精度dtype 管理经过验证。升级关注该 DISABLED 测试启用即代表修复。十、小结DISABLED test_redispatch_nn_functional_linear_cross_entropy_cuda_float16#186203的本质是nn.functional.linearnn.functional.cross_entropy在 CUDA float16 下当调用被__torch_function__机制拦截并「重分发redispatch」时dtype尤其 float16→float32 的 promote/ 设备传播在重分发路径上有缺口导致数值错或崩溃测试被禁用。 应对原样 redispatch自定义 Tensor 拦截层不要改 dtype/设备让 PyTorch 内部自己 promote主动 promotecross_entropy前.float()绕开内部 promote 的缺口用 AMPtorch.amp.autocast标准混合精度dtype 管理经官方验证不踩重分发坑等升级关注该 DISABLED 测试启用即代表修复。 记住裸的 float16 linearCE 是稳定的PyTorch 内部会正确 promote坑只在「被 torch_function 重分发包裹」时——要么别拦这两个算子的组合要么在外部自己做好 promote。