Merge staging-realm into realm-mlm

megatron/data/indexed_dataset.py

@@ -42,6 +42,7 @@ def infer_dataset_impl(path):
             else:
                 return None
     else:
+        print(f"Dataset path does not exist: {path}")
         return None
 
 
@@ -61,6 +62,7 @@ def make_dataset(path, impl, skip_warmup=False):
         return IndexedCachedDataset(path)
     elif impl == 'mmap' and MMapIndexedDataset.exists(path):
         return MMapIndexedDataset(path, skip_warmup)
+    print(f"Unknown dataset implementation: {impl}")
     return None
 
 
@@ -466,9 +468,8 @@ class MMapIndexedDataset(torch.utils.data.Dataset):
     def __getitem__(self, idx):
         if isinstance(idx, int):
             ptr, size = self._index[idx]
-            np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=size, offset=ptr)
-            if self._index.dtype != np.int64:
-                np_array = np_array.astype(np.int64)
+            np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                     count=size, offset=ptr)
             return np_array
         elif isinstance(idx, slice):
             start, stop, step = idx.indices(len(self))
@@ -478,10 +479,25 @@ class MMapIndexedDataset(torch.utils.data.Dataset):
             sizes = self._index._sizes[idx]
             offsets = list(accumulate(sizes))
             total_size = sum(sizes)
-            np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=total_size, offset=ptr)
+            np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                     count=total_size, offset=ptr)
             sents = np.split(np_array, offsets[:-1])
             return sents
 
+    def get(self, idx, offset=0, length=None):
+        """ Retrieves 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.
+        """
+        ptr, size = self._index[idx]
+        if length is None:
+            length = size - offset
+        ptr += offset * np.dtype(self._index.dtype).itemsize
+        np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype,
+                                 count=length, offset=ptr)
+        return np_array
+
     @property
     def sizes(self):
         return self._index.sizes

megatron/data/test/test_indexed_dataset.py

+# This file isn't really a formal automated test, it's just a place to
+# put some code used during development and manual testing of
+# indexed_dataset.
+
 import argparse
 import os
 import sys
@@ -7,52 +11,90 @@ import torch
 script_dir = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.join(script_dir, "../../../"))
 
-from megatron.data import indexed_dataset, FullBertTokenizer, AlbertDataset
+from megatron.tokenizer import build_tokenizer
+from megatron.data import indexed_dataset
 
 def test_indexed_dataset(args):
     ds = indexed_dataset.make_dataset(args.data, args.dataset_impl)
-    tokenizer = FullBertTokenizer(args.vocab, do_lower_case=True)
+    tokenizer = build_tokenizer(args)
     print(len(ds.doc_idx))
     print(len(ds))
     print(ds.doc_idx[-1])
     if ds.supports_prefetch:
         # just prefetch the whole thing in test (so assume it is small)
         ds.prefetch(range(len(ds)))
-    for i in range(len(ds.doc_idx)-1):
+    if args.count > len(ds.doc_idx)-1:
+        args.count = len(ds.doc_idx)-1
+
+    for i in range(args.count):
         start = ds.doc_idx[i]
         end = ds.doc_idx[i+1]
         ids = ds[start:end]
+        print(f"Document {i}:")
+        print("--------------")
         for s in ids:
             assert len(s) > 0
             l = s.data.tolist()
