From b044dde38eb5cce10c2202d11210563dd303a407 Mon Sep 17 00:00:00 2001 From: Philip Monk Date: Mon, 27 Jul 2026 16:22:02 -0700 Subject: [PATCH 1/2] Add expert GTP topology regression tests Signed-off-by: Philip Monk --- tests/unit_tests/training/test_param_norm.py | 109 ++++++++++++++++++- 1 file changed, 106 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/training/test_param_norm.py b/tests/unit_tests/training/test_param_norm.py index d58eef225d3..fa37e017c4f 100644 --- a/tests/unit_tests/training/test_param_norm.py +++ b/tests/unit_tests/training/test_param_norm.py @@ -19,6 +19,8 @@ def _build_tiny_moe_gpt( tensor_parallel_size: int, expert_parallel_size: int, expert_tensor_parallel_size: int, + tensor_parallel_num_weight_shards: int | None = None, + expert_tensor_parallel_num_weight_shards: int | None = None, bf16: bool = False, ) -> GPTModel: config = TransformerConfig( @@ -34,6 +36,8 @@ def _build_tiny_moe_gpt( tensor_model_parallel_size=tensor_parallel_size, expert_model_parallel_size=expert_parallel_size, expert_tensor_parallel_size=expert_tensor_parallel_size, + tensor_parallel_num_weight_shards=tensor_parallel_num_weight_shards, + expert_tensor_parallel_num_weight_shards=expert_tensor_parallel_num_weight_shards, sequence_parallel=tensor_parallel_size > 1, use_cpu_initialization=True, add_bias_linear=False, @@ -116,19 +120,106 @@ def test_moe_param_norm_counts_each_logical_parameter_once( Utils.destroy_model_parallel() +def test_moe_param_norm_uses_expert_gtp_topology_when_it_differs_from_dense_gtp(monkeypatch): + """Expert parameters must use EGTP even when EP, TP, and ETP alone do not distinguish them.""" + from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( + GTPShardedParam, + reset_gtp_state, + ) + from megatron.core.tensor_parallel.gtp_api import HAVE_GTP + + if not HAVE_GTP: + pytest.skip("GTP requires TransformerEngine >= 2.19") + if Utils.world_size < 2 or Utils.world_size % 2 != 0: + pytest.skip("test requires an even world size") + + monkeypatch.setattr( + common_utils, "get_args", lambda: SimpleNamespace(use_megatron_fsdp=False, bf16=False) + ) + + try: + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=1, + ) + reference_model = _build_tiny_moe_gpt( + tensor_parallel_size=1, + expert_parallel_size=1, + expert_tensor_parallel_size=1, + ) + _fill_parameters_with_ones(reference_model) + expected_numel = sum(param.numel() for param in reference_model.parameters()) + expected_norm = math.sqrt(expected_numel) + del reference_model + + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + expert_model_parallel_size=1, + expert_tensor_parallel_size=1, + gtp_remat_size=1, + expert_gtp_remat_size=2, + ) + model = _build_tiny_moe_gpt( + tensor_parallel_size=1, + expert_parallel_size=1, + expert_tensor_parallel_size=1, + tensor_parallel_num_weight_shards=1, + expert_tensor_parallel_num_weight_shards=2, + ) + _fill_parameters_with_ones(model) + + expert_params = [ + param for name, param in model.named_parameters() if ".experts." in name + ] + assert any(isinstance(param, GTPShardedParam) for param in expert_params) + + actual_norm = common_utils.calc_params_l2_norm(model) + + assert actual_norm == pytest.approx(expected_norm) + finally: + reset_gtp_state() + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("use_distributed_optimizer", (False, True), ids=("optimizer", "distopt")) @pytest.mark.parametrize( - ("tensor_parallel_size", "expert_parallel_size", "expert_tensor_parallel_size"), - ((2, 2, 1), (2, 1, 2), (4, 1, 2), (2, 1, 4)), - ids=("expert-parallel", "expert-tensor-parallel", "tp-larger-than-etp", "etp-larger-than-tp"), + ( + "tensor_parallel_size", + "expert_parallel_size", + "expert_tensor_parallel_size", + "gtp_weight_remat_size", + "expert_gtp_weight_remat_size", + ), + ( + (2, 2, 1, 1, 1), + (2, 1, 2, 1, 1), + (4, 1, 2, 1, 1), + (2, 1, 4, 1, 1), + (1, 1, 1, 1, 2), + ), + ids=( + "expert-parallel", + "expert-tensor-parallel", + "tp-larger-than-etp", + "etp-larger-than-tp", + "expert-gtp-differs-from-dense-gtp", + ), ) def test_moe_grad_norm_and_clipping_count_each_logical_gradient_once( tensor_parallel_size: int, expert_parallel_size: int, expert_tensor_parallel_size: int, + gtp_weight_remat_size: int, + expert_gtp_weight_remat_size: int, use_distributed_optimizer: bool, ): """Gradient clipping should use each logical parameter's gradient exactly once.""" + if expert_gtp_weight_remat_size > 1: + from megatron.core.tensor_parallel.gtp_api import HAVE_GTP + + if not HAVE_GTP: + pytest.skip("GTP requires TransformerEngine >= 2.19") if Utils.world_size < 4 or Utils.world_size % 4 != 0: pytest.skip("test requires a world size divisible by four") @@ -149,11 +240,19 @@ def test_moe_grad_norm_and_clipping_count_each_logical_gradient_once( tensor_model_parallel_size=tensor_parallel_size, expert_model_parallel_size=expert_parallel_size, expert_tensor_parallel_size=expert_tensor_parallel_size, + gtp_remat_size=gtp_weight_remat_size, + expert_gtp_remat_size=expert_gtp_weight_remat_size, ) model = _build_tiny_moe_gpt( tensor_parallel_size=tensor_parallel_size, expert_parallel_size=expert_parallel_size, expert_tensor_parallel_size=expert_tensor_parallel_size, + tensor_parallel_num_weight_shards=( + tensor_parallel_size * gtp_weight_remat_size + ), + expert_tensor_parallel_num_weight_shards=( + expert_tensor_parallel_size * expert_gtp_weight_remat_size + ), bf16=True, ) ddp_config = DistributedDataParallelConfig( @@ -199,4 +298,8 @@ def test_moe_grad_norm_and_clipping_count_each_logical_gradient_once( grads_checked += 1 assert grads_checked > 0 finally: + if expert_gtp_weight_remat_size > 1: + from megatron.core.tensor_parallel.generalized_tensor_parallelism import reset_gtp_state + + reset_gtp_state() Utils.destroy_model_parallel() From 4f1ef5fef0af8c959ba1b714fea3e949060e5977 Mon Sep 17 00:00:00 2001 From: Philip Monk Date: Mon, 27 Jul 2026 18:03:26 -0700 Subject: [PATCH 2/2] Use expert process groups when EGTP differs from GTP Signed-off-by: Philip Monk --- .../core/distributed/param_and_grad_buffer.py | 5 ++- .../core/extensions/transformer_engine.py | 8 +++- .../core/post_training/modelopt/layers.py | 1 + .../generalized_tensor_parallelism.py | 15 +++++-- megatron/core/tensor_parallel/layers.py | 4 +- .../test_grad_sync_with_expert_parallel.py | 44 +++++++++++++++---- tests/unit_tests/training/test_param_norm.py | 43 +++++++++++------- 7 files changed, 85 insertions(+), 35 deletions(-) diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 15247c2be53..8923728d2c2 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -888,8 +888,9 @@ def group_params_for_buffers( - param_dtype: storage dtype (torch.uint8 for FP8/NVFP4 parameters, else param.dtype). - grad_dtype: gradient reduction dtype (torch.float if grad_reduce_in_fp32, else param.dtype). - is_expert_parallel: whether the parameter uses the expert topology (param.allreduce == False), - which requires a separate buffer for the expert data-parallel group. This is true for experts - when expert-parallelism > 1 or expert-tensor-parallelism != tensor-parallelism. + which requires a separate buffer for the expert data-parallel group. This is true for experts + when expert-parallelism > 1, expert-tensor-parallelism != tensor-parallelism, or expert-GTP + != GTP. The param_indices track each parameter's position among same-dtype params (using the "fake" high-precision dtype for FP8/NVFP4 params), needed for loading non-native-fp8 diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 8f0117c68ec..39a314625ef 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -93,7 +93,7 @@ def _set_expert_parameter_attributes( ) -> None: """Set process-group and tensor-partition metadata on an expert TE module. - ``allreduce=False`` selects EDP for gradient reduction. + ``allreduce=False`` selects the expert topology, including EDP for gradient reduction. Weights and biases, including TEGroupedLinear's numbered parameters, are also marked as TP-partitioned according to ``parallel_mode``; row-parallel biases remain replicated. @@ -104,7 +104,7 @@ def _set_expert_parameter_attributes( Args: module: Transformer Engine module whose direct parameters should be marked. parallel_mode: Tensor-parallel mode used by the module (``"column"``, ``"row"``, or None). - use_expert_pgs: Whether to use EP/ETP/EDP process groups instead of TP/CP/DP. + use_expert_pgs: Whether to use EP/ETP/EGTP/EDP process groups instead of TP/GTP/CP/DP. """ for name, param in module.named_parameters(recurse=False): param.allreduce = not use_expert_pgs @@ -982,6 +982,7 @@ def __init__( use_expert_pgs = is_expert and ( self.expert_parallel or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size + or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size ) if is_expert: rng_tracker_name = get_expert_parallel_rng_tracker_name() @@ -1475,6 +1476,7 @@ def __init__( use_expert_pgs = ( config.expert_model_parallel_size > 1 or config.expert_tensor_parallel_size != config.tensor_model_parallel_size + or config.expert_gtp_weight_remat_size != config.gtp_weight_remat_size ) _set_expert_parameter_attributes(self, "column", use_expert_pgs) @@ -1731,6 +1733,7 @@ def __init__( use_expert_pgs = ( config.expert_model_parallel_size > 1 or config.expert_tensor_parallel_size != config.tensor_model_parallel_size + or config.expert_gtp_weight_remat_size != config.gtp_weight_remat_size ) _set_expert_parameter_attributes(self, "row", use_expert_pgs) @@ -2163,6 +2166,7 @@ def __init__( use_expert_pgs = is_expert and ( self.expert_parallel or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size + or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size ) if is_expert: extra_kwargs["rng_tracker_name"] = get_expert_parallel_rng_tracker_name() diff --git a/megatron/core/post_training/modelopt/layers.py b/megatron/core/post_training/modelopt/layers.py index 5f1a746e95b..cd8919e50e7 100644 --- a/megatron/core/post_training/modelopt/layers.py +++ b/megatron/core/post_training/modelopt/layers.py @@ -163,6 +163,7 @@ def __init__( self.config.expert_model_parallel_size > 1 or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size + or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size ) setattr(param, "allreduce", not use_expert_groups) else: diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index 08f17e54996..768402f1ea4 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -501,11 +501,14 @@ def _gtp_slice_one_param(param, gtp_remat_group, *, name=""): shard = tensor[gtp_rank * shard_size : (gtp_rank + 1) * shard_size] gtp_shard = GTPShardedParam(shard.clone()) gtp_shard.pad_length = pad_length - # Preserve the source weight's TP attributes (dropped when wrapping into GTPShardedParam), - # so param_is_not_tensor_parallel_duplicate still classifies it without GTP-specific code. - from megatron.core.tensor_parallel import copy_tensor_model_parallel_attributes + # Preserve duplicate-filtering metadata dropped when wrapping into GTPShardedParam. + from megatron.core.tensor_parallel import ( + copy_gtp_attributes, + copy_tensor_model_parallel_attributes, + ) copy_tensor_model_parallel_attributes(gtp_shard, param) + copy_gtp_attributes(gtp_shard, param) return gtp_shard @@ -539,10 +542,14 @@ def _gtp_wrap_bf16_shard(module, name, param): :func:`_gtp_slice_one_param`, which slices a full weight — this only wraps it, no slicing. Returns the new param (also swapped into the module). """ - from megatron.core.tensor_parallel import copy_tensor_model_parallel_attributes + from megatron.core.tensor_parallel import ( + copy_gtp_attributes, + copy_tensor_model_parallel_attributes, + ) gtp_shard = GTPShardedParam(param.data) copy_tensor_model_parallel_attributes(gtp_shard, param) + copy_gtp_attributes(gtp_shard, param) delattr(module, name) module._parameters[name] = gtp_shard return gtp_shard diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index 3ff635b362b..a42cbd05841 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -100,7 +100,7 @@ def param_is_not_tensor_parallel_duplicate(param, tp_group=None, expert_tp_group """ if hasattr(param, "tensor_model_parallel") and param.tensor_model_parallel: return True - # allreduce=False marks parameters reduced over expert DP, so filter their duplicates over ETP. + # allreduce=False marks parameters using the expert topology, so filter duplicates over ETP. if not getattr(param, "allreduce", True) and expert_tp_group is not None: tp_group = expert_tp_group # Prefer provided tp_group when available (new explicit path). @@ -961,6 +961,7 @@ def __init__( use_expert_pgs = self.is_expert and ( self.expert_parallel or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size + or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size ) self.output_size_per_partition = divide(output_size, world_size) @@ -1379,6 +1380,7 @@ def __init__( use_expert_pgs = self.is_expert and ( self.expert_parallel or self.config.expert_tensor_parallel_size != self.config.tensor_model_parallel_size + or self.config.expert_gtp_weight_remat_size != self.config.gtp_weight_remat_size ) setattr(self.weight, "allreduce", not use_expert_pgs) diff --git a/tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py b/tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py index fe56317f39a..a088016d52b 100644 --- a/tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py +++ b/tests/unit_tests/distributed/test_grad_sync_with_expert_parallel.py @@ -142,26 +142,44 @@ def _build_expert_linear(implementation: str, config: TransformerConfig) -> torc @pytest.mark.parametrize( - ("tensor_model_parallel_size", "expert_tensor_parallel_size"), [(2, 1), (1, 2)] + ( + "tensor_model_parallel_size", + "expert_tensor_parallel_size", + "gtp_weight_remat_size", + "expert_gtp_weight_remat_size", + ), + [(2, 1, 1, 1), (1, 2, 1, 1), (1, 1, 1, 2)], ) @pytest.mark.parametrize( "implementation", ["native", "transformer_engine", "transformer_engine_grouped"] ) def test_expert_grad_sync_uses_expert_data_parallel_group( - implementation: str, tensor_model_parallel_size: int, expert_tensor_parallel_size: int + implementation: str, + tensor_model_parallel_size: int, + expert_tensor_parallel_size: int, + gtp_weight_remat_size: int, + expert_gtp_weight_remat_size: int, ): - """Expert gradients must not be reduced over ordinary DP when ETP differs from TP.""" + """Expert gradients must use expert DP when the expert and dense topologies differ.""" + if expert_gtp_weight_remat_size > 1: + from megatron.core.tensor_parallel.gtp_api import HAVE_GTP + + if not HAVE_GTP: + pytest.skip("GTP requires TransformerEngine >= 2.19") if Utils.world_size < 4 or Utils.world_size % 4 != 0: pytest.skip("Test requires a world size divisible by four") if Utils.world_size > 16: pytest.skip("Rank-encoded gradients are intended for small unit-test world sizes") - Utils.initialize_model_parallel( - tensor_model_parallel_size=tensor_model_parallel_size, - expert_model_parallel_size=1, - expert_tensor_parallel_size=expert_tensor_parallel_size, - ) try: + Utils.initialize_model_parallel( + tensor_model_parallel_size=tensor_model_parallel_size, + expert_model_parallel_size=1, + expert_tensor_parallel_size=expert_tensor_parallel_size, + gtp_remat_size=gtp_weight_remat_size, + expert_gtp_remat_size=expert_gtp_weight_remat_size, + ) + # Per-token loss leaves DDP's pre-collective gradient scaling at one. config = TransformerConfig( num_layers=1, @@ -174,6 +192,10 @@ def test_expert_grad_sync_uses_expert_data_parallel_group( tensor_model_parallel_size=tensor_model_parallel_size, expert_model_parallel_size=1, expert_tensor_parallel_size=expert_tensor_parallel_size, + tensor_parallel_num_weight_shards=(tensor_model_parallel_size * gtp_weight_remat_size), + expert_tensor_parallel_num_weight_shards=( + expert_tensor_parallel_size * expert_gtp_weight_remat_size + ), calculate_per_token_loss=True, gradient_accumulation_fusion=False, perform_initialization=False, @@ -194,7 +216,7 @@ def test_expert_grad_sync_uses_expert_data_parallel_group( ) expert_dp_group = parallel_state.get_expert_data_parallel_group( - partial_expert_data_parallel=True + with_gtp_remat=False, partial_expert_data_parallel=True ) ordinary_dp_group = parallel_state.get_data_parallel_group( with_context_parallel=True, partial_data_parallel=True @@ -222,6 +244,10 @@ def test_expert_grad_sync_uses_expert_data_parallel_group( assert len(model.expert_parallel_buffers) == 1 assert all(param.allreduce is False for param in model.parameters()) finally: + if expert_gtp_weight_remat_size > 1: + from megatron.core.tensor_parallel.generalized_tensor_parallelism import reset_gtp_state + + reset_gtp_state() Utils.destroy_model_parallel() diff --git a/tests/unit_tests/training/test_param_norm.py b/tests/unit_tests/training/test_param_norm.py index fa37e017c4f..a5088460b71 100644 --- a/tests/unit_tests/training/test_param_norm.py +++ b/tests/unit_tests/training/test_param_norm.py @@ -123,8 +123,10 @@ def test_moe_param_norm_counts_each_logical_parameter_once( def test_moe_param_norm_uses_expert_gtp_topology_when_it_differs_from_dense_gtp(monkeypatch): """Expert parameters must use EGTP even when EP, TP, and ETP alone do not distinguish them.""" from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( + GTP_CONFIG, GTPShardedParam, reset_gtp_state, + update_gtp_config, ) from megatron.core.tensor_parallel.gtp_api import HAVE_GTP @@ -136,6 +138,9 @@ def test_moe_param_norm_uses_expert_gtp_topology_when_it_differs_from_dense_gtp( monkeypatch.setattr( common_utils, "get_args", lambda: SimpleNamespace(use_megatron_fsdp=False, bf16=False) ) + # Keep the all-ones assertion focused on topology rather than physical GTP padding. + original_pad_for_alignment = GTP_CONFIG.pad_for_alignment + update_gtp_config(pad_for_alignment=0) try: Utils.initialize_model_parallel( @@ -144,9 +149,7 @@ def test_moe_param_norm_uses_expert_gtp_topology_when_it_differs_from_dense_gtp( expert_tensor_parallel_size=1, ) reference_model = _build_tiny_moe_gpt( - tensor_parallel_size=1, - expert_parallel_size=1, - expert_tensor_parallel_size=1, + tensor_parallel_size=1, expert_parallel_size=1, expert_tensor_parallel_size=1 ) _fill_parameters_with_ones(reference_model) expected_numel = sum(param.numel() for param in reference_model.parameters()) @@ -169,15 +172,14 @@ def test_moe_param_norm_uses_expert_gtp_topology_when_it_differs_from_dense_gtp( ) _fill_parameters_with_ones(model) - expert_params = [ - param for name, param in model.named_parameters() if ".experts." in name - ] + expert_params = [param for name, param in model.named_parameters() if ".experts." in name] assert any(isinstance(param, GTPShardedParam) for param in expert_params) actual_norm = common_utils.calc_params_l2_norm(model) assert actual_norm == pytest.approx(expected_norm) finally: + update_gtp_config(pad_for_alignment=original_pad_for_alignment) reset_gtp_state() Utils.destroy_model_parallel() @@ -191,13 +193,7 @@ def test_moe_param_norm_uses_expert_gtp_topology_when_it_differs_from_dense_gtp( "gtp_weight_remat_size", "expert_gtp_weight_remat_size", ), - ( - (2, 2, 1, 1, 1), - (2, 1, 2, 1, 1), - (4, 1, 2, 1, 1), - (2, 1, 4, 1, 1), - (1, 1, 1, 1, 2), - ), + ((2, 2, 1, 1, 1), (2, 1, 2, 1, 1), (4, 1, 2, 1, 1), (2, 1, 4, 1, 1), (1, 1, 1, 1, 2)), ids=( "expert-parallel", "expert-tensor-parallel", @@ -223,6 +219,17 @@ def test_moe_grad_norm_and_clipping_count_each_logical_gradient_once( if Utils.world_size < 4 or Utils.world_size % 4 != 0: pytest.skip("test requires a world size divisible by four") + original_pad_for_alignment = None + if expert_gtp_weight_remat_size > 1: + from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( + GTP_CONFIG, + update_gtp_config, + ) + + # Keep the all-ones assertion focused on topology rather than physical GTP padding. + original_pad_for_alignment = GTP_CONFIG.pad_for_alignment + update_gtp_config(pad_for_alignment=0) + try: Utils.initialize_model_parallel( tensor_model_parallel_size=1, @@ -247,9 +254,7 @@ def test_moe_grad_norm_and_clipping_count_each_logical_gradient_once( tensor_parallel_size=tensor_parallel_size, expert_parallel_size=expert_parallel_size, expert_tensor_parallel_size=expert_tensor_parallel_size, - tensor_parallel_num_weight_shards=( - tensor_parallel_size * gtp_weight_remat_size - ), + tensor_parallel_num_weight_shards=(tensor_parallel_size * gtp_weight_remat_size), expert_tensor_parallel_num_weight_shards=( expert_tensor_parallel_size * expert_gtp_weight_remat_size ), @@ -299,7 +304,11 @@ def test_moe_grad_norm_and_clipping_count_each_logical_gradient_once( assert grads_checked > 0 finally: if expert_gtp_weight_remat_size > 1: - from megatron.core.tensor_parallel.generalized_tensor_parallelism import reset_gtp_state + from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( + reset_gtp_state, + update_gtp_config, + ) + update_gtp_config(pad_for_alignment=original_pad_for_alignment) reset_gtp_state() Utils.destroy_model_parallel()