datatrove is all you need

Author: TJ-Solergibert

examples/config_nanoset.yaml

@@ -7,25 +7,25 @@ checkpoints:
 data_stages:
 - data:
     dataset:
-      dataset_path: datasets/testing_alpaca_small_input_ids.npy
+      dataset_folder: datasets/c4-es/tokenized
     num_loading_workers: 1
     seed: 42
   name: General purpose training (Single dataset)
   start_training_step: 1
 - data:
     dataset:
-      dataset_path:
-      - datasets/yelp_review_full_input_ids.npy
-      - datasets/testing_alpaca_small_input_ids.npy
+      dataset_folder:
+      - datasets/SlimPajama-6B/tokenized
+      - datasets/c4-es/tokenized
     num_loading_workers: 1
     seed: 42
   name: Second purpose training (> 1 dataset)
   start_training_step: 15
 - data:
     dataset:
-      dataset_path:
-        datasets/testing_alpaca_small_input_ids.npy: 0.8
-        datasets/yelp_review_full_input_ids.npy: 0.2
+      dataset_folder:
+        datasets/SlimPajama-6B/tokenized: 0.8
+        datasets/c4-es/tokenized: 0.2
     num_loading_workers: 1
     seed: 42
   name: Third purpose training (Blended dataset)
@@ -57,7 +57,7 @@ model:
     initializer_range: 0.02
     intermediate_size: 64
     is_llama_config: true
-    max_position_embeddings: 256
+    max_position_embeddings: 1024
     num_attention_heads: 4
     num_hidden_layers: 2
     num_key_value_heads: 4
@@ -67,7 +67,7 @@ model:
     rope_scaling: null
     tie_word_embeddings: true
     use_cache: true
-    vocab_size: 32000
+    vocab_size: 50257
 optimizer:
   accumulate_grad_in_fp32: true
   clip_grad: 1.0
@@ -88,11 +88,11 @@ optimizer:
   weight_decay: 0.01
   zero_stage: 0
 parallelism:
-  dp: 2
+  dp: 1
   expert_parallel_size: 1
   pp: 1
   pp_engine: 1f1b
-  tp: 2
+  tp: 1
   tp_linear_async_communication: true
   tp_mode: REDUCE_SCATTER
 profiler: null
@@ -105,6 +105,6 @@ tokens:
   limit_test_batches: 0
   limit_val_batches: 0
   micro_batch_size: 2
-  sequence_length: 128
+  sequence_length: 1024
   train_steps: 200
   val_check_interval: -1

pyproject.toml

