Rust in Everything
基于 candle 0.11.0

unsupported dtype BF16 for op matmul

主线位置：推理主线第三步「前向，logits 对上」，为了省内存把权重留在 bf16 / f16，还在 CPU 上跑。
版本：PyTorch 2.15.0 / transformers 5.22.0（CPU）、candle 0.11.0（CPU，accelerate 和不开 accelerate；Metal）
flowchart TB
  w["bf16 权重"] --> d{"VarBuilder dtype"}
  d -->|"BF16 CPU"| e1["matmul 报错"]
  d -->|"F16 CPU accelerate"| e2["f16 matmul 报错"]
  d -->|"F16 CPU gemm"| bad["能跑 差 0.43"]
  d -->|"F32 CPU"| ok1["差 3e-5"]
  d -->|"BF16 F16 Metal"| ok2["和 torch 同量级"]

实验代码：PyTorch 侧 half_cpu.py，candle 侧 crate pitb-exp 的 bin p19_half，同一份代码用 --no-default-features（不开 accelerate）另编一次。release 构建，平台 M2 Max、macOS 27、rustc 1.99.0。权重是本地的 SmolLM2-135M-Instruct；PyTorch 固定 torch.set_num_threads(8)。各设备、各 dtype 的完整表格见半精度。

PyTorch 的写法

模型在 CPU 上，整个 .to(dtype)：

model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.float32, attn_implementation="eager").eval()
with torch.inference_mode():
  ref = model(ids).logits[0, -1]
  ref_gen = model.generate(ids, max_new_tokens=20, do_sample=False)[0, 5:]
  for dt in [torch.bfloat16, torch.float16]:
    m = model.to(dt)  # 习惯写法：整个模型换 dtype，还在 CPU 上
    lg = m(ids).logits[0, -1].float()
    gen = m.generate(ids, max_new_tokens=20, do_sample=False)[0, 5:]

cpu torch.bfloat16  max|diff| vs f32 = 4.336e-01 (max|ref| 20.6), greedy same prefix 18/20
cpu torch.float16   max|diff| vs f32 = 6.842e-02 (max|ref| 20.6), greedy same prefix 20/20

两种都能跑，误差是半精度该有的量级。（这两个数和半精度那页 from_pretrained(dtype=torch.bfloat16) 的 5.34e-1、4.86e-2 不同：.to(dtype) 把 RoPE 的 inv_freq buffer 也转成了 bf16 / f16，from_pretrained(dtype=...) 让它留在 float32。两种写法算的不是同一个模型。）

搬到 candle 的现象

candle 里选 dtype 的地方是 VarBuilder，Cache 也要同一个 dtype：

let model = Llama::load(st_vb(dt, &dev)?, &cfg)?;
println!("loaded {dt:?} accelerate={}", cfg!(feature = "accelerate"));
let (logits, greedy) = run(&model, &cfg, dt, &dev)?; // Cache::new(true, dt, ..)，prefill + 19 步
println!("{}", check(&logits, &greedy)?);

st_vb 是 VarBuilder::from_mmaped_safetensors(&[".../model.safetensors"], dt, &dev)，check 和 transformers f32 的参照比。CPU，开 accelerate（默认）：

== acc cpu bf16
loaded BF16 accelerate=true
Error: unsupported dtype BF16 for op matmul
== acc cpu f16
loaded F16 accelerate=true
Error: the accelerate backend does not support f16 matmul

CPU，不开 accelerate：

== noacc cpu bf16
loaded BF16 accelerate=false
Error: unsupported dtype BF16 for op matmul
== noacc cpu f16
loaded F16 accelerate=false
logits max|diff|=4.28e-1 (max|ref| 20.6), greedy same prefix 20/20
  • 加载都成功（loaded 那一行），报错在第一次前向。PyTorch 里 dtype 不支持通常在 .to() 时就知道，这里要等真正做矩阵乘。
  • 不开 accelerate 的 F16 不报错，误差 0.428，是 PyTorch CPU f16（0.068）的 6 倍，和 bf16 一个量级。这是三种情况里最难发现的一种。

原因

CPU 的矩阵乘是 impl Map2 for MatMul 里按 feature 二选一的 f：

  • 不开 accelerate 时是 gemm 那一支，开头一个白名单：

          match T::DTYPE {
          DType::F16 | DType::F32 | DType::F64 => {}
          _ => Err(Error::UnsupportedDTypeForOp(T::DTYPE, "matmul").bt())?,
      }

    没有 BF16。F16 交给 gemm crate，它在 aarch64 上选 f16 的乘加指令，累加器本身是 f16，1536 项加下来误差就到了这个量级（单次矩阵乘的对照见半精度的「CPU 上的 F16」一节）。

  • 开 accelerate 时是另一个 f，只有 F32（sgemm）和 F64（dgemm），F16 单独 bail!("the accelerate backend does not support f16 matmul")，其余落到同一个 UnsupportedDTypeForOp。

Storage::matmul 在分派前只查两边 dtype 一致，不查后端支不支持，所以错误出现在 CPU 后端里面。全表和 Metal 一侧的情况见源码卷的坑 unsupported dtype BF16 for op matmul。origin/main 上 cpu_backend/mod.rs 另有几处修改（包括批维 stride 为 0 的 matmul，#3960），这两处 dtype 分支没有变。

改法

CPU 上载成 F32；要半精度就放到 Metal：

== acc cpu f32
loaded F32 accelerate=true
logits max|diff|=2.98e-5 (max|ref| 20.6), greedy same prefix 20/20
== acc metal bf16
loaded BF16 accelerate=true
logits max|diff|=6.02e-1 (max|ref| 20.6), greedy same prefix 20/20
== acc metal f16
loaded F16 accelerate=true
logits max|diff|=4.52e-2 (max|ref| 20.6), greedy same prefix 20/20
  • CPU F32 就是 A3 的结果：差 2.98e-5，容差 1e-4。文件是 bf16，载成 F32 时内存翻倍（峰值 RSS 见什么时候选 candle）。
  • Metal BF16 差 0.602、F16 差 0.045，和 PyTorch 同 dtype 的误差同一量级；bf16 下的贪心是否一致受近平局影响，不说明谁更准（半精度）。
  • 不要用「CPU + F16 + 关掉 accelerate」来绕开报错：能跑，但误差是 f16 该有的 6 倍以上。
  • 想要 CPU 上的小内存，用量化（GGUF），不是 bf16（GGUF）。

怎么检查

  • dtype 按设备选，写在一处：let dt = if dev.is_cpu() { DType::F32 } else { DType::BF16 };，VarBuilder 和 Cache::new 用同一个 dt。
  • 加载后先跑一次很短的前向（一个 token）当冒烟测试，dtype 的问题在这里暴露，不要等到服务收到第一个请求。
  • 看 Cargo.toml 里 candle-core 有没有 accelerate（或 mkl）feature：同一份代码在两种构建下对 F16 的反应不同，一个报错、一个静默变差。

小结

  • PyTorch CPU 上 model.to(torch.bfloat16) / float16 能跑，logits 差 0.43 / 0.068。
  • candle CPU：BF16 前向报 unsupported dtype BF16 for op matmul；F16 开 accelerate 报 the accelerate backend does not support f16 matmul，不开时能跑但用 f16 累加，差 0.428。加载都成功，错在第一次 matmul。
  • CPU 用 F32（差 2.98e-5），半精度放到 Metal（BF16 0.602、F16 0.045）。

💬 关联讨论

发起讨论
加载中...