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()。