matmul: dtype mismatch
主线位置：部署主线的「加载」和「推理」之间。权重文件来自 PyTorch，加载这一步不报错，问题留到第一次 forward。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB file["F64 BF16 文件"] --> pt["load_state_dict"] pt -->|"拷进已有参数"| ptok["参数仍是 F32"] file --> bs["burn-store load_from"] bs -->|"保留文件 dtype"| p64["参数变成 F64"] p64 --> mm["F32 输入 matmul"] mm --> boom["panic dtype mismatch"] bs -->|"FloatCastAdapter"| ok["参数 F32"]
实验代码：PyTorch 侧 pit_store.py（dtype 模式），burn 侧 crate pitb-store-io 的 bin pit（default-features = false，feature std、flex、autodiff、optim、store、safetensors、pytorch，debug 和 release 各跑一次）。模型是主线的两层 MLP。平台 M2 Max、macOS 27、rustc 1.99.0。
PyTorch 的写法
float64 的权重文件比想象的常见：权重是 numpy 算出来的（numpy 默认 float64，torch.from_numpy 原样保留），或者训练时调过 model.double()。半精度文件则是为了省一半体积存的。PyTorch 里这都不是问题，load_state_dict 把值拷进模型已有的参数，dtype 跟着参数走：
w = np.sin(np.arange(1, HID * IN + 1)).reshape(HID, IN) * 0.5
p = nn.Parameter(torch.from_numpy(w))
print("np default dtype:", w.dtype, "| nn.Parameter(torch.from_numpy(w)).dtype:", p.dtype)
save_file(a1_model().double().state_dict(), f"{out}/a1_f64.safetensors")
save_file(a1_model().half().state_dict(), f"{out}/a1_f16.safetensors")
save_file(a1_model().bfloat16().state_dict(), f"{out}/a1_bf16.safetensors")
for name in ["a1_f64", "a1_f16", "a1_bf16"]:
path = f"{out}/{name}.safetensors"
with safe_open(path, "pt") as f:
print(name, {k: f.get_slice(k).get_dtype() for k in f.keys()})
m = Mlp()
print(name, m.load_state_dict(load_file(path)), "->", m.fc1.weight.dtype)
np default dtype: float64 | nn.Parameter(torch.from_numpy(w)).dtype: torch.float64
a1_f64 {'fc1.bias': 'F64', 'fc1.weight': 'F64', 'fc2.bias': 'F64', 'fc2.weight': 'F64'}
a1_f64 <All keys matched successfully> -> torch.float32
a1_f16 {'fc1.bias': 'F16', 'fc1.weight': 'F16', 'fc2.bias': 'F16', 'fc2.weight': 'F16'}
a1_f16 <All keys matched successfully> -> torch.float32
a1_bf16 {'fc1.bias': 'BF16', 'fc1.weight': 'BF16', 'fc2.bias': 'BF16', 'fc2.weight': 'BF16'}
a1_bf16 <All keys matched successfully> -> torch.float32
三个文件读完，参数都是 float32。PyTorch 也不做 matmul 的 dtype 提升，只是这个习惯让人碰不到：真把一个 float64 的模型喂 float32 的输入，同样报错。
f64 model, f32 input: RuntimeError: mat1 and mat2 must have the same dtype, but got Float and Double
搬到 burn 的现象
burn 侧按导入 PyTorch 权重的写法读同一个文件，加载后用一个访问者列出每个参数的 dtype（和参数与遍历里的 Named 一样，只是收集 param.val().dtype()），再做一次 forward：
let mut model = Mlp::new(&device);
let mut st = SafetensorsStore::from_file(&path).with_from_adapter(PyTorchToBurnAdapter);
match model.load_from(&mut st) {
Ok(r) => println!("{name}: load Ok, errors = {}", r.errors.len()),
Err(e) => println!("{name}: load Err\n{e}"),
}
let mut v = DTypes::default();
model.visit(&mut v);
println!("{name}: param dtypes {:?}", v.found);
let logits = model.forward(a1_input_f64(&device));
a1_f64: load Ok, errors = 0
a1_f64: param dtypes [("fc1.weight", F64), ("fc1.bias", F64), ("fc2.weight", F64), ("fc2.bias", F64)]
a1_f64: non-F32 params = 4
thread 'main' (8775600) panicked at .../burn-flex-0.22.1/src/ops/matmul.rs:84:5:
assertion `left == right` failed: matmul: dtype mismatch
left: F32
right: F64
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
半精度的两个文件一样，只是 right 换了：
a1_f16: param dtypes [("fc1.weight", F16), ("fc1.bias", F16), ("fc2.weight", F16), ("fc2.bias", F16)]
...
assertion `left == right` failed: matmul: dtype mismatch
left: F32
right: F16
...
a1_bf16: param dtypes [("fc1.weight", BF16), ("fc1.bias", BF16), ("fc2.weight", BF16), ("fc2.bias", BF16)]
...
assertion `left == right` failed: matmul: dtype mismatch
left: F32
right: BF16
（... 处是和 F64 相同的行。）release 构建读 F64 文件，输出除了线程号逐行相同，在同一个位置 panic：这不是只在 debug 里查的断言。加载报告 errors = 0，ApplyResult 里看不出任何异常；panic 出在后端的 matmul 里，不在加载那一行。
原因
两处行为叠在一起。
第一处是 burn-store 加载浮点参数时保留文件里的 dtype。TargetDType 上面的注释写明了这是设计：
/// A loaded float or int parameter keeps whatever dtype the source declared. For floats that
/// is what makes a half-precision save loadable into a module written for full precision
/// (`dtype_preservation_f64` below pins it); for ints it is what keeps a PyTorch checkpoint's
/// int64 buffers at int64 rather than narrowing them to the backend's element (pinned by
/// `pytorch-tests`' `integer_full_precision`, added deliberately in #4854).
它只检查「浮点进浮点、整数进整数」，所以 F64、F16、BF16 进 F32 的参数都算合法，加载成功。这个行为是怎么来的（之前 F16 文件会被读成 F32），见导入 PyTorch 权重末尾的冷知识。
第二处是 burn 没有 dtype 提升。输入是设备默认的 F32，参数是 F64，Flex 的 matmul 第一行就是一个 assert_eq!：
pub fn matmul(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor {
assert_eq!(lhs.dtype(), rhs.dtype(), "matmul: dtype mismatch");
assert_eq! 在 release 里也在，所以两种构建都 panic。dtype 不一致时别的运算会怎样，见 dtype 与设备和没有 dtype 提升。origin/main 上这两处都没有变：applier.rs 只换了内部的张量类型，测试仍然断言「should preserve F64 dtype」；matmul 的断言原样在第 84 行。
问： 既然 FloatCastAdapter 一行就能转，burn-store 为什么不像 load_state_dict 那样默认转成模型的 dtype？
答： 上面那段注释给了两个理由：保留文件的 dtype，半精度存的模型才能读进一个按全精度写的模块；PyTorch checkpoint 里的 int64 缓冲区也才不会被收窄成后端默认的整数类型（Flex 上是 i32）。默认转换会让后一种情况静默丢精度，而 dtype 不符至少会在 forward 时 panic。想要 load_state_dict 的语义，就自己显式串上转换，见下一节。
改法
在适配器链上串一个 FloatCastAdapter，把所有浮点参数转成模型要的 dtype：
let mut st = SafetensorsStore::from_file(&path)
.with_from_adapter(PyTorchToBurnAdapter.chain(FloatCastAdapter::to(DType::F32)));
a1_f64: load Ok, errors = 0
a1_f64: param dtypes [("fc1.weight", F32), ("fc1.bias", F32), ("fc2.weight", F32), ("fc2.bias", F32)]
a1_f64: non-F32 params = 0
a1_f64: logits dtype F32
和 PyTorch 读同一组权重的前向比：
logits n= 12 max|diff|=1.49e-08
worst=1.49e-08 tol=1e-06 OK
这些 f64 值本来就是从 f32 扩上去的，转回来没有损失，1.49e-8 是两边 matmul 求和顺序不同留下的（见怎么读这一卷）。BF16 文件也这样转成 F32，先和原来的 f32 模型的 PyTorch 输出比：
logits n= 12 max|diff|=1.10e-03
worst=1.10e-03 tol=1e-06 MISMATCH
再和 PyTorch 用 load_state_dict 读同一个 BF16 文件之后的输出比：
logits n= 12 max|diff|=5.96e-08
worst=5.96e-08 tol=1e-06 OK
1.10e-3 是存成 BF16 时丢掉的精度（BF16 只有 8 位有效位），两边丢的一样；读进来以后两边一致到 5.96e-8。半精度文件的对照对象要选对：和原始的 f32 模型比，差的是存储精度，不是 burn 算错了。
如果模型本来就打算在 F64 上跑，不要转参数，改为让输入也是 F64（设备默认 dtype 的设置见 dtype 与设备）。
怎么检查
- 加载前在 Python 里看一眼文件：
safe_open(path, "pt")打开后对每个键f.get_slice(k).get_dtype()，就是上面第二行输出的写法。 - 加载后在 burn 里走一遍参数，数一数不是 F32 的有几个（上面的
non-F32 params）。ApplyResult不报 dtype，errors = 0不代表能跑。 - 加载之后马上用一个假输入跑一次 forward，让 dtype 问题在加载的地方暴露，而不是在第一个真实请求里。
- 从 PyTorch 来的权重，一律在适配器链末尾带上
FloatCastAdapter::to(DType::F32)；文件本来就是 F32 时它什么也不改。
小结
- PyTorch 的
load_state_dict把 F64 / F16 / BF16 的权重转成参数原有的 dtype；burn-store 保留文件的 dtype，加载报告 0 错误。 - 第一次 matmul 才 panic：
assertion `left == right` failed: matmul: dtype mismatch，debug 和 release 都一样。 - 改法：
PyTorchToBurnAdapter.chain(FloatCastAdapter::to(DType::F32))；F64 文件转回后和 PyTorch 差 1.49e-8，BF16 文件和 PyTorch 读同一文件差 5.96e-8。 - 检查：加载后遍历参数看 dtype，再跑一次假输入。