change loss return format so that it can work with calculate_per_token_loss (NVIDIA-NeMo#12459)

* loss upscaling has been moved to MCore, no need to handle it in model level any more

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* return loss_sum and num_valid_tokens separately

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* change num_tokens dtype to int

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix a return type

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* clean masked_token_loss

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* Apply isort and black reformatting

Signed-off-by: xrennvidia <xrennvidia@users.noreply.github.com>

* minor fix

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* Apply isort and black reformatting

Signed-off-by: xrennvidia <xrennvidia@users.noreply.github.com>

* minor fix

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* Apply isort and black reformatting

Signed-off-by: xrennvidia <xrennvidia@users.noreply.github.com>

* bug fix

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* remove one unused import

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* fix pylint error

Signed-off-by: Xiaowei Ren <xren@nvidia.com>

* Apply isort and black reformatting

Signed-off-by: xrennvidia <xrennvidia@users.noreply.github.com>

---------

Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: xrennvidia <xrennvidia@users.noreply.github.com>
Co-authored-by: xrennvidia <xrennvidia@users.noreply.github.com>

Author: xrennvidia

nemo/collections/speechlm/models/speech_to_text_llm_model.py

@@ -46,7 +46,6 @@
 from nemo.collections.llm import fn
 from nemo.collections.llm.gpt.model.base import (
     GPTConfig,
-    GPTModel,
     get_batch_on_this_context_parallel_rank,
     get_packed_seq_params,
 )
@@ -883,7 +882,7 @@ def inference_step(self, batch, mode):
 
         if isinstance(forward_output, tuple):
             # reduce validation loss
-            loss = self.validation_loss_reduction.forward(batch=batch, forward_out=forward_output)[1]['avg']
+            loss = self.validation_loss_reduction.forward(batch=batch, forward_out=forward_output)[-1]['avg']
         else:
             # no labels provided, use a dummy loss value
             loss = 0.0
@@ -915,8 +914,14 @@ def inference_step(self, batch, mode):
                 labels_text = clean_end_string(labels_text, self.tokenizer, data_cfg.end_string)
 
             if data_cfg.get("remove_text_pc", False):
-                preds_text = [remove_punctuations(p.lower(), data_cfg.get("punctuations", None)) for p in preds_text]
-                labels_text = [remove_punctuations(l.lower(), data_cfg.get("punctuations", None)) for l in labels_text]
+                preds_text = [
+                    remove_punctuations(pred_text.lower(), data_cfg.get("punctuations", None))
+                    for pred_text in preds_text
+                ]
+                labels_text = [
+                    remove_punctuations(label_text.lower(), data_cfg.get("punctuations", None))
+                    for label_text in labels_text
+                ]
 
             if data_cfg.get("log_every_n_steps", None) is not None:
                 if batch_idx % data_cfg.log_every_n_steps == 0:

nemo/lightning/megatron_parallel.py

@@ -1742,7 +1742,7 @@ def __init__(self, validation_step: bool = False, val_drop_last: bool = True) ->
 
     def forward(
         self, batch: Dict[str, torch.Tensor], forward_out: torch.Tensor
-    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]:
         """Taken from: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py#L951-L976 ."""  # pylint: disable=line-too-long
         from megatron.core import parallel_state
 
@@ -1752,33 +1752,30 @@ def forward(
         if isinstance(forward_out, tuple):
             forward_out, loss_mask = forward_out
             batch["loss_mask"] = loss_mask
+
         cp_size = parallel_state.get_context_parallel_world_size()
-        if cp_size == 1:
-            loss_for_ub = masked_token_loss(forward_out, batch["loss_mask"])
+        loss_sum_for_ub = masked_token_loss(forward_out, batch["loss_mask"], cp_size)
+        if cp_size == 1 or batch['num_valid_tokens_in_ub'] is None:
+            num_valid_tokens_in_ub = batch["loss_mask"].sum()
         else:
-            loss_for_ub = masked_token_loss_context_parallel(
-                forward_out, batch["loss_mask"], batch['num_valid_tokens_in_ub']
-            )
+            num_valid_tokens_in_ub = batch['num_valid_tokens_in_ub']
+        if num_valid_tokens_in_ub < 0.5:  # no valid tokens
+            num_valid_tokens_in_ub += 1.0
+        num_valid_tokens_in_ub = num_valid_tokens_in_ub.clone().detach().to(torch.int)
 
         if self.validation_step and not self.val_drop_last:
-            num_valid_tokens_in_ub = batch["loss_mask"].sum()
-            if loss_for_ub.isnan():
-                assert batch["loss_mask"].count_nonzero() == 0, "Got NaN loss with non-empty input"
+            if loss_sum_for_ub.isnan():
+                assert num_valid_tokens_in_ub == 0, "Got NaN loss with non-empty input"
                 loss_sum_for_ub = torch.zeros_like(num_valid_tokens_in_ub)
-            else:
-                loss_sum_for_ub = num_valid_tokens_in_ub * loss_for_ub
 
             loss_sum_and_ub_size_all_gpu = torch.cat(
-                [
-                    loss_sum_for_ub.clone().detach().view(1),
-                    torch.tensor([num_valid_tokens_in_ub], device=torch.cuda.current_device()).clone().detach(),
-                ]
+                [loss_sum_for_ub.clone().detach().view(1), num_valid_tokens_in_ub]
             )
             torch.distributed.all_reduce(loss_sum_and_ub_size_all_gpu, group=parallel_state.get_data_parallel_group())
-            return loss_for_ub * cp_size, {"loss_sum_and_ub_size": loss_sum_and_ub_size_all_gpu}
+            return loss_sum_for_ub, num_valid_tokens_in_ub, {"loss_sum_and_ub_size": loss_sum_and_ub_size_all_gpu}
 
-        reduced_loss = average_losses_across_data_parallel_group([loss_for_ub])
-        return loss_for_ub * cp_size, {"avg": reduced_loss}
+        reduced_loss = average_losses_across_data_parallel_group([loss_sum_for_ub / num_valid_tokens_in_ub])
+        return loss_sum_for_ub, num_valid_tokens_in_ub, {"avg": reduced_loss}
 
     def reduce(self, losses_reduced_per_micro_batch) -> torch.Tensor:
         """Taken from: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py#L535-L552 ."""  # pylint: disable=line-too-long
@@ -1818,34 +1815,17 @@ def forward(
         return super().forward(batch, forward_out)
 
 
-def masked_token_loss(tensor: Tensor, mask: Tensor):
+def masked_token_loss(tensor: Tensor, mask: Tensor, cp_size: int = 1):
     """
     The function takes as input per-token loss and masks non-required values.
     """
     losses = tensor.float()
     loss_mask = mask.view(-1).float()
-    num_valid_tokens = loss_mask.sum()
-    if num_valid_tokens < 0.5:  # no valid tokens
-        num_valid_tokens += 1.0
-    loss = torch.sum(losses.view(-1) * loss_mask) / num_valid_tokens  # sequence level nll
-
-    return loss
-
-
-def masked_token_loss_context_parallel(tensor: Tensor, mask: Tensor, num_valid_tokens_in_ub: int):
-    """
-    masked token loss for CP > 1 as a separate function for readability.
-    """
-    from megatron.core import parallel_state
+    loss = torch.sum(losses.view(-1) * loss_mask)  # sequence level nll
+    if cp_size > 1:
+        from megatron.core import parallel_state
 
-    losses = tensor.float()
-    loss_mask = mask.view(-1).float()
-    if num_valid_tokens_in_ub is None:
-        num_valid_tokens_in_ub = loss_mask.sum()
-    if num_valid_tokens_in_ub < 0.5:  # no valid tokens
-        num_valid_tokens_in_ub += 1.0
-    loss = torch.sum(losses.view(-1) * loss_mask) / num_valid_tokens_in_ub  # sequence level nll
-    torch.distributed.all_reduce(loss, group=parallel_state.get_context_parallel_group())
+        torch.distributed.all_reduce(loss, group=parallel_state.get_context_parallel_group())
 
     return loss