Metal 上越界下标不报错
主线位置：推理主线「前向」的第一步。token id 进
Embedding，以及训练主线 loss 里按标签gather。
版本：PyTorch 2.15.0（CPU 和 MPS）、candle 0.11.0（CPU，accelerate；Metal）
flowchart TB id["id 5，词表 4"] -->|"PyTorch CPU MPS"| e1["报错"] id -->|"candle CPU"| e2["invalid index 5"] id -->|"candle Metal"| clamp["夹到第 3 行"] g["gather 类别 3"] -->|"candle Metal"| nb["读到下一行"] clamp --> silent["静默继续"] nb --> silent id -->|"主机上先查"| ok["拒绝"]
本页的对照：PyTorch 侧是实验代码里的 p09.py，candle 侧是实验 crate pita-t 的 p09_oob（features metal accelerate，release）。平台 M2 Max、macOS 27、rustc 1.99.0。
PyTorch 的写法
给分词器加了新 token 却忘了 resize_token_embeddings，新 token 的 id 就超出了词嵌入的行数；标签里混进一个不存在的类别也是同一类问题：
emb = torch.nn.Embedding(4, 2).to(dev) # weight = arange(8).reshape(4, 2)
emb(torch.tensor([1, 5], device=dev))
logp = torch.arange(6, dtype=torch.float32, device=dev).reshape(2, 3)
logp.gather(1, torch.tensor([[3], [0]], device=dev))
cpu emb(ids) IndexError index out of range in self
cpu gather RuntimeError index 3 is out of bounds for dimension 1 with size 3
mps emb(ids) AcceleratorError index 5 is out of bounds for dimension with size 4
mps gather AcceleratorError gather: index 3 is out of bounds for dimension with size 3
PyTorch 在 MPS 上也检查下标，越界抛异常。
搬到 candle 的现象
let emb = Embedding::new(Tensor::arange(0f32, 8., &dev)?.reshape((4, 2))?, 2);
emb.forward(&Tensor::new(&[1u32, 5], &dev)?);
let logp = Tensor::arange(0f32, 6., &dev)?.reshape((2, 3))?;
logp.gather(&Tensor::new(&[[3u32], [0]], &dev)?, 1);
logp.gather(&Tensor::new(&[[0u32], [3]], &dev)?, 1);
== cpu
emb.forward([1, 5]) Err: index-select invalid index 5 with dim size 4
logp.gather([[3],[0]], 1) Err: gather invalid index 3 with dim size 3
logp.gather([[0],[3]], 1) Err: gather invalid index 3 with dim size 3
== metal
emb.forward([1, 5]) F32 [2, 2] [2.0, 3.0, 6.0, 7.0]
logp.gather([[3],[0]], 1) F32 [2, 1] [3.0, 3.0]
logp.gather([[0],[3]], 1) F32 [2, 1] [0.0, 0.0]
Metal 上三行都不报错：
- id 5 取到了第 3 行
[6.0, 7.0]，也就是词表的最后一个 token； - 第 0 行的类别 3 读到了第 1 行的第 0 个元素 3.0；
- 最后一行的类别 3 读过了缓冲区末尾，这次得到 0.0。
这些数和正常的嵌入、正常的 log 概率没有区别，模型会照样跑完。
原因
CPU 后端的 index_select 和 gather 逐个比较下标和维度长度，越界返回 InvalidIndex。Metal 内核里没有可以返回错误的通道，indexing.metal 的 index 把下标夹到最后一个：
const INDEX_TYPENAME input_i = min(input_ids[id_i], (INDEX_TYPENAME)(src_dim_size - 1));
const size_t right_rank_i = tid % right_size;
const size_t left_rank_i = tid / right_size / ids_size;
/*
// Force prevent out of bounds indexing
// since there doesn't seem to be a good way to force crash
// No need to check for zero we're only allowing unsized.
*/
gather 连这一步都没有，直接算 src_i = (left_rank_i * src_dim_size + input_i) * right_size + right_rank_i 去读，越界的下标落到相邻的行。PyTorch 的 MPS 后端在 GPU 上也查出了越界并报回来（上面 AcceleratorError 那两行），candle 的 Metal 后端没有对应的机制。origin/main 上 indexing.metal 没有改。
归约与索引讲过这两个内核的行为；BERT 的位置嵌入超过 512 时在 Metal 上被夹到 511，是同一个机制在真实模型里的表现，见源码卷的 index-select invalid index 512 with dim size 512。负数下标在 Metal 上的情况见 invalid index 18446744073709551516。
改法
下标在上设备之前就是主机上的 Vec<u32>（分词器的输出），在那里检查：
let host_ids = vec![1u32, 5];
match host_ids.iter().find(|&&i| i as usize >= vocab) {
Some(bad) => println!("id {bad} >= vocab {vocab}: refuse"),
None => { /* build the tensor and run */ }
}
let max_id = ids.max(0)?.to_scalar::<u32>()?; // when only the device tensor is at hand
id 5 >= vocab 4: refuse
device-side check: ids.max(0) = 5, ok = false
两个设备输出相同。第二种写法在设备上做一次归约、读回一个数，适合下标是模型自己算出来的情况（比如 argmax 得到的类别再去 gather）。分词器和词表的配套关系，在加载时比一次 tokenizer.get_vocab_size(true) 和 config.vocab_size 就能发现。
怎么检查
- 新模型、新分词器第一次跑，先在 CPU 上跑一遍：CPU 会对越界下标返回
Err，Metal 不会。 - 加过 token 的分词器，确认
get_vocab_size(true)不超过词嵌入的行数（SmolLM2 的vocab_size是 49152）。 - 输出在 CPU 和 Metal 上差得远、而数值又都「像样」，下标越界是要排查的一项。
小结
- PyTorch 在 CPU 和 MPS 上都检查下标；candle 只有 CPU 返回
Err。 - candle Metal 的
index_select把越界下标夹到最后一行，gather读到相邻行或缓冲区之外，都不报错。 - 改法：在主机上检查
id < vocab_size，或在设备上max一次再读回比较。