diff --git a/docs/source/reference/llms.rst b/docs/source/reference/llms.rst index 05eecbb2e57..4dc15f702ec 100644 --- a/docs/source/reference/llms.rst +++ b/docs/source/reference/llms.rst @@ -544,6 +544,19 @@ SFT SFTLoss SFTLossOutput +Reward Model Training +~~~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: torchrl.objectives.llm + +.. autosummary:: + :toctree: generated/ + :template: rl_template.rst + + RewardModelLoss + RewardModelLossOutput + reward_model_loss + .. currentmodule:: torchrl.data.llm .. autosummary:: diff --git a/docs/source/reference/llms_objectives.rst b/docs/source/reference/llms_objectives.rst index e93c4d1e04c..b0a1e35c416 100644 --- a/docs/source/reference/llms_objectives.rst +++ b/docs/source/reference/llms_objectives.rst @@ -29,3 +29,14 @@ SFT SFTLoss SFTLossOutput + +Reward Model Training +--------------------- + +.. autosummary:: + :toctree: generated/ + :template: rl_template.rst + + RewardModelLoss + RewardModelLossOutput + reward_model_loss diff --git a/test/llm/test_llm_objectives.py b/test/llm/test_llm_objectives.py index 125cea91102..6a0a8ce5b64 100644 --- a/test/llm/test_llm_objectives.py +++ b/test/llm/test_llm_objectives.py @@ -14,6 +14,7 @@ import torch from tensordict import lazy_stack, MetaData, TensorDict +from tensordict.nn import TensorDictModule from torchrl._utils import logger from torchrl.data import History, LazyStackStorage, ReplayBuffer from torchrl.data.llm.history import _CHAT_TEMPLATES @@ -29,6 +30,11 @@ MCAdvantageSelector, RayMCAdvantage, ) +from torchrl.objectives.llm.reward import ( + reward_model_loss, + RewardModelLoss, + RewardModelLossOutput, +) from torchrl.objectives.llm.sft import SFTLoss _has_transformers = importlib.util.find_spec("transformers") is not None @@ -1123,6 +1129,158 @@ def test_grpo_loss_with_real_models( assert torch.isfinite(result.loss_objective) +class _Scorer(torch.nn.Module): + """A tiny, dependency-free score network mapping input_ids -> scalar score.""" + + def __init__(self, vocab_size: int = 128, embed_dim: int = 8): + super().__init__() + self.embed = torch.nn.Embedding(vocab_size, embed_dim) + self.head = torch.nn.Linear(embed_dim, 1) + + def forward(self, input_ids): + return self.head(self.embed(input_ids).float().mean(-2)) + + +class TestRewardModel: + """Tests for the model-agnostic Bradley-Terry :class:`RewardModelLoss`.""" + + vocab_size = 128 + seq_len = 16 + batch_size = 4 + + def _score_network(self): + return TensorDictModule( + _Scorer(self.vocab_size), in_keys=["input_ids"], out_keys=["score"] + ) + + def _make_data(self, chosen_key="chosen", rejected_key="rejected"): + def _ids(): + return torch.randint(0, self.vocab_size, (self.batch_size, self.seq_len)) + + return TensorDict( + { + chosen_key: TensorDict(input_ids=_ids(), batch_size=[self.batch_size]), + rejected_key: TensorDict( + input_ids=_ids(), batch_size=[self.batch_size] + ), + }, + batch_size=[self.batch_size], + ) + + def test_reward_model_loss_fn_matches_formula(self): + """The bare loss helper matches -log_sigmoid(chosen - rejected).""" + chosen = torch.randn(8) + rejected = torch.randn(8) + expected = -torch.nn.functional.logsigmoid(chosen - rejected) + torch.testing.assert_close( + reward_model_loss(chosen, rejected, "none"), expected + ) + torch.testing.assert_close( + reward_model_loss(chosen, rejected, "mean"), expected.mean() + ) + torch.testing.assert_close( + reward_model_loss(chosen, rejected, "sum"), expected.sum() + ) + + def test_forward_output_type_and_backward(self): + loss_fn = RewardModelLoss(score_network=self._score_network()) + out = loss_fn(self._make_data()) + assert isinstance(out, RewardModelLossOutput) + assert out.loss_reward_model.shape == () + assert torch.isfinite(out.loss_reward_model) + assert out.loss_center is None + # accuracy is a detached metric in [0, 1] + assert 0.0 <= out.accuracy.item() <= 1.0 + assert not out.accuracy.requires_grad + # gradients flow to the score network + out.loss_reward_model.backward() + grads = [ + p.grad for p in loss_fn.score_network.parameters() if p.grad is not None + ] + assert grads, "expected gradients to flow into the score network" + + @pytest.mark.parametrize("reduction", ["mean", "sum", "none"]) + def test_reduction(self, reduction): + loss_fn = RewardModelLoss( + score_network=self._score_network(), reduction=reduction + ) + out = loss_fn(self._make_data()) + if reduction == "none": + assert out.loss_reward_model.shape == (self.batch_size,) + else: + assert out.loss_reward_model.shape == () + + def test_precomputed_scores_no_network(self): + """With score_network=None the scores are read directly from the inputs.""" + loss_fn = RewardModelLoss(score_network=None) + chosen = torch.randn(self.batch_size) + rejected = torch.randn(self.batch_size) + td = TensorDict( + { + "chosen": TensorDict(score=chosen, batch_size=[self.batch_size]), + "rejected": TensorDict(score=rejected, batch_size=[self.batch_size]), + }, + batch_size=[self.batch_size], + ) + out = loss_fn(td) + expected = -torch.nn.functional.logsigmoid(chosen - rejected).mean() + torch.testing.assert_close(out.loss_reward_model, expected) + expected_acc = (chosen > rejected).float().mean() + torch.testing.assert_close(out.accuracy, expected_acc) + + def test_center_coeff(self): + loss_fn = RewardModelLoss(score_network=self._score_network(), center_coeff=0.1) + out = loss_fn(self._make_data()) + assert out.loss_center is not None + assert out.loss_center.requires_grad + # total differentiable loss is the sum of the loss_ terms + (out.loss_reward_model + out.loss_center).backward() + + def test_nested_keys(self): + """Exercise NestedKey inputs via set_keys.""" + loss_fn = RewardModelLoss(score_network=self._score_network()) + loss_fn.set_keys(chosen=("data", "chosen"), rejected=("data", "rejected")) + inner = self._make_data() + td = TensorDict({"data": inner}, batch_size=[self.batch_size]) + out = loss_fn(td) + assert torch.isfinite(out.loss_reward_model) + assert ("data", "chosen") in loss_fn.in_keys + + def test_nested_score_key(self): + """Exercise NestedKey score lookup via set_keys.""" + loss_fn = RewardModelLoss(score_network=None) + loss_fn.set_keys(score=("metrics", "score")) + chosen = torch.randn(self.batch_size) + rejected = torch.randn(self.batch_size) + td = TensorDict( + { + "chosen": TensorDict( + {"metrics": TensorDict(score=chosen, batch_size=[self.batch_size])}, + batch_size=[self.batch_size], + ), + "rejected": TensorDict( + { + "metrics": TensorDict( + score=rejected, batch_size=[self.batch_size] + ) + }, + batch_size=[self.batch_size], + ), + }, + batch_size=[self.batch_size], + ) + out = loss_fn(td) + expected = -torch.nn.functional.logsigmoid(chosen - rejected).mean() + torch.testing.assert_close(out.loss_reward_model, expected) + + def test_missing_key_raises(self): + loss_fn = RewardModelLoss(score_network=self._score_network()) + td = self._make_data() + del td["rejected"] + with pytest.raises(KeyError): + loss_fn(td) + + if __name__ == "__main__": args, unknown = argparse.ArgumentParser().parse_known_args() pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown) diff --git a/torchrl/objectives/llm/__init__.py b/torchrl/objectives/llm/__init__.py index 0170d0399c6..a268e97ff92 100644 --- a/torchrl/objectives/llm/__init__.py +++ b/torchrl/objectives/llm/__init__.py @@ -16,6 +16,7 @@ MCAdvantageSelector, RayMCAdvantage, ) +from .reward import reward_model_loss, RewardModelLoss, RewardModelLossOutput from .sft import SFTLoss, SFTLossOutput __all__ = [ @@ -29,6 +30,9 @@ "MCAdvantage", "MCAdvantageSelector", "RayMCAdvantage", + "RewardModelLoss", + "RewardModelLossOutput", + "reward_model_loss", "SFTLoss", "SFTLossOutput", ] diff --git a/torchrl/objectives/llm/reward.py b/torchrl/objectives/llm/reward.py new file mode 100644 index 00000000000..a27dd38d0ac --- /dev/null +++ b/torchrl/objectives/llm/reward.py @@ -0,0 +1,276 @@ +# 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 + +import torch +from tensordict import NestedKey, TensorClass, TensorDictBase +from tensordict.nn import TensorDictModule +from torchrl.objectives.common import LossModule + + +def reward_model_loss( + chosen_scores: torch.Tensor, + rejected_scores: torch.Tensor, + reduction: Literal["mean", "sum", "none"], +) -> torch.Tensor: + r"""Compute the Bradley-Terry pairwise reward-model loss. + + The loss is computed as ``-log_sigmoid(chosen_scores - rejected_scores)``. It is + small when the reward model assigns a higher score to the chosen response than to + the rejected one, and large otherwise. + + .. math:: + + \text{loss} = -\log\sigma(r_\theta(x, y_c) - r_\theta(x, y_r)) + + Args: + chosen_scores (torch.Tensor): the scalar scores assigned to the chosen + responses. Must have shape ``[B]``. + rejected_scores (torch.Tensor): the scalar scores assigned to the rejected + responses. Must have shape ``[B]``. + reduction (Literal["mean", "sum", "none"]): the reduction to apply to the loss. + + Returns: + The Bradley-Terry loss. + + References: + - Ralph Allan Bradley, Milton E. Terry, 1952. "Rank Analysis of Incomplete Block + Designs: I. The Method of Paired Comparisons". + """ + loss = -torch.nn.functional.logsigmoid(chosen_scores - rejected_scores) + if reduction == "mean": + return loss.mean() + if reduction == "sum": + return loss.sum() + if reduction == "none": + return loss + raise ValueError(f"Invalid reduction: {reduction}.") + + +class RewardModelLossOutput(TensorClass["nocast"]): + """Reward-model loss output. + + Attributes: + loss_reward_model (torch.Tensor): the Bradley-Terry pairwise loss. + loss_center (torch.Tensor | None): the score-centering regularization loss. + Only present when ``center_coeff`` is set on the loss. Defaults to ``None``. + accuracy (torch.Tensor | None): the fraction of pairs for which the chosen + score exceeds the rejected score. This is a detached metric for logging + and is not differentiable. Defaults to ``None``. + + .. note:: + The differentiable total loss is the sum of the ``loss_`` prefixed fields, + i.e. ``loss_reward_model`` (plus ``loss_center`` when set). ``accuracy`` is a + detached diagnostic and should not be backpropagated: + + >>> loss_fn = RewardModelLoss(score_network) + >>> loss_output = loss_fn(td) + >>> loss = loss_output.loss_reward_model + >>> if loss_output.loss_center is not None: + ... loss = loss + loss_output.loss_center + >>> loss.backward() + """ + + loss_reward_model: torch.Tensor + loss_center: torch.Tensor | None = None + accuracy: torch.Tensor | None = None + + +class RewardModelLoss(LossModule): + r"""Bradley-Terry reward-model training loss for RLHF. + + Trains a scalar reward model from pairwise human preference data. Given a prompt and + two responses (a ``chosen`` one preferred by an annotator and a ``rejected`` one), + the model is encouraged to assign a higher score to the chosen response. This is the + reward-modelling stage that precedes policy optimization in RLHF pipelines. + + The loss is model-agnostic: ``score_network`` can wrap any backbone that maps a + tokenized response to a single scalar score, for example a Hugging Face + ``AutoModelForSequenceClassification`` with ``num_labels=1`` or + :class:`~torchrl.modules.models.llm.GPT2RewardModel`. + + Args: + score_network (TensorDictModule, optional): a module mapping a (chosen or + rejected) sub-tensordict to a per-sequence scalar score written under the + ``score`` key. The same network is applied to the chosen and rejected + inputs (weight sharing). If ``None``, the chosen and rejected scores are + expected to already be present under the ``score`` key of their respective + sub-tensordicts. Defaults to ``None``. + + Keyword Args: + reduction (Literal["mean", "sum", "none"], optional): the reduction to apply to + the loss. Defaults to ``"mean"``. + center_coeff (float, optional): if set, adds a centering regularization term + ``center_coeff * (chosen_score**2 + rejected_score**2)`` that discourages + the reward model from drifting to large magnitudes. Defaults to ``None`` + (disabled). + device (torch.device or str, optional): if provided, chosen/rejected + sub-tensordicts are moved to this device before scoring and loss + computation. The ``score_network`` is expected to already be on this + device. Defaults to ``None``. + + .. note:: + The input tensordict is expected to contain the following keys by default: + - ``"chosen"``: a sub-tensordict with the chosen response inputs (e.g. + ``input_ids`` / ``attention_mask``), or the chosen ``score`` directly + when ``score_network`` is ``None``. + - ``"rejected"``: the corresponding sub-tensordict for the rejected + response. + + These keys (and the ``score`` output key) can be customized using the + :meth:`~torchrl.objectives.common.LossModule.set_keys` method. + + .. seealso:: :class:`~torchrl.modules.models.llm.GPT2RewardModel` for a ready-made + GPT2-based reward-model backbone, and + :class:`~torchrl.data.llm.reward.PairwiseDataset` for a pairwise preference + dataset. + + References: + - Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea + Voss, Alec Radford, Dario Amodei, Paul Christiano, 2020. + `"Learning to summarize from human feedback" `_ + - Long Ouyang et al., 2022. + `"Training language models to follow instructions with human feedback" `_ + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictModule + >>> from torchrl.objectives.llm.reward import RewardModelLoss + >>> + >>> # A toy score network mapping input_ids to a scalar score per sequence. + >>> class Scorer(torch.nn.Module): + ... def __init__(self, vocab_size=128, embed_dim=8): + ... super().__init__() + ... self.embed = torch.nn.Embedding(vocab_size, embed_dim) + ... self.head = torch.nn.Linear(embed_dim, 1) + ... + ... def forward(self, input_ids): + ... return self.head(self.embed(input_ids).mean(-2)) + >>> + >>> score_network = TensorDictModule( + ... Scorer(), in_keys=["input_ids"], out_keys=["score"] + ... ) + >>> loss_fn = RewardModelLoss(score_network=score_network) + >>> data = TensorDict( + ... chosen=TensorDict( + ... input_ids=torch.randint(0, 128, (4, 16)), batch_size=[4] + ... ), + ... rejected=TensorDict( + ... input_ids=torch.randint(0, 128, (4, 16)), batch_size=[4] + ... ), + ... batch_size=[4], + ... ) + >>> loss_vals = loss_fn(data) + >>> print(f"Reward model loss: {loss_vals.loss_reward_model.item():.4f}") + >>> print(f"Accuracy: {loss_vals.accuracy.item():.4f}") + >>> loss_vals.loss_reward_model.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: + chosen (NestedKey): The input tensordict key where the chosen response + sub-tensordict is expected. Defaults to ``"chosen"``. + rejected (NestedKey): The input tensordict key where the rejected response + sub-tensordict is expected. Defaults to ``"rejected"``. + score (NestedKey): The key (within each chosen/rejected sub-tensordict) + where the per-sequence scalar score is written by ``score_network`` or + read from when ``score_network`` is ``None``. Defaults to ``"score"``. + """ + + chosen: NestedKey = "chosen" + rejected: NestedKey = "rejected" + score: NestedKey = "score" + + default_keys = _AcceptedKeys + tensor_keys: _AcceptedKeys + + def __init__( + self, + score_network: TensorDictModule | None = None, + *, + reduction: Literal["mean", "sum", "none"] = "mean", + center_coeff: float | None = None, + device: torch.device | str | None = None, + ) -> None: + super().__init__() + self.score_network = score_network + self.reduction = reduction + self.center_coeff = center_coeff + self.device = torch.device(device) if device is not None else None + self._set_in_keys() + + def _set_in_keys(self) -> None: + """Sets the input keys for the loss module.""" + self.in_keys = [self.tensor_keys.chosen, self.tensor_keys.rejected] + self.out_keys = [] # Loss modules typically don't have out_keys + + def _score(self, tensordict: TensorDictBase, key: NestedKey) -> torch.Tensor: + """Run the score network on a sub-tensordict and return a per-sequence score.""" + sub_td = tensordict.get(key, default=None) + if sub_td is None: + raise KeyError( + f"Could not find the sub-tensordict at key {key!r} in the input " + f"tensordict with keys {set(tensordict.keys())}." + ) + if self.device is not None: + sub_td = sub_td.to(self.device) + if self.score_network is not None: + with ( + torch.device(self.device) + if self.device is not None + else contextlib.nullcontext() + ): + sub_td = self.score_network(sub_td) + score = sub_td.get(self.tensor_keys.score, default=None) + if score is None: + raise KeyError( + f"Could not find the score at key {self.tensor_keys.score!r} under " + f"{key!r}. If score_network is None, the scores must be precomputed." + ) + # reduce a trailing singleton dimension (e.g. [B, 1] -> [B]) + if score.ndim > 1 and score.shape[-1] == 1: + score = score.squeeze(-1) + return score + + def forward(self, tensordict: TensorDictBase) -> RewardModelLossOutput: + chosen_score = self._score(tensordict, self.tensor_keys.chosen) + rejected_score = self._score(tensordict, self.tensor_keys.rejected) + if chosen_score.shape != rejected_score.shape: + raise ValueError( + f"Chosen and rejected scores have different shapes: " + f"{chosen_score.shape=} vs {rejected_score.shape=}." + ) + + loss = reward_model_loss(chosen_score, rejected_score, self.reduction) + + loss_center = None + if self.center_coeff is not None: + center = chosen_score.pow(2) + rejected_score.pow(2) + if self.reduction == "mean": + center = center.mean() + elif self.reduction == "sum": + center = center.sum() + loss_center = self.center_coeff * center + + with torch.no_grad(): + accuracy = (chosen_score > rejected_score).float().mean() + + return RewardModelLossOutput( + loss_reward_model=loss, + loss_center=loss_center, + accuracy=accuracy, + )