Rust in Everything
基于 burn 0.22.1

Metal 上越界的下标不报错，读出 0

主线位置：训练主线的 forward → loss：Embedding 用 token id 取行，交叉熵用标签取 log 概率，两处都是「按下标读」。

版本：PyTorch 2.15.0（CPU 和 MPS）、burn 0.22.1（Flex 设备和 Metal 设备）
flowchart TB
  id["越界的 id 或标签"] --> dev{"设备"}
  dev -->|"torch CPU"| e1["当场 IndexError"]
  dev -->|"torch MPS"| e2["同步时报错"]
  dev -->|"burn Flex"| e3["当场 panic"]
  dev -->|"burn Metal"| z["读出 0"]
  z --> l["loss 偏小 梯度为 0"]

实验代码：PyTorch 侧 oob.py，burn 侧 crate pitb-exec-gpu（burn = { version = "=0.22.1", features = ["flex", "metal"] }，打开默认 feature，所以 Metal 上有 Fusion；release）。每个程序跑两遍，结果相同。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

词表 4 个词，有一个 token id 写成了 4（多了一个特殊符号、忘了扩词表）；3 分类，有一个标签是 3（标签从 1 开始编号）。这在 PyTorch 里不会走远：

emb = nn.Embedding(4, 3).to(dev)
ids = torch.tensor([[1, 4]], device=dev)
out = emb(ids)
print(f"{dev}: embedding([[1, 4]]) returned")
print(f"{dev}: values {out.cpu().flatten().tolist()}")

logits = torch.tensor([[2.0, 0.5, -1.0], [0.1, 0.2, 0.3]], device=dev)
for t in ([0, 2], [0, 3], [0, 300]):
  loss = F.cross_entropy(logits, torch.tensor(t, device=dev))
  print(f"{dev}: cross_entropy targets {t} returned")
  print(f"{dev}: loss {loss.item()}")

（实验代码里每次调用都包在 try 里，打印异常类型和第一行信息。）

cpu: embedding raised IndexError: index out of range in self
cpu: cross_entropy targets [0, 2] returned
cpu: loss 0.6216270923614502
cpu: cross_entropy targets [0, 3] raised IndexError: Target 3 is out of bounds.
cpu: cross_entropy targets [0, 300] raised IndexError: Target 300 is out of bounds.
mps: embedding([[1, 4]]) returned
mps: embedding raised AcceleratorError: index 4 is out of bounds for dimension with size 4
mps: cross_entropy targets [0, 2] returned
mps: loss 0.6216270327568054
mps: cross_entropy targets [0, 3] returned
mps: cross_entropy targets [0, 3] raised AcceleratorError: Target 3 is out of bounds.
mps: cross_entropy targets [0, 300] returned
mps: cross_entropy targets [0, 300] raised AcceleratorError: Target 300 is out of bounds.

CPU 在调用那一行报。MPS 上调用先返回，读结果（.cpu()、.item()）时报 AcceleratorError，信息里有越界的下标。两种设备都不会让一个越界的下标悄悄算下去。

搬到 burn 的现象

同样的模型和数据，EmbeddingConfig::new(4, 3)，权重第 k 行是 [10k, 10k+1, 10k+2]：

let mut emb = EmbeddingConfig::new(4, 3).init(&device);
let w: Vec<f32> = (0..4).flat_map(|k| (0..3).map(move |c| (10 * k + c) as f32)).collect();
emb.weight = Param::from_tensor(Tensor::<1>::from_floats(w.as_slice(), &device).reshape([4, 3]));
let ids = Tensor::<2, Int>::from_ints([[1, 4]], &device);
let out = emb.forward(ids);

let logits = Tensor::<2>::from_floats([[2.0, 0.5, -1.0], [0.1, 0.2, 0.3]], &device);
let loss_fn = CrossEntropyLossConfig::new().init(&device);
for t in [[0, 2], [0, 3], [0, 300]] {
  let targets = Tensor::<1, Int>::from_ints(t, &device);
  let l = loss_fn.forward(logits.clone(), targets);
  // print l
}

Flex：

