PyTorch训练时loss突然变NaN，换TensorFlow就没事，是我哪里写错了吗？

这种“同模型同数据换框架就爆”的情况我也遇到过，大概率不是autograd本身有bug，而是两边默认行为不一样。PyTorch的ViT里如果用了LayerNorm或者softmax，某些实现会在half精度或极端值下更容易溢出，但你说了没用AMP，那就要看是不是梯度裁剪没加、或者优化器eps默认值不同。TensorFlow的Adam默认eps是1e-7，PyTorch也是1e-8，差一个量级在ViT这种深网络里到第三epoch完全可能积累成NaN。另外检查一下PyTorch这边有没有用torch.nn.functional.scaled_dot_product_attention，它的flash实现和TF的attention在mask处理上边界行为不同，pad位置如果没完全mask掉会慢慢把梯度推爆。建议你先把lr降到5e-5跑一遍，如果还炸就逐层打印grad norm，看是哪一层先出问题。也可以试试在PyTorch里开torch.autograd.set_detect_anomaly(True)，虽然慢但能定位到具体op。别太怀疑自己写错了，框架间的数值差异有时候就是这么玄学。

ViT在PyTorch里第3个epoch炸成NaN，这个时间点其实挺有信息量的，说明不是一开始就崩，而是训练到某个临界点后数值开始失控。你提到换TensorFlow就没事，我第一反应是两边默认的epsilon或者reduce方式可能不一样，比如LayerNorm里PyTorch默认eps是1e-5，TF的Keras实现有时候是1e-6或者反过来，这种小差别在浅层看不出，深了以后梯度累积就会放大。另外ViT里attention的softmax如果logits稍微大一点，PyTorch的autograd在反向时对极端值的处理确实可能比TF更激进，尤其是你用了learnable temperature或者没做qk norm的话。建议你先在PyTorch里把torch.autograd.set_detect_anomaly(True)打开，看看到底是哪个op先出的inf，大概率能定位到attention或者FFN里的某个矩阵乘。还有一个偏门但常见的坑：PyTorch的AdamW实现里weight decay和lr的耦合方式跟TF的AdamW不一样，如果你复现的代码里optimizer超参没对齐，前期看不出来，几个epoch后权重范数涨上去就炸了。可以试试把lr再降一个数量级，或者加grad clip到1.0，如果这样就不NaN了，那基本就是优化器层面的问题而不是模型结构。最后，PyTorch和TF在默认dtype的累加精度上也有区别，特别是你如果用了bf16或者fp16的amp但没开loss scaling，不过你说没用混合精度，那可能还是得从optimizer和norm层查起。

这个坑我也踩过，而且不止一次。PyTorch和TensorFlow在默认数值行为上确实有差异，最典型的就是softmax、layernorm这类op的epsilon默认值不同，TF通常给得更大方一些，PyTorch相对激进，ViT里attention的scale又小，几个epoch累积下来很容易在某次前向里直接溢出。你可以先试试在LayerNorm和softmax里手动把eps调到1e-5甚至1e-4，看NaN还出不出现。

另外别忽略优化器状态，Adam系的exp_avg_sq在前期如果遇到个别大梯度，会把更新量推得特别猛，PyTorch这边默认没有做梯度裁剪，TF的Keras有时在fit里会隐式帮你兜一下。建议加一句clip_grad_norm_(model.parameters(), 1.0)，再打开torch.autograd.set_detect_anomaly(True)跑一小段，它会直接告诉你哪个op先出的inf。

还有个容易被忽略的点是初始化，同样的结构在不同框架里默认初始化方案可能不一样，ViT对patch embedding和cls token的初始化挺敏感的，PyTorch这边如果照搬别人的代码但漏了trunc_normal_，前期数值就会飘。你可以对比一下两边第一层权重的std，差距大的话基本就能定位了。

我一般遇到这种情况会先把lr降到1e-5跑两个epoch确认不是数据问题，然后逐层打印activation的max值，通常能在变NaN前一个step就看出是哪一层开始不对劲。TF没事不代表PyTorch写错了，更多是两边默认行为不同，调一调就能对上。

ViT里attention那块的softmax在PyTorch默认是fp32算的，但如果你中间某些中间变量被广播成fp16或者bfloat16就可能溢出，TF那边可能默认帮你做了clamp。建议先开torch.autograd.set_detect_anomaly(True)跑一遍，看是哪个op先出的问题。另外检查下LayerNorm的eps，PyTorch默认1e-5，TF有时是1e-3，小eps在深层网络里容易放大数值。我之前也遇到过类似情况，最后发现是学习率warmup没加，第三个epoch刚好进入高lr阶段直接炸了。

ViT里LayerNorm的eps设了多少？PyTorch默认1e-5，TF可能是1e-3，这差距在深层很容易炸。

我之前也遇到过几乎一样的情况，最后发现是某个自定义norm层在PyTorch里除方差时没加eps，前期没事，跑到几个epoch后数值一抖就炸了。TF那边默认eps给得比较宽松，所以看起来“没事”，其实不代表模型没问题。建议你把每个epoch的梯度范数打出来，定位到具体是哪一层先爆的，大概率不是autograd的锅。

这种情况其实挺常见的，未必是你写错了什么。PyTorch和TensorFlow在底层算子实现上确实有差异，尤其是LayerNorm、Softmax这类涉及除法和指数的op，eps的默认值或者计算顺序稍有不同，数值稳定性就会差很多。ViT里attention那块QK^T除以sqrt(d)之后如果数值偏大，softmax很容易溢出，PyTorch可能在某些版本里没有做足够的保护。你可以试试在LayerNorm里把eps调大一点，比如从1e-5改成1e-6甚至1e-4看看有没有改善。另外检查一下有没有哪些地方用了in-place操作，autograd在in-place场景下偶尔会给出奇怪的梯度。还有个偏方是开torch.autograd.set_detect_anomaly(True)跑一遍，虽然慢但能定位到具体是哪个op先出的问题。TensorFlow没事不代表PyTorch有问题，只是两边的数值路径不一样，换个种子或者换个PyTorch版本可能结果又不同了。