logits 里有 NaN，贪心选了别的 token

主线位置：推理主线的「生成」一步。对最后一个位置的 logits 取 argmax。

版本：PyTorch 2.15.0（CPU 和 MPS）、candle 0.11.0（CPU，accelerate；Metal）

flowchart TB
  lg["logits 含 NaN"] -->|PyTorch| pn["返回 NaN 的位置"]
  lg -->|"candle CPU"| c{"NaN 在第 0 位"}
  c -->|是| c0["返回 0"]
  c -->|否| cs["跳过 NaN"]
  lg -->|"candle Metal"| ms["跳过 NaN"]
  lg -->|"取样"| we["取样报错"]
  lg -->|"先查 ne self"| ok["报出 NaN"]

本页的对照：PyTorch 侧是实验代码里的 p08.py，candle 侧是实验 crate pita-t 的 p08_nan（features metal accelerate，release）。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

半精度溢出、错误的掩码值、除以 0 的归一化，都会让 logits 里出现 NaN。两行 logits，NaN 分别在第 1 位和第 0 位：

logits = torch.tensor([[1.0, nan, 2.0, 0.5], [nan, 0.0, 5.0, 0.5]])
logits.argmax(-1); logits.max(-1).values
torch.multinomial(torch.softmax(logits, -1), 1)

cpu logits.argmax(-1) [1, 0] max [nan, nan]
cpu multinomial: RuntimeError probability tensor contains either `inf`, `nan` or element < 0
mps logits.argmax(-1) [1, 0] max [nan, nan]

PyTorch 把 NaN 当作最大值传播：max 是 NaN，argmax 指向 NaN 的位置，CPU 和 MPS 相同。贪心解码会生成那个 token，取样直接报错，问题马上能看见。

搬到 candle 的现象

let logits = Tensor::new(&[[1f32, n, 2., 0.5], [n, 0., 5., 0.5]], &dev)?;
logits.argmax(D::Minus1); logits.max(D::Minus1);
let mut greedy = LogitsProcessor::new(42, None, None);
let mut sample = LogitsProcessor::new(42, Some(0.8), None);
greedy.sample(&logits.i(r)?); sample.sample(&logits.i(r)?);

== cpu
logits.argmax(D::Minus1)           U32 [2] [2, 0]
logits.max(D::Minus1)              F32 [2] [2.0, NaN]
row 0: greedy Ok(2)  sample Err("A weight is negative, too large or not a valid number")
row 1: greedy Ok(0)  sample Err("A weight is negative, too large or not a valid number")
== metal
logits.argmax(D::Minus1)           U32 [2] [2, 2]
logits.max(D::Minus1)              F32 [2] [2.0, 5.0]
row 0: greedy Ok(2)  sample Err("A weight is negative, too large or not a valid number")
row 1: greedy Ok(2)  sample Err("A weight is negative, too large or not a valid number")

同样的 logits，三种结果：PyTorch 选 [1, 0]，candle CPU 选 [2, 0]，candle Metal 选 [2, 2]。贪心解码在 candle 上不报错，选出一个看起来正常的 token，生成继续往下走；Metal 上 max 甚至不返回 NaN，后面的检查也看不到它。取样那条路三边都报错，candle 的报错是 rand 的 WeightedIndex 给的，不提 NaN。

原因

• CPU 的 max、argmax 由 ReduceIndex::fold_impl 实现：val 从第 0 个元素开始，f(val, s) 为真才替换，max 的 f 是 |x, y| x < y。和 NaN 的比较永远是假，所以 NaN 在第 0 位时换不出去（第二行得到 0 和 NaN），不在第 0 位时换不进来（第一行得到 2）。
• Metal 的 max 用 Max 里的 fast::max(a, b)，argmax 用 indexed 的 operator>（lhs.val > rhs.val || ...），归约从最小值开始，NaN 永远比不过别的值，所以两行都跳过 NaN。
• LogitsProcessor 的贪心就是 sample_argmax 里的 logits.argmax(D::Minus1)?，继承上面两种行为。取样路径先 softmax，NaN 让整行概率都是 NaN，WeightedIndex::new 拒绝；这条报错的另一个来源（温度为 0）见源码卷的 A weight is negative, too large or not a valid number。

归约与索引的「动脑筋」用同一组数据讲过 max 的这三种结果。origin/main 上 CPU 的 ReduceIndex 和 Metal 的 Max 没有改。

改法

NaN 是上游出了问题的信号，修的是产生 NaN 的地方；这里要做的是让它不被吞掉。在 argmax 之前数一下 NaN：

let bad = logits.ne(&logits)?.to_dtype(DType::U32)?.sum_keepdim(D::Minus1)?;

nan count per row (x.ne(&x))       U32 [2, 1] [1, 1]

CPU 和 Metal 都数出每行一个。x.ne(&x) 只在 NaN 处为 1（IEEE 规定 NaN 不等于自己），结果是 U8，转成 U32 再求和，避开 U8 求和回绕。生成循环里每步读回一个标量，代价是一次很小的归约：

let n_nan = logits.ne(&logits)?.to_dtype(DType::U32)?.sum_all()?.to_scalar::<u32>()?;
if n_nan > 0 {
  anyhow::bail!("{n_nan} NaN in logits at step {step}");
}

只在调试构建里做也可以。和 PyTorch 对齐 logits 时（搬一个 LLM），比较之前也先做这一步。bin p08b 里 b = [1, 2.00001, NaN, 4]，和 a = [1, 2, 3, 4] 比：

(&a - &b)?.abs()?.flatten_all()?.max(0)?.to_scalar::<f32>()?

cpu max|a - b| = 1.001358e-5
metal max|a - b| = 1.001358e-5

NaN 不在第 0 位，CPU 和 Metal 都把它跳过，报出一个看起来「对上了」的差。

怎么检查

• 生成结果在 CPU 和 Metal 上从某一步开始不同，而 logits 的差远小于第一名和第二名的差距，先查 NaN。
• 比较两组输出时，不要只用 max：Metal 上它看不见 NaN，CPU 上只有 NaN 在第 0 位时才看得见。数 ne(self) 的个数，或者读回主机用 f32::is_nan。
• 半精度（半精度）和自己写的掩码（-inf 加 -inf、0 × inf）是 NaN 最常见的来源。

小结

• PyTorch 的 max / argmax 传播 NaN；candle CPU 只在 NaN 位于第 0 位时返回它，Metal 总是跳过它。
• LogitsProcessor 的贪心继承这一点，生成不报错地继续；取样路径报 A weight is negative, too large or not a valid number。
• 改法：argmax 之前用 x.ne(&x) 数 NaN，发现就报错。