跳到主要内容

数据截至 (上游 commit 932e1f2f4c5a)

02 · 模型抽象与三种请求

这一章讲什么: 「同一套任务跑在 vLLM、transformers、nanotron、远程 API 上」是怎么做到的;以及评测分数里最不可见的变量——分词与批处理细节——在代码里是怎么被处理的。读完你能自己写一个后端,也明白为什么换 harness 分数会变。


1. 它要解决的小问题

任务层(第 1 章)产出的是统一答题卡 Doc。但要让模型「答题」,有三种本质不同的问法:

问法适用题型需要模型给出什么
生成式(GENERATIVE)开放问答、数学、写作一段文字
选项打分(LOGPROBS)多选题、分类每个候选答案的 log 概率
整段困惑度(PERPLEXITY)语料困惑度、bits-per-byte一整段文字的 log 概率

难点有三层:

  1. 后端能力不齐:vLLM 没有「打分」API,远程 API 不给你 logprobs,nanotron 的模型是训练中分片加载的。
  2. 分词是隐形炸弹:context 和 continuation 拼起来分词,和分开分词,结果可能不同;结尾空格归谁、EOS 剥不剥,都会改分数。
  3. 吞吐:几万条样本,长度差 10 倍,batch 怎么组、OOM 怎么办、重复跑怎么缓存。

2. 直觉:一个契约 + 一个信封

lighteval 的答案是两层抽象:

契约——LightevalModel(src/lighteval/models/abstract_model.py:220)是抽象基类,每个后端只需实现三个抽象方法,签名全是「list[Doc] 进、list[ModelResponse] 出」(abstract_model.py:250-282):

  • greedy_until(docs) — 生成;
  • loglikelihood(docs) — 给每个 Doc 的每个 choice 打分;
  • loglikelihood_rolling(docs) — 给整段文本打分(困惑度)。

信封——ModelResponse(src/lighteval/models/model_output.py:28)是模型回答的统一信封:text(生成的文字)、logprobs(各候选的分数)、argmax_logits_eq_gold(贪心输出是否恰好等于 gold)、input_tokens/output_tokens 等。不同问法只填信封的不同格子,指标层按需取格(第 3 章)。

直觉一句话:契约规定「怎么问」,信封规定「怎么答」;后端爱用 vLLM 还是 API 随意,只要答题卡进、信封出。


3. 图示:一条 Doc 在后端里的旅程

list[Doc](同一生成参数的一桶)


① DynamicBatchDataset 按长度降序排序,记住原始下标 (data.py:44)


② PromptManager Doc → prompt 文本(见第 1 章)


③ tok_encode_pair context/continuation 分别分词并对齐
(仅 loglikelihood 路) (abstract_model.py:297)


④ 引擎调用 vLLM generate / transformers forward /
动态 batch + 截断 nanotron 分布式 forward


⑤ 装 ModelResponse 文本 / logprobs / argmax 是否命中 gold


⑥ get_original_order 按①记下的下标还原顺序 (data.py:88)

怎么读这张图:①⑥ 是一对——排序是为了吞吐,还原是为了让「第 i 个回答对应第 i 道题」永远成立。


4. 原理演示:多选题怎么变成一次前向

loglikelihood 路的核心想法:多选题不是让模型「回答」,而是分别计算每个选项接在题面后的概率,谁高选谁。示意代码:

# 示意,非源码
def score_multiple_choice(model, context, choices):
scores = []
for choice in choices:
# 关键:context 与 choice 的分词要在交界处对齐(见 §5.1)
ctx_ids, cont_ids = tok_encode_pair(context, [choice])
logits = model(ctx_ids[0] + cont_ids[0]).logits
# logits[i] 预测第 i+1 个 token,所以要错一位取
lp = sum(log_softmax(logits)[i, t] for i, t in enumerate(cont_ids[0]))
scores.append(lp)
return int(np.argmax(scores)) # 分数最高的选项下标

重点看注释那行:错位一位是这类代码最容易错的地方,真实实现里有专门的修正和一段注释说明(§5.3)。


5. 真实实现

5.1 分词对齐:tok_encode_pair