-            tokens = tokenizer.convert_ids_to_tokens(l)
-            for t in tokens:
-                if '\n' in t:
-                    print("Newline in string!")
-        print(i)
-
-def test_albert_dataset(args):
-    # tokenizer = FullBertTokenizer(args.vocab, do_lower_case=True)
-    # idataset = indexed_dataset.make_dataset(args.data, args.dataset_impl)
-    # ds = AlbertDataset(idataset, tokenizer)
-    ds = AlbertDataset.from_paths(args.vocab, args.data, args.dataset_impl,
-                                  args.epochs, args.max_num_samples,
-                                  args.masked_lm_prob, args.seq_length,
-                                  args.short_seq_prob, args.seed)
-    truncated = 0
-    total = 0
-    for s in ds:
-        ids = s['text']
-        tokens = ds.tokenizer.convert_ids_to_tokens(ids)
-        print(tokens)
-        exit()
+            text = tokenizer.detokenize(l)
+            print(text)
+            print("---")
+
+def test_indexed_dataset_get(args):
+    ds = indexed_dataset.make_dataset(args.data, args.dataset_impl)
+    tokenizer = build_tokenizer(args)
+    size = ds.sizes[0]
+    print(f"size: {size}")
+    full = ds.get(0)
+    print(full)
+    #print(tokenizer.detokenize(full.data.tolist()))
+    print("---")
+    end = ds.get(0, offset=size-10)
+    print(end)
+    #print(tokenizer.detokenize(end.data.tolist()))
+
+    start = ds.get(0, length=10)
+    print(start)
+    #print(tokenizer.detokenize(start.data.tolist()))
+
+    part = ds.get(0, offset=2, length=8)
+    print(part)
+    #print(tokenizer.detokenize(part.data.tolist()))
+
+# def test_albert_dataset(args):
+#     # tokenizer = FullBertTokenizer(args.vocab, do_lower_case=True)
+#     # idataset = indexed_dataset.make_dataset(args.data, args.dataset_impl)
+#     # ds = AlbertDataset(idataset, tokenizer)
+#     ds = AlbertDataset.from_paths(args.vocab, args.data, args.dataset_impl,
+#                                   args.epochs, args.max_num_samples,
+#                                   args.masked_lm_prob, args.seq_length,
+#                                   args.short_seq_prob, args.seed)
+#     truncated = 0
+#     total = 0
+#     for i, s in enumerate(ds):
+#         ids = s['text']
+#         tokens = ds.tokenizer.convert_ids_to_tokens(ids)
+#         print(tokens)
+#         if i >= args.count-1:
+#             exit()
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--data', type=str, help='prefix to data files')
-    parser.add_argument('--vocab', type=str, help='Path to vocab.txt')
     parser.add_argument('--dataset-impl', type=str, default='infer',
                         choices=['lazy', 'cached', 'mmap', 'infer'])
+    parser.add_argument('--count', type=int, default=10,
+                        help='Number of samples/documents to print')
+
+    group = parser.add_argument_group(title='tokenizer')
+    group.add_argument('--tokenizer-type', type=str, required=True,
+                       choices=['BertWordPieceLowerCase',
+                                'GPT2BPETokenizer'],
+                       help='What type of tokenizer to use.')
+    group.add_argument('--vocab-file', type=str, default=None,
+                       help='Path to the vocab file')
+    group.add_argument('--merge-file', type=str, default=None,
+                       help='Path to the BPE merge file (if necessary).')
+
     parser.add_argument('--epochs', type=int, default=5,
                         help='Number of epochs to plan for')
     parser.add_argument('--max-num-samples', type=int, default=None,
@@ -66,12 +108,15 @@ def main():
     parser.add_argument('--seed', type=int, default=1234,
                         help='random seed')
     args = parser.parse_args()
+    args.rank = 0
+    args.make_vocab_size_divisible_by = 128
+    args.model_parallel_size = 1
 
     if args.dataset_impl == "infer":
         args.dataset_impl = indexed_dataset.infer_dataset_impl(args.data)
 
-    test_albert_dataset(args)
-#    test_indexed_dataset(args)
+#    test_albert_dataset(args)
+    test_indexed_dataset_get(args)
 
 if __name__ == "__main__":
     main()

megatron/deprecated_data_utils/datasets.py

@@ -918,10 +918,10 @@ class InverseClozeDataset(data.Dataset):
 
         sample = {
             'input_text': np.array(input_tokens),
-            'input_types': np.array(input_token_types),
+            'query_types': np.array(input_token_types),
             'input_pad_mask': np.array(input_pad_mask),
             'context_text': np.array(context_tokens),
-            'context_types': np.array(context_token_types),
+            'block_types': np.array(context_token_types),
             'context_pad_mask': np.array(context_pad_mask)
         }
 

megatron/global_vars.py

@@ -65,7 +65,7 @@ def set_global_variables(extra_args_provider=None, args_defaults={}):
     """Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers."""
     args = _parse_args(extra_args_provider=extra_args_provider,
                        defaults=args_defaults)
-    _build_tokenizer(args)
+    _ = _build_tokenizer(args)
     _set_tensorboard_writer(args)
     _set_adlr_autoresume(args)
     _set_timers()
@@ -85,6 +85,13 @@ def _build_tokenizer(args):
     global _GLOBAL_TOKENIZER
     _ensure_var_is_not_initialized(_GLOBAL_TOKENIZER, 'tokenizer')
     _GLOBAL_TOKENIZER = build_tokenizer(args)
+    return _GLOBAL_TOKENIZER
+
+
+def rebuild_tokenizer(args):
+    global _GLOBAL_TOKENIZER
+    _GLOBAL_TOKENIZER = None
+    return _build_tokenizer(args)
 
 
 def _set_tensorboard_writer(args):

megatron/initialize.py

@@ -59,6 +59,7 @@ def _initialize_distributed():
     """Initialize torch.distributed and mpu."""
     args = get_args()
 
+    device_count = torch.cuda.device_count()
     if torch.distributed.is_initialized():
 
         if args.rank == 0:
@@ -66,8 +67,9 @@ def _initialize_distributed():
                   'skipping initialization ...', flush=True)
         args.rank = torch.distributed.get_rank()
         args.world_size = torch.distributed.get_world_size()
