added timing for count-zeros & inner-step.

megatron/optimizer/optimizer.py

@@ -435,11 +435,15 @@ class MixedPrecisionOptimizer(MegatronOptimizer):
         timers('optimizer-clip-main-grad').stop()
 
         # count the zeros in the grads
+        timers('optimizer-count-zeros').start()
         num_zeros_in_grad = self.count_zeros() if \
                             self.log_num_zeros_in_grad else None
+        timers('optimizer-count-zeros').stop()
 
         # Step the optimizer.
+        timers('optimizer-inner-step').start()
         self.optimizer.step()
+        timers('optimizer-inner-step').stop()
 
         # Update params from main params.
         timers('optimizer-copy-main-to-model-params').start()

megatron/training.py

@@ -523,6 +523,8 @@ def training_log(loss_dict, total_loss_dict, learning_rate, iteration,
     add_to_logging('optimizer-copy-to-main-grad')
     add_to_logging('optimizer-unscale-and-check-inf')
     add_to_logging('optimizer-clip-main-grad')
+    add_to_logging('optimizer-count-zeros')
+    add_to_logging('optimizer-inner-step')
     add_to_logging('optimizer-copy-main-to-model-params')
     add_to_logging('optimizer')
     add_to_logging('batch-generator')