diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 4439f123852..1b13cdbe3cc 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -894,8 +894,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 27193ebf827..12dbfba6992 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, add_bias_linear: bool = False, ) -> GPTModel: @@ -36,6 +38,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=add_bias_linear, @@ -119,22 +123,116 @@ 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 ( + GTP_CONFIG, + GTPShardedParam, + reset_gtp_state, + update_gtp_config, + ) + 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) + ) + # 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, + 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: + update_gtp_config(pad_for_alignment=original_pad_for_alignment) + 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_gradient_stats_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 norm, clipping, and zero count should include each logical gradient 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") + 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, @@ -152,11 +250,17 @@ def test_moe_gradient_stats_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( @@ -211,6 +315,14 @@ def test_moe_gradient_stats_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, + update_gtp_config, + ) + + update_gtp_config(pad_for_alignment=original_pad_for_alignment) + reset_gtp_state() Utils.destroy_model_parallel()