diff --git a/mindformers/dataset/blended_datasets/Makefile b/mindformers/dataset/blended_datasets/Makefile new file mode 100644 index 0000000000000000000000000000000000000000..8f9db7686696fbea6c94b998db4b40ef426c748d --- /dev/null +++ b/mindformers/dataset/blended_datasets/Makefile @@ -0,0 +1,9 @@ +CXXFLAGS += -O3 -Wall -shared -std=c++11 -fPIC -fdiagnostics-color +CPPFLAGS += $(shell python3 -m pybind11 --includes) +LIBNAME = helpers +LIBEXT = $(shell python3-config --extension-suffix) + +default: $(LIBNAME)$(LIBEXT) + +%$(LIBEXT): %.cpp + $(CXX) $(CXXFLAGS) $(CPPFLAGS) $< -o $@ diff --git a/mindformers/dataset/blended_datasets/blended_dataset.py b/mindformers/dataset/blended_datasets/blended_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..6cee62c0197ad1ea1e1f451405009e151ab132cc --- /dev/null +++ b/mindformers/dataset/blended_datasets/blended_dataset.py @@ -0,0 +1,186 @@ +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + +import hashlib +import json +import os +import time +from collections import OrderedDict +from typing import Dict, List, Optional, Tuple, Union + +import numpy + +from mindformers.dataset.blended_datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig +from mindformers.dataset.blended_datasets.megatron_dataset import MegatronDataset +from mindformers.dataset.blended_datasets.utils import normalize +from mindformers.tools.logger import logger +from mindformers.tools.utils import get_rank_info + + +_VERBOSE = False + + +class BlendedDataset(): + """Conjugating class for a set of MegatronDataset instances + + Args: + datasets (List[MegatronDataset]): The MegatronDataset instances to blend + + weights (List[Union[int, float]]): The weights that determine the dataset blend ratios + + size (Optional[int]): The number of samples to draw from the blend. If None, for each dataset index idx draw exactly weights[idx] samples from datasets[idx]. + + config (BlendedMegatronDatasetConfig): The config + + Raises: + RuntimeError: When the dataset has fewer or more samples than 'size' post-initialization + """ + + def __init__( + self, + datasets: List[MegatronDataset], + weights: List[Union[int, float]], + size: Optional[int], + config: BlendedMegatronDatasetConfig, + ) -> None: + assert len(datasets) == len(weights) + assert len(datasets) < 32767 + assert all(map(lambda _: type(_) == type(datasets[0]), datasets)) + assert all(map(lambda _: _.index_split == datasets[0].index_split, datasets)) + assert all(map(lambda _: _ > 0, weights)) + assert all(map(lambda _: type(_) == type(weights[0]), weights)) + if size is None and isinstance(weights[0], float): + assert all(map(lambda _: _ == int(_), weights)) + + # Alert user to unnecessary blending + if len(datasets) == 1: + logger.warning(f"Building a BlendedDataset for a single MegatronDataset") + + if size is not None: + weights = normalize(weights) + + self.datasets = datasets + self.split = self.datasets[0].index_split + self.weights = weights + self.size = size + self.config = config + + unique_identifiers = OrderedDict() + unique_identifiers["class"] = type(self).__name__ + unique_identifiers["datasets"] = [dataset.unique_identifiers for dataset in self.datasets] + unique_identifiers["split"] = self.split.name + unique_identifiers["weights"] = self.weights + unique_identifiers["size"] = self.size + + self.unique_description = json.dumps( + unique_identifiers, indent=4, default=lambda obj: obj.unique_identifiers + ) + self.unique_description_hash = hashlib.md5( + self.unique_description.encode("utf-8") + ).hexdigest() + + self.built_anew_on_cache_miss = False + + self.dataset_index, self.dataset_sample_index = self._build_indices() + + logger.info(f"build dataset indices, self.dataset_index.shape = {self.dataset_index.shape}, self.dataset_sample_index.shape = {self.dataset_sample_index.shape}") + + def __len__(self) -> int: + return self.dataset_index.shape[0] + + def __getitem__(self, idx: int) -> Dict[str, Union[int, numpy.ndarray]]: + dataset_id = self.dataset_index[idx] + dataset_sample_id = self.dataset_sample_index[idx] + return self.datasets[dataset_id][dataset_sample_id] + + def _build_indices(self) -> Tuple[numpy.ndarray, numpy.ndarray]: + """Build and optionally cache the dataset index and the dataset sample index + + The dataset index is a 1-D mapping which determines the dataset to query. The dataset + sample index is a 1-D mapping which determines the sample to request from the queried + dataset. + + Returns: + Tuple[numpy.ndarray, numpy.ndarray]: The dataset index and the dataset sample index + """ + path_to_cache = self.config.path_to_cache + + if path_to_cache: + get_path_to = lambda suffix: os.path.join( + path_to_cache, + f"{self.unique_description_hash}-{type(self).__name__}-{self.split.name}-{suffix}", + ) + path_to_description = get_path_to("description.txt") + path_to_dataset_index = get_path_to("dataset_index.npy") + path_to_dataset_sample_index = get_path_to("dataset_sample_index.npy") + cache_hit = all( + map( + os.path.isfile, + [path_to_description, path_to_dataset_index, path_to_dataset_sample_index], + ) + ) + else: + cache_hit = False + + # if not path_to_cache or (not cache_hit and torch.distributed.get_rank() == 0): + rank_id, _ = get_rank_info() + if not path_to_cache or (not cache_hit and rank_id == 0): + logger.info(f"Build and save the {type(self).__name__} indices") + self.built_anew_on_cache_miss = True + + # Build the dataset and dataset sample indexes + logger.info(f"\tBuild and save the dataset and dataset sample indexes") + t_beg = time.time() + from mindformers.dataset.blended_datasets import helpers + + if self.size is not None: + dataset_index = numpy.zeros(self.size, dtype=numpy.int16) + dataset_sample_index = numpy.zeros(self.size, dtype=numpy.int64) + helpers.build_blending_indices( + dataset_index, + dataset_sample_index, + self.weights, + len(self.datasets), + self.size, + _VERBOSE, + ) + else: + size = sum(self.weights) + dataset_index = numpy.zeros(size, dtype=numpy.int16) + dataset_sample_index = numpy.zeros(size, dtype=numpy.int64) + helpers.build_exhaustive_blending_indices( + dataset_index, dataset_sample_index, self.weights, len(self.datasets) + ) + + if path_to_cache: + os.makedirs(path_to_cache, exist_ok=True) + # Write the description + with open(path_to_description, "wt") as writer: + writer.write(self.unique_description) + # Save the indexes + numpy.save(path_to_dataset_index, dataset_index, allow_pickle=True) + numpy.save(path_to_dataset_sample_index, dataset_sample_index, allow_pickle=True) + else: + logger.info(f"Unable to save the {type(self).__name__} indexes because path_to_cache is None") + + t_end = time.time() + logger.info(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + return dataset_index, dataset_sample_index + + logger.info(f"Load the {type(self).__name__} indices") + + logger.info(f"\tLoad the dataset index from {path_to_dataset_index}") + t_beg = time.time() + dataset_index = numpy.load(path_to_dataset_index, allow_pickle=True, mmap_mode='r') + t_end = time.time() + logger.info(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + logger.info(f"\tLoad the dataset sample index from {path_to_dataset_sample_index}") + t_beg = time.time() + dataset_sample_index = numpy.load( + path_to_dataset_sample_index, allow_pickle=True, mmap_mode='r' + ) + t_end = time.time() + logger.info(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + return dataset_index, dataset_sample_index diff --git a/mindformers/dataset/blended_datasets/blended_megatron_dataset_builder.py b/mindformers/dataset/blended_datasets/blended_megatron_dataset_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..a22a98cae23528442a171ba1014de384c69aba7a --- /dev/null +++ b/mindformers/dataset/blended_datasets/blended_megatron_dataset_builder.py @@ -0,0 +1,505 @@ +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + +import math +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, Iterable, List, Optional, Type, Union + +import numpy + +from mindspore.communication.comm_func import barrier + +from mindformers.dataset.blended_datasets.blended_dataset import BlendedDataset +from mindformers.dataset.blended_datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig +from mindformers.dataset.blended_datasets.megatron_dataset import LowLevelDataset, MegatronDataset +from mindformers.dataset.blended_datasets.utils import Split, normalize +from mindformers.tools.logger import logger +from mindformers.tools.utils import get_real_group_size, get_real_rank + +MidLevelDataset = MegatronDataset + +TopLevelDataset = Union[BlendedDataset, MidLevelDataset] + +DistributedDataset = Union[TopLevelDataset, MidLevelDataset, LowLevelDataset] + + +class BlendedMegatronDatasetBuilder(object): + """Builder class for the BlendedDataset and MegatronDataset classes + + Args: + cls (Type[MegatronDataset]): The class to instantiate, must inherit from MegatronDataset + + sizes (List[Optional[int]]): The minimum total number of samples to draw, or None, per split + + is_built_on_rank (Callable): A callable which returns True if the dataset should be built on the current rank and False otherwise. It should be Megatron Core parallelism aware i.e. global rank, local group rank, and virtual rank may inform its return value. + + config (BlendedMegatronDatasetConfig): The config object which informs dataset creation + """ + + def __init__( + self, + cls: Type[MidLevelDataset], + sizes: List[int], + is_built_on_rank: Callable, + config: BlendedMegatronDatasetConfig, + ): + self.cls = cls + self.sizes = sizes + self.is_built_on_rank = is_built_on_rank + self.config = config + + logger.info( + f"Building dataset splits with cls={cls}, sizes={self.sizes}, and config={self.config}" + ) + + if not self.config.mock: + for split in Split: + size_is_none = self.sizes[split.value] is None + if self.config.blend_per_split is None: + weights_are_none = self.config.blend[1] is None + else: + if self.config.blend_per_split[split.value] is None: + continue + weights_are_none = self.config.blend_per_split[split.value][1] is None + if size_is_none: + assert ( + weights_are_none + ), f"size_is_none => weights_are_none fails for {split.name} split" + + def build(self) -> List[Optional[TopLevelDataset]]: + """Build all dataset splits according to the provided blend(s) + + This method is distributed-aware and must be called on all ranks. + + The dataset splits returned can vary according to the config. Supply config.blend and + config.split to build BlendedDataset and/or MegatronDataset splits from the same + distribution. Supply config.blend_per_split to build BlendedDataset and/or MegatronDataset + splits from separate distributions. In either case, for each split, handle the following + cases: + + (1) The split is None + - do nothing + + (2) The split has one contributing dataset, and... + + (a) 'size' is not None + - Build a mid-level dataset with low-level dataset sampling in proportion to the size + + (b) 'size' is None + - Build mid-level datasets with no excess low-level dataset sampling + + (3) The split has multiple contributing datasets, and... + + (a) 'weights' is not None and 'size' is not None + - Build mid-level datasets with low-level dataset sampling in proportion to their weights and the size + - Build a top-level dataset of length marginally greater than 'size' with mid-level dataset sampling in proportion to their weights and the size + + (b) 'weights' is not None and 'size' is None + - Error + + (c) 'weights' is None and 'size' is not None + - Build mid-level datasets with no excess low-level dataset sampling + - Build a top-level dataset of length 'size' with mid-level dataset sampling in proportion to their lengths and the size + + - The 'size' of the top-level dataset is capped at the sum of the mid-level dataset lengths + + (d) 'weights' is None and 'size' is None + - Build mid-level datasets with no excess low-level dataset sampling + - Build a top-level dataset with no excess mid-level dataset sampling + + Returns: + List[Optional[TopLevelDataset]]: A list containing a dataset instance (or None) per split + """ + datasets = self._build_blended_dataset_splits() + + for dataset in datasets: + if dataset is not None and len(dataset) > 0: + if isinstance(dataset, BlendedDataset): + if dataset.built_anew_on_cache_miss or any( + x.built_anew_on_cache_miss for x in dataset.datasets + ): + logger.info( + f"Verifying NumPy indices for {type(dataset).__name__} {dataset.split.name} split" + ) + else: + logger.info( + f"NumPy indices for {type(dataset).__name__} {dataset.split.name} split are fully cached, skipping verification" + ) + continue + # Check blend size + assert dataset.size is None or dataset.size == dataset.dataset_index.shape[0] + # Check blend access of mid-level datasets + _, sizes = numpy.unique(dataset.dataset_index, return_counts=True) + for i, dataset_and_size in enumerate(zip(dataset.datasets, sizes)): + if len(dataset_and_size[0]) < dataset_and_size[1]: + raise IndexError( + f"{type(dataset).__name__} blend goes out of bounds for {type([dataset_and_size[0]]).__name__} {i} for {dataset.split.name} split" + ) + + return datasets + + def _build_blended_dataset_splits( + self, + ) -> List[Optional[TopLevelDataset]]: + """Build all dataset splits according to the provided blend(s) + + See the BlendedMegatronDatasetBuilder.build alias for more information. + + Returns: + List[Optional[TopLevelDataset]]: A list containing a dataset instance (or None) per split + """ + ## + # Return fake "mock" datasets + ## + if self.config.mock: + split = self.config.split_matrix + try: + return self._build_megatron_dataset_splits(None, split, self.sizes) + except Exception as error: + raise Exception( + f"{self.cls.__name__} failed to build as a mock data generator" + ) from error + + ## + # All splits come from the same distribution + ## + elif self.config.blend: + prefixes, weights = self.config.blend + if weights is not None: + weights = normalize(weights) + + split = self.config.split_matrix + + # Blend consists of a single prefix + if len(prefixes) == 1 and weights is None: + return self._build_megatron_dataset_splits(prefixes[0], split, self.sizes) + + # Build the mid-level datasets + if weights is None: + sizes_per_dataset = [[None for split in Split] for prefix in prefixes] + else: + sizes_per_dataset = _get_size_per_split_per_dataset(weights, self.sizes) + + # build each dataset in parallel + megatron_datasets = self._build_megatron_datasets_parallel( + prefixes, split, sizes_per_dataset + ) + + # Build the top-level datasets + blended_datasets = [None] * len(Split) + for i in range(len(Split)): + if split[i] is not None: + weights_i = weights + if weights_i is not None and self.sizes[i] is not None: + size_i = sum(list(zip(*sizes_per_dataset))[i]) + elif weights_i is None: + try: + weights_i = [ + len(megatron_dataset) for megatron_dataset in megatron_datasets[i] + ] + except TypeError: + weights_i = [0 for _ in prefixes] + if self.sizes[i] is not None: + size_i = min(self.sizes[i], sum(weights_i)) + else: + size_i = None # => the size will be sum(weights_i) + else: + raise RuntimeError + blended_datasets[i] = self.build_generic_dataset( + BlendedDataset, + self.is_built_on_rank, + True, # synchronize_ranks, default behavior to build on rank-0 first + megatron_datasets[i], + weights_i, + size_i, + self.config, + ) + + return blended_datasets + + ## + # Each split comes from a separate distribution + ## + else: + blended_datasets = [None] * len(Split) + for i in range(len(Split)): + split_spoof = [None] * len(Split) + split_spoof[i] = (0.0, 1.0) + sizes_spoof = [0] * len(Split) + sizes_spoof[i] = self.sizes[i] + + # Blend is provided for the split + blend = self.config.blend_per_split[i] + if blend is not None: + prefixes, weights = blend + if weights is not None: + weights = normalize(weights) + + # Blend consists of a sigle prefix + if len(prefixes) == 1: + blended_datasets[i] = self._build_megatron_dataset_splits( + prefixes[0], split_spoof, sizes_spoof + )[i] + continue + + # Build mid-level datasets + if weights is None: + sizes_per_dataset = [[None for split in Split] for prefix in prefixes] + else: + sizes_per_dataset = _get_size_per_split_per_dataset(weights, sizes_spoof) + + # build each dataset in parallel + megatron_datasets = self._build_megatron_datasets_parallel( + prefixes, split_spoof, sizes_per_dataset + )[i] + + # Build top-level dataset + if weights is not None and self.sizes[i] is not None: + size = list(map(sum, zip(*sizes_per_dataset)))[i] + elif weights is None: + try: + weights = [ + len(megatron_dataset) for megatron_dataset in megatron_datasets + ] + except TypeError: + weights = [0 for _ in prefixes] + if self.sizes[i] is not None: + size = min(self.sizes[i], sum(weights)) + else: + size = None # => the size will be sum(weights) + else: + raise RuntimeError + blended_datasets[i] = self.build_generic_dataset( + BlendedDataset, + self.is_built_on_rank, + True, # synchronize_ranks, default behavior to build on rank-0 first + megatron_datasets, + weights, + size, + self.config, + ) + + return blended_datasets + + def _build_megatron_datasets_parallel( + self, + prefixes: List[str], + split: List[float], + sizes_per_dataset: List[List[int]], + ) -> List[List[Optional[MegatronDataset]]]: + """Build the megatron datasets for a list of prefixes in parallel + + Args: + prefixes (List[str]): The list of prefix strings + + split (List[float]): The dataset split ratios (must sum to 1.00) + + sizes_per_dataset (List[List[int]]): The number of samples to request + per MegatronDataset per spilt + + Returns: + List[List[Optional[MegatronDataset]]]: For each split, have a list of + MegatronDataset per prefix + """ + + # Helper function to wrap the threading logic + def _threading_helper( + megatron_datasets: List[List[Optional[MegatronDataset]]], + num_workers: int, + prefixes: List[str], + split: List[float], + sizes_per_dataset: List[List[int]], + ) -> None: + with ThreadPoolExecutor(max_workers=num_workers) as executor: + all_futures = [] + for i in range(len(prefixes)): + all_futures.append( + executor.submit( + self._build_megatron_dataset_splits, + prefixes[i], + split, + sizes_per_dataset[i], + False, # synchronize_ranks, barrier is called in this function + ) + ) + for future in all_futures: + try: + megatron_datasets_split = future.result() + for j in range(len(megatron_datasets_split)): + megatron_datasets[j].append(megatron_datasets_split[j]) + except Exception as err: + raise err + + megatron_datasets = [[] for _ in range(len(Split))] + num_dataset_builder_threads = self.config.num_dataset_builder_threads + + if get_real_group_size() > 1: + rank = get_real_rank() + # First, build on rank 0 + if rank == 0: + num_workers = num_dataset_builder_threads + if num_workers > 1: + # since only rank 0 is running, scale up the thread count + # but not too much to avoid overloading storage on miss path. + # if user set num_dataset_builder_threads to 1, + # i.e. meant for serial build, do not scale up. + num_workers *= min(2, max(1, get_real_group_size())) + _threading_helper( + megatron_datasets, + num_workers, + prefixes, + split, + sizes_per_dataset, + ) + + barrier() + + # Then, build on other ranks; guaranteed to be data_cache hit + if rank != 0: + _threading_helper( + megatron_datasets, + num_dataset_builder_threads, + prefixes, + split, + sizes_per_dataset, + ) + else: + _threading_helper( + megatron_datasets, num_dataset_builder_threads, prefixes, split, sizes_per_dataset + ) + + return megatron_datasets + + def _build_megatron_dataset_splits( + self, + dataset_path: Optional[str], + split: List[float], + sizes: List[int], + synchronize_ranks: bool = True, + ) -> List[Optional[MidLevelDataset]]: + """Build each MidLevelDataset split from a single LowLevelDataset + + Args: + dataset_path (Optional[str]): The path on disk which defines the underlying LowLevelDataset, or None for mock dataset classes + + split (List[Tuple[float, float]]): The dataset split matrix + + sizes (List[int]): The number of total samples to draw from each split + + synchronize_ranks (bool): Whether to call barrier for rank-0 / barrier / other-ranks behavior. Set to False when we enforce this behavior at higher level. + + Returns: + List[Optional[MidLevelDataset]]: The MidLevelDataset (or None) per split + """ + # Build the low level dataset + low_level_dataset = self.cls.build_low_level_dataset(dataset_path, self.config) + + # Build the split indices for the low level dataset + num_elements = self.cls.numel_low_level_dataset(low_level_dataset) + split_indices = [] + for i, _ in enumerate(Split): + if split[i] is not None: + beg = int(round(split[i][0] * float(num_elements))) + end = int(round(split[i][1] * float(num_elements))) + split_indices.append(numpy.arange(start=beg, stop=end, step=1, dtype=numpy.int32)) + else: + split_indices.append(None) + + # Build the mid level dataset + mid_level_datasets = [] + for i, _split in enumerate(Split): + if split[i] is None: + mid_level_datasets.append(None) + else: + mid_level_datasets.append( + self.build_generic_dataset( + self.cls, + self.is_built_on_rank, + synchronize_ranks, + low_level_dataset, + dataset_path, + split_indices[i], + sizes[i], + _split, + self.config, + ) + ) + + return mid_level_datasets + + @staticmethod + def build_generic_dataset( + cls: Union[Type[DistributedDataset], Callable], + is_built_on_rank: Callable, + synchronize_ranks: bool, + *args: Any, + ) -> Optional[Union[DistributedDataset, Iterable]]: + """Build the DistributedDataset + + Return None if and only if the underlying dataset class is not built on the current rank + and torch.distributed is initialized. + + Args: + cls (Union[Type[DistributedDataset], Callable]): The DistributedDataset class to be built. In special cases, e.g. when we are building the low level dataset for a RawMegatronDataset instance, we can accept a Callable which returns an Iterable. + + synchronize_ranks (bool): Whether to call barrier for rank-0 / barrier / other-ranks behavior. Set to False when we enforce this behavior at higher level. + + args (Tuple[Any]): The positional arguments used to build the provided DistributedDataset class + + Raises: + Exception: When the dataset constructor raises an OSError + + Returns: + Optional[Union[DistributedDataset, Iterable]]: The DistributedDataset instantion, the Iterable instantiation, or None + """ + if get_real_group_size() > 1: + rank = get_real_rank() + + dataset = None + + # First, build on rank 0 + if rank == 0 and is_built_on_rank(): + try: + dataset = cls(*args) + except OSError as err: + log = ( + f"Failed to write dataset materials to the data cache directory. " + + f"Please supply a directory to which you have write access via " + + f"the path_to_cache attribute in BlendedMegatronDatasetConfig and " + + f"retry. Refer to the preserved traceback above for more information." + ) + raise Exception(log) from err + + if synchronize_ranks: + barrier() + + # After, build on other ranks + if rank != 0 and is_built_on_rank(): + dataset = cls(*args) + + return dataset + + return cls(*args) + + +def _get_size_per_split_per_dataset( + normalized_weights: List[float], target_size_per_split: List[int] +) -> List[List[int]]: + """Determine the contribution of the MegatronDataset splits to the BlendedDataset splits + + Args: + normalized_weights (List[float]): e.g. [0.3, 0.7] + + target_size_per_split (List[int]): The number of samples to target for each BlendedDataset split + + Returns: + List[List[int]]: The number of samples to request per MegatronDataset per split + """ + assert numpy.isclose(sum(normalized_weights), 1.0) + + # Use 0.5% target margin to ensure we satiate the request + sizes_per_dataset = [ + [int(math.ceil(target_size * weight * 1.005)) for target_size in target_size_per_split] + for weight in normalized_weights + ] + + return sizes_per_dataset diff --git a/mindformers/dataset/blended_datasets/blended_megatron_dataset_config.py b/mindformers/dataset/blended_datasets/blended_megatron_dataset_config.py new file mode 100644 index 0000000000000000000000000000000000000000..ed684f7bd1ca0c22ddbdcbff20c0d56e585a1aa9 --- /dev/null +++ b/mindformers/dataset/blended_datasets/blended_megatron_dataset_config.py @@ -0,0 +1,175 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +import functools +import re +from dataclasses import dataclass, field +from typing import List, Optional, Tuple, Union + +from mindformers.dataset.blended_datasets.utils import Split, normalize +from mindformers.models.tokenization_utils_base import PreTrainedTokenizerBase +from mindformers.tools.logger import logger + + +@dataclass +class BlendedMegatronDatasetConfig: + """Configuration object for Megatron Core datasets""" + + random_seed: int + """The seed for all RNG during dataset creation.""" + + sequence_length: int + """The sequence length.""" + + blend: Optional[Tuple[List[str], Optional[List[float]]]] = None + """The blend, consisting of a list of dataset prefixes and optionally a list of dataset + weights. For example, [["dataset-path1", "dataset-path2"], [0.3, 0.7]]. When the weights are + None, they are inferred from the lengths of the contributing datasets. Not to be used with + 'blend_per_split'. Defaults to None. + """ + + blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]] = None + """A set of blends, as defined above, one for each split distribution. Not to be used with + 'blend'. Defauls to None. + """ + + split: Optional[str] = None + """The split string, a comma separated weighting for the dataset splits when drawing samples + from a single distribution. Not to be used with 'blend_per_split'. Defaults to None. + """ + + split_matrix: Optional[List[Tuple[float, float]]] = field(init=False, default=None) + """The split matrix consisting of non-overlapping book-ends of each split in order. For more + information, refer to 'convert_split_vector_to_split_matrix'. Created automatically from + 'split'. Not to be passed in to the constructor. + """ + + num_dataset_builder_threads: int = 1 + """The number of threads to use for dataset building.""" + + path_to_cache: Optional[str] = None + """Where all re-useable dataset indices are to be cached.""" + + mmap_bin_files: bool = True + """Whether to mmap the .bin files or use file pointers.""" + + mock: bool = field(init=False, default=False) + """Whether to bypass real data loading and validation in favor of mock data generation. + Created automatically from 'blend' and 'blend_per_split'. Not to be passed in to the + constructor. + """ + + tokenizer: Optional[Union[PreTrainedTokenizerBase]] = None + """The PreTrainedTokenizerBase instance or None. Required for datasets which do online tokenization.""" + + eod: int = -1 + """The end of document token id.""" + + pad: int = -1 + """The pad token id.""" + + def __post_init__(self) -> None: + """Do asserts and set fields post init + """ + if self.blend_per_split is not None and any(self.blend_per_split): + assert self.blend is None, "blend and blend_per_split are incompatible" + assert self.split is None, "split and blend_per_split are incompatible" + assert len(self.blend_per_split) == len( + Split + ), f"blend_per_split must contain {len(Split)} blends" + for split in Split: + if self.blend_per_split[split.value] is None: + logger.info( + f"blend not provided for {split.name} split" + ) + else: + assert self.blend_per_split[split.value][1] is None or len( + self.blend_per_split[split.value][0] + ) == len( + self.blend_per_split[split.value][1] + ), "blend per split prefixes and weights must be equal in number" + else: + if self.blend is not None: + assert self.blend[1] is None or len(self.blend[0]) == len( + self.blend[1] + ), "blend prefixes and weights must be equal in number" + assert self.split is not None, "split must be provided when blend is not None" + else: + self.mock = True + logger.warning( + f"Let mock = True, as both blend and blend_per_split are None" + ) + self.split = "1," + logger.warning( + f"Let split = {self.split}, an arbitrarily even split, as mock is True" + ) + split_vector = parse_and_normalize_split(self.split) + self.split_matrix = convert_split_vector_to_split_matrix(split_vector) + logger.warning(f"Let split_matrix = {self.split_matrix}") + + +def parse_and_normalize_split(split: str) -> List[float]: + """Parse the dataset split ratios from a string + + Args: + split (str): The train valid test split string e.g. "99,1,0" + + Returns: + List[float]: The trian valid test split ratios e.g. [0.99, 0.01, 0.0] + """ + split = list(map(float, re.findall(r"[.0-9]+", split))) + split = split + [0.0 for _ in range(len(Split) - len(split))] + + print(len(split)) + print(len(Split)) + + assert len(split) == len(Split) + assert all(map(lambda _: _ >= 0.0, split)) + + split = normalize(split) + + return split + + +def convert_split_vector_to_split_matrix( + vector_a: List[float], vector_b: Optional[List[float]] = None +) -> List[Optional[Tuple[float, float]]]: + """Build the split matrix from one or optionally two contributing split vectors. + + Ex. a standard conversion: + + [0.99, 0.01, 0.0] -> [(0, 0.99), (0.99, 1.0), None] + + Ex. a conversion for Retro when Retro pretraining uses a [0.99, 0.01, 0.0] split and Retro + preprocessing used a [0.98, 0.02, 0.0] split: + + [0.99, 0.01, 0.0], [0.98, 0.02, 0.0] -> [(0, 0.98), (0.99, 1.0), None] + + Args: + vector_a (List[float]): The primary split vector + + vector_b (Optional[List[float]]): An optional secondary split vector which constrains the primary split vector. Defaults to None. + + Returns: + List[Tuple[float, float]]: The split matrix consisting of book-ends of each split in order + """ + if vector_b is None: + vector_b = vector_a + + # [.900, .090, .010] -> [0.00, .900, .990, 100] + expansion_a = functools.reduce(lambda a, b: a + [a[len(a) - 1] + b], [[0], *vector_a]) + expansion_b = functools.reduce(lambda a, b: a + [a[len(a) - 1] + b], [[0], *vector_b]) + + # [0.00, .900, .990, 100.0] -> [(0.00, .900), (.900, .990), (.990, 100)] + bookends_a = list(zip(expansion_a[:-1], expansion_a[1:])) + bookends_b = list(zip(expansion_b[:-1], expansion_b[1:])) + + # gather per-split overlap or None + matrix = [] + for bookend_a, bookend_b in zip(bookends_a, bookends_b): + if min(bookend_a[1], bookend_b[1]) <= max(bookend_a[0], bookend_b[0]): + overlap = None + else: + overlap = (max(bookend_a[0], bookend_b[0]), min(bookend_a[1], bookend_b[1])) + matrix.append(overlap) + + return matrix diff --git a/mindformers/dataset/blended_datasets/gpt_dataset.py b/mindformers/dataset/blended_datasets/gpt_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..83a1b68285e84420148a5c0eb5aad77299fd0c3b --- /dev/null +++ b/mindformers/dataset/blended_datasets/gpt_dataset.py @@ -0,0 +1,771 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +import os +import time +from dataclasses import dataclass +from typing import Dict, Optional, Tuple + +import numpy + +from mindformers.dataset.blended_datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig +from mindformers.dataset.blended_datasets.indexed_dataset import IndexedDataset +from mindformers.dataset.blended_datasets.megatron_dataset import MegatronDataset +from mindformers.dataset.blended_datasets.utils import Split +from mindformers.dataset.blended_datasets.utils_s3 import S3Config, is_s3_path +from mindformers.models.tokenization_utils_base import PreTrainedTokenizerBase +from mindformers.tools.logger import logger +from mindformers.tools.utils import get_rank_info + + +_PAD_TOKEN_ID = -1 + + +@dataclass +class GPTDatasetConfig(BlendedMegatronDatasetConfig): + """Configuration object for Megatron Core GPT datasets""" + + reset_position_ids: bool = None + """Option to reset the position IDs in the dataset at an interval""" + + reset_attention_mask: bool = None + """Option to reset the attention mask from the dataset""" + + eod_mask_loss: bool = None + """Option to enable the EOD mask loss""" + + create_attention_mask: bool = True + """Option to enable the attention masks generation. Can be disabled if attention kernel + generates masks by itself. + """ + + drop_last_partial_validation_sequence: bool = True + """Option to drop the last partial validation sequence""" + + add_extra_token_to_sequence: bool = True + """Option to draw sequences with one extra token to ensure the sample input tokens and sample + output tokens are both of the desired sequence length + """ + + s3_cache_path: str = None + """Path for caching indices for s3 dataloading.""" + + def __post_init__(self) -> None: + """Do asserts and set fields post init""" + super().__post_init__() + + # assert self.tokenizer is not None + + assert self.reset_position_ids is not None + assert self.reset_attention_mask is not None + assert self.eod_mask_loss is not None + + +class GPTDataset(MegatronDataset): + """The base GPT dataset + + Args: + indexed_dataset (IndexedDataset): The IndexedDataset around which to build the GPTDataset + + dataset_path (Optional[str]): The real path on disk to the dataset, for bookkeeping + + indexed_indices (numpy.ndarray): The set of the documents indices to expose + + num_samples (Optional[int]): The number of samples to draw from the indexed dataset. When None, build as many samples as correspond to one epoch. + + index_split (Split): The indexed_indices Split + + config (GPTDatasetConfig): The config + """ + + def __init__( + self, + indexed_dataset: IndexedDataset, + dataset_path: Optional[str], + indexed_indices: numpy.ndarray, + num_samples: Optional[int], + index_split: Split, + config: GPTDatasetConfig, + ) -> None: + super().__init__( + indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config + ) + self.masks_and_position_ids_are_cacheable = not any( + [ + self.config.reset_position_ids, + self.config.reset_attention_mask, + self.config.eod_mask_loss, + ] + ) + self.masks_and_position_ids_are_cached = False + self.cached_attention_mask = None + self.cached_loss_mask = None + self.cached_position_ids = None + + try: + self._pad_token_id = self.config.tokenizer.pad + except: + self._pad_token_id = self.config.pad + + try: + self._eod_token_id = self.config.tokenizer.eod + except: + self._eod_token_id = self.config.eod + + ( + self.document_index, + self.sample_index, + self.shuffle_index, + ) = self._build_document_sample_shuffle_indices() + + @staticmethod + def numel_low_level_dataset(low_level_dataset: IndexedDataset) -> int: + """Abstract method implementation + + For GPT, the underlying IndexedDataset should be split by sequence, as opposed to, say, + BERT, which should be split by document + + Args: + low_level_dataset (IndexedDataset): The underlying IndexedDataset + + Returns: + int: The number of unique elements in the underlying IndexedDataset + """ + return low_level_dataset.sequence_lengths.shape[0] + + @staticmethod + def build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> IndexedDataset: + """Abstract method implementation + + Args: + dataset_path (str): The real path prefix to the IndexedDataset .bin and .idx files + + config (GPTDatasetConfig): The config + + Returns: + IndexedDataset: The underlying IndexedDataset + """ + if is_s3_path(dataset_path): + return IndexedDataset( + dataset_path, + multimodal=False, + mmap=config.mmap_bin_files, + s3_config=S3Config(path_to_idx_cache=config.s3_cache_path), + ) + return IndexedDataset(dataset_path, multimodal=False, mmap=config.mmap_bin_files) + + def __len__(self) -> int: + """Abstract method implementation + + Returns: + int: The length of the dataset + """ + return self.sample_index.shape[0] - 1 + + def __getitem__(self, idx: Optional[int]) -> Dict[str, numpy.ndarray]: + """Abstract method implementation + + Args: + idx (Optioal[int]): The index into the dataset + + Returns: + Dict[str, numpy.ndarray]: The sample information wrapped in a dictionary + """ + if idx is None: + # Batch padding sequence so the index does not matter + text, _ = self._query_document_sample_shuffle_indices(0) + else: + text, _ = self._query_document_sample_shuffle_indices(idx) + + if self.config.add_extra_token_to_sequence: + tokens = text[:-1] + labels = text[1:] + else: + tokens = text + # labels = torch.roll(text, shifts=-1, dims=0) + labels = numpy.roll(text, shift=-1, axis=0) + labels[-1] = self._pad_token_id + + if ( + not self.masks_and_position_ids_are_cacheable + or not self.masks_and_position_ids_are_cached + ): + attention_mask, loss_mask, position_ids = _get_ltor_masks_and_position_ids( + tokens, + self._eod_token_id, + self.config.reset_position_ids, + self.config.reset_attention_mask, + self.config.eod_mask_loss, + self.config.create_attention_mask, + ) + if self.masks_and_position_ids_are_cacheable: + self.cached_attention_mask = attention_mask + self.cached_loss_mask = loss_mask + self.cached_position_ids = position_ids + self.masks_and_position_ids_are_cached = True + else: + attention_mask = self.cached_attention_mask + loss_mask = self.cached_loss_mask + position_ids = self.cached_position_ids + + # For padded sequences, mask the loss + loss_mask[labels == self._pad_token_id] = 0.0 + + # For padded sequences, ensure the embedding layer can map the token ID + tokens[tokens == self._pad_token_id] = 0 + labels[labels == self._pad_token_id] = 0 + + # Batch padding sequence so we mask the loss + if idx is None: + # loss_mask = torch.zeros_like(loss_mask) + loss_mask = numpy.zeros_like(loss_mask) + + if self.config.create_attention_mask: + # return { + # "tokens": tokens, + # "labels": labels, + # "attention_mask": attention_mask, + # "loss_mask": loss_mask, + # "position_ids": position_ids, + # } + return tokens, labels, loss_mask, position_ids, attention_mask + else: + # return { + # "tokens": tokens, + # "labels": labels, + # "loss_mask": loss_mask, + # "position_ids": position_ids, + # } + return tokens, labels, loss_mask, position_ids + + def _query_document_sample_shuffle_indices( + self, idx: int + ) -> Tuple[numpy.ndarray, numpy.ndarray]: + """Get the text (token ids) and document ids for a given index + + Args: + idx (int): The index into the dataset + + Returns: + Tuple[numpy.ndarray, numpy.ndarray]: The text ids and document ids + """ + # Do the shuffle mapping + idx = self.shuffle_index[idx] + + # Get the beginning and end documents and offsets + doc_index_beg, doc_index_beg_offset = self.sample_index[idx] + doc_index_end, doc_index_end_offset = self.sample_index[idx + 1] + + document_ids = [] + sample_parts = [] + + # Sample spans a single document + if doc_index_beg == doc_index_end: + # Add the document id + document_ids.append(self.document_index[doc_index_beg]) + + # Add the entire sample + sample_parts.append( + self.dataset.get( + self.document_index[doc_index_beg], + offset=doc_index_beg_offset, + length=doc_index_end_offset + - doc_index_beg_offset + + self.config.add_extra_token_to_sequence, + ) + ) + + # Sample spans multiple documents + else: + for i in range(doc_index_beg, doc_index_end + 1): + # Add the document id + document_ids.append(self.document_index[i]) + + # Add the sample part + offset = 0 if i > doc_index_beg else doc_index_beg_offset + length = ( + None + if i < doc_index_end + else doc_index_end_offset + self.config.add_extra_token_to_sequence + ) + sample_parts.append( + self.dataset.get(self.document_index[i], offset=offset, length=length) + ) + assert len(document_ids) == len( + sample_parts + ), f"len(document_ids) ({len(document_ids)}) != len(sample_parts) ({len(sample_parts)})" + + length = sum(map(len, sample_parts)) + + # Pad the sample if necessary + if length < (self.config.sequence_length + self.config.add_extra_token_to_sequence): + sample_parts.append( + [self._pad_token_id] + * (self.config.sequence_length + self.config.add_extra_token_to_sequence - length) + ) + + # import pdb;pdb.set_trace() + + return ( + numpy.concatenate(sample_parts, dtype=numpy.int64), + numpy.array(document_ids, dtype=numpy.int64), + ) + + def _build_document_sample_shuffle_indices( + self, + ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: + """Build the document index, the sample index, and the shuffle index + + The document index: + -- 1-D + -- An ordered array of document ids + + The sample index: + -- 2-D + -- The document indices and offsets which mark the start of every sample + + The shuffle index: + -- 1-D + -- A random permutation of index range of the sample index + + Returns: + Tuple[numpy.ndarray, numpy.ndarray]: The document index, the sample index, and the shuffle index + """ + path_to_cache = self.config.path_to_cache + if path_to_cache is None and not self.config.mock: + path_to_cache = os.path.join( + self.dataset.path_prefix, "cache", f"{type(self).__name__}_indices" + ) + + if path_to_cache: + get_path_to = lambda suffix: os.path.join( + path_to_cache, + f"{self.unique_description_hash}-{type(self).__name__}-{self.index_split.name}-{suffix}", + ) + path_to_description = get_path_to("description.txt") + path_to_document_index = get_path_to("document_index.npy") + path_to_sample_index = get_path_to("sample_index.npy") + path_to_shuffle_index = get_path_to("shuffle_index.npy") + cache_hit = all( + map( + os.path.isfile, + [ + path_to_description, + path_to_document_index, + path_to_sample_index, + path_to_shuffle_index, + ], + ) + ) + else: + cache_hit = False + """ + if not path_to_cache or ( + not cache_hit + and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0) + ): + """ + if not path_to_cache or ( + not cache_hit + ): + + logger.warning(f"Build and save the {type(self).__name__} {self.index_split.name} indices") + + self.built_anew_on_cache_miss = True + t_beg = time.time() + + sequence_length = self.config.sequence_length + num_tokens_per_epoch = self._get_num_tokens_per_epoch() + num_epochs = self._get_num_epochs(num_tokens_per_epoch) + + if num_epochs == 1: + separate_final_epoch = False + else: + # Get the number of samples for the last epoch + num_samples_sans_final_epoch = ( + (num_epochs - 1) * num_tokens_per_epoch + - self.config.add_extra_token_to_sequence + ) // sequence_length + num_samples_from_final_epoch = self.num_samples - num_samples_sans_final_epoch + num_samples_per_epoch = ( + num_tokens_per_epoch - self.config.add_extra_token_to_sequence + ) // sequence_length + + # num_samples_from_final_epoch should be non-negative + assert num_samples_from_final_epoch >= 0 + + # num_samples_from_final_epoch should not exceed max value + assert num_samples_from_final_epoch <= num_samples_per_epoch + 1 + + # Separate the final epoch if it falls below the threshold + threshold = 0.80 + separate_final_epoch = num_samples_from_final_epoch < int( + threshold * num_samples_per_epoch + ) + + logger.debug(f"> num_samples_from_final_epoch: {num_samples_from_final_epoch}") + logger.debug(f"> threshold: {threshold}") + logger.debug(f"> num_samples_from_final_epoch: {num_samples_from_final_epoch}") + + logger.debug(f"> separate_final_epoch: {separate_final_epoch}") + + numpy_random_state = numpy.random.RandomState(self.config.random_seed) + + # Build the document index + document_index = _build_document_index( + self.indices, num_epochs, numpy_random_state, separate_final_epoch + ) + + drop_last_partial_sequence = True + if self.index_split == Split.valid: + drop_last_partial_sequence = self.config.drop_last_partial_validation_sequence + + # Build the sample index + from mindformers.dataset.blended_datasets import helpers + + if self.index_split == Split.valid: + drop_last_partial_sequence = self.config.drop_last_partial_validation_sequence + else: + drop_last_partial_sequence = True + + assert document_index.dtype == numpy.int32 + assert self.dataset.sequence_lengths.dtype == numpy.int32 + if len(document_index) * 2 > len(self.dataset.sequence_lengths): + # Heuristic: if "access density" of sequence_lengths is relatively high, + # force loading the mmap-ed array into memory by taking a copy. + # System performance benefits come from two aspects: + # 1. **sequentially** pre-loading the whole file if we're gonna read a large fraction anyways. + # 2. GIL is held when calling into c++ code; making the c++ func faster improves parallelism. + sequence_lengths_for_cpp = self.dataset.sequence_lengths.copy() + else: + sequence_lengths_for_cpp = self.dataset.sequence_lengths + sample_index = helpers.build_sample_idx( + sequence_lengths_for_cpp, + document_index, + sequence_length, + num_epochs, + num_tokens_per_epoch, + drop_last_partial_sequence, + self.config.add_extra_token_to_sequence, + ) + + # Build the shuffle index + if separate_final_epoch: + shuffle_index = _build_shuffle_index( + num_samples_sans_final_epoch, sample_index.shape[0] - 1, numpy_random_state + ) + else: + shuffle_index = _build_shuffle_index( + sample_index.shape[0] - 1, sample_index.shape[0] - 1, numpy_random_state + ) + + # only rank 0 save indices file. + rank, _ = get_rank_info() + if path_to_cache and rank == 0: + os.makedirs(path_to_cache, exist_ok=True) + # Write the description + with open(path_to_description, "wt") as writer: + writer.write(self.unique_description) + numpy.save(path_to_document_index, document_index, allow_pickle=True) + numpy.save(path_to_sample_index, sample_index, allow_pickle=True) + numpy.save(path_to_shuffle_index, shuffle_index, allow_pickle=True) + else: + logger.warning( + f"Unable to save the {type(self).__name__} indexes because path_to_cache is None", + ) + + t_end = time.time() + logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + logger.info(f"> total number of samples: {sample_index.shape[0] - 1}") + logger.info(f"> total number of epochs: {num_epochs}") + + return document_index, sample_index, shuffle_index + + logger.info( + f"Load the {type(self).__name__} {self.index_split.name} indices" + ) + + logger.info( + f"\tLoad the document index from {os.path.basename(path_to_document_index)}", + ) + t_beg = time.time() + document_index = numpy.load(path_to_document_index, allow_pickle=True, mmap_mode='r') + t_end = time.time() + logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + logger.info( + f"\tLoad the sample index from {os.path.basename(path_to_sample_index)}" + ) + t_beg = time.time() + sample_index = numpy.load(path_to_sample_index, allow_pickle=True, mmap_mode='r') + t_end = time.time() + logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + logger.info( + f"\tLoad the shuffle index from {os.path.basename(path_to_shuffle_index)}", + ) + t_beg = time.time() + shuffle_index = numpy.load(path_to_shuffle_index, allow_pickle=True, mmap_mode='r') + t_end = time.time() + logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + logger.info( + f"> total number of samples: {sample_index.shape[0] - 1}" + ) + + return document_index, sample_index, shuffle_index + + def _get_num_tokens_per_epoch(self) -> int: + """Calculate the number of tokens in a single epoch + + Returns: + int: The number of tokens in a single epoch + """ + return int(numpy.sum(self.dataset.sequence_lengths[self.indices])) + + def _get_num_epochs(self, num_tokens_per_epoch: int) -> int: + """Calculate the number of epochs + + Args: + num_tokens_per_epoch (int): The number of tokens in a single epoch + + Returns: + int: The number of epochs + """ + num_epochs = 1 + num_tokens = num_tokens_per_epoch + if self.num_samples is None: + return num_epochs + else: + num_tokens_requested = ( + self.num_samples * self.config.sequence_length + ) + self.config.add_extra_token_to_sequence + while num_tokens < num_tokens_requested: + num_epochs += 1 + num_tokens += num_tokens_per_epoch + return num_epochs + + +def _build_document_index( + documents: numpy.ndarray, + num_epochs: int, + numpy_random_state: numpy.random.RandomState, + separate_final_epoch: bool, +) -> numpy.ndarray: + """Build an array with length = num epochs * num documents + + Args: + documents (numpy.ndarray): the subset of exposed document indices + + num_epochs (int): The number of epochs + + numpy_random_state (numpy.random.RandomState): The NumPy random state + + separate_final_epoch (bool): Whether to exclude the last epoch from the global shuffle + + Returns: + numpy.ndarray: The document index + """ + if not separate_final_epoch or num_epochs == 1: + document_index = numpy.mgrid[0:num_epochs, 0 : len(documents)][1] + document_index[:] = documents + document_index = document_index.reshape(-1) + document_index = document_index.astype(numpy.int32) + numpy_random_state.shuffle(document_index) + return document_index + + doc_idx_first = _build_document_index(documents, num_epochs - 1, numpy_random_state, False) + doc_idx_last = _build_document_index(documents, 1, numpy_random_state, False) + return numpy.concatenate((doc_idx_first, doc_idx_last)) + + +def _build_shuffle_index( + num_samples: int, total_size: int, numpy_random_state: numpy.random.RandomState +) -> numpy.ndarray: + """Build the range [0, size) and shuffle + + Args: + num_samples (int): The size of the first shuffle range [0, num_samples) + + total_size (int): The size of the entire index. If larger than 'num_samples', it defines the second shuffle range [num_samples, total_size) + + numpy_random_state (numpy.random.RandomState): The NumPy random state + + Returns: + numpy.ndarray: The shuffle index + """ + dtype_ = numpy.uint32 + if total_size >= (numpy.iinfo(numpy.uint32).max - 1): + dtype_ = numpy.int64 + + shuffle_idx_first = numpy.arange(start=0, stop=num_samples, step=1, dtype=dtype_) + numpy_random_state.shuffle(shuffle_idx_first) + if num_samples == total_size: + return shuffle_idx_first + + shuffle_idx_last = numpy.arange(start=num_samples, stop=total_size, step=1, dtype=dtype_) + numpy_random_state.shuffle(shuffle_idx_last) + + return numpy.concatenate((shuffle_idx_first, shuffle_idx_last)) + + +def _get_ltor_masks_and_position_ids( + # data: torch.Tensor, + data: numpy.ndarray, + eod_token: int, + reset_position_ids: bool, + reset_attention_mask: bool, + eod_mask_loss: bool, + create_attention_mask: bool, +): + """Build masks and position id for left to right model. + + Args: + data (numpy.ndarray): The data tenor that holds the tokens from the dataset + + eod_token (int): ID of the token to that is considered the EOD + + reset_position_ids (bool): Switch to reset the document position ID's + + reset_attention_mask (bool): Switch to reset the attention mask + + eod_mask_loss (bool): Switch to enable the EOD mask loss + + create_attention_mask (bool): Switch to enable the attention masks generation. Can be disabled if attention kernel generates masks by itself. + + Returns: + numpy.ndarray: Attention mask needed to be used for Attention + + numpy.ndarray: The mask used for loss value during training + + numpy.ndarray: The position ID's of the token + """ + seq_length = numpy.size(data) + + if create_attention_mask: + attention_mask = numpy.expand_dims(numpy.tril(numpy.ones((seq_length, seq_length))), axis=0) + else: + attention_mask = None + + # Loss mask. + loss_mask = numpy.ones(seq_length, dtype=numpy.float32) + if eod_mask_loss: + loss_mask[data == eod_token] = 0.0 + + # Position ids. + position_ids = numpy.arange(seq_length, dtype=numpy.float32) + # We need to clone as the ids will be modifed based on batch index. + import copy + if reset_position_ids: + position_ids = copy.deepcopy(position_ids) + + if reset_position_ids or reset_attention_mask: + # Find indices where EOD token is. + eod_index = position_ids[data == eod_token] + # Detach indices from positions if going to modify positions. + + if reset_position_ids: + eod_index = copy.deepcopy(eod_index) + + # Loop through EOD indices: + prev_index = 0 + for j in range(numpy.size(eod_index)): + i = int(eod_index[j]) + # Mask attention loss. + if reset_attention_mask and attention_mask is not None: + attention_mask[0, (i + 1) :, : (i + 1)] = 0 + # Reset positions. + if reset_position_ids: + position_ids[(i + 1) :] -= i + 1 - prev_index + prev_index = i + 1 + + if attention_mask is not None: + # Convert attention mask to binary: + attention_mask = attention_mask < 0.5 + + return attention_mask, loss_mask, position_ids + + +class MockGPTLowLevelDataset: + + seed: int = 0 + size: int = 100000 + max_sequence_length: int = 4096 + + def __init__(self, tokenizer: PreTrainedTokenizerBase) -> None: + self.tokenizer = tokenizer + rng = numpy.random.default_rng(seed=self.seed) + self.sequence_lengths = rng.integers(low=1, high=self.max_sequence_length, size=self.size, dtype=numpy.int32) + + def __len__(self) -> int: + return self.size + + def __getitem__(self, idx: int) -> numpy.number: + length = self.sequence_lengths[idx] + sample = numpy.int64(numpy.concatenate([numpy.arange(length - 1) + 1, [self.tokenizer.eod]])) + return sample + + def get(self, idx: int, offset: int = 0, length: Optional[int] = None) -> numpy.ndarray: + if length is None: + length = self.sequence_lengths[idx] - offset + return self[idx][offset : offset + length] + + +class MockGPTDataset(GPTDataset): + """The mock GPT dataset + + Args: + indexed_dataset (MockGPTLowLevelDataset): The MockGPTLowLevelDataset around which to build the MockGPTDataset + + dataset_path (Optional[str]): This argument is of no consequence for the MockGPTDataset + + indices (numpy.ndarray): The set of the dataset indices to expose + + num_samples (int): The number of samples to draw from the dataset + + index_split (Split): The indices Split + + config (GPTDatasetConfig): The config + """ + + def __init__( + self, + dataset: MockGPTLowLevelDataset, + dataset_path: Optional[str], + indices: numpy.ndarray, + num_samples: int, + index_split: Split, + config: GPTDatasetConfig, + ) -> None: + config.mock = True + assert config.mock + + super().__init__(dataset, dataset_path, indices, num_samples, index_split, config) + + @staticmethod + def numel_low_level_dataset(low_level_dataset: MockGPTLowLevelDataset) -> int: + """Abstract method implementation + + Args: + low_level_dataset (MockGPTLowLevelDataset): The underlying MockGPTLowLevelDataset + + Returns: + int: The number of unique elements in the underlying MockGPTLowLevelDataset + """ + return len(low_level_dataset) + + @staticmethod + def build_low_level_dataset( + dataset_path: Optional[str], config: GPTDatasetConfig + ) -> MockGPTLowLevelDataset: + """Abstract method implementation + + Args: + dataset_path (Optional[str]): This argument is of no consequence for the MockGPTLowLevelDataset + + config (GPTDatasetConfig): The config + + Returns: + MockGPTLowLevelDataset: The underlying MockGPTLowLevelDataset + """ + return MockGPTLowLevelDataset(config.tokenizer) diff --git a/mindformers/dataset/blended_datasets/helpers.cpp b/mindformers/dataset/blended_datasets/helpers.cpp new file mode 100644 index 0000000000000000000000000000000000000000..8a1e317c6896f428997e4afa2eb43118415f9e4d --- /dev/null +++ b/mindformers/dataset/blended_datasets/helpers.cpp @@ -0,0 +1,840 @@ +/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. */ +/* Copyright 2024 Huawei Technologies Co., Ltd*/ +/* Helper methods for fast index mapping builds */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; +using namespace std; + +const int32_t LONG_SENTENCE_LEN = 512; + + +void build_exhaustive_blending_indices(py::array_t &dataset_index, py::array_t &dataset_sample_index, const py::array_t &sizes, const int32_t num_datasets) { + /* + Build blending indices by sampling exactly as many samples from dataset[i] + as is requested by sizes[i] for all i in the range [0, num_datasets). + */ + auto dataset_index_ptr = dataset_index.mutable_unchecked<1>(); + auto dataset_sample_index_ptr = dataset_sample_index.mutable_unchecked<1>(); + auto sizes_ptr = sizes.unchecked<1>(); + + int64_t total_size = 0; + int64_t dataset_sample_counts[num_datasets]; + std::set dataset_unspent_indices; + for (int32_t i = 0; i < num_datasets; ++i) { + total_size += sizes_ptr[i]; + dataset_sample_counts[i] = 0; + dataset_unspent_indices.insert(i); + } + + // still need fractional weights to sample in proportion to sizes + double weights[num_datasets]; + for (int32_t i = 0; i < num_datasets; ++i) { + weights[i] = sizes_ptr[i] / static_cast(total_size); + } + + int64_t index_sample = 0; + while (dataset_unspent_indices.size() > 0) { + double index_sample_double = std::max(static_cast(index_sample), 1.0); + + int64_t error_argmax; + double error_max = std::numeric_limits::lowest(); + + for (int32_t index_dataset : dataset_unspent_indices) { + double error = weights[index_dataset] * index_sample_double - static_cast(dataset_sample_counts[index_dataset]); + if (error > error_max) { + error_argmax = index_dataset; + error_max = error; + } + } + + // Populate the indices. + dataset_index_ptr[index_sample] = static_cast(error_argmax); + dataset_sample_index_ptr[index_sample] = dataset_sample_counts[error_argmax]; + + // Update the total samples. + dataset_sample_counts[error_argmax] += 1; + + if (sizes_ptr[error_argmax] - static_cast(dataset_sample_counts[error_argmax]) == 0) { + dataset_unspent_indices.erase(error_argmax); + } + + index_sample += 1; + } +} + +void build_blending_indices(py::array_t &dataset_index, + py::array_t &dataset_sample_index, + const py::array_t &weights, + const int32_t num_datasets, + const int64_t size, const bool verbose) +{ + /* Given multiple datasets and a weighting array, build samples + such that it follows those wieghts.*/ + + if (verbose) + { + std::cout << "> building indices for blended datasets ..." << std::endl; + } + + // Get the pointer access without the checks. + auto dataset_index_ptr = dataset_index.mutable_unchecked<1>(); + auto dataset_sample_index_ptr = dataset_sample_index.mutable_unchecked<1>(); + auto weights_ptr = weights.unchecked<1>(); + + // Initialize buffer for number of samples used for each dataset. + int64_t current_samples[num_datasets]; + for (int64_t i = 0; i < num_datasets; ++i) + { + current_samples[i] = 0; + } + + // For each sample: + for (int64_t sample_idx = 0; sample_idx < size; ++sample_idx) + { + + // Determine where the max error in sampling is happening. + auto sample_idx_double = std::max(static_cast(sample_idx), 1.0); + int64_t max_error_index = 0; + double max_error = weights_ptr[0] * sample_idx_double - + static_cast(current_samples[0]); + for (int64_t dataset_idx = 1; dataset_idx < num_datasets; ++dataset_idx) + { + double error = weights_ptr[dataset_idx] * sample_idx_double - + static_cast(current_samples[dataset_idx]); + if (error > max_error) + { + max_error = error; + max_error_index = dataset_idx; + } + } + + // Populate the indices. + dataset_index_ptr[sample_idx] = static_cast(max_error_index); + dataset_sample_index_ptr[sample_idx] = current_samples[max_error_index]; + + // Update the total samples. + current_samples[max_error_index] += 1; + } + + // print info + if (verbose) + { + std::cout << " > sample ratios:" << std::endl; + for (int64_t dataset_idx = 0; dataset_idx < num_datasets; ++dataset_idx) + { + auto ratio = static_cast(current_samples[dataset_idx]) / + static_cast(size); + std::cout << " dataset " << dataset_idx << ", input: " << weights_ptr[dataset_idx] << ", achieved: " << ratio << std::endl; + } + } +} + +py::array build_sample_idx(const py::array_t &sizes_, + const py::array_t &doc_idx_, + const int32_t seq_length, + const int32_t num_epochs, + const int64_t tokens_per_epoch, + const bool drop_last_partial_sequence = true, + const int add_extra_token_to_sequence = 1) +{ + /* Sample index (sample_idx) is used for gpt2 like dataset for which + the documents are flattened and the samples are built based on this + 1-D flatten array. It is a 2D array with sizes [number-of-samples + 1, 2] + where [..., 0] contains the index into `doc_idx` and [..., 1] is the + starting offset in that document.*/ + + // Consistency checks. + assert(seq_length > 1); + assert(num_epochs > 0); + assert(tokens_per_epoch > 1); + + // Remove bound checks. + auto sizes = sizes_.unchecked<1>(); + auto doc_idx = doc_idx_.unchecked<1>(); + + // Mapping and it's length (1D). + int64_t num_samples = 0; + if (drop_last_partial_sequence == true) + { + num_samples = (num_epochs * tokens_per_epoch - add_extra_token_to_sequence) / seq_length; + } + else + { + num_samples = ceil(float(num_epochs * tokens_per_epoch - add_extra_token_to_sequence) / seq_length); + } + int64_t *sample_idx = new int64_t[2 * (num_samples + 1)]; + + // Index into sample_idx. + int64_t sample_index = 0; + // Index into doc_idx. + int64_t doc_idx_index = 0; + // Begining offset for each document. + int32_t doc_offset = 0; + // Start with first document and no offset. + sample_idx[2 * sample_index] = doc_idx_index; + sample_idx[2 * sample_index + 1] = doc_offset; + ++sample_index; + + while (sample_index <= num_samples) + { + // Start with a fresh sequence. + int32_t remaining_seq_length = seq_length + add_extra_token_to_sequence; + while (remaining_seq_length != 0) + { + // Get the document length. + auto doc_id = doc_idx[doc_idx_index]; + auto doc_length = sizes[doc_id] - doc_offset; + // And add it to the current sequence. + remaining_seq_length -= doc_length; + // If we have more than a full sequence, adjust offset and set + // remaining length to zero so we return from the while loop. + // Note that -1 here is for the same reason we have -1 in + // `_num_epochs` calculations. + if (remaining_seq_length <= 0) + { + doc_offset += (remaining_seq_length + doc_length - add_extra_token_to_sequence); + remaining_seq_length = 0; + } + else + { + // Otherwise, start from the begining of the next document. + if (doc_idx_index == (doc_idx_.shape(0) - 1)) + { + // If we have reached the end of the documents, break. + assert(sample_index == num_samples); + doc_offset = sizes[doc_idx[doc_idx_index]] - add_extra_token_to_sequence; + break; + } + ++doc_idx_index; + doc_offset = 0; + } + } + // Record the sequence. + sample_idx[2 * sample_index] = doc_idx_index; + sample_idx[2 * sample_index + 1] = doc_offset; + ++sample_index; + } + + // Method to deallocate memory. + py::capsule free_when_done(sample_idx, [](void *mem_) + { + int64_t *mem = reinterpret_cast(mem_); + delete[] mem; }); + + // Return the numpy array. + const auto byte_size = sizeof(int64_t); + return py::array(std::vector{num_samples + 1, 2}, // shape + {2 * byte_size, byte_size}, // C-style contiguous strides + sample_idx, // the data pointer + free_when_done); // numpy array references +} + +inline int32_t get_target_sample_len(const int32_t short_seq_ratio, + const int32_t max_length, + std::mt19937 &rand32_gen) +{ + /* Training sample length. */ + if (short_seq_ratio == 0) + { + return max_length; + } + const auto random_number = rand32_gen(); + if ((random_number % short_seq_ratio) == 0) + { + return 2 + random_number % (max_length - 1); + } + return max_length; +} + +template +py::array build_mapping_impl(const py::array_t &docs_, + const py::array_t &sizes_, + const int32_t num_epochs, + const uint64_t max_num_samples, + const int32_t max_seq_length, + const double short_seq_prob, + const int32_t seed, + const bool verbose, + const int32_t min_num_sent) +{ + /* Build a mapping of (start-index, end-index, sequence-length) where + start and end index are the indices of the sentences in the sample + and sequence-length is the target sequence length. + */ + + // Consistency checks. + assert(num_epochs > 0); + assert(max_seq_length > 1); + assert(short_seq_prob >= 0.0); + assert(short_seq_prob <= 1.0); + assert(seed > 0); + + // Remove bound checks. + auto docs = docs_.unchecked<1>(); + auto sizes = sizes_.unchecked<1>(); + + // For efficiency, convert probability to ratio. Note: rand() generates int. + int32_t short_seq_ratio = 0; + if (short_seq_prob > 0) + { + short_seq_ratio = static_cast(round(1.0 / short_seq_prob)); + } + + if (verbose) + { + const auto sent_start_index = docs[0]; + const auto sent_end_index = docs[docs_.shape(0) - 1]; + const auto num_sentences = sent_end_index - sent_start_index; + cout << " using:" << endl + << std::flush; + cout << " number of documents: " << docs_.shape(0) - 1 << endl + << std::flush; + cout << " sentences range: [" << sent_start_index << ", " << sent_end_index << ")" << endl + << std::flush; + cout << " total number of sentences: " << num_sentences << endl + << std::flush; + cout << " number of epochs: " << num_epochs << endl + << std::flush; + cout << " maximum number of samples: " << max_num_samples << endl + << std::flush; + cout << " maximum sequence length: " << max_seq_length << endl + << std::flush; + cout << " short sequence probability: " << short_seq_prob << endl + << std::flush; + cout << " short sequence ration (1/prob): " << short_seq_ratio << endl + << std::flush; + cout << " seed: " << seed << endl + << std::flush; + } + + // Mapping and it's length (1D). + int64_t num_samples = -1; + DocIdx *maps = NULL; + + // Perform two iterations, in the first iteration get the size + // and allocate memory and in the second iteration populate the map. + bool second = false; + for (int32_t iteration = 0; iteration < 2; ++iteration) + { + + // Set the seed so both iterations produce the same results. + std::mt19937 rand32_gen(seed); + + // Set the flag on second iteration. + second = (iteration == 1); + + // Counters: + uint64_t empty_docs = 0; + uint64_t one_sent_docs = 0; + uint64_t long_sent_docs = 0; + + // Current map index. + uint64_t map_index = 0; + + // For each epoch: + for (int32_t epoch = 0; epoch < num_epochs; ++epoch) + { + if (map_index >= max_num_samples) + { + if (verbose && (!second)) + { + cout << " reached " << max_num_samples << " samples after " + << epoch << " epochs ..." << endl + << std::flush; + } + break; + } + // For each document: + for (int32_t doc = 0; doc < (docs.shape(0) - 1); ++doc) + { + + // Document sentences are in [sent_index_first, sent_index_last) + const auto sent_index_first = docs[doc]; + const auto sent_index_last = docs[doc + 1]; + + // At the begining of the document previous index is the + // start index. + auto prev_start_index = sent_index_first; + + // Remaining documents. + auto num_remain_sent = sent_index_last - sent_index_first; + + // Some bookkeeping + if ((epoch == 0) && (!second)) + { + if (num_remain_sent == 0) + { + ++empty_docs; + } + if (num_remain_sent == 1) + { + ++one_sent_docs; + } + } + + // Detect documents with long sentences. + bool contains_long_sentence = false; + if (num_remain_sent > 1) + { + for (auto sent_index = sent_index_first; + sent_index < sent_index_last; ++sent_index) + { + if (sizes[sent_index] > LONG_SENTENCE_LEN) + { + if ((epoch == 0) && (!second)) + { + ++long_sent_docs; + } + contains_long_sentence = true; + break; + } + } + } + + // If we have more than two sentences. + if ((num_remain_sent >= min_num_sent) && (!contains_long_sentence)) + { + + // Set values. + auto seq_len = int32_t{0}; + auto num_sent = int32_t{0}; + auto target_seq_len = get_target_sample_len(short_seq_ratio, + max_seq_length, + rand32_gen); + + // Loop through sentences. + for (auto sent_index = sent_index_first; + sent_index < sent_index_last; ++sent_index) + { + + // Add the size and number of sentences. + seq_len += sizes[sent_index]; + ++num_sent; + --num_remain_sent; + + // If we have reached the target length. + // and if not only one sentence is left in the document. + // and if we have at least two sentneces. + // and if we have reached end of the document. + if (((seq_len >= target_seq_len) && + (num_remain_sent > 1) && + (num_sent >= min_num_sent)) || + (num_remain_sent == 0)) + { + + // Check for overflow. + if ((3 * map_index + 2) > + std::numeric_limits::max()) + { + cout << "number of samples exceeded maximum " + << "allowed by type int64: " + << std::numeric_limits::max() + << endl; + throw std::overflow_error("Number of samples"); + } + + // Populate the map. + if (second) + { + const auto map_index_0 = 3 * map_index; + maps[map_index_0] = static_cast(prev_start_index); + maps[map_index_0 + 1] = static_cast(sent_index + 1); + maps[map_index_0 + 2] = static_cast(target_seq_len); + } + + // Update indices / counters. + ++map_index; + prev_start_index = sent_index + 1; + target_seq_len = get_target_sample_len(short_seq_ratio, + max_seq_length, + rand32_gen); + seq_len = 0; + num_sent = 0; + } + + } // for (auto sent_index=sent_index_first; ... + } // if (num_remain_sent > 1) { + } // for (int doc=0; doc < num_docs; ++doc) { + } // for (int epoch=0; epoch < num_epochs; ++epoch) { + + if (!second) + { + if (verbose) + { + cout << " number of empty documents: " << empty_docs << endl + << std::flush; + cout << " number of documents with one sentence: " << one_sent_docs << endl + << std::flush; + cout << " number of documents with long sentences: " << long_sent_docs << endl + << std::flush; + cout << " will create mapping for " << map_index << " samples" << endl + << std::flush; + } + assert(maps == NULL); + assert(num_samples < 0); + maps = new DocIdx[3 * map_index]; + num_samples = static_cast(map_index); + } + + } // for (int iteration=0; iteration < 2; ++iteration) { + + // Shuffle. + // We need a 64 bit random number generator as we might have more + // than 2 billion samples. + std::mt19937_64 rand64_gen(seed + 1); + for (auto i = (num_samples - 1); i > 0; --i) + { + const auto j = static_cast(rand64_gen() % (i + 1)); + const auto i0 = 3 * i; + const auto j0 = 3 * j; + // Swap values. + swap(maps[i0], maps[j0]); + swap(maps[i0 + 1], maps[j0 + 1]); + swap(maps[i0 + 2], maps[j0 + 2]); + } + + // Method to deallocate memory. + py::capsule free_when_done(maps, [](void *mem_) + { + DocIdx *mem = reinterpret_cast(mem_); + delete[] mem; }); + + // Return the numpy array. + const auto byte_size = sizeof(DocIdx); + return py::array(std::vector{num_samples, 3}, // shape + {3 * byte_size, byte_size}, // C-style contiguous strides + maps, // the data pointer + free_when_done); // numpy array references +} + +py::array build_mapping(const py::array_t &docs_, + const py::array_t &sizes_, + const int num_epochs, + const uint64_t max_num_samples, + const int max_seq_length, + const double short_seq_prob, + const int seed, + const bool verbose, + const int32_t min_num_sent) +{ + + if (sizes_.size() > std::numeric_limits::max()) + { + if (verbose) + { + cout << " using uint64 for data mapping..." << endl + << std::flush; + } + return build_mapping_impl(docs_, sizes_, num_epochs, + max_num_samples, max_seq_length, + short_seq_prob, seed, verbose, + min_num_sent); + } + else + { + if (verbose) + { + cout << " using uint32 for data mapping..." << endl + << std::flush; + } + return build_mapping_impl(docs_, sizes_, num_epochs, + max_num_samples, max_seq_length, + short_seq_prob, seed, verbose, + min_num_sent); + } +} + +template +py::array build_blocks_mapping_impl(const py::array_t &docs_, + const py::array_t &sizes_, + const py::array_t &titles_sizes_, + const int32_t num_epochs, + const uint64_t max_num_samples, + const int32_t max_seq_length, + const int32_t seed, + const bool verbose, + const bool use_one_sent_blocks) +{ + /* Build a mapping of (start-index, end-index, sequence-length) where + start and end index are the indices of the sentences in the sample + and sequence-length is the target sequence length. + */ + + // Consistency checks. + assert(num_epochs > 0); + assert(max_seq_length > 1); + assert(seed > 0); + + // Remove bound checks. + auto docs = docs_.unchecked<1>(); + auto sizes = sizes_.unchecked<1>(); + auto titles_sizes = titles_sizes_.unchecked<1>(); + + if (verbose) + { + const auto sent_start_index = docs[0]; + const auto sent_end_index = docs[docs_.shape(0) - 1]; + const auto num_sentences = sent_end_index - sent_start_index; + cout << " using:" << endl + << std::flush; + cout << " number of documents: " << docs_.shape(0) - 1 << endl + << std::flush; + cout << " sentences range: [" << sent_start_index << ", " << sent_end_index << ")" << endl + << std::flush; + cout << " total number of sentences: " << num_sentences << endl + << std::flush; + cout << " number of epochs: " << num_epochs << endl + << std::flush; + cout << " maximum number of samples: " << max_num_samples << endl + << std::flush; + cout << " maximum sequence length: " << max_seq_length << endl + << std::flush; + cout << " seed: " << seed << endl + << std::flush; + } + + // Mapping and its length (1D). + int64_t num_samples = -1; + DocIdx *maps = NULL; + + // Acceptable number of sentences per block. + int min_num_sent = 2; + if (use_one_sent_blocks) + { + min_num_sent = 1; + } + + // Perform two iterations, in the first iteration get the size + // and allocate memory and in the second iteration populate the map. + bool second = false; + for (int32_t iteration = 0; iteration < 2; ++iteration) + { + + // Set the flag on second iteration. + second = (iteration == 1); + + // Current map index. + uint64_t map_index = 0; + + uint64_t empty_docs = 0; + uint64_t one_sent_docs = 0; + uint64_t long_sent_docs = 0; + // For each epoch: + for (int32_t epoch = 0; epoch < num_epochs; ++epoch) + { + // assign every block a unique id + int32_t block_id = 0; + + if (map_index >= max_num_samples) + { + if (verbose && (!second)) + { + cout << " reached " << max_num_samples << " samples after " + << epoch << " epochs ..." << endl + << std::flush; + } + break; + } + // For each document: + for (int32_t doc = 0; doc < (docs.shape(0) - 1); ++doc) + { + + // Document sentences are in [sent_index_first, sent_index_last) + const auto sent_index_first = docs[doc]; + const auto sent_index_last = docs[doc + 1]; + const auto target_seq_len = max_seq_length - titles_sizes[doc]; + + // At the begining of the document previous index is the + // start index. + auto prev_start_index = sent_index_first; + + // Remaining documents. + auto num_remain_sent = sent_index_last - sent_index_first; + + // Some bookkeeping + if ((epoch == 0) && (!second)) + { + if (num_remain_sent == 0) + { + ++empty_docs; + } + if (num_remain_sent == 1) + { + ++one_sent_docs; + } + } + // Detect documents with long sentences. + bool contains_long_sentence = false; + if (num_remain_sent >= min_num_sent) + { + for (auto sent_index = sent_index_first; + sent_index < sent_index_last; ++sent_index) + { + if (sizes[sent_index] > LONG_SENTENCE_LEN) + { + if ((epoch == 0) && (!second)) + { + ++long_sent_docs; + } + contains_long_sentence = true; + break; + } + } + } + // If we have enough sentences and no long sentences. + if ((num_remain_sent >= min_num_sent) && (!contains_long_sentence)) + { + + // Set values. + auto seq_len = int32_t{0}; + auto num_sent = int32_t{0}; + + // Loop through sentences. + for (auto sent_index = sent_index_first; + sent_index < sent_index_last; ++sent_index) + { + + // Add the size and number of sentences. + seq_len += sizes[sent_index]; + ++num_sent; + --num_remain_sent; + + // If we have reached the target length. + // and there are an acceptable number of sentences left + // and if we have at least the minimum number of sentences. + // or if we have reached end of the document. + if (((seq_len >= target_seq_len) && + (num_remain_sent >= min_num_sent) && + (num_sent >= min_num_sent)) || + (num_remain_sent == 0)) + { + + // Populate the map. + if (second) + { + const auto map_index_0 = 4 * map_index; + // Each sample has 4 items: the starting sentence index, ending sentence index, + // the index of the document from which the block comes (used for fetching titles) + // and the unique id of the block (used for creating block indexes) + + maps[map_index_0] = static_cast(prev_start_index); + maps[map_index_0 + 1] = static_cast(sent_index + 1); + maps[map_index_0 + 2] = static_cast(doc); + maps[map_index_0 + 3] = static_cast(block_id); + } + + // Update indices / counters. + ++map_index; + ++block_id; + prev_start_index = sent_index + 1; + seq_len = 0; + num_sent = 0; + } + } // for (auto sent_index=sent_index_first; ... + } // if (num_remain_sent > 1) { + } // for (int doc=0; doc < num_docs; ++doc) { + } // for (int epoch=0; epoch < num_epochs; ++epoch) { + + if (!second) + { + if (verbose) + { + cout << " number of empty documents: " << empty_docs << endl + << std::flush; + cout << " number of documents with one sentence: " << one_sent_docs << endl + << std::flush; + cout << " number of documents with long sentences: " << long_sent_docs << endl + << std::flush; + cout << " will create mapping for " << map_index << " samples" << endl + << std::flush; + } + assert(maps == NULL); + assert(num_samples < 0); + maps = new DocIdx[4 * map_index]; + num_samples = static_cast(map_index); + } + + } // for (int iteration=0; iteration < 2; ++iteration) { + + // Shuffle. + // We need a 64 bit random number generator as we might have more + // than 2 billion samples. + std::mt19937_64 rand64_gen(seed + 1); + for (auto i = (num_samples - 1); i > 0; --i) + { + const auto j = static_cast(rand64_gen() % (i + 1)); + const auto i0 = 4 * i; + const auto j0 = 4 * j; + // Swap values. + swap(maps[i0], maps[j0]); + swap(maps[i0 + 1], maps[j0 + 1]); + swap(maps[i0 + 2], maps[j0 + 2]); + swap(maps[i0 + 3], maps[j0 + 3]); + } + + // Method to deallocate memory. + py::capsule free_when_done(maps, [](void *mem_) + { + DocIdx *mem = reinterpret_cast(mem_); + delete[] mem; }); + + // Return the numpy array. + const auto byte_size = sizeof(DocIdx); + return py::array(std::vector{num_samples, 4}, // shape + {4 * byte_size, byte_size}, // C-style contiguous strides + maps, // the data pointer + free_when_done); // numpy array references +} + +py::array build_blocks_mapping(const py::array_t &docs_, + const py::array_t &sizes_, + const py::array_t &titles_sizes_, + const int num_epochs, + const uint64_t max_num_samples, + const int max_seq_length, + const int seed, + const bool verbose, + const bool use_one_sent_blocks) +{ + + if (sizes_.size() > std::numeric_limits::max()) + { + if (verbose) + { + cout << " using uint64 for data mapping..." << endl + << std::flush; + } + return build_blocks_mapping_impl(docs_, sizes_, titles_sizes_, + num_epochs, max_num_samples, max_seq_length, seed, verbose, use_one_sent_blocks); + } + else + { + if (verbose) + { + cout << " using uint32 for data mapping..." << endl + << std::flush; + } + return build_blocks_mapping_impl(docs_, sizes_, titles_sizes_, + num_epochs, max_num_samples, max_seq_length, seed, verbose, use_one_sent_blocks); + } +} + +PYBIND11_MODULE(helpers, m) +{ + m.def("build_mapping", &build_mapping); + m.def("build_blocks_mapping", &build_blocks_mapping); + m.def("build_sample_idx", &build_sample_idx); + m.def("build_blending_indices", &build_blending_indices); + m.def("build_exhaustive_blending_indices", &build_exhaustive_blending_indices); +} + diff --git a/mindformers/dataset/blended_datasets/indexed_dataset.py b/mindformers/dataset/blended_datasets/indexed_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..ce7b463c5dc70ab7a31968e870ee4e6ea46a0a50 --- /dev/null +++ b/mindformers/dataset/blended_datasets/indexed_dataset.py @@ -0,0 +1,860 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# Essentially re-written in entirety + +# import logging +import os +import shutil +import struct +import time +from abc import ABC, abstractmethod +from enum import Enum +from functools import lru_cache +from itertools import accumulate +from types import TracebackType +from typing import List, Optional, Tuple, Type, Union + +try: + import boto3 +except ModuleNotFoundError: + pass +import numpy +# import torch + +from mindformers.dataset.blended_datasets.utils_s3 import ( + S3Config, + is_s3_path, + maybe_download_file, + object_exists, + parse_s3_path, +) +from mindformers.tools.logger import logger +# logger = logging.getLogger(__name__) + +_INDEX_HEADER = b"MMIDIDX\x00\x00" + + +class DType(Enum): + """The NumPy data type Enum for writing/reading the IndexedDataset indices""" + + uint8 = 1 + int8 = 2 + int16 = 3 + int32 = 4 + int64 = 5 + float64 = 6 + float32 = 7 + uint16 = 8 + + @classmethod + def code_from_dtype(cls, value: Type[numpy.number]) -> int: + """Get the code from the dtype + + Args: + value (Type[numpy.number]): The dtype + + Returns: + int: The code + """ + return cls[value.__name__].value + + @classmethod + def dtype_from_code(cls, value: int) -> Type[numpy.number]: + """Get the dtype from the code + + Args: + value (int): The code + + Returns: + Type[numpy.number]: The dtype + """ + return getattr(numpy, cls(value).name) + + @staticmethod + def size(key: Union[int, Type[numpy.number]]) -> int: + """Get the size of the dtype/code in bytes + + Args: + key (Union[int, Type[numpy.number]]): The dtype or code + + Raises: + ValueError: If the key is neither dtype nor integer code + + Returns: + int: The size of the dtype/code in in bytes + """ + if isinstance(key, int): + return DType.dtype_from_code(key)().itemsize + elif numpy.number in key.__mro__: + return key().itemsize + else: + raise ValueError + + @staticmethod + def optimal_dtype(cardinality: Optional[int]) -> Type[numpy.number]: + """Get the dtype to use for an index of a certain cardinality + + Args: + cardinality (Optional[int]): The number of elements to be indexed + + Returns: + Type[numpy.number]: The dtype to use for the index + """ + if cardinality is not None and cardinality < 65500: + return numpy.uint16 + else: + return numpy.int32 + + +class _IndexWriter(object): + """Object class to write the index (.idx) file + + Args: + idx_path (str): The path to the index file + + dtype (Type[numpy.number]): The dtype of the index file + """ + + def __init__(self, idx_path: str, dtype: Type[numpy.number]) -> None: + self.idx_path = idx_path + self.dtype = dtype + + def __enter__(self) -> "_IndexWriter": + """Enter the context introduced by the 'with' keyword + + Returns: + _IndexWriter: The instance + """ + self.idx_writer = open(self.idx_path, "wb") + # fixed, vestigial practice + self.idx_writer.write(_INDEX_HEADER) + # fixed, vestigial practice + self.idx_writer.write(struct.pack(" Optional[bool]: + """Exit the context introduced by the 'with' keyword + + Args: + exc_type (Optional[Type[BaseException]]): Exception type + + exc_val (Optional[BaseException]): Exception value + + exc_tb (Optional[TracebackType]): Exception traceback object + + Returns: + Optional[bool]: Whether to silence the exception + """ + self.idx_writer.close() + + def write( + self, + sequence_lengths: List[int], + sequence_modes: Optional[List[int]], + document_indices: List[int], + ) -> None: + """Write the index (.idx) file + + Args: + sequence_lengths (List[int]): The length of each sequence + + sequence_modes (Optional[List[int]]): The mode of each sequences + + document_indices (List[int]): The seqyebce indices demarcating the end of each document + """ + sequence_pointers = self._sequence_pointers(sequence_lengths) + + # the number of sequences in the dataset + sequence_count = len(sequence_lengths) + self.idx_writer.write(struct.pack(" List[int]: + """Build the sequence pointers per the sequence lengths and dtype size + + Args: + sequence_lengths (List[int]): The length of each sequence + + Returns: + List[int]: The pointer to the beginning of each sequence + """ + itemsize = DType.size(self.dtype) + curr_ptr = 0 + list_ptr = [] + for length in sequence_lengths: + list_ptr.append(curr_ptr) + curr_ptr += length * itemsize + return list_ptr + + +class _IndexReader(object): + """Object class to read the index (.idx) file + + Args: + idx_path (str): The path to the index file + + multimodal (bool): Whether the dataset is multimodal + """ + + def __init__(self, idx_path: str, multimodal: bool) -> None: + + logger.info(f"Load the {type(self).__name__} from {idx_path}") + + with open(idx_path, "rb") as stream: + header = stream.read(9) + assert header == _INDEX_HEADER, f"bad header, cannot read: {idx_path}" + + version = struct.unpack(" time elapsed: {t_end - t_beg:4f} seconds") + + logger.info(f"\tExtract the sequence pointers") + t_beg = time.time() + self.sequence_pointers = numpy.frombuffer( + self.bin_buffer, + dtype=numpy.int64, + count=self.sequence_count, + offset=offset + self.sequence_lengths.nbytes, + ) + t_end = time.time() + logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + logger.info(f"\tExtract the document indices") + t_beg = time.time() + self.document_indices = numpy.frombuffer( + self.bin_buffer, + dtype=numpy.int64, + count=self.document_count, + offset=offset + self.sequence_lengths.nbytes + self.sequence_pointers.nbytes, + ) + t_end = time.time() + logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + self.sequence_modes = None + if multimodal: + logger.info(f"\tExtract the sequence modes") + t_beg = time.time() + self.sequence_modes = numpy.frombuffer( + self.bin_buffer, + dtype=numpy.int8, + count=self.sequence_count, + offset=offset + + self.sequence_lengths.nbytes + + self.sequence_pointers.nbytes + + self.document_indices.nbytes, + ) + t_end = time.time() + logger.debug(f"\t> time elapsed: {t_end - t_beg:4f} seconds") + + assert self.sequence_lengths.shape[0] == len(self) + assert self.sequence_lengths.shape[0] == self.sequence_count + assert self.sequence_lengths.shape[0] == self.document_indices[-1] + + logger.info(f"> total number of sequences: {len(self)}") + logger.info(f"> total number of documents: {self.document_indices.shape[0] - 1}") + + def __del__(self) -> None: + """Clean up the object""" + if hasattr(self, "bin_buffer_mmap"): + self.bin_buffer_mmap._mmap.close() + del self.bin_buffer_mmap + + def __len__(self) -> int: + """Return the length of the dataset + + Returns: + int: The length of the dataset + """ + return self.sequence_count + + @lru_cache(maxsize=8) + def __getitem__(self, idx: int) -> Tuple[numpy.int32, numpy.int64, Optional[numpy.int8]]: + """Return the pointer, length, and mode at the index + + Args: + idx (int): The index into the dataset + + Returns: + Tuple[numpy.int32, numpy.int64, Optional[numpy.int8]]: The pointer, length and mode at the index + """ + return ( + self.sequence_pointers[idx], + self.sequence_lengths[idx], + self.sequence_modes[idx] if self.sequence_modes is not None else None, + ) + + +class _BinReader(ABC): + """Abstract class to read the data (.bin) file""" + + @abstractmethod + def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: + """Read bytes into a numpy array. + + Args: + dtype (Type[numpy.number]): Data-type of the returned array. + + count (int): Number of items to read. + + offset (int): Start reading from this offset (in bytes). + + Returns: + numpy.ndarray: An array with `count` items and data-type `dtype` constructed from reading bytes from the data file starting at `offset`. + """ + pass + + +class _MMapBinReader(_BinReader): + """A _BinReader that memory maps the data (.bin) file + + Args: + bin_path (str): bin_path (str): The path to the data (.bin) file. + """ + + def __init__(self, bin_path: str) -> None: + self._bin_buffer_mmap = numpy.memmap(bin_path, mode="r", order="C") + self._bin_buffer = memoryview(self._bin_buffer_mmap) + + def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: + """Read bytes into a numpy array. + + Args: + dtype (Type[numpy.number]): Data-type of the returned array. + + count (int): Number of items to read. + + offset (int): Start reading from this offset (in bytes). + + Returns: + numpy.ndarray: An array with `count` items and data-type `dtype` constructed from reading bytes from the data file starting at `offset`. + """ + return numpy.frombuffer( + self._bin_buffer, + dtype=dtype, + count=count, + offset=offset, + ) + + def __del__(self) -> None: + """Clean up the object.""" + if self._bin_buffer_mmap is not None: + self._bin_buffer_mmap._mmap.close() + del self._bin_buffer_mmap + + +class _FileBinReader(_BinReader): + """A _BinReader that reads from the data (.bin) file using a file pointer + + Args: + bin_path (str): bin_path (str): The path to the data (.bin) file. + """ + + def __init__(self, bin_path: str) -> None: + self._bin_path = bin_path + + def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: + """Read bytes into a numpy array. + + Args: + dtype (Type[numpy.number]): Data-type of the returned array. + + count (int): Number of items to read. + + offset (int): Start reading from this offset (in bytes). + + Returns: + numpy.ndarray: An array with `count` items and data-type `dtype` constructed from reading bytes from the data file starting at `offset`. + """ + sequence = numpy.empty(count, dtype=dtype) + with open(self._bin_path, mode='rb', buffering=0) as bin_buffer_file: + bin_buffer_file.seek(offset) + bin_buffer_file.readinto(sequence) + return sequence + + +class _S3BinReader(_BinReader): + """A _BinReader that reads from the data (.bin) file from S3 + + Args: + bin_path (str): bin_path (str): The path to the data (.bin) file. + + bin_chunk_nbytes (int, optional): If not None, then maintain an in-memory cache to speed up calls to the `read` method. Furthermore, on a cache miss, download this number of bytes to refresh the cache. Otherwise (None), do not maintain an in-memory cache. A class that inherits from _BinReader may not implement caching in which case it should assert that `bin_chunk_nbytes` is None at initialization. + """ + + def __init__(self, bin_path: str, bin_chunk_nbytes: int) -> None: + assert bin_chunk_nbytes > 0 + self._client = boto3.client("s3") + self._s3_bucket, self._s3_key = parse_s3_path(bin_path) + self._cache = None + self._cache_bytes_start = None + self._cache_bytes_end = None + self._cache_nbytes = bin_chunk_nbytes + + def _extract_from_cache(self, offset: int, size: int) -> bytes: + """Extract `size` bytes starting at `offset` bytes into the cache""" + start = offset - self._cache_bytes_start + assert start >= 0 + end = start + size + assert end <= len(self._cache) + return self._cache[start:end] + + def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray: + """Read bytes into a numpy array. + + Let `size` be the `count` * `DType.size(dtype)`. If the requested span of bytes [`offset`, + `offset` + `size`) is covered by the in-memory cache maintained by this class, then this + function extracts the requested span from that cache and returns it. Otherwise, this + function first refreshes the cache and then extracts the requested span from the refreshed + cache and returns it. + + The cache is refreshed based on `offset` and `size`. In particular, we divide all the bytes + in an S3 object into blocks, where each block contains `bin_chunk_nbytes` bytes. We assign + each block an index starting from 0. We take the block with index (`offset` // + `bin_chunk_nbytes`) to refresh the cache. If this new block still does not cover the + requested span, we extend it just enough to include `offset` + `size`. + + Args: + dtype (Type[numpy.number]): Data-type of the returned array. + + count (int): Number of items to read. + + offset (int): Start reading from this offset (in bytes). + + Returns: + numpy.ndarray: An array with `count` items and data-type `dtype` constructed from reading bytes from the data file starting at `offset`. + """ + size = count * DType.size(dtype) + if ( + self._cache is not None + and offset >= self._cache_bytes_start + and offset + size <= self._cache_bytes_end + ): + return numpy.frombuffer(self._extract_from_cache(offset, size), dtype=dtype) + + bytes_start = (offset // self._cache_nbytes) * self._cache_nbytes + assert bytes_start >= 0 + assert offset >= bytes_start + bytes_end = max(bytes_start + self._cache_nbytes, offset + size) + assert bytes_end >= 1 + self._cache = self._client.get_object( + Bucket=self._s3_bucket, + Key=self._s3_key, + # Subtract 1, because the end of Range is inclusive. + Range=f'bytes={bytes_start}-{bytes_end-1}', + )['Body'].read() + self._cache_bytes_start = bytes_start + self._cache_bytes_end = bytes_end + return numpy.frombuffer(self._extract_from_cache(offset, size), dtype=dtype) + + def __del__(self) -> None: + """Clean up the object""" + self._client.close() + + +# class IndexedDataset(torch.utils.data.Dataset): +class IndexedDataset(): + """The low-level interface dataset class + + Args: + path_prefix (str): The index (.idx) and data (.bin) prefix + + multimodal (bool): Whether the dataset is multimodal. Defaults to False. + + mmap (bool): Whether to mmap the .bin files. Defaults to True. + + s3_config (Optional[S3Config]): Supplied only for data stored on S3. IndexedDataset downloads the index (.idx) file to `s3_config.path_to_idx_cache` and streams data from the data (.bin) file in `s3_config.bin_chunk_nbytes` blocks. Note that `mmap` must be disabled for S3 data loading. Defaults to None. + """ + + def __init__( + self, + path_prefix: str, + multimodal: bool = False, + mmap: bool = True, + s3_config: Optional[S3Config] = None, + ) -> None: + super().__init__() + self.path_prefix = None + self.multimodal = None + self.mmap = None + self.s3_config = None + + self.index = None + self.bin_reader = None + + if is_s3_path(path_prefix) and s3_config is not None: + idx_path = get_idx_path(path_prefix) + cache_idx_path = os.path.join(s3_config.path_to_idx_cache, os.path.basename(idx_path)) + maybe_download_file(idx_path, cache_idx_path) + + self.initialize(path_prefix, multimodal, mmap, s3_config) + + def initialize( + self, path_prefix: str, multimodal: bool, mmap: bool, s3_config: Optional[S3Config] + ) -> None: + """Initialize the dataset + + This method is called by IndexedDataset.__init__ during object creation and by + IndexedDataset.__setstate__ during un-pickling + + Args: + path_prefix (str): The index (.idx) and data (.bin) prefix + + multimodal (bool): Whether the dataset is multimodal + + mmap (bool): Whether to mmap the .bin file + + s3_config (Optional[S3Config]): See IndexedDataset docstring for details. + """ + idx_path = get_idx_path(path_prefix) + bin_path = get_bin_path(path_prefix) + if s3_config is None: + assert os.path.exists(idx_path) and os.path.exists( + bin_path + ), f"One or both of the .idx and .bin files cannot be found at the path prefix {path_prefix}" + self.path_prefix = path_prefix + self.multimodal = multimodal + self.mmap = mmap + self.s3_config = s3_config + if mmap: + assert not s3_config + self.bin_reader = _MMapBinReader(bin_path) + elif s3_config: + assert not mmap + self.bin_reader = _S3BinReader(bin_path, s3_config.bin_chunk_nbytes) + idx_path = os.path.join( + s3_config.path_to_idx_cache, os.path.basename(get_idx_path(path_prefix)) + ) + else: + self.bin_reader = _FileBinReader(bin_path) + self.index = _IndexReader(idx_path, self.multimodal) + + def __getstate__(self) -> Tuple[str, bool, bool, Optional[S3Config]]: + """Get the state during pickling + + Returns: + Tuple[str, bool, bool, Optional[S3Config]]: The state tuple + """ + return self.path_prefix, self.multimodal, self.mmap, self.s3_config + + def __setstate__(self, state: Tuple[str, bool, bool, Optional[S3Config]]) -> None: + """Set the state during un-pickling + + Args: + state (Tuple[str, bool, bool, Optional[S3Config]]): The state tuple + """ + path_prefix, multimodal, mmap, s3_config = state + self.initialize(path_prefix, multimodal, mmap, s3_config) + + def __del__(self) -> None: + """Clean up the object""" + del self.bin_reader + del self.index + + def __len__(self) -> int: + """Return the length of the dataset i.e. the number of sequences in the index + + Returns: + int: The length of the dataset + """ + return len(self.index) + + def __getitem__( + self, idx: Union[int, numpy.integer, slice] + ) -> Union[numpy.ndarray, Tuple[numpy.ndarray, numpy.ndarray]]: + """Return from the dataset + + Args: + idx (Union[int, numpy.integer, slice]): The index or index slice into the dataset + + Raises: + ValueError: When the index slice is non-contiguous + + TypeError: When the index is of an unexpected type + + Returns: + Union[numpy.ndarray, Tuple[numpy.ndarray, numpy.ndarray]]: The sequence tokens and modes at the index or index slice + """ + if isinstance(idx, (int, numpy.integer)): + sequence_pointer, sequence_length, sequence_mode = self.index[idx] + sequence = self.bin_reader.read( + dtype=self.index.dtype, + count=sequence_length, + offset=sequence_pointer, + ) + return (sequence, sequence_mode) if sequence_mode is not None else sequence + elif isinstance(idx, slice): + start, stop, step = idx.indices(len(self)) + if step != 1: + raise ValueError("Slices into indexed_dataset must be contiguous") + sequence_lengths = self.index.sequence_lengths[idx] + sequence_modes = self.index.sequence_modes[idx] if self.multimodal else None + sequence_offsets = list(accumulate(sequence_lengths)) + sequences = numpy.split( + self.bin_reader.read( + dtype=self.index.dtype, + count=sum(sequence_lengths), + offset=self.index.sequence_pointers[start], + ), + sequence_offsets[:-1], + ) + return (sequences, sequence_modes) if sequence_modes is not None else sequences + else: + raise TypeError("Unexpected type received for idx: {}".format(type(idx))) + + def get(self, idx: int, offset: int = 0, length: Optional[int] = None) -> numpy.ndarray: + """Retrieve a single item from the dataset with the option to only + return a portion of the item. + + get(idx) is the same as [idx] but get() does not support slicing. + + Args: + idx (Union[int, numpy.integer]): The index into the dataset + + offset (int): The integer token offset in the sequence + + length (int): The number of tokens to grab from the sequence + + Returns: + Union[numpy.ndarray, Tuple[numpy.ndarray, numpy.ndarray]]: The sequence tokens and modes at the index + """ + sequence_pointer, sequence_length, sequence_mode = self.index[idx] + if length is None: + length = sequence_length - offset + sequence_pointer += offset * DType.size(self.index.dtype) + sequence = self.bin_reader.read( + dtype=self.index.dtype, count=length, offset=sequence_pointer + ) + return (sequence, sequence_mode) if sequence_mode is not None else sequence + + @property + def sequence_lengths(self) -> numpy.ndarray: + """Get the sequence lengths + + Returns: + numpy.ndarray: The sequence lengths + """ + return self.index.sequence_lengths + + @property + def document_indices(self) -> numpy.ndarray: + """Get the document indices + + Returns: + numpy.ndarray: The document indices + """ + return self.index.document_indices + + def get_document_indices(self) -> numpy.ndarray: + """Get the document indices + + This method is slated for deprecation. + + Returns: + numpy.ndarray: The document indices + """ + return self.index.document_indices + + def set_document_indices(self, document_indices: numpy.ndarray) -> None: + """Set the document indices + + This method is slated for deprecation. + + Args: + document_indices (numpy.ndarray): The document indices + """ + self.index.document_indices = document_indices + + @property + def sequence_modes(self) -> numpy.ndarray: + """Get the sequence modes + + Returns: + numpy.ndarray: The sequence modes + """ + return self.index.sequence_modes + + @staticmethod + def exists(path_prefix: str) -> bool: + """Return whether the IndexedDataset exists on disk at the prefix + + Args: + path_prefix (str): The prefix to the index (.idx) and data (.bin) files + + Returns: + bool: Whether the IndexedDataset exists on disk at the prefix + """ + if is_s3_path(path_prefix): + s3_client = boto3.client("s3") + return object_exists(s3_client, get_idx_path(path_prefix)) and object_exists( + s3_client, get_bin_path(path_prefix) + ) + return os.path.exists(get_idx_path(path_prefix)) and os.path.exists( + get_bin_path(path_prefix) + ) + + +class IndexedDatasetBuilder(object): + """Builder class for the IndexedDataset class + + Args: + bin_path (str): The path to the data (.bin) file + + dtype (Type[numpy.number], optional): The dtype of the index file. Defaults to numpy.int32. + + multimodal (bool, optional): Whether the dataset is multimodal. Defaults to False. + """ + + def __init__( + self, bin_path: str, dtype: Type[numpy.number] = numpy.int32, multimodal: bool = False + ) -> None: + self.data_file = open(bin_path, "wb") + self.dtype = dtype + self.multimodal = multimodal + + self.sequence_lengths = [] + self.document_indices = [0] + self.sequence_modes = [] if self.multimodal else None + + def add_item(self, np_array: numpy.ndarray, mode: int = 0) -> None: + """Add a single item to the dataset + + Args: + tensor (numpy.ndarray): The item to add to the data file + + mode (int, optional): The mode for the item. Defaults to 0. + """ + # np_array = numpy.array(tensor.numpy(), dtype=self.dtype) + self.data_file.write(np_array.tobytes(order="C")) + self.sequence_lengths.append(np_array.size) + if self.multimodal: + self.sequence_modes.append(mode) + + def add_document( + self, np_array: numpy.ndarray, lengths: List[int], modes: Optional[List[int]] = None + ) -> None: + """Add an entire document to the dataset + + Args: + tensor (numpy.ndarray): The document to add + + lengths (List[int]): The lengths of each item in the document + + modes (Optional[List[int]], optional): The modes for each item in the document. Defaults to None. + """ + # np_array = numpy.array(tensor, dtype=self.dtype) + self.data_file.write(np_array.tobytes(order="C")) + self.sequence_lengths.extend(lengths) + self.document_indices.append(len(self.sequence_lengths)) + if self.multimodal: + self.sequence_modes.extend(modes if modes is not None else [0] * lengths) + + def end_document(self) -> None: + """Finalize the document, for use with IndexedDatasetBuilder.add_item""" + self.document_indices.append(len(self.sequence_lengths)) + + def add_index(self, path_prefix: str) -> None: + """Add an entire IndexedDataset to the dataset + + Args: + path_prefix (str): The index (.idx) and data (.bin) prefix + """ + # Concatenate index + index = _IndexReader(get_idx_path(path_prefix), multimodal=self.multimodal) + assert index.dtype == self.dtype + + offset = len(self.sequence_lengths) + self.sequence_lengths.extend(index.sequence_lengths) + self.document_indices.extend((offset + index.document_indices)[1:]) + + if self.multimodal: + self.sequence_modes.extend(index.sequence_modes) + + # Concatenate data + with open(get_bin_path(path_prefix), "rb") as f: + shutil.copyfileobj(f, self.data_file) + + def finalize(self, idx_path: str) -> None: + """Clean up and write the index (.idx) file + + Args: + idx_path (str): The path to the index file + """ + self.data_file.close() + with _IndexWriter(idx_path, self.dtype) as writer: + writer.write(self.sequence_lengths, self.sequence_modes, self.document_indices) + + +def get_idx_path(path_prefix: str) -> str: + """Get the path to the index file from the prefix + + Args: + path_prefix (str): The prefix + + Returns: + str: The path to the index file + """ + return path_prefix + ".idx" + + +def get_bin_path(path_prefix: str) -> str: + """Get the path to the data file from the prefix + + Args: + path_prefix (str): The prefix + + Returns: + str: The path to the data file + """ + return path_prefix + ".bin" diff --git a/mindformers/dataset/blended_datasets/megatron_dataset.py b/mindformers/dataset/blended_datasets/megatron_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..bcb3b6379e2b54b60da9af74fe8e4ad065489fae --- /dev/null +++ b/mindformers/dataset/blended_datasets/megatron_dataset.py @@ -0,0 +1,140 @@ +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + +import hashlib +import json +from abc import ABC, abstractmethod +from collections import OrderedDict +from typing import Any, Dict, Iterable, List, Optional, Union + +import numpy + +from mindformers.dataset.blended_datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig +from mindformers.dataset.blended_datasets.indexed_dataset import IndexedDataset +from mindformers.dataset.blended_datasets.utils import Split + +LowLevelDataset = Union[IndexedDataset, Iterable] + + +# class MegatronDataset(ABC, torch.utils.data.Dataset): +class MegatronDataset(ABC): + """The highest level wrapper class from which all dataset classes should inherit + + Args: + dataset (LowLevelDataset): The dataset around which to build the MegatronDataset + + dataset_path (Optional[str]): The real path on disk to the dataset, for bookkeeping + + indices (numpy.ndarray): The set of the documents indices to expose + + num_samples (Optional[int]): The minimum number of samples to build from the indexed dataset. When None, build as many samples as correspond to one epoch. + + index_split (Split): The indices Split + + config (BlendedMegatronDatasetConfig): The config + """ + + def __init__( + self, + dataset: LowLevelDataset, + dataset_path: Optional[str], + indices: numpy.ndarray, + num_samples: Optional[int], + index_split: Split, + config: BlendedMegatronDatasetConfig, + ) -> None: + self.dataset = dataset + self.dataset_path = dataset_path + self.indices = indices + self.num_samples = num_samples + self.index_split = index_split + self.config = config + + self.unique_identifiers = OrderedDict() + + self.unique_identifiers["class"] = type(self).__name__ + self.unique_identifiers["dataset_path"] = self.dataset_path + self.unique_identifiers["num_samples"] = self.num_samples + self.unique_identifiers["index_split"] = self.index_split.name + for attr in self._key_config_attributes(): + self.unique_identifiers[attr] = getattr(self.config, attr) + + self.unique_description = json.dumps( + self.unique_identifiers, indent=4, default=lambda obj: obj.unique_identifiers + ) + self.unique_description_hash = hashlib.md5( + self.unique_description.encode("utf-8") + ).hexdigest() + + self.built_anew_on_cache_miss = False + + @staticmethod + def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int: + """Return the number of elements in the underlying low level dataset for the purpose of + segregating the train/valid/test split indices + + It may be that the low level dataset can be split any number of ways, depending on the mid + level dataset it supports, which is why we define the "number of elements" function + separately from the __len__ function here in the mid level dataset class + + Args: + low_level_dataset (LowLevelDataset): The underlying low level dataset + + Returns: + int: The number of elements in the underlying low level dataset + """ + raise NotImplementedError + + @staticmethod + def build_low_level_dataset( + dataset_path: str, config: BlendedMegatronDatasetConfig + ) -> LowLevelDataset: + """Build the low level dataset via a function to be called from within + BlendedMegatronDatasetBuilder.build_generic_dataset + + It may be that the low level dataset spans any subset of train/valid/test splits, which is + why we define a static "build" function separately from the constructor in the mid level + dataset class + + Args: + dataset_path (str): The real path on disk to the dataset + + config (BlendedMegatronDatasetConfig): The dataset config + + Returns: + LowLevelDataset: The low level dataset + """ + raise NotImplementedError + + @staticmethod + def _key_config_attributes() -> List[str]: + """Return all config attributes which contribute to uniquely identifying the dataset. + + These attributes will be used to build a uniquely identifying string and MD5 hash which + will be used to cache/load dataset resources from run to run. + + Returns: + List[str]: The key config attributes + """ + return ["random_seed", "sequence_length", "split", "split_matrix", "tokenizer"] + + @abstractmethod + def __len__(self) -> int: + """Return the length of the dataset + + Returns: + int: See abstract implementation + """ + pass + + @abstractmethod + # def __getitem__(self, idx: int) -> Dict[str, Union[torch.Tensor, numpy.ndarray]]: + def __getitem__(self, idx: int) -> Dict[str, numpy.ndarray]: + """Return from the dataset + + Args: + idx (int): The index into the dataset + + Returns: + Dict[str, numpy.ndarray]: See abstract implementation + """ + pass diff --git a/mindformers/dataset/blended_datasets/utils.py b/mindformers/dataset/blended_datasets/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3373c35477c1a733c5e317d209403115472f8cc2 --- /dev/null +++ b/mindformers/dataset/blended_datasets/utils.py @@ -0,0 +1,87 @@ +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. + +import logging +from enum import Enum +from typing import List, Optional, Tuple + +import numpy +from mindformers.tools import logger + + +logger = logging.getLogger(__name__) + + +class Split(Enum): + train = 0 + valid = 1 + test = 2 + + +def compile_helpers(): + """Compile C++ helper functions at runtime. Make sure this is invoked on a single process. + """ + import os + import subprocess + + command = ["make", "-C", os.path.abspath(os.path.dirname(__file__))] + if subprocess.run(command).returncode != 0: + import sys + + logger.error("Failed to compile the C++ dataset helper functions") + sys.exit(1) + + +def normalize(weights: List[float]) -> List[float]: + """Do non-exponentiated normalization + + Args: + weights (List[float]): The weights + + Returns: + List[float]: The normalized weights + """ + w = numpy.array(weights, dtype=numpy.float64) + w_sum = numpy.sum(w) + w = (w / w_sum).tolist() + return w + + +def get_blend_from_list( + blend: Optional[List[str]], +) -> Optional[Tuple[List[str], Optional[List[float]]]]: + """Get the megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig blend from the blend list + + Args: + blend (Optional[List[str]]): The blend list, which can be either (1) a list of prefixes, e.g. ["path/to/dataset_1_prefix", "path/to/dataset_2_prefix"], or (2) a flattened, zipped list of weights and prefixes, e.g. ["30", "path/to/dataset_1_prefix", "70", "path/to/dataset_2_prefix"] + + Returns: + Optional[Tuple[List[str], Optional[List[float]]]]: The blend, consisting of a list of dataset prefixes and optionally a list of dataset weights, e.g. [["path/to/dataset_1_prefix", "path/to/dataset_2_prefix"], [30.0, 70.0]]. + """ + if blend is None: + return None + + if len(blend) % 2 == 1: + weight_per_dataset = None + raw_prefix_per_dataset = blend + else: + raw_weight_per_dataset, raw_prefix_per_dataset = zip( + *[(blend[i], blend[i + 1]) for i in range(0, len(blend), 2)] + ) + + weight_per_dataset = [] + for rwpd in raw_weight_per_dataset: + try: + weight = float(rwpd) + except ValueError: + weight = None + weight_per_dataset.append(weight) + + is_none = map(lambda _: _ is None, weight_per_dataset) + if any(is_none): + assert all(is_none) + weight_per_dataset = None + raw_prefix_per_dataset = blend + + prefix_per_dataset = [rppd.strip() for rppd in raw_prefix_per_dataset] + + return prefix_per_dataset, weight_per_dataset diff --git a/mindformers/dataset/blended_datasets/utils_s3.py b/mindformers/dataset/blended_datasets/utils_s3.py new file mode 100644 index 0000000000000000000000000000000000000000..d557360a1bcd97566ef908d0021178301198e1d2 --- /dev/null +++ b/mindformers/dataset/blended_datasets/utils_s3.py @@ -0,0 +1,165 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +import os +from typing import Any, Dict, NamedTuple, Protocol, Tuple + +try: + import boto3 + import botocore.exceptions as exceptions +except ModuleNotFoundError: + pass + +from mindspore.communication.comm_func import barrier +from mindformers.tools.utils import get_real_group_size, get_real_rank + +S3_PREFIX = "s3://" + + +class S3Config(NamedTuple): + """Config when the data (.bin) file and the index (.idx) file are in S3 + + TODO: These parameters are few and can be consolidated with parameters specific to bin reader + classes - @jkamalu + + Attributes: + + path_to_idx_cache (str): The local directory where we will store the index (.idx) file + + bin_chunk_nbytes (int): If the number of bytes is too small, then we send a request to S3 at each call of the `read` method in _S3BinReader, which is slow, because each request has a fixed cost independent of the size of the byte range requested. If the number of bytes is too large, then we only rarely have to send requests to S3, but it takes a lot of time to complete the request when we do, which can block training. We've found that 256 * 1024 * 1024 (i.e., 256 MiB) has worked well (though we have not put that much effort into tuning it), so we default to it. + """ + + path_to_idx_cache: str + + bin_chunk_nbytes: int = 256 * 1024 * 1024 + + +class S3Client(Protocol): + """The protocol which all s3 clients should abide by""" + + def download_file(self, Bucket: str, Key: str, Filename: str) -> None: ... + + def upload_file(self, Filename: str, Bucket: str, Key: str) -> None: ... + + def head_object(self, Bucket: str, Key: str) -> Dict[str, Any]: ... + + def get_object(self, Bucket: str, Key: str, Range: str) -> Dict[str, Any]: ... + + def close(self) -> None: ... + + +def is_s3_path(path: str) -> bool: + """Ascertain whether a path is in S3 + + Args: + path (str): The path + + Returns: + bool: True if the path is in S3, False otherwise + """ + return path.startswith(S3_PREFIX) + + +def parse_s3_path(path: str) -> Tuple[str, str]: + """Parses the given S3 path returning correspsonding bucket and key. + + Args: + path (str): The S3 path + + Returns: + Tuple[str, str]: A (bucket, key) tuple + """ + assert is_s3_path(path) + parts = path.replace(S3_PREFIX, "").split("/") + bucket = parts[0] + if len(parts) > 1: + key = "/".join(parts[1:]) + assert S3_PREFIX + bucket + "/" + key == path + else: + key = "" + return bucket, key + + +def object_exists(client: S3Client, path: str) -> bool: + """Ascertain whether the object at the given S3 path exists in S3 + + Args: + client (S3Client): The S3 client + + path (str): The S3 path + + Raises: + botocore.exceptions.ClientError: The error code is 404 + + Returns: + bool: True if the object exists in S3, False otherwise + """ + parsed_s3_path = parse_s3_path(path) + try: + response = client.head_object(bucket=parsed_s3_path[0], key=parsed_s3_path[1]) + except exceptions.ClientError as e: + if e.response["Error"]["Code"] != "404": + raise e + return True + + +def _download_file(client: S3Client, s3_path: str, local_path: str) -> None: + """Download the object at the given S3 path to the given local file system path + + Args: + client (S3Client): The S3 client + + s3_path (str): The S3 source path + + local_path (str): The local destination path + """ + dirname = os.path.dirname(local_path) + os.makedirs(dirname, exist_ok=True) + parsed_s3_path = parse_s3_path(s3_path) + client.download_file(parsed_s3_path[0], parsed_s3_path[1], local_path) + + +def maybe_download_file(s3_path: str, local_path: str) -> None: + """Download the object at the given S3 path to the given local file system path + + In a distributed setting, downloading the S3 object proceeds in stages in order + to try to have the minimum number of processes download the object in order for + all the ranks to have access to the downloaded object. + + Args: + s3_path (str): The S3 source path + + local_path (str): The local destination path + """ + + if get_real_group_size() > 1: + rank = get_real_rank() + local_rank = rank % get_real_group_size() + else: + rank = 0 + local_rank = 0 + + s3_client = boto3.client("s3") + + if (not os.path.exists(local_path)) and (rank == 0): + _download_file(s3_client, s3_path, local_path) + + if get_real_group_size() > 1: + barrier() + + # If the `local_path` is in a file system that is not + # shared across all the ranks, then we assume it's in the + # host file system and each host needs to download the file. + if (not os.path.exists(local_path)) and (local_rank == 0): + _download_file(s3_client, s3_path, local_path) + + if get_real_group_size() > 1: + barrier() + + # If the `local_path` still does not exist, then we assume + # each rank is saving to a separate location. + if not os.path.exists(local_path): + _download_file(s3_client, s3_path, local_path) + + if get_real_group_size() > 1: + barrier() + + assert os.path.exists(local_path)