+        if device_count > 0:
             device = torch.cuda.current_device()
-        local_rank = args.rank % torch.cuda.device_count()
+            local_rank = args.rank % device_count
             assert local_rank == device, \
                 'expected local-rank to be the same as rank % device-count.'
 
@@ -76,7 +78,8 @@ def _initialize_distributed():
         if args.rank == 0:
             print('> initializing torch distributed ...', flush=True)
         # Manually set the device ids.
-        device = args.rank % torch.cuda.device_count()
+        if device_count > 0:
+            device = args.rank % device_count
             if args.local_rank is not None:
                 assert args.local_rank == device, \
                     'expected local-rank to be the same as rank % device-count.'
@@ -94,6 +97,7 @@ def _initialize_distributed():
             init_method=init_method)
 
     # Set the model-parallel / data-parallel communicators.
+    if device_count > 0:
         mpu.initialize_model_parallel(args.model_parallel_size)
 
 
@@ -112,6 +116,7 @@ def _set_random_seed(seed):
         random.seed(seed)
         np.random.seed(seed)
         torch.manual_seed(seed)
+        if torch.cuda.device_count() > 0:
             mpu.model_parallel_cuda_manual_seed(seed)
     else:
         raise ValueError('Seed ({}) should be a positive integer.'.format(seed))

megatron/model/bert_model.py

@@ -219,6 +219,7 @@ class BertModel(MegatronModule):
 
 
 class ICTBertModel(MegatronModule):
+    """Bert-based module for Inverse Cloze task."""
     def __init__(self,
                  ict_head_size,
                  num_tokentypes=2,
@@ -231,41 +232,38 @@ class ICTBertModel(MegatronModule):
             parallel_output=parallel_output
         )
 
-        self.question_model = BertModel(**bert_args)
-        self._question_key = 'question_model'
-        self.context_model = BertModel(**bert_args)
-        self._context_key = 'context_model'
+        # this model embeds (pseudo-)queries - Embed_input in the paper
+        self.query_model = BertModel(**bert_args)
+        self._query_key = 'question_model'
 
-    def forward(self, input_tokens, input_attention_mask, input_types,
-                context_tokens, context_attention_mask, context_types, return_logits=False):
+        # this model embeds evidence blocks - Embed_doc in the paper
+        self.block_model = BertModel(**bert_args)
+        self._block_key = 'context_model'
 
-        question_ict_logits, _ = self.question_model.forward(input_tokens, 1 - input_attention_mask, input_types)
-        context_ict_logits, _ = self.context_model.forward(context_tokens, 1 - context_attention_mask, context_types)
+    def forward(self, query_tokens, query_attention_mask, query_types,
+                block_tokens, block_attention_mask, block_types):
+        """Run a forward pass for each of the models and compute the similarity scores."""
 
-        # [batch x h] * [h x batch]
-        retrieval_scores = question_ict_logits.matmul(torch.transpose(context_ict_logits, 0, 1))
-
-        if return_logits:
-            return question_ict_logits, context_ict_logits, retrieval_scores
-
-        return retrieval_scores
+        query_logits, _ = self.query_model.forward(query_tokens, 1 - query_attention_mask, query_types)
+        block_logits, _ = self.block_model.forward(block_tokens, 1 - block_attention_mask, block_types)
 
+        return query_logits, block_logits
 
     def state_dict_for_save_checkpoint(self, destination=None, prefix='',
                                        keep_vars=False):
+        """Save dict with state dicts of each of the models."""
         state_dict_ = {}
