Function compute_logprobs

src/post_training/rollout.py:233–265

Teacher-forced recomputation of per-token log-probs of ``sequences`` under ``model``. Mirrors the model's training shift: ``logits[:, t]`` predicts ``sequences[:, t+1]``, so the returned tensors have length ``T-1`` and are aligned to target positions ``1..T-1``. Args:

(
    model,
    sequences: torch.Tensor,
    response_mask: torch.Tensor,
    *,
    temperature: float = 1.0,
    requires_grad: bool = True,
)

def compute_logprobs(
    model,
    sequences: torch.Tensor,
    response_mask: torch.Tensor,
    *,
    temperature: float = 1.0,
    requires_grad: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Teacher-forced recomputation of per-token log-probs of ``sequences`` under ``model``.

    Mirrors the model's training shift: ``logits[:, t]`` predicts ``sequences[:, t+1]``,
    so the returned tensors have length ``T-1`` and are aligned to target positions
    ``1..T-1``.

    Args:
        sequences:     (B, T) token ids (prompt + response).
        response_mask: (B, T) bool over response positions (as from :class:`RolloutBatch`).
        temperature:   divide logits by this before log-softmax, matching sampling.
        requires_grad: if False, runs under ``no_grad`` (use for ref / old-policy).

    Returns:
        (logprobs, mask) each (B, T-1). ``logprobs`` is the log-prob of the realized
        next token; ``mask`` is the response mask shifted to align with those targets.
    """
    ctx = torch.enable_grad() if requires_grad else torch.no_grad()
    with ctx:
        logits = _logits_from(model, sequences)[:, :-1, :]  # predict tokens 1..T-1
        logprobs_all = F.log_softmax(logits.float() / max(temperature, 1e-6), dim=-1)
        targets = sequences[:, 1:].unsqueeze(-1)
        logprobs = logprobs_all.gather(-1, targets).squeeze(-1)  # (B, T-1)
    mask = response_mask[:, 1:].to(torch.bool)
    return logprobs, mask


def sequence_logprobs(