获取 Logprobs
在 TRIO 中,有两种常用方式可以获取 token 级别的 logprobs:
| 方式 | 适合场景 | 能拿到什么 |
|---|---|---|
sample() | 让模型继续生成,并记录这次采样过程 | prompt token 的 logprobs、生成 token 的 logprobs |
compute_logprobs() | 已经有一段完整文本,用指定模型重新打分 | 输入全文中每个实际 token 的 logprobs |
简单来说,sample() 回答的是“模型这次采样时,每个 prompt token 和生成 token 的概率是多少”;
compute_logprobs() 回答的是“给定这一整段文本,指定模型对其中每个 token 有多认可”。
方式一:通过 sample 获取 logprobs
sample() 会让模型从 prompt 继续生成。它默认会在 sequences[*].logprobs 中返回生成 token
的 logprobs;如果设置 include_prompt_logprobs=True,还会返回 prompt token 的 logprobs。
import pytrio as trio
client = trio.ServiceClient()
sampling_client = client.create_sampling_client(base_model="Qwen/Qwen3.5-4B")
tokenizer = sampling_client.get_tokenizer()
messages = [
{"role": "user", "content": "1 + 1 等于多少?"},
]
prompt_text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
prompt_ids = tokenizer.encode(prompt_text, add_special_tokens=False)
response = sampling_client.sample(
prompt=trio.ModelInput.from_ints(prompt_ids),
sampling_params=trio.SamplingParams(max_tokens=8, temperature=0.7),
include_prompt_logprobs=True,
).result()
sequence = response.sequences[0]
prompt_logprobs = response.prompt_logprobs
completion_tokens = list(sequence.tokens)
completion_logprobs = list(sequence.logprobs)
print(type(prompt_logprobs), len(prompt_logprobs))
print(type(completion_logprobs), len(completion_logprobs))
print(completion_tokens)
print(completion_logprobs)
print(sequence.text)这段代码里真正可以直接拿来用的是两个列表:
response.prompt_logprobs:list[float | None],长度和prompt_ids一致。sequence.logprobs:list[float | None],长度和sequence.tokens一致,对应模型本次生成出来的 token。
因此返回结果的数据结构可以理解为:
response.prompt_logprobs
[
None,
-0.20,
-0.00,
...
]
response.sequences[0].tokens
[17, 151645]
response.sequences[0].logprobs
[-0.42, -0.02]
response.sequences[0].text
'2'如果希望看得更直观,可以把 token id 和 logprob zip 到一起:
completion = [
{
"token_id": token_id,
"token": tokenizer.decode([token_id]),
"logprob": logprob,
}
for token_id, logprob in zip(sequence.tokens, sequence.logprobs)
]得到的数据会类似这样:
[
{"token_id": 17, "token": "2", "logprob": -0.42},
{"token_id": 151645, "token": "<|im_end|>", "logprob": -0.02},
]上面的 token id 和数值只是输出形态示例,实际结果会随模型、chat template、采样参数和生成内容变化。 重点是:prompt logprobs 和生成 token logprobs 都是可以直接拿来做 RL、重打分或日志分析的列表。
方式二:通过 compute_logprobs 获取全文 logprobs
compute_logprobs() 不会生成新 token,而是对你传入的一整段 ModelInput 做前向评分:
def compute_logprobs(self, prompt: ModelInput) -> APIFuture[list[float | None]]返回结果和输入 token 一一对齐:
- 第
i个返回值表示第i个 token 在前文 token 条件下的 log probability。 - 第一个 token 没有前文条件,因此可能返回
None。 - 返回的是每个实际 token 的对数概率,不是整个词表的概率分布。
import pytrio as trio
client = trio.ServiceClient()
sampling_client = client.create_sampling_client(base_model="Qwen/Qwen3.5-4B")
tokenizer = sampling_client.get_tokenizer()
messages = [
{"role": "user", "content": "1 + 1 等于多少?"},
{"role": "assistant", "content": "2"},
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
)
tokens = tokenizer.encode(text, add_special_tokens=False)
logprobs = sampling_client.compute_logprobs(
prompt=trio.ModelInput.from_ints(tokens),
).result()
scored_tokens = [
{
"token_id": token_id,
"token": tokenizer.decode([token_id]),
"logprob": logprob,
}
for token_id, logprob in zip(tokens, logprobs)
]compute_logprobs() 返回的是 list[float | None],长度和输入的 tokens 一致。上面的
scored_tokens 会类似这样:
[
{"token_id": 151644, "token": "<|im_start|>", "logprob": None},
{"token_id": 872, "token": "user", "logprob": -0.20},
{"token_id": 198, "token": "\n", "logprob": -0.00},
{"token_id": 16, "token": "1", "logprob": -2.18},
{"token_id": 488, "token": " +", "logprob": -5.86},
{"token_id": 220, "token": " 1", "logprob": -1.09},
{"token_id": 151645, "token": "<|im_end|>", "logprob": -0.00},
{"token_id": 198, "token": "\n", "logprob": -0.00},
{"token_id": 151644, "token": "<|im_start|>", "logprob": 0.0},
{"token_id": 77091, "token": "assistant", "logprob": -0.01},
{"token_id": 198, "token": "\n", "logprob": -0.00},
{"token_id": 17, "token": "2", "logprob": -7.80},
{"token_id": 151645, "token": "<|im_end|>", "logprob": -0.02},
]这也是 compute_logprobs() 和 sample() 的核心区别:sample() 只对一次生成过程返回
prompt 与生成 token 的 logprobs;compute_logprobs() 可以对一段已经确定的完整文本重新打分。
在 OPD 中计算 teacher logprobs
OPD(on-policy distillation)里,student 先用当前策略生成 completion,然后 teacher 对
student 实际生成的同一段 token 轨迹计算 logprobs。这样可以得到逐 token 的
student_logprobs - teacher_logprobs,再把 reverse KL 转成 importance_sampling 所需的
advantage。
在 OPD 和类似的模型蒸馏流程中,teacher 通常会选择能力更强、更稳定,或者已经对目标任务表现更好的模型。 teacher logprobs 不是用来继续生成文本,而是用来回答一个训练问题:更强的 teacher 对 student 实际采样出来的每个 token 有多认可。这个逐 token 信号随后会进入 reverse KL 或 advantage 计算,指导 student 向 teacher 偏好的 token 分布靠近。
关键点是:teacher 看到的输入必须是 prompt + completion,然后只截取 completion 区间的
logprobs。
import numpy as np
import pytrio as trio
def completion_teacher_logprobs(
teacher_client,
prompt_ids: list[int],
completion_ids: list[int],
) -> list[float]:
all_ids = prompt_ids + completion_ids
all_logprobs = teacher_client.compute_logprobs(
prompt=trio.ModelInput.from_ints(all_ids),
).result()
completion_logprobs = all_logprobs[len(prompt_ids):]
if len(completion_logprobs) != len(completion_ids):
raise ValueError("teacher logprobs and completion tokens are not aligned")
if any(value is None for value in completion_logprobs):
raise ValueError("completion logprobs should not contain None")
return [float(value) for value in completion_logprobs]
# student 采样时返回的 seq.logprobs 是旧策略对 completion token 的 logprobs。
student_logprobs = [float(value) for value in seq.logprobs]
teacher_logprobs = completion_teacher_logprobs(
teacher_client=teacher_client,
prompt_ids=prompt_ids,
completion_ids=list(seq.tokens),
)
reverse_kl = np.asarray(student_logprobs) - np.asarray(teacher_logprobs)
advantages = -kl_penalty_coef * reverse_kl这段逻辑对应 OPD 里的核心训练信号:
reverse_kl = student_logprobs - teacher_logprobs
advantages = -kl_penalty_coef * reverse_kl后续再把 completion_ids、student 采样时的 logprobs 和这里得到的 advantages 右移对齐,
放进 Datum.loss_fn_inputs,即可用 loss_fn="importance_sampling" 训练。
和常见生成 API 的区别
很多云端生成 API 默认只返回最终文本。部分 API 可以额外返回生成阶段的 token logprobs,例如
DeepSeek 的 logprobs 参数会返回输出 token 的对数概率,top_logprobs 会返回每个输出位置的
top-N token 及其对数概率。OpenAI Chat Completions 和 Gemini 也提供类似的生成阶段输出
token logprobs 能力。
这类接口通常回答的是“模型这次生成出来的 token 概率是多少”。compute_logprobs 解决的是另一个
问题:给定一段已经确定的 token 序列,计算指定模型对这段序列中每个实际 token 的 logprobs。
因此,在 OPD 这样的流程里,你可以让 student 负责采样,再让 teacher 对同一条完整轨迹重新打分。
需要注意的是,逐 token 相减要求 token 对齐。实践中建议 student 和 teacher 使用同一 tokenizer
或至少兼容的 tokenizer;如果两者 tokenization 不一致,student_logprobs - teacher_logprobs
就不能直接按位置相减。