tok_encode_pair(src/lighteval/models/abstract_model.py:297-354)是 loglikelihood 路的地基,做三件不显然的事:

  1. 挪动结尾空格(默认开):context 以空格结尾时,把空格挪到 continuation 头上——「 Paris」会被分词成单个 token ĠParis,概率质量集中,分数更真实(abstract_model.py:321-325)。做 answer-only perplexity / bits-per-byte 的后端可以用 move_trailing_context_space=False 关掉,让 continuation 严格等于 gold 字符串。
  2. 剥 context 末尾的 EOS(pairwise 模式):BOS/EOS 若夹在中间,模型会「无视」前文(abstract_model.py:335-336 的注释与代码)。
  3. 两种切法可选:pairwise=True 时 context 和 continuation 完全分开编码再拼;否则整段编码后按长度差切回(abstract_model.py:342-352)。前者对中文这类无空格语言更稳(docstring 里明说)。

5.2 vLLM 后端:用「假生成」换 prompt_logprobs

vLLM 只有生成接口,没有打分接口。lighteval 的做法(VLLMModel._generate,src/lighteval/models/vllm/vllm_model.py:422-481):

# vllm_model.py:443-449(generate=False 分支)
sampling_params = SamplingParams(
temperature=0.0,
prompt_logprobs=1, # 让引擎顺带吐出 prompt 各位置的 logprob
max_tokens=1, # 只生成 1 个 token——我们根本不想要生成
detokenize=False,
)

即:把「context + continuation」当 prompt 喂进去,让它假生成 1 个 token,真正的产物是 prompt_logprobs——prompt 里每个位置的概率分布。

随后 _loglikelihood_tokens(vllm_model.py:487-553)从 prompt_logprobs 尾部抠出 continuation 对应的几个 token:

  • 逐个取 logprobs_at_position[token],求和得该选项总分(vllm_model.py:530-537);
  • 同时检查每个位置是否 rank 1,得到「贪心输出是否恰好是 gold」的布尔分 argmax_logits_eq_gold(vllm_model.py:534-539)。

多选题 N 个选项 = N 次 prefill,_loglikelihood_tokens 把同一 Doc 的 N 个选项摊平进同一批,再按 len(doc.choices) 切回(vllm_model.py:505-551)。

生成路 _greedy_until(vllm_model.py:337-420)相对常规,但有两个决策值得知道:

  • 左截断:context + max_new_tokens 超过模型上限时,从左边砍 context,砍完还不够直接抛错(vllm_model.py:374-397)。代码注释明说了取舍:「宁可砍 prompt,也不让生成被压到过短」。
  • 数据并行:data_parallel_size > 1 时不开一个 vLLM 实例,而是起 N 个 ray remote 各起一个 LLM,请求交错分配(interleave)以均衡各 worker 的 context 长度(vllm_model.py:451-472)。

VLLMModel 的生成/打分入口都带 @cached 装饰器(vllm_model.py:322:483),缓存见 §5.4。

5.3 transformers 后端:排序 + 自动 batch + 分布式收集

transformers 后端(TransformersModel,src/lighteval/models/transformers/transformers_model.py:200)是功能最全的本地后端,它的吞吐来自三件套:

第一件:按长度降序排。 LoglikelihoodDataset._sorting_criteria 返回负的「题面长度 + 最长选项长度」(src/lighteval/data.py:161-183),降序排。注释给出三个理由:时间估计只会高估不会低估;每个 batch 的第一个元素就是该 batch 的 padding 长度;OOM 会在一开始就炸,而不是跑到 90% 才炸(data.py:163-173)。生成路 GenerativeTaskDataset 更进一步,先按 (num_samples, use_logits, stop_sequences, gen_length) 分组再按长度排,保证同组样本生成参数一致(data.py:228-251)。

第二件:OOM 自动减半找 batch。 _get_batch_size(transformers_model.py:505-530)拿当前最长输入造一个全 1 假 batch 试跑,find_executable_batch_size(src/lighteval/utils/parallelism.py:57)捕获 OOM 就把 batch 减半重试。定下来后,loglikelihood 路还会按「选项数」折算——batch 实际是 batch_size × max_choices 条序列,并向上取整到 8 的倍数(transformers_model.py:947-953);每跑完一个 split,下一个 split 的起始 batch 翻倍(因为样本更短了,transformers_model.py:946)。

