Sample packing for ConcatDataset (#2278)

Author: ebsmothers

recipes/dev/early_exit_finetune_distributed.py

@@ -653,7 +653,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

recipes/full_finetune_distributed.py

@@ -646,7 +646,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

recipes/full_finetune_single_device.py

@@ -558,7 +558,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

recipes/knowledge_distillation_distributed.py

@@ -652,7 +652,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

recipes/knowledge_distillation_single_device.py

@@ -531,7 +531,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

recipes/lora_finetune_distributed.py

@@ -591,7 +591,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

recipes/lora_finetune_single_device.py

@@ -528,7 +528,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

recipes/qat_distributed.py

@@ -606,7 +606,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

recipes/qat_lora_finetune_distributed.py

@@ -633,7 +633,7 @@ def _setup_data(
                 for single_cfg_dataset in cfg_dataset
             ]
             ds = ConcatDataset(datasets=datasets)
-            packed = False
+            packed = getattr(ds, "packed", False)
         else:
             ds = config.instantiate(cfg_dataset, self._tokenizer)
             packed = cfg_dataset.get("packed", False)

tests/torchtune/datasets/test_concat_dataset.py

@@ -80,7 +80,7 @@ def test_invalid_index_type(self, datasets):
         with pytest.raises(TypeError):
             multi_dataset["invalid_type"]  # Non-integer index
 
-    def test_packed_dataset(self, torch_datasets):
+    def test_single_packed_dataset(self, torch_datasets):
         torch_datasets[0] = PackedDataset(
             torch_datasets[0],
             max_seq_len=25,
@@ -90,3 +90,33 @@ def test_packed_dataset(self, torch_datasets):
 
         with pytest.raises(ValueError):
             concated_dataset = ConcatDataset(torch_datasets)
+
+    def test_all_packed_datasets(self, torch_datasets):
+        for i in range(len(torch_datasets)):
+            torch_datasets[i] = PackedDataset(
+                torch_datasets[i],
+                max_seq_len=2000,
+                max_packs=16,
+                split_across_pack=True,
+            )
+        concated_dataset = ConcatDataset(torch_datasets)
+        assert concated_dataset.packed
+
+        # 2k tokens per pack
+        # 1st ds has 4k tokens, 2nd ds has 8k tokens, 3rd ds has 15k tokens
+        # 4th ds has 16k tokens, 5th ds has 23k tokens, 6th ds has 42k tokens
+
+        assert concated_dataset[0]["seq_lens"][0] == 4
+        # 2nd packed ds starts at idx 2
+        assert concated_dataset[2]["seq_lens"][0] == 8
+        # 3rd packed ds starts at idx 6
+        assert concated_dataset[6]["seq_lens"][0] == 15
+        # 4th packed ds starts at idx 14
+        assert concated_dataset[14]["seq_lens"][0] == 16
+        # 5th packed ds starts at idx 22
+        assert concated_dataset[22]["seq_lens"][0] == 23
+        # 6th packed ds starts at idx 34
+        assert concated_dataset[34]["seq_lens"][0] == 42
+
+        # Total length is 2 + 4 + 8 + 8 + 12 + 16 (because of max_packs) = 50
+        assert len(concated_dataset) == 50

torchtune/datasets/_concat.py

@@ -67,12 +67,12 @@ class ConcatDataset(Dataset):
     def __init__(self, datasets: List[Dataset]):
         self._datasets: List[Dataset] = datasets
 
-        for dataset in self._datasets:
-            if isinstance(dataset, PackedDataset):
-                raise ValueError(
-                    "ConcatDataset can't process instances of PackedDataset."
-                )
-
+        is_packed = [isinstance(dataset, PackedDataset) for dataset in datasets]
+        if any(is_packed) and not all(is_packed):
+            raise ValueError(
+                "ConcatDataset can't process a mix of packed and non-packed datasets."
+            )
+        self.packed = all(is_packed)
         self._len: int = sum(len(dataset) for dataset in datasets)
         self._indexes: List[Tuple[int, int, int]] = []