-        state_dict_[self._question_key] \
-            = self.question_model.state_dict_for_save_checkpoint(
+        state_dict_[self._query_key] \
+            = self.query_model.state_dict_for_save_checkpoint(
             destination, prefix, keep_vars)
-        state_dict_[self._context_key] \
-            = self.context_model.state_dict_for_save_checkpoint(
+        state_dict_[self._block_key] \
+            = self.block_model.state_dict_for_save_checkpoint(
             destination, prefix, keep_vars)
         return state_dict_
 
     def load_state_dict(self, state_dict, strict=True):
-        """Customized load."""
-
-        self.question_model.load_state_dict(
-            state_dict[self._question_key], strict=strict)
-        self.context_model.load_state_dict(
-            state_dict[self._context_key], strict=strict)
+        """Load the state dicts of each of the models"""
+        self.query_model.load_state_dict(
+            state_dict[self._query_key], strict=strict)
+        self.block_model.load_state_dict(
+            state_dict[self._block_key], strict=strict)

megatron/model/transformer.py

@@ -102,6 +102,7 @@ class ParallelSelfAttention(MegatronModule):
                  output_layer_init_method, layer_number):
         super(ParallelSelfAttention, self).__init__()
         args = get_args()
+        self.fp16 = args.fp16
 
         self.attention_mask_func = attention_mask_func
         self.apply_query_key_layer_scaling = args.apply_query_key_layer_scaling
@@ -244,7 +245,7 @@ class ParallelSelfAttention(MegatronModule):
             query_layer, key_layer)
 
         # fp32 conversion.
-        if self.attention_softmax_in_fp32:
+        if self.fp16 and self.attention_softmax_in_fp32:
             attention_scores = attention_scores.float()
 
         # Apply attention mask. [b, np, s, s]
@@ -267,7 +268,7 @@ class ParallelSelfAttention(MegatronModule):
         attention_probs = self._get_attention_probs(attention_scores)
 
         # fp16 conversion
-        if self.attention_softmax_in_fp32:
+        if self.fp16 and self.attention_softmax_in_fp32:
             attention_probs = attention_probs.half()
 
         # Context layer. [b, s, hp]

megatron/training.py

@@ -37,11 +37,12 @@ from megatron.learning_rates import AnnealingLR
 from megatron.model import DistributedDataParallel as LocalDDP
 from megatron.model import get_params_for_weight_decay_optimization
 from megatron.utils import check_adlr_autoresume_termination
+from megatron.utils import make_data_loader
 from megatron.utils import report_memory
 
 