@@ -49,7 +49,7 @@ fast-modeling = [
 
 nanosets = [
      "transformers",
-     "datasets",
+     "datatrove[io,processing]",
      "numba",
 ]
 

run_train.py

@@ -143,17 +143,17 @@ def get_dataloader_from_data_stage(
     elif isinstance(data.dataset, NanosetDatasetsArgs):
         # Get tokenizer cardinality
         tokenizer = AutoTokenizer.from_pretrained(trainer.config.tokenizer.tokenizer_name_or_path)
-        token_dtype = np.int32 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else np.uint16
+        token_size = 4 if len(tokenizer) > np.iinfo(np.uint16).max + 1 else 2
         del tokenizer
         # Create Nanoset
         from nanotron.data.nanoset import Nanoset
 
         with main_rank_first(trainer.parallel_context.world_pg):
             train_dataset = Nanoset(
-                dataset_paths=data.dataset.dataset_path,
+                dataset_folders=data.dataset.dataset_folder,
                 dataset_weights=data.dataset.dataset_weights,
                 sequence_length=trainer.sequence_length,
-                token_dtype=token_dtype,
+                token_size=token_size,
                 train_split_num_samples=trainer.config.tokens.train_steps * trainer.global_batch_size,
                 random_seed=data.seed,
             )

src/nanotron/config/config.py

@@ -93,18 +93,18 @@ def __post_init__(self):
 
 @dataclass
 class NanosetDatasetsArgs:
-    dataset_path: Union[str, dict, List[str]]
+    dataset_folder: Union[str, dict, List[str]]
 
     def __post_init__(self):
-        if isinstance(self.dataset_path, str):  # Case 1: 1 Dataset file
-            self.dataset_path = [self.dataset_path]
+        if isinstance(self.dataset_folder, str):  # Case 1: 1 Dataset file
+            self.dataset_folder = [self.dataset_folder]
             self.dataset_weights = [1]
-        elif isinstance(self.dataset_path, List):  # Case 2: > 1 Dataset file
+        elif isinstance(self.dataset_folder, List):  # Case 2: > 1 Dataset file
             self.dataset_weights = None  # Set to None so we consume all the samples randomly
-        elif isinstance(self.dataset_path, dict):  # Case 3: dict with > 1 dataset_path and weights
-            tmp_dataset_path = self.dataset_path.copy()
-            self.dataset_path = list(tmp_dataset_path.keys())
-            self.dataset_weights = list(tmp_dataset_path.values())
+        elif isinstance(self.dataset_folder, dict):  # Case 3: dict with > 1 dataset_folder and weights
+            tmp_dataset_folder = self.dataset_folder.copy()
+            self.dataset_folder = list(tmp_dataset_folder.keys())
+            self.dataset_weights = list(tmp_dataset_folder.values())
 
 
 @dataclass

src/nanotron/data/collator.py

@@ -0,0 +1,80 @@
+import dataclasses
+from typing import Dict, List, Union
+
+import numpy as np
+import torch
+from nanotron import distributed as dist
+from nanotron.parallel.context import ParallelContext
+from nanotron.parallel.pipeline_parallel.tensor_pointer import TensorPointer
+
+
+@dataclasses.dataclass
+class NanosetDataCollatorForCLM:
+    """
+    Data collator used for causal language modeling with Nanosets dataset.
+
+    - input_pp_rank: Discards last input id token
+    - output_pp_rank: Discards first label id token
+    - other pp ranks: Don't have data. Instead, we use `TensorPointer` to point to the rank having the data.
+    """
+
+    sequence_length: int
+    input_pp_rank: int
+    output_pp_rank: int
+    parallel_context: ParallelContext
+
+    def __call__(self, examples: List[Dict[str, List[np.ndarray]]]) -> Dict[str, Union[torch.Tensor, TensorPointer]]:
+        # Process the case when current rank doesn't require data. We return `TensorPointer` that points to ranks having the data.
+        current_pp_rank = dist.get_rank(self.parallel_context.pp_pg)
+        if current_pp_rank not in [
+            self.input_pp_rank,
+            self.output_pp_rank,
+        ]:
+            assert all(len(example) == 0 for example in examples)
+            return {
+                "input_ids": TensorPointer(group_rank=self.input_pp_rank),
+                "input_mask": TensorPointer(group_rank=self.input_pp_rank),
+                "label_ids": TensorPointer(group_rank=self.output_pp_rank),
+                "label_mask": TensorPointer(group_rank=self.output_pp_rank),
+            }
+
+        # Make sure we load only what's necessary, ie we only load a `input_ids` column.
+        assert all(list(example.keys()) == ["input_ids"] for example in examples)
+
+        # TODO @nouamanetazi: Is it better to have examples as np.array or torch.Tensor?
+        input_ids = torch.vstack([examples[i]["input_ids"] for i in range(len(examples))])  # (b, s)
+        batch_size, expanded_input_length = input_ids.shape
+
+        result: Dict[str, Union[torch.LongTensor, TensorPointer]] = {}
+
+        result["input_ids"] = TensorPointer(group_rank=self.input_pp_rank)
+        result["input_mask"] = TensorPointer(group_rank=self.input_pp_rank)
+        result["label_ids"] = TensorPointer(group_rank=self.output_pp_rank)
+        result["label_mask"] = TensorPointer(group_rank=self.output_pp_rank)
+
+        assert (
+            expanded_input_length == self.sequence_length + 1
+        ), f"Samples should be of length {self.sequence_length + 1} (seq_len+1), but got {expanded_input_length}"
+
+        # Process inputs: last token is the label
+        if current_pp_rank == self.input_pp_rank:
+            result["input_ids"] = input_ids[:, :-1]
+            result["input_mask"] = torch.ones((batch_size, self.sequence_length), dtype=torch.bool)
+
+        # Process labels: shift them to the left
+        if current_pp_rank == self.output_pp_rank:
+            result["label_ids"] = input_ids[:, 1:]
+            result["label_mask"] = torch.ones((batch_size, self.sequence_length), dtype=torch.bool)
+
+        if isinstance(result["input_ids"], torch.Tensor) and result["input_ids"].shape[-1] != self.sequence_length:
+            raise ValueError(
+                f"`labels` are incorrectly preprocessed. `labels` length is {result['input_ids'].shape[-1]}, but should be"
+                f" {self.sequence_length}."
+            )
+        if isinstance(result["label_ids"], torch.Tensor) and result["label_ids"].shape[-1] != self.sequence_length:
+            raise ValueError(
+                f"`labels` are incorrectly preprocessed. `labels` length is {result['label_ids'].shape[-1]}, but should be"
+                f" {self.sequence_length}."
+            )
+
+        return result

src/nanotron/data/dataloader_builder.py

@@ -1,7 +1,7 @@
 import nanotron.distributed as dist
 from nanotron import logging
+from nanotron.data.collator import NanosetDataCollatorForCLM
 from nanotron.dataloader import (
-    DataCollatorForCLM,
     EmptyInfiniteDataset,
     get_dataloader_worker_init,
     get_sampler,
@@ -32,7 +32,7 @@ def build_nanoset_dataloader(
         # No need to spawn a lot of workers, we can just use main
         dataloader_num_workers = 0
 
-    data_collator = DataCollatorForCLM(
+    data_collator = NanosetDataCollatorForCLM(
         sequence_length=sequence_length,
         input_pp_rank=input_pp_rank,
         output_pp_rank=output_pp_rank,

src/nanotron/data/nanoset.py

@@ -1,7 +1,10 @@
+import os
+import warnings
 from typing import Dict, List, Tuple, Union
 
 import numpy as np
 import torch
+from datatrove.utils.dataset import DatatroveFolderDataset
 from nanotron import logging
 from nanotron.data.utils import count_dataset_indexes, normalize
 from nanotron.logging import log_rank
@@ -15,49 +18,61 @@ class Nanoset(torch.utils.data.Dataset):
     The Nanoset dataset
 
     Args:
-        dataset_paths (List[str]): List of paths to tokenized datasets
+        dataset_folders (List[str]): List of folders with tokenized datasets
         dataset_weights (List[float]): List with the weights for weighted datasets. If None, consume all samples from all datasets without weighting. Weights are normalized in __init__
         sequence_length (int): Sequence length of the built samples
-        token_dtype (Union[np.uint16, np.int32]): dtype of the tokens stored in the processed dataset files. np.uin16 for vocab sizes < 65535, np.int32 otherwise
+        token_size (int): Number of bytes for the tokens stored in the processed dataset files. 2 for vocab sizes < 65535, 4 otherwise
         train_split_num_samples (int): Number of samples the dataset needs. It's the training steps * global batch size
     """
 
     def __init__(
         self,
-        dataset_paths: List[str],
+        dataset_folders: List[str],
         dataset_weights: Union[List[float], None],
         sequence_length: int,
-        token_dtype: Union[np.uint16, np.int32],
+        token_size: int,
         train_split_num_samples: int,
         random_seed: int = 1234,
     ) -> None:
 
+        # Assertions
+        if isinstance(dataset_folders, str):
+            warnings.warn("dataset_folders should be of type List[str] but str was provided. Converting to List[str]")
+            dataset_folders = [dataset_folders]
+
         # Init
-        self.dataset_paths = dataset_paths
+        self.dataset_folders = dataset_folders
         self.dataset_weights = dataset_weights
         self.sequence_length = sequence_length
-        self.token_dtype = token_dtype
+        self.token_size = token_size
         self.train_split_num_samples = train_split_num_samples
         self.random_seed = random_seed
+        self.datatrove_datasets = []
+        for dataset_folder in self.dataset_folders:
+            self.datatrove_datasets.append(
+                DatatroveFolderDataset(
+                    folder_path=dataset_folder,
+                    filename_pattern=os.path.join(dataset_folder, "*.ds"),
+                    seq_len=sequence_length,
+                    recursive=False,
+                    token_size=token_size,
+                    shuffle=True,
+                )
+            )
 
         # Build Nanoset Index
         ## To build the index we need the length of each dataset
-        self.dataset_lengths = []
-        for dataset_path in self.dataset_paths:
-            self.dataset_buffer_mmap = np.memmap(dataset_path, mode="r", order="C", dtype=self.token_dtype)
-            self.dataset_buffer = memoryview(self.dataset_buffer_mmap)
-            dataset_number_of_tokens = int(len(self.dataset_buffer))
-            number_of_samples = int(
-                (dataset_number_of_tokens - 1) / sequence_length
-            )  # Discard last sample if length < sequence_length
-            self.dataset_lengths.append(number_of_samples)
+        self.dataset_lengths = [len(datatrove_dataset) for datatrove_dataset in self.datatrove_datasets]
         ## Set dataset weights
         if (
             self.dataset_weights is None
         ):  # Case of training with > 1 datasets without weighting them: Consume both datasets entirely on each epoch
             self.dataset_weights = normalize(self.dataset_lengths)
         else:
             self.dataset_weights = normalize(dataset_weights)
+        assert len(dataset_folders) == len(
+            self.dataset_weights
+        ), f"Specified {len(self.dataset_weights)} weights but {len(dataset_folders)} datasets were provided."
         ## Build dataset index and dataset sample index
         self.dataset_index, self.dataset_sample_index = self.build_nanoset_index()
 
@@ -79,25 +94,12 @@ def __getitem__(self, idx: int) -> Dict[str, np.ndarray]:
             idx (int): The index into the dataset
 
         Returns:
-            Dict[str, numpy.ndarray]: The input ids wrapped in a dictionary
+            Dict[str, torch.LongTensor]: The input ids wrapped in a dictionary
         """
-
         dataset = self.dataset_index[idx]
         dataset_sample = self.dataset_sample_index[idx]
 
-        # Rebuild the memmap in every access to free memory
-        # https://stackoverflow.com/a/61472122
-        self.dataset_buffer_mmap = np.memmap(self.dataset_paths[dataset], mode="r", order="C", dtype=self.token_dtype)
-        self.dataset_buffer = memoryview(self.dataset_buffer_mmap)
-
-        # uint16 -> 2 bytes per token, int32 -> 4 bytes per token
-        offset = dataset_sample * self.sequence_length * (np.iinfo(self.token_dtype).bits / 8)
-        input_ids_tokens = np.frombuffer(
-            self.dataset_buffer, dtype=self.token_dtype, count=(self.sequence_length + 1), offset=int(offset)
-        )
-
-        # Return tokens as np.int32 as Torch can't handle uint16
-        return {"input_ids": input_ids_tokens.astype(np.int32)}
+        return self.datatrove_datasets[dataset][dataset_sample]
 
     def build_nanoset_index(self) -> np.ndarray:
         """
@@ -124,15 +126,6 @@ def build_nanoset_index(self) -> np.ndarray:
 
         return dataset_index, dataset_sample_index
 
-    def __del__(self) -> None:
-        """
-        Clean up Nanoset
-        """
-
-        if hasattr(self, "dataset_buffer_mmap"):
-            self.dataset_buffer_mmap._mmap.close()
-        del self.dataset_buffer_mmap
-
     def print_nanoset_info(self):
 
         log_rank(f"> Total number of samples: {len(self)}", logger=logger, level=logging.INFO, rank=0)
@@ -141,10 +134,10 @@ def print_nanoset_info(self):
         )
 
         # Print samples from each dataset + weight
-        dataset_sample_count = count_dataset_indexes(self.dataset_index, len(self.dataset_paths))
+        dataset_sample_count = count_dataset_indexes(self.dataset_index, len(self.dataset_folders))
         for index, sample_count in enumerate(dataset_sample_count):
             log_rank(
-                f">   Total number of samples from the {self.dataset_paths[index].rsplit('/', 1)[-1]} dataset: {sample_count} ({round(normalize(dataset_sample_count).tolist()[index], 2)})",
+                f">   Total number of samples from the {self.dataset_folders[index]} dataset: {sample_count} ({round(normalize(dataset_sample_count).tolist()[index], 2)})",
                 logger=logger,
                 level=logging.INFO,
                 rank=0,