cerebras.modelzoo.data.common.h5_map_dataset.dataset.HDF5Dataset

class cerebras.modelzoo.data.common.h5_map_dataset.dataset.HDF5Dataset(*args, **kwargs)

Bases: torch.utils.data.Dataset

Dynamically read samples from disk for using mapping paradigms.

It supports two different data formats on disk. The first is data stored in an H5 file in the shape (num_tokens,), i.e. a series of documents tokenized and concatenated together. We call this format the ‘corpus’ format The second format is H5 data of shape (num_sequences, …), i.e. data has already been tokenized and split into sequences. We call this format the ‘sample’ format.

The corpus format supports flexible choice of MSL backed by a single copy of the data on disk. Both formats support deterministic restart, and a data order that is independent of the configuration of the cluster you are running on. I.e. you can pause a run, increase or decrease the number of systems you are running on, and restart the run with no change in data order.

When used in combination with shuffling, this implementation relies on random access reads to disk to dynamically split samples into sequences and shuffle. Users with unusually slow storage should look out for data loading bottlenecks and might consider using use_worker_cache=True if disk access is indeed a bottleneck.

Parameters

params (dict) –

a dictionary containing the following fields: - “data_dir” (str or list[str]): the path to the HDF5 files.

Exactly one of “data_dir” or “mixture” must be specified.

• ”batch_size” (int): batch size

• ”shuffle” (bool): whether or not to shuffle the dataset. Defaults

  to False

• ”shuffle_seed” (int): seed used for deterministic shuffling.

  Defaults to 0.

• ”use_worker_cache” (bool): whether or not to copy data to storage

  that is directly attached to each individual worker node. Useful when your network storage is unusually slow, but otherwise discouraged.

• ”max_sequence_length” (int): the sequence length of samples

  produced by the dataloader. When using the ‘corpus’ data format, the same preprocessed data will work with any max sequence length, so this may be set at runtime. When using the ‘sample’ format this must be set to None.

• ”data_subset” (str): an optional specification to only consider a

  subset of the full dataset, useful for sequence length scheduling and multi-epoch testing. Expected to be a comma separated list of ranges, e.g. ‘0.0-0.5’ or ‘0.1-0.3,0.7-1.0’. Specifying ‘0.0-0.5’ creates a dataset from the first half of the data on disk and disregards the second half.

• ”mixture” list[dict]: an optional specification of multiple

  datasets to mix over to create one single weighted combination. Each element must be a dictionary containing keys data_dir and weight. data_dir serves the same purpose as mentioned above. weight defines the probability with which this dataset should be sampled from. Weights are normalized to sum to 1. Optionally, the dictionary may also contain a data_subset field which functions the same as the data_subset argument above.

• ”drop_last” (bool): similar to the PyTorch drop_last setting

  except that samples that when set to True, samples that would have been dropped at the end of one epoch are yielded at the start of the next epoch so that there is no data loss. This is necessary for a data ordering that is independent of the distributed setup being used.

• ”num_samples” (int): the number of samples to shuffle over (if

  shuffling is enabled). In multi-epoch training, it is common to set this to the total number of samples that you plan to train on so that epochs are not sequential but instead shuffled together for potentially improved convergence.

• ”sort_files” (bool): whether or not the reader should sort the input

  files. This is included for backwards compatibility and should almost always be set to True.

• ”use_vsl” (bool): Flag to enable variable sequence length training.

  It requires the dataset to have two extra features: the attention_span of keys and the position_ids of tokens. Defaults to False.

• ”pad_last” (bool): Flag to enable padding of the last batch so

  that the last batch has the same batch size as the rest of the batches. Defaults to False.

Methods

generate_sample	Generates an empty tensor with the same shape and dtype as a sample from its dataset.
load_state_dict
map
state_dict

Attributes

by_sample

generate_sample()

Generates an empty tensor with the same shape and dtype as a sample from its dataset.

previous

cerebras.modelzoo.data.common.h5_map_dataset.dataset

next

cerebras.modelzoo.data.common.h5_map_dataset.dataset.MLMHDF5Dataset