-def pretrain(train_val_test_data_provider, model_provider, forward_step_func,
-             extra_args_provider=None, args_defaults={}):
+def pretrain(train_valid_test_dataset_provider, model_provider,
+             forward_step_func, extra_args_provider=None, args_defaults={}):
     """Main training program.
 
     This function will run the followings in the order provided:
@@ -51,8 +52,8 @@ def pretrain(train_val_test_data_provider, model_provider, forward_step_func,
         4) train the modle using the forward_step_func.
 
     Arguments:
-        train_val_test_data_provider: a function that builds datasets
-            and returns `train, val, test` dataloaders.
+        train_valid_test_dataset_provider: a function that takes the size of
+            train/valid/test dataset and returns `train, valid, test` datasets.
         model_provider: a function that returns a vanilla version of the
             model. By vanilla we mean a simple model on cpu with no fp16 or ddp.
         forward_step_func: a function that takes a `data iterator` and `model`,
@@ -78,35 +79,28 @@ def pretrain(train_val_test_data_provider, model_provider, forward_step_func,
     timers('model and optimizer').stop()
 
     # Data stuff.
-    timers('train/valid/test dataset').start()
-    train_data, val_data, test_data = train_val_test_data_provider()
-    timers('train/valid/test dataset').stop()
-
-    # Train, validation, and test data.
-    timers('train/valid/test dataloader').start()
-    train_data_iterator, val_data_iterator, \
-        test_data_iterator = get_train_val_test_data_iterators(train_data,
-                                                               val_data,
-                                                               test_data)
-    timers('train/valid/test dataloader').stop()
+    timers('train/valid/test data iterators').start()
+    train_data_iterator, valid_data_iterator, test_data_iterator \
+        = build_train_valid_test_data_iterators(
+            train_valid_test_dataset_provider)
+    timers('train/valid/test data iterators').stop()
 
     # Print setup timing.
     print_rank_0('done with setups ...')
-    timers.log(['model and optimizer', 'train/valid/test dataset',
-                'train/valid/test dataloader'])
+    timers.log(['model and optimizer', 'train/valid/test data iterators'])
     print_rank_0('training ...')
 
     iteration = 0
     if args.do_train and args.train_iters > 0:
         iteration, _ = train(forward_step_func,
                              model, optimizer, lr_scheduler,
-                             train_data_iterator, val_data_iterator)
+                             train_data_iterator, valid_data_iterator)
 
 
     if args.do_valid:
         prefix = 'the end of training for val data'
         evaluate_and_print_results(prefix, forward_step_func,
-                                   val_data_iterator, model,
+                                   valid_data_iterator, model,
                                    iteration, False)
 
     if args.save and iteration != 0:
@@ -152,7 +146,6 @@ def get_model(model_provider_func):
 
     raise NotImplementedError('Unknown DDP implementation specified: {}. '
                               'Exiting.'.format(args.DDP_impl))
-    sys.exit()
 
 
 def get_optimizer(model):
@@ -269,19 +262,16 @@ def train_step(forward_step_func, data_iterator,
     timers('forward').start()
     loss, loss_reduced = forward_step_func(data_iterator, model)
     timers('forward').stop()
-    torch.cuda.synchronize()
 
     # Calculate gradients, reduce across processes, and clip.
     timers('backward').start()
     backward_step(optimizer, model, loss)
     timers('backward').stop()
-    torch.cuda.synchronize()
 
     # Update parameters.
     timers('optimizer').start()
     optimizer.step()
     timers('optimizer').stop()
-    torch.cuda.synchronize()
 
     # Update learning rate.
     skipped_iter = 0
@@ -354,7 +344,7 @@ def training_log(loss_dict, total_loss_dict, learning_rate, iteration,
 
 
 def train(forward_step_func, model, optimizer, lr_scheduler,
-          train_data_iterator, val_data_iterator):
+          train_data_iterator, valid_data_iterator):
     """Train the model function."""
     args = get_args()
     timers = get_timers()
@@ -381,9 +371,12 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
         iteration += 1
 
         # Logging.
+        loss_scale = None
+        if args.fp16:
+            loss_scale = optimizer.loss_scale
         report_memory_flag = training_log(loss_dict, total_loss_dict,
                                           optimizer.param_groups[0]['lr'],
-                                          iteration, optimizer.loss_scale,
+                                          iteration, loss_scale,
                                           report_memory_flag)
 
         # Autoresume
@@ -402,7 +395,7 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
            args.do_valid:
             prefix = 'iteration {}'.format(iteration)
             evaluate_and_print_results(prefix, forward_step_func,
-                                       val_data_iterator, model,
+                                       valid_data_iterator, model,
                                        iteration, False)
 
         if args.exit_interval and iteration % args.exit_interval == 0:
@@ -471,37 +464,87 @@ def evaluate_and_print_results(prefix, forward_step_func,
     print_rank_0('-' * length)
 
 
-def get_train_val_test_data_iterators(train_data, val_data, test_data):
-    """Build train/validation/test iterators"""
+def build_train_valid_test_data_iterators(
+        build_train_valid_test_datasets_provider):
+    """XXX"""
     args = get_args()
 
+    (train_dataloader, valid_dataloader, test_dataloader) = (None, None, None)
+
+    print_rank_0('> building train, validation, and test datasets ...')
+    # Data loader only on rank 0 of each model parallel group.
+    if mpu.get_model_parallel_rank() == 0:
+        # Rank, size, and global batch size.
+        data_parallel_size = mpu.get_data_parallel_world_size()
+        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]))
+
+        # Build the datasets.
+        train_ds, valid_ds, test_ds = build_train_valid_test_datasets_provider(
+            train_val_test_num_samples)
+
+        # Build dataloders.
+        train_dataloader = make_data_loader(train_ds)
+        valid_dataloader = make_data_loader(valid_ds)
+        test_dataloader = make_data_loader(test_ds)
+
+        # Flags to know if we need to do training/validation/testing.
+        do_train = train_dataloader is not None and args.train_iters > 0
+        do_valid = valid_dataloader is not None and args.eval_iters > 0
+        do_test = test_dataloader 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])
+
+    # Broadcast num tokens.
+    torch.distributed.broadcast(flags,
+                                mpu.get_model_parallel_src_rank(),
+                                group=mpu.get_model_parallel_group())
+    args.do_train = flags[0].item()
+    args.do_valid = flags[1].item()
+    args.do_test = flags[2].item()
+
     # Shift the start iterations.
-    if train_data is not None:
-        train_data.batch_sampler.start_iter = args.iteration % \
-                                              len(train_data)
+    if train_dataloader is not None:
+        train_dataloader.batch_sampler.start_iter = args.iteration % \
+                                                    len(train_dataloader)
         print_rank_0('setting training data start iteration to {}'.
-                     format(train_data.batch_sampler.start_iter))
-    if val_data is not None:
+                     format(train_dataloader.batch_sampler.start_iter))
+    if valid_dataloader is not None:
         start_iter_val = (args.iteration // args.eval_interval) * \
                          args.eval_iters
-        val_data.batch_sampler.start_iter = 0
-
+        valid_dataloader.batch_sampler.start_iter = start_iter_val % \
+                                                    len(valid_dataloader)
         print_rank_0('setting validation data start iteration to {}'.
-                     format(val_data.batch_sampler.start_iter))
+                     format(valid_dataloader.batch_sampler.start_iter))
 
-    if train_data is not None:
-        train_data_iterator = iter(train_data)
+    # Build iterators.
+    if train_dataloader is not None:
+        train_data_iterator = iter(train_dataloader)
     else:
         train_data_iterator = None
 
-    if val_data is not None:
-        val_data_iterator = iter(val_data)
+    if valid_dataloader is not None:
+        valid_data_iterator = iter(valid_dataloader)
     else:
-        val_data_iterator = None
+        valid_data_iterator = None
 
-    if test_data is not None:
-        test_data_iterator = iter(test_data)
+    if test_dataloader is not None:
+        test_data_iterator = iter(test_dataloader)
     else:
         test_data_iterator = None
 
-    return train_data_iterator, val_data_iterator, test_data_iterator
+    return train_data_iterator, valid_data_iterator, test_data_iterator

openwebtext/make_gpt2_dataset.py

-# coding=utf-8
-# Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-import json
-import numpy as np
-import time
-import os
-import sys
-
-from tokenizer import Tokenizer
-
-
-def tokenize_corpus(filename, np_filename, print_interval=10000):
-
-    print(' > tokenizing {}'.format(filename))
-
-    tokenizer = Tokenizer(cache_dir='./cache')
-
-    tokenized_docs = []
-    num_docs = 0
-    num_tokens = 0
-    start_time = time.time()
-    with open(filename, 'r') as f:
-        for line in f:
-            try:
-                myjson = json.loads(line)
-                url = myjson['url']
-                sample = myjson['text']
-                tokens = tokenizer.tokenize_document(sample)
-                tokenized_docs.append(np.array(tokens, dtype=np.uint16))
-                num_docs += 1
-                num_tokens += len(tokens)
-                if num_docs % print_interval == 0:
-                    print('    processed {:9d} documents in {:.2f} (s) so far'.
-                          format(num_docs, time.time() - start_time),
-                          flush=True)
-            except Exception as e:
-                print('    skipping ', line, e)
-
-    print('  >> processed {} document with total of {} tokens ...'.format(
-        num_docs, num_tokens))
-
-    tokenized_docs = np.array(tokenized_docs, dtype=object)
-    np.save(np_filename, tokenized_docs, allow_pickle=True)
-    print('  >> saved the tokenzed document to {} ...'.format(np_filename))
-
-
-if __name__ == '__main__':
-
-    print('building gpt2 dataset ...')
-
-    path = sys.argv[1]
-    shard = sys.argv[2]
-
-    input_filename = os.path.join(path,
-                                  'shards/shard_{:04d}'.format(int(shard)))
-    output_filename = os.path.join(path,
-                                  'npys/shard_{:04d}.npy'.format(int(shard)))
-    print('will be reading {}'.format(input_filename))
-    print('and will write the results to {}'.format(output_filename))
-
-    tokenize_corpus(input_filename, output_filename)
-
-

openwebtext/make_gpt2_sizes.py

-
-import glob
-import json
-import os
-import time
-import sys
-
-import numpy as np
-
-
-if __name__ == '__main__':
-
-    print('building the shard sizes ...')
-
-    path = sys.argv[1]
-    print('> reading numpy files from {}'.format(path))
-
-    npy_files = glob.glob(path + '/*.npy')
-    npy_files.sort()
-    print('  found {} numpy files'.format(len(npy_files)))
-
-    size_dict = {}
-    counter = 0
-    start_time = time.time()
-    for filename in npy_files:
-        data = np.load(filename, allow_pickle=True)
-        size = np.hstack(data).size
-        np_filename = os.path.basename(filename)
-        size_dict[np_filename] = size
-        counter += 1
-        if counter % 10 == 0:
-            print('   processed {} files in {:.2f} seconds'.format(
-                counter, time.time() - start_time))
-
-    output_filename = os.path.join(path, 'sizes.txt')
-    with open(output_filename, 'w') as f:
-        json.dump(size_dict, f)
-    print('> wrote sizes to {}'.format(output_filename))

openwebtext/run_make_gpt2_dataset.sh

-#!/bin/bash
-
-echo "processing gpt2 data ..."
-DIR="/raid/mpatwary/redownload_v0/0-21"
-
-for thread in {0..3}; do
-    echo " launching thread "$thread && python make_gpt2_dataset.py $DIR $thread > $DIR/logs/shard_$thread.log 2>&1 &
-done

openwebtext/tokenizer.py

-# coding=utf-8
-# Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import sys
-sys.path.append('..')
-
-from megatron.data_utils.tokenization_gpt2 import GPT2Tokenizer
-
-
-class Tokenizer:
-
-    def __init__(self, cache_dir=None):
-        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2',
-                                                       cache_dir=cache_dir)
-        self.tokenizer.max_len = int(1e12)
-        self.eod_token = self.tokenizer.encoder['<|endoftext|>']
-        assert self.eod_token < 65535, 'vocab size will not fit in uint16'
-        print('> GPT2 tokenizer with {} vocab size and eod token {} ...'.format(
-            len(self.tokenizer.encoder), self.eod_token))
-
-    def tokenize_document(self, document):
-        tokens = self.tokenizer.encode(document)
-        tokens.append(self.eod_token)
-        return tokens

pretrain_bert.py

@@ -25,13 +25,11 @@ from megatron import print_rank_0
 from megatron.data.bert_dataset import build_train_valid_test_datasets
 from megatron.model import BertModel
 from megatron.training import pretrain
-from megatron.utils import make_data_loader
 from megatron.utils import reduce_losses
 
 
 def model_provider():
     """Build the model."""
-    args = get_args()
 
     print_rank_0('building BERT model ...')
 
@@ -44,6 +42,7 @@ def model_provider():
 
 
 def get_batch(data_iterator):
+    """Build the batch."""
 
     # Items and their type.
     keys = ['text', 'types', 'labels', 'is_random', 'loss_mask', 'padding_mask']
@@ -96,33 +95,12 @@ def forward_step(data_iterator, model):
     return loss, {'lm loss': reduced_losses[0], 'sop loss': reduced_losses[1]}
 
 
-def get_train_val_test_data():
-    """Load the data on rank zero and boradcast number of tokens to all GPUS."""
+def train_valid_test_datasets_provider(train_val_test_num_samples):
+    """Build train, valid, and test datasets."""
     args = get_args()
 
-    (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:
     print_rank_0('> building train, validation, and test datasets '
                  'for BERT ...')
-
-        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,
@@ -135,31 +113,10 @@ def get_train_val_test_data():
         skip_warmup=(not args.mmap_warmup))
     print_rank_0("> finished creating BERT 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])
-
-    # Broadcast num tokens.
-    torch.distributed.broadcast(flags,
-                                mpu.get_model_parallel_src_rank(),
-                                group=mpu.get_model_parallel_group())
-    args.do_train = flags[0].item()
-    args.do_valid = flags[1].item()
-    args.do_test = flags[2].item()
-
-    return train_data, valid_data, test_data
+    return train_ds, valid_ds, test_ds
 
 
 if __name__ == "__main__":
 
-    pretrain(get_train_val_test_data, model_provider, forward_step,
+    pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
              args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})

pretrain_bert_ict.py

@@ -25,7 +25,6 @@ from megatron import print_rank_0
 from megatron.data.bert_dataset import build_train_valid_test_datasets
 from megatron.model import ICTBertModel
 from megatron.training import pretrain
-from megatron.utils import make_data_loader
 from megatron.utils import reduce_losses
 
 num_batches = 0
@@ -46,8 +45,8 @@ def model_provider():
 def get_batch(data_iterator):
 
     # Items and their type.
-    keys = ['input_text', 'input_types', 'input_pad_mask',
-            'context_text', 'context_types', 'context_pad_mask']
+    keys = ['query_tokens', 'query_types', 'query_pad_mask',
+            'block_tokens', 'block_types', 'block_pad_mask']
     datatype = torch.int64
 
     # Broadcast data.
@@ -58,16 +57,15 @@ def get_batch(data_iterator):
     data_b = mpu.broadcast_data(keys, data, datatype)
 
     # Unpack.
-    input_tokens = data_b['input_text'].long()
-    input_types = data_b['input_types'].long()
-    input_pad_mask = data_b['input_pad_mask'].long()
-    context_tokens = data_b['context_text'].long()
-    context_types = data_b['context_types'].long()
-    context_pad_mask = data_b['context_pad_mask'].long()
-    context_indices = data_b['context_indices'].long()
+    query_tokens = data_b['query_tokens'].long()
+    query_types = data_b['query_types'].long()
+    query_pad_mask = data_b['query_pad_mask'].long()
+    block_tokens = data_b['block_tokens'].long()
+    block_types = data_b['block_types'].long()
+    block_pad_mask = data_b['block_pad_mask'].long()
 
-    return input_tokens, input_types, input_pad_mask,\
-           context_tokens, context_types, context_pad_mask, context_indices
+    return query_tokens, query_types, query_pad_mask,\
+           block_tokens, block_types, block_pad_mask
 
 
 def forward_step(data_iterator, model):
@@ -76,15 +74,18 @@ def forward_step(data_iterator, model):
 
     # Get the batch.
     timers('batch generator').start()
-    input_tokens, input_types, input_pad_mask,\
-    context_tokens, context_types, context_pad_mask = get_batch(data_iterator)
+    query_tokens, query_types, query_pad_mask,\
+    block_tokens, block_types, block_pad_mask = get_batch(data_iterator)
     timers('batch generator').stop()
 
     # Forward model.
-    retrieval_scores = model(input_tokens, input_pad_mask, input_types,
-                             context_tokens, context_pad_mask, context_types).float()
+    query_logits, block_logits = model(query_tokens, query_pad_mask, query_types,
+                                       block_tokens, block_pad_mask, block_types).float()
 
+    # [batch x h] * [h x batch]
+    retrieval_scores = query_logits.matmul(torch.transpose(block_logits, 0, 1))
     softmaxed = F.softmax(retrieval_scores, dim=1)
+
     top5_vals, top5_indices = torch.topk(softmaxed, k=5, sorted=True)
     batch_size = softmaxed.shape[0]
 
@@ -99,33 +100,12 @@ def forward_step(data_iterator, model):
                             'top5_acc': reduced_losses[2]}
 
 
-def get_train_val_test_data():
-    """Load the data on rank zero and boradcast number of tokens to all GPUS."""
+def train_valid_test_datasets_provider(train_val_test_num_samples):
+    """Build train, valid and test datasets."""
     args = get_args()
-
-    (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:
     print_rank_0('> building train, validation, and test datasets '
                  'for BERT ...')
 
-        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_iters + 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,
@@ -139,31 +119,10 @@ def get_train_val_test_data():
         ict_dataset=True)
     print_rank_0("> finished creating BERT ICT 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])
-
-    # Broadcast num tokens.
-    torch.distributed.broadcast(flags,
-                                mpu.get_model_parallel_src_rank(),
-                                group=mpu.get_model_parallel_group())
-    args.do_train = flags[0].item()
-    args.do_valid = flags[1].item()
-    args.do_test = flags[2].item()
-
-    return train_data, valid_data, test_data
+    return train_ds, valid_ds, test_ds
 
 
 if __name__ == "__main__":
 
-    pretrain(get_train_val_test_data, model_provider, forward_step,
+    pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
              args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})

pretrain_gpt2.py

@@ -15,8 +15,6 @@
 
 """Pretrain GPT2"""
 
-import os
-
 import torch
 
 from megatron import get_args
@@ -24,17 +22,15 @@ from megatron import get_timers
 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
-from megatron.utils import make_data_loader
 from megatron.utils import reduce_losses
 
 
 def model_provider():
     """Build the model."""
-    args = get_args()
 
     print_rank_0('building GPT2 model ...')
     model = GPT2Model(num_tokentypes=0, parallel_output=True)
@@ -98,71 +94,26 @@ def forward_step(data_iterator, model):
     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."""
+def train_valid_test_datasets_provider(train_val_test_num_samples):
+    """Build train, valid, and test datasets."""
     args = get_args()
 
-    (train_data, val_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)])
-    else:
-        flags = torch.cuda.LongTensor([0, 0, 0])
-
-    # Broadcast num tokens.
-    torch.distributed.broadcast(flags,
-                                mpu.get_model_parallel_src_rank(),
-                                group=mpu.get_model_parallel_group())
-    args.do_train = flags[0].item()
-    args.do_valid = flags[1].item()
-    args.do_test = flags[2].item()
+    print_rank_0('> building train, validation, and test datasets '
+                 'for GPT2 ...')
+    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 ...")
 
-    return train_data, val_data, test_data
+    return train_ds, valid_ds, test_ds
 
 
 if __name__ == "__main__":
 
-    pretrain(get_train_val_test_data, model_provider, forward_step,
+    pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
              args_defaults={'tokenizer_type': 'GPT2BPETokenizer'})

requirements.txt

-nltk>=3.4
-numpy>=1.15.4
-pandas>=0.24.0
-sentencepiece>=0.1.8
-tensorflow>=1.12.0