Sample packing for ConcatDataset (#2278)

ebsmothers · web-flow · commit 103609545fec · 2025-01-17T16:26:59.000-08:00
diff --git a/recipes/dev/early_exit_finetune_distributed.py b/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)
diff --git a/recipes/full_finetune_distributed.py b/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)
diff --git a/recipes/full_finetune_single_device.py b/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)
diff --git a/recipes/knowledge_distillation_distributed.py b/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)
diff --git a/recipes/knowledge_distillation_single_device.py b/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)
diff --git a/recipes/lora_finetune_distributed.py b/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)
diff --git a/recipes/lora_finetune_single_device.py b/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)
diff --git a/recipes/qat_distributed.py b/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)
diff --git a/recipes/qat_lora_finetune_distributed.py b/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)
diff --git a/tests/torchtune/datasets/test_concat_dataset.py b/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
diff --git a/torchtune/datasets/_concat.py b/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]] = []
 

Original file line number	Diff line number	Diff line change
`@@ -653,7 +653,7 @@ def _setup_data(`
`653`	`653`	`for single_cfg_dataset in cfg_dataset`
`654`	`654`	`]`
`655`	`655`	`ds = ConcatDataset(datasets=datasets)`
`656`		`- packed = False`
	`656`	`+ packed = getattr(ds, "packed", False)`
`657`	`657`	`else:`
`658`	`658`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`659`	`659`	`packed = cfg_dataset.get("packed", False)`
Original file line number	Diff line number	Diff line change
`@@ -646,7 +646,7 @@ def _setup_data(`
`646`	`646`	`for single_cfg_dataset in cfg_dataset`
`647`	`647`	`]`
`648`	`648`	`ds = ConcatDataset(datasets=datasets)`
`649`		`- packed = False`
	`649`	`+ packed = getattr(ds, "packed", False)`
`650`	`650`	`else:`
`651`	`651`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`652`	`652`	`packed = cfg_dataset.get("packed", False)`
Original file line number	Diff line number	Diff line change
`@@ -558,7 +558,7 @@ def _setup_data(`
`558`	`558`	`for single_cfg_dataset in cfg_dataset`
`559`	`559`	`]`
`560`	`560`	`ds = ConcatDataset(datasets=datasets)`
`561`		`- packed = False`
	`561`	`+ packed = getattr(ds, "packed", False)`
`562`	`562`	`else:`
`563`	`563`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`564`	`564`	`packed = cfg_dataset.get("packed", False)`
Original file line number	Diff line number	Diff line change
`@@ -652,7 +652,7 @@ def _setup_data(`
`652`	`652`	`for single_cfg_dataset in cfg_dataset`
`653`	`653`	`]`
`654`	`654`	`ds = ConcatDataset(datasets=datasets)`
`655`		`- packed = False`
	`655`	`+ packed = getattr(ds, "packed", False)`
`656`	`656`	`else:`
`657`	`657`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`658`	`658`	`packed = cfg_dataset.get("packed", False)`
Original file line number	Diff line number	Diff line change
`@@ -531,7 +531,7 @@ def _setup_data(`
`531`	`531`	`for single_cfg_dataset in cfg_dataset`
`532`	`532`	`]`
`533`	`533`	`ds = ConcatDataset(datasets=datasets)`
`534`		`- packed = False`
	`534`	`+ packed = getattr(ds, "packed", False)`
`535`	`535`	`else:`
`536`	`536`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`537`	`537`	`packed = cfg_dataset.get("packed", False)`
Original file line number	Diff line number	Diff line change
`@@ -591,7 +591,7 @@ def _setup_data(`
`591`	`591`	`for single_cfg_dataset in cfg_dataset`
`592`	`592`	`]`
`593`	`593`	`ds = ConcatDataset(datasets=datasets)`
`594`		`- packed = False`
	`594`	`+ packed = getattr(ds, "packed", False)`
`595`	`595`	`else:`
`596`	`596`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`597`	`597`	`packed = cfg_dataset.get("packed", False)`
Original file line number	Diff line number	Diff line change
`@@ -528,7 +528,7 @@ def _setup_data(`
`528`	`528`	`for single_cfg_dataset in cfg_dataset`
`529`	`529`	`]`
`530`	`530`	`ds = ConcatDataset(datasets=datasets)`
`531`		`- packed = False`
	`531`	`+ packed = getattr(ds, "packed", False)`
`532`	`532`	`else:`
`533`	`533`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`534`	`534`	`packed = cfg_dataset.get("packed", False)`
Original file line number	Diff line number	Diff line change
`@@ -606,7 +606,7 @@ def _setup_data(`
`606`	`606`	`for single_cfg_dataset in cfg_dataset`
`607`	`607`	`]`
`608`	`608`	`ds = ConcatDataset(datasets=datasets)`
`609`		`- packed = False`
	`609`	`+ packed = getattr(ds, "packed", False)`
`610`	`610`	`else:`
`611`	`611`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`612`	`612`	`packed = cfg_dataset.get("packed", False)`
Original file line number	Diff line number	Diff line change
`@@ -633,7 +633,7 @@ def _setup_data(`
`633`	`633`	`for single_cfg_dataset in cfg_dataset`
`634`	`634`	`]`
`635`	`635`	`ds = ConcatDataset(datasets=datasets)`
`636`		`- packed = False`
	`636`	`+ packed = getattr(ds, "packed", False)`
`637`	`637`	`else:`
`638`	`638`	`ds = config.instantiate(cfg_dataset, self._tokenizer)`
`639`	`639`	`packed = cfg_dataset.get("packed", False)`