diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 3ee4fc30488..5c282c9c82d 100755 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1743,7 +1743,8 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool: self.te_quant_params, self.training, is_context_quantized ) - def forward(self, x, m_splits, m_splits_gpu=None): + def forward(self, x, m_splits, m_splits_gpu=None, actual_m_splits=None, + unpad_output=False): """Forward.""" _is_first_microbatch = ( None if self.disable_parameter_transpose_cache else self.is_first_microbatch @@ -1754,8 +1755,18 @@ def forward(self, x, m_splits, m_splits_gpu=None): # Added in TE commit: https://github.com/ROCm/TransformerEngine/commit/2776c33 if is_te_min_version("2.7.0", check_equality=False): if class_has_method_param(te.pytorch.GroupedLinear, "forward", "m_splits_tensor"): + fused_padding_kwargs = {} + if actual_m_splits is not None and class_has_method_param( + te.pytorch.GroupedLinear, "forward", "actual_m_splits" + ): + fused_padding_kwargs["actual_m_splits"] = actual_m_splits + if unpad_output and class_has_method_param( + te.pytorch.GroupedLinear, "forward", "unpad_output" + ): + fused_padding_kwargs["unpad_output"] = unpad_output with quant_context: - out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch, m_splits_tensor=m_splits_gpu) + out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch, + m_splits_tensor=m_splits_gpu, **fused_padding_kwargs) else: warnings.warn( "Transformer Engine is missing `m_splits_tensor` parameter in GroupedLinear.forward; " diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 21efa6e7390..25992e98839 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -61,6 +61,16 @@ logger = logging.getLogger(__name__) +is_rocm = hasattr(torch.version, 'hip') and torch.version.hip is not None +HAVE_FUSED_PADDING_MCT = False +if HAVE_TE and is_rocm: + HAVE_FUSED_PADDING_MCT = hasattr(Fp8Padding, 'compute_padded_splits') + if not HAVE_FUSED_PADDING_MCT: + logger.warning( + "Fp8Padding.compute_padded_splits not found; " + "using unfused BF16 padding path for MoE." + ) + class GroupedMLP(MegatronModule): """An efficient implementation of the Experts layer using GroupedGEMM. @@ -643,14 +653,20 @@ def forward( output (torch.Tensor): The output of the local experts. """ tokens_per_expert_gpu = None + actual_tokens_per_expert = None if isinstance(tokens_per_expert, tuple): tokens_per_expert, tokens_per_expert_gpu = tokens_per_expert tokens_per_expert = tokens_per_expert.tolist() if self.config.fp8 or self.config.fp4: actual_tokens_per_expert = tokens_per_expert - permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( - permuted_local_hidden_states, tokens_per_expert - ) + if HAVE_FUSED_PADDING_MCT: + # ROCm: fuse padding into cast_transpose — compute padded + # splits without allocating a BF16 intermediate. + tokens_per_expert = self.quantization_padding.compute_padded_splits(tokens_per_expert) + else: + permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( + permuted_local_hidden_states, tokens_per_expert + ) permuted_probs, _ = self.quantization_padding( permuted_probs.unsqueeze(-1), actual_tokens_per_expert ) @@ -667,11 +683,12 @@ def forward( # Probs already applied, so reset to 1. permuted_probs = torch.ones_like(permuted_probs) + _fused_kwargs = {"actual_m_splits": actual_tokens_per_expert} if HAVE_FUSED_PADDING_MCT and actual_tokens_per_expert is not None else {} with off_interface( self.offload_expert_fc1, permuted_local_hidden_states, "expert_fc1" ) as permuted_local_hidden_states: fc1_output, bias_parallel = self.linear_fc1( - permuted_local_hidden_states, tokens_per_expert, tokens_per_expert_gpu + permuted_local_hidden_states, tokens_per_expert, tokens_per_expert_gpu, **_fused_kwargs, ) if self.offload_expert_fc1: fc1_output = off_interface.group_commit( @@ -748,7 +765,18 @@ def glu(x): with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: bias_act_output = bias_act_func(fc1_output, bias_parallel, permuted_probs) - output, output_bias = self.linear_fc2(bias_act_output, tokens_per_expert, tokens_per_expert_gpu) + _can_unpad_fc2 = HAVE_FUSED_PADDING_MCT and actual_tokens_per_expert is not None and not self.config.add_bias_linear + if _can_unpad_fc2: + _fc2_fused_kwargs = { + "actual_m_splits": actual_tokens_per_expert, + "unpad_output": True, + } + else: + _fc2_fused_kwargs = {} + output, output_bias = self.linear_fc2( + bias_act_output, tokens_per_expert, tokens_per_expert_gpu, + **_fc2_fused_kwargs, + ) if self.activation_recompute: self.activation_checkpoint.discard_output_and_register_recompute(output) @@ -761,7 +789,7 @@ def glu(x): output = self._apply_bias(output, output_bias, tokens_per_expert, permuted_probs) # upad and concat the output - if self.config.fp8 or self.config.fp4: + if (self.config.fp8 or self.config.fp4) and not _can_unpad_fc2: output = self.quantization_unpadding(output, actual_tokens_per_expert) output_bias = None