Error message on packed=True for stack exchange dataset

tests/torchtune/datasets/test_hh_rlhf_helpful_dataset.py

@@ -107,3 +107,14 @@ def test_dataset_get_item(self, mock_load_dataset, train_on_input):
         else:
             # Check that the input is masked
             assert sample["rejected_labels"].count(CROSS_ENTROPY_IGNORE_IDX) == 16
+
+    def test_dataset_fails_with_packed(self):
+        with pytest.raises(
+            ValueError,
+            match="Packed is currently not supported for preference datasets",
+        ):
+            hh_rlhf_helpful_dataset(
+                tokenizer=DummyTokenizer(),
+                train_on_input=True,
+                packed=True,
+            )

tests/torchtune/datasets/test_preference_dataset.py

@@ -155,3 +155,17 @@ def test_load_local_json(self):
 
         assert expected_chosen_labels[0] == ds[0]["chosen_labels"]
         assert expected_rejected_labels[0] == ds[0]["rejected_labels"]
+
+    def test_dataset_fails_with_packed(self):
+        with pytest.raises(
+            ValueError,
+            match="Packed is currently not supported for preference datasets.",
+        ):
+            preference_dataset(
+                tokenizer=DummyTokenizer(),
+                source="json",
+                data_files=str(ASSETS / "hh_rlhf_tiny.json"),
+                train_on_input=False,
+                split="train",
+                packed=True,
+            )

tests/torchtune/datasets/test_stack_exchange_paired_dataset.py

@@ -100,6 +100,16 @@ def test_dataset_get_item(self, mock_load_dataset, train_on_input):
             # Check that the input is masked
             assert sample["rejected_labels"].count(CROSS_ENTROPY_IGNORE_IDX) == 52
 
+    def test_dataset_fails_with_packed(self):
+        with pytest.raises(
+            ValueError,
+            match="Packed is currently not supported for preference datasets",
+        ):
+            stack_exchange_paired_dataset(
+                tokenizer=DummyTokenizer(),
+                packed=True,
+            )
+
 
 class TestStackExchangePairedToMessages:
     @pytest.fixture

torchtune/datasets/_preference.py

@@ -89,9 +89,14 @@ class requires the dataset to have "chosen" and "rejected" model responses. Thes
         filter_fn (Optional[Callable]): callable used to filter the dataset prior to any pre-processing. See
             the Hugging Face `docs <https://huggingface.co/docs/datasets/v2.20.0/process#select-and-filter>`_ for more
             details.
+        packed (bool): Whether or not to pack the dataset to ``max_seq_len`` prior to training. Default is False. Packed is
+            currently not supported for ``PreferenceDataset`` and a ``ValueError`` will be raised if this is set to True.
         **load_dataset_kwargs (Dict[str, Any]): additional keyword arguments to pass to ``load_dataset``. See Hugging
             Face's `API ref <https://huggingface.co/docs/datasets/en/package_reference/loading_methods#datasets.load_dataset>`_
             for more details.
+
+    Raises:
+        ValueError: If ``packed`` is True, this feature is not supported for ``PreferenceDataset``.
     """
 
     def __init__(
@@ -101,8 +106,14 @@ def __init__(
         message_transform: Transform,
         tokenizer: ModelTokenizer,
         filter_fn: Optional[Callable] = None,
+        packed: bool = False,
         **load_dataset_kwargs: Dict[str, Any],
     ) -> None:
+        if packed:
+            raise ValueError(
+                "Packed is currently not supported for preference datasets."
+            )
+
         self._tokenizer = tokenizer
         self._message_transform = message_transform
         self._data = load_dataset(source, **load_dataset_kwargs)