Class ReferenceModelWrapper

rl_engine/alignment/model_wrappers.py:113–160

Standard adapter for the frozen reference model used by KL penalties.

class ReferenceModelWrapper(PolicyModelWrapper):
    """Standard adapter for the frozen reference model used by KL penalties."""

    def __init__(
        self,
        model: torch.nn.Module,
        *,
        freeze: bool = True,
        eval_mode: bool = True,
    ):
        super().__init__(model)
        if freeze:
            self.freeze()
        if eval_mode:
            self.eval()

    def freeze(self) -> "ReferenceModelWrapper":
        for parameter in self.model.parameters():
            parameter.requires_grad_(False)
        return self

    def forward_logits(self, input_ids: torch.Tensor, **model_kwargs: Any) -> torch.Tensor:
        with torch.no_grad():
            return super().forward_logits(input_ids, **model_kwargs)

    def selected_logprobs(
        self,
        input_ids: torch.Tensor,
        token_ids: torch.Tensor,
        *,
        mask: Optional[torch.Tensor] = None,
        logits_start: Optional[int] = None,
        logits_end: Optional[int] = None,
        temperature: float = 1.0,
        output_dtype: torch.dtype = torch.float32,
        **model_kwargs: Any,
    ) -> torch.Tensor:
        with torch.no_grad():
            return super().selected_logprobs(
                input_ids,
                token_ids,
                mask=mask,
                logits_start=logits_start,
                logits_end=logits_end,
                temperature=temperature,
                output_dtype=output_dtype,
                **model_kwargs,
            )