From 0d0ddf3ea273959bc3cefd80b6e1722202a76491 Mon Sep 17 00:00:00 2001 From: Li Tao Date: Sun, 2 Nov 2025 23:36:48 -0800 Subject: [PATCH 1/2] Remove redundant reduce in aux_loss logging Signed-off-by: Li Tao --- megatron/core/transformer/moe/moe_utils.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 5a0793ef5b9..a8de09944fc 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -770,12 +770,15 @@ def reduce_aux_losses_tracker_across_ranks(track_names: Optional[List[str]] = No torch.distributed.all_reduce( values, group=tracker[name]['avg_group'], op=torch.distributed.ReduceOp.AVG ) - # This ensures proper loss averaging across all ranks including CP ranks - torch.distributed.all_reduce( - values, - group=parallel_state.get_data_parallel_group(with_context_parallel=True), - op=torch.distributed.ReduceOp.AVG, - ) + # Average aux losses across data parallel ranks. + # The `global_load_balancing_loss` already uses `tp_dp_cp_group` in `reduce_group`, + # so we don't need to reduce it again. Others use tp_cp_group in `reduce_group`. + if name != "global_load_balancing_loss": + torch.distributed.all_reduce( + values, + group=parallel_state.get_data_parallel_group(with_context_parallel=False), + op=torch.distributed.ReduceOp.AVG, + ) def track_moe_metrics( From d0d5c497f25fc29e72d46990069e26c36d1244a6 Mon Sep 17 00:00:00 2001 From: lit Date: Tue, 4 Nov 2025 00:47:03 -0800 Subject: [PATCH 2/2] Fix comment Signed-off-by: lit --- megatron/core/transformer/moe/moe_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index a8de09944fc..e8c6a05340c 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -772,7 +772,7 @@ def reduce_aux_losses_tracker_across_ranks(track_names: Optional[List[str]] = No ) # Average aux losses across data parallel ranks. # The `global_load_balancing_loss` already uses `tp_dp_cp_group` in `reduce_group`, - # so we don't need to reduce it again. Others use tp_cp_group in `reduce_group`. + # so we don't need to reduce it again. Others use `tp_cp_group` in `reduce_group`. if name != "global_load_balancing_loss": torch.distributed.all_reduce( values,