diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index a132e1be3a1..376f9a1f1c0 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -514,6 +514,13 @@ def __init__( opt, config, None, init_state_fn_list[i] if init_state_fn_list else None ) + self.tp_group = self.pg_collection.tp + self.expert_tp_group = getattr(self.pg_collection, 'expt_tp', self.tp_group) + for optimizer in optimizers: + # Child optimizers perform TP duplicate filtering when collecting gradients. + optimizer.tp_group = self.tp_group + optimizer.expert_tp_group = self.expert_tp_group + super().__init__(optimizers) # Assign self.model_chunks AFTER super().__init__: ChainedOptimizer.__init__ @@ -857,6 +864,8 @@ def count_zeros(self): params, grad_stats_parallel_group=None, use_decoupled_grad=self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8, + tp_group=self.tp_group, + expert_tp_group=self.expert_tp_group, ) def start_param_sync_for_bucket_group_subset(self) -> None: diff --git a/tests/unit_tests/training/test_param_norm.py b/tests/unit_tests/training/test_param_norm.py index d58eef225d3..fbeb3d1e765 100644 --- a/tests/unit_tests/training/test_param_norm.py +++ b/tests/unit_tests/training/test_param_norm.py @@ -20,6 +20,7 @@ def _build_tiny_moe_gpt( expert_parallel_size: int, expert_tensor_parallel_size: int, bf16: bool = False, + add_bias_linear: bool = False, ) -> GPTModel: config = TransformerConfig( num_layers=1, @@ -28,7 +29,8 @@ def _build_tiny_moe_gpt( ffn_hidden_size=16, num_moe_experts=2, moe_ffn_hidden_size=16, - moe_shared_expert_intermediate_size=16, + # Shared experts do not support linear biases. + moe_shared_expert_intermediate_size=None if add_bias_linear else 16, moe_router_topk=1, moe_router_pre_softmax=True, tensor_model_parallel_size=tensor_parallel_size, @@ -36,7 +38,7 @@ def _build_tiny_moe_gpt( expert_tensor_parallel_size=expert_tensor_parallel_size, sequence_parallel=tensor_parallel_size > 1, use_cpu_initialization=True, - add_bias_linear=False, + add_bias_linear=add_bias_linear, normalization="RMSNorm", moe_grouped_gemm=True, bf16=bf16, @@ -51,7 +53,8 @@ def _build_tiny_moe_gpt( max_sequence_length=8, position_embedding_type="rope", ) - assert any(".shared_experts." in name for name, _ in model.named_parameters()) + if not add_bias_linear: + assert any(".shared_experts." in name for name, _ in model.named_parameters()) return model.cuda() @@ -200,3 +203,86 @@ def test_moe_grad_norm_and_clipping_count_each_logical_gradient_once( assert grads_checked > 0 finally: Utils.destroy_model_parallel() + + +def test_layer_wise_muon_grad_norm_uses_expert_tp_group_for_row_parallel_bias(): + """LayerWise Muon must deduplicate replicated expert FC2 bias grads over ETP. + + With TP=2, EP=2, and ETP=1, every rank is ETP rank zero. The two EP ranks own + distinct row-parallel expert biases, so both gradients must contribute to the global + norm. Falling back to the regular TP rank drops the expert on TP rank one and + undercounts the squared norm by a factor of two. + """ + from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer + from megatron.core.process_groups_config import ProcessGroupCollection + + if Utils.world_size < 4 or Utils.world_size % 4 != 0: + pytest.skip("test requires a world size divisible by four") + + tensor_parallel_size = 2 + expert_parallel_size = 2 + expert_tensor_parallel_size = 1 + + try: + Utils.initialize_model_parallel( + tensor_model_parallel_size=tensor_parallel_size, + expert_model_parallel_size=expert_parallel_size, + expert_tensor_parallel_size=expert_tensor_parallel_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, + bf16=True, + add_bias_linear=True, + ) + + expert_fc2_biases = [ + param + for name, param in model.named_parameters() + if ".experts." in name and ".linear_fc2.bias" in name + ] + assert len(expert_fc2_biases) == model.config.num_moe_experts // expert_parallel_size + for parameter in expert_fc2_biases: + assert parameter.ndim == 1 + assert parameter.allreduce is False + assert parameter.tensor_model_parallel is False + + model = DistributedDataParallel( + model.config, DistributedDataParallelConfig(use_distributed_optimizer=False), model + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + optimizer = get_megatron_optimizer( + OptimizerConfig( + optimizer="muon", + lr=0.0, + weight_decay=0.0, + bf16=True, + use_distributed_optimizer=False, + use_layer_wise_distributed_optimizer=True, + muon_tp_mode="duplicated", + ), + [model], + use_gloo_process_groups=False, + pg_collection=pg_collection, + ) + + assert isinstance(optimizer, LayerWiseDistributedOptimizer) + assert pg_collection.tp.size() == tensor_parallel_size + assert pg_collection.expt_tp.size() == expert_tensor_parallel_size + + for parameter in model.parameters(): + parameter.main_grad.zero_() + for parameter in expert_fc2_biases: + parameter.main_grad.fill_(1.0) + assert optimizer.prepare_grads() is False + + actual_norm = optimizer.get_grad_norm() + actual_norm_value = ( + actual_norm.item() if isinstance(actual_norm, torch.Tensor) else actual_norm + ) + expected_norm = math.sqrt(model.config.num_moe_experts * model.config.hidden_size) + + assert actual_norm_value == pytest.approx(expected_norm) + finally: + Utils.destroy_model_parallel()