半精度

主线位置：推理主线的第三步「前向，logits 对上」换成半精度。权重文件本来就是 bf16，f32 那一页（对齐 logits）是加载时转成了 f32；本页不转，看两边在哪些设备上能跑、logits 离 f32 多远、贪心生成还剩几个 token 一样。

版本：PyTorch 2.15.0 / transformers 5.22.0（CPU、MPS）、candle 0.11.0（CPU，accelerate 和不开 accelerate；Metal）

flowchart TB
  st["bf16 权重文件"] --> vb["VarBuilder dtype"]
  vb --> f32["F32"]
  vb --> bf["BF16"]
  vb --> h["F16"]
  f32 --> ok1["CPU Metal 都能跑"]
  bf --> cpu1["CPU matmul 报错"]
  bf --> mt1["Metal 能跑"]
  h --> acc["CPU 开 accelerate 报错"]
  h --> gemm["CPU gemm 能跑 f16 累加"]
  h --> mt2["Metal 能跑"]

实验代码：Python 侧 pt_logits.py、pt_tf.py、pt_attn32.py、pt_margin.py、pt_f16mm.py；candle 侧 crate llm-exp 的 bin run、tf、f16mm、cachedt、tie，同一份代码另外用 --no-default-features（不开 accelerate）编一次。release 构建，平台 M2 Max、macOS 27、rustc 1.99.0。模型是本地的 SmolLM2-135M-Instruct 权重；PyTorch 固定 torch.set_num_threads(8)。参照一律是 transformers eager、f32、CPU 的输出（就是 A3 的参照）。

怎么切到半精度

transformers 在 from_pretrained 时给 dtype，要上 GPU 再 .to("mps")：

dtype = {"f32": torch.float32, "bf16": torch.bfloat16, "f16": torch.float16}[dt]
model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=dtype, attn_implementation=attn).to(dev).eval()
ids = torch.tensor([[504, 3575, 282, 4649, 314]], device=dev)  # "The capital of France is"
with torch.inference_mode():
  logits = model(ids).logits[0, -1].float().cpu()
  out = model.generate(ids, max_new_tokens=20, do_sample=False)

candle 在建 VarBuilder 时给 dtype，权重在加载时转换；Cache::new 也要同一个 dtype：

let dt = dtype(&args[2]); // DType::F32 / BF16 / F16
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[format!("{model_dir}/model.safetensors")], dt, &dev)? };
let model = Llama::load(vb, &cfg)?;
let mut cache = Cache::new(true, dt, &cfg, &dev)?;
let logits = model.forward(&input, 0, &mut cache)?.squeeze(0)?;

Cache 里装着 RoPE 的 cos / sin 表（KV 缓存），按这个 dtype 存。模型用 BF16、Cache 还用 F32 建（cachedt，Metal）：

Err(dtype mismatch in rope BF16 F32 F32)

PyTorch 一侧没有对应的坑：transformers 的 LlamaRotaryEmbedding 每次按输入的 dtype 转换 cos / sin。

两边返回的 logits 都换回 f32 再比：transformers 是 .float()，candle 的 Llama::forward 最后一行本来就是 to_dtype(DType::F32)。

哪些组合能跑

candle 的两个报错都发生在第一次前向，加载是成功的：

rope_theta=100000 dtype=BF16
Error: unsupported dtype BF16 for op matmul
rope_theta=100000 dtype=F16
Error: the accelerate backend does not support f16 matmul

CPU 后端按 dtype 分派 matmul，没有 BF16 那一支；开 accelerate 时 F16 也没有。全表和出处见源码卷的坑 unsupported dtype BF16 for op matmul。PyTorch 用户的习惯是「CPU 上也能 model.to(torch.bfloat16) 省一半内存」；在 candle 的 CPU 上要么载成 F32（什么时候选 candle 量过：bf16 文件载成 F32，峰值 RSS 860 MiB），要么换到 Metal。

logits 离 f32 多远

最后一个位置的 49152 个 logits 和 f32 参照比（最大 |ref| 20.6），再看 20 个贪心 token 从第几个开始和参照不同：

拿什么当尺子：在 16 到 32 之间，bf16 相邻两个可表示数差 0.125，f16 差 0.0156。f16 的几行误差在 0.04 到 0.06，约 2.5 到 4 个 f16 间距；bf16 的几行在 0.47 到 0.60，约 4 到 5 个 bf16 间距。candle Metal 的两行和 torch 同 dtype 的几行是同一个量级：F16 的 4.52e-2 在 torch 的 3.88e-2 到 5.97e-2 之间，BF16 的 6.02e-1 比 torch 最大的 5.69e-1 高一点。例外是 CPU 上的 F16，下面单独说。

