moving dataset save fn to shared module

Author: djsaunde

src/axolotl/utils/data/lock.py

@@ -11,6 +11,7 @@
 
 LOCK_FILE_NAME = "datasets_prep.lock"
 READY_FILE_NAME = "datasets_ready.flag"
+PROCESS_COUNTER_FILE_NAME = "process_counter.txt"
 
 
 class FileLockLoader:
@@ -27,10 +28,14 @@ def __init__(self, cfg: DictDefault):
         )
         self.lock_file_path = Path(self.dataset_prepared_path) / LOCK_FILE_NAME
         self.ready_flag_path = Path(self.dataset_prepared_path) / READY_FILE_NAME
-        self.counter_path = Path(self.dataset_prepared_path) / "process_counter.txt"
+        self.counter_path = Path(self.dataset_prepared_path) / PROCESS_COUNTER_FILE_NAME
 
     def load(self, load_fn: Callable[[], Any]) -> Any:
+        import torch.distributed as dist
+
         with FileLock(str(self.lock_file_path)):
+            print(f"FileLock acquired by rank {dist.get_rank()}")
+
             # Increment process counter
             self._increment_counter()
 

src/axolotl/utils/data/rl.py

@@ -1,7 +1,6 @@
 """Data handling specific to RL trainers."""
 
 import inspect
-import os
 from functools import partial
 from typing import Any, Callable, Literal
 
@@ -13,15 +12,15 @@
 from axolotl.prompt_strategies.kto import load as load_kto
 from axolotl.prompt_strategies.orpo import load as load_orpo
 from axolotl.utils.data.lock import FileLockLoader
