Update

[ghstack-poisoned]

Author: vmoens

docs/source/reference/data.rst

@@ -1133,6 +1133,9 @@ efficient sampling.
     get_dataloader
     ConstantKLController
     AdaptiveKLController
+    LLMData
+    LLMInput
+    LLMOutput
 
 
 Utils

examples/rlhf/data/__init__.py

@@ -1,3 +1,3 @@
-from torchrl.data.rlhf.prompt import get_prompt_dataloader_tldr
+from torchrl.data.llm.prompt import get_prompt_dataloader_tldr
 
 __all__ = ["get_prompt_dataloader_tldr"]

examples/rlhf/models/reward.py

@@ -8,7 +8,7 @@
 from tensordict.nn import TensorDictModule
 from torchrl._utils import logger as torchrl_logger
 
-from torchrl.modules.models.rlhf import GPT2RewardModel
+from torchrl.modules.models.llm import GPT2RewardModel
 
 
 def init_reward_model(

examples/rlhf/train.py

@@ -17,8 +17,8 @@
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torchrl._utils import logger as torchrl_logger
 
-from torchrl.data.rlhf.dataset import get_dataloader
-from torchrl.data.rlhf.prompt import PromptData
+from torchrl.data.llm.dataset import get_dataloader
+from torchrl.data.llm.prompt import PromptData
 from utils import get_file_logger, resolve_name_or_path, setup
 
 

examples/rlhf/train_reward.py

@@ -9,8 +9,8 @@
 from models.reward import init_reward_model
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torchrl._utils import logger as torchrl_logger
-from torchrl.data.rlhf.dataset import get_dataloader
-from torchrl.data.rlhf.reward import PairwiseDataset
+from torchrl.data.llm.dataset import get_dataloader
+from torchrl.data.llm.reward import PairwiseDataset
 from utils import get_file_logger, resolve_name_or_path, setup
 
 

examples/rlhf/train_rlhf.py

@@ -6,7 +6,7 @@
 import hydra
 import torch
 from models.actor_critic import init_actor_critic
-from torchrl.data.rlhf.utils import AdaptiveKLController, RolloutFromModel
+from torchrl.data.llm.utils import AdaptiveKLController, RolloutFromModel
 
 from torchrl.record.loggers import get_logger
 

examples/rlhf/utils.py

@@ -22,9 +22,9 @@
     TensorDictReplayBuffer,
     TensorStorage,
 )
+from torchrl.data.llm.dataset import get_dataloader
+from torchrl.data.llm.prompt import PromptData
 from torchrl.data.replay_buffers import SamplerWithoutReplacement
-from torchrl.data.rlhf.dataset import get_dataloader
-from torchrl.data.rlhf.prompt import PromptData
 from torchrl.objectives import ClipPPOLoss
 from torchrl.objectives.value import GAE
 

test/assets/generate.py

@@ -5,6 +5,7 @@
 
 """Script used to generate the mini datasets."""
 import multiprocessing as mp
+import pathlib
 
 try:
     mp.set_start_method("spawn")
@@ -14,8 +15,8 @@
 
 from datasets import Dataset, DatasetDict, load_dataset
 
-from torchrl.data.rlhf.dataset import get_dataloader
-from torchrl.data.rlhf.prompt import PromptData
+from torchrl.data.llm.dataset import get_dataloader
+from torchrl.data.llm.prompt import PromptData
 
 
 def generate_small_dataset(comparison=True):
@@ -42,7 +43,7 @@ def get_minibatch():
             batch_size=16,
             block_size=33,
             tensorclass_type=PromptData,
-            dataset_name="../datasets_mini/openai_summarize_tldr",
+            dataset_name=f"{pathlib.Path(__file__).parent}/../datasets_mini/openai_summarize_tldr",
             device="cpu",
             num_workers=2,
             infinite=False,
@@ -52,9 +53,12 @@ def get_minibatch():
             root_dir=tmpdir,
         )
         for data in dl:
-            data = data.clone().memmap_("test/datasets_mini/tldr_batch/")
+            data = data.clone().memmap_(
+                f"{pathlib.Path(__file__).parent}/../datasets_mini/tldr_batch/"
+            )
             break
 
 
 if __name__ == "__main__":
