进阶

获取 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_logprobslist[float | None],长度和 prompt_ids 一致。
  • sequence.logprobslist[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 就不能直接按位置相减。

本页目录