Rust in Everything
基于 burn 0.22.1

出错的 op 不报错,后面读别的张量才报

主线位置:算子主线的 Fusion 一层。a.matmul(b) 这样的调用在 Metal 上只登记一条 OperationIr,真正执行在 fusion 服务线程上、读数据或流满的时候(见 Fusion 装饰器)。执行时的错误回不到登记那一行。

flowchart TB
  reg["登记 huge = zeros + 1"] --> ok1["立刻返回 dims 正常"]
  ok1 --> read1["读 small"]
  read1 --> drain["服务线程执行整条流"]
  drain --> fail["huge 的块分配失败"]
  fail --> mark["huge 记成 Errored"]
  read1 --> okr["small 读到 Ok"]
  mark --> read2["读 huge"]
  read2 --> errr["这时才拿到 Err"]

实验 crate pc-gpu,Metal(Device::metal(DeviceKind::DefaultDevice),dispatch<fusion<cubecl<wgpu<msl>>>>),release 构建。M2 Max、macOS 27、rustc 1.99.0。

现象

一个放不进显存的张量([2^18, 2^18] 个 f32,256 GiB),后面跟一个不相干的小计算:

let huge = Tensor::<2>::zeros([1 << 18, 1 << 18], &device) + 1.0;
println!("--- registered, dims {:?}", huge.dims());
let small = Tensor::<1>::from_floats([1.0, 2.0], &device) * 2.0;
println!("{:?}", small.try_to_data().map(|d| d.try_to_vec::<f32>().unwrap()));
println!("{:?}", device.sync());
match huge.slice([0..1, 0..2]).try_to_data() {
  Ok(d) => println!("Ok {:?}", d.try_to_vec::<f32>()),
  Err(e) => println!("Err: {e}"),
}
--- huge = zeros([1 << 18, 1 << 18]) + 1  (256 GiB of f32)
--- registered, dims [262144, 262144]
--- small.try_to_data()
Ok([2.0, 4.0])
--- device.sync()
Ok(())
--- huge.slice([0..1, 0..2]).try_to_data()
Err: An error happened during execution
Caused by:
  Several failures at once
Caused by:
  The bytes were never written (failure #3, still claiming 2 buffer(s))
  skipped `SliceKernel`: it needed memory ManagedMemoryId { value: 4 }, which carried the failure
Caused by:
  An IO error happened
Caused by:
  can't allocate buffer of size: 274877906944
No backtrace available
Asked at:
No backtrace available
  • 建 huge 那一行没有任何反应,dims() 也正常:形状是登记时算的。
  • 读 small 时整条流被执行,huge 的分配就在这时失败,但 small 拿到的是 Ok。
  • device.sync() 也是 Ok(())。
  • 直到读 huge(或由它算出的东西)才看到 can't allocate buffer of size: 274877906944。

换成 into_data(),它在同一位置 panic,位置是读的那一行(into_data 带 #[track_caller],第 216 行就是 huge.slice(..).into_data()),不是建 huge 的那一行:

--- huge.slice([0..1, 0..2]).into_data()

thread 'main' (3534530) panicked at src/main.rs:216:47:
Error while reading data: use `try_into_data` instead to catch the error at runtime: An error happened during execution
Caused by:
  Several failures at once
...
  can't allocate buffer of size: 274877906944

训练里,这一行可能是几十个 op 之后的 loss.into_scalar()。

出处

fusion 服务线程执行一个块失败时,不把错误交给「正好在 drain 这条流的人」,而是把这个块要写的每个张量标成 Handle::Errored(TensorError);后面依赖它的块被跳过,输出同样标错,并记下离源头隔了几层。读张量时 read_float 先取出这个错误,变成 ExecutionError。机制和设计理由在 Fusion 装饰器的「错误挂在张量上」一节。

sync 为什么是 Ok:Fusion 的 sync 经 try_sync 把当前流 drain 掉,再调内层后端的 sync;它的 Err 只表示「服务线程没能跑这个任务」(panic 了或者没了)。块失败已经被写进张量,drain 本身是成功的,内层的 cubecl 也没有待完成的工作出错。Device::sync 的文档写的是「Returns an [ExecutionError] if an operation failed to execute」,0.22.1 在 Metal + fusion 上实测两种失败都返回 Ok(()),和文档对不上。

同一种延迟也出现在跨设备的 cat:Device::default()(wgpu Auto)和 Device::metal(..) 是两个 Cube 设备,TensorCheck::cat 不查设备(见检查与报错,这个坑本身见 input tensor is on the wrong backend),错误在执行时才出现。这里关心的是错误什么时候、在哪个张量上出现:

metal == default: false
--- bad = cat([a, b]); worse = bad + 1; good = a * 10
--- good.try_to_data()

thread 'DSU-100-16384' (3528032) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-ir-0.22.1/src/handle.rs:333:32:
Should have handle for tensor TensorId { value: 3 }
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
Ok([10.0, 20.0])
--- bad.try_to_data()
Err: An error happened during execution
Caused by:
  Should have handle for tensor TensorId { value: 3 }
--- worse.try_to_data()
Err: An error happened during execution
Caused by:
  Should have handle for tensor TensorId { value: 3 }

这一种在服务线程 DSU-100-16384 上 panic,panic 信息会打到 stderr,算是唯一的「当场」线索;分配失败那种连这一行都没有。

修法

找出错的那一步,靠两件事:读得早,读得近。

调试时在可疑的 op 后面立刻读它的结果(try_to_data 或 try_into_scalar),错误就出现在离源头最近的地方。光 sync 不够,上面已经看到它返回 Ok。在 cat 之后插一次 metal.sync(),结果是:

--- bad = cat([a, b]); worse = bad + 1; good = a * 10

thread 'DSU-100-16384' (3540600) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-ir-0.22.1/src/handle.rs:333:32:
Should have handle for tensor TensorId { value: 3 }
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
--- metal.sync() right after cat: Ok(())
--- good.try_to_data()
Ok([10.0, 20.0])
--- bad.try_to_data()
Err: An error happened during execution
Caused by:
  Should have handle for tensor TensorId { value: 3 }
--- worse.try_to_data()
Err: An error happened during execution
Caused by:
  the work writing it was skipped 1 operation(s) below a failure: Should have handle for tensor TensorId { value: 3 }

sync 把 cat 单独执行掉了(panic 出现在 sync 返回之前),返回值仍是 Ok(())。这一次 worse 在另一个块里,报错写成「skipped 1 operation(s) below a failure」:错误信息里的层数(TensorError 的 Display)说明它离出错的那一步隔了几个 op,层数是 0(没有 skipped 字样)的那个张量就是出错的 op 写的。

至于两个具体原因:跨设备时先 to_device 到同一设备再 cat;

cat after to_device: [2.0, 3.0, 4.0, 5.0]

显存不够则要减小张量或批大小,这一点和延迟无关,只是它也要等到读的时候才知道。

怎么确认没中招

  • 生产代码读结果用 try_into_data / try_to_data,把 Err 打出来,「Caused by」链的最后一行才是根因。
  • 不要把 device.sync() 的 Ok 当成「前面都成功了」。
  • 看 stderr 里有没有 DSU- 开头的线程 panic,那是 fusion 服务线程。

小结

  • GPU 上的 op 登记时不执行,执行失败的错误挂在它写出的张量上,读到它们才交出来。
  • 同一时刻读别的张量照常成功;0.22.1 的 device.sync() 在两种失败下都返回 Ok(()),和文档不符。
  • 错误信息里「skipped N operation(s)」给出离源头的距离;调试时在可疑 op 后立刻 try_to_data。

💬 关联讨论

发起讨论
加载中...