switch to mcore dataset [with FIM support]

examples/nlp/language_modeling/conf/megatron_gpt_config.yaml

@@ -9,15 +9,15 @@ trainer:
   devices: 1
   num_nodes: 1
   accelerator: gpu
-  precision: 16
+  precision: bf16
   logger: False # logger provided by exp_manager
   enable_checkpointing: False
   use_distributed_sampler: False
   max_epochs: -1 # PTL default. In practice, max_steps will be reached first.
-  max_steps: 100000 # consumed_samples = global_step * micro_batch_size * data_parallel_size * accumulate_grad_batches
+  max_steps: 5000 # consumed_samples = global_step * micro_batch_size * data_parallel_size * accumulate_grad_batches
   log_every_n_steps: 10
-  val_check_interval: 100
-  limit_val_batches: 50
+  val_check_interval: 250
+  limit_val_batches: 25
   limit_test_batches: 500
   accumulate_grad_batches: 1 # do not modify, grad acc is automatic for training megatron models
   gradient_clip_val: 1.0
@@ -28,17 +28,17 @@ exp_manager:
   explicit_log_dir: null
   exp_dir: null
   name: megatron_gpt
-  create_wandb_logger: False
+  create_wandb_logger: True
   wandb_logger_kwargs:
-    project: null
-    name: null
+    project: mcore_ds_test
+    name: mcore_ds_new
   resume_if_exists: True
   resume_ignore_no_checkpoint: True
   resume_from_checkpoint: ${model.resume_from_checkpoint}
   create_checkpoint_callback: True
   checkpoint_callback_params:
     monitor: val_loss
-    save_top_k: 10
+    save_top_k: 5
     mode: min
     always_save_nemo: False # saves nemo file during validation, not implemented for model parallel
     save_nemo_on_train_end: False # not recommended when training large models on clusters with short time limits
@@ -47,7 +47,7 @@ exp_manager:
 
 model:
   # use GPTModel from megatron.core
-  mcore_gpt: False
+  mcore_gpt: True
 
   # specify micro_batch_size, global_batch_size, and model parallelism
   # gradient accumulation will be done automatically based on data_parallel_size
@@ -210,7 +210,7 @@ model:
     # Dictionary: can override from CLI "model.data.data_prefix"={"train":[1.0, /path/to/data], "validation":/path/to/data, "test":/path/to/test}
     # Or see example below:
     # "model.data.data_prefix: {train:[1.0,/path/to/data], validation:[/path/to/data], test:[/path/to/test]}"
-    data_prefix: ???
+    data_prefix: [1.0, /home/data/test_text_document]
     index_mapping_dir: null # path to save index mapping .npy files, by default will save in the same location as data_prefix
     data_impl: mmap
     splits_string: 900,50,50
@@ -236,16 +236,16 @@ model:
     gen_shape: False # Generate model and kernel details including input shapes
 
   optim:
-    name: fused_adam
+    name: distributed_fused_adam
     lr: 2e-4
     weight_decay: 0.01
     betas:
     - 0.9
     - 0.98
     sched:
       name: CosineAnnealing
-      warmup_steps: 500
-      constant_steps: 0
+      warmup_steps: 50
+      constant_steps: 500
       min_lr: 2e-5
 
   gc_interval: 0

nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py

@@ -75,7 +75,9 @@
     HAVE_APEX = False
 
 try:
-    from megatron.core import InferenceParams, parallel_state
+    from megatron.core import InferenceParams, mpu, parallel_state, tensor_parallel
+    from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
+    from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig
     from megatron.core.models.gpt import GPTModel as MCoreGPTModel
     from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
     from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
@@ -111,6 +113,15 @@
     return name_spec_dict[spec_name]
 
 
