Commit dd76c914 authored by Raul Puri's avatar Raul Puri
Browse files

gpt2 mmap sample map

parent 2b47c5e7

megatron/data/gpt2_dataset.py

+3 −3
Original line number Diff line number Diff line
@@ -212,13 +212,13 @@ def _build_index_mappings(name, data_prefix, documents, sizes, 
    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)

    doc_idx = np.load(doc_idx_filename, allow_pickle=True, mmap_mode='r')

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

        sample_idx_filename))

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

    sample_idx = np.load(sample_idx_filename, allow_pickle=True, mmap_mode='r')

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

        shuffle_idx_filename))

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

    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(