Source code for torchrl.envs.llm.transforms.kl
# 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 warnings
from collections.abc import Callable
from contextlib import nullcontext
from copy import copy
from typing import Any, Literal, TYPE_CHECKING
import torch
from tensordict import NestedKey, set_list_to_stack, TensorDictBase, unravel_key
from tensordict.utils import _zip_strict, is_seq_of_nested_key, logger as torchrl_logger
from torch.nn.utils.rnn import pad_sequence
from torchrl.data import Composite, Unbounded
from torchrl.data.tensor_specs import DEVICE_TYPING
from torchrl.envs import EnvBase, Transform
from torchrl.envs.llm.transforms.ray_service import _RayServiceMetaClass, RayTransform
from torchrl.envs.transforms.transforms import Compose
from torchrl.envs.transforms.utils import _set_missing_tolerance
from torchrl.modules.llm.policies.common import LLMWrapperBase
if TYPE_CHECKING:
import transformers
class RayKLRewardTransform(RayTransform):
"""A Ray-based implementation of :class:`~torchrl.envs.llm.transforms.kl.KLRewardTransform`.
This class creates a Ray remote actor from KLRewardTransform that can be shared across multiple workers.
All method calls are delegated to the remote actor, ensuring that multiple environments can
share the same KL computation resources.
To avoid serialization issues with large models, this class supports model factories
that create models on the remote actor rather than passing full models through Ray channels.
Args:
ref_model (LLMWrapperBase, optional): the reference model. Prefer using a model factory instead
to avoid serialization issues.
Keyword Args:
ref_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a reference model.
This allows for explicit resource control and avoids serialization issues.
num_cpus (int, optional): Number of CPUs to allocate to the Ray actor. Defaults to 1.
num_gpus (int, optional): Number of GPUs to allocate to the Ray actor. Defaults to 0.
device (torch.device, optional): Device to use on the remote Ray actor for tensor operations.
The local Ray transform will handle CPU serialization and device restoration automatically.
Defaults to None.
actor_name (str, optional): Name of the Ray actor to use. If provided, the actor will be reused if it already exists.
**kwargs: Additional keyword arguments to pass to KLRewardTransform.
Note:
When using model factories, the corresponding model argument (ref_model) should be None.
Model factories are preferred for large models to avoid serialization overhead.
Examples:
>>> # Option 1: Using model factory for explicit resource control
>>> def create_ref_model():
... return TransformersWrapper(ref_model, tokenizer=tokenizer, generate=False, return_log_probs=True)
>>> transform = RayKLRewardTransform(
... ref_model=None,
... ref_model_factory=create_ref_model,
... num_gpus=1,
... device=torch.device("cuda")
... )
>>> # Option 2: Pass model directly (Ray handles serialization)
>>> transform = RayKLRewardTransform(ref_model=ref_model, device=torch.device("cuda"))
"""
def __init__(
self,
ref_model: LLMWrapperBase | None = None,
*,
ref_model_factory: Callable[[], LLMWrapperBase] | None = None,
num_cpus: int | None = None,
num_gpus: int = 0,
device: DEVICE_TYPING | None = None,
actor_name: str | None = None,
**kwargs,
):
# Validate arguments: model and factory should not both be provided
if ref_model is not None and ref_model_factory is not None:
raise ValueError(
"Cannot provide both 'ref_model' and 'ref_model_factory'. Choose one."
)
if ref_model is None and ref_model_factory is None:
raise ValueError(
"Must provide exactly one of 'ref_model' or 'ref_model_factory'."
)
# Store creation parameters for actor creation
self._ref_model = ref_model
self._ref_model_factory = ref_model_factory
self._creation_kwargs = kwargs
# Store device separately for passing to remote actor
self._remote_device = device
# Default num_cpus
if num_cpus is None:
num_cpus = 1
# Call parent constructor without device (Ray transform handles CPU/device mapping)
super().__init__(
num_cpus=num_cpus,
num_gpus=num_gpus,
device=None, # Don't store device locally
actor_name=actor_name,
**kwargs,
)
def _create_actor(self, **kwargs):
"""Create the remote KLRewardTransform actor."""
# Create the remote KLRewardTransform with resource specifications
RemoteKLRewardTransform = self._ray.remote(
num_cpus=self._num_cpus, num_gpus=self._num_gpus
)(KLRewardTransform)
if self._actor_name is not None:
RemoteKLRewardTransform = RemoteKLRewardTransform.options(
name=self._actor_name
)
# Determine how to create model on the remote actor
ref_model_arg = self._ref_model
# If we have factory, we'll pass it and set model to None
creation_kwargs = self._creation_kwargs.copy()
if self._ref_model_factory is not None:
creation_kwargs["ref_model_factory"] = self._ref_model_factory
ref_model_arg = None
# Pass device to the remote actor
if self._remote_device is not None:
creation_kwargs["device"] = self._remote_device
# Create the shared actor
actor = RemoteKLRewardTransform.remote(
ref_model=ref_model_arg, **creation_kwargs
)
return actor
def __repr__(self):
"""String representation."""
try:
if hasattr(self, "_actor") and self._actor is not None:
return self._ray.get(self._actor.__repr__.remote())
else:
return "RayKLRewardTransform(actor=None)"
except Exception:
return f"RayKLRewardTransform(actor={getattr(self, '_actor', 'None')})"
[docs]class KLRewardTransform(Transform, metaclass=_RayServiceMetaClass):
"""A legacy transform for computing KL divergence-based rewards.
**Deprecated**: This transform is maintained for backward compatibility but is no longer
the recommended approach. Use :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL` instead,
which provides better modularity and integration with the new wrapper design.
**Recent Changes:**
- **Legacy Status**: This transform is now considered legacy and may not work optimally
with the new modular wrapper design.
- **ChatHistory Integration**: Limited support for the new :class:`~torchrl.modules.llm.policies.ChatHistory` objects.
- **Input Mode Support**: May not handle all input modes (`"history"`, `"text"`, `"tokens"`) consistently.
**Recommendation**:
Use :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL` for new code, which provides:
- Better integration with the new wrapper design
- Consistent support for all input modes
- Proper handling of ChatHistory objects
- More modular and composable architecture
Args:
ref_model (LLMWrapperBase): the reference model.
Keyword Args:
ref_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a reference model.
assistant_only (bool): whether to only compute KL on assistant tokens. Defaults to `True`.
tokenizer (transformers.AutoTokenizer): the tokenizer to use. Defaults to `None`.
detach (bool): whether to detach the KL from the computation graph. Defaults to `True`.
device (torch.device): the device to cast the tensors to. This is not the device of the specs, but the device
onto which the tensors will be moved. It allows to keep the model on a different device
than the upcoming data. When using Ray service, this device will be used on the remote actor.
Defaults to `None`.
padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
use_ray_service (bool, optional): whether to use Ray service. Defaults to `False`.
actor_name (str, optional): the name of the Ray actor to use. Defaults to `None`.
Examples:
>>> # Legacy usage (not recommended for new code)
>>> transform = KLRewardTransform(gen_model, ref_model)
>>>
>>> # Recommended approach using RetrieveKL
>>> from torchrl.envs.llm.transforms.kl import RetrieveKL
>>> transform = RetrieveKL(gen_model, ref_model, assistant_only=True)
.. seealso::
:class:`~torchrl.envs.llm.transforms.kl.RetrieveKL`: The recommended transform for KL divergence computation.
:class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb`: Base transform for retrieving log-probabilities.
:class:`~torchrl.envs.llm.transforms.kl.KLComputation`: Transform for computing KL divergence between log-prob tensors.
"""
DEFAULT_IN_KEYS = ["reward"]
_RayServiceClass = RayKLRewardTransform
def __init__(
self,
ref_model: LLMWrapperBase | None = None,
*,
ref_model_factory: Callable[[], LLMWrapperBase] | None = None,
coef=1.0,
in_keys=None,
out_keys=None,
log_prob_key: NestedKey = ("log_probs", "full"),
device: torch.device | None = None,
add_to_reward: bool = True,
tokenizer: transformers.AutoTokenizer | None = None,
assistant_only: bool = True,
padding_side: str = "left",
use_ray_service: bool = False,
):
# Handle model factory - create model if factory is provided
if ref_model_factory is not None:
if ref_model is not None:
raise ValueError(
"Cannot provide both 'ref_model' and 'ref_model_factory'. Choose one."
)
ref_model = ref_model_factory()
elif ref_model is None:
raise ValueError(
"Must provide exactly one of 'ref_model' or 'ref_model_factory'."
)
if in_keys is None:
in_keys = self.DEFAULT_IN_KEYS
if out_keys is None:
out_keys = copy(in_keys)
if len(out_keys) == len(in_keys):
out_keys = out_keys + ["kl_penalty", "ref_log_prob"]
elif len(out_keys) != len(in_keys) + 2:
raise ValueError(
"The out_keys must have the same length as the in_keys (plus two additional optional kl entries for logging)."
)
super().__init__(in_keys=in_keys, out_keys=out_keys)
if not is_seq_of_nested_key(self.in_keys) or not is_seq_of_nested_key(
self.out_keys
):
raise ValueError(
f"invalid in_keys / out_keys:\nin_keys={self.in_keys} \nout_keys={self.out_keys}"
)
if len(self.in_keys) != 1 or len(self.out_keys) != 3:
raise ValueError(
f"Only one in_key/out_key is allowed, got in_keys={self.in_keys}, out_keys={self.out_keys}."
)
self._out_keys = [unravel_key(out_key) for out_key in self._out_keys]
if getattr(ref_model, "generate", False):
raise ValueError(
"The actor is configured to generate text, not compute the log-probs."
)
# update the in_keys for dispatch etc
self.in_keys = self.in_keys + ref_model.in_keys
self.in_keys = [unravel_key(in_key) for in_key in self.in_keys]
self.add_to_reward = add_to_reward
# check that the model has parameters
self.__dict__["ref_model"] = ref_model
# self._buffers["actor_params"] = params.clone().detach()
self.device = device
# find the sample log-prob key
self.log_prob_full_key = log_prob_key
self._tokenizer = tokenizer
self.assistant_only = assistant_only
self.padding_side = padding_side
if not isinstance(coef, torch.Tensor):
coef = torch.as_tensor(coef)
self.register_buffer("coef", coef)
# sanity check for the ref_model
if not getattr(ref_model, "input_mode", "tokens") == "tokens":
raise ValueError(
"The ref_model must be configured to use tokens as input. Please set the `input_mode` argument to `tokens`."
)
@property
def pad_output(self):
# We need pad_output to match the pad_output of the inference model
return self.ref_model.pad_output
@property
def tokenizer(self):
tokenizer = self._tokenizer
if tokenizer is not None:
return tokenizer
try:
return self.ref_model.tokenizer
except AttributeError:
raise AttributeError(
"The ref_model does not have a tokenizer. Please pass the tokenizer to the constructor."
)
def set_container(self, container: Transform | EnvBase) -> None:
result = super().set_container(container)
if self.action_key is None:
parent = getattr(self, "parent", None)
if parent is not None:
action_keys = parent.action_keys
if len(action_keys) != 1:
raise ValueError(
f"More than one action_key found. Please pass the `action_key` argument directly to {type(self).__name__}."
)
action_key = action_keys[0]
self.action_key = action_key
return result
def _reset(
self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
) -> TensorDictBase:
with _set_missing_tolerance(self, True):
tensordict_reset = self._step(tensordict_reset, tensordict_reset)
return tensordict_reset
@property
def action_key(self) -> NestedKey:
# Get the action from the base env (a ChatEnv).
if self.parent.base_env.input_mode == "history":
return ("history", "full")
if self.parent.base_env.input_mode == "text":
return ("text", "full")
if self.parent.base_env.input_mode == "tokens":
return ("tokens", "full")
raise ValueError(f"Invalid input mode: {self.parent.base_env.input_mode}")
def _step(
self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
) -> TensorDictBase:
original_device = None
if self.device is not None:
original_device = tensordict.device
tensordict = tensordict.to(self.device)
next_tensordict = next_tensordict.to(self.device)
# tensordict = self._get_text_response(tensordict, next_tensordict)
response = tensordict.get(self.action_key, None)
if response is None:
if not self.missing_tolerance:
raise RuntimeError(
f"Action with key {self.action_key} not found data {tensordict}"
)
# being called after reset or without action, skipping
if self.out_keys[0] != "reward" and self.parent is not None:
next_tensordict.set(self.out_keys[0], self.parent.reward_spec.zero())
return next_tensordict
# We use the ("tokens", "full") key to get the log-probs of the reference model
with torch.device(self.device) if self.device is not None else nullcontext():
td_input = tensordict.copy()
ref_log_prob_td = self.ref_model(td_input)
if self.pad_output:
ref_log_prob_padded = ref_log_prob_td.get(self.log_prob_full_key)
else:
ref_log_prob_unpadded = ref_log_prob_td.get(
self.log_prob_full_key, as_list=True # type: ignore[misc]
)
if self.assistant_only:
# Get the assistant mask
mask = tensordict.get(("masks", "all_assistant_mask"))
# mask will often be None - fall back on prompt / response separation
if mask is None:
if self.pad_output:
# simple case: just take the prompt length
prompt_length = tensordict.get(("tokens", "prompt")).shape[-1]
mask = tensordict.get(("masks", "all_attention_mask")).clone()
mask[..., :prompt_length] = False
else:
# simple case: just take the prompt length
prompt_length = [
t.size(-1)
for t in tensordict.get(("tokens", "prompt"), as_list=True) # type: ignore[misc]
]
mask = tensordict.get(("masks", "all_attention_mask"), as_list=True) # type: ignore[misc]
for i in range(len(prompt_length)):
mask[i] = mask[i].clone()
mask[i][..., : prompt_length[i]] = False
# we want to keep the batch dimension
ref_log_prob_list = []
if self.pad_output:
for i in range(ref_log_prob_padded.size(0)):
ref_log_prob_list.append(
ref_log_prob_padded[i].masked_fill(~mask[i], 0)
)
else:
for i in range(len(ref_log_prob_unpadded)):
ref_log_prob_list.append(
ref_log_prob_unpadded[i].masked_fill(~mask[i], 0)
)
if self.pad_output:
ref_log_prob = pad_sequence(
ref_log_prob_list,
batch_first=True,
padding_value=0,
padding_side=self.padding_side,
)
else:
ref_log_prob = torch.nested.nested_tensor(
ref_log_prob_list, layout=torch.strided
)
# we obtain the current log-probs (already computed) from the current tensordict
if self.pad_output:
curr_log_prob_padded = tensordict.get(self.log_prob_full_key)
else:
curr_log_prob_unpadded = tensordict.get(
self.log_prob_full_key, as_list=True # type: ignore[misc]
)
if self.assistant_only:
# we want to keep the batch dimension
curr_log_prob_list = []
if self.pad_output:
for i in range(curr_log_prob_padded.size(0)):
curr_log_prob_list.append(
curr_log_prob_padded[i].masked_fill(~mask[i], 0)
)
else:
for i in range(len(curr_log_prob_unpadded)):
curr_log_prob_list.append(
curr_log_prob_unpadded[i].masked_fill(~mask[i], 0)
)
if self.pad_output:
curr_log_prob = pad_sequence(
curr_log_prob_list,
batch_first=True,
padding_value=0,
padding_side=self.padding_side,
)
else:
curr_log_prob = torch.nested.nested_tensor(
curr_log_prob_list, layout=torch.strided
)
ref_log_prob = ref_log_prob.to(curr_log_prob.device)
# We want the log-probs to have a similar dim to the reward
curr_log_prob = curr_log_prob.unsqueeze(-1)
ref_log_prob = ref_log_prob.unsqueeze(-1)
for i in range(ref_log_prob.size(0)):
if ref_log_prob[i].shape != curr_log_prob[i].shape:
# Don't check shapes if nested
raise ValueError(
f"the log-probability tensor shapes must match, got cur_log_prob.shape={curr_log_prob[i].shape} and log_prob.shape={ref_log_prob[i].shape}. "
f"One possible reason is that the padding token is identical to the eos token, which means that the eos_token log_prob is truncated from the "
f"reference model output."
)
kl = curr_log_prob - ref_log_prob
if self.add_to_reward:
reward_key = self.in_keys[0]
reward = next_tensordict.get(reward_key)
# we use the unbiased consistent estimator of the KL: log_p(x) - log_q(x) when x ~ p(x)
if not reward.is_nested and ref_log_prob.is_nested:
reward = torch.nested.nested_tensor(
[rew.expand(lp.shape) for rew, lp in zip(reward, ref_log_prob)],
layout=torch.strided,
)
if reward is not None and reward.ndim != curr_log_prob.ndim:
raise ValueError(
"The number of dimensions of reward must be the same as the number of dimensions of the KL "
f"term. Got ndim={reward.ndim} and {curr_log_prob.ndim} respectively."
)
if reward is None:
reward = 0
reward = reward - self.coef * kl
next_tensordict.set(self.out_keys[0], reward)
next_tensordict.set(self.out_keys[1], kl)
next_tensordict.set(self.out_keys[2], ref_log_prob)
if original_device is not None:
next_tensordict = next_tensordict.to(original_device)
return next_tensordict
[docs] def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
next_td = tensordict.pop("next")
next_td = self._step(tensordict, next_td)
return tensordict.set("next", next_td)
[docs] def transform_output_spec(self, output_spec: Composite) -> Composite:
in_key = unravel_key(self.in_keys[0])
out_key = unravel_key(self.out_keys[0])
if "full_observation_spec" in output_spec.keys():
observation_spec = output_spec["full_observation_spec"]
else:
observation_spec = Composite(
shape=output_spec.shape, device=output_spec.device
)
output_spec["full_observation_spec"] = observation_spec
if in_key == "reward" and out_key == "reward":
parent = self.parent
reward_keys = parent.reward_keys
if len(reward_keys) == 1:
reward_key = reward_keys[0]
shape = output_spec["full_reward_spec"].shape
elif "reward" in reward_keys:
reward_key = "reward"
shape = output_spec["full_reward_spec"].shape
else:
shape = output_spec.shape
reward_key = "reward"
# For LLMs, the shape of the reward is (batch, -1, 1)
shape = torch.Size((*shape, -1, 1))
reward_spec = Unbounded(
device=output_spec.device,
shape=shape,
)
output_spec["full_reward_spec"] = Composite(
{reward_key: reward_spec},
shape=output_spec["full_reward_spec"].shape,
)
elif in_key == "reward":
# TODO: we should at least allow to make this a component of the reward specs, to avoid a call during reset
parent = self.parent
reward_spec = output_spec["full_reward_spec"][parent.reward_key]
shape = output_spec["full_reward_spec"].shape
# For LLMs, the shape of the reward is (batch, -1, 1)
shape = torch.Size((*shape, -1, 1))
reward_spec = reward_spec.clone()
reward_spec.shape = shape
# then we need to populate the output keys
observation_spec[out_key] = reward_spec
else:
observation_spec = output_spec["full_observation_spec"]
reward_spec = observation_spec[in_key]
shape = observation_spec.shape
shape = torch.Size((*shape, -1, 1))
reward_spec = reward_spec.clone()
reward_spec.shape = shape
# then we need to populate the output keys
observation_spec[out_key] = reward_spec
observation_spec[self.out_keys[1]] = reward_spec.clone()
return output_spec
[docs]class RetrieveLogProb(Transform):
"""A transform to retrieve log-probabilities from a model for KL divergence computation.
This transform computes log-probabilities from a reference model, which can then be used
to compute KL divergence with another model's log-probabilities. It's designed to work
with the :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL` and :class:`~torchrl.envs.llm.transforms.kl.KLComputation` transforms.
Args:
model (LLMWrapperBase): the model to use to compute the log-probs.
Keyword Args:
log_probs_full_key (NestedKey): the key where the log-probs are stored.
If not provided, the key will be retrieved from the model's `log_probs_key` attribute
(i.e., `(model.log_probs_key, "full")`).
assistant_only (bool): whether to zero out the log-probs of the non-assistant tokens (i.e., steps of history
where the role is not `"assistant"`). Defaults to `True`.
.. note:: When `assistant_only=True`, the model must have `input_mode='history'` to properly identify
assistant tokens. For other input modes (`"text"` or `"tokens"`), set `assistant_only=False`.
This ensures users are conscious of the limitation that assistant token identification requires
structured conversation history.
tokenizer_kwargs (dict): the keyword arguments to pass to the tokenizer to be used to apply the chat template to the history when `assistant_only` is `True`.
To control the tokenization in the ref_model, pass the tokenizer kwargs to the ref_model constructor.
Defaults to `{"return_assistant_tokens_mask": True, "tokenize": True, "return_dict": True, "padding": False, "add_generation_prompt": False}`.
tokenizer (transformers.AutoTokenizer): the tokenizer to be used to tokenize the input and compute the assitant mask. If not provided, the tokenizer will be inferred from the `ref_model`.
detach (bool): whether to exclude the log-probs from the gradient computation. Defaults to `True`.
device (torch.device): the device to use for tensor creation. Defaults to `None`.
padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
Examples:
>>> from torchrl.data.llm import History
>>> from torchrl.modules.llm import TransformersWrapper
>>> from torchrl.modules.llm.policies import ChatHistory
>>> from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
>>> from tensordict import TensorDict, set_list_to_stack
>>> import torch
>>>
>>> # Set up list to stack for History
>>> set_list_to_stack(True).set()
>>>
>>> # Create chat data
>>> 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)
>>> print(f"Created history with shape: {history.shape}")
Created history with shape: torch.Size([2, 3])
>>>
>>> # Setup tokenizer and model
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
>>> tokenizer.pad_token = tokenizer.eos_token
>>> model = OPTForCausalLM(OPTConfig()).eval()
>>>
>>> # Create reference model
>>> ref_model = TransformersWrapper(
... model,
... tokenizer=tokenizer,
... input_mode="history",
... generate=False,
... return_log_probs=True,
... pad_output=True,
... )
>>>
>>> # Create the RetrieveLogProb transform
>>> transform = RetrieveLogProb(
... ref_model,
... assistant_only=True,
... tokenizer=tokenizer,
... )
>>>
>>> # Prepare data using ChatHistory
>>> chat_history = ChatHistory(full=history)
>>> data = TensorDict(history=chat_history, batch_size=(2,))
>>>
>>> # Apply the transform to get reference log probabilities
>>> result = transform(data)
>>> log_probs_key = (ref_model.log_probs_key, "full")
>>> ref_log_probs = result.get(log_probs_key)
>>> print(f"Log-probs shape: {ref_log_probs.shape}")
Log-probs shape: torch.Size([2, 26])
.. note::
By default, the log-probabilities are stored as a list of tensors (one per sample, with variable length).
Use `as_padded_tensor=True` in `.get()` to obtain a batchable tensor (with padding).
The reference log probabilities are computed only for assistant tokens when `assistant_only=True`.
**Input Mode Compatibility:**
- When `assistant_only=True` (default), the model must have `input_mode='history'` to properly identify assistant tokens.
- When `assistant_only=False`, the transform works with any input mode (`"history"`, `"text"`, or `"tokens"`).
- This design ensures users are conscious of the limitation that assistant token identification requires structured conversation history.
.. seealso::
:class:`~torchrl.envs.llm.transforms.kl.RetrieveKL`: A higher-level transform that combines two `RetrieveLogProb` instances with `KLComputation`.
:class:`~torchrl.envs.llm.transforms.kl.KLComputation`: A transform that computes KL divergence between two log-prob tensors.
:class:`~torchrl.envs.llm.transforms.kl.KLRewardTransform`: A legacy transform for KL reward computation (use `RetrieveKL` instead).
"""
def __init__(
self,
model: LLMWrapperBase,
*,
log_probs_full_key: NestedKey | None = None,
assistant_only: bool = True,
tokenizer_kwargs: dict | None = None,
detach: bool = True,
device: torch.device | None = None,
tokenizer: transformers.AutoTokenizer | None = None,
padding_side: str = "left",
):
# Set up keys
if log_probs_full_key is None:
log_probs_full_key = (model.log_probs_key, "full")
elif (
not isinstance(log_probs_full_key, tuple)
or log_probs_full_key[-1] != "full"
):
warnings.warn(
f"The log_probs_full_key {log_probs_full_key} is not a tuple or does not end with 'full'. "
"This may cause issues with the KL computation. "
"Please use a tuple with the log_probs_key and 'full' as the last element."
)
self.log_probs_full_key = log_probs_full_key
# Set up input/output keys
in_keys = list(model.in_keys)
out_keys = [self.log_probs_full_key]
super().__init__(in_keys=in_keys, out_keys=out_keys)
# Store model and configuration
self.model = model
self.assistant_only = assistant_only
self.detach = detach
self.device = device
self.tokenizer = tokenizer
self.padding_side = padding_side
# Set up tokenizer kwargs
if tokenizer_kwargs is None:
tokenizer_kwargs = {}
tokenizer_kwargs.setdefault("return_assistant_tokens_mask", True)
tokenizer_kwargs.setdefault("tokenize", True)
tokenizer_kwargs.setdefault("return_dict", True)
tokenizer_kwargs.setdefault("padding", False)
tokenizer_kwargs.setdefault("add_generation_prompt", False)
self.tokenizer_kwargs = tokenizer_kwargs
# Validate model configuration (after setting assistant_only)
self._validate_model_config(model)
def _validate_model_config(self, model: LLMWrapperBase):
"""Validate model configuration."""
if not getattr(model, "return_log_probs", True):
raise ValueError(
"The model must have `return_log_probs=True` to use the `RetrieveLogProb` transform."
)
if getattr(model, "generate", True):
raise ValueError(
"The model must have `generate=False` to use the `RetrieveLogProb` transform."
)
# Check input mode compatibility with assistant_only
input_mode = getattr(model, "input_mode", "history")
if self.assistant_only and input_mode != "history":
raise ValueError(
f"The model must have `input_mode='history'` when `assistant_only=True`. "
f"Current input_mode is '{input_mode}'. "
f"To use input_mode '{input_mode}', set `assistant_only=False`."
)
[docs] def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
next_td = tensordict.get("next")
next_is_none = False
if next_td is None:
next_is_none = True
next_td = tensordict
output = self._step(tensordict, next_td)
if next_is_none:
return output
return tensordict.set("next", output)
def _mask_assistant_tokens(
self, td: TensorDictBase, lp_key: NestedKey
) -> torch.Tensor:
"""Mask log-probs to only include assistant tokens.
Args:
td: TensorDict containing the data
lp_key: Key for log-probs in the TensorDict
Returns:
Masked log-probs tensor
"""
with torch.device(self.device) if self.device is not None else nullcontext():
# Get assistant mask
assistant_masks = td.get(("masks", "all_assistant_mask"), as_list=True) # type: ignore[misc]
log_probs = td.get(lp_key, as_list=True) # type: ignore[misc]
log_probs = [
torch.masked_fill(lp, ~mask, 0.0)
for lp, mask in _zip_strict(log_probs, assistant_masks)
]
if self.model.pad_output:
log_probs = pad_sequence(
log_probs,
batch_first=True,
padding_value=0.0,
padding_side=self.padding_side,
)
else:
log_probs = torch.nested.as_nested_tensor(
log_probs, layout=self.model.layout
)
return log_probs
@set_list_to_stack(True)
def _step(
self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
) -> TensorDictBase:
# Compute log-probs using the model
# Use tensordict since we want to process the "full" entry
ref_td = self.model(tensordict.copy())
tmp_log_probs_key = (self.model.log_probs_key, "full")
# Apply assistant masking if requested
if self.assistant_only:
log_probs = self._mask_assistant_tokens(ref_td, tmp_log_probs_key)
ref_td.set(tmp_log_probs_key, log_probs)
# Rename and store the log-probs
if tmp_log_probs_key != self.log_probs_full_key:
ref_td.rename_key_(tmp_log_probs_key, self.log_probs_full_key)
next_tensordict.update(ref_td, keys_to_update=(self.log_probs_full_key,))
return next_tensordict
[docs] def transform_observation_spec(self, observation_spec: Composite) -> Composite:
# Add kl to observation spec
observation_spec["kl_penalty"] = Unbounded(
device=observation_spec.device,
shape=observation_spec.shape,
)
return observation_spec
class RayRetrieveKL(RayTransform):
"""A Ray-based implementation of :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL`.
This class creates a Ray remote actor from RetrieveKL that can be shared across multiple workers.
All method calls are delegated to the remote actor, ensuring that multiple environments can
share the same KL computation resources.
To avoid serialization issues with large models, this class supports model factories
that create models on the remote actor rather than passing full models through Ray channels.
Args:
gen_model (LLMWrapperBase | Literal["from_collector"]): the generation model, or "from_collector" for lazy initialization.
Prefer using a model factory instead to avoid serialization issues.
ref_model (LLMWrapperBase | None): the reference model. Prefer using a model factory instead
to avoid serialization issues.
Keyword Args:
gen_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a generation model.
This allows for explicit resource control and avoids serialization issues.
ref_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a reference model.
This allows for explicit resource control and avoids serialization issues.
num_cpus (int, optional): Number of CPUs to allocate to the Ray actor. Defaults to 1.
num_gpus (int, optional): Number of GPUs to allocate to the Ray actor. Defaults to 0.
device (torch.device, optional): Device to use on the remote Ray actor for tensor operations.
The local Ray transform will handle CPU serialization and device restoration automatically.
Defaults to None.
actor_name (str, optional): Name of the Ray actor to use. If provided, the actor will be reused if it already exists.
**kwargs: Additional keyword arguments to pass to RetrieveKL.
Note:
When using model factories, the corresponding model arguments (gen_model, ref_model) should be None.
Model factories are preferred for large models to avoid serialization overhead.
Examples:
>>> # Option 1: Using model factories for explicit resource control
>>> def create_gen_model():
... return TransformersWrapper(model, tokenizer=tokenizer, generate=False, return_log_probs=True)
>>> def create_ref_model():
... return TransformersWrapper(ref_model, tokenizer=tokenizer, generate=False, return_log_probs=True)
>>> transform = RayRetrieveKL(
... gen_model=None, ref_model=None,
... gen_model_factory=create_gen_model,
... ref_model_factory=create_ref_model,
... num_gpus=1,
... device=torch.device("cuda")
... )
>>> # Option 2: Pass models directly (Ray handles serialization)
>>> transform = RayRetrieveKL(gen_model=gen_model, ref_model=ref_model, device=torch.device("cuda"))
"""
def __init__(
self,
gen_model: LLMWrapperBase | Literal["from_collector"] | None = "from_collector",
ref_model: LLMWrapperBase | None = None,
*,
gen_model_factory: Callable[[], LLMWrapperBase] | None = None,
ref_model_factory: Callable[[], LLMWrapperBase] | None = None,
num_cpus: int | None = None,
num_gpus: int = 0,
device: DEVICE_TYPING | None = None,
actor_name: str | None = None,
**kwargs,
):
# Validate arguments: models and factories should not both be provided
if gen_model is not None and gen_model_factory is not None:
raise ValueError(
"Cannot provide both 'gen_model' and 'gen_model_factory'. Choose one."
)
if ref_model is not None and ref_model_factory is not None:
raise ValueError(
"Cannot provide both 'ref_model' and 'ref_model_factory'. Choose one."
)
# Store creation parameters for actor creation
self._gen_model = gen_model
self._ref_model = ref_model
self._gen_model_factory = gen_model_factory
self._ref_model_factory = ref_model_factory
self._creation_kwargs = kwargs
# Store device separately for passing to remote actor
self._remote_device = device
# Default num_cpus
if num_cpus is None:
num_cpus = 1
# Call parent constructor without device (Ray transform handles CPU/device mapping)
super().__init__(
num_cpus=num_cpus,
num_gpus=num_gpus,
device=None, # Don't store device locally
actor_name=actor_name,
**kwargs,
)
def _create_actor(self, **kwargs):
"""Create the remote RetrieveKL actor."""
# Create the remote RetrieveKL with resource specifications
RemoteRetrieveKL = self._ray.remote(
num_cpus=self._num_cpus, num_gpus=self._num_gpus
)(RetrieveKL)
if self._actor_name is not None:
RemoteRetrieveKL = RemoteRetrieveKL.options(name=self._actor_name)
# Determine how to create models on the remote actor
gen_model_arg = self._gen_model
ref_model_arg = self._ref_model
# If we have factories, we'll pass them and set models to None
creation_kwargs = self._creation_kwargs.copy()
if self._gen_model_factory is not None:
creation_kwargs["gen_model_factory"] = self._gen_model_factory
gen_model_arg = None
if self._ref_model_factory is not None:
creation_kwargs["ref_model_factory"] = self._ref_model_factory
ref_model_arg = None
# Pass device to the remote actor
if self._remote_device is not None:
creation_kwargs["device"] = self._remote_device
# Create the shared actor
actor = RemoteRetrieveKL.remote(
gen_model=gen_model_arg, ref_model=ref_model_arg, **creation_kwargs
)
return actor
def __repr__(self):
"""String representation."""
try:
if hasattr(self, "_actor") and self._actor is not None:
return self._ray.get(self._actor.__repr__.remote())
else:
return "RayRetrieveKL(actor=None)"
except Exception:
return f"RayRetrieveKL(actor={getattr(self, '_actor', 'None')})"
[docs]class RetrieveKL(Compose, metaclass=_RayServiceMetaClass):
"""A transform to retrieve the KL divergence between two models' log-probabilities.
This transform combines two :class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb` instances
with a :class:`~torchrl.envs.llm.transforms.kl.KLComputation` to compute KL divergence
between a generation model and a reference model.
.. note::
Both gen_model and ref_model must use the same pad_output value (True or False), otherwise KL computation will fail.
Args:
gen_model (LLMWrapperBase): the generation model, wrapped in such a way that it does not generate but computes the log-probs.
In cases where the transform is used within a :class:`~torchrl.collectors.llm.LLMCollector` run on a remote worker, the
policy may not be available ahead of time. In this case, the `gen_model` can be set to `"from_collector"` (default) to retrieve the
policy from the collector. See :meth:`~torchrl.modules.llm.policies.LLMWrapperBase.get_new_version` for more details
about generating a new version of the policy to gather the log-probs.
ref_model (LLMWrapperBase): the reference model, wrapped in such a way that it does not generate but computes the log-probs.
Keyword Args:
gen_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a generation model.
This allows for explicit resource control and avoids serialization issues when using Ray.
ref_model_factory (Callable[[], LLMWrapperBase], optional): A callable that returns a reference model.
This allows for explicit resource control and avoids serialization issues when using Ray.
assistant_only (bool): whether to only retrieve the log-probs of the assistant tokens (i.e., steps of history
where the role is `"assistant"`). Defaults to `True`.
.. note:: When `assistant_only=True`, both models must have `input_mode='history'` to properly identify assistant tokens.
For other input modes (`"text"` or `"tokens"`), set `assistant_only=False`.
This ensures users are conscious of the limitation that assistant token identification requires structured conversation history.
gen_log_probs_full_key (str): the key where the log-probs of the generation model are stored. Defaults to `("log_probs", "full")`.
ref_log_probs_full_key (str): the key where the log-probs of the reference model are stored. Defaults to `("ref_log_probs", "full")`.
history_key (str): the key where the history is stored. Defaults to `"history"`.
tokenizer_kwargs (dict): the keyword arguments to pass to the tokenizer to be used to apply the chat template to the history when `assistant_only` is `True`.
To control the tokenization in the actor, pass the tokenizer kwargs to the actor constructor.
Defaults to `{"return_assistant_tokens_mask": True, "tokenize": True, "return_tensors": "pt", "padding": True, "add_generation_prompt": False}`.
detach (bool): whether to exclude the log-probs from the gradient computation. Defaults to `True`.
device (torch.device): the device to cast the tensors to. This is not the device of the specs, but the device
onto which the tensors will be moved. It allows to keep the model on a different device
than the upcoming data itself. When using Ray service, this device will be used on the remote actor.
Defaults to `None`.
tokenizer (transformers.AutoTokenizer): the tokenizer to be used to tokenize the input and compute the assitant mask. If not provided, the tokenizer will be inferred from the `actor`.
padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
kl_key (NestedKey): the key where the KL divergence is stored. Defaults to `"kl_penalty"`.
add_to_reward (bool): whether to add the KL divergence to the reward. Defaults to `True`.
coeff (float): the coefficient for the KL term when adding to reward. Defaults to `1.0`.
padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
use_ray_service (bool, optional): if ``True``, returns a :class:`RayRetrieveKL` instance instead,
which creates a Ray actor for shared KL computation across multiple environments.
Defaults to ``False``.
actor_name (str, optional): the name of the Ray actor to use. Defaults to `None`.
**kwargs: additional arguments to pass to the `RetrieveLogProb` transform.
Examples:
>>> from torchrl.data.llm import History
>>> from torchrl.modules.llm import TransformersWrapper
>>> from torchrl.modules.llm.policies import ChatHistory
>>> from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
>>> from tensordict import TensorDict, set_list_to_stack
>>> import torch
>>>
>>> # Set up list to stack for History
>>> set_list_to_stack(True).set()
>>>
>>> # Create chat data
>>> 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)
>>> print(f"Created history with shape: {history.shape}")
Created history with shape: torch.Size([2, 3])
>>>
>>> # Setup tokenizer and model
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
>>> tokenizer.pad_token = tokenizer.eos_token
>>> model = OPTForCausalLM(OPTConfig()).eval()
>>>
>>> # Create generation and reference models
>>> gen_model = TransformersWrapper(
... model,
... tokenizer=tokenizer,
... input_mode="history",
... generate=False,
... return_log_probs=True,
... pad_output=True,
... log_probs_key="gen_log_probs",
... )
>>> ref_model = TransformersWrapper(
... model,
... tokenizer=tokenizer,
... input_mode="history",
... generate=False,
... return_log_probs=True,
... pad_output=True,
... log_probs_key="ref_log_probs",
... )
>>>
>>> # Create RetrieveKL transform
>>> transform = RetrieveKL(
... gen_model=gen_model,
... ref_model=ref_model,
... assistant_only=True,
... tokenizer=tokenizer,
... )
>>>
>>> # Prepare data with next tensordict using ChatHistory
>>> chat_history = ChatHistory(full=history)
>>> next_td = TensorDict(history=chat_history, batch_size=(2,))
>>> data = TensorDict(history=chat_history, next=next_td, batch_size=(2,))
>>>
>>> # Apply transform
>>> result = transform(data)
>>> kl = result["next"].get("kl_penalty")
>>> print(f"KL shape: {kl.shape}")
KL shape: torch.Size([2, 26])
Note:
**Input Mode Compatibility:**
- When `assistant_only=True`, both models must have `input_mode='history'` to properly identify assistant tokens.
- When `assistant_only=False`, the transform works with any input mode (`"history"`, `"text"`, or `"tokens"`).
- This design ensures users are conscious of the limitation that assistant token identification requires structured conversation history.
.. seealso::
:class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb`: The base transform for retrieving log-probabilities from a single model.
:class:`~torchrl.envs.llm.transforms.kl.KLComputation`: The transform that computes KL divergence between two log-prob tensors.
:class:`~torchrl.envs.llm.transforms.kl.KLRewardTransform`: A legacy transform for KL reward computation (use `RetrieveKL` instead).
"""
_RayServiceClass = RayRetrieveKL
def __init__(
self,
gen_model: LLMWrapperBase | Literal["from_collector"] = "from_collector",
ref_model: LLMWrapperBase | None = None,
*,
gen_model_factory: Callable[[], LLMWrapperBase] | None = None,
ref_model_factory: Callable[[], LLMWrapperBase] | None = None,
assistant_only: bool = True,
history_key: str = "history",
tokenizer_kwargs: dict[str, Any] | None = None,
detach: bool = True,
device: torch.device | None = None,
tokenizer: transformers.AutoTokenizer | None = None,
padding_side: str = "left",
gen_log_probs_full_key: NestedKey = ("log_probs", "full"),
ref_log_probs_full_key: NestedKey = ("ref_log_probs", "full"),
kl_key: NestedKey = "kl_penalty",
add_to_reward: bool = True,
coeff: float = 1.0,
use_ray_service: bool = False,
**kwargs,
):
# Handle model factories - create models if factories are provided
if gen_model_factory is not None:
if gen_model is not None and gen_model != "from_collector":
raise ValueError(
"Cannot provide both 'gen_model' and 'gen_model_factory'. Choose one."
)
gen_model = gen_model_factory()
if ref_model_factory is not None:
if ref_model is not None:
raise ValueError(
"Cannot provide both 'ref_model' and 'ref_model_factory'. Choose one."
)
ref_model = ref_model_factory()
if isinstance(gen_model, str) and gen_model == "from_collector":
# Lazy init
self._initialized = False
self._init_params = {
"ref_model": ref_model,
"gen_model_factory": gen_model_factory,
"ref_model_factory": ref_model_factory,
"assistant_only": assistant_only,
"history_key": history_key,
"tokenizer_kwargs": tokenizer_kwargs,
"detach": detach,
"device": device,
"tokenizer": tokenizer,
"gen_log_probs_full_key": gen_log_probs_full_key,
"ref_log_probs_full_key": ref_log_probs_full_key,
"kl_key": kl_key,
"add_to_reward": add_to_reward,
"coeff": coeff,
"padding_side": padding_side,
**kwargs,
}
super().__init__()
return
self._initialized = True
# Check pad_output consistency if both models are provided
if hasattr(gen_model, "pad_output") and hasattr(ref_model, "pad_output"):
if gen_model.pad_output != ref_model.pad_output:
raise ValueError(
f"pad_output mismatch: gen_model.pad_output={gen_model.pad_output}, "
f"ref_model.pad_output={ref_model.pad_output}. "
"Both models must use the same padding strategy for KL computation."
)
if not getattr(gen_model, "return_log_probs", True):
raise ValueError(
"The generation model must have `return_log_probs=True` to use the `RetrieveKL` transform."
)
elif getattr(gen_model, "generate", False):
raise ValueError(
"The generation model must have `generate=False` to use the `RetrieveKL` transform."
)
if not getattr(ref_model, "return_log_probs", True):
raise ValueError(
"The reference model must have `return_log_probs=True` to use the `RetrieveKL` transform."
)
elif getattr(ref_model, "generate", False):
raise ValueError(
"The reference model must have `generate=False` to use the `RetrieveKL` transform."
)
if getattr(gen_model, "log_probs_key", "gen_log_probs") == getattr(
ref_model, "log_probs_key", "log_probs"
):
raise ValueError(
"The generation and reference models must have different `log_prob_key` values to use the `RetrieveKL` transform."
)
if gen_model is None:
raise ValueError("gen_model cannot be None when not using 'from_collector'")
if ref_model is None:
raise ValueError("ref_model cannot be None")
t1 = RetrieveLogProb(
gen_model,
log_probs_full_key=gen_log_probs_full_key,
assistant_only=assistant_only,
tokenizer_kwargs=tokenizer_kwargs,
detach=detach,
device=device,
tokenizer=tokenizer,
padding_side=padding_side,
**kwargs,
)
t2 = RetrieveLogProb(
ref_model,
log_probs_full_key=ref_log_probs_full_key,
assistant_only=assistant_only,
tokenizer_kwargs=tokenizer_kwargs,
detach=detach,
device=device,
tokenizer=tokenizer,
padding_side=padding_side,
**kwargs,
)
t3 = KLComputation(
gen_log_probs_full_key=gen_log_probs_full_key,
ref_log_probs_full_key=ref_log_probs_full_key,
kl_key=kl_key,
add_to_reward=add_to_reward,
coeff=coeff,
)
super().__init__(t1, t2, t3)
def _init_deferred(self):
torchrl_logger.info("Initializing RetrieveKL transform")
container = self.container
if container is None:
# also logging, since this will be sometimes hidden within the AttributeError
torchrl_logger.warning(
"The container is not set. Please set the container before calling this method."
)
raise ValueError(
"The container is not set. Please set the container before calling this method."
)
container.empty_cache()
self.empty_cache()
collector = self.collector
if collector is None:
# also logging, since this will be sometimes hidden within the AttributeError
torchrl_logger.warning(
"The collector is not set. Please set the collector before calling this method."
)
raise ValueError(
"The collector is not set. Please set the collector before calling this method."
)
ref_model = self._init_params["ref_model"]
pad_output = getattr(ref_model, "pad_output", None)
gen_log_probs_full_key = self._init_params["gen_log_probs_full_key"]
if (
not isinstance(gen_log_probs_full_key, tuple)
or gen_log_probs_full_key[-1] != "full"
):
raise ValueError(
f"The gen_log_probs_full_key {gen_log_probs_full_key} is not a tuple or does not end with 'full'. "
"This may cause issues with the KL computation. "
"Please use a tuple with the log_probs_key and 'full' as the last element."
)
log_probs_key = gen_log_probs_full_key[:-1]
gen_model = collector.policy.get_new_version(
generate=False,
return_log_probs=True,
log_probs_key=log_probs_key,
input_mode=ref_model.input_mode,
input_key=(ref_model.input_mode, "full"),
pad_output=pad_output, # Pass pad_output from ref_model
)
# Create the transforms manually instead of calling __init__
t1 = RetrieveLogProb(
gen_model,
log_probs_full_key=gen_log_probs_full_key,
assistant_only=self._init_params["assistant_only"],
tokenizer_kwargs=self._init_params["tokenizer_kwargs"],
detach=self._init_params["detach"],
device=self._init_params["device"],
tokenizer=self._init_params["tokenizer"],
padding_side=self._init_params["padding_side"],
)
ref_log_probs_full_key = self._init_params["ref_log_probs_full_key"]
if (
not isinstance(ref_log_probs_full_key, tuple)
or ref_log_probs_full_key[-1] != "full"
):
raise ValueError(
f"The ref_log_probs_full_key {ref_log_probs_full_key} is not a tuple or does not end with 'full'. "
"This may cause issues with the KL computation. "
"Please use a tuple with the log_probs_key and 'full' as the last element."
)
t2 = RetrieveLogProb(
ref_model,
log_probs_full_key=ref_log_probs_full_key,
assistant_only=self._init_params["assistant_only"],
tokenizer_kwargs=self._init_params["tokenizer_kwargs"],
detach=self._init_params["detach"],
device=self._init_params["device"],
tokenizer=self._init_params["tokenizer"],
padding_side=self._init_params["padding_side"],
)
t3 = KLComputation(
gen_log_probs_full_key=gen_log_probs_full_key,
ref_log_probs_full_key=ref_log_probs_full_key,
kl_key=self._init_params["kl_key"],
add_to_reward=self._init_params["add_to_reward"],
coeff=self._init_params["coeff"],
)
# Replace the transforms in the Compose
self.transforms.extend([t1, t2, t3])
del self._init_params
self._initialized = True
torchrl_logger.info("Successfully initialized")
def _step(
self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
) -> TensorDictBase:
if not self._initialized:
self._init_deferred()
return super()._step(tensordict, next_tensordict)
def _reset(
self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
) -> TensorDictBase:
if not self._initialized:
self._init_deferred()
return super()._reset(tensordict, tensordict_reset)
[docs] def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
if not self._initialized:
self._init_deferred()
return super().forward(tensordict)
[docs] def transform_observation_spec(self, observation_spec: Composite) -> Composite:
if not self._initialized:
self._init_deferred()
return super().transform_observation_spec(observation_spec)
[docs] def transform_reward_spec(self, reward_spec: Composite) -> Composite:
if not self._initialized:
self._init_deferred()
return super().transform_reward_spec(reward_spec)
def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
if not self._initialized:
self._init_deferred()
return super()._inv_call(tensordict)
[docs] def transform_action_spec(self, action_spec: Composite) -> Composite:
if not self._initialized:
self._init_deferred()
return super().transform_action_spec(action_spec)
[docs] def transform_input_spec(self, input_spec: Composite) -> Composite:
if not self._initialized:
self._init_deferred()
return super().transform_input_spec(input_spec)
[docs] def transform_output_spec(self, output_spec: Composite) -> Composite:
if not self._initialized:
self._init_deferred()
return super().transform_output_spec(output_spec)
[docs] def transform_state_spec(self, state_spec: Composite) -> Composite:
if not self._initialized:
self._init_deferred()
return super().transform_state_spec(state_spec)
[docs]class KLComputation(Transform):
"""A transform to compute KL divergence between two log-prob tensors and optionally add it to the reward.
This transform computes KL divergence between generation and reference log-probabilities
and can optionally subtract it from the reward (for KL penalty). It's designed to work
with the :class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb` and :class:`~torchrl.envs.llm.transforms.kl.RetrieveKL` transforms.
.. note::
Both input log-prob tensors must use the same padding strategy (pad_output) for correct KL computation.
Args:
gen_log_probs_full_key (NestedKey): the key where the generation model log-probs are stored.
Defaults to `("gen_log_probs", "full")`.
ref_log_probs_full_key (NestedKey): the key where the reference model log-probs are stored.
Defaults to `("ref_log_probs", "full")`.
kl_key (NestedKey): the key where the KL divergence is stored. Defaults to `"kl_penalty"`.
add_to_reward (bool): whether to add the KL divergence to the reward. Defaults to `True`.
coeff (float): the coefficient for the KL term when adding to reward. Defaults to `1.0`.
padding_side (str): the side of the padding when using pad_sequence. Defaults to `"left"`.
Examples:
>>> from tensordict import TensorDict
>>> import torch
>>>
>>> # Create sample log-probs
>>> gen_log_probs = torch.randn(2, 10) # 2 samples, 10 tokens each
>>> ref_log_probs = torch.randn(2, 10)
>>>
>>> # Create data with next tensordict
>>> next_td = TensorDict(
... {
... ("gen_log_probs", "full"): gen_log_probs,
... ("ref_log_probs", "full"): ref_log_probs,
... "reward": torch.randn(2, 10, 1),
... },
... batch_size=(2,)
... )
>>> data = TensorDict(next=next_td, batch_size=(2,))
>>>
>>> # Create KLComputation transform
>>> kl_transform = KLComputation(
... gen_log_probs_key=("gen_log_probs", "full"),
... ref_log_probs_key=("ref_log_probs", "full"),
... kl_key="kl_penalty",
... add_to_reward=True,
... coef=1.0,
... )
>>>
>>> # Apply transform
>>> result = kl_transform(data)
>>> kl = result["next"].get("kl_penalty")
>>> print(f"KL shape: {kl.shape}")
KL shape: torch.Size([2, 10])
.. seealso::
:class:`~torchrl.envs.llm.transforms.kl.RetrieveLogProb`: The base transform for retrieving log-probabilities from a single model.
:class:`~torchrl.envs.llm.transforms.kl.RetrieveKL`: A higher-level transform that combines two `RetrieveLogProb` instances with `KLComputation`.
:class:`~torchrl.envs.llm.transforms.kl.KLRewardTransform`: A legacy transform for KL reward computation (use `RetrieveKL` instead).
"""
def __init__(
self,
gen_log_probs_full_key: NestedKey = ("log_probs", "full"),
ref_log_probs_full_key: NestedKey = ("ref_log_probs", "full"),
*,
kl_key: NestedKey = "kl_penalty",
add_to_reward: bool = True,
coeff: float = 1.0,
padding_side: str = "left",
):
in_keys = [gen_log_probs_full_key, ref_log_probs_full_key]
if add_to_reward:
in_keys.append("reward")
out_keys = [kl_key]
if add_to_reward:
out_keys.append("reward")
super().__init__(in_keys=in_keys, out_keys=out_keys)
self.gen_log_probs_full_key = gen_log_probs_full_key
self.ref_log_probs_full_key = ref_log_probs_full_key
self.kl_key = kl_key
self.add_to_reward = add_to_reward
self.coeff = coeff
self.padding_side = padding_side
[docs] def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
next_td = tensordict.get("next")
has_next_td = True
if next_td is None:
next_td = tensordict
has_next_td = False
next_td = self._step(tensordict, next_td)
if has_next_td:
return tensordict.set("next", next_td)
return next_td
def _step(
self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
) -> TensorDictBase:
# Get log-probs
gen_log_probs = next_tensordict.get(self.gen_log_probs_full_key, as_list=True) # type: ignore[misc]
ref_log_probs = next_tensordict.get(self.ref_log_probs_full_key, as_list=True) # type: ignore[misc]
if gen_log_probs is None or ref_log_probs is None:
raise ValueError(
f"Log-probs not found. Expected keys: {self.gen_log_probs_key}, {self.ref_log_probs_key}"
)
# Debug: Check lengths and shapes
if len(gen_log_probs) != len(ref_log_probs):
raise ValueError(
f"Batch size mismatch: gen_log_probs has {len(gen_log_probs)} samples, ref_log_probs has {len(ref_log_probs)} samples"
)
# Check individual sequence lengths
for i, (gen_lp, ref_lp) in enumerate(_zip_strict(gen_log_probs, ref_log_probs)):
if gen_lp.shape != ref_lp.shape:
raise ValueError(
f"Sample {i} has different shapes: gen_log_probs[{i}].shape={gen_lp.shape}, ref_log_probs[{i}].shape={ref_lp.shape}"
)
# Compute KL divergence: KL(p||q) = E_p[log p - log q]
# Here gen_log_probs = log p, ref_log_probs = log q
kl = [
gen_lp - ref_lp
for gen_lp, ref_lp in _zip_strict(gen_log_probs, ref_log_probs)
]
kl = torch.nested.as_nested_tensor(kl, layout=torch.strided)
next_tensordict.set(self.kl_key, kl)
# Add to reward if requested
if self.add_to_reward:
reward = next_tensordict.get("reward", as_list=True) # type: ignore[misc]
if reward is not None:
if isinstance(reward, list):
if reward[0].ndim != kl[0].ndim + 1:
raise ValueError(
f"The rewards have shape {reward[0].shape} but the kl has shape {kl[0].shape}. "
f"The rewards should have one more dimension than the KL."
)
reward = [
r - self.coeff * k.unsqueeze(-1)
for r, k in _zip_strict(reward, kl)
]
next_tensordict.set(
"reward",
torch.nested.as_nested_tensor(reward, layout=torch.strided),
)
else:
if reward.ndim != kl.ndim + 1:
raise ValueError(
f"The rewards have shape {reward.shape} but the kl has shape {kl.shape}. "
f"The rewards should have one more dimension than the KL."
)
reward = reward - self.coeff * kl.unsqueeze(-1)
next_tensordict.set("reward", reward)
return next_tensordict
[docs] def transform_observation_spec(self, observation_spec: Composite) -> Composite:
# Add kl to observation spec
observation_spec[self.kl_key] = Unbounded(
device=observation_spec.device,
shape=observation_spec.shape,
)
return observation_spec
[docs] def transform_reward_spec(self, reward_spec: Composite) -> Composite:
# Optionally adjust reward spec if KL is added to reward
if self.add_to_reward:
shape = reward_spec["reward"].shape
# For LLMs, the shape of the reward is (batch, -1, 1)
shape = torch.Size((*shape, -1, 1))
reward_spec["reward"] = reward_spec["reward"].clone()
reward_spec["reward"].shape = shape
return reward_spec