+    generate_small_dataset(False)
     get_minibatch()

test/assets/tldr_batch.zip

@@ -14,7 +14,7 @@
 
 from torch import distributions as dist, nn
 from torchrl.data import Binary, Bounded, Categorical, Composite, MultiOneHot, OneHot
-from torchrl.data.rlhf.dataset import _has_transformers
+from torchrl.data.llm.dataset import _has_transformers
 from torchrl.modules import MLP, SafeModule, TanhDelta, TanhNormal
 from torchrl.modules.tensordict_module.actors import (
     _process_action_space_spec,

test/test_actors.py

@@ -61,7 +61,7 @@
 from torchrl.envs.libs.dm_control import _has_dmc, DMControlEnv
 from torchrl.envs.libs.gym import _has_gym, gym_backend, GymEnv, GymWrapper
 from torchrl.envs.transforms import Compose, StepCounter, TransformedEnv
-from torchrl.envs.transforms.rlhf import as_padded_tensor
+from torchrl.envs.transforms.llm import as_padded_tensor
 from torchrl.envs.transforms.transforms import (
     AutoResetEnv,
     AutoResetTransform,

test/test_env.py

@@ -21,17 +21,17 @@
     TensorDictBase,
 )
 from tensordict.nn import TensorDictModule
-from torchrl.data.rlhf import TensorDictTokenizer
-from torchrl.data.rlhf.dataset import (
+from torchrl.data.llm import TensorDictTokenizer
+from torchrl.data.llm.dataset import (
     _has_datasets,
     _has_transformers,
     get_dataloader,
     TokenizedDatasetLoader,
 )
-from torchrl.data.rlhf.prompt import PromptData, PromptTensorDictTokenizer
-from torchrl.data.rlhf.reward import PairwiseDataset, pre_tokenization_hook
-from torchrl.data.rlhf.utils import RolloutFromModel
-from torchrl.modules.models.rlhf import GPT2RewardModel
+from torchrl.data.llm.prompt import PromptData, PromptTensorDictTokenizer
+from torchrl.data.llm.reward import PairwiseDataset, pre_tokenization_hook
+from torchrl.data.llm.utils import RolloutFromModel
+from torchrl.modules.models.llm import GPT2RewardModel
 
 if os.getenv("PYTORCH_TEST_FBCODE"):
     from pytorch.rl.test._utils_internal import get_default_devices

test/test_rlhf.py

@@ -117,8 +117,8 @@
 from torchrl.envs.libs.gym import _has_gym, GymEnv, set_gym_backend
 from torchrl.envs.libs.unity_mlagents import _has_unity_mlagents
 from torchrl.envs.transforms import VecNorm
+from torchrl.envs.transforms.llm import KLRewardTransform
 from torchrl.envs.transforms.r3m import _R3MNet
-from torchrl.envs.transforms.rlhf import KLRewardTransform
 from torchrl.envs.transforms.transforms import (
     _has_tv,
     ActionDiscretizer,

test/test_transforms.py

@@ -3,6 +3,22 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+from .llm import (
+    AdaptiveKLController,
+    ConstantKLController,
+    create_infinite_iterator,
+    get_dataloader,
+    LLMData,
+    LLMInput,
+    LLMOutput,
+    PairwiseDataset,
+    PromptData,
+    PromptTensorDictTokenizer,
+    RewardData,
+    RolloutFromModel,
+    TensorDictTokenizer,
+    TokenizedDatasetLoader,
+)
 from .map import (
     BinaryToDecimal,
     HashToInt,
@@ -56,19 +72,6 @@
     Writer,
     WriterEnsemble,
 )
-from .rlhf import (
-    AdaptiveKLController,
-    ConstantKLController,
-    create_infinite_iterator,
-    get_dataloader,
-    PairwiseDataset,
-    PromptData,
-    PromptTensorDictTokenizer,
-    RewardData,
-    RolloutFromModel,
-    TensorDictTokenizer,
-    TokenizedDatasetLoader,
-)
 from .tensor_specs import (
     Binary,
     BinaryDiscreteTensorSpec,
@@ -125,6 +128,9 @@
     "H5StorageCheckpointer",
     "HashToInt",
     "ImmutableDatasetWriter",
+    "LLMData",
+    "LLMInput",
+    "LLMOutput",
     "LazyMemmapStorage",
     "LazyStackStorage",
     "LazyStackedCompositeSpec",

torchrl/data/__init__.py

@@ -11,18 +11,21 @@
 )
 from .prompt import PromptData, PromptTensorDictTokenizer
 from .reward import PairwiseDataset, RewardData
-from .utils import AdaptiveKLController, ConstantKLController, RolloutFromModel
+from .utils import AdaptiveKLController, ConstantKLController, RolloutFromModel, LLMData, LLMOutput, LLMInput
 
 __all__ = [
-    "create_infinite_iterator",
-    "get_dataloader",
-    "TensorDictTokenizer",
-    "TokenizedDatasetLoader",
+    "AdaptiveKLController",
+    "ConstantKLController",
+    "LLMData",
+    "LLMInput",
+    "LLMOutput",
+    "PairwiseDataset",
     "PromptData",
     "PromptTensorDictTokenizer",
-    "PairwiseDataset",
     "RewardData",
-    "AdaptiveKLController",
-    "ConstantKLController",
     "RolloutFromModel",
+    "TensorDictTokenizer",
+    "TokenizedDatasetLoader",
+    "create_infinite_iterator",
+    "get_dataloader",
 ]

torchrl/data/rlhf/__init__.py renamed to torchrl/data/llm/__init__.py

@@ -31,7 +31,7 @@ class TokenizedDatasetLoader:
         max_length (int): the maximum sequence length.
         dataset_name (str): the name of the dataset.
         tokenizer_fn (callable): the tokeinizing method constructor, such as
-            :class:`torchrl.data.rlhf.TensorDictTokenizer`. When called,
+            :class:`torchrl.data.llm.TensorDictTokenizer`. When called,
             it should return a :class:`tensordict.TensorDict` instance
             or a dictionary-like structure with the tokenized data.
         pre_tokenization_hook (callable, optional): called on
@@ -62,8 +62,8 @@ class TokenizedDatasetLoader:
     The dataset will be stored in ``<root_dir>/<split>/<max_length>/``.
 
     Examples:
-        >>> from torchrl.data.rlhf import TensorDictTokenizer
-        >>> from torchrl.data.rlhf.reward import  pre_tokenization_hook
+        >>> from torchrl.data.llm import TensorDictTokenizer
+        >>> from torchrl.data.llm.reward import  pre_tokenization_hook
         >>> split = "train"
         >>> max_length = 550
         >>> dataset_name = "CarperAI/openai_summarize_comparisons"
@@ -359,7 +359,7 @@ def get_dataloader(
             Defaults to ``max(os.cpu_count() // 2, 1)``.
 
     Examples:
-        >>> from torchrl.data.rlhf.reward import PairwiseDataset
+        >>> from torchrl.data.llm.reward import PairwiseDataset
         >>> dataloader = get_dataloader(
         ...     batch_size=256, block_size=550, tensorclass_type=PairwiseDataset, device="cpu")
         >>> for d in dataloader:

torchrl/data/rlhf/dataset.py renamed to torchrl/data/llm/dataset.py

@@ -7,7 +7,7 @@
 import torch
 from tensordict import tensorclass, TensorDict
 
-from torchrl.data.rlhf.dataset import TensorDictTokenizer, TokenizedDatasetLoader
+from torchrl.data.llm.dataset import TensorDictTokenizer, TokenizedDatasetLoader
 
 DEFAULT_DATASET = "CarperAI/openai_summarize_tldr"
 

torchrl/data/rlhf/prompt.py renamed to torchrl/data/llm/prompt.py

@@ -8,7 +8,7 @@
 
 import torch
 from tensordict import tensorclass
-from torchrl.data.rlhf.dataset import TensorDictTokenizer, TokenizedDatasetLoader
+from torchrl.data.llm.dataset import TensorDictTokenizer, TokenizedDatasetLoader
 
 DEFAULT_DATASET = "CarperAI/openai_summarize_comparisons"
 _has_datasets = importlib.util.find_spec("datasets") is not None