BERT

主线位置：推理主线的 safetensors 那条路。VarBuilder::from_mmaped_safetensors 按名字取 f32 张量（见 VarBuilder、safetensors 与 mmap），BertModel::load 用它们搭出 embedding 和 12 层 encoder，forward 一次算完整个句子，没有 KV 缓存，也不取样。代码在 candle-transformers/src/models/bert.rs（625 行）和 with_tracing.rs（195 行）。

flowchart TB
  ids["input_ids 和 token_type_ids"] --> emb["BertEmbeddings"]
  emb --> enc["BertEncoder 12 层"]
  mask["attention_mask"] --> ext["扩展成加性 mask"]
  ext --> enc
  enc --> seq["B x T x 768"]
  seq -->|mean pooling| vec["句向量"]
  seq -->|BertForMaskedLM| mlm["词表 logits"]

本页的实验在 models-exp 的 bert 里（release），模型是本地的 bert-base-chinese（model.safetensors 412 MB，F32），平台 M2 Max、macOS 27、rustc 1.99.0。

结构对到代码

bert.rs 是照着 Hugging Face transformers 的 modeling_bert.py 一个类一个类翻译的，每个结构体上面都留着对应 Python 代码的链接。权重名字由 vb.pp(..) 一层层拼出来：

几处和 Python 版不同的地方：

• 没有 pooler。 Python 的 BertModel 还有一个 pooler（取 [CLS] 位置过 dense 和 tanh）。candle 的 BertModel::forward 只返回 (B, T, 768) 的序列输出，句向量怎么取由调用方决定。bert-base-chinese 文件里的 bert.pooler.dense.* 不会被读。
• Dropout 是空的。 Dropout::forward 里只有一行 // TODO 和 Ok(x.clone())：这个模型只能做推理。
• 位置 id 总是 0..seq_len。 BertEmbeddings::forward 里自己生成，旁边一条注释「TODO: Proper absolute positions?」。
• softmax 用的是拼出来的 ops::softmax（max_keepdim、exp、sum_keepdim、除法四个 op），不是融合的 softmax_last_dim。LayerNorm 走 with_tracing::LayerNorm → candle_nn::LayerNorm，在连续输入上用融合的 layer_norm（见常用层）。
• GELU 是精确的 erf 版本。 config.json 里的 "hidden_act": "gelu" 反序列化成 HiddenAct::Gelu，对应 gelu_erf；GeluApproximate 是 tanh 近似。后面会看到这个选择在 CPU 上很贵。

attention_mask 在 get_extended_attention_mask 里变成加到注意力分数上的值：

    let attention_mask = attention_mask.to_dtype(dtype)?;
    // torch.finfo(dtype).min
    (attention_mask.ones_like()? - &attention_mask)?.broadcast_mul(
        &Tensor::try_from(f32::MIN)?
            .to_device(attention_mask.device())?
            .to_dtype(dtype)?,
    )

mask 为 1 的位置加 0，为 0 的位置（padding）加 f32::MIN，softmax 之后权重为 0。不传 mask 时用全 1，即每个位置都能看所有位置。

加载 bert-base-chinese：两处名字对不上

直接照 candle-examples 的写法加载（load）：

plain load: Err(cannot find tensor embeddings.word_embeddings.weight)

