[Bug已解决] torch.compile 加 PyTorch 大于等于 2.12 下梯度非确定性解决方案

[Bug已解决] torch.compile 在 PyTorch 2.12 下产生非确定性梯度Non-deterministic gradients解决方案一、现象长什么样你开了torch.compile做加速并且期望可复现同样输入、同样种子跑两次梯度一致。但升级到PyTorch 2.12后发现import torch def run(): torch.manual_seed(0) model MyModel() x torch.randn(8, 16) out model(x).sum() out.backward() return model[0].weight.grad.clone() g torch.compile(run) a g() b g() print(torch.allclose(a, b)) # 期望 True实际可能是 False两次跑出来的梯度不一致non-deterministic而升级前 2.12是一致的。即 pytorch/pytorch#188890torch.compile在 PyTorch 2.12 下产生非确定性梯度。 本文聚焦为什么编译后会梯度不确定、2.12 改了什么导致、怎么恢复确定性设确定性算法 / 关某些优化 / 升级修复。二、背景确定性训练与 torch.compile确定性determinism相同输入 相同随机种子两次前向/反向得到完全相同的数值。训练可复现的基础。 很多 GPU 算子本身非确定尤其用到原子加、并行归约顺序不固定atomicAdd类归约如 embedding 累加、某些 backward使用 Tensor Core 的fp16/bf16矩阵乘不同算法路径结果末位不同并发 kernel 的归约顺序随调度变化。 PyTorch 提供开关强制确定性torch.use_deterministic_algorithms(True) # 强制确定性非确定算子会报错或选确定路径 torch.backends.cuda.deterministic True # 让 cuDNN 选确定性算法 torch.backends.cudnn.deterministic True torch.manual_seed(0) # 固定种子torch.compile的问题它生成的 Inductor 内核可能用了非确定的归约/原子操作或在某些优化如 autotuning 选了不同 kernel 变体下两次运行走了不同代码路径 → 梯度不同。三、为什么 PyTorch 2.12 会引入非确定性2.12 引入/改动了某些 Inductor 行为导致默认编译产物不再保证确定性新的归约代码生成Inductor 的某个 PR 改变了sum/mean/softmaxbackward 的归约内核生成方式用了「基于原子加的非确定归约」而非「确定性的多级归约」autotuning 选择不稳定编译期选中的 kernel 变体在不同进程/不同输入下不同导致数值路径不同Triton 内核的并行归约顺序升级后 Triton 生成的归约 kernel 并行度更高原子加顺序不固定 → 末位不确定默认关闭了确定性保护某些原本默认确定性的内核在 2.12 改为「性能优先」默认非确定。 本质2.12 的 Inductor 默认行为更偏向性能牺牲了「编译产物确定性」这一隐式保证导致之前可复现的训练变得不可复现。四、最小可运行复现带守卫结构示意下面演示「期望确定性但得到非确定」的检查结构实际触发需 GPU 2.12import torch def set_deterministic(): torch.manual_seed(0) torch.use_deterministic_algorithms(True) torch.backends.cuda.deterministic True def demo_grad_determinism(): if not torch.cuda.is_available(): print([skip] 无 GPU仅说明确定性检查) return import torch.nn as nn model nn.Linear(16, 8).cuda() x torch.randn(8, 16, devicecuda) def step(): model.zero_grad() out model(x).sum() out.backward() return model.weight.grad.clone() set_deterministic() s1 step() s2 step() print(两次梯度一致:, torch.allclose(s1, s2)) # 若开启 torch.compile 后不一致则命中共性问题 compiled_step torch.compile(step) c1 compiled_step() c2 compiled_step() print(compile 后两次一致:, torch.allclose(c1, c2)) if __name__ __main__: demo_grad_determinism()要点torch.compile默认不保证确定性2.12 更明显。下面给恢复确定性的做法。五、解决方案一开启确定性算法开关最关键显式强制确定性让 Inductor 选确定路径而非原子加非确定路径import torch # 强制确定性 torch.use_deterministic_algorithms(True) torch.backends.cuda.deterministic True torch.backends.cudnn.deterministic True # 允许不确定的算子报错而非静默非确定便于发现 # torch.use_deterministic_algorithms(True, warn_onlyFalse) torch.manual_seed(0) model MyModel().cuda() compiled torch.compile(model)torch.use_deterministic_algorithms(True)会让 PyTorch拒绝/替换非确定的算子编译产物也会走确定归约恢复一致性。六、解决方案二关掉引入非确定性的 Inductor 优化若开启确定性开关后仍不稳某些 Inductor 优化绕过该开关关掉相关优化import torch._inductor.config as cfg # 关闭可能产生非确定归约的优化 cfg.triton.use_block_ptr False # 某些 block_ptr 归约非确定 cfg.coordinate_descent_tuning False # autotuning 选择不稳定 cfg.max_autotune False # 见 064 节autotune 选变体不稳 # 或编译时关闭可能非确定的选项 compiled torch.compile(model, options{triton.use_block_ptr: False})这些开关减少「不同运行走不同 kernel 变体」的可能提升确定性。七、解决方案三固定 autotuning 结果本地缓存Inductor 的 autotuning 在不同机器/进程可能选不同变体。用本地 autotuning 缓存固定选择import torch._inductor.config as cfg cfg.autotune_local_cache True # 把 autotuning 结果存到本地复用同一变体 cfg.autotune_remote_cache False # 关远程缓存避免取到不同变体 # 第一次跑会选并缓存之后复用同一选择数值稳定 compiled torch.compile(model, modemax-autotune)固定变体后两次运行走完全相同的内核梯度一致。八、解决方案四升级到修复版本 / 降级#188890 是 2.12 的回归官方会修让torch.compile默认尊重确定性开关或恢复确定的归约生成。处理# 方案 A升级到已修复的 nightly / 新版 pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu124 # 方案 B若修复未出且必须确定性临时降级到 2.12 pip install torch2.11.0 --index-url https://download.pytorch.org/whl/cu124判断修复同样代码开启use_deterministic_algorithms(True)后torch.compile两次梯度allclose为 True。九、排查清单升级 2.12 后torch.compile梯度不可复现 → 确认 #188890 类回归。开确定性开关torch.use_deterministic_algorithms(True)backends.cuda.deterministicTrue。关非确定优化triton.use_block_ptrFalse/coordinate_descent_tuningFalse/max_autotuneFalse。固定 autotuning开autotune_local_cache复用同一 kernel 变体。升级到修复版修复未出则临时降级 2.12。验证torch.allclose(grad_run1, grad_run2)必须 True 才算确定性恢复。十、小结Non-deterministic gradients when using torch.compile and PyTorch 2.12#188890的本质是PyTorch 2.12 的 Inductor 默认行为更偏向性能改动/新增了某些归约内核生成基于原子加的非确定归约与 autotuning 变体选择导致torch.compile的编译产物不再隐式保证确定性两次运行走不同数值路径梯度不一致。 应对开确定性开关torch.use_deterministic_algorithms(True)torch.backends.cuda.deterministicTrue强制走确定归约关非确定优化triton.use_block_ptrFalse/ 关 autotuning减少变体不稳定固定 autotuning开autotune_local_cache复用同一内核变体升级/降级升级到修复版或临时回退 2.12 保确定性。 记住torch.compile默认不保证确定性2.12 更明显。要可复现训练必须显式开use_deterministic_algorithms(True)并对 Inductor 的非确定优化保持警惕。