直接比两个半精度运行。candle Metal F16 对 torch MPS f16 eager，容差 0.1（约 6 个 f16 间距）：

logits                   n=49152  max|diff|=5.47e-02
greedy                   n= 20  max|diff|=0.00e+00
worst=5.47e-02 tol=1e-01 OK

candle Metal BF16 对 torch MPS bf16 eager，和 torch CPU bf16 eager 对 torch MPS bf16 eager，各取 compare.py 输出的第一行（两次的 greedy 都不同，所以都报 MISMATCH，原因见下一节）：

logits                   n=49152  max|diff|=3.12e-01
...
logits                   n=49152  max|diff|=3.09e-01
...

同 dtype 下，candle 对 torch 和 torch 自己从 CPU 换到 MPS 差得一样多。

贪心生成：4/20 不代表更差

bf16 那几行里 candle 是 20/20，torch 有三种是 4/20。这不说明 candle 的 bf16 算得更准。看 f32 参照里每一步第一名和第二名 logit 的差（pt_margin.py，沿参考序列一次前向取 20 个位置）：

margins: 0:0.569 1:0.251 2:0.447 3:3.325 4:0.094 5:0.492 6:1.603 7:2.955 8:1.650 9:0.887 10:0.309 11:0.089 12:3.122 13:1.773 14:0.424 15:0.577 16:1.133 17:0.061 18:0.042 19:3.398

第 4 步（ Paris. Paris is 之后选 the 还是 a）第一名只领先 0.094，第 17、18 步只领先 0.061 和 0.042，都比一个 bf16 间距还小，而 bf16 的误差是 0.5 上下。在这几步上选哪个 token 取决于舍入落在哪边，和实现好坏无关；第 4 步一旦选了 a，后面就是另一句话（ Paris is a major city in France, known for its historical landmarks, ...）。

为了把「一步分叉、后面全变」和「每一步对不对」分开，沿参考序列逐步取 logits（teacher forcing，每一步都喂参考序列里的 token，方法见 KV 缓存），比每一步的 argmax，20 步合起来最大 |ref| 32.73：

bf16 不同的步全落在 4、11、17、18 这四个领先不到 0.1 的位置上。torch MPS bf16 eager 逐步看是 20/20，自由生成却在第 4 步分叉：自由生成走的是 generate 的带缓存 decode，逐步看是一次 24 个 token 的前向，矩阵形状不同，第 4 步的舍入就落到了另一边。

不是注意力的精度

一个自然的猜测：candle 的 llama.rs 在注意力里先把 q、k、v 转成 F32（CausalSelfAttention::forward），transformers eager 只在 softmax 里用 f32，两次 matmul 都是 bf16：

            let in_dtype = q.dtype();
            let q = q.to_dtype(DType::F32)?;
            let k = k.to_dtype(DType::F32)?;
            let v = v.to_dtype(DType::F32)?;
            let att = (q.matmul(&k.t()?)? / (self.head_dim as f64).sqrt())?;
            // ...
            let att = candle_nn::ops::softmax_last_dim(&att)?;
            // Convert to contiguous as matmul doesn't support strided vs for now.
            att.matmul(&v.contiguous()?)?.to_dtype(in_dtype)?

attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
if attention_mask is not None:
  attn_weights = attn_weights + attention_mask
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
attn_output = torch.matmul(attn_weights, value_states)

把 transformers 的 eager_attention_forward 换成 candle 的做法再跑（pt_attn32.py）：

orig = ml.eager_attention_forward

def eager_f32(module, query, key, value, attention_mask, scaling, dropout=0.0, **kwargs):
  out, w = orig(module, query.float(), key.float(), value.float(),
                None if attention_mask is None else attention_mask.float(), scaling, dropout, **kwargs)
  return out.to(query.dtype), w

ml.eager_attention_forward = eager_f32

pt_attn32_bf16_cpu.json      max|diff|=3.816e-01  greedy prefix  4/20  same positions  7/20
pt_attn32_bf16_mps.json      max|diff|=5.374e-01  greedy prefix  4/20  same positions  4/20
pt_attn32_f16_cpu.json       max|diff|=4.058e-02  greedy prefix 20/20  same positions 20/20
pt_attn32_f16_mps.json       max|diff|=4.418e-02  greedy prefix 20/20  same positions 20/20

CPU bf16 的 logits 误差从 0.534 降到 0.382，贪心反而从 17/20 变成 4/20；MPS 不变。注意力用不用 f32 改变的是舍入落在哪，不决定第 4 步那个领先 0.094 的 token 能不能保住。半精度下比较两份实现，看逐步的 logits 误差和 argmax，不要看自由生成 20 个 token 的前缀长度。

CPU 上的 F16：累加也是 f16