thread 'main' (8781369) panicked at .../burn-flex-0.22.1/src/ops/gather_scatter.rs:144:5:
index 4 out of bounds for dimension of size 4
flex: embedding([[1, 4]]) panicked
flex: cross_entropy targets [0, 2] -> loss [0.6216271]

thread 'main' (8781399) panicked at .../burn-flex-0.22.1/src/ops/gather_scatter.rs:144:5:
index 3 out of bounds for dimension of size 3
flex: cross_entropy targets [0, 3] panicked

thread 'main' (8781399) panicked at .../burn-flex-0.22.1/src/ops/gather_scatter.rs:144:5:
index 300 out of bounds for dimension of size 3
flex: cross_entropy targets [0, 300] panicked

Metal：

metal: embedding([[1, 4]]) returned, dims [1, 2, 3], values [10.0, 11.0, 12.0, 0.0, 0.0, 0.0]
metal: cross_entropy targets [0, 2] -> loss [0.62162685]
metal: cross_entropy targets [0, 3] -> loss [0.12065555]
metal: cross_entropy targets [0, 300] -> loss [0.12065555]
  • Flex 和 PyTorch CPU 一样，在调用时就停下来。
  • Metal 上没有任何报错，读结果也没有。越界的那个 token 拿到一行 0，形状完全正确。
  • 越界的标签让 loss 变小了：0.12065555 正好是第一个样本的交叉熵除以 2（PyTorch 算 cross_entropy([[2.0, 0.5, -1.0]], [0]) / 2 是 0.12065566331148148，差在 f32 的最后几位）。第二个样本取到的 log 概率是 0，等于「模型对这个标签有 100% 把握」，它对 loss 的贡献是 0。
  • 标签是 3 还是 300，loss 一样。

这个样本也不会产生梯度。logits 建在 Device::metal(..).autodiff() 上，require_grad() 之后 backward，看第二行的梯度（bin cegrad）：

metal: targets [0, 2] grad row 1 [0.15030481, 0.16611251, -0.31641728]
metal: targets [0, 3] grad row 1 [0.0, 0.0, 0.0]

训练照常进行，数据里标错的样本对模型没有任何作用，也没有任何提示。

读出来的也不保证总是 0。x = [1, 2, 3] 在 Metal 上 select 不同的越界下标（bin neighbor，之前先建了一个 64 个 7.0 的张量）：

metal: x.select(0, [3]) -> [0.0]
metal: x.select(0, [4]) -> [0.0]
metal: x.select(0, [5]) -> [0.0]
metal: x.select(0, [64]) -> [9e-44]
metal: x.select(0, [1000]) -> [0.0]
metal: x.select(0, [1000000]) -> [0.0]

下标 64 读到的是一个非零的极小数，是内存里碰巧有的位模式。读到什么取决于那块显存里有什么。

原因

burn 的 Embedding 没有自己的内核：Embedding::forward 调 embedding，后端的 ModuleOps::embedding 默认实现把下标展平后调 float_select。交叉熵在 forward_default 里用 log_softmax(logits, 1).gather(1, target_indices) 取每个样本标签处的 log 概率。

Metal 走 cubecl 的内核。select_kernel 和 gather_kernel 都声明成 #[cube(launch_unchecked, ...)]，只检查输出位置有没有超出输出张量，不检查读出来的下标值：

#[cube(launch_unchecked, address_type = "dynamic")]
fn gather_kernel<T: Numeric, I: Numeric>(
    // ...
) {
    if !indices.is_in_bounds(ABSOLUTE_POS) {
        terminate!();
    }
    // ...
    offset += usize::cast_from(indices.read(ABSOLUTE_POS)) * input.stride(dim);

    output.write(ABSOLUTE_POS, input[offset]);
}

input[offset] 越界时读到什么，cubecl 没有规定。launch_unchecked 和 launch 的区别、#[cube(launch)] 实际上也不截断越界访问，见 burn 卷的自定义内核越界读不报错。

Flex 在 checked_index 里逐个检查，越界就 panic（上面 gather_scatter.rs:144 的那条信息来自它调用的 index_oob）。所以同一份代码，CPU 上报错、GPU 上不报。

