From 1fcee8c48430aec3ebc15b2c110268e1ab39e7ac Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Tue, 11 Nov 2025 10:05:32 +0000 Subject: [PATCH 1/4] reduce overhead in timer Signed-off-by: Xin Yao --- megatron/core/timers.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/megatron/core/timers.py b/megatron/core/timers.py index 75387c33d3a..14aab032a9f 100644 --- a/megatron/core/timers.py +++ b/megatron/core/timers.py @@ -302,16 +302,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() 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 From b6cd128ff84dc9f97fa2206e7834f07561dcd924 Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Tue, 11 Nov 2025 10:39:47 +0000 Subject: [PATCH 2/4] divide timers_to_log according to log level Signed-off-by: Xin Yao --- megatron/core/timers.py | 6 +++- megatron/training/training.py | 57 +++++++++++++++++++---------------- 2 files changed, 36 insertions(+), 27 deletions(-) diff --git a/megatron/core/timers.py b/megatron/core/timers.py index 14aab032a9f..cd82928a8b1 100644 --- a/megatron/core/timers.py +++ b/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 @@ -303,7 +307,7 @@ def _get_global_min_max_time(self, names, reset, barrier, normalizer): 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() + 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): # filter out the ones we did not have any timings for diff --git a/megatron/training/training.py b/megatron/training/training.py index d8b57326f67..b82871a956d 100644 --- a/megatron/training/training.py +++ b/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', From 290d1aeb84756548530e403e9fbff728ee02a7f1 Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Wed, 12 Nov 2025 06:39:40 +0000 Subject: [PATCH 3/4] lint Signed-off-by: Xin Yao --- megatron/core/timers.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/megatron/core/timers.py b/megatron/core/timers.py index cd82928a8b1..03d1f2bd25b 100644 --- a/megatron/core/timers.py +++ b/megatron/core/timers.py @@ -297,7 +297,6 @@ 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 @@ -307,7 +306,9 @@ def _get_global_min_max_time(self, names, reset, barrier, normalizer): 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 + 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): # filter out the ones we did not have any timings for @@ -432,7 +433,7 @@ def log( if rank is None: rank = torch.distributed.get_world_size() - 1 if rank == torch.distributed.get_rank() and output_string is not None: - print(output_string, flush=True) + print(output_string, flush=True) # pylint: disable=W0141 def write( self, From 4d27836abc0289424bf58fcdb0c1abe24278a61c Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Tue, 18 Nov 2025 02:47:46 +0000 Subject: [PATCH 4/4] replace print with log Signed-off-by: Xin Yao --- megatron/core/timers.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/megatron/core/timers.py b/megatron/core/timers.py index 03d1f2bd25b..95365cc2802 100644 --- a/megatron/core/timers.py +++ b/megatron/core/timers.py @@ -2,6 +2,7 @@ """Megatron timers.""" +import logging import time from abc import ABC, abstractmethod from typing import List @@ -28,6 +29,8 @@ except: dist_all_gather_func = torch.distributed._all_gather_base +logger = logging.getLogger(__name__) + class TimerBase(ABC): """Timer base class.""" @@ -433,7 +436,7 @@ def log( if rank is None: rank = torch.distributed.get_world_size() - 1 if rank == torch.distributed.get_rank() and output_string is not None: - print(output_string, flush=True) # pylint: disable=W0141 + logger.info(output_string) def write( self,