+global is_dataset_built_on_rank
+
+
+def is_dataset_built_on_rank():
+    return (
+        mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage()
+    ) and mpu.get_tensor_model_parallel_rank() == 0
+
+
 class MegatronGPTExportableModel(torch.nn.Module, Exportable):
     """
     Megatron GPT Wrapper for ONNX export
@@ -231,6 +242,10 @@
             self.if_first_step = 0
             self.prev_global_batch_size = None
 
+        self.reset_position_ids = cfg.data.get('reset_position_ids', False)
+        self.reset_attention_mask = cfg.data.get('reset_attention_mask', False)
+        self.eod_mask_loss = cfg.data.get('eod_mask_loss', False)
+
         if not self.megatron_amp_O2 and self.cfg.get('virtual_pipeline_model_parallel_size', None):
             raise ValueError('Virtual pipeline model parallel is only supported when using megatron_amp_O2')
 
@@ -837,6 +852,102 @@
         # TODO @tmoon: Use once available in Megatron-LM
         # return DataIteratorList(iters)
 
+    def get_ltor_masks_and_position_ids(
+        self, data, eod_token, reset_position_ids, reset_attention_mask, eod_mask_loss
+    ):
+        """Build masks and position id for left to right model."""
+
+        # Extract batch size and sequence length.
+        micro_batch_size, seq_length = data.size()
+
+        # Attention mask (lower triangular).
+        if reset_attention_mask:
+            att_mask_batch = micro_batch_size
+        else:
+            att_mask_batch = 1
+        attention_mask = torch.tril(torch.ones((att_mask_batch, seq_length, seq_length), device=data.device)).view(
+            att_mask_batch, 1, seq_length, seq_length
+        )
+
+        # Loss mask.
+        loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device)
+        if eod_mask_loss:
+            loss_mask[data == eod_token] = 0.0
+
+        # Position ids.
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)
+        position_ids = position_ids.unsqueeze(0).expand_as(data)
+        # We need to clone as the ids will be modifed based on batch index.
+        if reset_position_ids:
+            position_ids = position_ids.clone()
+
+        if reset_position_ids or reset_attention_mask:
+            # Loop through the batches:
+            for b in range(micro_batch_size):
+
+                # Find indecies where EOD token is.
+                eod_index = position_ids[b, data[b] == eod_token]
+                # Detach indecies from positions if going to modify positions.
+                if reset_position_ids:
+                    eod_index = eod_index.clone()
+
+                # Loop through EOD indecies:
+                prev_index = 0
+                for j in range(eod_index.size()[0]):
+                    i = eod_index[j]
+                    # Mask attention loss.
+                    if reset_attention_mask:
+                        attention_mask[b, 0, (i + 1) :, : (i + 1)] = 0
+                    # Reset positions.
+                    if reset_position_ids:
+                        position_ids[b, (i + 1) :] -= i + 1 - prev_index
+                        prev_index = i + 1
+
+        # Convert attention mask to binary:
+        attention_mask = attention_mask < 0.5
+
+        return attention_mask, loss_mask, position_ids
+
+    def get_batch(self, data_iterator):
+        """Generate a batch."""
+
+        # TODO: this is pretty hacky, find a better way
+        if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()):
+            return None, None, None, None, None
+
+        # Items and their type.
+        keys = ['text']
+        datatype = torch.int64
+
+        # Broadcast data.
+        if data_iterator is not None:
+            data = next(data_iterator)
+        else:
+            data = None
+        data_b = tensor_parallel.broadcast_data(keys, data, datatype)
+
+        # Unpack.
+        tokens_ = data_b['text'].long()
+        labels = tokens_[:, 1:].contiguous()
+        tokens = tokens_[:, :-1].contiguous()
+
+        # Get the masks and postition ids.
+        attention_mask, loss_mask, position_ids = self.get_ltor_masks_and_position_ids(
+            tokens, self.tokenizer.eos_id, self.reset_position_ids, self.reset_attention_mask, self.eod_mask_loss
+        )
+
+        batch = {
+            'tokens': tokens,
+            'labels': labels,
+            'loss_mask': loss_mask,
+            'attention_mask': attention_mask,
+            'position_ids': position_ids,
+        }
+        # slice batch along sequence dimension for context parallelism
+        batch = self.get_batch_on_this_context_parallel_rank(batch)
+
+        return batch
+
     def get_batch_on_this_context_parallel_rank(self, batch):
         cp_size = self.cfg.get('context_parallel_size', 1)
         num_valid_tokens_in_ub = None
@@ -867,7 +978,8 @@
         def fwd_output_and_loss_func(dataloader_iter, model, checkpoint_activations_all_layers=None):
 
             # Get data batch
-            batch = next(dataloader_iter)
+            # batch = next(dataloader_iter)
+            batch = self.get_batch(dataloader_iter)
 
             # Transfer needed data to GPU
             required_keys = set()
@@ -1094,18 +1206,33 @@
                 1
             ] = 1  # This is to make sure we only have one epoch on every validation iteration
 
-        self._train_ds, self._validation_ds, self._test_ds = build_train_valid_test_datasets(
-            cfg=self.cfg,
-            trainer=self.trainer,
-            data_prefix=self.cfg.data.data_prefix,
-            data_impl=self.cfg.data.data_impl,
-            splits_string=self.cfg.data.splits_string,
-            train_valid_test_num_samples=train_valid_test_num_samples,
-            seq_length=self.cfg.data.seq_length,
-            seed=self.cfg.seed,
-            skip_warmup=self.cfg.data.get('skip_warmup', True),
-            tokenizer=self.tokenizer,
+        # self._train_ds, self._validation_ds, self._test_ds = build_train_valid_test_datasets(
+        #     cfg=self.cfg,
+        #     trainer=self.trainer,
+        #     data_prefix=self.cfg.data.data_prefix,
+        #     data_impl=self.cfg.data.data_impl,
+        #     splits_string=self.cfg.data.splits_string,
+        #     train_valid_test_num_samples=train_valid_test_num_samples,
+        #     seq_length=self.cfg.data.seq_length,
+        #     seed=self.cfg.seed,
+        #     skip_warmup=self.cfg.data.get('skip_warmup', True),
+        #     tokenizer=self.tokenizer,
+        # )
+
+        dataset_config = GPTDatasetConfig(
+            is_built_on_rank=is_dataset_built_on_rank,
+            random_seed=self.cfg.seed,
+            sequence_length=self.cfg.data.seq_length,
+            blend=self.cfg.data.data_prefix,
+            blend_per_split=None,
+            split=self.cfg.data.splits_string,
+            path_to_cache=self.cfg.data.index_mapping_dir,
         )
+
+        self._train_ds, self._validation_ds, self._test_ds = BlendedMegatronDatasetBuilder(
+            GPTDataset, train_valid_test_num_samples, dataset_config,
+        ).build()
+
         if self._train_ds is not None:
             logging.info(f'Length of train dataset: {len(self._train_ds)}')
         if self._validation_ds is not None: