Reduce Overhead in Timers

megatron/core/timers.py

@@ -270,6 +270,9 @@ def _get_elapsed_time_all_ranks(self, names, reset, barrier):
             torch.tensor: Tensor of size [world_size, len(names)] with times in float.
         """
 
+        if len(names) == 0:
+            return None
+
         # First make sure all the callers are in sync.
         if barrier:
             torch.distributed.barrier()
@@ -294,6 +297,7 @@ def _get_elapsed_time_all_ranks(self, names, reset, barrier):
                 rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset)
 
         # See the note above for why we are not using gather.
+
         dist_all_gather_func(rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1))
 
         return rank_name_to_time
@@ -302,16 +306,17 @@ def _get_global_min_max_time(self, names, reset, barrier, normalizer):
         """Report only min and max times across all ranks."""
 
         rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
+        # Using Python built-in methods to avoid the overhead of PyTorch operations.
+        rank_name_to_time = rank_name_to_time.permute(1, 0).tolist() if rank_name_to_time is not None else None
         name_to_min_max_time = {}
         for i, name in enumerate(names):
-            rank_to_time = rank_name_to_time[:, i]
             # filter out the ones we did not have any timings for
-            rank_to_time = rank_to_time[rank_to_time > 0.0]
+            rank_to_time = list(filter(lambda x: x > 0.0, rank_name_to_time[i]))
             # If the timer exists:
-            if rank_to_time.numel() > 0:
+            if len(rank_to_time) > 0:
                 name_to_min_max_time[name] = (
-                    rank_to_time.min().item() / normalizer,
-                    rank_to_time.max().item() / normalizer,
+                    min(rank_to_time) / normalizer,
+                    max(rank_to_time) / normalizer,
                 )
         return name_to_min_max_time
 

megatron/training/training.py

@@ -1419,32 +1419,37 @@ def training_log(
     total_loss_dict[nan_iters_key] = total_loss_dict.get(nan_iters_key, 0) + int(got_nan)
 
     # Logging.
-    timers_to_log = [
-        'forward-backward',
-        'forward-compute',
-        'backward-compute',
-        'batch-generator',
-        'forward-recv',
-        'forward-send',
-        'backward-recv',
-        'backward-send',
-        'forward-send-forward-recv',
-        'forward-send-backward-recv',
-        'backward-send-forward-recv',
-        'backward-send-backward-recv',
-        'forward-backward-send-forward-backward-recv',
-        'layernorm-grads-all-reduce',
-        'embedding-grads-all-reduce',
-        'all-grads-sync',
-        'params-all-gather',
-        'optimizer-copy-to-main-grad',
-        'optimizer-unscale-and-check-inf',
-        'optimizer-clip-main-grad',
-        'optimizer-count-zeros',
-        'optimizer-inner-step',
-        'optimizer-copy-main-to-model-params',
-        'optimizer',
-    ]
+    timers_to_log = []
+    if args.timing_log_level >= 1:
+        timers_to_log.extend([
+            'forward-backward',
+            'layernorm-grads-all-reduce',
+            'embedding-grads-all-reduce',
+            'all-grads-sync',
+            'params-all-gather',
+            'optimizer-copy-to-main-grad',
+            'optimizer-unscale-and-check-inf',
+            'optimizer-clip-main-grad',
+            'optimizer-count-zeros',
+            'optimizer-inner-step',
+            'optimizer-copy-main-to-model-params',
+            'optimizer',
+        ])
+    if args.timing_log_level >= 2:
+        timers_to_log.extend([
+            'batch-generator',
+            'forward-compute',
+            'backward-compute',
+            'forward-recv',
+            'forward-send',
+            'backward-recv',
+            'backward-send',
+            'forward-send-forward-recv',
+            'forward-send-backward-recv',
+            'backward-send-forward-recv',
+            'backward-send-backward-recv',
+            'forward-backward-send-forward-backward-recv',
+        ])
     # Add timers from RL loop if needed.
     if getattr(args, 'perform_rl_step', False):
         timers_to_log.extend(['rollout-collection', 'inference-setup', 'collect-rollouts', 'postrollout-gc-collect',