torchrl-nightly 2025.6.20__cp313-cp313-manylinux1_x86_64.whl → 2025.6.22__cp313-cp313-manylinux1_x86_64.whl

torchrl/objectives/llm/sft.py

@@ -0,0 +1,465 @@
+# 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
+import warnings
+
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+from tensordict import NestedKey, TensorClass, TensorDictBase
+from tensordict.nn import TensorDictModule
+from tensordict.utils import _zip_strict
+from torchrl.data import History
+from torchrl.modules.llm.policies.transformers_wrapper import TransformersWrapper
+from torchrl.objectives.common import LossModule
+
+
+def sft_loss(summed_log_probs: torch.Tensor, reduction: str) -> torch.Tensor:
+    """Compute the SFT loss."""
+    if reduction == "mean":
+        loss = -summed_log_probs.mean()
+    elif reduction == "sum":
+        loss = -summed_log_probs.sum()
+    elif reduction == "none":
+        loss = -summed_log_probs
+    else:
+        raise ValueError(f"Invalid reduction: {reduction}.")
+    return loss
+
+
+def minor_sft_loss(
+    log_probs: torch.Tensor,
+    ref_log_probs: torch.Tensor,
+    beta: float,
+    reduction: str,
+) -> torch.Tensor:
+    """Compute the MinorSFT loss.
+
+    This loss is inspired by DPO and is designed to be less aggressive than standard SFT.
+    It computes ``-log_sigmoid(beta * (log_probs - ref_log_probs))``.
+
+    Args:
+        log_probs (torch.Tensor): The log probabilities from the model being trained.
+        ref_log_probs (torch.Tensor): The log probabilities from the reference model.
+        beta (float): The beta parameter from DPO.
+        reduction (str): The reduction to apply to the loss.
+
+    Returns:
+        The MinorSFT loss.
+
+    References:
+        - Shiming Xie, Hong Chen, Fred Yu, Zeye Sun, Xiuyu Wu, 2024.
+          `"Minor SFT loss for LLM fine-tune to increase performance and reduce model deviation" <https://arxiv.org/abs/2408.10642>`_
+    """
+    if log_probs.shape != ref_log_probs.shape:
+        raise ValueError(
+            f"Current log probabilities and reference log probabilities have different shapes: {log_probs.shape=} vs {ref_log_probs.shape=}."
+        )
+    loss = -torch.nn.functional.logsigmoid(beta * (log_probs - ref_log_probs))
+    if reduction == "mean":
+        return loss.mean()
+    if reduction == "sum":
+        return loss.sum()
+    if reduction == "none":
+        return loss
+    raise ValueError(f"Invalid reduction: {reduction}")
+
+
+class SFTLossOutput(TensorClass["nocast"]):
+    """SFT Loss Output.
+
+    Attributes:
+        loss_sft (torch.Tensor): The loss for the SFT objective.
+        loss_kl_to_ref (torch.Tensor | None): The loss for the KL divergence to the reference model.
+        kl_to_ref (torch.Tensor | None): The KL divergence to the reference model.
+
+    .. note::
+        The loss components are kept separate to allow for logging and visualization.
+        Before backpropagation, the loss components are to be summed together. Since non-loss components are not differentiable
+        when the loss is constructed via :class:`~torchrl.objectives.llm.sft.SFTLoss`, summing
+        the :class:`~torchrl.objectives.llm.sft.SFTLossOutput` directly is a proper way of obtaining the total loss.
+
+            >>> loss_fn = SFTLoss(...)
+            >>> loss_output = loss_fn(td)
+            >>> loss = loss_output.loss_sft + loss_output.loss_kl_to_ref
+            >>> loss.backward()
+            >>> # or equivalently
+            >>> loss = loss_fn(td)
+            >>> loss.sum(reduce=True).backward()
+    """
+
+    loss_sft: torch.Tensor
+    loss_kl_to_ref: torch.Tensor | None = None
+    kl_to_ref: torch.Tensor | None = None
+
+
+class SFTLoss(LossModule):
+    r"""Supervised fine-tuning loss.
+
+    Args:
+        actor_network (TensorDictModule): the actor network. Usually a :class:`~torchrl.modules.llm.TransformersWrapper` instance,
+            with `return_log_prob=True` and `from_text=True`.
+        tokenizer (`Tokenizer`): 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_network`.
+        tokenizer_kwargs (dict, optional): keyword arguments to pass to the tokenizer during :meth:`~torchrl.data.llm.chat.History.apply_chat_template`.
+            This can be used to override arguments such as the `chat_template` or `chat_template_name`.
+        reduction (Literal["mean", "sum", "none"], optional): the reduction to apply to the loss. Defaults to `"mean"`.
+        normalize_by_seq_length (bool, optional): whether to normalize the loss by the sequence length. Defaults to `True`.
+        kl_to_ref_coeff (float | None, optional): coefficient for KL divergence to reference model. Defaults to `None`.
+        loss_function (Literal["sft", "minor_sft"], optional): The loss function to use. Defaults to `"sft"`.
+        beta (float, optional): The beta parameter for MinorSFT loss. This is only used when `loss_function` is `"minor_sft"`.
+            Higher values of beta make the loss more aggressive (pushes the model to generate responses further from the reference model):
+
+            .. math::
+                \text{loss} = -\log\sigma(\beta \cdot (\text{log_probs} - \text{ref_log_probs}))
+
+            Defaults to `0.1`.
+        device (torch.device | None, optional): the device to use for the loss, when tokenizing the input. Defaults to `None`.
+
+    .. note::
+        The input tensordict is expected to contain the following keys by default:
+            - ``("next", "history")``: The chat history
+            - ``("next", "ref_log_prob")`` (optional): Reference model log probabilities, required if kl_to_ref_coeff is set
+
+        These keys can be customized using the ``set_keys()`` method.
+
+    .. seealso:: :class:`~torchrl.envs.llm.transforms.RetrieveLogProb` for the KL divergence computation.
+
+    References:
+        - Shiming Xie, Hong Chen, Fred Yu, Zeye Sun, Xiuyu Wu, 2024.
+          `"Minor SFT loss for LLM fine-tune to increase performance and reduce model deviation" <https://arxiv.org/abs/2408.10642>`_
+
+    Examples:
+        >>> from torchrl.data.llm.chat import History, _CHAT_TEMPLATES
+        >>> from torchrl.modules.llm import TransformersWrapper
+        >>> from torchrl.objectives.llm.sft import SFTLoss
+        >>> from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
+        >>> from tensordict import TensorDict, lazy_stack
+        >>> import torch
+        >>>
+        >>> # 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)
+        >>>
+        >>> # Setup tokenizer and model
+        >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
+        >>> tokenizer.pad_token = tokenizer.eos_token
+        >>> tokenizer.chat_template = _CHAT_TEMPLATES["chatml_format"]
+        >>> model = OPTForCausalLM(OPTConfig()).eval()
+        >>>
+        >>> # Create training and reference policies
+        >>> policy_train = TransformersWrapper(
+        ...     model,
+        ...     tokenizer=tokenizer,
+        ...     generate=False,
+        ...     from_text=True,
+        ...     chat_template_name="qwen",
+        ... )
+        >>> policy_ref = TransformersWrapper(
+        ...     model,
+        ...     tokenizer=tokenizer,
+        ...     generate=False,
+        ...     from_text=True,
+        ...     return_log_probs=True,
+        ...     chat_template_name="qwen",
+        ... )
+        >>>
+        >>> # Create the RetrieveLogProb transform
+        >>> transform = RetrieveLogProb(
+        ...     policy_ref,
+        ...     assistant_only=True,
+        ...     tokenizer_kwargs={"chat_template_name": "qwen"},
+        ...     tokenizer=tokenizer,
+        ... )
+        >>>
+        >>> # Prepare data
+        >>> text = history[:, :-1].apply_chat_template(
+        ...     tokenizer=tokenizer, chat_template_name="qwen", add_generation_prompt=True
+        ... )
+        >>> text_response = history.apply_chat_template(
+        ...     tokenizer=tokenizer, chat_template_name="qwen", add_generation_prompt=False
+        ... )
+        >>> text_response = [
+        ...     txt[len(txt_start):] for txt, txt_start in zip(text_response, text)
+        ... ]
+        >>> td = TensorDict(
+        ...     text=text,
+        ...     text_response=text_response,
+        ...     history=history,
+        ...     next=TensorDict(
+        ...         reward=torch.randn(2, 1),
+        ...         done=torch.zeros(2, dtype=torch.bool),
+        ...         history=history,
+        ...     ),
+        ...     batch_size=(2,),
+        ... )
+        >>> data = lazy_stack(list(td.unbind(0)))
+        >>>
+        >>> # Apply the transform to get reference log probabilities
+        >>> data = transform(data)
+        >>> assert "ref_log_prob" in data["next"].keys()
+        >>>
+        >>> # Use with SFTLoss for KL regularization
+        >>> loss = SFTLoss(
+        ...     actor_network=policy_train,
+        ...     tokenizer=tokenizer,
+        ...     reduction="mean",
+        ...     normalize_by_seq_length=True,
+        ...     kl_to_ref_coeff=0.1,
+        ...     tokenizer_kwargs={"chat_template_name": "qwen"},
+        ...     loss_function="sft",
+        ... )
+        >>> loss_vals = loss(data)
+        >>> print(f"SFT Loss: {loss_vals.loss_sft.item():.4f}")
+        >>> print(f"KL to Reference Loss: {loss_vals.loss_kl_to_ref.item():.4f}")
+
+    """
+
+    @dataclass
+    class _AcceptedKeys:
+        """Maintains default values for all configurable tensordict keys.
+
+        This class defines which tensordict keys can be set using '.set_keys(key_name=key_value)' and their
+        default values.
+
+        Attributes:
+            history (NestedKey): The input tensordict key where the chat history is expected.
+                Defaults to ``("next", "history")``.
+            ref_log_prob (NestedKey): The input tensordict key where the reference model log probabilities are expected.
+                Only used when kl_to_ref_coeff is set. Defaults to ``("next", "ref_log_prob")``.
+            log_probs (NestedKey): The output tensordict key where the model's log probabilities will be written.
+                Defaults to ``"log_probs"``.
+        """
+
+        history: NestedKey = ("next", "history")
+        ref_log_prob: NestedKey = ("next", "ref_log_prob")
+        log_probs: NestedKey = "log_probs"
+
+    default_keys = _AcceptedKeys
+    tensor_keys: _AcceptedKeys
+
+    def __init__(
+        self,
+        actor_network: TensorDictModule | TransformersWrapper,
+        tokenizer: transformers.AutoTokenizer | None = None,  # noqa: F821
+        tokenizer_kwargs: dict | None = None,
+        reduction: Literal["mean", "sum", "none"] = "mean",
+        normalize_by_seq_length: bool = True,
+        kl_to_ref_coeff: float | None = None,
+        loss_function: Literal["sft", "minor_sft"] = "sft",
+        beta: float = 0.1,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+        self.in_keys = []
+        self.actor_network = actor_network
+        if tokenizer is None:
+            tokenizer = actor_network.tokenizer
+        self.tokenizer = tokenizer
+        if tokenizer_kwargs is None:
+            tokenizer_kwargs = {}
+        if tokenizer is None:
+            raise ValueError("Tokenizer must be provided.")
+        tokenizer_kwargs.setdefault("return_assistant_tokens_mask", True)
+        tokenizer_kwargs.setdefault("tokenize", True)
+        tokenizer_kwargs.setdefault("return_tensors", "pt")
+        tokenizer_kwargs.setdefault("padding", False)
+        tokenizer_kwargs.setdefault("add_generation_prompt", False)
+        self.tokenizer_kwargs = tokenizer_kwargs
+        self.reduction = reduction
+        self.normalize_by_seq_length = normalize_by_seq_length
+        self.kl_to_ref_coeff = kl_to_ref_coeff
+        self.loss_function = loss_function
+        if self.loss_function == "minor_sft" and kl_to_ref_coeff:
+            warnings.warn(
+                "kl_to_ref_coeff should not be set when using minor_sft loss, as KL regularization is implicit. Setting kl_to_ref_coeff to 0.0."
+            )
+            self.kl_to_ref_coeff = 0.0
+        self.beta = beta
+        self._set_in_keys()
+        self.device = device
+
+    def _set_in_keys(self) -> None:
+        """Sets the input keys for the loss module."""
+        in_keys = [self.tensor_keys.history]
+        if self.kl_to_ref_coeff is not None or self.loss_function == "minor_sft":
+            in_keys.append(self.tensor_keys.ref_log_prob)
+        self.in_keys = in_keys
+        self.out_keys = []  # Loss modules typically don't have out_keys
+
+    def _kl_to_ref(
+        self,
+        cur_log_prob: list[torch.Tensor],
+        ref_log_prob: list[torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute KL divergence to reference model.
+
+        Args:
+            cur_log_prob (List[torch.Tensor]): Log probabilities from current model. Must have shape [T] where T is the number of tokens in the assistant response.
+            ref_log_prob (List[torch.Tensor]): Log probabilities from reference model. Must have shape [T] where T is the number of tokens in the assistant response.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]: (KL loss term, KL penalty for logging)
+        """
+        # Apply mask
+        ref_log_prob = torch.cat(ref_log_prob)
+        cur_log_prob = torch.cat(cur_log_prob)
+        # ref_log_prob = ref_log_prob[mask]
+        # cur_log_prob = cur_log_prob[mask].squeeze()
+        if cur_log_prob.shape != ref_log_prob.shape:
+            raise ValueError(
+                f"Current log probabilities and reference log probabilities have different shapes: {cur_log_prob.shape=} vs {ref_log_prob.shape=}."
+            )
+        # Compute KL using same approximation as GRPO
+        diff = ref_log_prob - cur_log_prob
+
+        kl_penalty = (diff.expm1() - diff).mean()
+        return self.kl_to_ref_coeff * kl_penalty, kl_penalty
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        # Gather history
+        history: History = tensordict[self.tensor_keys.history]
+
+        # Apply tokenizer to history and gather mask
+        with torch.device(
+            self.device
+        ) if self.device is not None else contextlib.nullcontext():
+            token_struct = history.apply_chat_template(
+                tokenizer=self.tokenizer, **self.tokenizer_kwargs
+            )
+        if "assistant_masks" not in token_struct:
+            raise ValueError(
+                f"Assistant masks are not present in the token structure: {token_struct=}."
+            )
+        assistant_masks = token_struct.get(
+            "assistant_masks",
+            as_list=True,
+        )
+        assistant_masks = [mask.bool() for mask in assistant_masks]
+        attention_mask = token_struct.get("attention_mask", as_list=True)
+        attention_mask = [mask.bool() for mask in attention_mask]
+        assistant_masks = [
+            mask & a_mask for mask, a_mask in zip(assistant_masks, attention_mask)
+        ]
+
+        if not any(mask.any(-1).all() for mask in assistant_masks):
+            raise ValueError("Some inputs have no valid assistant masks.")
+        input_loss = tensordict.select(self.tensor_keys.history)
+        if (
+            isinstance(self.tensor_keys.history, tuple)
+            and self.tensor_keys.history[0] == "next"
+        ):
+            input_loss = input_loss["next"]
+
+        with torch.device(
+            self.device
+        ) if self.device is not None else contextlib.nullcontext():
+            output_loss = self.actor_network(input_loss)
+
+        # get log-probs
+        log_probs = output_loss.get(
+            self.tensor_keys.log_probs,
+            as_list=True,
+        )
+        # apply mask
+        if not all(
+            mask.shape == lp.shape
+            for mask, lp in _zip_strict(assistant_masks, log_probs)
+        ):
+            raise ValueError(
+                f"Assistant masks and log_probs have different shapes: {[mask.shape for mask in assistant_masks]} vs {[lp.shape for lp in log_probs]}. Tokens from current template: {[inp.shape for inp in token_struct.get('input_ids', as_padded_tensor=True)]}"
+            )
+
+        log_probs_masked = [
+            lp.masked_fill(~mask, 0.0)
+            for lp, mask in _zip_strict(log_probs, assistant_masks)
+        ]
+
+        # Sum log probs, optionally normalize by sequence length
+        summed_log_probs = torch.stack(
+            [lp.sum(tensordict.ndim - 1) for lp in log_probs_masked]
+        )
+        seq_lengths = torch.stack(
+            [mask.sum(tensordict.ndim - 1) for mask in assistant_masks]
+        )
+        if self.normalize_by_seq_length:
+            # Compute sequence lengths for normalization (number of assistant tokens)
+            summed_log_probs = summed_log_probs / seq_lengths.clamp(min=1)
+
+        # Compute main loss
+        if self.loss_function == "sft":
+            loss = sft_loss(summed_log_probs, self.reduction)
+            # Add KL divergence loss if reference model is provided
+            if self.kl_to_ref_coeff is not None:
+                ref_log_probs = tensordict.get(
+                    self.tensor_keys.ref_log_prob,
+                    default=None,
+                    as_list=True,
+                )
+                if ref_log_probs is None:
+                    raise ValueError(
+                        "Reference log probs not found in tensordict but kl_to_ref_coeff was set"
+                    )
+
+                loss_kl, kl_penalty = self._kl_to_ref(
+                    [lp[mask] for lp, mask in _zip_strict(log_probs, assistant_masks)],
+                    ref_log_probs,
+                )
+                output = SFTLossOutput(
+                    loss_sft=loss,
+                    loss_kl_to_ref=loss_kl,
+                    kl_to_ref=kl_penalty.detach(),
+                )
+            else:
+                output = SFTLossOutput(loss_sft=loss)
+        elif self.loss_function == "minor_sft":
+            ref_log_probs = tensordict.get(self.tensor_keys.ref_log_prob, as_list=True)
+            if ref_log_probs is None:
+                raise ValueError(
+                    f"Reference log probs not found at {self.tensor_keys.ref_log_prob=} in tensordict but loss_function is 'minor_sft'"
+                )
+
+            # we need to re-sum ref_log_probs as they are not summed per-sequence
+            summed_ref_log_probs = torch.stack([lp.sum() for lp in ref_log_probs]).to(
+                summed_log_probs.device
+            )
+            if self.normalize_by_seq_length:
+                summed_ref_log_probs = summed_ref_log_probs / seq_lengths.clamp(min=1)
+            loss = minor_sft_loss(
+                summed_log_probs, summed_ref_log_probs, self.beta, self.reduction
+            )
+            if self.kl_to_ref_coeff is not None:
+                with torch.no_grad():
+                    loss_kl, kl_penalty = self._kl_to_ref(
+                        [
+                            lp[mask]
+                            for lp, mask in _zip_strict(log_probs, assistant_masks)
+                        ],
+                        ref_log_probs,
+                    )
+                output = SFTLossOutput(
+                    loss_sft=loss,
+                    loss_kl_to_ref=loss_kl,
+                    kl_to_ref=kl_penalty.detach(),
+                )
+            else:
+                output = SFTLossOutput(loss_sft=loss)
+        else:
+            raise ValueError(f"Invalid loss function: {self.loss_function}")
+
+        return output

torchrl/version.py

@@ -1,2 +1,2 @@
-__version__ = '2025.6.20'
-git_version = 'e1e15d692a6df69bfdc80e85f45d37a4b967e625'
+__version__ = '2025.6.22'
+git_version = '77dbc6c9ffbce3d2ce3f26b659355cd46d8132c3'

{torchrl_nightly-2025.6.20.dist-info → torchrl_nightly-2025.6.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchrl-nightly
-Version: 2025.6.20
+Version: 2025.6.22
 Home-page: https://github.com/pytorch/rl
 Author: torchrl contributors
 Author-email: vmoens@fb.com

{torchrl_nightly-2025.6.20.dist-info → torchrl_nightly-2025.6.22.dist-info}/RECORD

@@ -3,11 +3,11 @@ build_tools/setup_helpers/__init__.py,sha256=7l8TvVqxKezgzKCLuRv20mvGLloprFVZYm8
 build_tools/setup_helpers/extension.py,sha256=4-PDLr-pw40bJnd9SfxnTaSjUyuXU_Tg8yOg69Kl0o4,5914
 torchrl/__init__.py,sha256=mhDBx2UIuBKc0gmi8dVNHokQ6tCbIovruZmyAxcSsy8,2938
 torchrl/_extension.py,sha256=z7wQ8i1iYWYcnygq_j0nq9sT-koY13tfHhTLNbMk17Q,2353
-torchrl/_torchrl.cpython-313-x86_64-linux-gnu.so,sha256=K2mDnBxDxSj0qflvFxCBKYjue5HjA5oOtzfr9s26sVk,21451208
+torchrl/_torchrl.cpython-313-x86_64-linux-gnu.so,sha256=BCB4sN_0LejjRpv5bAtaI_d135o_BAHu_XLlQdHGNQg,21451176
 torchrl/_utils.py,sha256=Cw5EG6x5oSZF1iE3YCs1a32VUKp0rTXIs2u67q9zKUI,41078
-torchrl/version.py,sha256=tueh-ZkmvytwgeWC2kBvr7G0Iu6dTR_gZ03kkRW7v8Q,83
+torchrl/version.py,sha256=-iU3qRfg2kpxfWLGxmIBA7z2dNl2hOX09zklpdgeA4Y,83
 torchrl/collectors/__init__.py,sha256=hJ3JD6shRku0BL6SzJQq44FZ5Q1RGR8LealFyU3FRn4,799
-torchrl/collectors/collectors.py,sha256=5yY-WMEAzj2vKBA0RxijFjBq7y5faQGY2BwkPfI5qSU,177457
+torchrl/collectors/collectors.py,sha256=CdTerIwhCTr6n5OoJLNad0bNQ5OLliPZFWkU18QBKSA,177625
 torchrl/collectors/utils.py,sha256=MlXrkYuDmV0Em-tVNQiLL32FWgPNDgceYYG_GgpiviA,11320
 torchrl/collectors/weight_update.py,sha256=nSIfs8ALsfggLoC2ylg1oOAqdGku1tt4e-50JCZJBww,21073
 torchrl/collectors/distributed/__init__.py,sha256=_24P0ALFunLhL-ls7EsssGUhJkZ_m3nw7krfMTwPqS0,705
@@ -18,12 +18,12 @@ torchrl/collectors/distributed/rpc.py,sha256=0xQDqKlvLmCb_2wL9oZojt4rONaSq09abPL
 torchrl/collectors/distributed/sync.py,sha256=oZW3nUYrUK52N6pMYX9M0WWhMeTzLl25maxM7X2G8Ec,27272
 torchrl/collectors/distributed/utils.py,sha256=MuxSeb4TkiyWJYyMyXWLgyCDgtgbGU6g8nNVf59xqCE,6464
 torchrl/collectors/llm/__init__.py,sha256=rx9DktowQ-gvFleb07US9d9WFc4aNG6zKpiOPSW4A7U,355
-torchrl/collectors/llm/base.py,sha256=wyZmNIZ_92lUkfZKgbCh8OXDUoHNhVn_7s4qizmrH58,20388
-torchrl/collectors/llm/ray_collector.py,sha256=cc1oZ1zh322lJL21bgJd4b6w9-QUQiCOx2i63zzgMLo,10948
+torchrl/collectors/llm/base.py,sha256=G6n2_U7CIr0BBUMbrOSv-AIRSxFFMqft_Ia3Ir3Ggks,20671
+torchrl/collectors/llm/ray_collector.py,sha256=1o9rbQtoJ48Ovo_YP76KQ-dLlKJt-bGdH2VMaK_-olg,11101
 torchrl/collectors/llm/utils.py,sha256=-KRSlOmjj34M0c3msP7yS_0DlLmqCijEbf_bADLjzuM,1148
 torchrl/collectors/llm/weight_update/__init__.py,sha256=bKjvD7yZG5VnHgvYc4EmKI1seK4FyMBKTqeLzkqR_3s,272
 torchrl/collectors/llm/weight_update/vllm.py,sha256=81ShmKzNjVIg7hxlPvLHhF-YqeXv98cIk0l6ByD-MDU,11276
-torchrl/data/__init__.py,sha256=RuBnwrzJqJZxU1drtdzUHdWTrZpL6z4SPLYBYM2AMqc,4769
+torchrl/data/__init__.py,sha256=oowsio6ZUOZnJV8JV43xgs17B37XO1yKAYIQPdk8yt0,4819
 torchrl/data/rlhf.py,sha256=JUmdYBWgkN229DwpXuDrhy9ddjduNvU2kyHzHR6MoA0,963
 torchrl/data/tensor_specs.py,sha256=rfuYM9WLUnF4vHwM4opvypShZ3RN7954WhiPMyG3CSU,247841
 torchrl/data/utils.py,sha256=attuNwzfgjszyp0lJSrV06f2peX3r0qTjRZWEwfl6Yg,12108
@@ -39,12 +39,13 @@ torchrl/data/datasets/openx.py,sha256=QXjJPZHoRhefVux00iAL-g4spynrWjmI_M2IuaQ8TA
 torchrl/data/datasets/roboset.py,sha256=rLPdyEQI9yEibXU6SZFA0YD79EGFKY8o5oyUlvcn4aM,16648
 torchrl/data/datasets/utils.py,sha256=nAFDTlBIPyEoPoJC-Hc_fcOhzE7UZQE4BwKxq15Vhvk,339
 torchrl/data/datasets/vd4rl.py,sha256=z90MqrxKzod8TPGK0uzkC6vw5wQIE4cgrDAC4e72jyk,18262
-torchrl/data/llm/__init__.py,sha256=FWApDyEPlyE7jA7CRMmGpxnLYQ4ZKKAdCbbpmOMf-OU,908
-torchrl/data/llm/chat.py,sha256=aGnrefzwxvEhuPlTHmLPHyzqG2xWbMnzwjWBfoldNIM,30817
+torchrl/data/llm/__init__.py,sha256=By2FWnjqADPmHnNXh6DVLQ9CYPj51gn3HxPW_DYPMyc,971
+torchrl/data/llm/chat.py,sha256=K5Cuw4GHSJWGg5vXwGyV9oqS7X0ddcx1FA1sUNEvjKY,33174
 torchrl/data/llm/common.py,sha256=CYBaAop8QETotOCBGTw_pfKjxFYlsoSGElki6wBx5jo,2135
 torchrl/data/llm/dataset.py,sha256=t-41hAzQcjrdoKwpHIMbcrT7pRcQ7DHl2a1-lr6E7W4,20703
 torchrl/data/llm/prompt.py,sha256=bg5LzJfwOq5Ns72KQMciIprMWAmDDinzdopwdopU04c,8380
 torchrl/data/llm/reward.py,sha256=FbPchNXG3smJV9NCbB5Yk4grsCa2Se4KZ_tojVLKWQM,8404
+torchrl/data/llm/topk.py,sha256=SZq89yeFr8rNbpVR-S5vC7AVoeb6JKYZPeSS-n4FwKE,8353
 torchrl/data/llm/utils.py,sha256=axe3wSovfWBm5YmR_uJYpfAmYtd__2i9SCKgUSezkBk,23600
 torchrl/data/map/__init__.py,sha256=1IB8lWApscQOOscsCEhQrUDy_AE1wWV51Tcl1Segsqk,555
 torchrl/data/map/hash.py,sha256=29cKgYjd5vVeR2bu2kI5BwtOq9FeZD41RA7Q3UxP9vo,7252
@@ -56,11 +57,11 @@ torchrl/data/postprocs/__init__.py,sha256=Z9JpRKMGsuFGpB3ro4R9Y_hYTBqkkzbkWZR79T
 torchrl/data/postprocs/postprocs.py,sha256=h8LO8zBosRm7iLmUOxdtPxZ84yavkv9usYtLSBq9tC4,15118
 torchrl/data/replay_buffers/__init__.py,sha256=v_oKflSohims6uw40XhLkjDX7vZM9UwXrWAeZfftogw,2360
 torchrl/data/replay_buffers/checkpointers.py,sha256=VF18DlRiy361gecbT2HL5VLTQU4Faxq7mULsownjYiQ,14790
-torchrl/data/replay_buffers/ray_buffer.py,sha256=joYh_ypj4Zk2CNUEdfgNiAybvA8vNJTmrAlwC1bhejg,9043
-torchrl/data/replay_buffers/replay_buffers.py,sha256=c0yoPBzQjfWfTLOXbPj7VqA-ZnTP_7oS0OVpwo5gaPk,88639
-torchrl/data/replay_buffers/samplers.py,sha256=HZguztdX2tvDwIPKT0STAKuncJi8l005FCuuNKnAVoE,105427
+torchrl/data/replay_buffers/ray_buffer.py,sha256=at8rYXxtlctoPCnL5oJRNoEkjEASHoXjPIt6UH16OCA,9854
+torchrl/data/replay_buffers/replay_buffers.py,sha256=lKTcEQOooT_MY4cuVuAdaYPKN9Ob9v3o46FGdnCyOS8,90459
+torchrl/data/replay_buffers/samplers.py,sha256=Kp48OPzvEWeTbPS8LNMRiGaYwUdrMgVVc3OaRIkNIR4,110296
 torchrl/data/replay_buffers/scheduler.py,sha256=SRZf_FJLUEIBz684W9RlLt3In158s9N5h4xb_MWnBgY,10152
-torchrl/data/replay_buffers/storages.py,sha256=WnNopbDT3DbjiN7QRdb-Iet-sLBmt7hSfEnKDx1-VEI,59579
+torchrl/data/replay_buffers/storages.py,sha256=9h2iyLv9jnKG7kB1925SRlcxly-IABqGjPhoMGov-6Y,60704
 torchrl/data/replay_buffers/utils.py,sha256=tU98Nc_j9bMrWBs96gFUTDXLmWEZCvHRYjSXjPMc_lY,38603
 torchrl/data/replay_buffers/writers.py,sha256=p9b8k89u-JrqoObT4aCLa0qCkKWdM__l7lGUQDKSdsU,27727
 torchrl/envs/__init__.py,sha256=c-_VtMuAcRdg0hBmltn6AbTU7B1X-ARBEfqOQoPFEZk,5817
@@ -96,8 +97,8 @@ torchrl/envs/libs/smacv2.py,sha256=i0TRHuZ9S9v0NfufPgQAcTlvAjf6JKv8hHvOzjSgsaw,2
 torchrl/envs/libs/unity_mlagents.py,sha256=Z3qSU0H3o2NXbS2lNvQ7OmYxkr3AWAMyRHfxeCtNZrk,49667
 torchrl/envs/libs/utils.py,sha256=RgiR16KJWFEtQim44-AIcHByGTq_NrtpjWoYIC13aYA,5207
 torchrl/envs/libs/vmas.py,sha256=a71_jU4r627hFXcMsT5wNSb4TMpyd3punLdOF3Cc8O0,36297
-torchrl/envs/llm/__init__.py,sha256=DiYt8YjoxmwoM62XPtNUPMYaqZyf1UXY6dAD_vcBIfE,1221
-torchrl/envs/llm/chat.py,sha256=DT_kcsfpM0W3bayRVk3rdtNKyv3pjoOsicw56LG6fp8,17619
+torchrl/envs/llm/__init__.py,sha256=o8uAVGHYngy_k6xM5qIkqgHaz__S1HyG7QjLd78gtaA,1265
+torchrl/envs/llm/chat.py,sha256=mVLjmBTwd6IWdlKJMRcynDJNVVbiHjCop5EVUXpaaAA,17794
 torchrl/envs/llm/envs.py,sha256=Er-ahjgvtYG4LB7_EWOMbdobiUV5DOHPBQYkVTu80r4,34677
 torchrl/envs/llm/datasets/__init__.py,sha256=FFethtv8unJWzphGLPQVC5QD9NMdaygEjx25O1DHHZk,473
 torchrl/envs/llm/datasets/gsm8k.py,sha256=wTntpV-bi0gbyvJ-JnuHQmPXjXgV4hEssGFed8GRGGc,15299
@@ -105,18 +106,18 @@ torchrl/envs/llm/datasets/ifeval.py,sha256=fVbMSVjpnlZR36B0yDUgDcM1Ye-EP6ui7g9nP
 torchrl/envs/llm/libs/__init__.py,sha256=vhEm5Fhz1sLWt107zfZLy5pzGmfQi0fNBGazTq1m7dU,266
 torchrl/envs/llm/libs/mlgym.py,sha256=ECnkrNoPV73L1fIO05SlTTXuTSNOM2pdX6aJcEYJVlo,31372
 torchrl/envs/llm/reward/__init__.py,sha256=a-Xsye29z2LugO1cOCFM2FNsqNwEp-5XwQk4saVQlu8,370
-torchrl/envs/llm/reward/gsm8k.py,sha256=6y6I8UdPanS6g7skWFStNm_nXP0nS5ctcAHFWEkFup0,7702
+torchrl/envs/llm/reward/gsm8k.py,sha256=2pUXYkCw6_arM6HCZJcrEYwRZMDntsFAzdpf3QXNthI,7862
 torchrl/envs/llm/reward/ifeval/__init__.py,sha256=g5NtrwfwqK22hRcoIdz8-KWBh5Ogre9J-Bf3uGWE9Pg,314
 torchrl/envs/llm/reward/ifeval/_instructions.py,sha256=rAoTdwG42smCLJgwW7kAwJrNonjIS6OwdohDE70oMOA,61696
 torchrl/envs/llm/reward/ifeval/_instructions_main.py,sha256=CofKXvG0J2H-1ZXP1fL6UZI8ArNCIO2w5R_37drRIW8,4117
 torchrl/envs/llm/reward/ifeval/_instructions_registry.py,sha256=3_guc8LZ0mWQc-n6E4cQgYMgZRYa6xfgvXgrze9aO_w,3814
 torchrl/envs/llm/reward/ifeval/_instructions_util.py,sha256=aA3fupO8MvqBCqD7Y_Qk6y32toWF1lZGAflWON1ruXM,26042
 torchrl/envs/llm/reward/ifeval/_scorer.py,sha256=zJHBgaGlluEv6czsI6ZtLqArV_J_W9zY7UPAJhT5YIo,14563
-torchrl/envs/llm/transforms/__init__.py,sha256=fpcS83ud3OC2NWnkFeTdE8r4Mtlbcp_OiITzDM03aes,748
+torchrl/envs/llm/transforms/__init__.py,sha256=roEOZVFOs1PhC1cGF-LIXQt5DlXZx6mgIJ-1k0JDTfI,788
 torchrl/envs/llm/transforms/browser.py,sha256=zF7jHHHrdpxUCjFFtiYK-vhw-p1YqsqwP8_b4SiK0Rs,10423
 torchrl/envs/llm/transforms/dataloading.py,sha256=dv4IV3OWEa6-evxBk3WAZjkBi1_yKUs2NQ2gGmL2lKQ,24533
 torchrl/envs/llm/transforms/format.py,sha256=ESn0S9k5G4FQPBICq9h6ZsLKXZqiU71tYW8UnW4rgLI,2519
-torchrl/envs/llm/transforms/kl.py,sha256=lRWW1Gf8bu71jMatAlk91Eeuh50mmPedjKrnXKUm5D0,11721
+torchrl/envs/llm/transforms/kl.py,sha256=N68378chSx54X5a7YLJzIV6d870H5xrBb5-qWqzpX1U,22744
 torchrl/envs/llm/transforms/policy_version.py,sha256=by2TjsZLwVjQbq7ggBoAco2Iq_2aEYgyxh9asTXL1vk,6893
 torchrl/envs/llm/transforms/tokenizer.py,sha256=CcuKRu33YnyDgLtQtyxTGDFC6iI3b3fUA6Nb1Lnh7h8,13953
 torchrl/envs/llm/transforms/tools.py,sha256=I-HR0zjH4tFMp9xPH556H5Q5JqmqXdsAXwElAR93e5U,29498
@@ -147,8 +148,8 @@ torchrl/modules/llm/backends/__init__.py,sha256=WdVy9EdiAfk8i5zFa49TEkRvcUd0L4Un
 torchrl/modules/llm/backends/vllm.py,sha256=x57Xop1xd5ZShicsh47ZFmz4VpfZ3eCzVx7k0COvpqQ,9387
 torchrl/modules/llm/policies/__init__.py,sha256=rVQwVhSTS1hLcSynvPXKq9_9gGC6gC1SyOz5DNg1qcc,426
 torchrl/modules/llm/policies/common.py,sha256=m76rSjgYbf-ZMEUFZNbjBbyXNHbR8BXt1z5o9honJOM,3019
-torchrl/modules/llm/policies/transformers_wrapper.py,sha256=yn_qVpFqjr41HrkkxGhLDfIjtS9PCgklnbkAecu4Evc,22615
-torchrl/modules/llm/policies/vllm_wrapper.py,sha256=g3eaQSNti6NQBpKcokeLL9b0K3Kt38ltaPv8qlIIqDo,29782
+torchrl/modules/llm/policies/transformers_wrapper.py,sha256=M0Drk7MFY596Ek8_duNTXFpc4c2Ar94Jy3viXnhRS2M,25370
+torchrl/modules/llm/policies/vllm_wrapper.py,sha256=1vwfoIYxOL2IwBMVZUFrwOexIwS7x1xbhBVdru6gYxY,31487
 torchrl/modules/models/__init__.py,sha256=DrOG-7hynjjUh_tc2EqysiUiNMRiDR0WLtZql9TPNcI,1743
 torchrl/modules/models/batchrenorm.py,sha256=TojpTUluIcFdTSemIVRLGtB2O5q54mRHy3vJP6DuI5I,4750
 torchrl/modules/models/decision_transformer.py,sha256=Lttf_wZMNqXbB_vpxMYgEp18gEzOvm3NvMnxQkHkH4M,6604
@@ -194,8 +195,9 @@ torchrl/objectives/sac.py,sha256=Oq9Iq90s9KFbnM4KSRUd2onU1JfW6aW80LWGdtO0CY8,639
 torchrl/objectives/td3.py,sha256=RnlkGzBBTY0KrfRKytsFbNyoVUy2HLfwSL4_9YQRep8,23190
 torchrl/objectives/td3_bc.py,sha256=jHGwCzPuCbN37zAxsiDQIe92yR1UE7rjcnJoy8b_NjE,25950
 torchrl/objectives/utils.py,sha256=nhB7a2gLoZMLgYSWTpSgQqZWEGRBkvVoa8yszTlecm4,24001
-torchrl/objectives/llm/__init__.py,sha256=LnYwAuaG-ylQQcu2BRQWavaDhjMPikXNT6YaH_3QoEU,328
+torchrl/objectives/llm/__init__.py,sha256=SXYwry5YoDp5m0QRFmOYzz60siJQmofcTvCOmC1DlXw,396
 torchrl/objectives/llm/grpo.py,sha256=rsPVvfE_2Bbl8K1aq_LIry1ViDnibfGYWexfSIbJx80,16788
+torchrl/objectives/llm/sft.py,sha256=zAdVT1CmXJJPjEwPt4SPJNzFUC2m-flcfOsejIuAFkg,20107
 torchrl/objectives/multiagent/__init__.py,sha256=CHxWmq5_3kveLcAdyB7cgSVYVIald7EZo81RRgozxo0,237
 torchrl/objectives/multiagent/qmixer.py,sha256=JyDcZeV2zv2MqKsyJ-ql9ISYHJ58e3pzb5-0BThswhI,16973
 torchrl/objectives/value/__init__.py,sha256=AdluF370wYzOAcP_yglUAFnNByKVZzivBYJafkDQbJA,561
@@ -221,8 +223,8 @@ torchrl/trainers/helpers/losses.py,sha256=qH-2YJwMtDAYAPXTTYy3cOPiq4ILC6xTjfnGUU
 torchrl/trainers/helpers/models.py,sha256=ihTERG2c96E8cS3Tnul6a_ys6iDEEJmHh05p9blQTW8,21807
 torchrl/trainers/helpers/replay_buffer.py,sha256=ZUZHOa0TILyeWJ3iahzTJ6UvMl_0FdxuZfJEja94Bn8,2001
 torchrl/trainers/helpers/trainers.py,sha256=j6B5XA7_FFHMQeOIQwjNcO0CGE_4mZKUC9_jH_iqqh4,12071
-torchrl_nightly-2025.6.20.dist-info/LICENSE,sha256=xdjS4_xk-IwnLuIFCvTYTl9Y8aXRejqpmke3dGam_nI,1098
-torchrl_nightly-2025.6.20.dist-info/METADATA,sha256=fR7h0Sq2FvApealDc_cnX5Aj5QtIL4acVWv-Cz60FTk,39023
-torchrl_nightly-2025.6.20.dist-info/WHEEL,sha256=HRqO1yy0EkQFVSOPjhgaTzf773tbWecKJXRlZH64XT8,104
-torchrl_nightly-2025.6.20.dist-info/top_level.txt,sha256=JeTJ1jV7QJwLcUS1nr21aPn_wb-XlAZ9c-z_EH472JA,20
-torchrl_nightly-2025.6.20.dist-info/RECORD,,
+torchrl_nightly-2025.6.22.dist-info/LICENSE,sha256=xdjS4_xk-IwnLuIFCvTYTl9Y8aXRejqpmke3dGam_nI,1098
+torchrl_nightly-2025.6.22.dist-info/METADATA,sha256=p6Lp-DEGEipD6Ak0XOcCM_dRoqAan8l4PSSAtgyr7K4,39023
+torchrl_nightly-2025.6.22.dist-info/WHEEL,sha256=HRqO1yy0EkQFVSOPjhgaTzf773tbWecKJXRlZH64XT8,104
+torchrl_nightly-2025.6.22.dist-info/top_level.txt,sha256=JeTJ1jV7QJwLcUS1nr21aPn_wb-XlAZ9c-z_EH472JA,20
+torchrl_nightly-2025.6.22.dist-info/RECORD,,