# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
import contextlib
from dataclasses import dataclass
from typing import Literal, TYPE_CHECKING
import torch
from tensordict import NestedKey, TensorClass, TensorDictBase
from tensordict.nn import TensorDictModule
from tensordict.utils import _zip_strict
from torchrl.data import History
from torchrl.modules.llm.policies.transformers_wrapper import TransformersWrapper
from torchrl.objectives.common import LossModule
from torchrl.objectives.llm._utils import _runtime_device
if TYPE_CHECKING:
import transformers
[docs]
def k3_kl_token_estimate(
target_log_prob: torch.Tensor, log_prob: torch.Tensor
) -> torch.Tensor:
"""Per-token k3 estimate of the KL divergence to a target distribution.
Uses the k3 estimator ``exp(d) - 1 - d`` with
``d = target_log_prob - log_prob``, the same approximation used by
:class:`~torchrl.objectives.llm.GRPOLoss` and
:class:`~torchrl.objectives.llm.SFTLoss` for their KL regularizers, but
returned per token instead of reduced, so callers can apply their own
masking and reduction.
The estimator is direction-neutral: it estimates
``KL(sampling distribution || target distribution)``, and the direction
of the resulting KL is determined by which distribution's
log-probabilities are passed as which argument. The estimate is unbiased
only when the scored tokens are actual samples from the distribution
behind ``log_prob``; on tokens drawn from any other source, the (masked)
sum or mean of the result is still a nonnegative divergence between the
two distributions, but it no longer estimates either KL divergence.
Args:
target_log_prob (torch.Tensor): per-token log-probabilities of the
target (non-sampling) distribution, evaluated on the same tokens
as ``log_prob``.
log_prob (torch.Tensor): per-token log-probabilities of the
distribution the tokens were sampled from.
Returns:
A tensor of nonnegative per-token KL contributions with the same
shape as the inputs.
References:
- John Schulman, 2020. `"Approximating KL Divergence" <http://joschu.net/blog/kl-approx.html>`_
Examples:
>>> import torch
>>> from torchrl.objectives.llm import k3_kl_token_estimate
>>> log_prob = torch.full((4,), -1.0)
>>> k3_kl_token_estimate(log_prob, log_prob)
tensor([0., 0., 0., 0.])
"""
if target_log_prob.shape != log_prob.shape:
raise ValueError(
f"target_log_prob and log_prob must have the same shape, got "
f"{target_log_prob.shape=} and {log_prob.shape=}."
)
diff = target_log_prob - log_prob
return diff.expm1() - diff
def _distillation_loss(summed_kl: torch.Tensor, reduction: str) -> torch.Tensor:
"""Reduce per-sequence KL estimates into the distillation loss."""
if reduction == "mean":
return summed_kl.mean()
if reduction == "sum":
return summed_kl.sum()
if reduction == "none":
return summed_kl
raise ValueError(f"Invalid reduction: {reduction}.")
[docs]
class DistillationLossOutput(TensorClass["nocast"]):
"""Distillation Loss Output.
Attributes:
loss_distill (torch.Tensor): The differentiable distillation loss.
kl_to_teacher (torch.Tensor): The detached mean per-sequence KL
estimate, for logging.
.. note::
``kl_to_teacher`` is not differentiable; summing the output class
directly (``loss_output.sum(reduce=True)``) is not the intended way
to obtain the total loss here. Backpropagate through
``loss_distill`` instead.
"""
loss_distill: torch.Tensor
kl_to_teacher: torch.Tensor
[docs]
class DistillationLoss(LossModule):
r"""Token-level knowledge-distillation loss for LLM policies.
Distills a student policy toward a (typically stronger, frozen) teacher by
minimizing a KL divergence estimated from the per-token log-probabilities
the two models assign to the same tokens. The student log-probabilities
are computed by running ``actor_network`` on the input history; the
teacher log-probabilities are read from the input tensordict, where a
:class:`~torchrl.envs.llm.transforms.RetrieveLogProb` transform wrapping
the teacher (or an offline scoring pass) has written them.
Both distillation directions are supported through ``kl_direction``:
- ``"reverse"`` (default): minimizes ``KL(student || teacher)``. This is
the on-policy setting: the scored tokens must be sampled from the
student for the estimate to be that KL.
- ``"forward"``: minimizes ``KL(teacher || student)``. Here the scored
tokens must be sampled from the teacher (teacher-generated completions)
for the estimate to be that KL.
In both cases the KL is estimated per token with the k3 estimator
(:func:`~torchrl.objectives.llm.k3_kl_token_estimate`), the same
approximation used by the KL regularizers of
:class:`~torchrl.objectives.llm.GRPOLoss` and
:class:`~torchrl.objectives.llm.SFTLoss`. As in those losses, gradients
flow through the student log-probabilities only; the score-function term
of the sampling distribution is ignored.
.. note::
The k3 estimator is an unbiased KL estimate only when the scored
tokens are samples from the corresponding sampling distribution: the
student for ``"reverse"``, the teacher for ``"forward"``. On data
that comes from neither model (e.g. a plain text corpus), the
objective is still a valid nonnegative divergence pulling the student
toward the teacher, but its value is not the advertised KL and it can
have very high variance wherever the student assigns much more
probability than the teacher. For off-policy distillation, prefer
teacher-generated completions with ``kl_direction="forward"``.
.. note::
Because the score-function term of the sampling distribution is
dropped, the expected gradient of both directions matches the
gradient of the forward (mass-covering) KL. The usual mode-seeking
(reverse) vs mass-covering (forward) distinction therefore applies to
the logged value (``kl_to_teacher``), not to the direction the
optimization actually descends.
Args:
actor_network (TensorDictModule): the student network to train.
Usually a :class:`~torchrl.modules.llm.TransformersWrapper`
instance with ``generate=False``.
tokenizer (`Tokenizer`, optional): the tokenizer used to re-tokenize
the history when token masks are not present in the input.
Defaults to the tokenizer of ``actor_network``, when available.
tokenizer_kwargs (dict, optional): keyword arguments passed to
:meth:`~torchrl.data.llm.History.apply_chat_template` when
re-tokenizing.
kl_direction (Literal["reverse", "forward"], optional): the KL
direction to minimize (see above). Defaults to ``"reverse"``.
reduction (Literal["mean", "sum", "none"], optional): the reduction to
apply across sequences. Defaults to ``"mean"``.
normalize_by_seq_length (bool, optional): whether to normalize each
sequence's KL by its number of scored tokens. Defaults to ``True``.
assistant_only (bool, optional): if ``True`` (default), only assistant
(response) tokens contribute to the loss. If ``False``, all
attended tokens contribute. This flag must be paired consistently
with the ``assistant_only`` flag of the
:class:`~torchrl.envs.llm.transforms.RetrieveLogProb` transform
that wrote the teacher log-probabilities:
``RetrieveLogProb(assistant_only=True)`` zero-fills the teacher
log-probabilities of non-assistant tokens and can only be paired
with ``DistillationLoss(assistant_only=True)``; use
``RetrieveLogProb(assistant_only=False)`` when distilling on all
attended tokens.
device (torch.device | None, optional): fallback device used when neither
the input nor the student parameters and buffers provide one. This
does not move the student network; move the complete loss with
:meth:`~torch.nn.Module.to`. Defaults to ``None``.
.. note::
The input tensordict is expected to contain the following keys by
default:
- ``("history", "full")``: the chat history;
- ``("next", "teacher_log_probs", "full")``: the teacher per-token
log-probabilities on the same tokens, as written by a
:class:`~torchrl.envs.llm.transforms.RetrieveLogProb` constructed
with ``log_probs_full_key=("teacher_log_probs", "full")``.
These keys can be customized using the :meth:`set_keys` method.
The token masks are read from the fixed keys
``("masks", "all_assistant_mask")`` and
``("masks", "all_attention_mask")`` when present; otherwise the
history is re-tokenized with the tokenizer to recover them.
.. seealso:: :class:`~torchrl.envs.llm.transforms.RetrieveLogProb` to
compute the teacher log-probabilities, and
:class:`~torchrl.objectives.llm.SFTLoss` for supervised fine-tuning
with an optional KL regularizer.
References:
- Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk,
Sabela Ramos, Matthieu Geist, Olivier Bachem, 2024.
`"On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes" <https://arxiv.org/abs/2306.13649>`_
Examples:
>>> from torchrl.data.llm.history import History, _CHAT_TEMPLATES
>>> from torchrl.envs.llm.transforms import RetrieveLogProb
>>> from torchrl.modules.llm import TransformersWrapper
>>> from torchrl.modules.llm.policies.common import ChatHistory
>>> from torchrl.objectives.llm import DistillationLoss
>>> from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
>>> from tensordict import TensorDict, lazy_stack
>>> import torch
>>>
>>> chats = [
... [
... {"role": "system", "content": "You are a helpful assistant."},
... {"role": "user", "content": "Hello, how are you?"},
... {"role": "assistant", "content": "I'm doing well, thank you!"},
... ],
... [
... {"role": "system", "content": "You are a helpful assistant."},
... {"role": "user", "content": "What's the weather like?"},
... {"role": "assistant", "content": "I can't check the weather for you."},
... ],
... ]
>>> history = History.from_chats(chats)
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
>>> tokenizer.pad_token = tokenizer.eos_token
>>> tokenizer.chat_template = _CHAT_TEMPLATES["chatml_format"]
>>> student_model = OPTForCausalLM(OPTConfig())
>>> teacher_model = OPTForCausalLM(OPTConfig()).eval()
>>> student = TransformersWrapper(
... student_model,
... tokenizer=tokenizer,
... generate=False,
... chat_template_name="qwen",
... input_mode="history",
... pad_output=False,
... )
>>> teacher = TransformersWrapper(
... teacher_model,
... tokenizer=tokenizer,
... generate=False,
... return_log_probs=True,
... chat_template_name="qwen",
... input_mode="history",
... pad_output=False,
... )
>>> transform = RetrieveLogProb(
... teacher,
... log_probs_full_key=("teacher_log_probs", "full"),
... assistant_only=True,
... tokenizer_kwargs={"chat_template_name": "qwen"},
... tokenizer=tokenizer,
... )
>>> td = TensorDict(
... history=ChatHistory(
... full=history, prompt=history[..., :-1], response=history[..., -1:]
... ),
... next=TensorDict(
... done=torch.zeros(2, dtype=torch.bool),
... history=ChatHistory(prompt=history),
... ),
... batch_size=(2,),
... )
>>> data = lazy_stack(list(td.unbind(0)))
>>> with torch.no_grad():
... data = transform(data)
>>> loss = DistillationLoss(
... actor_network=student,
... tokenizer=tokenizer,
... kl_direction="reverse",
... tokenizer_kwargs={"chat_template_name": "qwen"},
... )
>>> loss_vals = loss(data)
>>> loss_vals.loss_distill.backward()
"""
@dataclass
class _AcceptedKeys:
"""Maintains default values for all configurable tensordict keys.
This class defines which tensordict keys can be set using
'.set_keys(key_name=key_value)' and their default values.
Attributes:
history (NestedKey): The input tensordict key where the chat
history is expected. Defaults to ``("history", "full")``.
teacher_log_prob (NestedKey): The input tensordict key where the
teacher per-token log-probabilities are expected.
Defaults to ``("next", "teacher_log_probs", "full")``.
log_probs (NestedKey): The key where the student model writes its
per-token log-probabilities. Defaults to ``("log_probs", "full")``.
"""
history: NestedKey = ("history", "full")
teacher_log_prob: NestedKey = ("next", "teacher_log_probs", "full")
log_probs: NestedKey = ("log_probs", "full")
default_keys = _AcceptedKeys
tensor_keys: _AcceptedKeys
def __init__(
self,
actor_network: TensorDictModule | TransformersWrapper,
tokenizer: transformers.AutoTokenizer | None = None, # noqa: F821
tokenizer_kwargs: dict | None = None,
kl_direction: Literal["reverse", "forward"] = "reverse",
reduction: Literal["mean", "sum", "none"] = "mean",
normalize_by_seq_length: bool = True,
assistant_only: bool = True,
device: torch.device | None = None,
):
super().__init__()
self.in_keys = []
self.actor_network = actor_network
if tokenizer is None:
tokenizer = getattr(actor_network, "tokenizer", None)
self.tokenizer = tokenizer
if tokenizer_kwargs is None:
tokenizer_kwargs = {}
tokenizer_kwargs.setdefault("return_assistant_tokens_mask", True)
tokenizer_kwargs.setdefault("tokenize", True)
tokenizer_kwargs.setdefault("return_tensors", "pt")
tokenizer_kwargs.setdefault("padding", False)
tokenizer_kwargs.setdefault("add_generation_prompt", False)
self.tokenizer_kwargs = tokenizer_kwargs
if kl_direction not in ("reverse", "forward"):
raise ValueError(
f"kl_direction must be 'reverse' or 'forward', got {kl_direction!r}."
)
self.kl_direction = kl_direction
if reduction not in ("mean", "sum", "none"):
raise ValueError(
f"reduction must be 'mean', 'sum' or 'none', got {reduction!r}."
)
self.reduction = reduction
self.normalize_by_seq_length = normalize_by_seq_length
self.assistant_only = assistant_only
self._set_in_keys()
self.register_buffer(
"_device_fallback", torch.empty(0, device=device), persistent=False
)
def _set_in_keys(self) -> None:
"""Sets the input keys for the loss module."""
self.in_keys = [self.tensor_keys.history, self.tensor_keys.teacher_log_prob]
self.out_keys = []
def _get_masks(
self, tensordict: TensorDictBase, history: History
) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor] | None]:
"""Returns the loss masks, the attention masks and, when recoverable, the assistant masks."""
assistant_masks = tensordict.get(("masks", "all_assistant_mask"), as_list=True)
attention_mask = tensordict.get(("masks", "all_attention_mask"), as_list=True)
device = _runtime_device(
self.actor_network,
tensordict,
assistant_masks,
attention_mask,
fallback=self._device_fallback,
)
if (self.assistant_only and assistant_masks is None) or attention_mask is None:
if self.tokenizer is None:
raise ValueError(
"A tokenizer is required to recover the token masks because the "
"input tensordict has no ('masks', ...) entries. Pass a tokenizer "
"to the loss constructor or provide the masks in the input."
)
with torch.device(
device
) if device is not None else contextlib.nullcontext():
token_struct = history.apply_chat_template(
tokenizer=self.tokenizer, **self.tokenizer_kwargs
)
if assistant_masks is None and "assistant_masks" in token_struct:
assistant_masks = token_struct.get("assistant_masks", as_list=True)
if self.assistant_only and assistant_masks is None:
raise ValueError(
f"Assistant masks are not present in the token structure: {token_struct=}."
)
if attention_mask is None:
attention_mask = token_struct.get("attention_mask", as_list=True)
attention_mask = [mask.bool() for mask in attention_mask]
if assistant_masks is not None:
assistant_masks = [mask.bool() for mask in assistant_masks]
if not self.assistant_only:
return attention_mask, attention_mask, assistant_masks
return (
[
mask & a_mask
for mask, a_mask in _zip_strict(assistant_masks, attention_mask)
],
attention_mask,
assistant_masks,
)
@staticmethod
def _check_teacher_log_probs_consistency(
teacher_log_probs: list[torch.Tensor],
attention_masks: list[torch.Tensor],
assistant_masks: list[torch.Tensor],
) -> None:
"""Detects zero-filled teacher log-probs on non-assistant tokens.
``RetrieveLogProb(assistant_only=True)`` fills the log-probs of
non-assistant tokens with exact zeros. Feeding those into the loss
with ``assistant_only=False`` would distill the student toward a
distribution that assigns probability one to every non-assistant
token, and the k3 estimator can explode there.
"""
zeros = 0
total = 0
for tlp, att, ast in _zip_strict(
teacher_log_probs, attention_masks, assistant_masks
):
if tlp.shape != att.shape or tlp.shape != ast.shape:
return
non_assistant = att & ~ast
total += non_assistant.sum().item()
zeros += ((tlp == 0.0) & non_assistant).sum().item()
if total > 0 and zeros > 0.9 * total:
raise RuntimeError(
f"{zeros} out of {total} attended non-assistant tokens have a teacher "
"log-probability of exactly 0.0. This is the signature of teacher "
"log-probs written by RetrieveLogProb(assistant_only=True), which "
"zero-fills non-assistant positions, while this loss was constructed "
"with assistant_only=False and therefore includes those positions. "
"Either construct the loss with assistant_only=True, or recompute the "
"teacher log-probs with RetrieveLogProb(assistant_only=False)."
)
[docs]
def forward(self, tensordict: TensorDictBase) -> DistillationLossOutput:
history: History = tensordict[self.tensor_keys.history]
masks, attention_masks, assistant_masks = self._get_masks(tensordict, history)
token_dim = tensordict.ndim - 1
if any((~mask.any(dim=token_dim)).any() for mask in masks):
raise ValueError(
"Some sequences have no tokens selected for distillation. "
"Check the assistant/attention masks of the input."
)
input_loss = tensordict.select(self.tensor_keys.history)
device = _runtime_device(
self.actor_network,
tensordict,
masks,
attention_masks,
fallback=self._device_fallback,
)
with torch.device(device) if device is not None else contextlib.nullcontext():
output_loss = self.actor_network(input_loss)
log_probs = output_loss.get(self.tensor_keys.log_probs, as_list=True)
teacher_log_probs = tensordict.get(
self.tensor_keys.teacher_log_prob, default=None, as_list=True
)
if teacher_log_probs is None:
raise KeyError(
f"Teacher log-probs not found at {self.tensor_keys.teacher_log_prob} in the "
f"input tensordict (keys: {set(tensordict.keys(include_nested=True, leaves_only=True))}). "
"Use RetrieveLogProb(teacher, log_probs_full_key=('teacher_log_probs', 'full')) "
"to compute them, or point the teacher_log_prob key to your data with set_keys()."
)
if not all(
lp.shape == tlp.shape
for lp, tlp in _zip_strict(log_probs, teacher_log_probs)
):
raise ValueError(
f"Student and teacher log-probs have different shapes: "
f"{[lp.shape for lp in log_probs]} vs {[tlp.shape for tlp in teacher_log_probs]}. "
"Both models must score the same tokenized sequences."
)
if not all(
mask.shape == lp.shape for mask, lp in _zip_strict(masks, log_probs)
):
raise ValueError(
f"Masks and log-probs have different shapes: "
f"{[mask.shape for mask in masks]} vs {[lp.shape for lp in log_probs]}."
)
if not self.assistant_only and assistant_masks is not None:
self._check_teacher_log_probs_consistency(
teacher_log_probs, attention_masks, assistant_masks
)
kl_tokens = []
for lp, tlp, mask in _zip_strict(log_probs, teacher_log_probs, masks):
lp = lp.masked_fill(~mask, 0.0)
tlp = tlp.detach().masked_fill(~mask, 0.0)
if self.kl_direction == "reverse":
kl = k3_kl_token_estimate(tlp, lp)
else:
kl = k3_kl_token_estimate(lp, tlp)
kl_tokens.append(kl)
summed_kl = torch.stack([kl.sum(token_dim) for kl in kl_tokens])
if self.normalize_by_seq_length:
seq_lengths = torch.stack([mask.sum(token_dim) for mask in masks])
summed_kl = summed_kl / seq_lengths.clamp(min=1)
loss = _distillation_loss(summed_kl, self.reduction)
return DistillationLossOutput(
loss_distill=loss, kl_to_teacher=summed_kl.detach().mean()
)
