wehub-resource-sync
712 lines
27 KiB
Python
712 lines
27 KiB
Python
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
|
|
|
|
"""GPT style dataset."""
|
|
import hashlib
|
|
import math
|
|
import os
|
|
import time
|
|
|
|
import numpy as np
|
|
import paddle
|
|
|
|
from paddlenlp.data.blendable_dataset import BlendableDataset
|
|
from paddlenlp.data.indexed_dataset import make_dataset as make_indexed_dataset
|
|
|
|
local_rank = int(os.getenv("PADDLE_RANK_IN_NODE", 0))
|
|
|
|
|
|
# class FakeHCG:
|
|
# def get_data_parallel_group(self):
|
|
# return None
|
|
|
|
# def get_pipe_parallel_group(self):
|
|
# return None
|
|
|
|
# def get_model_parallel_group(self):
|
|
# return None
|
|
|
|
|
|
def check_data_split(splits_string, do_train, do_eval, do_predict):
|
|
splits = []
|
|
if splits_string.find(",") != -1:
|
|
splits = [float(s) for s in splits_string.split(",")]
|
|
elif splits_string.find("/") != -1:
|
|
splits = [float(s) for s in splits_string.split("/")]
|
|
else:
|
|
splits = [float(splits_string)]
|
|
while len(splits) < 3:
|
|
splits.append(0.0)
|
|
splits = splits[:3]
|
|
splits_sum = sum(splits)
|
|
data_flag = True
|
|
assert splits_sum > 0.0, "sum of splits should larger than 0.0!"
|
|
if (do_train and splits[0] == 0) or (do_eval and splits[1] == 0) or (do_predict and splits[2] == 0):
|
|
data_flag = False
|
|
if not data_flag:
|
|
raise ValueError("If do_train/do_eval/do_predict is True, the corresponding dataset split should not be 0!")
|
|
|
|
|
|
def get_train_valid_test_split_(splits_string, size):
|
|
"""Get dataset splits from comma or '/' separated string list."""
|
|
|
|
splits = []
|
|
if splits_string.find(",") != -1:
|
|
splits = [float(s) for s in splits_string.split(",")]
|
|
elif splits_string.find("/") != -1:
|
|
splits = [float(s) for s in splits_string.split("/")]
|
|
else:
|
|
splits = [float(splits_string)]
|
|
while len(splits) < 3:
|
|
splits.append(0.0)
|
|
splits = splits[:3]
|
|
splits_sum = sum(splits)
|
|
assert splits_sum > 0.0
|
|
splits = [split / splits_sum for split in splits]
|
|
splits_index = [0]
|
|
for index, split in enumerate(splits):
|
|
splits_index.append(splits_index[index] + int(round(split * float(size))))
|
|
diff = splits_index[-1] - size
|
|
for index in range(1, len(splits_index)):
|
|
splits_index[index] -= diff
|
|
assert len(splits_index) == 4
|
|
assert splits_index[-1] == size
|
|
return splits_index
|
|
|
|
|
|
def get_datasets_weights_and_num_samples(data_prefix, train_val_test_num_samples):
|
|
|
|
# The data prefix should be in the format of:
|
|
# weight-1, data-prefix-1, weight-2, data-prefix-2, ..
|
|
assert len(data_prefix) % 2 == 0
|
|
num_datasets = len(data_prefix) // 2
|
|
weights = [0] * num_datasets
|
|
prefixes = [0] * num_datasets
|
|
for i in range(num_datasets):
|
|
weights[i] = float(data_prefix[2 * i])
|
|
prefixes[i] = (data_prefix[2 * i + 1]).strip()
|
|
# Normalize weights
|
|
weight_sum = 0.0
|
|
for weight in weights:
|
|
weight_sum += weight
|
|
assert weight_sum > 0.0
|
|
weights = [weight / weight_sum for weight in weights]
|
|
|
|
# Add 0.5% (the 1.005 factor) so in case the blending dataset does
|
|
# not uniformly distribute the number of samples, we still have
|
|
# samples left to feed to the network.
|
|
# (NOTE, yujun06): This is a workaround to avoid issues with indexing in the blending dataset. Therefore, we need to add 20 samples to each dataset.
|
|
datasets_train_valid_test_num_samples = []
|
|
for weight in weights:
|
|
datasets_train_valid_test_num_samples.append(
|
|
[int(math.ceil(val * weight * 1.005)) + 20 for val in train_val_test_num_samples]
|
|
)
|
|
|
|
return prefixes, weights, datasets_train_valid_test_num_samples
|
|
|
|
|
|
def print_rank_0(*args, **kwargs):
|
|
if paddle.distributed.get_rank() == 0:
|
|
print(*args, **kwargs)
|
|
|
|
|
|
def build_train_valid_test_datasets(
|
|
data_prefix,
|
|
data_impl,
|
|
splits_string,
|
|
train_val_test_num_samples,
|
|
seq_length,
|
|
seed,
|
|
skip_warmup,
|
|
train_data_prefix=None,
|
|
valid_data_prefix=None,
|
|
test_data_prefix=None,
|
|
return_doc_ids=False,
|
|
share_folder=False,
|
|
*,
|
|
data_cache_path=None,
|
|
need_data=True,
|
|
):
|
|
"""Build train, valid, and test datasets."""
|
|
|
|
# Single dataset.
|
|
if len(data_prefix) == 1:
|
|
return _build_train_valid_test_datasets(
|
|
data_prefix[0],
|
|
data_impl,
|
|
splits_string,
|
|
train_val_test_num_samples,
|
|
seq_length,
|
|
seed,
|
|
skip_warmup,
|
|
share_folder=share_folder,
|
|
data_cache_path=data_cache_path,
|
|
need_data=need_data,
|
|
)
|
|
|
|
# Blending dataset.
|
|
# Parse the values.
|
|
output = get_datasets_weights_and_num_samples(data_prefix, train_val_test_num_samples)
|
|
prefixes, weights, datasets_train_valid_test_num_samples = output
|
|
# NOTE: megatron/gpt_dataset.py has been updated. When creating BlendableDataset, we will use the raw train_val_test_num_samples instead of the expanded ones.
|
|
# Please refer to https://github.com/NVIDIA/NeMo/blob/72f630d087d45655b1a069dc72debf01dfdbdb2d/nemo/collections/nlp/data/language_modeling/megatron/gpt_dataset.py#L74-L80 for more information
|
|
train_num_samples, valid_num_samples, test_num_samples = train_val_test_num_samples
|
|
|
|
# Build individual datasets.
|
|
train_datasets = []
|
|
valid_datasets = []
|
|
test_datasets = []
|
|
for i in range(len(prefixes)):
|
|
train_ds, valid_ds, test_ds = _build_train_valid_test_datasets(
|
|
prefixes[i],
|
|
data_impl,
|
|
splits_string,
|
|
datasets_train_valid_test_num_samples[i],
|
|
seq_length,
|
|
seed,
|
|
skip_warmup,
|
|
return_doc_ids,
|
|
share_folder=share_folder,
|
|
data_cache_path=data_cache_path,
|
|
need_data=need_data,
|
|
)
|
|
if train_ds:
|
|
train_datasets.append(train_ds)
|
|
if valid_ds:
|
|
valid_datasets.append(valid_ds)
|
|
if test_ds:
|
|
test_datasets.append(test_ds)
|
|
|
|
blending_train_dataset = None
|
|
if train_datasets:
|
|
blending_train_dataset = BlendableDataset(
|
|
train_datasets, weights, train_num_samples, share_folder, data_cache_path=data_cache_path
|
|
)
|
|
blending_valid_dataset = None
|
|
if valid_datasets:
|
|
blending_valid_dataset = BlendableDataset(
|
|
valid_datasets, weights, valid_num_samples, share_folder, data_cache_path=data_cache_path
|
|
)
|
|
blending_test_dataset = None
|
|
if test_datasets:
|
|
blending_test_dataset = BlendableDataset(
|
|
test_datasets,
|
|
weights,
|
|
test_num_samples,
|
|
share_folder,
|
|
data_cache_path=data_cache_path,
|
|
)
|
|
|
|
return (blending_train_dataset, blending_valid_dataset, blending_test_dataset)
|
|
|
|
|
|
def _build_train_valid_test_datasets(
|
|
data_prefix,
|
|
data_impl,
|
|
splits_string,
|
|
train_val_test_num_samples,
|
|
seq_length,
|
|
seed,
|
|
skip_warmup,
|
|
return_doc_ids=False,
|
|
share_folder=False,
|
|
*,
|
|
data_cache_path=None,
|
|
need_data=True,
|
|
):
|
|
"""Build train, valid, and test datasets."""
|
|
|
|
# Indexed dataset.
|
|
if need_data:
|
|
indexed_dataset = get_indexed_dataset_(data_prefix, data_impl, skip_warmup)
|
|
|
|
total_num_of_documents = indexed_dataset.sizes.shape[0]
|
|
splits = get_train_valid_test_split_(splits_string, total_num_of_documents)
|
|
|
|
# Print stats about the splits.
|
|
print_rank_0(" > dataset split:")
|
|
|
|
def print_split_stats(name, index):
|
|
print_rank_0(" {}:".format(name))
|
|
print_rank_0(
|
|
" document indices in [{}, {}) total of {} "
|
|
"documents".format(splits[index], splits[index + 1], splits[index + 1] - splits[index])
|
|
)
|
|
|
|
print_split_stats("train", 0)
|
|
print_split_stats("validation", 1)
|
|
print_split_stats("test", 2)
|
|
|
|
if paddle.distributed.get_world_size() > 1:
|
|
paddle.distributed.barrier()
|
|
|
|
def build_dataset(index, name):
|
|
documents = np.arange(splits[index], splits[index + 1], 1, np.int32) if need_data else None
|
|
dataset = GPTDataset(
|
|
name,
|
|
data_prefix,
|
|
documents,
|
|
indexed_dataset if need_data else None,
|
|
splits_string,
|
|
train_val_test_num_samples[index],
|
|
seq_length,
|
|
seed,
|
|
return_doc_ids,
|
|
share_folder,
|
|
data_cache_path=data_cache_path,
|
|
need_data=need_data,
|
|
)
|
|
if need_data:
|
|
return dataset if splits[index + 1] > splits[index] else None
|
|
else:
|
|
return None
|
|
|
|
train_dataset = build_dataset(0, "train")
|
|
valid_dataset = build_dataset(1, "valid")
|
|
test_dataset = build_dataset(2, "test")
|
|
|
|
return (train_dataset, valid_dataset, test_dataset)
|
|
|
|
|
|
def get_indexed_dataset_(data_prefix, data_impl, skip_warmup):
|
|
"""Build indexed dataset."""
|
|
print_rank_0(" > building dataset index ...")
|
|
|
|
start_time = time.time()
|
|
indexed_dataset = make_indexed_dataset(data_prefix, data_impl, skip_warmup)
|
|
print_rank_0(" > finished creating indexed dataset in {:4f} " "seconds".format(time.time() - start_time))
|
|
print_rank_0(" number of documents: {}".format(indexed_dataset.sizes.shape[0]))
|
|
|
|
return indexed_dataset
|
|
|
|
|
|
class GPTDataset(paddle.io.Dataset):
|
|
def __init__(
|
|
self,
|
|
name,
|
|
data_prefix,
|
|
documents,
|
|
indexed_dataset,
|
|
splits_string,
|
|
num_samples,
|
|
seq_length,
|
|
seed,
|
|
return_doc_ids=False,
|
|
share_folder=False,
|
|
*,
|
|
data_cache_path=None,
|
|
need_data=True,
|
|
):
|
|
|
|
self.name = name
|
|
self.indexed_dataset = indexed_dataset
|
|
self.return_doc_ids = return_doc_ids
|
|
|
|
# Build index mappings.
|
|
if need_data and len(documents) > 0:
|
|
assert np.min(documents) >= 0
|
|
assert np.max(documents) < indexed_dataset.sizes.shape[0]
|
|
|
|
(
|
|
doc_idx_filename,
|
|
sample_idx_filename,
|
|
shuffle_idx_filename,
|
|
self.desc,
|
|
self.desc_hash,
|
|
num_epochs,
|
|
) = _build_index_mappings(
|
|
self.name,
|
|
data_prefix,
|
|
documents,
|
|
self.indexed_dataset.sizes,
|
|
splits_string,
|
|
num_samples,
|
|
seq_length,
|
|
seed,
|
|
share_folder,
|
|
data_cache_path=data_cache_path,
|
|
)
|
|
|
|
if paddle.distributed.get_world_size() > 1:
|
|
paddle.distributed.barrier()
|
|
|
|
# Load mappings.
|
|
if need_data and len(documents) > 0:
|
|
start_time = time.time()
|
|
print_rank_0(f" > loading doc-idx mapping from {doc_idx_filename}")
|
|
self.doc_idx = np.load(doc_idx_filename, allow_pickle=True, mmap_mode="r")
|
|
|
|
print_rank_0(f" > loading sample-idx mapping from {sample_idx_filename}")
|
|
self.sample_idx = np.load(sample_idx_filename, allow_pickle=True, mmap_mode="r")
|
|
|
|
print_rank_0(f" > loading shuffle-idx mapping from {shuffle_idx_filename}")
|
|
self.shuffle_idx = np.load(shuffle_idx_filename, allow_pickle=True, mmap_mode="r")
|
|
|
|
print_rank_0(" loaded indexed file in {:3.3f} seconds".format(time.time() - start_time))
|
|
print_rank_0(" total number of samples: {}".format(self.sample_idx.shape[0]))
|
|
print_rank_0(" total number of epochs: {}".format(num_epochs))
|
|
|
|
if paddle.distributed.get_world_size() > 1:
|
|
paddle.distributed.barrier()
|
|
|
|
def __len__(self):
|
|
# -1 is due to data structure used to retrieve the index:
|
|
# sample i --> [sample_idx[i], sample_idx[i+1])
|
|
return self.sample_idx.shape[0] - 1
|
|
|
|
def __getitem__(self, idx):
|
|
# Get the shuffled index.
|
|
idx = self.shuffle_idx[idx]
|
|
# Start and end documents and offsets.
|
|
doc_index_f = self.sample_idx[idx][0]
|
|
doc_index_l = self.sample_idx[idx + 1][0]
|
|
offset_f = self.sample_idx[idx][1]
|
|
offset_l = self.sample_idx[idx + 1][1]
|
|
# If we are within the same document, just extract the chunk.
|
|
doc_ids = []
|
|
if doc_index_f == doc_index_l:
|
|
doc_ids.append(self.doc_idx[doc_index_f])
|
|
|
|
sample, mask = self.indexed_dataset.get(
|
|
self.doc_idx[doc_index_f], offset=offset_f, length=offset_l - offset_f + 1
|
|
)
|
|
else:
|
|
# Otherwise, get the rest of the initial document.
|
|
doc_ids.append(self.doc_idx[doc_index_f])
|
|
sample, mask = self.indexed_dataset.get(self.doc_idx[doc_index_f], offset=offset_f)
|
|
append_mask = True
|
|
if mask is None:
|
|
append_mask = False
|
|
|
|
sample_list = [sample]
|
|
mask_list = []
|
|
mask_list = [mask]
|
|
# Loop over all in between documents and add the entire document.
|
|
for i in range(doc_index_f + 1, doc_index_l):
|
|
doc_ids.append(self.doc_idx[i])
|
|
sample, mask = self.indexed_dataset.get(self.doc_idx[i])
|
|
sample_list.append(sample)
|
|
if append_mask:
|
|
mask_list.append(mask)
|
|
|
|
# And finally add the relevant portion of last document.
|
|
doc_ids.append(self.doc_idx[doc_index_l])
|
|
sample, mask = self.indexed_dataset.get(self.doc_idx[doc_index_l], length=offset_l + 1)
|
|
sample_list.append(sample)
|
|
if append_mask:
|
|
mask_list.append(mask)
|
|
sample = np.concatenate(sample_list)
|
|
if append_mask:
|
|
mask = np.concatenate(mask_list)
|
|
# print(sample)
|
|
if self.return_doc_ids: # for retro preprocessing
|
|
if mask is None:
|
|
return {"text": np.array(sample, dtype=np.int64), "doc_ids": np.array(doc_ids, dtype=np.int64)}
|
|
else:
|
|
return {
|
|
"text": np.array(sample, dtype=np.int64),
|
|
"doc_ids": np.array(doc_ids, dtype=np.int64),
|
|
"mask": np.array(mask, dtype=np.int64),
|
|
}
|
|
else:
|
|
if mask is None:
|
|
return {"text": np.array(sample, dtype=np.int64)}
|
|
else:
|
|
return {"text": np.array(sample, dtype=np.int64), "mask": np.array(mask, dtype=np.int64)}
|
|
|
|
|
|
def _build_index_mappings(
|
|
name, data_prefix, documents, sizes, splits_string, num_samples, seq_length, seed, share_folder, *, data_cache_path
|
|
):
|
|
"""Build doc-idx, sample-idx, and shuffle-idx.
|
|
doc-idx: is an array (ordered) of documents to be used in training.
|
|
sample-idx: is the start document index and document offset for each
|
|
training sample.
|
|
shuffle-idx: maps the sample index into a random index into sample-idx.
|
|
"""
|
|
|
|
# Number of tokens in each epoch and number of required epochs.
|
|
tokens_per_epoch = _num_tokens(documents, sizes)
|
|
num_epochs = _num_epochs(tokens_per_epoch, seq_length, num_samples)
|
|
|
|
# rng state
|
|
np_rng = np.random.RandomState(seed=seed)
|
|
# Filename of the index mappings.
|
|
desc = "GPT Dataset\n\n"
|
|
desc += f"Data prefix {data_prefix}\n"
|
|
desc += f"Dataset name {name}\n"
|
|
desc += f"Number of samples {num_samples}\n"
|
|
desc += f"Sequence length {seq_length}\n"
|
|
desc += f"Random seed {seed}\n"
|
|
desc += f"Split {splits_string}\n"
|
|
desc_hash = hashlib.md5(desc.encode("utf-8")).hexdigest()
|
|
desc_filename = desc_hash + ".dsc"
|
|
doc_idx_filename = desc_hash + "_doc_idx.npy"
|
|
sample_idx_filename = desc_hash + "_sample_idx.npy"
|
|
shuffle_idx_filename = desc_hash + "_shuffle_idx.npy"
|
|
|
|
# Look for cache in main data dir first to avoid unnecessary
|
|
# duplication, then look in data-cache-path if specified,
|
|
# If nothing is found, use the last path looked in
|
|
build_indices = True
|
|
prefixes = [os.path.join(os.path.dirname(data_prefix), "index-cache")]
|
|
if data_cache_path is not None:
|
|
prefixes.append(data_cache_path)
|
|
for prefix in prefixes:
|
|
idx_path = {
|
|
"desc": os.path.join(prefix, desc_filename),
|
|
"doc": os.path.join(prefix, doc_idx_filename),
|
|
"sample": os.path.join(prefix, sample_idx_filename),
|
|
"shuffle": os.path.join(prefix, shuffle_idx_filename),
|
|
}
|
|
for f in idx_path.values():
|
|
if not os.path.isfile(f):
|
|
break
|
|
else:
|
|
# Found our files!
|
|
build_indices = False
|
|
break
|
|
data_cache_dir = os.path.dirname(idx_path["desc"])
|
|
# data_cache_success = True
|
|
# Build the indexed mapping if not exist.
|
|
check_rank_flag = build_indices and local_rank == 0
|
|
if share_folder:
|
|
check_rank_flag = build_indices and paddle.distributed.get_rank() == 0
|
|
|
|
# if build_indices and paddle.distributed.get_rank() == 0:
|
|
|
|
print(
|
|
f"searching for causal dataset, build_indices={build_indices}, share_folder {share_folder}, check_rank_flag {check_rank_flag}",
|
|
flush=True,
|
|
)
|
|
if check_rank_flag:
|
|
print_rank_0(" > WARNING: could not find index map files, building " "the indices on rank 0 ...")
|
|
|
|
# For the last epoch, decide whether include the entire epoch
|
|
# in the global shuffle or not.
|
|
|
|
# If we need only one epoch, then separating last epoch does
|
|
# not mean anything.
|
|
if num_epochs == 1:
|
|
separate_last_epoch = False
|
|
print(" > only one epoch required, setting " "separate_last_epoch to False", flush=True)
|
|
|
|
else:
|
|
# Get the number of samples for the last epoch
|
|
num_samples_from_epochs_minus_one = ((num_epochs - 1) * tokens_per_epoch - 1) // seq_length
|
|
last_epoch_num_samples = num_samples - num_samples_from_epochs_minus_one
|
|
assert last_epoch_num_samples >= 0, "last epoch number of samples should be non-negative."
|
|
num_samples_per_epoch = (tokens_per_epoch - 1) // seq_length
|
|
assert last_epoch_num_samples <= (
|
|
num_samples_per_epoch + 1
|
|
), "last epoch number of samples exceeded max value."
|
|
# If we have less than 80% of the samples for the last epoch,
|
|
# separate out the epoch and treat it differently.
|
|
# Note: the 80% number is just based on common sense and can
|
|
# be adjusted if needed.
|
|
separate_last_epoch = last_epoch_num_samples < int(0.80 * num_samples_per_epoch)
|
|
if separate_last_epoch:
|
|
string = (
|
|
" > last epoch number of samples ({}) is smaller "
|
|
"than 80% of number of samples per epoch ({}), "
|
|
"setting separate_last_epoch to True"
|
|
)
|
|
else:
|
|
string = (
|
|
" > last epoch number of samples ({}) is larger "
|
|
"than 80% of number of samples per epoch ({}), "
|
|
"setting separate_last_epoch to False"
|
|
)
|
|
print(string.format(last_epoch_num_samples, num_samples_per_epoch), flush=True)
|
|
|
|
try:
|
|
os.makedirs(data_cache_dir, exist_ok=True)
|
|
|
|
# description
|
|
with open(idx_path["desc"], "wt") as fd:
|
|
fd.write(desc)
|
|
|
|
# doc-idx.
|
|
start_time = time.time()
|
|
doc_idx = _build_doc_idx(documents, num_epochs, np_rng, separate_last_epoch)
|
|
np.save(idx_path["doc"], doc_idx, allow_pickle=True)
|
|
print_rank_0(
|
|
" > elapsed time to build and save doc-idx mapping "
|
|
"(seconds): {:4f}".format(time.time() - start_time)
|
|
)
|
|
# sample-idx.
|
|
start_time = time.time()
|
|
# Use C++ implementation for speed.
|
|
# First compile and then import.
|
|
# from megatron.data import helpers
|
|
from fast_dataindex import helpers
|
|
|
|
assert doc_idx.dtype == np.int32
|
|
assert sizes.dtype == np.int32
|
|
sample_idx = helpers.build_sample_idx(sizes, doc_idx, seq_length, num_epochs, tokens_per_epoch)
|
|
np.save(idx_path["sample"], sample_idx, allow_pickle=True)
|
|
print_rank_0(
|
|
" > elapsed time to build and save sample-idx mapping "
|
|
"(seconds): {:4f}".format(time.time() - start_time)
|
|
)
|
|
# shuffle-idx.
|
|
start_time = time.time()
|
|
# -1 is due to data structure used to retrieve the index:
|
|
# sample i --> [sample_idx[i], sample_idx[i+1])
|
|
if separate_last_epoch:
|
|
num_samples_ = num_samples_from_epochs_minus_one
|
|
else:
|
|
num_samples_ = sample_idx.shape[0] - 1
|
|
shuffle_idx = _build_shuffle_idx(num_samples_, sample_idx.shape[0] - 1, np_rng)
|
|
np.save(idx_path["shuffle"], shuffle_idx, allow_pickle=True)
|
|
print_rank_0(
|
|
" > elapsed time to build and save shuffle-idx mapping"
|
|
" (seconds): {:4f}".format(time.time() - start_time)
|
|
)
|
|
except OSError:
|
|
print(f"There was an error trying to create the data cache directory ({data_cache_dir})")
|
|
print('or a file in it. This defaults to a directory "index-cache" within the directory')
|
|
print("the data files are in and can be set with the --data-cache-path argument. Please")
|
|
print("ensure you have write access to this directory or specify one that you do have")
|
|
print("write access to.")
|
|
# data_cache_success = False
|
|
else:
|
|
while True:
|
|
if (
|
|
(not os.path.isfile(idx_path["doc"]))
|
|
or (not os.path.isfile(idx_path["sample"]))
|
|
or (not os.path.isfile(idx_path["shuffle"]))
|
|
):
|
|
print("building indices on rank 0 ...", flush=True)
|
|
time.sleep(3)
|
|
else:
|
|
try:
|
|
np.load(idx_path["shuffle"], allow_pickle=True, mmap_mode="r")
|
|
print("build success", flush=True)
|
|
break
|
|
except Exception:
|
|
print("%s file is still writing or damaged, please wait for a moment." % idx_path["shuffle"])
|
|
time.sleep(3)
|
|
# try:
|
|
# hcg = paddle.distributed.fleet.get_hybrid_communicate_group()
|
|
# except:
|
|
# hcg = FakeHCG()
|
|
|
|
# counts = paddle.to_tensor([data_cache_success], dtype="int64")
|
|
# paddle.distributed.all_reduce(counts, group=hcg.get_data_parallel_group())
|
|
# paddle.distributed.all_reduce(counts, group=hcg.get_pipe_parallel_group())
|
|
# if counts[0].item() != (
|
|
# paddle.distributed.get_world_size() // paddle.distributed.get_world_size(group=hcg.get_model_parallel_group())
|
|
# ):
|
|
# print_rank_0("Data index creation unsuccessful, exiting.")
|
|
# exit()
|
|
# paddle.distributed.barrier()
|
|
|
|
return idx_path["doc"], idx_path["sample"], idx_path["shuffle"], desc, desc_hash, num_epochs
|
|
|
|
|
|
def _num_tokens(documents, sizes):
|
|
"""Total number of tokens in the dataset."""
|
|
return np.sum(sizes[documents])
|
|
|
|
|
|
def _num_epochs(tokens_per_epoch, seq_length, num_samples):
|
|
"""Based on number of samples and sequence length, calculate how many
|
|
epochs will be needed."""
|
|
num_epochs = 0
|
|
total_tokens = 0
|
|
while True:
|
|
num_epochs += 1
|
|
total_tokens += tokens_per_epoch
|
|
# -1 is because we need to retrieve seq_length + 1 token each time
|
|
# but the last token will overlap with the first token of the next
|
|
# sample except for the last sample.
|
|
if ((total_tokens - 1) // seq_length) >= num_samples:
|
|
return num_epochs
|
|
|
|
|
|
def _build_doc_idx(documents, num_epochs, np_rng, separate_last_epoch):
|
|
"""Build an array with length = number-of-epochs * number-of-documents.
|
|
Each index is mapped to a corresponding document."""
|
|
if not separate_last_epoch or num_epochs == 1:
|
|
doc_idx = np.mgrid[0:num_epochs, 0 : len(documents)][1]
|
|
doc_idx[:] = documents
|
|
doc_idx = doc_idx.reshape(-1)
|
|
doc_idx = doc_idx.astype(np.int32)
|
|
np_rng.shuffle(doc_idx)
|
|
return doc_idx
|
|
|
|
doc_idx_first = _build_doc_idx(documents, num_epochs - 1, np_rng, False)
|
|
doc_idx_last = _build_doc_idx(documents, 1, np_rng, False)
|
|
return np.concatenate((doc_idx_first, doc_idx_last))
|
|
|
|
|
|
def _build_sample_idx(sizes, doc_idx, seq_length, num_epochs, tokens_per_epoch):
|
|
"""Sample index mapping 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."""
|
|
|
|
# Total number of samples. For -1 see comments in `_num_epochs`.
|
|
num_samples = (num_epochs * tokens_per_epoch - 1) // seq_length
|
|
sample_idx = np.zeros([num_samples + 1, 2], dtype=np.int32)
|
|
|
|
# Index into sample_idx.
|
|
sample_index = 0
|
|
# Index into doc_idx.
|
|
doc_idx_index = 0
|
|
# Beginning offset for each document.
|
|
doc_offset = 0
|
|
# Start with first document and no offset.
|
|
sample_idx[sample_index][0] = doc_idx_index
|
|
sample_idx[sample_index][1] = doc_offset
|
|
sample_index += 1
|
|
while sample_index <= num_samples:
|
|
# Start with a fresh sequence.
|
|
remaining_seq_length = seq_length + 1
|
|
while remaining_seq_length != 0:
|
|
# Get the document length.
|
|
doc_id = doc_idx[doc_idx_index]
|
|
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 - 1
|
|
remaining_seq_length = 0
|
|
else:
|
|
# Otherwise, start from the beginning of the next document.
|
|
doc_idx_index += 1
|
|
doc_offset = 0
|
|
# Record the sequence.
|
|
sample_idx[sample_index][0] = doc_idx_index
|
|
sample_idx[sample_index][1] = doc_offset
|
|
sample_index += 1
|
|
|
|
return sample_idx
|
|
|
|
|
|
def _build_shuffle_idx(num_samples, total_size, np_rng):
|
|
"""Build the range [0, size) and shuffle."""
|
|
print(
|
|
" > building shuffle index with split [0, {}) and [{}, {}) "
|
|
"...".format(num_samples, num_samples, total_size),
|
|
flush=True,
|
|
)
|
|
|
|
dtype_ = np.uint32
|
|
if total_size >= (np.iinfo(np.uint32).max - 1):
|
|
dtype_ = np.int64
|
|
|
|
shuffle_idx_first = np.arange(start=0, stop=num_samples, step=1, dtype=dtype_)
|
|
np_rng.shuffle(shuffle_idx_first)
|
|
if num_samples == total_size:
|
|
return shuffle_idx_first
|
|
|
|
shuffle_idx_last = np.arange(start=num_samples, stop=total_size, step=1, dtype=dtype_)
|
|
np_rng.shuffle(shuffle_idx_last)
|
|
|
|
return np.concatenate((shuffle_idx_first, shuffle_idx_last))
|