Dataset wrapper classes (#37)

Author: jlamypoirier

fast_llm/data/config.py

@@ -151,14 +151,6 @@ class Dataset(abc.ABC):
     A generic dataset class compatible with torch.utils.data.Dataset but with a slightly different signature.
     """
 
-    @abc.abstractmethod
-    def __getitem__(self, index: int):
-        pass
-
-    @abc.abstractmethod
-    def __len__(self):
-        pass
-
     @property
     @abc.abstractmethod
     def name(self):
@@ -167,17 +159,16 @@ def name(self):
         """
 
 
-class RawDataset(Dataset):  # noqa
-    """
-    A raw dataset class containing a list of unsampled, unprocessed samples, i.e., matching what is stored on disk.
-    (Excluding off-line processing prior to training.)
-    Functionally identical to a `Dataset`, but renamed for clarity.
-    """
-
-
 class SampledDataset(Dataset):  # noqa
     """
     A sampled dataset class containing a prepared list of samples to be indexed sequentially (as-is) during training.
     (See the `Sampler` class below.)
-    Functionally identical to a `Dataset`, but renamed for clarity.
     """
+
+    @abc.abstractmethod
+    def __getitem__(self, index: int):
+        pass
+
+    @abc.abstractmethod
+    def __len__(self):
+        pass

fast_llm/data/gpt/concatenated.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+from fast_llm.data.gpt.dataset import GPTIndexedDataset
+from fast_llm.utils import padded_cumsum
+
+
+class GPTConcatenatedDataset(GPTIndexedDataset):
+
+    def __init__(
+        self,
+        name: str,
+        datasets: list[GPTIndexedDataset],
+    ):
+        self._name = name
+        self._datasets = datasets
+        sizes = [dataset.num_documents for dataset in self._datasets]
+        self._dataset_splits = padded_cumsum(sizes)
+        self._num_documents = sum(sizes)
+
+    @property
+    def num_tokens(self):
+        return sum(dataset.num_tokens for dataset in self._datasets)
+
+    def num_documents(self):
+        return sum(dataset.num_documents for dataset in self._datasets)
+
+    def get_document_sizes(self) -> "np.ndarray":
+        # TODO: This can be really big.
+        return np.concatenate([dataset.get_document_sizes() for dataset in self._datasets])
+
+    def get(self, document: int, offset: int = 0, length: int | None = None):
+        """
+        Get the sample (document) with the given index (in the dataset slice),
+        optionally sub-sampled to a specific offset (starting point) and maximum length
+        (end = min(offset + length, sample_length).
+        """
+        dataset = np.searchsorted(self._dataset_splits[1:], document, side="right")
+        return self._datasets[dataset].get(document - self._dataset_splits[dataset], offset, length)
+
+    @property
+    def name(self):
+        return self._name

fast_llm/data/gpt/data.py

@@ -11,10 +11,10 @@
 from fast_llm.data.blended import BlendedDataset
 from fast_llm.data.config import Data, DatasetSource, SampledDataset
 from fast_llm.data.gpt.config import DataConfig
-from fast_llm.data.gpt.dataset import GPTDataset
 from fast_llm.data.gpt.dummy import DummyGPTDataset
 from fast_llm.data.gpt.memmap import GPTMemmapDataset
-from fast_llm.data.gpt.sampled import GPTSampledDataset
+from fast_llm.data.gpt.sampled import GPTSampledIndexedDataset
+from fast_llm.data.gpt.slice import GPTDatasetSlice
 from fast_llm.data.iterator import SampledDatasetIterator
 from fast_llm.data.tokenizer import Tokenizer
 from fast_llm.engine.config_utils.run import get_run, log_main_rank
@@ -197,13 +197,15 @@ def get_iterator(
         )
 
     def _build_and_sample_gpt_dataset(self, name: str, dataset_samples_per_phase: dict[PhaseType, int]):
-        dataset_split = GPTDataset.from_splits(name, GPTMemmapDataset(self._dataset_prefixes[name]), self._phase_split)
+        dataset_split = GPTDatasetSlice.from_splits(
+            GPTMemmapDataset(name, self._dataset_prefixes[name]), self._phase_split
+        )
 
         sampled_datasets = {}
         for phase, num_samples in dataset_samples_per_phase.items():
             if num_samples == 0:
                 continue
-            sampled_datasets[phase] = GPTSampledDataset(
+            sampled_datasets[phase] = GPTSampledIndexedDataset(
                 dataset_split[phase],
                 num_samples=num_samples,
                 sequence_length=self._max_sequence_length,

fast_llm/data/gpt/dataset.py

@@ -1,12 +1,11 @@
+import abc
 import math
 
 import numpy as np
 import numpy.random
 
-from fast_llm.data.config import RawDataset
-from fast_llm.data.gpt.memmap import GPTMemmapDataset
-from fast_llm.engine.distributed.config import PhaseType
-from fast_llm.utils import Assert, padded_cumsum
+from fast_llm.data.config import Dataset
+from fast_llm.utils import Assert
 
 try:
     from fast_llm.csrc.data import build_sample_idx  # noqa
@@ -16,99 +15,70 @@
     _extension_available = False
 
 
-class GPTDataset(RawDataset):
+class GPTIndexedDataset(Dataset):
     """
-    A GPT dataset, which reads samples from (a split of) a `MMapIndexedDataset` pointing to a GPT dataset.
+    A GPT dataset containing a list of unsampled, unprocessed samples.
+    TODO: Move sampling responsibility here?
     """
 
-    def __init__(
-        self,
-        name: str,
-        indexed_dataset: GPTMemmapDataset,
-        split_begin: int | None = None,
-        split_end: int | None = None,
-    ):
-        self._name = name
-        self._indexed_dataset = indexed_dataset
-
-        self._split_begin = 0 if split_begin is None else split_begin
-        self._split_end = len(indexed_dataset) if split_end is None else split_end
-
-        # Checks
-        try:
-            Assert.geq(self._split_begin, 0)
-            Assert.in_range_incl(self._split_end, self._split_begin + 1, len(indexed_dataset))
-        except Exception as e:
-            raise AssertionError(
-                f"Invalid document indices for dataset {name} with length {len(indexed_dataset)}"
-            ) from e
-
-    def __len__(self):
-        return self._split_end - self._split_begin
-
-    def __getitem__(self, index: int):
-        """
-        Get the sample (document) with the given index (in the split dataset).
-        """
-        return self.get(index)
+    def get(self, document: int, offset: int = 0, length: int | None = None):
+        pass
 
-    def get(self, idx, offset=0, length=None):
+    @property
+    def num_documents(self) -> int:
         """
-        Get the sample (document) with the given index (in the split dataset),
-        optionally sub-sampled to a specific offset (starting point) and maximum length
-        (end = min(offset + length, sample_length).
+        Number of documents in the dataset.
+        Can be calculated from document sizes but may be overridden if there is a better method.
         """
-        return self._indexed_dataset.get(idx, offset, length)
+        return len(self.get_document_sizes())
 
     @property
-    def name(self):
-        return self._name
+    def num_tokens(self) -> int:
+        """
+        Number of tokens in the dataset.
+        Can be calculated from document sizes but may be overridden if there is a better method.
+        """
+        return self.get_document_sizes().sum()
 
-    @classmethod
-    def from_splits(cls, name: str, indexed_dataset: GPTMemmapDataset, phase_split: dict[PhaseType, float]):
+    @abc.abstractmethod
+    def get_document_sizes(self) -> "np.ndarray":
         """
-        Create a set of GPT datasets from a MMapIndexedDataset,
-        each containing approximately the requested proportion of the total tokens.
+        The size of each document in the dataset.
+        The resulting array could be very large, so this method should be called cautiously,
+        and derived classes should try to avoid holding the whole array im memory.
         """
-        split_probs = list(phase_split.values())
-        Assert.eq(sum(split_probs), 1)
-        num_documents = indexed_dataset.sizes.shape[0]
-        splits = [round(x) for x in padded_cumsum(split_probs) * num_documents]
-        return {
-            phase: GPTDataset(f"{name}_{phase.value}", indexed_dataset, split_begin, split_end)
-            for phase, split_begin, split_end in zip(phase_split, splits[:-1], splits[1:])
-        }
 
     def sample(self, num_samples: int, sequence_length: int, np_rng: numpy.random.RandomState, verbose: bool):
         """
         Create a `GPTSampledDataset` with the requested parameters.
         """
-        tokens_per_epoch = np.sum(self._indexed_dataset.sizes[self._split_begin : self._split_end])
-        num_epochs = math.ceil((sequence_length * num_samples + 1) / tokens_per_epoch)
+        document_sizes = self.get_document_sizes()
+        num_documents = len(document_sizes)
+        num_tokens = document_sizes.sum()
+
+        num_epochs = math.ceil((sequence_length * num_samples + 1) / num_tokens)
         # For the last epoch, decide whether include the entire epoch
         # in the global shuffle or not.
         # Get the number of samples for the last epoch
-        main_epochs_samples = ((num_epochs - 1) * tokens_per_epoch - 1) // sequence_length
+        main_epochs_samples = ((num_epochs - 1) * num_tokens - 1) // sequence_length
         last_epoch_samples = num_samples - main_epochs_samples
-        samples_per_epoch = (tokens_per_epoch - 1) // sequence_length
+        samples_per_epoch = (num_tokens - 1) // sequence_length
         # If we have less than 80% of the samples for the last epoch, separate out the epoch and treat it differently.
         # Note: the 80% number is just based on common sense and can be adjusted if needed.
         separate_last_epoch = num_epochs > 1 and last_epoch_samples < 0.8 * samples_per_epoch
 
-        doc_idx = np.tile(np.arange(self._split_begin, self._split_end, dtype=np.int32), num_epochs)
+        doc_idx = np.tile(np.arange(num_documents, dtype=np.int32), num_epochs)
         if separate_last_epoch:
-            np_rng.shuffle(doc_idx[: -len(self)])
-            np_rng.shuffle(doc_idx[-len(self) :])
+            np_rng.shuffle(doc_idx[:-num_documents])
+            np_rng.shuffle(doc_idx[-num_documents:])
         else:
             np_rng.shuffle(doc_idx)
 
         assert _extension_available, (
             "The C++ extension for dataset sampling is missing." " Please make sure Fast-LLM is installed correctly."
         )
 
-        sample_idx = build_sample_idx(
-            self._indexed_dataset.sizes, doc_idx, sequence_length, num_epochs, tokens_per_epoch, verbose
-        )
+        sample_idx = build_sample_idx(document_sizes, doc_idx, sequence_length, num_epochs, num_tokens, verbose)
 
         # shuffle-idx.
         # -1 is due to data structure used to retrieve the index:

fast_llm/data/gpt/memmap.py

@@ -3,10 +3,11 @@
 
 import numpy as np
 
+from fast_llm.data.gpt.dataset import GPTIndexedDataset
 from fast_llm.utils import Assert, div, padded_cumsum
 
 
-class GPTMemmapDataset:
+class GPTMemmapDataset(GPTIndexedDataset):
     """
     A memory map dataset, which handles lazy loading of a pre-processed dataset in the Megatron-LM format,
     i.e. a pair of numpy file containing
@@ -27,11 +28,12 @@ class GPTMemmapDataset:
     }
     _INDEX_HEADER = b"MMIDIDX\x00\x00"
 
-    def __init__(self, prefix: pathlib.Path | str):
-        self._init(prefix)
+    def __init__(self, name: str, prefix: pathlib.Path | str):
+        self._init(name, prefix)
 
-    def _init(self, prefix: pathlib.Path | str):
+    def _init(self, name: str, prefix: pathlib.Path | str):
         super().__init__()
+        self._name = name
         self._prefix = pathlib.Path(prefix)
 
         with self._prefix.with_suffix(".idx").open("rb") as stream:
@@ -45,48 +47,58 @@ def _init(self, prefix: pathlib.Path | str):
 
         self._index_bin_buffer_mmap = np.memmap(self._prefix.with_suffix(".idx"), mode="r", order="C")
         self._index_bin_buffer = memoryview(self._index_bin_buffer_mmap)
-        self.sizes = np.frombuffer(self._index_bin_buffer, dtype=np.int32, count=self._num_documents, offset=offset)
+        self._document_sizes = np.frombuffer(
+            self._index_bin_buffer, dtype=np.int32, count=self._num_documents, offset=offset
+        )
         self._pointers = np.frombuffer(
-            self._index_bin_buffer, dtype=np.int64, count=self._num_documents, offset=offset + self.sizes.nbytes
+            self._index_bin_buffer,
+            dtype=np.int64,
+            count=self._num_documents,
+            offset=offset + self._document_sizes.nbytes,
         )
 
         self._bin_buffer_mmap = np.memmap(self._prefix.with_suffix(".bin"), mode="r", order="C")
         self._bin_buffer = memoryview(self._bin_buffer_mmap)
 
     def __getstate__(self):
-        return self.prefix
+        return (self._name, self._prefix)
 
     def __setstate__(self, state):
-        self._init(state)
+        self._init(*state)
 
     def __del__(self):
         self._bin_buffer_mmap._mmap.close()  # noqa
         del self._bin_buffer_mmap
         self._index_bin_buffer_mmap._mmap.close()  # noqa
         del self._index_bin_buffer_mmap
 
-    def __len__(self):
-        return self._num_documents
-
     def get(self, idx, offset=0, length=None):
         return np.frombuffer(
             self._bin_buffer,
             dtype=self._dtype,
-            count=self.sizes[idx] - offset if length is None else length,
+            count=self._document_sizes[idx] - offset if length is None else length,
             offset=self._pointers[idx] + offset * np.dtype(self._dtype).itemsize,
         )
 
     @property
-    def num_documents(self):
+    def name(self):
+        return self._name
+
+    @property
+    def num_documents(self) -> int:
         return self._num_documents
 
     @property
-    def num_tokens(self):
+    def num_tokens(self) -> int:
         return div(self._bin_buffer_mmap.size, np.dtype(self._dtype).itemsize)
 
-    @property
-    def prefix(self):
-        return self._prefix
+    def get_document_sizes(self) -> "np.ndarray":
+        """
+        The size of each document in the dataset.
+        The resulting array could be very large, so this method should be called cautiously,
+        and derived classes should try to avoid holding the whole array im memory.
+        """
+        return self._document_sizes
 
     @classmethod
     def write_dataset(cls, prefix: pathlib.Path | str, documents: list[np.ndarray]):