moved current gpt2 dataloader to old

# See the License for the specific language governing permissions and

# limitations under the License.

"""GPT2 dataset."""

"""GPT2 style dataset."""

import json

import os

import numpy as np

import time

import numpy as np

import torch

from torch.utils.data import Dataset

from megatron import print_rank_0

from megatron import mpu

from megatron.data.bert_dataset import get_train_valid_test_split_

from megatron.data.indexed_dataset import make_dataset as make_indexed_dataset

class GPT2Dataset(Dataset):

    def __init__(self, data_path, sizes_filename, seq_length,

                 initial_seed, max_epochs=100):

        # Input parameters.

        self.data_path = data_path

        self.sizes_filename = sizes_filename

        self.seq_length = seq_length

        self.initial_seed = initial_seed

        self.max_epochs = max_epochs

        # Shard stuff.

        # Dictionary from shard nameto its size (number of element).

        self.master_shard_size_dict = None

        # Dictionary from shard name to modified size so it is

        # divisible by self.seq_length.

        self.shard_size_dict = None

        # Long array (self.max_epochs * num-shards) populated

        # randomly with shard names.

        self.shards_name = None

        # Start index of the data for a shard.

        self.shards_start_index = None

        self.build_shard_mappings_()

        self.data_length = self.shards_start_index[-1]

        # Data.

        self.shards_data = [None]*self.shards_name.size

        self.shards_sample_index = [None]*self.shards_name.size

