bge-reranker-v2-m3 批量重排显存翻倍 怎么切 batch

这个坑值得记一笔。24G 的 4090 上用 FlagEmbedding 跑 bge-reranker-v2-m3，一次把 64 条候选和问题配对送进 rerank，显存从 9G 冲到 23G 后 OOM，batch_size 用的还是默认值。

刚开始盯着「FlagEmbedding 默认 batch_size 偏大，长文档下显存近似平方增长」，折腾了半个多小时，排掉了。

这类报错背后常见的原因有几处，逐个排除就行。rerank 的输入是 query 与文档成对展开，条数等于候选数；打分时序列长度取批次内最长，padding 把短句也撑满；模型以 fp32 加载，权重本身先占了两倍。

量当前进程的显存占用做基线

python -c "import torch;print(round(torch.cuda.memory_allocated()/1024**3,2))"

显式给 batch_size 并开 fp16

python -c "from FlagEmbedding import FlagReranker;r=FlagReranker('BAAI/bge-reranker-v2-m3',use_fp16=True);print(r.compute_score([['什么是RAG','检索增强生成']],batch_size=8,normalize=True))"

然后把候选截到 512 token 再送重排

python -c "from transformers import AutoTokenizer;t=AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-m3');print(len(t('文本'*2000)['input_ids']))"

压测一批看峰值显存落点

python -c "import torch;print(round(torch.cuda.max_memory_allocated()/1024**3,2))"

中途换过一次思路，把范围缩到最小的那一层，反而快了。

改完先别撤，确认这几项。64 条候选全部返回分数且排序稳定，nvidia-smi 峰值显存低于 18G，单次重排耗时不超过 400 毫秒。

提醒一点。重排收益在候选 20 到 50 条之间最明显，再往上加候选只是线性涨延迟，命中率提升很小。这类问题最好一次改到位，反复小改反而容易把原因搅混。