Rust in Everything
基于 candle-core 0.11.0

index-select invalid index 512 with dim size 512

主线位置:推理主线 safetensors 那条路,BERT 前向的第一步,BertEmbeddings::forward 把词嵌入、类型嵌入、位置嵌入加起来。BERT 的结构见 BERT,Embedding 就是 index_select,见 常用层。代码在 candle-transformers 的 models/bert.rs。

flowchart TB
  text["一页中文 540 字"] --> tok["分词 542 个 token"]
  tok --> pos["位置 0 到 541"]
  pos --> dev{"设备"}
  dev -->|CPU| err["invalid index 512"]
  dev -->|Metal| clamp["下标夹到 511 静默"]
  tok --> fix["截断到 512 或分段"]

实验 crate exp-candle/pb-pit,bin bert512,本地的 bert-base-chinese,release,M2 Max、macOS 27、rustc 1.99.0。加载时的 gamma / beta 改名见 上一页。

现象

超过 512 个 token

一页 540 个汉字的文本,按分词器原样编码后送进模型:

let page = "张量的数据放在一个读写锁后面。".repeat(36);
let enc = tok.encode(page.as_str(), true).map_err(E::msg)?;
println!("{} chars -> {} tokens", page.chars().count(), enc.get_ids().len());
let out = embed(&model, enc.get_ids())?;

CPU:

540 chars -> 542 tokens
Error: index-select invalid index 512 with dim size 512
exit=1

Metal(METAL=1):

540 chars -> 542 tokens
[1, 542, 768]
exit=0

CPU 报的是一个索引错误,不提 BERT、不提长度上限;Metal 上则没有错误,出来一个 542 行的结果。单独测 Metal 的 index_select,512 行的表(第 i 行存 i),取下标 510 到 513:

metal index_select rows Ok([510.0, 511.0, 511.0, 511.0])
cpu   index_select rows Err("index-select invalid index 512 with dim size 512")

Metal 把越界的 512、513 都换成了 511:第 512 个 token 之后的位置嵌入全是最后一个位置的,模型照算,结果不对也不报错。

英文全成了 [UNK]

"Rust 的 Tensor 在 Metal 上" -> ["[UNK]", "的", "[UNK]", "在", "[UNK]", "上"]
"rust 的 tensor 在 metal 上" -> ["ru", "##st", "的", "te", "##ns", "##or", "在", "metal", "上"]

出处

BertEmbeddings::forward 不看长度,直接为每个 token 造位置下标:

        if let Some(position_embeddings) = &self.position_embeddings {
            // TODO: Proper absolute positions?
            let position_ids = (0..seq_len as u32).collect::<Vec<_>>();
            let position_ids = Tensor::new(&position_ids[..], input_ids.device())?;
            embeddings = embeddings.broadcast_add(&position_embeddings.forward(&position_ids)?)?
        }

位置嵌入表是 config.json 里的 "max_position_embeddings": 512 行。CPU 的 IndexSelect 对每个下标检查 index >= src_dim,返回 Error::InvalidIndex,于是有了标题那句话。Metal 的 index 内核 则是有意夹住:

      const INDEX_TYPENAME input_i = min(input_ids[id_i], (INDEX_TYPENAME)(src_dim_size - 1));
      const size_t right_rank_i = tid % right_size;
      const size_t left_rank_i = tid / right_size / ids_size;
      /*
      // Force prevent out of bounds indexing
      // since there doesn't seem to be a good way to force crash
      // No need to check for zero we're only allowing unsized.
      */

注释说得很清楚:GPU 内核没法让进程崩溃,就退而求其次不越界读。代价是同一个 index_select 在两个设备上一个报错、一个静默给出别的行。这一行从 2023 年 Metal 后端加入索引算子时就在。

分词器那一边:bert-base-chinese 的 tokenizer.json 里 BertNormalizer 的 "lowercase": false(tokenizer_config.json 也是 "do_lower_case": false),而 vocab.txt 的 21128 行里带大写字母的只有 [PAD]、[UNK]、<S> 这类 7 个特殊记号,大写开头的词整个变成 [UNK]。这是模型文件的设定,不是 candle 的问题,换哪个分词器实现读同样的文件都一样;candle 只是不会替你转小写。

修法

让分词器截断,或者按 510 个 token(留出 [CLS] 和 [SEP])分段分别算再合并;英文先转小写:

tok
  .with_truncation(Some(TruncationParams { max_length: 512, ..Default::default() }))
  .map_err(E::msg)?;
let enc = tok.encode(page.as_str(), true).map_err(E::msg)?;
println!("truncated: {} tokens, {} overflow pieces", enc.get_ids().len(), enc.get_overflowing().len());
println!("{:?}", embed(&model, enc.get_ids())?.shape());
let s = "Rust 的 Tensor 在 Metal 上".to_lowercase();
truncated: 512 tokens, 1 overflow pieces
[1, 512, 768]
"rust 的 tensor 在 metal 上" -> ["ru", "##st", "的", "te", "##ns", "##or", "在", "metal", "上"]
  • 截断之后多出来的 30 个 token 在 get_overflowing() 里(1 段),不想丢就对它们再算一次,再把几段的向量平均或分别检索。在本站中 的 chunks 方案就是按 510 个 token 切段、按 token 数加权平均成一页的向量。
  • max_length 从 config.max_position_embeddings 读,不要写死;不同的 BERT 变体不一样。
  • 转小写要在分词之前做,对整段文本 to_lowercase()。

怎么确认没中招

  • 进模型之前断言 ids.len() <= cfg.max_position_embeddings。这一条在 Metal 上尤其重要:那里越界不会报错。
  • 统计一下 [UNK] 的比例(tok.token_to_id("[UNK]") 在 get_ids() 里出现的次数)。中文文本里夹着英文时,没转小写的 [UNK] 比例会明显偏高。
  • 同一段输入在 CPU 和 Metal 上各跑一次:一个报错一个不报,就是长度超了。

小结

  • BertModel 不截断输入;超过 512 个 token 时,CPU 上位置嵌入的 index_select 报 index-select invalid index 512 with dim size 512。
  • Metal 的索引内核把越界下标夹到最后一行,不报错,512 之后的 token 都用第 511 个位置的嵌入。
  • 分词器开截断(或按 510 分段),长度上限从配置读;bert-base-chinese 的分词器不转小写,自己 to_lowercase()。

💬 关联讨论

发起讨论
加载中...