Apply fixes used during training

Author: odelalleau

nemo/collections/nlp/data/language_modeling/megatron/gpt_sft_chat_dataset.py

@@ -51,7 +51,9 @@ def _get_header_conversation_type_mask_role(source, special_tokens):
         if TYPE_INSTRUCTION[data_type] != '':
             conversation = conversation + '\n' + TYPE_INSTRUCTION[data_type]
     mask_role = source.get('mask', 'User')
-    header = f"{special_tokens['system_turn_start']}{SYSTEM_TOKEN}{END_NAME_SIGNAL}{conversation}{END_SIGNAL}"
+    system_token = source.get("system_token", SYSTEM_TOKEN)
+    header = f"{special_tokens['system_turn_start']}{system_token}{END_NAME_SIGNAL}{conversation}{END_SIGNAL}"
+    # logging.info(f"DBG HEADER:\n```{header}```")
     conversation = _add_speaker_and_signal(header, source['conversations'], mask_role, data_type, special_tokens)
     return header, conversation, data_type, mask_role
 
@@ -60,13 +62,14 @@ def get_prompt_template_example(special_tokens):
     source = {
         'system': '{system message}',
         'conversations': [
-            {'from': 'User', 'value': '{turn 1 user message}', 'label': None},
-            {'from': 'Assistant', 'value': '{turn 1 assistant message}', 'label': '{turn 1 assistant label}'},
-            {'from': 'User', 'value': '{turn 2 user message}', 'label': None},
-            {'from': 'Assistant', 'value': '{turn 2 assistant message}', 'label': '{turn 2 assistant label}'},
+            {'from': '{user role}', 'value': '{turn 1 user message}', 'label': None},
+            {'from': '{assistant role}', 'value': '{turn 1 assistant message}', 'label': '{turn 1 assistant label}'},
+            {'from': '{user role}', 'value': '{turn 2 user message}', 'label': None},
+            {'from': '{assistant role}', 'value': '{turn 2 assistant message}', 'label': '{turn 2 assistant label}'},
         ],
-        "mask": "User",
+        "mask": "{user role}",
         "type": "VALUE_TO_TEXT",
+        "system_token": '{system token}',
     }
     _, conversation, _, _ = _get_header_conversation_type_mask_role(source, special_tokens)
     return conversation
@@ -273,6 +276,7 @@ def preprocess(
         id1 = tokenizer.text_to_ids(PREFIX_STR + s["value"])
         id2 = tokenizer.text_to_ids(PREFIX_STR)
         tokenized_sentence = id1[len(id2) :]
+        # logging.info(f"CONV DBG: {tokenized_sentence[0:20]} ... {tokenized_sentence[-20:]}")
         ids.append(torch.tensor(tokenized_sentence))
         tokenized_lens.append(len(tokenized_sentence))
     speakers = [sentence["from"] for sentence in source['conversations']]
@@ -326,6 +330,8 @@ def _build_samples_mapping(self):
         id2 = self.tokenizer.text_to_ids(PREFIX_STR)
         self.num_turn_start_tokens = len(id1) - len(id2)
 
+        # logging.info(f"DATASET DBG:\n{self.special_tokens=}\n{self.label_start_tokens=}, {self.name_end_token_ids=}, {self.num_turn_start_tokens=}")
+
     def _process_example(self, example):
         """
         Create an example by concatenating text and answer.

nemo/collections/nlp/data/language_modeling/megatron/gpt_sft_dataset.py

@@ -228,6 +228,7 @@ def __getitem__(self, idx):
         else:
             auto_gen_idx = False
         try:
+            idx = int(idx)
             example = self.indexed_dataset[idx]
             if auto_gen_idx:
                 example['__AUTOGENERATED__'] = True
@@ -542,6 +543,7 @@ def __getitem__(self, idx):
             # assert idx < len(self.samples_mapping)
             idx = self.samples_mapping[idx]
 
+        idx = int(idx)
         input_ids = self.indexed_dataset[idx]['input_ids']
         seq_boundaries = self.indexed_dataset[idx]['seq_start_id'] + [len(input_ids)]
         loss_mask = self.indexed_dataset[idx]['loss_mask']

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

@@ -1296,6 +1296,9 @@ def fwd_output_and_loss_func(dataloader_iter, model, checkpoint_activations_all_
                 'loss_mask': batch['loss_mask'],
             }
 
+            # if torch.distributed.get_rank() == 0:
+            #     logging.info(f"*****DEBUG OUTPUT*****\nTOKENS:\n{batch['tokens'][0].tolist()}\nPOSITION_IDS:\n{batch['position_ids'][0].tolist()}\nLABELS:\n{batch['labels'][0].tolist()}\nLOSS_MASK:\n{batch['loss_mask'][0].tolist()}\nATTENTION_MASK:\n{None if batch['attention_mask'] is None else batch['attention_mask'][0].tolist()}\n")
+
             if not self.mcore_gpt:
                 forward_args['checkpoint_activations_all_layers'] = checkpoint_activations_all_layers
                 if not self.use_loss_mask:
@@ -1592,7 +1595,7 @@ def loss_func(self, loss_mask, num_valid_tokens_in_ub, output_tensor):
         losses = output_tensor.float()
         loss_mask = loss_mask.view(-1).float()
         # TODO: add nemo version here
-        loss = torch.sum(losses.view(-1) * loss_mask) / num_valid_tokens_in_ub  # sequence level nll
+        loss = torch.sum(losses.view(-1) * loss_mask) / max(1, num_valid_tokens_in_ub)  # sequence level nll
         if parallel_state.get_context_parallel_world_size() > 1:
             torch.distributed.all_reduce(loss, group=parallel_state.get_context_parallel_group())
         return loss