diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index ef868ebbdb8..54ad4f35611 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -437,10 +437,25 @@ def custom_forward(hidden_states, intermediate_tensors): def backward_dw(self, routed_experts: bool = True, shared_experts: bool = False): """Compute weight gradients for experts and shared experts.""" + # TODO(Wohox): replace the "routed_experts" and "shared_experts" arguments with better + # naming to better explain that they are actually from different fine-grained callables, + # or use scanning to decide which backward_dw should be called. if routed_experts: self.experts.backward_dw() - if shared_experts and self.use_shared_expert and not self.shared_expert_overlap: - self.shared_experts.backward_dw() + if self.config.moe_latent_size: + # TODO(Wohox): fc2_latent_proj forward and backward are executed in comm stream, + # so we execute its backward_dw in the comm stream too. But this may harm the + # EP overlap performance. Better to check if there is a better way to handle this. + from megatron.core.pipeline_parallel.utils import get_comm_stream + + comm_stream = get_comm_stream() + with torch.cuda.stream(comm_stream): + self.fc2_latent_proj.backward_dw() + if shared_experts: + if self.use_shared_expert and not self.shared_expert_overlap: + self.shared_experts.backward_dw() + if self.config.moe_latent_size: + self.fc1_latent_proj.backward_dw() def set_for_recompute_pre_mlp_layernorm(self): """Set the MoE layer for recompute pre_mlp_layernorm. Only needed for fp8/fp4.""" diff --git a/megatron/training/training.py b/megatron/training/training.py index fbc267fba82..bec74ca7c36 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -364,6 +364,7 @@ def transformer_flops(): if args.moe_ffn_hidden_size is not None else args.ffn_hidden_size ) + moe_latent_size = args.moe_latent_size shared_expert_ffn_hidden_size = ( 0 if args.moe_shared_expert_intermediate_size is None @@ -545,7 +546,20 @@ def transformer_flops(): (args.ffn_hidden_size * ffn_expansion_factor) * num_dense_layers # routed experts - + (moe_ffn_hidden_size * num_experts_routed_to * ffn_expansion_factor) + + ( + (moe_ffn_hidden_size * num_experts_routed_to * ffn_expansion_factor) + if moe_latent_size is None + else ( + ( + moe_ffn_hidden_size + * num_experts_routed_to + * ffn_expansion_factor + * moe_latent_size + / args.hidden_size + ) # Routed experts run on moe_latent_size. + + 2 * moe_latent_size # Up proj and down proj. + ) + ) * num_moe_layers # Shared Experts. + (shared_expert_ffn_hidden_size * ffn_expansion_factor)