Tensor::gather 的文档写着：

    /// # Warning
    /// Not all backends have runtime bound checks for the indices, so make sure the they are valid.
    /// Otherwise, out of bounds indices could lead to unexpected results instead of panicking.

Tensor::select 的文档没有这一段，Embedding 和 CrossEntropyLoss 的文档也没有提。origin/main（2026-10-28）上 select_kernel、gather_kernel 没有改动，select 的文档仍然没有这条警告。这是上游的问题：越界下标在 GPU 后端上是未定义的读，没有任何检查可以打开。

📝NOTE
冷知识：select 的越界警告在一次搬家里丢了

PR #2462（「Add warning in docstring for indices bound checks」，首次发布于 v0.16.0）给 gather、scatter、select、select_assign 四个方法加了上面那段「# Warning」。十个月后的 PR #3681（「Add NumPy-like take operation with multi-dimensional index support」，首次发布于 v0.19.0）让 select 和 select_assign 支持负数的 dim，把这两个方法从 numeric.rs 挪到 base.rs。diff 里旧的 select 连同它的警告一起删掉，新写的 select 文档只有 Arguments 和 Example，select_assign 则带着警告重新写了一遍。从那以后，Embedding 底下用的 select 就不再提醒「下标要自己保证在范围内」。

改法

在 CPU 上先跑一遍。 同样的数据用 Device::flex() 跑一个 epoch（或者只跑数据管道和 loss），越界的下标当场 panic，信息里有下标和维度大小。

在 GPU 上自己检查。 下标是数据的一部分，在送进 Embedding、select 或 loss 之前检查一次范围：

fn checked_select(x: Tensor<2>, ids: Tensor<1, Int>) -> Result<Tensor<2>, String> {
  let n = x.dims()[0] as i64;
  let lo: i64 = ids.clone().min().into_scalar();
  let hi: i64 = ids.clone().max().into_scalar();
  if lo < 0 || hi >= n {
    return Err(format!("index range [{lo}, {hi}] outside 0..{n}"));
  }
  Ok(x.select(0, ids))
}

flex: checked_select [1, 0] -> Ok [10.0, 11.0, 12.0, 0.0, 1.0, 2.0]
flex: checked_select [1, 2] -> Err index range [1, 2] outside 0..2
metal: checked_select [1, 0] -> Ok [10.0, 11.0, 12.0, 0.0, 1.0, 2.0]
metal: checked_select [1, 2] -> Err index range [1, 2] outside 0..2

into_scalar 要把结果读回主机，等于一次同步，每一步都做会打断 Metal 上的攒批（eager 和延迟）。下标通常在 CPU 上就有（从文件读出来的 token id、标签），更便宜的做法是在 Batcher 里、建张量之前检查原始的 Vec，或者在数据集准备阶段检查一次。

不要用 clamp 把下标夹进范围：那样不报错，越界的 token 变成词表最后一个词，问题照样被藏起来。

怎么检查

  • 词表大小和标签数是模型配置的一部分，在数据管道里断言 max(id) < vocab_size、max(label) < num_classes、min >= 0。
  • 训练前用 Flex 跑一遍数据和 loss；有越界 panic 时信息形如 index 4 out of bounds for dimension of size 4。
  • Metal 上 loss 比预期低、某些样本的梯度全是 0、Embedding 的输出里有整行的 0，都值得回头查下标。

小结

  • PyTorch 的越界下标 CPU 当场报 IndexError，MPS 在读结果时报 AcceleratorError；burn Flex 当场 panic，burn Metal 不报，读出 0 或别的值。
  • Embedding（经 select）和 CrossEntropyLoss（经 gather）都受影响：越界标签的样本 loss 贡献为 0、梯度为 0，loss 看着反而更好。
  • 原因是 cubecl 的 select_kernel / gather_kernel 用 launch_unchecked，不检查下标值；gather 的文档有警告，select 的警告在 #3681 时丢了；origin/main 上没有变。
  • 改法：先在 Flex 上跑一遍；GPU 上在建张量之前检查下标范围，不要用 clamp 掩盖。

💬 关联讨论

发起讨论
加载中...