From 3e72b96672b1d0719619be8ca3cab0af636ac1ba Mon Sep 17 00:00:00 2001 From: guapisolo Date: Tue, 11 Aug 2026 23:32:26 +0000 Subject: [PATCH] fix(optimizer): use expert TP group for grad stats Expert parameters use expert tensor-parallel ownership, but optimizer gradient filtering only received the dense TP group. This omitted valid expert shards when TP and ETP ownership differed and underestimated clipping norms. Preserve allreduce metadata on optimizer parameter copies and pass both TP groups through gradient-norm and zero-count filtering. Signed-off-by: guapisolo --- megatron/core/optimizer/__init__.py | 9 ++-- megatron/core/optimizer/clip_grads.py | 5 +- megatron/core/optimizer/optimizer.py | 9 +++- tests/unit_tests/test_optimizer.py | 78 +++++++++++++++++++++++++++ 4 files changed, 95 insertions(+), 6 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 27b675d1b8d..99c644fa3a6 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -686,11 +686,12 @@ def init_state_fn(opt, config=None): setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group) if pg_collection is None or not hasattr(pg_collection, 'tp'): - tp_group = parallel_state.get_tensor_model_parallel_group() - else: - tp_group = pg_collection.tp - # TODO(M4): plumb tp_group through optimizer constructors so this setattr disappears. + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + tp_group = pg_collection.tp + expert_tp_group = getattr(pg_collection, 'expt_tp', tp_group) + # TODO(M4): plumb TP groups through optimizer constructors so these setattrs disappear. setattr(optimizer, 'tp_group', tp_group) + setattr(optimizer, 'expert_tp_group', expert_tp_group) return optimizer diff --git a/megatron/core/optimizer/clip_grads.py b/megatron/core/optimizer/clip_grads.py index 3c5491d39a1..aafbb629bf7 100644 --- a/megatron/core/optimizer/clip_grads.py +++ b/megatron/core/optimizer/clip_grads.py @@ -201,6 +201,7 @@ def count_zeros_fp32( grad_stats_parallel_group: torch.distributed.ProcessGroup, use_decoupled_grad: bool = False, tp_group: Optional[torch.distributed.ProcessGroup] = None, + expert_tp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> float: """Counts the number of zero values in the gradients of the given parameters. @@ -245,7 +246,9 @@ def count_zeros_fp32( total_num_zeros += num_zeros continue is_not_shared = param_is_not_shared(param) - is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param, tp_group=tp_group) + is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate( + param, tp_group=tp_group, expert_tp_group=expert_tp_group + ) if grad_not_none and is_not_shared and is_not_tp_duplicate: grad_obj = getattr(param, grad_attr) data_parallel_group = get_data_parallel_group_if_dtensor(grad_obj, data_parallel_group) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index d045f64dcc3..b72b8a08f18 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -124,6 +124,8 @@ def _is_separate_grad_norm_group(grad_norm_group: Optional[str]) -> bool: def copy_optimizer_param_metadata(destination: torch.Tensor, source: torch.Tensor) -> None: """Copy optimizer-relevant metadata when creating param views/copies.""" + if hasattr(source, 'allreduce'): + destination.allreduce = source.allreduce if hasattr(source, 'shared'): destination.shared = source.shared if hasattr(source, GRAD_NORM_GROUP_ATTR): @@ -213,7 +215,9 @@ def _filter_grads_for_norm( grad_not_none = grad is not None is_not_shared = param_is_not_shared(param) is_not_tp_duplicate = tensor_parallel.param_is_not_tensor_parallel_duplicate( - param, getattr(self, 'tp_group', None) + param, + tp_group=getattr(self, 'tp_group', None), + expert_tp_group=getattr(self, 'expert_tp_group', None), ) if grad_not_none and is_not_shared and is_not_tp_duplicate: grads_for_norm.append(grad) @@ -380,6 +384,7 @@ def count_zeros(self) -> float: and getattr(params[0], "__fsdp_param__", False) ), tp_group=getattr(self, 'tp_group', None), + expert_tp_group=getattr(self, 'expert_tp_group', None), ) @abstractmethod @@ -1695,6 +1700,8 @@ def count_zeros(self): self.config.use_precision_aware_optimizer and getattr(params[0], "__fsdp_param__", False) ), + tp_group=getattr(self.chained_optimizers[0], 'tp_group', None), + expert_tp_group=getattr(self.chained_optimizers[0], 'expert_tp_group', None), ) else: num_zeros_in_grad = 0 diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 5b3e69c23b8..c3bc5fcca95 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import math import os from unittest.mock import patch @@ -12,6 +13,7 @@ # FP8 recipe will be used to test precision-aware-optimizer. from transformer_engine.pytorch.fp8 import fp8_autocast +from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig from megatron.core.optimizer import ( ChainedOptimizer, @@ -23,6 +25,7 @@ get_megatron_optimizer, get_standard_config_overrides, ) +from megatron.core.optimizer.optimizer import copy_optimizer_param_metadata from megatron.core.optimizer_param_scheduler import ParamGroupOverride from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer import TransformerConfig @@ -69,6 +72,81 @@ def forward(self, x): return x +def test_copy_optimizer_param_metadata_preserves_allreduce(): + source = torch.empty(1) + destination = torch.empty_like(source) + source.allreduce = False + + copy_optimizer_param_metadata(destination, source) + + assert destination.allreduce is False + + +@pytest.mark.skipif( + int(os.getenv('WORLD_SIZE', '1')) < 2, reason="test requires at least two distributed ranks" +) +@pytest.mark.parametrize("use_distributed_optimizer", (False, True), ids=("optimizer", "distopt")) +def test_expert_grad_stats_use_expert_tp_group(use_distributed_optimizer: bool): + """Expert grad stats must deduplicate over ETP, not dense TP.""" + world_size = int(os.environ['WORLD_SIZE']) + rank = int(os.environ['RANK']) + _init_distributed(world_size, rank) + + class DenseAndExpertParameters(nn.Module): + def __init__(self): + super().__init__() + self.dense = nn.Parameter(torch.ones(1, dtype=torch.bfloat16, device='cuda')) + self.expert = nn.Parameter(torch.ones(1, dtype=torch.bfloat16, device='cuda')) + self.expert.allreduce = False + + try: + Utils.initialize_model_parallel( + tensor_model_parallel_size=world_size, + expert_model_parallel_size=world_size, + expert_tensor_parallel_size=1, + ) + model = DenseAndExpertParameters() + transformer_config = TransformerConfig(num_attention_heads=1, num_layers=1) + model = DistributedDataParallel( + transformer_config, + DistributedDataParallelConfig(use_distributed_optimizer=use_distributed_optimizer), + model, + ) + optimizer = get_megatron_optimizer( + OptimizerConfig( + optimizer='adam', + lr=0.0, + bf16=True, + clip_grad=0.0, + log_num_zeros_in_grad=True, + use_distributed_optimizer=use_distributed_optimizer, + ), + [model], + use_gloo_process_groups=False, + ) + + dense_optimizer, expert_optimizer = optimizer.chained_optimizers + assert dense_optimizer.tp_group is parallel_state.get_tensor_model_parallel_group() + assert expert_optimizer.expert_tp_group is parallel_state.get_expert_tensor_parallel_group() + + for param in model.parameters(): + param.main_grad.fill_(1.0) + assert optimizer.prepare_grads() is False + + expected_grad_norm = math.sqrt(1 + world_size) + actual_grad_norm = optimizer.get_grad_norm() + if isinstance(actual_grad_norm, torch.Tensor): + actual_grad_norm = actual_grad_norm.item() + assert actual_grad_norm == pytest.approx(expected_grad_norm) + + for param in model.parameters(): + param.main_grad.zero_() + assert optimizer.prepare_grads() is False + assert optimizer.count_zeros() == 1 + world_size + finally: + Utils.destroy_model_parallel() + + @patch('torch.distributed.get_world_size', return_value=1) @patch( 'torch.distributed.all_gather_object', lambda output_list, obj: output_list.__setitem__(0, obj)