第三件:错位修正与跨进程收集。 _loglikelihood_tokens(transformers_model.py:919-1141)里,选项打分的切片起点是 start = max(input_length - continuation_length - 1, 0)(transformers_model.py:1012)——旁边注释记录了一个真实 bug 史:rolling 模式曾用未错位的 logits,导致每个 token 错一位、困惑度被虚高。多进程(accelerate)下,各 rank 的选项数/长度不一,代码先把每维 pad 到全局最大再 gather_for_metrics(transformers_model.py:1046-1104)。

模型已在内存的场景(比如训练脚本里)走 TransformersModel.from_model(transformers_model.py:254)直接包装,不重新加载权重——README 的 Python API 例子用的就是它。

5.4 缓存:模型回答按样本落盘

@cached(src/lighteval/utils/cache_management.py:364-449)包在每个后端的三个入口上,机制:

  1. :(模型配置 sha256 前 16 位, 任务名, 任务配置 sha256 前 16 位, 采样方式, 样本 id)(get_model_hash/_get_task_hash,cache_management.py:141-183)。改任何生成参数、改任务定义,hash 变,缓存自动失效。
  2. :parquet 文件,目录层级就是键的层级(SampleCache 类 docstring,cache_management.py:64-75)。
  3. 粒度:样本级——1000 道题里命中 800 道,只重跑 200 道,再合并(get_samples_to_process_and_cache,cache_management.py:243-269)。

一个分布式细节:accelerate 多进程下每个 rank 都握着收集后的全量结果,但只有主进程写 parquet,其余 rank 在 barrier 等着读——注释指向 issue #1102:并发写同一个 parquet 会写坏(cache_management.py:417-431)。


6. 坑

  • vLLM 后端没有 loglikelihood_rolling,直接 NotImplementedError(vllm_model.py:555-557);困惑度任务请换 transformers/nanotron。
  • 异步 vLLM 路径较新。 AsyncVLLMModel(vllm_model.py:560-719)走 vLLM v1 的 AsyncLLM,采样类指标(num_samples>1)在它上面有「可能有意外行为」的警告(vllm_model.py:626-630)。
  • 同步 vLLM 非 chat 模式吞吐受限。 它假设同 batch 样本 stop/max_tokens 一致,注释明说因此只按 batch size 1 处理(vllm_model.py:354-358);缓解靠 §5.3 说的按生成参数分组。
  • pad token 被设为 eos。 vLLM 分词器初始化里 tokenizer.pad_token = tokenizer.eos_token(vllm_model.py:319)——常见做法,但对「pad 与 eos 语义不同」的模型要留意。
  • 缓存合并时,已缓存样本的新结果会被丢弃并打 warning(cache_management.py:343-347)——换模型不换配置 hash 的场景不会触发,但手动改 parquet 会。
  • enforce_eager=True 被硬编码进 vLLM 引擎参数(vllm_model.py:279),关掉了 CUDA graph——换来的是兼容性和启动速度,损失的是峰值推理吞吐(inferred:注释未解释原因)。

7. 代码地图(本章)

主题文件路径符号名
后端契约src/lighteval/models/abstract_model.pyLightevalModeltok_encode_pairtok_encode
后端分发src/lighteval/models/model_loader.pyload_modelload_model_with_accelerate_or_defaultload_custom_model
回答信封src/lighteval/models/model_output.pyModelResponseBatch
vLLM 后端src/lighteval/models/vllm/vllm_model.pyVLLMModel._generate_loglikelihood_tokens_greedy_untilAsyncVLLMModel
transformers 后端src/lighteval/models/transformers/transformers_model.pyTransformersModel._loglikelihood_tokens_get_batch_size_padded_greedy_untilfrom_model
长度排序与分组src/lighteval/data.pyDynamicBatchDatasetLoglikelihoodDatasetGenerativeTaskDatasetGenDistributedSampler
OOM 减半src/lighteval/utils/parallelism.pyfind_executable_batch_sizetest_all_gather
回答缓存src/lighteval/utils/cache_management.pySampleCachecachedTaskID
nanotron 后端src/lighteval/models/nanotron/nanotron_model.pyNanotronLightevalModel(第 4 章展开)