def build_train_valid_test_datasets(data_prefix, data_impl, splits_string,

                                    train_valid_test_num_samples,

                                    seq_length, seed, skip_warmup):

    """Build train, valid, and test datasets."""

    # Indexed dataset.

    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)

    def build_dataset(index, name):

        dataset = None

        if splits[index + 1] > splits[index]:

            documents = np.arange(start=splits[index], stop=splits[index+1],

                                  step=1, dtype=np.int32)

            dataset = GPT2Dataset(name, data_prefix,

                                  documents, indexed_dataset,

                                  train_valid_test_num_samples[index],

                                  seq_length, seed)

        return dataset

    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 GPT2Dataset(torch.utils.data.Dataset):

    def __init__(self, name, data_prefix, documents, indexed_dataset,

                 num_samples, seq_length, seed):

        self.name = name

        self.indexed_dataset = indexed_dataset

        # Checks

        assert np.min(documents) >= 0

        assert np.max(documents) < indexed_dataset.sizes.shape[0]

        # Build index mappings.

        self.doc_idx, self.sample_idx, self.shuffle_idx = _build_index_mappings(

            self.name, data_prefix, documents, self.indexed_dataset.sizes,

            num_samples, seq_length, seed)

    def __len__(self):

        return self.data_length

        return self.sample_idx.shape[0]

    def __getitem__(self, idx):

        # Find which shard we need.

        shard_index = np.searchsorted(self.shards_start_index,

                                      idx, side='right') - 1

        # data index in the shard.

        data_idx = idx - self.shards_start_index[shard_index]

        # Load the shard if it is not in memory.

        if self.shards_data[shard_index] is None:

            print('global rank {} is building data for shard index {} ...'.

                  format(torch.distributed.get_rank(), shard_index))

            self.build_dataset_(shard_index)

        #assert self.shards_data[shard_index] is not None

        # Start index.

        start_index = self.shards_sample_index[shard_index][data_idx]

        # Add one for label shift.

        end_index = start_index + self.seq_length + 1

        data = self.shards_data[shard_index][start_index:end_index]

        return {'text': np.array(data, dtype=np.int64)}

    def build_dataset_(self, shard_index):

        # Garbage collect so we don't use a lot of memory.

        # Leave the last one in case other threads have not catche up yet.

        #for i in range(shard_index - 1):

        for i in range(shard_index):

            self.shards_data[i] = None

            self.shards_sample_index[i] = None

        # Read the shard.

        filename = os.path.join(self.data_path, self.shards_name[shard_index])

        print('loading {}'.format(filename))

        data = np.load(filename, allow_pickle=True)

        # Shuffle the data

        rng = np.random.RandomState(self.initial_seed + shard_index)

        rng.shuffle(data)

        # Flatten.

        data = np.hstack(data)

        size = (data.shape[0] - 1) // self.seq_length

        last_index = size * self.seq_length + 1

        data = data[0:last_index]

        self.shards_data[shard_index] = data

        indices = np.arange(size) * self.seq_length

        rng.shuffle(indices)

        self.shards_sample_index[shard_index] = indices

    def build_shard_mappings_(self):

        # Load the sizes file.

        sizes_filename = os.path.join(self.data_path, self.sizes_filename)

        if torch.distributed.get_rank() == 0:

            print(' > loading sizes from {}'.format(sizes_filename))

        with open(sizes_filename, 'r') as f:

            self.master_shard_size_dict = json.load(f)

        if torch.distributed.get_rank() == 0:

            print('   found {} shards'.format(len(self.master_shard_size_dict)))

        # Adjust sizes to be a multiple of seq_length.

        self.shard_size_dict = self.master_shard_size_dict.copy()

        total_samples = 0

        for shard in self.shard_size_dict:

            size = self.shard_size_dict[shard]

            size = ((size - 1) // self.seq_length) * self.seq_length

            total_samples += size // self.seq_length

            self.shard_size_dict[shard] = size

        # 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.

        if doc_index_f == doc_index_l:

            sample = 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.

            sample_list = [self.indexed_dataset.get(self.doc_idx[doc_index_f],

                                                    offset=offset_f)]

            # Loop over all in between documents and add the entire document.

            for i in range(doc_index_f+1, doc_index_l):

                sample_list.append(self.indexed_dataset.get(self.doc_idx[i]))

            # And finally add the relevant portion of last document.

            sample_list.append(self.indexed_dataset.get(

                self.doc_idx[doc_index_l],

                length=offset_l+1))

            sample = np.concatenate(sample_list)

        return {'text': np.array(sample, dtype=np.int64)}

def _build_index_mappings(name, data_prefix, documents, sizes,

                          num_samples, seq_length, seed):

    """doc-idx, sample-idx, and shuffle-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.

    _filename = data_prefix

    _filename += '_{}_indexmap'.format(name)

    _filename += '_{}ns'.format(num_samples)

    _filename += '_{}sl'.format(seq_length)

    _filename += '_{}s'.format(seed)

    doc_idx_filename = _filename + '_doc_idx.npy'

    sample_idx_filename = _filename + '_sample_idx.npy'

    shuffle_idx_filename = _filename + '_shuffle_idx.npy'

    # Build the indexed mapping if not exist.

    if torch.distributed.get_rank() == 0:

            print('   found {} samples in the dataset'.format(total_samples))

        # Build a list of shards.

        shards_ = np.sort(np.array(list(self.shard_size_dict.keys())))

        rng = np.random.RandomState(self.initial_seed)

        self.shards_name = np.copy(shards_)

        rng.shuffle(self.shards_name)

        for i in range(1, self.max_epochs):

            shards_c = np.copy(shards_)

            rng.shuffle(shards_c)

            self.shards_name = np.append(self.shards_name, shards_c)

        # Build the global indexing.

        self.shards_start_index = np.zeros(self.shards_name.size, dtype=np.int)

        self.shards_start_index[0] = 0

        for i in range(1, self.shards_name.size):

            shard = str(self.shards_name[i-1])

            size = self.shard_size_dict[shard]

            self.shards_start_index[i] = self.shards_start_index[i-1] + \

                                         size // self.seq_length

        if (not os.path.isfile(doc_idx_filename)) or \

           (not os.path.isfile(sample_idx_filename)) or \

           (not os.path.isfile(shuffle_idx_filename)):

            print_rank_0(' > WARNING: could not find index map files, building '

                         'the indices on rank 0 ...')

            # doc-idx.

            start_time = time.time()

            doc_idx = _build_doc_idx(documents, num_epochs, np_rng)

            np.save(doc_idx_filename, doc_idx, allow_pickle=True)

            print_rank_0(' > elasped 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.

            from megatron.data 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)

            #sample_idx = _build_sample_idx(sizes, doc_idx, seq_length,

            #                               num_epochs, tokens_per_epoch)

            np.save(sample_idx_filename, sample_idx, allow_pickle=True)

            print_rank_0(' > elasped time to build and save sample-idx mapping '

                         '(seconds): {:4f}'.format(time.time() - start_time))

            # shuffle-idx.

            start_time = time.time()

            shuffle_idx = _build_shuffle_idx(sample_idx.shape[0], np_rng)

            np.save(shuffle_idx_filename, shuffle_idx, allow_pickle=True)

            print_rank_0(' > elasped time to build and save shuffle-idx mapping'

                         ' (seconds): {:4f}'.format(time.time() - start_time))

    # This should be a barrier but nccl barrier assumes

    # device_index=rank which is not the case for model

    # parallel case

    counts = torch.cuda.LongTensor([1])

    torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())

    assert counts[0].item() == torch.distributed.get_world_size(

        group=mpu.get_data_parallel_group())

    # Load mappings.

    start_time = time.time()

    print_rank_0(' > loading doc-idx mapping from {}'.format(

        doc_idx_filename))

    doc_idx = np.load(doc_idx_filename, allow_pickle=True)

    print_rank_0(' > loading sample-idx mapping from {}'.format(

        sample_idx_filename))

    sample_idx = np.load(sample_idx_filename, allow_pickle=True)

    print_rank_0(' > loading shuffle-idx mapping from {}'.format(

        shuffle_idx_filename))

    shuffle_idx = np.load(shuffle_idx_filename, allow_pickle=True)

    print_rank_0('    loaded indexed file in {:3.3f} seconds'.format(

        time.time() - start_time))

    print_rank_0('    total number of samples: {}'.format(

        sample_idx.shape[0]))

    print_rank_0('    total number of epochs: {}'.format(num_epochs))

    return doc_idx, sample_idx, shuffle_idx

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 lenght, 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):

    """Build an array with length = number-of-epochs * number-of-dcuments.

    Each index is mapped to a corresponding document."""

    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

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 [..., 0] 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

    # Begining 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 begining 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(size, np_rng):

    """Build the range [0, size) and shuffle."""

    dtype_ = np.uint32

    if size >= (np.iinfo(np.uint32).max - 1):

        dtype_ = np.int64

    shuffle_idx = np.arange(start=0, stop=size, step=1, dtype=dtype_)

    np_rng.shuffle(shuffle_idx)

    return shuffle_idx

'''

class IndexedDataset:

    def __init__(self, num_docs, min_doc_length, max_doc_length, seq_length):

        self.seq_length = seq_length

        assert min_doc_length > 0

        self.tokens = []

        self.sizes = np.zeros(num_docs, dtype=np.int32)

        for i in range(num_docs):

            size = np.random.randint(low=min_doc_length, high=max_doc_length,

                                     size=1, dtype=np.uint32)[0]

            tokens_ = np.random.randint(low=1, high=60000,

                                        size=size, dtype=np.uint32)

            tokens_[-1] = 0

            self.sizes[i] = size

            self.tokens.append(tokens_)

        self.tokens_flat = None

    def get(self, doc_idx, offset=None, length=None):

        if length is None:

            if offset is None:

                return self.tokens[doc_idx]

            else:

                return self.tokens[doc_idx][offset:]

        if offset is None:

            return self.tokens[doc_idx][0:length]

        return self.tokens[doc_idx][offset:(offset+length)]

    def get_sample(self, index):

        start = index * self.seq_length

        end = start + self.seq_length + 1

        return self.tokens_flat[start:end]

    def build_tokens_flat(self, doc_idx):

        self.tokens_flat = np.concatenate([self.tokens[i] for i in doc_idx])

def test(seed, data_prefix, seq_length, num_samples,

         num_docs, min_doc_length, max_doc_length):

    print('testing for seed: {}, seq-length: {}, num-samples: {}, '

          'num-docs: {}, min-doc-length: {}, max-doc-length: {}'.format(

              seed, seq_length, num_samples,

              num_docs, min_doc_length, max_doc_length))

    np.random.seed(seed)

    indexed_dataset = IndexedDataset(num_docs, min_doc_length,

                                     max_doc_length, seq_length)

    indices = np.random.randint(indexed_dataset.sizes.shape[0]-2, size=2)

    documents = np.arange(np.min(indices), np.max(indices)+1)

    dataset = GPT2Dataset('gpt2', data_prefix, documents, indexed_dataset,

                          num_samples, seq_length, seed)

    print(' > number of epochs:', dataset.num_epochs)

    indexed_dataset.build_tokens_flat(dataset.doc_idx)

    for idx in range(num_samples):

        a = dataset[idx]

        b = indexed_dataset.get_sample(idx)

        assert np.sum(a - b) == 0

    print('passed')

if __name__ == '__main__':

    print('gpt2 dataset ...')

    import random

    data_prefix = 'junk/'

    for seed in range(1234, 1245):

        random.seed(seed)

        num_docs = random.randint(1, 999)

        min_doc_length = random.randint(1, 99)

        max_doc_length = random.randint(100, 9999)

        num_samples = random.randint(num_docs, 100*num_docs)

        seq_length = random.randint(min_doc_length, max_doc_length)

        test(seed, data_prefix, seq_length, num_samples,

             num_docs, min_doc_length, max_doc_length)

'''

from megatron import get_tokenizer

from megatron import mpu

from megatron import print_rank_0

from megatron.data.gpt2_dataset import GPT2Dataset

from megatron.data.gpt2_dataset import build_train_valid_test_datasets

from megatron.model import GPT2Model

from megatron.training import pretrain

from megatron.utils import get_ltor_masks_and_position_ids

    return loss, {'lm loss': reduced_loss[0]}

def make_gpt2_dataloaders():

    """Build gpt2 dataloders."""

    args = get_args()

    # Input parameters.

    input_data_sizes_file = args.input_data_sizes_file

    seq_length = args.seq_length

    initial_seed = args.seed

    # Build the datasets.

    def _build_dataset(name):

        return GPT2Dataset(os.path.join(args.data_path, name),

                           args.input_data_sizes_file,

                           args.seq_length, args.seed)

    train_ds = _build_dataset('train')

    valid_ds = _build_dataset('valid')

    test_ds = _build_dataset('test')

    # Dataloaders

    train = make_data_loader(train_ds)

    valid = make_data_loader(valid_ds)

    test = make_data_loader(test_ds)

    args.do_train = False

    args.do_valid = False

    args.do_test = False

    if train is not None:

        args.do_train = True

    if valid is not None:

        args.do_valid = True

    if test is not None:

        args.do_test = True

    return (train, valid, test)

def get_train_val_test_data():

    """Load the data on rank zero and boradcast number of tokens to all GPUS."""

    args = get_args()

    (train_data, val_data, test_data) = (None, None, None)

    (train_data, valid_data, test_data) = (None, None, None)

    # Data loader only on rank 0 of each model parallel group.

    if mpu.get_model_parallel_rank() == 0:

        (train_data, val_data, test_data) = make_gpt2_dataloaders()

        flags = torch.cuda.LongTensor([int(args.do_train),

                                       int(args.do_valid),

                                       int(args.do_test)])

        print_rank_0('> building train, validation, and test datasets '

                     'for GPT2 ...')

        data_parallel_size = mpu.get_data_parallel_world_size()

        data_parallel_rank = mpu.get_data_parallel_rank()

        global_batch_size = args.batch_size * data_parallel_size

        # Number of train/valid/test samples.

        train_iters = args.train_iters

        eval_iters = (train_iters // args.eval_interval + 1) * args.eval_iters

        test_iters = args.eval_iters

        train_val_test_num_samples = [train_iters * global_batch_size,

                                      eval_iters * global_batch_size,

                                      test_iters * global_batch_size]

        print_rank_0(' > datasets target sizes (minimum size):')

        print_rank_0('    train:      {}'.format(train_val_test_num_samples[0]))

        print_rank_0('    validation: {}'.format(train_val_test_num_samples[1]))

        print_rank_0('    test:       {}'.format(train_val_test_num_samples[2]))

        train_ds, valid_ds, test_ds = build_train_valid_test_datasets(

            data_prefix=args.data_path,

            data_impl=args.data_impl,

            splits_string=args.split,

            train_valid_test_num_samples=train_val_test_num_samples,

            seq_length=args.seq_length,

            seed=args.seed,

            skip_warmup=(not args.mmap_warmup))

        print_rank_0("> finished creating GPT2 datasets ...")

        train_data = make_data_loader(train_ds)

        valid_data = make_data_loader(valid_ds)

        test_data = make_data_loader(test_ds)

        do_train = train_data is not None and args.train_iters > 0

        do_valid = valid_data is not None and args.eval_iters > 0

        do_test = test_data is not None and args.eval_iters > 0

        # Need to broadcast num_tokens and num_type_tokens.

        flags = torch.cuda.LongTensor(

            [int(do_train), int(do_valid), int(do_test)])

    else:

        flags = torch.cuda.LongTensor([0, 0, 0])

    args.do_valid = flags[1].item()

    args.do_test = flags[2].item()

    return train_data, val_data, test_data

    return train_data, valid_data, test_data

if __name__ == "__main__":