Rust in Everything
基于 candle 0.11.0

数掩码里的 1：290 个加出 34

主线位置：张量主线的「计算」一步。用比较运算得到掩码，再求和数个数、做整数除法。
版本：PyTorch 2.15.0（CPU 和 MPS）、candle 0.11.0（CPU，accelerate；Metal；debug 和 release 构建）
flowchart TB
  ids["300 个 id"] --> ne["ne pad"]
  ne --> u8["U8 掩码"]
  u8 --> sum["sum_all"]
  sum -->|release| w["U8 34"]
  sum -->|"CPU debug"| p["panic overflow"]
  u8 -->|"先转 U32"| ok["290"]
  cnt["计数含 0"] --> div["整数除法"]
  div -->|CPU| dz["panic 除以 0"]
  div -->|Metal| ff["4294967295"]

本页的对照：PyTorch 侧是实验代码里的 p05.py，candle 侧是实验 crate pita-t 的 p05_intsum（features metal accelerate，debug 和 release）。平台 M2 Max、macOS 27、rustc 1.99.0。panic 行里的本地路径前缀换成了 ...，线程号保留原样。

PyTorch 的写法

一行 300 个 token，最后 10 个是填充（id 0），数真实 token 的个数；再拿总数除以个数：

ids = torch.full((1, 300), 7); ids[0, 290:] = 0
m = ids != 0
m.sum(); m.to(torch.uint8).sum()
tot = torch.tensor([290, 5]); cnt = torch.tensor([2, 0])
tot // cnt

cpu (ids != pad).sum() torch.int64 290 | uint8 sum torch.int64 290
cpu tot // cnt RuntimeError ZeroDivisionError
mps (ids != pad).sum() torch.int64 290 | uint8 sum torch.int64 290
mps tot // cnt [145, 0]

bool 和 uint8 求和都提升成 int64。整数除以 0 在 CPU 上抛异常，在 MPS 上静默给 0。

搬到 candle 的现象

let ids = Tensor::new(ids.as_slice(), &dev)?;   // 290 x 7u32, then 10 x 0
let m = ids.ne(0u32)?;
m.sum_all();
let tot = Tensor::new(&[290u32, 5], &dev)?;
let cnt = Tensor::new(&[2u32, 0], &dev)?;
&tot / &cnt;

求和，release（CPU 和 Metal 相同；Metal 的 debug 构建也是这样）：

mask dtype U8
m.sum_all()                        U8 [] [34]

求和，CPU debug：

mask dtype U8

thread 'main' (9167299) panicked at .../candle-core-0.11.0/src/cpu/kernels.rs:44:13:
attempt to add with overflow

除法，CPU（debug 和 release 相同）：

thread 'main' (9182722) panicked at .../candle-core-0.11.0/src/op.rs:459:28:
attempt to divide by zero

除法，Metal（debug 和 release 相同）：

tot / cnt                          U32 [2] [145, 4294967295]

290 mod 256 = 34。同一段代码，开发时在 CPU debug 下 panic，换成 release 或 Metal 后安静地得到 34。除以 0 的那一行在 CPU 上是 panic，&tot / &cnt 返回的 Result 根本没机会交给调用方，? 和 match 都接不住；Metal 上不报错，给出全 1 的位模式。

原因

  • 比较运算 ne、eq、ge 返回 U8（CPU 的 Cmp 把结果写成 u8::from(..)），candle 没有 bool。
  • sum_all 的结果保持输入的 dtype，没有 PyTorch 那样的「整数求和提升成 int64」。CPU 上的求和就是 vec_reduce_sum 里普通的 Rust 加法 *res += *xs.add(i)：Cargo 的 dev 配置对依赖也打开溢出检查，所以 debug 下 panic、release 下回绕；Metal 内核没有溢出检查，两种构建都回绕。
  • 整数除法是 bin_op!(Div, ..) 里的 v1 / v2。Rust 的整数除以 0 在任何构建下都 panic，这个检查不受溢出检查开关影响。

整数运算在两个设备、两种构建下的完整对照（u32 减法回绕、i64::MIN / -1 等）见整数运算的「溢出」和「除以 0」两节。origin/main 上 op.rs 只改了加法和乘法的向量化实现，Div 仍是 v1 / v2；kernels.rs 的 vec_reduce_sum 没有改。

改法

求和之前转成足够宽的类型；除数里可能有 0，就先把 0 换掉（或者在主机上先检查）：

m.to_dtype(DType::U32)?.sum_all();
cnt.maximum(1u32).and_then(|c| &tot / &c);

fix: m.to_dtype(U32).sum_all()     U32 [] [290]
fix: tot / cnt.maximum(1)          U32 [2] [145, 5]

CPU 和 Metal、debug 和 release 上都是这两行。U32 能数到 42 亿个，对 token 计数足够；要和 PyTorch 的 int64 对齐就转 DType::I64。maximum(1) 让空行的「平均」变成总数本身，这是否合理取决于调用方，关键是把除以 0 从 panic 变成一个能看见的值；要和 PyTorch CPU 一样报错，就在主机上检查 cnt 里有没有 0，自己返回 Err。

怎么检查

  • 凡是 sum_all / sum 的输入来自比较运算，先 to_dtype(DType::U32)。数 token、数正确预测（准确率）都是这种场合。
  • 测试至少跑一次 debug 构建：CPU debug 的 panic 会把回绕暴露出来，release 和 Metal 不会。
  • 整数张量做除数之前问一句「这里会不会是 0」；candle 的 CPU 后端在这里 panic，Metal 给出 0xFFFFFFFF，两种都不是 Err。

小结

  • candle 的掩码是 U8，sum_all 不提升：290 个 1 在 release 和 Metal 上得到 34，CPU debug 下 panic attempt to add with overflow。
  • 整数除以 0：CPU 上 panic attempt to divide by zero（release 也一样，不是 Err），Metal 上得到 4294967295。
  • 改法：求和前 to_dtype(DType::U32)；除数先 maximum(1) 或在主机上检查。

💬 关联讨论

发起讨论
加载中...