不开 accelerate 时 candle CPU 能跑 F16，但 logits 误差是 0.428，是其他 f16 运行的 10 倍左右，和 bf16 一个量级。只看一次矩阵乘（f16mm）：输入先舍入到 f16，参照在 f32 里乘，剩下的就只是累加和最后一次舍入的误差：

let a16 = a.to_dtype(DType::F16)?.to_dtype(DType::F32)?;
let b16 = b.to_dtype(DType::F16)?.to_dtype(DType::F32)?;
let want = a16.matmul(&b16)?;
// 再分别在 CPU 和 Metal 上用 F16 直接乘，和 want 比

cpu f16 matmul [8,1536]x[1536,576]: max|diff| vs f32-accumulated = 8.5829e-2 (max|ref| 6.531)
metal f16 matmul [8,1536]x[1536,576]: max|diff| vs f32-accumulated = 1.9474e-3 (max|ref| 6.531)

PyTorch 同样的形状：

cpu f16 matmul [8,1536]x[1536,576]: max|diff| vs f32-accumulated = 1.9045e-03 (max|ref| 7.225)
mps f16 matmul [8,1536]x[1536,576]: max|diff| vs f32-accumulated = 1.9045e-03 (max|ref| 7.225)

Metal 和 PyTorch 的 1.9e-3 约是 6 附近 f16 间距（0.0039）的一半，也就是只差最后一次舍入：它们在 f32 里累加。candle CPU 差 8.6e-2，是这个的 44 倍：不开 accelerate 时 CPU 的 matmul 交给 gemm crate（0.19.0），它在支持 fp16 指令的 aarch64 上选 neonfp16 微内核，乘加用的是 vfmaq_f16，累加器本身是 f16。1536 项加下来，误差就到了这个量级。CPU 上的 BF16 / F16 求和用半精度累加器，源码卷的坑 CPU 上 BF16 softmax 加起来不是 1 里还有一例；llama.rs 的 softmax 在 F32 里算，躲开了那一个，躲不开 matmul。

动脑筋答案

bf16 在 16 到 32 之间的间距是 16 / 128 = 0.125。22.371271 和 22.277479 夹在 22.25 和 22.375 之间，前者离 22.375 近，后者离 22.25 近，所以直接把 f32 参照转成 bf16， the 仍然领先 0.125。但半精度运行里的 logit 不是参照转出来的，是 30 层 bf16 计算的结果，每个都带着 0.5 上下的误差。实际几次运行第 4 步的这两个 logit：

tf_pt_eager_f32_cpu_full.json    the=22.371271 a=22.277479
tf_pt_eager_bf16_cpu_full.json   the=22.375000 a=22.250000
tf_pt_sdpa_bf16_cpu_full.json    the=22.250000 a=22.250000
tf_c_metal_bf16_cache.json       the=22.500000 a=22.375000
tf_c_metal_f16_cache.json        the=22.375000 a=22.281250

• bf16 的值都是 0.125 的整数倍，f16 是 1/64 的整数倍（22.28125）。
• torch CPU bf16 sdpa 两个数正好相等，都是 22.25。并列时 argmax 选下标小的那个： a 是 253， the 是 260，所以选了 a，第 4 步分叉。candle 也一样，CPU、Metal、F32、BF16 上 [.., 22.25 在 253, 22.25 在 260, ..] 的 argmax 都是 253，torch CPU、MPS 也是 253（tie）。
• candle Metal BF16 两个数都比参照高了一格（22.5 对 22.375），差距仍是一格，于是保住了 the。这是舍入的运气，不是精度。

小结

• 两边都是在加载时选 dtype：transformers from_pretrained(dtype=...)，candle VarBuilder 的 dtype；candle 的 Cache::new 也要同一个 dtype，否则报 dtype mismatch in rope BF16 F32 F32。
• candle CPU 不能用 BF16 跑（unsupported dtype BF16 for op matmul），开 accelerate 时 F16 也不行；Metal 两种都行。torch 的 CPU 和 MPS 两种都行。
• 对 f32 参照（最大 |ref| 20.6）：f16 差 0.04 到 0.06，candle Metal F16 是 0.045；bf16 差 0.47 到 0.60，candle Metal BF16 是 0.60。同 dtype 下 candle 对 torch 和 torch 换设备差得一样多。
• 贪心 20 个 token：bf16 的 torch 有三种 4/20，candle 20/20，原因是第 4 步第一名只领先 0.094，小于一个 bf16 间距（0.125）。逐步比 argmax，各家 bf16 都是 17 到 20 个相同，不同的步全在领先不到 0.1 的位置。
• 不开 accelerate 的 CPU F16 误差 0.43：gemm 在 aarch64 上用 f16 累加，单次矩阵乘的误差是 Metal 和 PyTorch 的 44 倍。