-from axolotl.utils.data.sft import _try_load_from_hub
 from axolotl.utils.data.shared import (
     create_train_validation_split,
     datasets_with_name_generator,
     generate_dataset_hash_from_config,
-    get_prepared_dataset_path,
     load_dataset_with_config,
     load_preprocessed_dataset,
     merge_datasets,
+    save_preprocessed_dataset,
+    try_load_from_hub,
 )
 from axolotl.utils.data.utils import (
     deduplicate_and_log_datasets,
@@ -82,22 +81,6 @@ def _load_datasets():
     return train_dataset, eval_dataset
 
 
-def _save_preprocessed_dataset(
-    cfg: DictDefault, dataset: Dataset, dataset_hash: str
-) -> None:
-    """Save preprocessed dataset to disk.
-
-    Args:
-        cfg: Configuration object.
-        dataset: Dataset to save.
-        dataset_hash: Hash identifying the dataset configuration.
-    """
-    prepared_ds_path = get_prepared_dataset_path(cfg, dataset_hash)
-    LOG.info(f"Saving prepared dataset to disk... {prepared_ds_path}")
-    os.makedirs(prepared_ds_path, exist_ok=True)
-    dataset.save_to_disk(str(prepared_ds_path))
-
-
 def _map_dataset(
     cfg: DictDefault,
     dataset: Dataset | DatasetDict,
@@ -265,7 +248,7 @@ def _load_split(cfg: DictDefault, split: Literal["train", "test"]) -> Dataset:
         dataset_hash = generate_dataset_hash_from_config(
             cfg, cfg.datasets, tokenizer.name_or_path
         )
-        _save_preprocessed_dataset(cfg, dataset, dataset_hash)
+        save_preprocessed_dataset(cfg, dataset, dataset_hash, split)
 
     return dataset
 
@@ -295,7 +278,7 @@ def _load_or_create_dataset_split(
     # Try loading from hub if push_dataset_to_hub is configured
     dataset = None
     if cfg.push_dataset_to_hub:
-        dataset = _try_load_from_hub(cfg, dataset_hash, split)
+        dataset = try_load_from_hub(cfg, dataset_hash, split)
 
     # Attempt to load preprocessed dataset
     if dataset is None:

src/axolotl/utils/data/sft.py

@@ -1,9 +1,8 @@
 """Data handling specific to SFT."""
 
 import functools
-import os
 import tempfile
-from typing import Any, Generator, Literal
+from typing import Literal
 
 from datasets import (
     Dataset,
@@ -20,10 +19,11 @@
     create_train_validation_split,
     datasets_with_name_generator,
     generate_dataset_hash_from_config,
-    get_prepared_dataset_path,
     load_dataset_with_config,
     load_preprocessed_dataset,
     merge_datasets,
+    save_preprocessed_dataset,
+    try_load_from_hub,
 )
 from axolotl.utils.data.utils import (
     deduplicate_and_log_datasets,
@@ -275,7 +275,7 @@ def _load_tokenized_prepared_datasets(
     # Try loading from hub if push_dataset_to_hub is configured
     dataset = None
     if cfg.push_dataset_to_hub:
-        dataset = _try_load_from_hub(cfg, dataset_hash, split)
+        dataset = try_load_from_hub(cfg, dataset_hash, split)
 
     # If not found on hub, try loading from disk
     if dataset is None:
@@ -296,71 +296,6 @@ def _load_tokenized_prepared_datasets(
     return dataset, prompters
 
 
-def _try_load_from_hub(
-    cfg: DictDefault, dataset_hash: str, split: str
-) -> Dataset | None:
-    """Try to load the prepared dataset from HuggingFace Hub."""
-    try:
-        LOG.info(
-            "Attempting to load prepared dataset from HuggingFace Hub at "
-            f"{cfg.push_dataset_to_hub} (version {dataset_hash})..."
-        )
-        dataset = load_dataset(
-            cfg.push_dataset_to_hub,
-            dataset_hash,
-            token=cfg.hf_use_auth_token,
-        )
-        return dataset[split]
-    except Exception:  # pylint: disable=broad-except # nosec
-        LOG.info("Unable to find prepared dataset in HuggingFace Hub")
-        return None
-
-
-def _generate_from_iterable_dataset(
-    dataset: IterableDataset, worker_id: list[int], num_workers: list[int]
-) -> Generator[Any, None, None]:
-    """Generator function to correctly split the dataset for each worker"""
-    for i, item in enumerate(dataset):
-        if i % num_workers[0] == worker_id[0]:
-            yield item
-
-
-def _save_preprocessed_dataset(
-    cfg: DictDefault,
-    dataset: Dataset,
-    dataset_hash: str,
-    split: str,
-) -> None:
-    prepared_ds_path = get_prepared_dataset_path(cfg, dataset_hash)
-    if isinstance(dataset, IterableDataset):
-        num_workers = cfg.dataset_processes
-
-        ds_from_iter = Dataset.from_generator(
-            functools.partial(_generate_from_iterable_dataset, dataset),
-            features=dataset.features,
-            num_proc=num_workers,
-            split=split,
-            gen_kwargs={
-                "worker_id": list(range(num_workers)),
-                "num_workers": [num_workers] * num_workers,
-            },
-        )
-        ds_from_iter.save_to_disk(str(prepared_ds_path))
-    else:
-        os.makedirs(prepared_ds_path, exist_ok=True)
-        dataset.save_to_disk(str(prepared_ds_path))
-    if cfg.push_dataset_to_hub:
-        LOG.info(
-            "Pushing merged prepared dataset to Huggingface hub at "
-            f"{cfg.push_dataset_to_hub} (version {dataset_hash})..."
-        )
-        dataset.push_to_hub(
-            cfg.push_dataset_to_hub,
-            dataset_hash,
-            private=True,
-        )
-
-
 def _load_raw_datasets(
     cfg: DictDefault,
     cfg_datasets: list,
@@ -370,7 +305,7 @@ def _load_raw_datasets(
     preprocess_iterable: bool = False,
 ) -> tuple[Dataset, list[Prompter | None]]:
     """Load, process, merge, and save raw datasets."""
-    LOG.info("Loading raw datasets...")
+    LOG.info("Loading raw datasets...", main_process_only=False)
     if not cfg.is_preprocess:
         LOG.warning(
             "Processing datasets during training can lead to VRAM instability. Please "
@@ -405,7 +340,7 @@ def _load_raw_datasets(
         dataset_hash = generate_dataset_hash_from_config(
             cfg, cfg.datasets, tokenizer.name_or_path
         )
-        _save_preprocessed_dataset(cfg, dataset, dataset_hash, split)
+        save_preprocessed_dataset(cfg, dataset, dataset_hash, split)
 
     return dataset, prompters
 

src/axolotl/utils/data/shared.py

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+import functools
+import os
 from pathlib import Path
-from typing import TYPE_CHECKING, Generator
+from typing import TYPE_CHECKING, Any, Generator
 
 from datasets import (
     Dataset,
@@ -391,6 +393,53 @@ def create_train_validation_split(
     return split_dataset["train"], split_dataset["test"]
 
 
+def _generate_from_iterable_dataset(
+    dataset: IterableDataset, worker_id: list[int], num_workers: list[int]
+) -> Generator[Any, None, None]:
+    """Generator function to correctly split the dataset for each worker"""
+    for i, item in enumerate(dataset):
+        if i % num_workers[0] == worker_id[0]:
+            yield item
+
+
+def save_preprocessed_dataset(
+    cfg: DictDefault,
+    dataset: Dataset,
+    dataset_hash: str,
+    split: str,
+) -> None:
+    """Save preprocessed dataset to disk and optionally push to the HF Hub."""
+    prepared_ds_path = get_prepared_dataset_path(cfg, dataset_hash)
+    if isinstance(dataset, IterableDataset):
+        num_workers = cfg.dataset_processes
+
+        ds_from_iter = Dataset.from_generator(
+            functools.partial(_generate_from_iterable_dataset, dataset),
+            features=dataset.features,
+            num_proc=num_workers,
+            split=split,
+            gen_kwargs={
+                "worker_id": list(range(num_workers)),
+                "num_workers": [num_workers] * num_workers,
+            },
+        )
+        ds_from_iter.save_to_disk(str(prepared_ds_path))
+    else:
+        os.makedirs(prepared_ds_path, exist_ok=True)
+        dataset.save_to_disk(str(prepared_ds_path))
+    if cfg.push_dataset_to_hub:
+        LOG.info(
+            "Pushing merged prepared dataset to Huggingface hub at "
+            f"{cfg.push_dataset_to_hub} (version {dataset_hash})...",
+            main_process_only=False,
+        )
+        dataset.push_to_hub(
+            cfg.push_dataset_to_hub,
+            dataset_hash,
+            private=True,
+        )
+
+
 def load_preprocessed_dataset(cfg: DictDefault, dataset_hash: str) -> Dataset | None:
     """Load preprocessed dataset from disk if available.
 
@@ -409,13 +458,39 @@ def load_preprocessed_dataset(cfg: DictDefault, dataset_hash: str) -> Dataset |
         and not cfg.skip_prepare_dataset
         and not cfg.is_preprocess
     ):
-        LOG.info(f"Loading prepared dataset from disk at {prepared_ds_path}...")
+        LOG.info(
+            f"Loading prepared dataset from disk at {prepared_ds_path}...",
+            main_process_only=False,
+        )
         return load_from_disk(str(prepared_ds_path))
 
-    LOG.info(f"Unable to find prepared dataset in {prepared_ds_path}")
+    LOG.info(
+        f"Unable to find prepared dataset in {prepared_ds_path}",
+        main_process_only=False,
+    )
     return None
 
 
+def try_load_from_hub(
+    cfg: DictDefault, dataset_hash: str, split: str
+) -> Dataset | None:
+    """Try to load the prepared dataset from HuggingFace Hub."""
+    try:
+        LOG.info(
+            "Attempting to load prepared dataset from HuggingFace Hub at "
+            f"{cfg.push_dataset_to_hub} (version {dataset_hash})..."
+        )
+        dataset = load_dataset(
+            cfg.push_dataset_to_hub,
+            dataset_hash,
+            token=cfg.hf_use_auth_token,
+        )
+        return dataset[split]
+    except Exception:  # pylint: disable=broad-except # nosec
+        LOG.info("Unable to find prepared dataset in HuggingFace Hub")
+        return None
+
+
 def generate_dataset_hash_from_config(
     cfg: DictDefault, cfg_datasets: list, tokenizer_name: str
 ) -> str:
@@ -451,13 +526,13 @@ def merge_datasets(datasets: list[Dataset], cfg: DictDefault) -> Dataset:
     if len(datasets) == 1:
         return datasets[0]
 
-    LOG.info("Merging datasets")
+    LOG.info("Merging datasets...")
     merged_dataset = concatenate_datasets(datasets)
 
     if cfg.shuffle_merged_datasets:
-        LOG.debug("Shuffle merged datasets")
+        LOG.debug("Shuffling merged datasets...")
         merged_dataset = merged_dataset.shuffle(seed=cfg.seed)
     else:
-        LOG.debug("NOT shuffling merged datasets")
+        LOG.debug("Not shuffling merged datasets.")
 
     return merged_dataset