add first_val_step for mcore schedules

Signed-off-by: jiemingz <jiemingz@nvidia.com>

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

@@ -224,6 +224,8 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):
 
         self.mcore_gpt = cfg.get('mcore_gpt', False)
         self.spec_name = cfg.get('name', '')
+        if cfg.get('fp8', False):
+            self.prev_step_training = True
 
         self.rampup_batch_size = self.cfg.get('rampup_batch_size', None)
         if self.rampup_batch_size:
@@ -485,7 +487,7 @@ def forward(self, tokens, text_position_ids, attention_mask, labels):
         output_tensor = self.model(tokens, text_position_ids, attention_mask, labels=labels)
         return output_tensor
 
-    def fwd_bwd_step(self, dataloader_iter, batch_idx, forward_only):
+    def fwd_bwd_step(self, dataloader_iter, batch_idx, forward_only, first_val_step=False):
 
         # handle asynchronous grad reduction
         no_sync_func = None
@@ -515,6 +517,7 @@ def fwd_bwd_step(self, dataloader_iter, batch_idx, forward_only):
             forward_only=forward_only,
             seq_length=self.cfg.encoder_seq_length,
             micro_batch_size=self.cfg.micro_batch_size,
+            first_val_step=first_val_step,
         )
 
         # only the last stages of the pipeline return losses
@@ -1012,9 +1015,11 @@ def validation_step(self, dataloader_iter, batch_idx):
         if isinstance(self.model, list):
             for model_module in self.model:
                 model_module.eval()
+        first_val_step = self.prev_step_training and not self.training
 
-        loss = self.fwd_bwd_step(dataloader_iter, batch_idx, True)
+        loss = self.fwd_bwd_step(dataloader_iter, batch_idx, True, first_val_step)
 
+        self.prev_step_training = self.training
         if isinstance(self.model, list):
             for model_module in self.model:
                 model_module.train()