diff --git a/.github/unittest/linux_sota/scripts/test_sota.py b/.github/unittest/linux_sota/scripts/test_sota.py index f169bd63757..c75f3161235 100644 --- a/.github/unittest/linux_sota/scripts/test_sota.py +++ b/.github/unittest/linux_sota/scripts/test_sota.py @@ -24,6 +24,19 @@ logger.eval_iter=2 \ logger.eval_episodes=4 \ checkpoint.save_iter=2 +""", + "reward_model_training": """python sota-implementations/reward_model_training/reward_model.py \ + model.name= \ + data.dataset_name= \ + data.synthetic_size=32 \ + data.batch_size=8 \ + data.max_length=32 \ + optim.max_iters=3 \ + logger.eval_iter=2 \ + logger.eval_iters=1 \ + logger.log_interval=1 \ + logger.backend= \ + checkpoint.save_iter=2 """, "diffusion_bc": """python sota-implementations/diffusion_bc/diffusion_bc.py \ optim.gradient_steps=55 \ diff --git a/docs/source/reference/llms.rst b/docs/source/reference/llms.rst index 05eecbb2e57..e0b640ebe37 100644 --- a/docs/source/reference/llms.rst +++ b/docs/source/reference/llms.rst @@ -544,6 +544,18 @@ SFT SFTLoss SFTLossOutput +Reward Model Training +~~~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: torchrl.objectives.llm + +.. autosummary:: + :toctree: generated/ + :template: rl_template.rst + + RewardModelLoss + RewardModelLossOutput + .. currentmodule:: torchrl.data.llm .. autosummary:: diff --git a/docs/source/reference/llms_objectives.rst b/docs/source/reference/llms_objectives.rst index e93c4d1e04c..33475759e4d 100644 --- a/docs/source/reference/llms_objectives.rst +++ b/docs/source/reference/llms_objectives.rst @@ -29,3 +29,13 @@ SFT SFTLoss SFTLossOutput + +Reward Model Training +--------------------- + +.. autosummary:: + :toctree: generated/ + :template: rl_template.rst + + RewardModelLoss + RewardModelLossOutput diff --git a/sota-check/README.md b/sota-check/README.md index 0c61daced87..f60dcbdb47c 100644 --- a/sota-check/README.md +++ b/sota-check/README.md @@ -26,7 +26,7 @@ export MUJOCO_GL=egl conda create -n rl-sota-bench python=3.10 -y conda install anaconda::libglu -y pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -pip3 install "gymnasium[atari,mujoco]" vmas tqdm wandb pygame "moviepy<2.0.0" imageio submitit hydra-core transformers +pip3 install "gymnasium[atari,mujoco]" vmas tqdm wandb pygame "moviepy<2.0.0" imageio submitit hydra-core transformers datasets cd /path/to/tensordict python setup.py develop diff --git a/sota-check/run_reward_model_training.sh b/sota-check/run_reward_model_training.sh new file mode 100644 index 00000000000..53ef77be849 --- /dev/null +++ b/sota-check/run_reward_model_training.sh @@ -0,0 +1,28 @@ +#!/bin/bash + +#SBATCH --job-name=reward_model_training +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --gres=gpu:1 +#SBATCH --output=slurm_logs/reward_model_training_%j.txt +#SBATCH --error=slurm_errors/reward_model_training_%j.txt + +current_commit=$(git rev-parse --short HEAD) +project_name="torchrl-example-check-$current_commit" +group_name="reward_model_training" +export PYTHONPATH=$(dirname $(dirname $PWD)) +python $PYTHONPATH/sota-implementations/reward_model_training/reward_model.py \ + model.name=gpt2 \ + model.device=cuda:0 \ + logger.backend=wandb \ + logger.project_name="$project_name" \ + logger.group_name="$group_name" + +# Capture the exit status of the Python command +exit_status=$? +# Write the exit status to a file +if [ $exit_status -eq 0 ]; then + echo "${group_name}_${SLURM_JOB_ID}=success" >> report.log +else + echo "${group_name}_${SLURM_JOB_ID}=error" >> report.log +fi diff --git a/sota-check/submitit-release-check.sh b/sota-check/submitit-release-check.sh index 8fd8fb3b072..e1b86a61f82 100755 --- a/sota-check/submitit-release-check.sh +++ b/sota-check/submitit-release-check.sh @@ -64,6 +64,7 @@ scripts=( run_ppo_atari.sh run_ppo_mujoco.sh run_rnd_mujoco.sh + run_reward_model_training.sh run_sac.sh run_td3.sh run_td3bc.sh diff --git a/sota-implementations/reward_model_training/README.md b/sota-implementations/reward_model_training/README.md new file mode 100644 index 00000000000..696e561de03 --- /dev/null +++ b/sota-implementations/reward_model_training/README.md @@ -0,0 +1,98 @@ +# RLHF reward-model training + +Trains a scalar **reward model** from pairwise human-preference data using the +Bradley-Terry objective +[`RewardModelLoss`](../../torchrl/objectives/llm/reward.py). Given a prompt and two +responses -- a `chosen` one preferred by an annotator and a `rejected` one -- the model +learns to assign a higher score to the chosen response: + +``` +loss = -log σ(r_θ(x, y_chosen) - r_θ(x, y_rejected)) +``` + +This is the reward-modelling stage that precedes policy optimization (e.g. +[GRPO](../grpo/)) in RLHF pipelines. + +This recipe is a minimal **single-node reference implementation**, not a large-scale +training stack. See [Scaling and integration](#scaling-and-integration) for how it is +meant to fit into serious training. + +## Install + +```bash +pip install -r requirements.txt +``` + +## Usage + +Train on the default summarization-comparisons dataset with a small Qwen backbone: + +```bash +python reward_model.py model.name=Qwen/Qwen2.5-0.5B +``` + +The backbone is **model-agnostic**: `model.name` accepts any Hugging Face model id +usable with `AutoModelForSequenceClassification` (a single-output head is attached +automatically), e.g. `facebook/opt-125m` or a larger instruct model. + +Use a different preference dataset (must expose `chosen`/`rejected`, and optionally +`prompt` fields): + +```bash +python reward_model.py model.name=Qwen/Qwen2.5-0.5B data.dataset_name=Anthropic/hh-rlhf \ + data.split_train=train data.split_val=test +``` + +Add the score-centering regularizer or freeze part of the backbone: + +```bash +python reward_model.py loss.center_coeff=0.01 optim.freeze_frac=0.7 +``` + +## Configuration + +Key fields in [`config.yaml`](config.yaml): + +| Group | Field | Description | +|-------|-------|-------------| +| `model` | `name` | HF backbone id; empty -> tiny from-scratch model (CI). | +| `model` | `device` | Training device; `null` auto-selects CUDA/MPS/CPU. | +| `data` | `dataset_name` | HF preference dataset; empty -> synthetic data (CI). | +| `data` | `max_length`, `batch_size`, `max_samples` | Tokenization / sampling. | +| `optim` | `max_iters`, `lr`, `freeze_frac`, `clip_grad` | Optimization. | +| `loss` | `reduction`, `center_coeff` | Bradley-Terry loss options. | +| `logger` | `backend`, `eval_iter`, `eval_iters` | Logging / evaluation. | + +## Hermetic smoke run + +Leaving `model.name` and `data.dataset_name` empty builds a tiny from-scratch model and +a synthetic preference dataset, so the recipe runs end-to-end with no download and +without the `datasets` dependency (this is what CI exercises): + +```bash +python reward_model.py model.name= data.dataset_name= \ + optim.max_iters=3 logger.eval_iter=2 logger.eval_iters=1 logger.backend= +``` + +## Scaling and integration + +This recipe is a minimal single-node baseline: a readable, single-process training loop +that exercises the loss and data format end to end. It is **not** a large-scale +reward-model training stack, and is not meant to become one. + +TorchRL owns the **contract**, not the parallelism. What this recipe provides: + +- a canonical TensorDict **preference-data format** (prompt / chosen / rejected), +- a canonical **reward-model loss** + ([`RewardModelLoss`](../../torchrl/objectives/llm/reward.py), Bradley-Terry), +- a small **reference trainer** (this recipe), +- and, in the future, **adapters/scorers** so an externally trained reward model plugs + back into TorchRL GRPO/PPO/SFT pipelines. + +For serious distributed reward-model training (parallelism, sharded checkpointing, +fault tolerance, packed/streamed datasets, mixed precision, activation checkpointing, +multi-node orchestration), use a dedicated backend such as TRL/Accelerate/FSDP, +TorchTitan/FSDP2, Megatron/NeMo, or OpenRLHF, while **preserving the TorchRL data and +scoring contract** so the resulting model plugs straight back into TorchRL's RLHF +pipelines. Interfacing with those backends, rather than reimplementing them, is the +intended direction for scaling this and the rest of the TorchRL LLM stack. diff --git a/sota-implementations/reward_model_training/config.yaml b/sota-implementations/reward_model_training/config.yaml new file mode 100644 index 00000000000..a634a458b7e --- /dev/null +++ b/sota-implementations/reward_model_training/config.yaml @@ -0,0 +1,51 @@ +seed: 42 + +# reward-model backbone +model: + # Any HF model id usable with AutoModelForSequenceClassification (e.g. + # Qwen/Qwen2.5-0.5B, facebook/opt-125m). Leave empty to build a tiny + # from-scratch model (used by CI). + name: Qwen/Qwen2.5-0.5B + # Training device. null = auto-select CUDA/MPS/CPU via get_available_device(). + device: null + +# preference data +data: + # HF preference dataset with prompt/chosen/rejected fields. Leave empty to use a + # hermetic synthetic dataset (no download, no `datasets` dependency). + dataset_name: CarperAI/openai_summarize_comparisons + split_train: train + split_val: valid1 + max_length: 550 + batch_size: 16 + max_samples: null # cap the number of pairs loaded (null = all) + synthetic_size: 64 # number of synthetic pairs when dataset_name is empty + +# optim +optim: + max_iters: 10000 + lr: 1.0e-5 + weight_decay: 0.0 + clip_grad: 1.0 + freeze_frac: 0.0 # fraction of backbone layers to freeze (0 = train all) + +# loss +loss: + reduction: mean + center_coeff: null # set a float to add the score-centering regularizer + +# logging / eval +logger: + backend: wandb + project_name: reward_model_training + group_name: null + exp_name: RewardModel_${model.name} + mode: online + eval_iter: 500 + eval_iters: 20 + log_interval: 10 + +# checkpointing +checkpoint: + save_dir: ./reward_model_ckpt + save_iter: 1000 diff --git a/sota-implementations/reward_model_training/requirements.txt b/sota-implementations/reward_model_training/requirements.txt new file mode 100644 index 00000000000..2e8765b55d9 --- /dev/null +++ b/sota-implementations/reward_model_training/requirements.txt @@ -0,0 +1,5 @@ +datasets +hydra-core +tqdm +transformers +wandb diff --git a/sota-implementations/reward_model_training/reward_model.py b/sota-implementations/reward_model_training/reward_model.py new file mode 100644 index 00000000000..d4934b12202 --- /dev/null +++ b/sota-implementations/reward_model_training/reward_model.py @@ -0,0 +1,159 @@ +# 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. +"""RLHF reward-model training recipe. + +Trains a scalar reward model from pairwise human-preference data using the +Bradley-Terry objective :class:`~torchrl.objectives.llm.RewardModelLoss`. This is the +reward-modelling stage that precedes policy optimization (e.g. GRPO/PPO) in RLHF +pipelines. + +The backbone is any Hugging Face ``AutoModelForSequenceClassification`` with a +single-output head. The helper functions live in the ``utils.py`` next to this script. + +Examples: + Train on the default summarization-comparisons dataset with a small Qwen backbone:: + + python reward_model.py model.name=Qwen/Qwen2.5-0.5B + + Run a quick hermetic smoke test (tiny from-scratch model, synthetic data):: + + python reward_model.py model.name= data.dataset_name= \\ + optim.max_iters=3 logger.backend= +""" +from __future__ import annotations + +import hydra +import torch +import tqdm +from torchrl._utils import get_available_device, logger as torchrl_logger, timeit +from torchrl.objectives.llm import RewardModelLoss +from torchrl.record.loggers import generate_exp_name, get_logger +from utils import ( + get_vocab_size, + log_metrics, + make_dataset, + make_optimizer, + make_replay_buffer, + make_reward_model, + make_tokenizer, +) + + +@hydra.main(config_path="", config_name="config", version_base="1.1") +def main(cfg): # noqa: F821 + # Logger + exp_name = generate_exp_name("RewardModel", cfg.logger.exp_name) + logger = None + if cfg.logger.backend: + logger = get_logger( + logger_type=cfg.logger.backend, + logger_name="reward_model_logging", + experiment_name=exp_name, + wandb_kwargs={ + "mode": cfg.logger.mode, + "config": dict(cfg), + "project": cfg.logger.project_name, + "group": cfg.logger.group_name, + }, + ) + + torch.manual_seed(int(cfg.seed)) + device = ( + torch.device(cfg.model.device) if cfg.model.device else get_available_device() + ) + + # Model + data + with timeit("setup/model"): + score_network = make_reward_model(cfg, device) + tokenizer = make_tokenizer(cfg) + vocab_size = get_vocab_size(tokenizer, score_network) + + with timeit("setup/data"): + train_data = make_dataset(cfg, tokenizer, cfg.data.split_train, vocab_size) + val_data = make_dataset(cfg, tokenizer, cfg.data.split_val, vocab_size) + train_rb = make_replay_buffer(train_data, cfg.data.batch_size, device) + val_rb = make_replay_buffer(val_data, cfg.data.batch_size, device) + + # Loss + optimizer + loss_module = RewardModelLoss( + score_network, + reduction=cfg.loss.reduction, + center_coeff=cfg.loss.center_coeff, + device=device, + ) + optimizer = make_optimizer(cfg, score_network) + + n_trainable = sum(p.numel() for p in score_network.parameters() if p.requires_grad) + torchrl_logger.info(f"Reward model with {n_trainable} trainable parameters.") + + @torch.no_grad() + def evaluate() -> tuple[float, float]: + score_network.eval() + losses, accuracies = [], [] + for _ in range(int(cfg.logger.eval_iters)): + batch = val_rb.sample() + out = loss_module(batch) + losses.append(out.loss_reward_model) + accuracies.append(out.accuracy) + score_network.train() + return ( + torch.stack(losses).mean().item(), + torch.stack(accuracies).mean().item(), + ) + + clip_grad = cfg.optim.clip_grad + max_iters = int(cfg.optim.max_iters) + pbar = tqdm.tqdm(range(1, max_iters + 1)) + for it in pbar: + timeit.printevery(num_prints=1000, total_count=max_iters, erase=True) + + with timeit("train/sample"): + batch = train_rb.sample() + + with timeit("train/forward"): + loss_out = loss_module(batch) + loss = loss_out.loss_reward_model + if loss_out.loss_center is not None: + loss = loss + loss_out.loss_center + + optimizer.zero_grad(set_to_none=True) + with timeit("train/backward"): + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_( + score_network.parameters(), clip_grad + ) + with timeit("train/optim_step"): + optimizer.step() + + metrics_to_log = { + "train/loss": loss.item(), + "train/accuracy": loss_out.accuracy.item(), + "train/grad_norm": grad_norm.item(), + } + pbar.set_postfix( + loss=f"{loss.item():.4f}", acc=f"{loss_out.accuracy.item():.3f}" + ) + + if it % int(cfg.logger.eval_iter) == 0: + val_loss, val_acc = evaluate() + metrics_to_log["eval/loss"] = val_loss + metrics_to_log["eval/accuracy"] = val_acc + torchrl_logger.info( + f"iter {it}: val_loss={val_loss:.4f}, val_acc={val_acc:.4f}" + ) + + if it % int(cfg.checkpoint.save_iter) == 0: + score_network.module.model.save_pretrained(cfg.checkpoint.save_dir) + + if logger is not None and it % int(cfg.logger.log_interval) == 0: + metrics_to_log.update(timeit.todict(prefix="time")) + log_metrics(logger, metrics_to_log, it) + + score_network.module.model.save_pretrained(cfg.checkpoint.save_dir) + pbar.close() + + +if __name__ == "__main__": + main() diff --git a/sota-implementations/reward_model_training/utils.py b/sota-implementations/reward_model_training/utils.py new file mode 100644 index 00000000000..69cd447b4e1 --- /dev/null +++ b/sota-implementations/reward_model_training/utils.py @@ -0,0 +1,226 @@ +# 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. +"""Helper functions for the RLHF reward-model training recipe. + +The recipe is model-agnostic: any Hugging Face ``AutoModelForSequenceClassification`` +with ``num_labels=1`` can be used as the reward-model backbone. A small from-config +model and a synthetic preference dataset are used when ``model.name`` / +``data.dataset_name`` are left empty, which keeps the CI smoke test hermetic (no +download, no ``datasets`` dependency). +""" +from __future__ import annotations + +import torch +from tensordict import TensorDict +from tensordict.nn import TensorDictModule +from torch import nn +from torchrl.data import ( + LazyTensorStorage, + SamplerWithoutReplacement, + TensorDictReplayBuffer, +) +from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer + + +class _RewardModel(nn.Module): + """Maps ``(input_ids, attention_mask)`` to a single scalar score per sequence.""" + + def __init__(self, hf_model: nn.Module) -> None: + super().__init__() + self.model = hf_model + + def forward( + self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None + ) -> torch.Tensor: + out = self.model(input_ids=input_ids, attention_mask=attention_mask) + return out.logits # shape [B, num_labels=1] + + +def make_reward_model(cfg, device: torch.device) -> TensorDictModule: + """Build the score network: an HF sequence-classification model with a 1-d head. + + When ``cfg.model.name`` is empty, a tiny GPT2-style model is built from scratch + (random weights, no download) so the recipe can run hermetically in CI. + """ + name = cfg.model.name + if name: + hf_model = AutoModelForSequenceClassification.from_pretrained( + name, num_labels=1 + ) + else: + config = AutoConfig.for_model( + "gpt2", + num_labels=1, + n_layer=2, + n_head=2, + n_embd=64, + vocab_size=256, + n_positions=max(int(cfg.data.max_length), 32), + # keep the special token ids within the tiny vocab + bos_token_id=0, + eos_token_id=0, + ) + hf_model = AutoModelForSequenceClassification.from_config(config) + + # Sequence-classification models need a pad token to locate the final token. + if hf_model.config.pad_token_id is None: + hf_model.config.pad_token_id = ( + hf_model.config.eos_token_id + if hf_model.config.eos_token_id is not None + else 0 + ) + + score_network = TensorDictModule( + _RewardModel(hf_model), + in_keys=["input_ids", "attention_mask"], + out_keys=["score"], + ) + return score_network.to(device) + + +def make_tokenizer(cfg): + """Return the tokenizer matching the model, or ``None`` in synthetic mode.""" + name = cfg.model.name + if not name: + return None + tokenizer = AutoTokenizer.from_pretrained(name) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + return tokenizer + + +def get_vocab_size(tokenizer, score_network: TensorDictModule) -> int: + if tokenizer is not None: + return len(tokenizer) + return score_network.module.model.config.vocab_size + + +def _pairwise_td( + chosen_ids: torch.Tensor, + rejected_ids: torch.Tensor, + chosen_mask: torch.Tensor | None = None, + rejected_mask: torch.Tensor | None = None, +) -> TensorDict: + """Pack tokenized chosen/rejected responses into the loss input layout.""" + n = chosen_ids.shape[0] + if chosen_mask is None: + chosen_mask = torch.ones_like(chosen_ids) + if rejected_mask is None: + rejected_mask = torch.ones_like(rejected_ids) + return TensorDict( + { + "chosen": TensorDict( + {"input_ids": chosen_ids, "attention_mask": chosen_mask}, + batch_size=[n], + ), + "rejected": TensorDict( + {"input_ids": rejected_ids, "attention_mask": rejected_mask}, + batch_size=[n], + ), + }, + batch_size=[n], + ) + + +def make_dataset(cfg, tokenizer, split: str, vocab_size: int) -> TensorDict: + """Build a pairwise preference dataset as a single batched ``TensorDict``. + + The returned tensordict has ``"chosen"`` and ``"rejected"`` sub-tensordicts, each + carrying ``input_ids`` / ``attention_mask`` -- exactly the keys + :class:`~torchrl.objectives.llm.RewardModelLoss` expects by default. + """ + dataset_name = cfg.data.dataset_name + max_length = int(cfg.data.max_length) + + if not dataset_name: + # Hermetic synthetic dataset (no download, no ``datasets`` dependency). + n = int(cfg.data.synthetic_size) + gen = torch.Generator().manual_seed(int(cfg.seed) + (split == "train")) + chosen_ids = torch.randint(0, vocab_size, (n, max_length), generator=gen) + rejected_ids = torch.randint(0, vocab_size, (n, max_length), generator=gen) + return _pairwise_td(chosen_ids, rejected_ids) + + # Real preference data. ``datasets`` is imported lazily so the synthetic/CI path + # never requires it. + from datasets import load_dataset + + ds = load_dataset(dataset_name, split=split) + max_samples = cfg.data.max_samples + if max_samples is not None: + ds = ds.select(range(min(int(max_samples), len(ds)))) + + chosen_texts, rejected_texts = [], [] + for sample in ds: + prompt = sample.get("prompt", "") + sep = "\n" if prompt else "" + chosen_texts.append(prompt + sep + sample["chosen"]) + rejected_texts.append(prompt + sep + sample["rejected"]) + + tok_kwargs = { + "max_length": max_length, + "padding": "max_length", + "truncation": True, + "return_tensors": "pt", + } + chosen_tok = tokenizer(chosen_texts, **tok_kwargs) + rejected_tok = tokenizer(rejected_texts, **tok_kwargs) + return _pairwise_td( + chosen_tok["input_ids"], + rejected_tok["input_ids"], + chosen_tok["attention_mask"], + rejected_tok["attention_mask"], + ) + + +def make_replay_buffer( + data: TensorDict, batch_size: int, device: torch.device +) -> TensorDictReplayBuffer: + rb = TensorDictReplayBuffer( + storage=LazyTensorStorage(data.shape[0], device=device), + sampler=SamplerWithoutReplacement(drop_last=True), + batch_size=batch_size, + ) + rb.extend(data) + return rb + + +def make_optimizer(cfg, score_network: TensorDictModule) -> torch.optim.Optimizer: + _maybe_freeze_backbone(score_network, cfg.optim.freeze_frac) + params = [p for p in score_network.parameters() if p.requires_grad] + return torch.optim.AdamW( + params, lr=cfg.optim.lr, weight_decay=cfg.optim.weight_decay + ) + + +def _maybe_freeze_backbone(score_network: TensorDictModule, freeze_frac: float) -> None: + """Best-effort freezing of the first ``freeze_frac`` of transformer layers. + + Freezing the lower layers of the backbone is a common efficiency trick for + reward-model fine-tuning. This is best-effort: if the backbone layer list cannot + be located for the given architecture, no layer is frozen. + """ + if not freeze_frac or freeze_frac <= 0: + return + base = score_network.module.model + # The base transformer is exposed via ``base_model`` on HF models. + transformer = getattr(base, "base_model", base) + layers = None + for attr in ("h", "layers", "layer", "block"): + candidate = getattr(transformer, attr, None) + if candidate is not None and len(candidate) > 0: + layers = candidate + break + if layers is None: + return + num_freeze = int(freeze_frac * len(layers)) + for layer in layers[:num_freeze]: + layer.requires_grad_(False) + + +def log_metrics(logger, metrics: dict, step: int) -> None: + if logger is None: + return + for key, value in metrics.items(): + logger.log_scalar(key, value, step) diff --git a/test/llm/test_llm_objectives.py b/test/llm/test_llm_objectives.py index 125cea91102..07f4c0fcdbb 100644 --- a/test/llm/test_llm_objectives.py +++ b/test/llm/test_llm_objectives.py @@ -29,6 +29,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 +1128,133 @@ 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): + from tensordict.nn import TensorDictModule + + 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 (CLAUDE.md requirement).""" + 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_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..300b9bdf891 100644 --- a/torchrl/objectives/llm/__init__.py +++ b/torchrl/objectives/llm/__init__.py @@ -16,6 +16,7 @@ MCAdvantageSelector, RayMCAdvantage, ) +from .reward import RewardModelLoss, RewardModelLossOutput from .sft import SFTLoss, SFTLossOutput __all__ = [ @@ -29,6 +30,8 @@ "MCAdvantage", "MCAdvantageSelector", "RayMCAdvantage", + "RewardModelLoss", + "RewardModelLossOutput", "SFTLoss", "SFTLossOutput", ] diff --git a/torchrl/objectives/llm/reward.py b/torchrl/objectives/llm/reward.py new file mode 100644 index 00000000000..3f36cd9678d --- /dev/null +++ b/torchrl/objectives/llm/reward.py @@ -0,0 +1,272 @@ +# 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, optional): the device on which to run ``score_network``. + 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 | None = None, + ) -> None: + super().__init__() + self.score_network = score_network + self.reduction = reduction + self.center_coeff = center_coeff + 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.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.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, + )