这条报错是误导。BertModel::load 会试两种前缀：

        let (embeddings, encoder) = match (
            BertEmbeddings::load(vb.pp("embeddings"), config),
            BertEncoder::load(vb.pp("encoder"), config),
        ) {
            (Ok(embeddings), Ok(encoder)) => (embeddings, encoder),
            (Err(err), _) | (_, Err(err)) => {
                if let Some(model_type) = &config.model_type {
                    if let (Ok(embeddings), Ok(encoder)) = (
                        BertEmbeddings::load(vb.pp(format!("{model_type}.embeddings")), config),
                        BertEncoder::load(vb.pp(format!("{model_type}.encoder")), config),
                    ) {
                        (embeddings, encoder)
                    } else {
                        return Err(err);
                    }
                // ...

先试不带前缀，失败了再用 config.model_type（这里是 "bert"）当前缀重试。重试也失败时返回的是第一次的错误，所以报的是没前缀的名字，而文件里本来就有 bert.embeddings.word_embeddings.weight。真正缺的要换成 BertForMaskedLM::load（它直接 vb.pp("bert")、不重试）才看得到：

BertForMaskedLM plain: Some("cannot find tensor bert.embeddings.LayerNorm.weight")

bert-base-chinese 是从 TensorFlow 转过来的老文件，LayerNorm 的参数叫 gamma / beta：

bert.embeddings.LayerNorm.beta F32 [768]
bert.embeddings.LayerNorm.gamma F32 [768]

而 candle_nn::layer_norm 要的是 weight / bias。Hugging Face transformers 加载时会替你改名（本机 uv 缓存里 transformers 4.57.1 的 modeling_utils.py 有一段注释「Rename LayerNorm beta & gamma params for some early models ported from Tensorflow (e.g. Bert)」），candle 不会。改名用 VarBuilder::rename_f（机制见 VarBuilder）：

fn renamed(vb: VarBuilder<'static>, log: Arc<Mutex<Vec<String>>>) -> VarBuilder<'static> {
  vb.rename_f(move |name: &str| {
    log.lock().unwrap().push(name.to_string());
    if let Some(p) = name.strip_suffix("LayerNorm.weight") {
      format!("{p}LayerNorm.gamma")
    } else if let Some(p) = name.strip_suffix("LayerNorm.bias") {
      format!("{p}LayerNorm.beta")
    } else {
      // ...（BertForMaskedLM 还要两条，见下文）
      name.to_string()
    }
  })
}

闭包拿到的是拼好前缀的完整名字。顺手把每次请求记下来，和文件里的名字对一下：

renamed load: ok; 199 names requested, file has 206 tensors
first requests: ["embeddings.word_embeddings.weight", "encoder.layer.0.attention.self.query.weight", "bert.embeddings.word_embeddings.weight"]
file tensors never requested (9): ["bert.pooler.dense.bias", "bert.pooler.dense.weight", "cls.predictions.bias", "cls.predictions.transform.LayerNorm.beta", "cls.predictions.transform.LayerNorm.gamma", "cls.predictions.transform.dense.bias", "cls.predictions.transform.dense.weight", "cls.seq_relationship.bias", "cls.seq_relationship.weight"]

• 前两次请求就是那两次失败的无前缀尝试：match 里的元组两边都会求值，BertEncoder::load 也试了一次。
• 197 次成功的请求 = embeddings 5 个 + 12 层 × 16 个。
• 没用到的 9 个：pooler 2 个，预训练用的两个头 cls.predictions（掩码语言模型）和 cls.seq_relationship（下一句预测）7 个。

用 BertForMaskedLM 验一下权重

句向量看不出权重加载得对不对，掩码预测可以。BertForMaskedLM 在 BertModel 上加了 cls.predictions 头，它的输出层要 cls.predictions.decoder.weight 和 .bias。bert-base-chinese 里没有这两个：输出层和词嵌入共享权重，偏置存在 cls.predictions.bias。只改 gamma / beta 时：

BertForMaskedLM gamma/beta only: Some("cannot find tensor cls.predictions.decoder.weight")

在改名闭包里再加两条，cls.predictions.decoder.weight → bert.embeddings.word_embeddings.weight，cls.predictions.decoder.bias → cls.predictions.bias。之后预测 [MASK]（mlm，CPU，前 3 个及概率）：

北京是中国的[MASK]都。 -> 首=0.905 古=0.054 国=0.006
我今天很高[MASK]。 -> 兴=0.815 興=0.184 傲=0.000
Rust 是一种编程语[MASK]。 -> 言=1.000 法=0.000 语=0.000

预测是对的，说明 199 个名字都落在了正确的张量上。改名闭包是共享权重的通用办法：两个名字指向文件里的同一个张量，VarBuilder 各读一份（不是零拷贝，见 safetensors 与 mmap）。

句向量：mean pooling

BertModel 输出每个 token 一个 768 维向量。句向量最常见的取法是对非 padding 的位置求平均（mean pooling），再做 L2 归一化，余弦相似度就是点积：

fn mean_pool(out: &Tensor, mask: &Tensor) -> Result<Tensor> {
  let m = mask.to_dtype(DType::F32)?.unsqueeze(2)?;
  let sum = out.broadcast_mul(&m)?.sum(1)?;
  let pooled = sum.broadcast_div(&m.sum(1)?)?;
  Ok(pooled.broadcast_div(&pooled.sqr()?.sum_keepdim(1)?.sqrt()?)?)
}

六个句子，tokenizer 开 BatchLongest 补齐成一个 (6, 18) 的批，attention_mask 来自 tokenizer（embed，CPU）：

mean pooling, cosine:
  s0 1.000 0.852 0.687 0.717 0.676 0.661
  s1 0.852 1.000 0.685 0.724 0.661 0.655
  s2 0.687 0.685 1.000 0.912 0.658 0.636
  s3 0.717 0.724 0.912 1.000 0.664 0.642
  s4 0.676 0.661 0.658 0.664 1.000 0.815
  s5 0.661 0.655 0.636 0.642 0.815 1.000
[CLS] vector, cosine:
  s0 1.000 0.739 0.665 0.659 0.599 0.605
  s1 0.739 1.000 0.704 0.696 0.637 0.624
  s2 0.665 0.704 1.000 0.931 0.667 0.646
  s3 0.659 0.696 0.931 1.000 0.664 0.631
  s4 0.599 0.637 0.667 0.664 1.000 0.814
  s5 0.605 0.624 0.646 0.631 0.814 1.000

• 每一行最大的非对角元素都是同一话题的另一句：天气（0.852）、难吃的饭（0.912）、Rust（0.815）。
• 但不相关的句子之间也有 0.63 到 0.72。bert-base-chinese 是用掩码语言模型目标训练的，没有为句子相似度训练过，向量都挤在一个窄锥里。排序可用，绝对值不能当成「相似的概率」。专门训练过的句向量模型（例如 sentence-transformers 系列）分得开得多，candle 的 bert 例子默认用的就是 all-MiniLM-L6-v2。
• 直接取 [CLS] 位置（不过 pooler）也能分出这三组，s1 和 s0 的差距小一些。

padding 有没有被正确遮住，可以拿单独编码的句子对比：

s1 alone (14 tokens) vs in batch: max abs diff 8.009374e-8
s1 in batch without attention_mask: max abs diff 2.3064047e-2

s1 单独算（14 个 token，不补齐）和放在补齐到 18 的批里算，句向量最大差 8e-8，只是浮点误差。同一个批不传 attention_mask（forward(.., None)，全 1），真实 token 也会注意到 4 个 [PAD]，句向量差到 0.02。

with_tracing：每个层带一个 span

bert.rs 的 Linear、LayerNorm 不是 candle_nn 的，而是 with_tracing 里的包装：

#[derive(Debug, Clone)]
pub struct Linear {
    inner: candle_nn::Linear,
    span: tracing::Span,
}

pub fn linear(d1: usize, d2: usize, vb: VarBuilder) -> Result<Linear> {
    let inner = candle_nn::linear(d1, d2, vb)?;
    let span = tracing::span!(tracing::Level::TRACE, "linear");
    Ok(Linear { inner, span })
}

impl Module for Linear {
    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        let _enter = self.span.enter();
        self.inner.forward(xs)
    }
}

每个包装都是「内层 + 一个 TRACE 级 span」，forward 进入 span 再转发。bert.rs 自己的结构体也各有一个 span：model、embeddings、encoder、layer、attn、self-attn、softmax、self-out、inter、hidden-act、out。span 在构造层的时候创建，forward 里只 enter。没有 subscriber 时，span! 返回一个禁用的 span，enter 什么也不做（见 tracing 卷的 Span 与 Entered）。candle-examples 的 bert 用 tracing-chrome 把这些 span 写成 Chrome trace 文件（--tracing）。

这里换一个更直接的看法：装一个自己写的 tracing_subscriber::Layer，在 on_enter / on_exit 记时间，按 span 名字汇总次数、总时间和自身时间（减去期间进入的其他 span）。层之间的嵌套按进入顺序算，因为这些 span 的「父 span」是在建模型时定的，和调用关系无关。一个 (6, 18) 的批，预热两轮后的第三轮（trace）：

同一轮的墙钟时间（forward 之后 synchronize）：CPU 87.2 ms，Metal 15.0 ms。encoder 一行和 layer 几乎一样，省略了。

• 次数对得上结构：linear 72 次 = 12 层 × 6 个（q、k、v、self-out、inter、out），layer-norm 25 次 = 12 × 2 + embedding 的 1 个。
• CPU 上一半以上的时间在 hidden-act，比 72 次矩阵乘加起来还多。gelu_erf 在 CPU 上是逐元素的标量循环：GeluErf 只有标量版本的 f32，每个元素调一次 erf_f32，单线程（逐元素运算单线程见 CPU 后端）。而 tanh 近似的 Gelu 在开了 accelerate feature 时有 f32_vec，整段交给 Accelerate 的 vs_gelu。矩阵乘本身走 Accelerate 的 sgemm，很快。
• Metal 上 span 只量到编码时间。 所有 span 加起来一两毫秒，墙钟时间却是十几毫秒：forward 只是把内核编进命令缓冲区，synchronize 时 GPU 才算完。span 告诉你 CPU 侧哪一层编码最慢，不告诉你 GPU 在哪一层花了时间。

把激活换成 tanh 近似（config.hidden_act = HiddenAct::GeluApproximate，gelu）：

各跑 6 轮，去掉每种的第一轮。

CPU 上一次前向省掉一半多；单看一个 (6, 18, 3072) 张量，gelu_erf 2.3 ms 对 gelu 0.33 ms。句向量最大差 1.9e-4，余弦相似度矩阵最大差 2.6e-4，上面的排序不受影响。Metal 上两种激活都是一个内核，几乎没有差别。candle-examples 的 bert 有 --approximate-gelu 开关，是 PR #1410（「Speed up bert with approx gelu」）加的，当时的 README 里 bge-large-zh-v1.5 一次前向从 176.7 ms 降到 116.8 ms。

CPU 与 Metal

同一个批，time：加载（含 rename_f）后连跑几轮 forward + mean pooling + to_vec2 读回；再把 6 个句子复制 8 份拼成 (48, 18) 的批：

这台机器在做这组实验时还在跑别的程序，同一设置在不同进程之间有 10% 到 20% 的波动，表里是几次运行的范围。

• Metal 快 4 到 7 倍。 12 层里 Linear 的权重一共约 8500 万个（每层 4 × 768² + 2 × 768 × 3072），每个 token 约 1.7 亿次浮点运算：小批约 18 GFLOP，大批约 147 GFLOP。Metal 上分别是每秒 1.1 和 1.8 TFLOP 左右，CPU 上不到 0.3 TFLOP，其中还有一半时间不在矩阵乘上（gelu_erf）。
• 批越大 Metal 越划算。 token 数乘 8，CPU 时间乘 7 左右，Metal 只乘 5 左右：小批时 Metal 每个内核都喂不饱，固定开销占比大。
• 进程里第一轮 Metal 要 40 多毫秒，之后 16 ms：第一次用到的内核要建管线、缓冲区要第一次分配。macOS 的 Metal 库缓存没命中时第一轮会长得多（这组实验里见过一次 406 ms、一次 4.1 s），机制见 Metal 后端。
• 结果和 CPU 一致：Metal 上的余弦相似度矩阵前三位小数和上面完全相同，s1 单独算与在批里算差 2.4e-7。
• 只算句向量、批又不大时，CPU 换 tanh 近似 GELU 之后约 36 ms，和 Metal 的差距缩到两倍多。本站的服务器没有 GPU，用 bert-base-chinese 给文章算向量的讨论见 10 在本站中。

动脑筋答案

trace-late，CPU：

round 0: wall (after synchronize) 94.38 ms
round 1: wall (after synchronize) 83.07 ms
round 2: wall (after synchronize) 86.38 ms
[Cpu last round] spans: name, enters, total ms, self ms

一行都没有。span! 在执行时问当前的 subscriber 要不要这个 span；那时还没有 subscriber，得到的是一个禁用的 span，存进了结构体的 span 字段。之后再装 subscriber，这些字段里的 span 也不会变成启用的：enter 对禁用的 span 是空操作，subscriber 一次 on_enter 都收不到。所以 candle-examples 的 bert 例子在 main 里先初始化 tracing，再加载模型。

小结

• bert.rs 一个结构体对一个 Python 类：BertEmbeddings → 12 个 BertLayer（自注意力、输出、中间层、输出）；没有 pooler，Dropout 是空操作，forward 返回 (B, T, 768)。
• 加载时先试无前缀、再试 bert. 前缀，失败时报的是第一次的名字；bert-base-chinese 的 LayerNorm 叫 gamma / beta，要用 rename_f 改名；BertForMaskedLM 还要把输出层指到词嵌入上。
• mean pooling 要用 tokenizer 的 attention_mask，既传给 forward，也用来求平均；不传 mask 时 padding 会改变句向量。
• bert-base-chinese 的句向量能排序、不能当绝对分数：无关句子之间余弦也有 0.6 以上。
• with_tracing 的包装在构造时建 span，subscriber 必须在加载模型之前装好。CPU 上 span 量出 gelu_erf 占一半时间，换成 tanh 近似快一倍、句向量差 2e-4；Metal 上 span 只量到编码时间。