diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 41dd0cda59e..0b66c226b4c 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -131,8 +131,12 @@ from transformer_engine.pytorch.attention.rope import apply_fused_qkv_rotary_pos_emb HAVE_FUSED_QKV_ROPE = True + HAVE_FUSED_QKV_ROPE_FP8_ATTENTION = ( + "defer_for_fp8_attention" in inspect.signature(apply_fused_qkv_rotary_pos_emb).parameters + ) except ImportError: HAVE_FUSED_QKV_ROPE = False + HAVE_FUSED_QKV_ROPE_FP8_ATTENTION = False class LinearQkvInterface(Protocol): @@ -1605,9 +1609,30 @@ def forward_pre_attn_and_core_attn( max_seqlen=rope_freqs_max_seqlen, ) else: - query, key, value = apply_fused_qkv_rotary_pos_emb( - mixed_qkv, q_pos_emb, k_pos_emb, qkv_split_arg_list + # The FP8 attention fusion hands TE placeholders that its core attention fills + # (or fuses with the FP8 quantization), so they must not be offloaded or + # recomputed from. + defer_for_fp8_attention = ( + self.config.fused_qkv_rope_fp8_attention + and not (self.checkpoint_core_attention and self.training) + and not (self.offload_core_attention and self.training) ) + if defer_for_fp8_attention: + assert HAVE_FUSED_QKV_ROPE_FP8_ATTENTION, ( + "fused_qkv_rope_fp8_attention requires a TransformerEngine version whose " + "apply_fused_qkv_rotary_pos_emb supports defer_for_fp8_attention." + ) + query, key, value = apply_fused_qkv_rotary_pos_emb( + mixed_qkv, + q_pos_emb, + k_pos_emb, + qkv_split_arg_list, + defer_for_fp8_attention=True, + ) + else: + query, key, value = apply_fused_qkv_rotary_pos_emb( + mixed_qkv, q_pos_emb, k_pos_emb, qkv_split_arg_list + ) # TODO, can apply positional embedding to value_layer so it has # absolute positional embedding. diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 8974e592a35..4ccaeb32645 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -591,6 +591,12 @@ class TransformerConfig(ModelParallelConfig): fused_single_qkv_rope: bool = False """If set, avoid splitting QKV before ROPE forward and avoid concatenating ROPE dgrads.""" + fused_qkv_rope_fp8_attention: bool = False + """If set, let TransformerEngine fuse the RoPE of the single QKV tensor with the FP8 + current-scaling quantization of Q/K/V inside core attention (falls back to the regular fused + RoPE whenever attention does not run FP8 current-scaling fused attention). Requires + fused_single_qkv_rope and is disabled for checkpointed or offloaded core attention.""" + fused_residual_rmsnorm: bool = False """If True, fuses residual connection and RMSNorm backward pass when TE is used.""" @@ -3104,6 +3110,9 @@ def __post_init__(self): if self.attention_output_gate: raise ValueError("fused_single_qkv_rope does not support gated attention for now.") + if self.fused_qkv_rope_fp8_attention and not self.fused_single_qkv_rope: + raise ValueError("fused_qkv_rope_fp8_attention requires fused_single_qkv_rope.") + if self.multi_latent_attention and self.rotary_interleaved: raise ValueError("rotary_interleaved does not work with multi_latent_attention.")