Add multiprocessing support to datasets (#150)

Useful where the data is sharded across files, as multiple files can be processed in parallel (when pre-downloaded).

Author: alan-cooney

pyproject.toml

@@ -180,7 +180,7 @@
             convention="google"
 
         [tool.ruff.lint.pylint]
-            max-args=10
+            max-args=15
 
 [tool.pyright]
     # All rules apart from base are shown explicitly below

sparse_autoencoder/optimizer/adam_with_reset.py

@@ -33,7 +33,7 @@ class AdamWithReset(Adam, AbstractOptimizerWithReset):
     The names of the parameters, so that we can find them later when resetting the state.
     """
 
-    def __init__(  # noqa: PLR0913 (extending existing implementation)
+    def __init__(  # (extending existing implementation)
         self,
         params: params_t,
         lr: float | Tensor = 1e-3,

sparse_autoencoder/source_data/abstract_dataset.py

@@ -111,6 +111,7 @@ def __init__(
         buffer_size: int = 1000,
         dataset_dir: str | None = None,
         dataset_files: str | Sequence[str] | Mapping[str, str | Sequence[str]] | None = None,
+        n_processes_preprocessing: int | None = None,
         preprocess_batch_size: int = 1000,
         *,
         pre_download: bool = False,
@@ -135,6 +136,7 @@ def __init__(
                 tokenized prompts once the preprocessing function has been applied.
             dataset_dir: Defining the `data_dir` of the dataset configuration.
             dataset_files: Path(s) to source data file(s).
+            n_processes_preprocessing: The number of processes to use for preprocessing.
             preprocess_batch_size: The batch size to use just for preprocessing the dataset (e.g.
                 tokenizing prompts).
             pre_download: Whether to pre-download the whole dataset.
@@ -146,43 +148,53 @@ def __init__(
 
         # Load the dataset
         should_stream = not pre_download
-        loaded_dataset = load_dataset(
+        dataset = load_dataset(
             dataset_path,
             streaming=should_stream,
             split=dataset_split,
             data_dir=dataset_dir,
             data_files=dataset_files,
         )
 
-        # Check the dataset is a Hugging Face Dataset or IterableDataset
-        if not isinstance(loaded_dataset, Dataset) and not isinstance(
-            loaded_dataset, IterableDataset
-        ):
-            error_message = (
-                f"Expected Hugging Face dataset to be a Dataset or IterableDataset, but got "
-                f"{type(loaded_dataset)}."
-            )
-            raise TypeError(error_message)
-
-        dataset: Dataset | IterableDataset = loaded_dataset
-
         # Setup preprocessing
         existing_columns: list[str] = list(next(iter(dataset)).keys())
-        mapped_dataset = dataset.map(
-            self.preprocess,
-            batched=True,
-            batch_size=preprocess_batch_size,
-            fn_kwargs={"context_size": context_size},
-            remove_columns=existing_columns,
-        )
 
         if pre_download:
+            if not isinstance(dataset, Dataset):
+                error_message = (
+                    f"Expected Hugging Face dataset to be a Dataset when pre-downloading, but got "
+                    f"{type(dataset)}."
+                )
+                raise TypeError(error_message)
+
             # Download the whole dataset
+            mapped_dataset = dataset.map(
+                self.preprocess,
+                batched=True,
+                batch_size=preprocess_batch_size,
+                fn_kwargs={"context_size": context_size},
+                remove_columns=existing_columns,
+                num_proc=n_processes_preprocessing,
+            )
             self.dataset = mapped_dataset.shuffle()
         else:
             # Setup approximate shuffling. As the dataset is streamed, this just pre-downloads at
             # least `buffer_size` items and then shuffles just that buffer.
             # https://huggingface.co/docs/datasets/v2.14.5/stream#shuffle
+            if not isinstance(dataset, IterableDataset):
+                error_message = (
+                    f"Expected Hugging Face dataset to be an IterableDataset when streaming, but "
+                    f"got {type(dataset)}."
+                )
+                raise TypeError(error_message)
+
+            mapped_dataset = dataset.map(
+                self.preprocess,
+                batched=True,
+                batch_size=preprocess_batch_size,
+                fn_kwargs={"context_size": context_size},
+                remove_columns=existing_columns,
+            )
             self.dataset = mapped_dataset.shuffle(buffer_size=buffer_size)  # type: ignore
 
     @final

sparse_autoencoder/source_data/text_dataset.py

@@ -72,6 +72,7 @@ def __init__(
         dataset_dir: str | None = None,
         dataset_files: str | Sequence[str] | Mapping[str, str | Sequence[str]] | None = None,
         dataset_split: str = "train",
+        n_processes_preprocessing: int | None = None,
         preprocess_batch_size: int = 1000,
         *,
         pre_download: bool = False,
@@ -94,6 +95,7 @@ def __init__(
             dataset_dir: Defining the `data_dir` of the dataset configuration.
             dataset_files: Path(s) to source data file(s).
             dataset_split: Dataset split (e.g., 'train').
+            n_processes_preprocessing: Number of processes to use for preprocessing.
             preprocess_batch_size: Batch size for preprocessing (tokenizing prompts).
             pre_download: Whether to pre-download the whole dataset.
         """
@@ -106,6 +108,7 @@ def __init__(
             dataset_files=dataset_files,
             dataset_path=dataset_path,
             dataset_split=dataset_split,
+            n_processes_preprocessing=n_processes_preprocessing,
             pre_download=pre_download,
             preprocess_batch_size=preprocess_batch_size,
         )

sparse_autoencoder/train/pipeline.py

@@ -79,7 +79,7 @@ class Pipeline:
     """Total number of activations trained on state."""
 
     @final
-    def __init__(  # noqa: PLR0913
+    def __init__(
         self,
         activation_resampler: AbstractActivationResampler | None,
         autoencoder: SparseAutoencoder,

sparse_autoencoder/train/sweep.py

@@ -141,12 +141,24 @@ def setup_source_data(hyperparameters: RuntimeHyperparameters) -> SourceDataset:
     Raises:
         ValueError: If the tokenizer name is not specified, but pre_tokenized is False.
     """
+    dataset_dir = (
+        hyperparameters["source_data"]["dataset_dir"]
+        if "dataset_dir" in hyperparameters["source_data"]
+        else None
+    )
+
+    dataset_files = (
+        hyperparameters["source_data"]["dataset_files"]
+        if "dataset_files" in hyperparameters["source_data"]
+        else None
+    )
+
     if hyperparameters["source_data"]["pre_tokenized"]:
         return PreTokenizedDataset(
             dataset_path=hyperparameters["source_data"]["dataset_path"],
             context_size=hyperparameters["source_data"]["context_size"],
-            dataset_dir=hyperparameters["source_data"]["dataset_dir"],
-            dataset_files=hyperparameters["source_data"]["dataset_files"],
+            dataset_dir=dataset_dir,
+            dataset_files=dataset_files,
         )
 
     if hyperparameters["source_data"]["tokenizer_name"] is None:
@@ -162,8 +174,9 @@ def setup_source_data(hyperparameters: RuntimeHyperparameters) -> SourceDataset:
         dataset_path=hyperparameters["source_data"]["dataset_path"],
         context_size=hyperparameters["source_data"]["context_size"],
         tokenizer=tokenizer,
-        dataset_dir=hyperparameters["source_data"]["dataset_dir"],
-        dataset_files=hyperparameters["source_data"]["dataset_files"],
+        dataset_dir=dataset_dir,
+        dataset_files=dataset_files,
+        n_processes_preprocessing=4,
     )
 
 

sparse_autoencoder/train/sweep_config.py

@@ -174,16 +174,16 @@ class SourceDataHyperparameters(NestedParameter):
     context_size: Parameter[int] = field(default=Parameter(DEFAULT_SOURCE_CONTEXT_SIZE))
     """Context size."""
 
-    dataset_dir: Parameter[str | None] = field(default=Parameter(None))
+    dataset_dir: Parameter[str] | None = field(default=None)
     """Dataset directory (within the HF dataset)"""
 
-    dataset_files: Parameter[str | None] = field(default=Parameter(None))
+    dataset_files: Parameter[str] | None = field(default=None)
     """Dataset files (within the HF dataset)."""
 
     pre_tokenized: Parameter[bool] = field(default=Parameter(value=True))
     """If the dataset is pre-tokenized."""
 
-    tokenizer_name: Parameter[str | None] = field(default=Parameter(None))
+    tokenizer_name: Parameter[str] | None = field(default=None)
     """Tokenizer name.
 
     Only set this if the dataset is not pre-tokenized.
@@ -195,11 +195,11 @@ def __post_init__(self) -> None:
         Raises:
             ValueError: If there is an error in the source data hyperparameters.
         """
-        if self.pre_tokenized.value is False and self.tokenizer_name.value is None:
+        if self.pre_tokenized.value is False and not isinstance(self.tokenizer_name, Parameter):
             error_message = "The tokenizer name must be specified, when `pre_tokenized` is False."
             raise ValueError(error_message)
 
-        if self.pre_tokenized.value is True and self.tokenizer_name.value is not None:
+        if self.pre_tokenized.value is True and isinstance(self.tokenizer_name, Parameter):
             error_message = "The tokenizer name must not be set, when `pre_tokenized` is True."
             raise ValueError(error_message)
 

sparse_autoencoder/train/utils/wandb_sweep_types.py

@@ -5,7 +5,7 @@
 from abc import ABC
 from dataclasses import asdict, dataclass, is_dataclass
 from enum import Enum, auto
-from typing import Any, Generic, TypeAlias, TypeVar, final
+from typing import Any, Generic, TypeVar, final
 
 from strenum import LowercaseStrEnum
 
@@ -264,11 +264,12 @@ def __repr__(self) -> str:
         return self.__str__()
 
 
-OptionalFloat: TypeAlias = float | None
-OptionalInt: TypeAlias = int | None
-OptionalStr: TypeAlias = str | None
-
-ParamType = TypeVar("ParamType", float, int, str, OptionalFloat, OptionalInt, OptionalStr)
+ParamType = TypeVar(
+    "ParamType",
+    float,
+    int,
+    str,
+)
 
 
 @dataclass(frozen=True)