Fix consumed_samples which was off by one batch

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

@@ -532,6 +532,10 @@ def compute_consumed_samples(self, steps_since_resume=0):
             )
         return int(consumed_samples)
 
+    def _compute_consumed_samples_after_training_step(self):
+        # Add +1 to account for the current batch, which is not counted yet in `trainer.global_step`.
+        return self.compute_consumed_samples(self.trainer.global_step + 1 - self.init_global_step)
+
     def _extract_consumed_samples_from_ckpt(self, ckpt_path):
         try:
             init_consumed_samples = int(float(re.findall(r"consumed_samples\=([0-9]+.[0-9]+)", ckpt_path)[0]))

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

@@ -382,10 +382,7 @@ def training_step(self, dataloader_iter, batch_idx):
             self.log('lr', lr, batch_size=1)
             self.log('global_step', self.trainer.global_step, prog_bar=True, batch_size=1)
             self.log(
-                'consumed_samples',
-                self.compute_consumed_samples(self.trainer.global_step - self.init_global_step),
-                prog_bar=True,
-                batch_size=1,
+                'consumed_samples', self._compute_consumed_samples_after_training_step(), prog_bar=True, batch_size=1,
             )
 
         return loss_mean[0]

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

@@ -624,7 +624,7 @@ def training_step(self, dataloader_iter, batch_idx):
             'global_step', self.trainer.global_step, prog_bar=True, rank_zero_only=True, batch_size=1,
         )
 
-        consumed_samples = self.compute_consumed_samples(self.trainer.global_step - self.init_global_step)
+        consumed_samples = self._compute_consumed_samples_after_training_step()
         # TODO: make sure compute_consumed_samples works for pipeline parallelism
         self.log(
             'consumed_samples', consumed_samples, prog_bar=True, rank_zero_only=True, batch_size=1,

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

@@ -416,7 +416,7 @@ def training_step(self, dataloader_iter, batch_idx):
         # TODO: make sure compute_consumed_samples works for pipeline parallelism
         self.log(
             'consumed_samples',
-            self.compute_consumed_samples(self.trainer.global_step - self.init_global_step),
+            self._compute_consumed_samples_after_training_step(),
             prog_bar=True,
             rank_zero_only=True,
             batch_size=1,

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

@@ -300,10 +300,7 @@ def training_step(self, batch, batch_idx):
             self.log('lr', lr, batch_size=1)
             self.log('global_step', self.trainer.global_step, prog_bar=True, batch_size=1)
             self.log(
-                'consumed_samples',
-                self.compute_consumed_samples(self.trainer.global_step - self.init_global_step),
-                prog_bar=True,
-                batch_size=1,
+                'consumed_samples', self._compute_consumed_samples_after_training_step(), prog_bar=True, batch_size=1,
             )
             self._reduced_loss_buffer = []
         return lm_loss