loss.backward()反向传播卡死挑战

当大家在CPU上做一些小的测试时，发现loss.backward()反向传播卡死，原因是AMD Zen2 CPU + PyTorch ≥1.8 CPU版（oneDNN后端）+ CrossEntropyLoss 在多线程下的原子锁死锁。
具体如下：
PyTorch CPU版后端是oneDNN（以前叫MKL-DNN）
从PyTorch 1.8开始，CPU版默认使用Intel oneDNN（即使你是AMD CPU也会走这个路径）。oneDNN在做CrossEntropyLoss的softmax+log+nll_loss三连计算时，会在forward阶段动态创建一个线程池（默认线程数=物理核数=16）。
Backward阶段又会再创建一个线程池
CrossEntropyLoss的backward实现里，oneDNN也会再动态创建一次线程池（同样默认16线程）。
→ 关键问题：这两个线程池的创建/销毁没有严格互斥保护，在某些CPU（尤其是AMD Zen2/Zen3架构的4800U）上会发生经典的“线程池初始化死锁”（thread pool bootstrapping deadlock）。
AMD CPU + oneDNN的兼容性坑
Intel oneDNN对AMD Zen2（你就是4800U）的zen2调度器支持一直不完美，表现为：
多线程时，内部的atomic锁竞争极端激烈
某些cache line false sharing特别严重
导致两个线程池在初始化/绑定CPU核心时互相死等（deadlock），但不是100%复现，所以你以前“有时能跑，有时卡死”。
这就是为什么：
OMP_NUM_THREADS=1 时 → 只创建一个线程，根本不会触发死锁 → 稳如老狗
OMP_NUM_THREADS=4 时 → 仍然会创建多个线程，触发了oneDNN内部的race condition → 死锁概率90%以上
只设OMP不设MKL也无效 → 因为PyTorch CPU现在关键听oneDNN的，不是老的MKL了

不是代码bug，而是：
AMD Zen2 CPU + PyTorch ≥1.8 CPU版（oneDNN后端）+ CrossEntropyLoss 在多线程下的原子锁死锁。
这个bug Intel官方也知道，但因为影响面小（主要是AMD用户+纯CPU训练），优先级一直很低，至今没彻底修（截至PyTorch 2.5还是存在）。

但是做分类任务不会卡死，原因如下：

最核心的一点：
PyTorch 的 oneDNN 后端只对「空间维度很大」的算子（比如 112x112 以上）才启用特殊优化路径。而 CrossEntropyLoss 的 backward 正好在「输出是 4D 张量且 H,W 很大」时，会走一个专门的 oneDNN kernel，这个 kernel 里面就有那个著名的线程池死锁。
你的 CIFAR-10 模型输出是 (N,10)，根本不满足条件，所以走的是安全的纯 Eigen/Tensor 路径 → 完全不卡。
这就是为什么：

所有分类任务（ImageNet、CIFAR、MNIST）几乎没人报这个死锁
所有分割/检测任务（UNet、DeepLab、FCN）在 AMD CPU 上纯 CPU 训练时，必报“卡在 backward”

解决方法：

OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 NUMEXPR_NUM_THREADS=1 VECLIB_MAXIMUM_THREADS=1 python train.py

同时
torch.softmax(pred, dim=1)
死锁就是这个 softmax 在 CPU 上对 (N, C, 224, 224) 这样的大 4D 张量执行时，同样会走 oneDNN 的多线程路径，所以该死锁的还