diff --git a/docs/source/reference/llms.rst b/docs/source/reference/llms.rst index 05eecbb2e57..c1178f911cf 100644 --- a/docs/source/reference/llms.rst +++ b/docs/source/reference/llms.rst @@ -15,7 +15,7 @@ Key Components 2. **LLM Wrappers**: Unified interfaces for Transformers, vLLM, SGLang, and async variants 3. **Environments**: ChatEnv, task-specific environments, and transforms 4. **Collectors**: LLMCollector and RayLLMCollector for data collection -5. **Objectives**: GRPOLoss, SFTLoss for training +5. **Objectives**: GRPOLoss, SFTLoss, DistillationLoss for training Quick Example ------------- @@ -544,6 +544,20 @@ SFT SFTLoss SFTLossOutput +Distillation +~~~~~~~~~~~~ + +.. currentmodule:: torchrl.objectives.llm + +.. autosummary:: + :toctree: generated/ + :template: rl_template.rst + + DistillationLoss + DistillationLossOutput + distillation_loss + reverse_kl_token_estimate + .. currentmodule:: torchrl.data.llm .. autosummary:: diff --git a/docs/source/reference/llms_objectives.rst b/docs/source/reference/llms_objectives.rst index e93c4d1e04c..e0d90f9f363 100644 --- a/docs/source/reference/llms_objectives.rst +++ b/docs/source/reference/llms_objectives.rst @@ -29,3 +29,15 @@ SFT SFTLoss SFTLossOutput + +Distillation +------------ + +.. autosummary:: + :toctree: generated/ + :template: rl_template.rst + + DistillationLoss + DistillationLossOutput + distillation_loss + reverse_kl_token_estimate diff --git a/test/llm/test_llm_objectives.py b/test/llm/test_llm_objectives.py index 125cea91102..6589fb9bee3 100644 --- a/test/llm/test_llm_objectives.py +++ b/test/llm/test_llm_objectives.py @@ -20,6 +20,11 @@ from torchrl.envs.llm.transforms.kl import RetrieveLogProb from torchrl.modules.llm import TransformersWrapper, vLLMWrapper from torchrl.modules.llm.policies.common import ChatHistory, Masks, Text, Tokens +from torchrl.objectives.llm.distillation import ( + distillation_loss, + DistillationLoss, + reverse_kl_token_estimate, +) from torchrl.objectives.llm.grpo import ( CISPOLoss, CISPOLossOutput, @@ -1006,6 +1011,313 @@ def test_sft_assistant_only(self, data): loss(td) +class TestDistillation: + """Tests for the on-policy reverse-KL distillation loss.""" + + class _StaticLogProbActor(torch.nn.Module): + def __init__(self, log_probs): + super().__init__() + self.log_probs = torch.nn.Parameter(log_probs) + + def forward(self, tensordict): + return TensorDict( + {("log_probs", "full"): self.log_probs}, + batch_size=tensordict.batch_size, + ) + + # ------------------------------------------------------------------ + # Pure-math unit tests: no model/tokenizer required, fast, CPU-only. + # These pin the KL estimator's contract before any LLM plumbing. + # ------------------------------------------------------------------ + def test_reverse_kl_zero_when_equal(self): + lp = torch.randn(16) + kl = reverse_kl_token_estimate(lp, lp.clone()) + torch.testing.assert_close(kl, torch.zeros_like(kl)) + + def test_reverse_kl_nonnegative(self): + torch.manual_seed(0) + student = torch.randn(2000) + teacher = torch.randn(2000) + kl = reverse_kl_token_estimate(student, teacher) + assert (kl >= -1e-6).all() + + def test_reverse_kl_matches_closed_form(self): + student = torch.tensor([-1.0, -2.0, -0.3]) + teacher = torch.tensor([-0.5, -2.5, -0.1]) + diff = teacher - student + expected = diff.expm1() - diff + torch.testing.assert_close( + reverse_kl_token_estimate(student, teacher), expected + ) + + def test_reverse_kl_shape_mismatch_raises(self): + with pytest.raises(ValueError, match="different shapes"): + reverse_kl_token_estimate(torch.randn(3), torch.randn(4)) + + @pytest.mark.parametrize("reduction", ["mean", "sum", "none"]) + def test_distillation_loss_reduction(self, reduction): + x = torch.tensor([1.0, 2.0, 3.0]) + out = distillation_loss(x, reduction) + if reduction == "mean": + torch.testing.assert_close(out, x.mean()) + elif reduction == "sum": + torch.testing.assert_close(out, x.sum()) + else: + torch.testing.assert_close(out, x) + + def test_distillation_loss_invalid_reduction(self): + with pytest.raises(ValueError, match="Invalid reduction"): + distillation_loss(torch.randn(3), "not-a-reduction") + + def test_missing_assistant_mask_in_one_sample_raises(self): + loss = DistillationLoss(actor_network=torch.nn.Identity(), tokenizer=object()) + assistant_mask = torch.tensor([[True, False, False], [False, False, False]]) + td = TensorDict( + { + ("masks", "all_assistant_mask"): assistant_mask, + ("masks", "all_attention_mask"): torch.ones_like(assistant_mask), + }, + batch_size=(2,), + ) + with pytest.raises(ValueError, match="Some inputs have no valid assistant"): + loss._get_assistant_masks(td, history=None) + + def test_teacher_log_probs_are_detached(self): + actor = self._StaticLogProbActor( + torch.tensor([[-1.0, -2.0, -3.0], [-0.7, -1.5, -2.4]]) + ) + teacher_log_probs = torch.tensor( + [[-0.5, -2.2, -2.8], [-1.0, -1.1, -2.0]], requires_grad=True + ) + mask = torch.ones(2, 3, dtype=torch.bool) + td = TensorDict( + { + ("history", "full"): torch.zeros(2, 1), + ("masks", "all_assistant_mask"): mask, + ("masks", "all_attention_mask"): mask, + ("next", "teacher_log_probs", "full"): teacher_log_probs, + }, + batch_size=(2,), + ) + loss = DistillationLoss(actor_network=actor, tokenizer=object()) + + loss(td).sum(reduce=True).backward() + + assert actor.log_probs.grad is not None + assert (actor.log_probs.grad != 0).any() + assert teacher_log_probs.grad is None + + def test_gradient_descends_toward_teacher(self): + # A gradient step on the student log-probs must reduce the reverse KL, + # confirming the loss sign pulls the student toward the teacher (and not + # away from it). This is the empirical check on the KL direction. + torch.manual_seed(0) + teacher = torch.randn(64) + student = torch.randn(64, requires_grad=True) + kl_before = reverse_kl_token_estimate(student, teacher).sum() + kl_before.backward() + with torch.no_grad(): + stepped = student - 0.1 * student.grad + kl_after = reverse_kl_token_estimate(stepped, teacher).sum() + assert kl_after < kl_before + + # ------------------------------------------------------------------ + # Integration tests: mirror TestSFT. Fixtures are function-scoped on + # purpose so the teacher transform (which mutates the tensordict in + # place) cannot leak teacher log-probs across tests. + # ------------------------------------------------------------------ + @pytest.fixture + def data(self): + from transformers import AutoTokenizer + + 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": "Hello, how are you?"}, + { + "role": "assistant", + "content": "I'm doing well, thank you very much!", + }, + ], + ] + history = History.from_chats(chats) + tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m") + tokenizer.pad_token = tokenizer.eos_token + text = history[:, :-1].apply_chat_template( + tokenizer=tokenizer, chat_template_name="qwen", add_generation_prompt=True + ) + full_text = history.apply_chat_template( + tokenizer=tokenizer, chat_template_name="qwen", add_generation_prompt=False + ) + text_response = [ + txt[len(txt_start) :] for txt, txt_start in zip(full_text, text) + ] + td = TensorDict( + text=Text(prompt=text, response=text_response, full=full_text), + history=ChatHistory( + full=history, prompt=history[..., :-1], response=history[..., -1:] + ), + next=TensorDict( + reward=torch.randn(2, 1), + done=torch.zeros(2, dtype=torch.bool), + history=ChatHistory(prompt=history), + ), + batch_size=(2,), + ) + return lazy_stack(list(td.unbind(0))) + + @staticmethod + def _tokenizer(): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained("facebook/opt-125m") + + @pytest.fixture + def policy_train(self): + from transformers import OPTConfig, OPTForCausalLM + + tokenizer = self._tokenizer() + model = OPTForCausalLM(OPTConfig()).eval() + policy_train = TransformersWrapper( + model, + tokenizer=tokenizer, + generate=False, + chat_template_name="qwen", + input_mode="history", + pad_output=False, + ) + return policy_train, tokenizer + + def _teacher_transform(self, model, tokenizer): + teacher = TransformersWrapper( + model, + tokenizer=tokenizer, + generate=False, + return_log_probs=True, + chat_template_name="qwen", + input_mode="history", + pad_output=False, + ) + return RetrieveLogProb( + teacher, + assistant_only=True, + tokenizer_kwargs={"chat_template_name": "qwen"}, + tokenizer=tokenizer, + log_probs_full_key=("teacher_log_probs", "full"), + ) + + @pytest.mark.skipif( + not _has_transformers, reason="transformers lib required to test distillation" + ) + @pytest.mark.parametrize("reduction", ["mean", "sum", "none"]) + @pytest.mark.parametrize("normalize_by_seq_length", [True, False]) + def test_distillation_forward( + self, reduction, normalize_by_seq_length, data, policy_train + ): + from transformers import OPTConfig, OPTForCausalLM + + policy_train, tokenizer = policy_train + # A separate teacher with different weights: KL > 0 and gradient flows. + teacher_model = OPTForCausalLM(OPTConfig()).eval() + transform = self._teacher_transform(teacher_model, tokenizer) + with torch.no_grad(): + transform(data) + + loss = DistillationLoss( + actor_network=policy_train, + tokenizer=tokenizer, + reduction=reduction, + normalize_by_seq_length=normalize_by_seq_length, + tokenizer_kwargs={"chat_template_name": "qwen"}, + ) + out = loss(data) + if reduction in ("mean", "sum"): + assert out.loss_distill.shape == () + else: + assert out.loss_distill.shape == (2,) + assert out.kl_to_teacher.shape == () + assert (out.kl_to_teacher >= -1e-4).all() + assert out.sum(reduce=True).shape == () + + # Gradient reaches the student, and not the detached teacher signal. + out.sum(reduce=True).backward() + grads = [p.grad for p in policy_train.model.parameters() if p.grad is not None] + assert len(grads) > 0 + assert any((g != 0).any() for g in grads) + + @pytest.mark.skipif( + not _has_transformers, reason="transformers lib required to test distillation" + ) + def test_distillation_zero_when_teacher_matches_student(self, data, policy_train): + policy_train, tokenizer = policy_train + # Teacher equals student: the reverse KL is (approximately) zero. + transform = self._teacher_transform(policy_train.model, tokenizer) + with torch.no_grad(): + transform(data) + loss = DistillationLoss( + actor_network=policy_train, + tokenizer=tokenizer, + tokenizer_kwargs={"chat_template_name": "qwen"}, + ) + out = loss(data) + torch.testing.assert_close( + out.kl_to_teacher, + torch.zeros((), dtype=out.kl_to_teacher.dtype), + atol=1e-4, + rtol=0, + ) + + @pytest.mark.skipif( + not _has_transformers, reason="transformers lib required to test distillation" + ) + def test_distillation_missing_teacher_log_probs_raises(self, data, policy_train): + policy_train, tokenizer = policy_train + loss = DistillationLoss( + actor_network=policy_train, + tokenizer=tokenizer, + tokenizer_kwargs={"chat_template_name": "qwen"}, + ) + with pytest.raises(ValueError, match="Teacher log-probs not found"): + loss(data) + + @pytest.mark.skipif( + not _has_transformers, reason="transformers lib required to test distillation" + ) + def test_distillation_set_keys_nested(self, data, policy_train): + policy_train, tokenizer = policy_train + teacher = TransformersWrapper( + policy_train.model, + tokenizer=tokenizer, + generate=False, + return_log_probs=True, + chat_template_name="qwen", + input_mode="history", + pad_output=False, + ) + transform = RetrieveLogProb( + teacher, + assistant_only=True, + tokenizer_kwargs={"chat_template_name": "qwen"}, + tokenizer=tokenizer, + log_probs_full_key=("a_teacher", "logp", "full"), + ) + with torch.no_grad(): + transform(data) + loss = DistillationLoss( + actor_network=policy_train, + tokenizer=tokenizer, + tokenizer_kwargs={"chat_template_name": "qwen"}, + ) + loss.set_keys(teacher_log_prob=("next", "a_teacher", "logp", "full")) + out = loss(data) + assert out.loss_distill.shape == () + + @pytest.mark.slow @pytest.mark.integration class TestGRPOLossIntegration: diff --git a/torchrl/objectives/llm/__init__.py b/torchrl/objectives/llm/__init__.py index 0170d0399c6..232175ac043 100644 --- a/torchrl/objectives/llm/__init__.py +++ b/torchrl/objectives/llm/__init__.py @@ -4,6 +4,12 @@ # LICENSE file in the root directory of this source tree. from __future__ import annotations +from .distillation import ( + distillation_loss, + DistillationLoss, + DistillationLossOutput, + reverse_kl_token_estimate, +) from .grpo import ( CISPOLoss, CISPOLossOutput, @@ -23,6 +29,8 @@ "CISPOLossOutput", "DAPO", "DAPOLossOutput", + "DistillationLoss", + "DistillationLossOutput", "GRPOLoss", "GRPOLossOutput", "LLMLossOutput", @@ -31,4 +39,6 @@ "RayMCAdvantage", "SFTLoss", "SFTLossOutput", + "distillation_loss", + "reverse_kl_token_estimate", ] diff --git a/torchrl/objectives/llm/distillation.py b/torchrl/objectives/llm/distillation.py new file mode 100644 index 00000000000..83196960aae --- /dev/null +++ b/torchrl/objectives/llm/distillation.py @@ -0,0 +1,410 @@ +# 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 + +if TYPE_CHECKING: + import transformers + + +def reverse_kl_token_estimate( + cur_log_prob: torch.Tensor, teacher_log_prob: torch.Tensor +) -> torch.Tensor: + r"""Per-token low-variance estimate of the reverse KL divergence to the teacher. + + Computes the per-token k3 estimator (Schulman, + `"Approximating KL Divergence" `_) of + :math:`\mathrm{KL}(\pi_\theta \,\|\, \pi_T)` using tokens drawn from the + student policy :math:`\pi_\theta`: + + .. math:: + \widehat{\mathrm{kl}}_t = e^{\Delta_t} - 1 - \Delta_t, + \qquad \Delta_t = \log \pi_T(a_t) - \log \pi_\theta(a_t) + + The estimator is non-negative and equals zero exactly when the two log-probs + match. With ``a_t`` sampled from the student, summing over a sequence + estimates the sequence-level reverse KL :math:`\mathrm{KL}(\pi_\theta\,\|\,\pi_T)`. + + Args: + cur_log_prob (torch.Tensor): student log-probabilities of the sampled tokens. + teacher_log_prob (torch.Tensor): teacher log-probabilities of the same tokens. + + Returns: + torch.Tensor: the per-token KL estimate, with the same shape as the inputs. + """ + if cur_log_prob.shape != teacher_log_prob.shape: + raise ValueError( + "Student and teacher log-probabilities have different shapes: " + f"{cur_log_prob.shape=} vs {teacher_log_prob.shape=}." + ) + diff = teacher_log_prob - cur_log_prob + return diff.expm1() - diff + + +def distillation_loss( + per_sequence_kl: torch.Tensor, reduction: Literal["mean", "sum", "none"] +) -> torch.Tensor: + """Reduce the per-sequence distillation KL into the final loss value. + + Args: + per_sequence_kl (torch.Tensor): a one-dimensional tensor holding one + (already token-reduced) KL value per sequence in the batch. + reduction (Literal["mean", "sum", "none"]): the batch reduction to apply. + + Returns: + torch.Tensor: the reduced loss (a scalar for ``"mean"``/``"sum"``, the + per-sequence tensor for ``"none"``). + """ + if reduction == "mean": + return per_sequence_kl.mean() + if reduction == "sum": + return per_sequence_kl.sum() + if reduction == "none": + return per_sequence_kl + raise ValueError(f"Invalid reduction: {reduction}.") + + +class DistillationLossOutput(TensorClass["nocast"]): + """Output of :class:`~torchrl.objectives.llm.DistillationLoss`. + + Attributes: + loss_distill (torch.Tensor): the coefficient-scaled on-policy reverse-KL + distillation loss. This is the only differentiable field. + kl_to_teacher (torch.Tensor | None): the unscaled mean per-token KL to the + teacher, detached, for logging. + + .. note:: + Only ``loss_distill`` carries a gradient; ``kl_to_teacher`` is detached. + Summing the output (``loss_fn(td).sum(reduce=True)``) therefore yields the + total differentiable loss, mirroring + :class:`~torchrl.objectives.llm.SFTLossOutput`. + """ + + loss_distill: torch.Tensor + kl_to_teacher: torch.Tensor | None = None + + +class DistillationLoss(LossModule): + r"""On-policy knowledge-distillation loss for LLM policies. + + Trains a student policy to imitate a (typically larger or stronger) teacher + by minimizing the reverse KL divergence + :math:`\mathrm{KL}(\pi_\theta \,\|\, \pi_T)` on the student's own generations. + Because the expectation is taken over the student's distribution, this is the + *on-policy* distillation regime: the student generates the sequences and the + teacher scores those sequences. The teacher's per-token log-probabilities are + expected to have been written to the input tensordict by an upstream + :class:`~torchrl.envs.llm.transforms.RetrieveLogProb` transform pointed at the + teacher model, exactly as reference log-probabilities are produced for + :class:`~torchrl.objectives.llm.SFTLoss` and + :class:`~torchrl.objectives.llm.GRPOLoss`. + + The divergence is estimated per token with the low-variance k3 estimator (see + :func:`~torchrl.objectives.llm.distillation.reverse_kl_token_estimate`), + restricted to assistant (response) tokens, then reduced per sequence and over + the batch. + + .. note:: + This is the *reverse-KL, token-log-probability* form of distillation. It + is a single-sample estimate of the sequence-level reverse KL on the + student's own samples, and its gradient is taken through the student's + re-scored log-probabilities only (the dependence of the sampling + distribution on the parameters is not differentiated, as is standard for + on-policy distillation surrogates). The exact per-token KL and the + forward-KL / generalized-JSD variants require the teacher's full + vocabulary distribution and are not covered here. + + Args: + actor_network (TensorDictModule): the student network. Usually a + :class:`~torchrl.modules.llm.TransformersWrapper` with + ``generate=False`` so that its log-probabilities are differentiable. + tokenizer (`Tokenizer`, optional): the tokenizer used to recompute the + assistant mask from the chat history when it is not already present in + the input tensordict. If not provided, it is taken from + ``actor_network``. + tokenizer_kwargs (dict, optional): keyword arguments forwarded to + :meth:`~torchrl.data.llm.chat.History.apply_chat_template`. + + Keyword Args: + distill_coeff (float, optional): scaling applied to the distillation loss. + Defaults to ``1.0``. + reduction (Literal["mean", "sum", "none"], optional): the batch reduction. + Defaults to ``"mean"``. + normalize_by_seq_length (bool, optional): whether to normalize the + per-sequence KL by the number of assistant tokens. Defaults to ``True``. + device (torch.device, optional): the device used when tokenizing the + input. Defaults to ``None``. + + .. note:: + The input tensordict is expected to contain the following keys by default: + - ``("history", "full")``: the chat history (used to run the student + and to recompute the assistant mask if needed). + - ``("next", "teacher_log_probs", "full")``: the teacher per-token + log-probabilities, produced upstream by + :class:`~torchrl.envs.llm.transforms.RetrieveLogProb`. + + These keys can be customized using the ``set_keys()`` method, e.g. + ``loss.set_keys(teacher_log_prob=("next", "my_teacher", "full"))``. + + .. seealso:: :class:`~torchrl.envs.llm.transforms.RetrieveLogProb` for the + teacher log-probability computation, and + :class:`~torchrl.objectives.llm.SFTLoss` for the supervised-fine-tuning + loss this module mirrors. + + References: + - Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, et al., 2023. + `"On-Policy Distillation of Language Models: Learning from Self-Generated + Mistakes" (GKD) `_ + - Yoon Kim, Alexander M. Rush, 2016. + `"Sequence-Level Knowledge Distillation" `_ + + Examples: + >>> from torchrl.data.llm.chat import History, _CHAT_TEMPLATES + >>> from torchrl.envs.llm.transforms import RetrieveLogProb + >>> from torchrl.modules.llm import TransformersWrapper + >>> 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!"}, + ... ], + ... ] + >>> 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"] + >>> # A small student and a separate (frozen) teacher. + >>> student = TransformersWrapper( + ... OPTForCausalLM(OPTConfig()).eval(), + ... tokenizer=tokenizer, + ... generate=False, + ... chat_template_name="qwen", + ... ) + >>> teacher = TransformersWrapper( + ... OPTForCausalLM(OPTConfig()).eval(), + ... tokenizer=tokenizer, + ... generate=False, + ... return_log_probs=True, + ... chat_template_name="qwen", + ... ) + >>> # Score the student's tokens under the teacher. + >>> transform = RetrieveLogProb( + ... teacher, + ... assistant_only=True, + ... tokenizer_kwargs={"chat_template_name": "qwen"}, + ... tokenizer=tokenizer, + ... log_probs_full_key=("teacher_log_probs", "full"), + ... ) + >>> text = history[:, :-1].apply_chat_template( + ... tokenizer=tokenizer, chat_template_name="qwen", add_generation_prompt=True + ... ) + >>> text_response = history.apply_chat_template( + ... tokenizer=tokenizer, chat_template_name="qwen", add_generation_prompt=False + ... ) + >>> text_response = [ + ... txt[len(txt_start):] for txt, txt_start in zip(text_response, text) + ... ] + >>> td = TensorDict( + ... text=text, + ... text_response=text_response, + ... history=history, + ... next=TensorDict( + ... reward=torch.randn(1, 1), + ... done=torch.zeros(1, dtype=torch.bool), + ... history=history, + ... ), + ... batch_size=(1,), + ... ) + >>> data = lazy_stack(list(td.unbind(0))) + >>> with torch.no_grad(): + ... data = transform(data) + >>> loss = DistillationLoss( + ... actor_network=student, + ... tokenizer=tokenizer, + ... distill_coeff=1.0, + ... tokenizer_kwargs={"chat_template_name": "qwen"}, + ... ) + >>> loss_vals = loss(data) + >>> print(f"distillation loss: {loss_vals.loss_distill.item():.4f}") + >>> loss_vals.sum(reduce=True).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 (produced upstream by + :class:`~torchrl.envs.llm.transforms.RetrieveLogProb`). Defaults to + ``("next", "teacher_log_probs", "full")``. + log_probs (NestedKey): The key under which the student log-probabilities + are read from the actor output. 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, + *, + distill_coeff: float = 1.0, + reduction: Literal["mean", "sum", "none"] = "mean", + normalize_by_seq_length: 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) + if tokenizer is None: + raise ValueError("Tokenizer must be provided.") + 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 + self.distill_coeff = distill_coeff + self.reduction = reduction + self.normalize_by_seq_length = normalize_by_seq_length + self.device = device + self._set_in_keys() + + 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 = [] # Loss modules typically don't have out_keys + + def _get_assistant_masks( + self, tensordict: TensorDictBase, history: History + ) -> tuple[list[torch.Tensor], object]: + """Return per-sequence assistant masks, recomputing them if necessary. + + Mirrors :class:`~torchrl.objectives.llm.SFTLoss`: prefer the masks already + stored on the tensordict, otherwise recompute them from the chat history. + """ + token_struct = None + assistant_masks = tensordict.get(("masks", "all_assistant_mask"), as_list=True) + attention_mask = tensordict.get(("masks", "all_attention_mask"), as_list=True) + if assistant_masks is None: + with torch.device( + self.device + ) if self.device is not None else contextlib.nullcontext(): + token_struct = history.apply_chat_template( + tokenizer=self.tokenizer, **self.tokenizer_kwargs + ) + if "assistant_masks" not in token_struct: + raise ValueError( + f"Assistant masks are not present in the token structure: {token_struct=}." + ) + assistant_masks = token_struct.get("assistant_masks", as_list=True) + attention_mask = token_struct.get("attention_mask", as_list=True) + assistant_masks = [mask.bool() for mask in assistant_masks] + attention_mask = [mask.bool() for mask in attention_mask] + assistant_masks = [ + mask & a_mask for mask, a_mask in zip(assistant_masks, attention_mask) + ] + if not all(mask.any(-1).all() for mask in assistant_masks): + raise ValueError("Some inputs have no valid assistant masks.") + return assistant_masks, token_struct + + def forward(self, tensordict: TensorDictBase) -> DistillationLossOutput: + history: History = tensordict[self.tensor_keys.history] + assistant_masks, token_struct = self._get_assistant_masks(tensordict, history) + + # Re-run the student so its log-probabilities are differentiable. + input_loss = tensordict.select(self.tensor_keys.history) + with torch.device( + self.device + ) if self.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-probabilities are computed upstream and arrive detached. + teacher_log_probs = tensordict.get( + self.tensor_keys.teacher_log_prob, + default=None, + as_list=True, + ) + if teacher_log_probs is None: + raise ValueError( + f"Teacher log-probs not found in tensordict at key " + f"{self.tensor_keys.teacher_log_prob}. Did you run a RetrieveLogProb " + f"transform pointed at the teacher? Existing keys: " + f"{set(tensordict.keys(include_nested=True, leaves_only=True))}" + ) + + if not all( + mask.shape == lp.shape + for mask, lp in _zip_strict(assistant_masks, log_probs) + ): + if token_struct is not None: + suffix = ( + "Tokens from current template: " + f"{[inp.shape for inp in token_struct.get('input_ids', as_padded_tensor=True)]}" + ) + else: + suffix = "" + raise ValueError( + f"Assistant masks and log_probs have different shapes: " + f"{[mask.shape for mask in assistant_masks]} vs " + f"{[lp.shape for lp in log_probs]}. {suffix}" + ) + + reduce_dim = tensordict.ndim - 1 + per_sequence_kl = [] + per_token_kl_mean = [] + for student_lp, teacher_lp, mask in _zip_strict( + log_probs, teacher_log_probs, assistant_masks + ): + teacher_lp = teacher_lp.detach().to(student_lp.device) + mask = mask.to(student_lp.device) + kl_token = reverse_kl_token_estimate(student_lp, teacher_lp) + kl_token = kl_token.masked_fill(~mask, 0.0) + summed = kl_token.sum(reduce_dim) + count = mask.sum(reduce_dim).clamp(min=1) + per_token_kl_mean.append(summed / count) + per_sequence_kl.append( + summed / count if self.normalize_by_seq_length else summed + ) + + per_sequence_kl = torch.stack(per_sequence_kl) + loss = self.distill_coeff * distillation_loss(per_sequence_kl, self.reduction) + kl_to_teacher = torch.stack(per_token_kl_mean).mean().detach() + return DistillationLossOutput(loss_distill=loss, kl_to_teacher=kl_to_teacher)