Samplable dataset (#40)

Author: jlamypoirier

fast_llm/data/config.py

@@ -1,5 +1,6 @@
 import abc
 import enum
+import pathlib
 import typing
 
 from fast_llm.config import Config, Field, FieldHint, check_field, config_class, skip_valid_if_none
@@ -159,7 +160,20 @@ def name(self):
         """
 
 
-class SampledDataset(Dataset):  # noqa
+@config_class
+class SamplingConfig(Config):
+    num_samples: int = Field(default=1, desc="Number of samples to generate.")
+    seed: int = Field(default=0, desc="Random seed.")
+    cache_directory: pathlib.Path | None = Field(default=None, desc="Path to the sampling cache directory.")
+    verbose: bool = Field(default=True, desc="Log sampling progress.")
+
+
+class SamplableDataset(Dataset):
+    def sample(self, config: SamplingConfig, data: Data):
+        pass
+
+
+class SampledDataset(Dataset):
     """
     A sampled dataset class containing a prepared list of samples to be indexed sequentially (as-is) during training.
     (See the `Sampler` class below.)

fast_llm/data/gpt/config.py

@@ -12,7 +12,7 @@
 
 
 @config_class()
-class DataConfig(DataConfig):
+class GPTDataConfig(DataConfig):
     """
     Configuration for the dataset(s), split and sampling.
     Currently hard-coded to a GPT dataset.

fast_llm/data/gpt/data.py

@@ -10,10 +10,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.config import GPTDataConfig
+from fast_llm.data.gpt.dataset import GPTSamplingConfig
 from fast_llm.data.gpt.dummy import DummyGPTDataset
 from fast_llm.data.gpt.memmap import GPTMemmapDataset
-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
@@ -40,10 +40,11 @@ class GPTData(Data):
     _cache_directory: pathlib.Path | None
     _samples_per_phase: dict[PhaseType, int]
     _phases: typing.ClassVar[tuple[PhaseType, ...]] = (PhaseType.training, PhaseType.validation, PhaseType.test)
+    _is_setup: bool = False
 
     def __init__(
         self,
-        config: DataConfig,
+        config: GPTDataConfig,
         distributed_config: DistributedConfig,
         vocab_size: int,
         max_sequence_length: int,
@@ -52,8 +53,8 @@ def __init__(
         Create the data and gather some basic information on the dataset(s).
         Should be `setup` before use.
         """
-        self._config = config.validate()
-        self._distributed_config = distributed_config.validate()
+        self._config = config
+        self._distributed_config = distributed_config
         self._vocab_size = vocab_size
         self._max_sequence_length = max_sequence_length
         Assert.eq(len(self._config.split), len(self._phases))
@@ -166,6 +167,20 @@ def setup(self, distributed: Distributed, samples_per_phase: dict[PhaseType, int
             )
             for phase, datasets in self._sampled_datasets.items()
         }
+        self._is_setup = True
+
+    @property
+    def config(self):
+        return self._config
+
+    @property
+    def tokenizer(self):
+        assert self._is_setup
+        return self._tokenizer
+
+    @property
+    def distributed(self):
+        return self._distributed
 
     def get_iterator(
         self,
@@ -176,6 +191,7 @@ def get_iterator(
         num_workers: int,
         prefetch_factor: int | None = None,
     ):
+        assert self._is_setup
         Assert.incl(phase, self._blended_datasets)
         Assert.in_range_incl(batch_config.sequence_length, 1, self._max_sequence_length)
         log_main_rank(f"Initializing {phase} data iterator from sample {consumed_samples}...")
@@ -205,25 +221,23 @@ def _build_and_sample_gpt_dataset(self, name: str, dataset_samples_per_phase: di
         for phase, num_samples in dataset_samples_per_phase.items():
             if num_samples == 0:
                 continue
-            sampled_datasets[phase] = GPTSampledIndexedDataset(
-                dataset_split[phase],
-                num_samples=num_samples,
-                sequence_length=self._max_sequence_length,
-                seed=self._distributed.config.seed,
-                group=self._distributed.world_group,
-                config=self._config,
-                tokenizer=self._tokenizer,
-                cache_directory=(
-                    self._dataset_prefixes[name].parent if self._cache_directory is None else self._cache_directory
+            sampled_datasets[phase] = dataset_split[phase].sample(
+                GPTSamplingConfig(
+                    num_samples=num_samples,
+                    sequence_length=self._max_sequence_length,
+                    seed=self._distributed_config.seed,
+                    cache_directory=(
+                        self._dataset_prefixes[name].parent if self._cache_directory is None else self._cache_directory
+                    ),
+                    verbose=self._num_datasets <= 5,
                 ),
-                verbose=self._num_datasets <= 5,
+                self,
             )
         return sampled_datasets
 
     def _build_and_sample_dummy_dataset(self, name: str, dataset_samples_per_phase: dict[PhaseType, int]):
         return {
             phase: DummyGPTDataset(
-                self._dataset_prefixes[name],
                 dataset_samples_per_phase[phase],
                 self._max_sequence_length,
                 self._vocab_size,

fast_llm/data/gpt/dataset.py

@@ -1,11 +1,14 @@
 import abc
-import math
+import typing
 
 import numpy as np
-import numpy.random
 
-from fast_llm.data.config import Dataset
-from fast_llm.utils import Assert
+from fast_llm.config import Field, config_class
+from fast_llm.data.config import SamplableDataset, SamplingConfig
+
+if typing.TYPE_CHECKING:
+    from fast_llm.data.gpt.data import GPTData
+
 
 try:
     from fast_llm.csrc.data import build_sample_idx  # noqa
@@ -15,7 +18,12 @@
     _extension_available = False
 
 
-class GPTIndexedDataset(Dataset):
+@config_class
+class GPTSamplingConfig(SamplingConfig):
+    sequence_length: int = Field(default=None, desc="Number of token in each sample.")
+
+
+class GPTIndexedDataset(SamplableDataset):
     """
     A GPT dataset containing a list of unsampled, unprocessed samples.
     TODO: Move sampling responsibility here?
@@ -48,52 +56,7 @@ def get_document_sizes(self) -> "np.ndarray":
         and derived classes should try to avoid holding the whole array im memory.
         """
 
-    def sample(self, num_samples: int, sequence_length: int, np_rng: numpy.random.RandomState, verbose: bool):
-        """
-        Create a `GPTSampledDataset` with the requested parameters.
-        """
-        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) * num_tokens - 1) // sequence_length
-        last_epoch_samples = num_samples - main_epochs_samples
-        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(num_documents, dtype=np.int32), num_epochs)
-        if separate_last_epoch:
-            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(document_sizes, doc_idx, sequence_length, num_epochs, num_tokens, verbose)
-
-        # shuffle-idx.
-        # -1 is due to data structure used to retrieve the index:
-        #    sample i --> [sample_idx[i], sample_idx[i+1])
-        total_size = sample_idx.shape[0] - 1
-        # TODO: Isn't the dataset already shuffled above?
-        shuffle_idx = np.arange(
-            0, total_size, dtype=np.int64 if total_size >= (np.iinfo(np.uint32).max - 1) else np.uint32
-        )
-        if separate_last_epoch:
-            np_rng.shuffle(shuffle_idx[:main_epochs_samples])
-            np_rng.shuffle(shuffle_idx[main_epochs_samples:])
-        else:
-            np_rng.shuffle(shuffle_idx)
-
-        Assert.geq(len(shuffle_idx), num_samples)
-        # TODO: The doc and sample idx are way bigger than needed when sampling for << 1 epoch.
-        return doc_idx, sample_idx, shuffle_idx[:num_samples]
+    def sample(self, config: GPTSamplingConfig, data: "GPTData"):
+        from fast_llm.data.gpt.sampled import GPTSampledIndexedDataset
+
+        return GPTSampledIndexedDataset(self, config, data)

fast_llm/data/gpt/dummy.py

@@ -1,10 +1,6 @@
-import pathlib
-
 import numpy as np
 
 from fast_llm.data.config import SampledDataset
-from fast_llm.engine.config_utils.run import log_main_rank
-from fast_llm.utils import Assert
 
 
 class DummyGPTDataset(SampledDataset):
@@ -13,19 +9,9 @@ class DummyGPTDataset(SampledDataset):
     The sample can be purely random, or read from a file to allow reproducing in other runs.
     """
 
-    def __init__(
-        self, prefix: pathlib.Path | None, num_samples: int, sequence_length: int, vocab_size: int, name: str = "dummy"
-    ):
+    def __init__(self, num_samples: int, sequence_length: int, vocab_size: int, name: str = "dummy"):
         self._num_samples = num_samples
-        if prefix is None:
-            self._dummy_sample = np.random.randint(0, vocab_size, size=(sequence_length + 1,), dtype=np.int64)
-        else:
-            log_main_rank(f"> Loading dummy dataset from file {prefix}")
-            self._dummy_sample = np.load(prefix, allow_pickle=True)[: sequence_length + 1]
-            Assert.eq(self._dummy_sample.shape, (sequence_length + 1,))
-            Assert.eq(self._dummy_sample.dtype, np.int64)
-            Assert.lt(self._dummy_sample.max(), vocab_size)
-            Assert.geq(self._dummy_sample.min(), 0)
+        self._dummy_sample = np.random.randint(0, vocab_size, size=(sequence_length + 1,), dtype=np.int64)
         self._name = name
 
     def __len__(self):