diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 895d46e9b3d..fc79cc9a7db 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -871,6 +871,29 @@ def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool return fp8_context + def get_fp8_disabled_context(config: TransformerConfig, is_init: bool = False): + """Return a context manager that disables TE quantization. + + Use this around submodule construction or execution that must stay in a higher + precision while its enclosing module uses an FP8 or FP4 context. + + Args: + config: Transformer configuration that controls quantization. + is_init: Whether to disable the parameter-initialization context instead of + the forward autocast context. + + Returns: + A disabled TE quantization context when quantization is active, otherwise a + no-op context. + """ + if is_init: + if not (config.fp8_param or config.fp4_param): + return nullcontext() + return transformer_engine.pytorch.fp8_model_init(enabled=False) + if not (config.fp8 or config.fp4): + return nullcontext() + return transformer_engine.pytorch.fp8_autocast(enabled=False) + else: def get_fp8_recipe(config: TransformerConfig): @@ -881,6 +904,10 @@ def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool """Returns dummy fp8 context manager since TE is not available.""" return nullcontext() + def get_fp8_disabled_context(config: TransformerConfig, is_init: bool = False): + """Return a no-op context manager since TE is not available.""" + return nullcontext() + if HAVE_TE: from transformer_engine.pytorch.fp8 import FP8GlobalStateManager diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 1fd5dcfae37..cf1c4a31fb0 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -29,17 +29,28 @@ @triton.jit def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): - token_idx = -1 - this_seq_len = 0 + # Cast ``pid_m`` and ``cu_seqlens`` loads to a single shared dtype so + # the loop-body reassignments don't surface as + # "initial value is int32 but redefined as int64" in newer Triton + # versions (which promote ``// Python_int`` to int64). + pid_m = pid_m.to(tl.int64) + token_idx = tl.full((), -1, dtype=tl.int64) + this_seq_len = tl.full((), 0, dtype=tl.int64) seq_idx = 0 - last_cum_seqlen = tl.load(cu_seqlens) // cp_size + last_cum_seqlen = tl.load(cu_seqlens).to(tl.int64) // cp_size while seq_idx < seq_num: - cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1) // cp_size + cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1).to(tl.int64) // cp_size if token_idx == -1 and cur_cum_seqlen > pid_m: token_idx = pid_m - last_cum_seqlen this_seq_len = cur_cum_seqlen - last_cum_seqlen last_cum_seqlen = cur_cum_seqlen seq_idx += 1 + # Padding tokens beyond cu_seqlens[-1] (from THD CUDA-graph padding) + # never match any sequence, leaving token_idx == -1. Clamp to 0 so + # the cos/sin table loads stay in-bounds; the wrong RoPE result is + # harmless because padding positions are excluded by loss_mask. + if token_idx == -1: + token_idx = tl.full((), 0, dtype=tl.int64) if cp_size > 1: if token_idx < this_seq_len // 2: token_idx = token_idx + cp_rank * this_seq_len // 2 @@ -65,29 +76,34 @@ def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): restore_value=["Q"], ) @triton.jit -def rotary_fwd_q_kernel( +def _mla_rope_fwd_inplace_kernel( Q, COS, SIN, - qk_head_dim, + nope_dim, emb_dim: tl.constexpr, head_num: tl.constexpr, batch_size, seq_num, cu_seqlens_q, + position_ids, stride_x_seq, stride_x_nheads, + stride_cos_seq, + stride_sin_seq, cp_rank, cp_size, + INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the forward pass for applying YARN RoPE to MLA's query. - This kernel inplace modifies the input tensor Q. + Forward pass: apply RoPE inplace to the trailing emb_dim elements. + Reads from interleaved layout, writes back to interleaved layout. Input: - Q: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + Q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] COS/SIN: [max_seq_len, emb_dim] batch_size: batch size for sbhd format, not used for thd format @@ -97,15 +113,24 @@ def rotary_fwd_q_kernel( pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: + if position_ids is not None: + token_idx = tl.load(position_ids + pid_m) + elif cu_seqlens_q is None: token_idx = pid_m // batch_size else: token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) @@ -113,7 +138,7 @@ def rotary_fwd_q_kernel( Q = Q + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads - x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads # x1 = t[..., 0::2], x2 = t[..., 1::2] x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 @@ -124,10 +149,14 @@ def rotary_fwd_q_kernel( x_left = x_1 * cos_left - x_2 * sin_left x_right = x_2 * cos_right + x_1 * sin_right - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - tl.store(Q + x_left_off, x_left, mask=mask) - tl.store(Q + x_right_off, x_right, mask=mask) + if REMOVE_INTERLEAVING: + tl.store(Q + x_1_off, x_left, mask=mask) + tl.store(Q + x_2_off, x_right, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + tl.store(Q + x_left_off, x_left, mask=mask) + tl.store(Q + x_right_off, x_right, mask=mask) @triton.autotune( @@ -145,29 +174,34 @@ def rotary_fwd_q_kernel( restore_value=["DO"], ) @triton.jit -def rotary_bwd_q_kernel( +def _mla_rope_bwd_inplace_kernel( DO, COS, SIN, - qk_head_dim, + nope_dim, emb_dim: tl.constexpr, head_num: tl.constexpr, batch_size, seq_num, cu_seqlens_q, + position_ids, stride_x_seq, stride_x_nheads, + stride_cos_seq, + stride_sin_seq, cp_rank, cp_size, + INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the backward pass for applying YARN RoPE to MLA's query. - This kernel inplace modifies the input tensor DO. + Backward pass: inverse RoPE inplace on the trailing emb_dim elements. + Reads from interleaved layout, writes to interleaved layout. Input: - DO: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + DO: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] COS/SIN: [max_seq_len, emb_dim] batch_size, seq_num, and cu_seqlens_q are the same as in the forward pass @@ -175,15 +209,24 @@ def rotary_bwd_q_kernel( pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: + if position_ids is not None: + token_idx = tl.load(position_ids + pid_m) + elif cu_seqlens_q is None: token_idx = pid_m // batch_size else: token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) @@ -191,25 +234,32 @@ def rotary_bwd_q_kernel( DO = DO + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads - x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(DO + x_left_off, mask=mask) - x_right = tl.load(DO + x_right_off, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(DO + x_1_off, mask=mask) + x_right = tl.load(DO + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 x_1 = x_left * cos_left + x_right * sin_right x_2 = -x_left * sin_left + x_right * cos_right - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 tl.store(DO + x_1_off, x_1, mask=mask) tl.store(DO + x_2_off, x_2, mask=mask) -class ApplyMLARotaryEmbQ(torch.autograd.Function): +class _FusedMLARoPEInplace(torch.autograd.Function): """ - Autograd function for applying YARN RoPE to MLA's query. + Autograd function for applying RoPE inplace to the trailing emb_dim + elements of a multi-head tensor (leaving the first nope_dim elements unchanged). """ @staticmethod @@ -218,22 +268,26 @@ def forward( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved=False, + inverse=False, + remove_interleaving=False, + position_ids=None, ): """ - Forward function for ApplyMLARotaryEmbQ. + Forward function for _FusedMLARoPEInplace. Args: - q: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + inverse: if True, negate sin inside the kernel to apply the inverse rotation """ assert not rotary_interleaved max_seqlen = None @@ -241,6 +295,7 @@ def forward( seq_num = None if cu_seqlens_q is None: # sbhd + assert position_ids is None max_seqlen, batch_size, nheads, headdim = q.shape q = q.view(-1, nheads, headdim) total_seqlen = q.shape[0] @@ -248,33 +303,43 @@ def forward( # thd total_seqlen, nheads, headdim = q.shape seq_num = len(cu_seqlens_q) - 1 + if position_ids is not None: + assert position_ids.shape == (total_seqlen,) assert q.stride(-1) == 1 - assert cos.is_contiguous() - assert sin.is_contiguous() - assert headdim == qk_head_dim + emb_dim + assert cos.stride(-1) == 1 + assert sin.stride(-1) == 1 + assert headdim == nope_dim + emb_dim assert emb_dim % 4 == 0 grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_q_kernel[grid]( + _mla_rope_fwd_inplace_kernel[grid]( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, nheads, batch_size, seq_num, cu_seqlens_q, + position_ids, q.stride(0), q.stride(1), + cos.stride(0), + sin.stride(0), cp_rank, cp_size, + INVERSE=inverse, + REMOVE_INTERLEAVING=remove_interleaving, ) - ctx.save_for_backward(cos, sin) - ctx.qk_head_dim = qk_head_dim + ctx.save_for_backward(cos, sin, *(() if position_ids is None else (position_ids,))) + ctx.has_position_ids = position_ids is not None + ctx.nope_dim = nope_dim ctx.emb_dim = emb_dim ctx.cu_seqlens_q = cu_seqlens_q ctx.rotary_interleaved = rotary_interleaved + ctx.inverse = inverse + ctx.remove_interleaving = remove_interleaving ctx.cp_rank = cp_rank ctx.cp_size = cp_size if cu_seqlens_q is None: @@ -284,13 +349,17 @@ def forward( @staticmethod def backward(ctx, grad): """ - Backward function for ApplyMLARotaryEmbQ. + Backward function for _FusedMLARoPEInplace. Args: - grad: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + grad: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] """ - cos, sin = ctx.saved_tensors + if ctx.has_position_ids: + cos, sin, position_ids = ctx.saved_tensors + else: + cos, sin = ctx.saved_tensors + position_ids = None max_seqlen = None batch_size = None seq_num = None @@ -300,65 +369,126 @@ def backward(ctx, grad): total_seqlen = grad.shape[0] else: seq_num = len(ctx.cu_seqlens_q) - 1 + if ctx.has_position_ids: + grad = grad.contiguous() total_seqlen, nheads, headdim = grad.shape assert grad.stride(-1) == 1 grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_q_kernel[grid]( + _mla_rope_bwd_inplace_kernel[grid]( grad, cos, sin, - ctx.qk_head_dim, + ctx.nope_dim, ctx.emb_dim, nheads, batch_size, seq_num, ctx.cu_seqlens_q, + position_ids, grad.stride(0), grad.stride(1), + cos.stride(0), + sin.stride(0), ctx.cp_rank, ctx.cp_size, + INVERSE=ctx.inverse, + REMOVE_INTERLEAVING=ctx.remove_interleaving, ) if ctx.cu_seqlens_q is None: grad = grad.view(max_seqlen, batch_size, nheads, headdim) - return grad, None, None, None, None, None, None, None, None + return grad, None, None, None, None, None, None, None, None, None, None, None -def fused_apply_mla_rope_for_q( +def fused_mla_rope_inplace( t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, - qk_head_dim: int, + nope_dim: int, emb_dim: int, cu_seqlens_q: Optional[torch.Tensor] = None, cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, -): + inverse: bool = False, + remove_interleaving: bool = False, + position_ids: Optional[torch.Tensor] = None, +) -> torch.Tensor: """ - Fused function for applying YARN RoPE to MLA's query. - This function inplace modifies the input tensor t. - Along the last dimension of t, the last emb_dim elements are applied with RoPE. - The first qk_head_dim elements are not modified. - It is an experimental feature and may change in future versions. + Fused RoPE applied inplace to the trailing emb_dim elements of a tensor, + leaving the first nope_dim elements unchanged. It supports both sbhd and thd input formats. + When ``inverse=True`` the rotation is reversed, which is useful for + undoing RoPE on the attention output. + For the notations below, seq_len is the length of the sequence per batch for sbhd format, total_seq_len is the total length of the sequences for thd format. max_seq_len is the maximum length of the sequences in the input tensor. Args: - t: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + t: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + inverse: if True, apply the inverse rotation + remove_interleaving: if True, output RoPE dims in non-interleaved layout + position_ids: optional THD row positions. When supplied, these positions + replace the built-in CP row-to-position mapping. Returns: t: inplace modified input tensor """ - return ApplyMLARotaryEmbQ.apply( - t, cos, sin, qk_head_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved + return _FusedMLARoPEInplace.apply( + t, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q, + cp_rank, + cp_size, + rotary_interleaved, + inverse, + remove_interleaving, + position_ids, + ) + + +def fused_mla_rope_out_of_place( + t: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + nope_dim: int, + emb_dim: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, + inverse: bool = False, + remove_interleaving: bool = False, + position_ids: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Apply the fused RoPE kernel without modifying the input tensor. + + Use this wrapper when an upstream autograd function may have retained its + output for backward. The underlying kernel remains in-place, so a private + copy is required to keep the retained tensor unchanged. + """ + return fused_mla_rope_inplace( + t.clone(), + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens_q, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, + inverse=inverse, + remove_interleaving=remove_interleaving, + position_ids=position_ids, ) @@ -376,7 +506,7 @@ def fused_apply_mla_rope_for_q( key=["emb_dim", "k_dim", "v_dim", "head_num"], ) @triton.jit -def rotary_fwd_kv_kernel( +def _mla_rope_fwd_kv_split_kernel( KV, K_POS_EMB, O_KEY, @@ -399,12 +529,12 @@ def rotary_fwd_kv_kernel( stride_v_nheads, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the forward pass for applying YARN RoPE to MLA's key and value. - It splits the input tensor KV into key and value, - and concatenates the processed RoPE to the key. + Forward pass: split KV into key and value, apply RoPE to k_pos_emb, + and concatenate the result onto key. Input: KV: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -460,14 +590,24 @@ def rotary_fwd_kv_kernel( x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_left_off = ( - tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads - + k_dim - + tl.arange(0, emb_dim // 2)[None, :] - ) - x_right_off = x_left_off + emb_dim // 2 - tl.store(K_ptr + x_left_off, x_left, mask=mask) - tl.store(K_ptr + x_right_off, x_right, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] * 2 + ) + x_2_off = x_1_off + 1 + tl.store(K_ptr + x_1_off, x_left, mask=mask) + tl.store(K_ptr + x_2_off, x_right, mask=mask) + else: + x_left_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 + tl.store(K_ptr + x_left_off, x_left, mask=mask) + tl.store(K_ptr + x_right_off, x_right, mask=mask) @triton.autotune( @@ -484,7 +624,7 @@ def rotary_fwd_kv_kernel( key=["emb_dim", "k_dim", "v_dim", "head_num"], ) @triton.jit -def rotary_bwd_kv_kernel( +def _mla_rope_bwd_kv_split_kernel( dK, dV, dKV, @@ -507,10 +647,11 @@ def rotary_bwd_kv_kernel( stride_demb_seq, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the backward pass for applying YARN RoPE to MLA's key and value. + Backward pass for the KV-split RoPE. Input: dK: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -555,10 +696,16 @@ def rotary_bwd_kv_kernel( dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim mask = x_off < head_num * stride_dk_nheads - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(dK_ptr + x_left_off, mask=mask) - x_right = tl.load(dK_ptr + x_right_off, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(dK_ptr + x_1_off, mask=mask) + x_right = tl.load(dK_ptr + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(dK_ptr + x_left_off, mask=mask) + x_right = tl.load(dK_ptr + x_right_off, mask=mask) x_left_accum += x_left x_right_accum += x_right x_left_accum = tl.sum(x_left_accum, axis=0) @@ -578,9 +725,10 @@ def rotary_bwd_kv_kernel( tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) -class ApplyMLARotaryEmbKV(torch.autograd.Function): +class _FusedMLARoPEKVSplit(torch.autograd.Function): """ - Autograd function for applying YARN RoPE to MLA's key and value. + Autograd function for applying RoPE to MLA's key and value. + Splits KV, applies RoPE to k_pos_emb, concatenates onto key. """ @staticmethod @@ -597,9 +745,10 @@ def forward( cp_rank, cp_size, rotary_interleaved=False, + remove_interleaving=False, ): """ - Forward function for ApplyMLARotaryEmbKV. + Forward function for _FusedMLARoPEKVSplit. Args: kv: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -634,7 +783,7 @@ def forward( o_value = kv.new_empty(total_seqlen, nheads, v_dim) grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_kv_kernel[grid]( + _mla_rope_fwd_kv_split_kernel[grid]( kv, k_pos_emb, o_key, @@ -657,8 +806,10 @@ def forward( o_value.stride(1), cp_rank, cp_size, + REMOVE_INTERLEAVING=remove_interleaving, ) ctx.save_for_backward(cos, sin) + ctx.remove_interleaving = remove_interleaving ctx.rotary_interleaved = rotary_interleaved ctx.emb_dim = emb_dim ctx.k_dim = k_dim @@ -674,7 +825,7 @@ def forward( @staticmethod def backward(ctx, dk, dv): """ - Backward function for ApplyMLARotaryEmbKV. + Backward function for _FusedMLARoPEKVSplit. Args: dk: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -702,7 +853,7 @@ def backward(ctx, dk, dv): d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim) grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_kv_kernel[grid]( + _mla_rope_bwd_kv_split_kernel[grid]( dk, dv, d_kv, @@ -725,14 +876,15 @@ def backward(ctx, dk, dv): d_emb.stride(0), ctx.cp_rank, ctx.cp_size, + REMOVE_INTERLEAVING=ctx.remove_interleaving, ) if ctx.cu_seqlens_kv is None: d_kv = d_kv.view(max_seqlen, batch_size, nheads, ctx.k_dim + ctx.v_dim) d_emb = d_emb.view(max_seqlen, batch_size, 1, ctx.emb_dim) - return d_kv, d_emb, None, None, None, None, None, None, None, None, None + return d_kv, d_emb, None, None, None, None, None, None, None, None, None, None -def fused_apply_mla_rope_for_kv( +def fused_mla_rope_kv_split( kv: torch.Tensor, k_pos_emb: torch.Tensor, cos: torch.Tensor, @@ -744,9 +896,10 @@ def fused_apply_mla_rope_for_kv( cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, -): + remove_interleaving: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: """ - Fused function for applying YARN RoPE to MLA's key and value. + Fused function for applying RoPE to MLA's key and value. It splits the input tensor kv into key and value, and concatenates the processed RoPE to the key. @@ -761,13 +914,14 @@ def fused_apply_mla_rope_for_kv( cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_kv: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + remove_interleaving: if True, output RoPE dims in non-interleaved layout Returns: key: [seq_len, batch_size, head_num, emb_dim + k_dim] or [total_seq_len, head_num, emb_dim + k_dim] value: [seq_len, batch_size, head_num, v_dim] or [total_seq_len, head_num, v_dim] """ - return ApplyMLARotaryEmbKV.apply( + return _FusedMLARoPEKVSplit.apply( kv, k_pos_emb, cos, @@ -779,4 +933,64 @@ def fused_apply_mla_rope_for_kv( cp_rank, cp_size, rotary_interleaved, + remove_interleaving, + ) + + +def fused_apply_mla_rope_for_q( + t: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + qk_head_dim: int, + emb_dim: int, + cu_seqlens_q: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, +) -> torch.Tensor: + """Backward-compatible in-place MLA query RoPE API. + + New callers should choose :func:`fused_mla_rope_inplace` or + :func:`fused_mla_rope_out_of_place` explicitly. This legacy name keeps + its original mutation behavior and does not add a clone to the hot path. + """ + return fused_mla_rope_inplace( + t, + cos, + sin, + qk_head_dim, + emb_dim, + cu_seqlens_q=cu_seqlens_q, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, + ) + + +def fused_apply_mla_rope_for_kv( + kv: torch.Tensor, + k_pos_emb: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + emb_dim: int, + k_dim: int, + v_dim: int, + cu_seqlens_kv: Optional[torch.Tensor] = None, + cp_rank: int = 0, + cp_size: int = 1, + rotary_interleaved: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Backward-compatible name for the MLA key/value split RoPE API.""" + return fused_mla_rope_kv_split( + kv, + k_pos_emb, + cos, + sin, + emb_dim, + k_dim, + v_dim, + cu_seqlens_kv=cu_seqlens_kv, + cp_rank=cp_rank, + cp_size=cp_size, + rotary_interleaved=rotary_interleaved, ) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index 9fab25a3fae..0468ddd14ae 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -95,6 +95,8 @@ def _apply_rotary_pos_emb_bshd( rotary_interleaved: bool = False, mla_rotary_interleaved: bool = False, mscale: float = 1.0, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, multi_latent_attention: Optional[bool] = None, ) -> Tensor: """Apply rotary positional embedding to input tensor T. @@ -118,6 +120,13 @@ def _apply_rotary_pos_emb_bshd( ) mla_rotary_interleaved = multi_latent_attention + # Some callers may pass freqs with an extra singleton axis, e.g. + # t: [s, b, d] and freqs: [s, 1, 1, d]. In that case, broadcasting would + # accidentally expand to [s, s, b, d]. Squeeze the extra singleton axis to + # keep freqs rank aligned with t. + if freqs.dim() == t.dim() + 1 and freqs.size(-2) == 1: + freqs = freqs.squeeze(-2) + rot_dim = freqs.shape[-1] # ideally t_pass is empty so rotary pos embedding is applied to all tensor t @@ -132,8 +141,18 @@ def _apply_rotary_pos_emb_bshd( # second part is sine component, need to change signs with _rotate_half method cos_ = (torch.cos(freqs) * mscale).to(t.dtype) sin_ = (torch.sin(freqs) * mscale).to(t.dtype) + if inverse: + sin_ = -sin_ t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_) + + # Fallback to original permutation + # DSv4 applies rope on V and O, so we need to uninterleave the tensor. + # The existing MLA code is safe because the dot product is permutation-invariant. + if mla_rotary_interleaved and mla_output_remove_interleaving: + x1, x2 = torch.chunk(t, 2, dim=-1) + t = torch.stack((x1, x2), dim=-1).flatten(start_dim=-2) + return torch.cat((t, t_pass), dim=-1) @@ -193,20 +212,28 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved: bool = False, mla_rotary_interleaved: bool = False, mscale: float = 1.0, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, cp_group: torch.distributed.ProcessGroup = None, multi_latent_attention: Optional[bool] = None, + max_seqlen: Optional[int] = None, ) -> Tensor: - """A baseline implementation of applying RoPE for `thd` format. + """Apply RoPE for `thd` format using vectorized CUDA operations. + + When ``max_seqlen`` is supplied, this path performs no GPU-to-CPU sync and is + compatible with CUDA Graph capture. The compatibility path for legacy callers + that omit ``max_seqlen`` retains one GPU-to-CPU sync. Args: - t (Tensor): Input tensor T is of shape [t, h, d] - cu_seqlens(Tensor): Cumulative sum of sequence lengths in a batch for `t`, - with shape [b + 1] and dtype torch.int32. - freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d] - cp_group (torch.distributed.ProcessGroup): The context parallel group + t (Tensor): Input tensor of shape [total_tokens, h, d] + cu_seqlens (Tensor): Cumulative sequence lengths, shape [num_seqs + 1], int32. + freqs (Tensor): RoPE frequencies, shape [max_s, 1, 1, d] or [total_tokens, 1, 1, d] + cp_group: Context parallel group + max_seqlen: Global max sequence length for this packed batch when known. Supplying it + avoids the compatibility-path host sync used by legacy callers. Returns: - Tensor: Shape [t, h, d]. The input tensor after applying RoPE. + Tensor: Shape [total_tokens, h, d]. Input with RoPE applied. """ if multi_latent_attention is not None: warnings.warn( @@ -219,53 +246,71 @@ def _apply_rotary_pos_emb_thd( raise ValueError("cp_group must be provided for THD format RoPE") cp_size = cp_group.size() cp_rank = cp_group.rank() - seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist() - sequence_splits = torch.split(t, seqlens) - total_seqlen = int(cu_seqlens[-1].item()) - has_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == total_seqlen - - # Handle two different frequency tensor formats: - # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains positions for the whole packed - # batch. Each sequence must therefore use its cu_seqlens offset when selecting the local CP - # front/back slices. For example, with cu_seqlens=[0, 4, 8], cp_size=2, rank 0 should use - # positions [0, 3, 4, 7], not [0, 3, 0, 3]. - # 2. Otherwise: freqs contains only max sequence length positions. Each packed sequence should - # reuse positions starting from 0, preserving the legacy THD behavior. - if has_packed_freqs: - # CASE 1: Exact mapping with offsets - local_freqs = [] - for i, x in enumerate(sequence_splits): - # cu_seqlens[i] is the starting offset of this sequence in the original batch - seq_start_offset = cu_seqlens[i].item() - local_freqs.append( - _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) - ) - freqs = torch.cat(local_freqs, dim=0) - return _apply_rotary_pos_emb_bshd( - t.unsqueeze(1), - freqs, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - ).squeeze(1) - - # CASE 2: Traditional mapping without offsets. Apply RoPE one sequence at a time so the second - # and later packed sequences do not look like continuations of the first sequence. - output = torch.empty_like(t) - output_offset = 0 - for x in sequence_splits: - freq_slice = _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs) - output_slice = _apply_rotary_pos_emb_bshd( - x.unsqueeze(1), - freq_slice, - rotary_interleaved=rotary_interleaved, - mla_rotary_interleaved=mla_rotary_interleaved, - mscale=mscale, - ).squeeze(1) - output.narrow(0, output_offset, x.size(0)).copy_(output_slice) - output_offset += x.size(0) + total_tokens = t.shape[0] + device = t.device + + token_pos = torch.arange(total_tokens, device=device, dtype=torch.int64) + + # `cu_seqlens` describes the global packed sequence. With CP, `t` is already + # CP-partitioned, so build a local cumulative-length view before assigning + # local tokens to packed sequences. + cu_seqlens_i64 = cu_seqlens.to(torch.int64) + global_seq_lens = cu_seqlens_i64[1:] - cu_seqlens_i64[:-1] + local_seq_lens = global_seq_lens // cp_size if cp_size > 1 else global_seq_lens + local_cu_seqlens = torch.zeros_like(cu_seqlens_i64) + local_cu_seqlens[1:] = torch.cumsum(local_seq_lens, dim=0) + + # `searchsorted(..., right=True) - 1` returns the local sequence index. The + # clamp guards padded tokens that sit beyond the final real local token; they + # get a harmless frequency and are later masked out. + seq_idx = torch.searchsorted(local_cu_seqlens, token_pos, right=True) - 1 + seq_idx = seq_idx.clamp(min=0, max=cu_seqlens.shape[0] - 2) + + local_seq_start = local_cu_seqlens[seq_idx] + local_pos = token_pos - local_seq_start + local_seq_len = local_seq_lens[seq_idx] + global_seq_start = cu_seqlens_i64[seq_idx] - return output + if cp_size > 1: + cp_seg = local_seq_len // 2 + full_seqlen = local_seq_len * cp_size + is_first_half = local_pos < cp_seg + freq_pos = torch.where( + is_first_half, + cp_rank * cp_seg + local_pos, + full_seqlen - (cp_rank + 1) * cp_seg + (local_pos - cp_seg), + ) + else: + freq_pos = local_pos.to(torch.int64) + + if max_seqlen is None: + # Backward compatibility for callers that predate ``max_seqlen``. This retains + # the old packed-frequency semantics at the cost of a GPU-to-CPU sync. Updated + # training paths pass ``max_seqlen`` and stay CUDA-graph safe. + exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) == int(cu_seqlens[-1].item()) + else: + exact_packed_freqs = freqs.dim() >= 1 and freqs.size(0) > max_seqlen + if exact_packed_freqs: + # `freqs` covers all positions across all sequences (used for non-1D + # RoPE / VLMs); shift by the per-sequence start offset so each token + # samples its absolute position. When `freqs` only spans one max-len + # sequence, no shift is needed. + freq_pos = freq_pos + global_seq_start + + # Padded positions can sit outside the frequency table. Clamp them into + # range; downstream padding masks exclude those positions from the result. + freq_pos = freq_pos.clamp(min=0, max=freqs.shape[0] - 1) + freqs_packed = freqs[freq_pos] + + return _apply_rotary_pos_emb_bshd( + t.unsqueeze(1), + freqs_packed, + rotary_interleaved=rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ).squeeze(1) def apply_rotary_pos_emb( @@ -276,6 +321,9 @@ def apply_rotary_pos_emb( mscale: float = 1.0, cp_group: torch.distributed.ProcessGroup = None, mla_rotary_interleaved: bool = False, + inverse: bool = False, + mla_output_remove_interleaving: bool = False, + max_seqlen: Optional[int] = None, ): """ Reroute to the appropriate apply_rotary_pos_emb function depending on @@ -312,6 +360,12 @@ def apply_rotary_pos_emb( "Using unfused implementation." ) use_unfused = True + if inverse: + warnings.warn( + "inverse RoPE is not supported by TE's fused RoPE. " + "Using unfused implementation." + ) + use_unfused = True if not use_unfused: assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available." return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved) @@ -333,6 +387,8 @@ def apply_rotary_pos_emb( rotary_interleaved=config.rotary_interleaved, mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ) else: return _apply_rotary_pos_emb_thd( @@ -343,6 +399,9 @@ def apply_rotary_pos_emb( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, cp_group=cp_group, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + max_seqlen=max_seqlen, ) diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py index 0e560f939f2..e591e4ff90d 100644 --- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py @@ -205,6 +205,47 @@ def forward( return emb + def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False, cp_group=None): + """Materialize cached cos/sin tensors for ``[seq_len, ..., dim]``.""" + self.max_seq_len_cached = seq_len + self.offset_cached = offset + self.dtype_cached = dtype + self.packed_seq_cached = packed_seq + + emb = self.forward(seq_len, offset, packed_seq=packed_seq, cp_group=cp_group) + self.register_buffer("cos_cached", emb.cos().to(dtype).contiguous(), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype).contiguous(), persistent=False) + + def get_cached_cos_sin( + self, + seq_len, + offset=0, + dtype=torch.get_default_dtype(), + packed_seq=False, + cp_group=None, + mscale=None, + ): + """Get cached cos and sin values. + + The cache is rebuilt on first use or whenever ``seq_len`` grows + beyond the cached length, or any of ``offset`` / ``dtype`` / + ``packed_seq`` changes from the previous call. + ``YarnRotaryEmbedding`` overrides this to also bake its + concentration factor into the cached cos/sin (controlled by + ``mscale``); for the base class without a concentration + factor the argument is accepted-and-ignored for API uniformity. + """ + del mscale # base class has no concentration factor + if ( + not hasattr(self, "max_seq_len_cached") + or seq_len > self.max_seq_len_cached + or offset != self.offset_cached + or dtype != self.dtype_cached + or packed_seq != self.packed_seq_cached + ): + self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq, cp_group) + return (self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]) + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): state_dict.pop(f'{prefix}inv_freq', None) return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) diff --git a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py index 166ef9b41e7..cb8a03d0b2b 100644 --- a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py @@ -186,13 +186,18 @@ def forward( emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) return emb, _mscale - def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False): + def _set_cos_sin_cache( + self, seq_len, offset, dtype, packed_seq=False, cp_group=None, mscale=None + ): self.max_seq_len_cached = seq_len self.offset_cached = offset self.dtype_cached = dtype self.packed_seq_cached = packed_seq + self.mscale_cached = mscale - emb, _mscale = self.forward(seq_len, offset, packed_seq) + emb, _mscale = self.forward(seq_len, offset, packed_seq=packed_seq, cp_group=cp_group) + if mscale is not None: + _mscale = mscale self.register_buffer( "cos_cached", (emb.cos() * _mscale).to(dtype).contiguous(), persistent=False ) @@ -201,16 +206,34 @@ def _set_cos_sin_cache(self, seq_len, offset, dtype, packed_seq=False): ) def get_cached_cos_sin( - self, seq_len, offset=0, dtype=torch.get_default_dtype(), packed_seq=False + self, + seq_len, + offset=0, + dtype=torch.get_default_dtype(), + packed_seq=False, + cp_group=None, + mscale=None, ): - """Get cached cos and sin values.""" + """Get cached cos and sin values. + + Args: + mscale: when ``None`` (default), the cached cos/sin are + multiplied by yarn's internal concentration factor (the + normal long-context behaviour). When a float is supplied, + that value is used in place of the internal factor — e.g. + the DSv4 hybrid model passes ``mscale=1.0`` to enforce + its "pure rotation" contract and keep the fused / + unfused rope paths bit-equivalent. + """ if ( - seq_len > self.max_seq_len_cached + not hasattr(self, "max_seq_len_cached") + or seq_len > self.max_seq_len_cached or offset != self.offset_cached or dtype != self.dtype_cached or packed_seq != self.packed_seq_cached + or mscale != getattr(self, "mscale_cached", None) ): - self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq) + self._set_cos_sin_cache(seq_len, offset, dtype, packed_seq, cp_group, mscale) return (self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]) diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index a76fe6e3a23..5189d264e59 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -10,6 +10,18 @@ AbsorbedMLASelfAttention, AbsorbedMLASelfAttentionSubmodules, ) +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, +) +from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + DSv4HybridSelfAttentionSubmodules, +) from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexer, DSAIndexerSubmodules, @@ -130,6 +142,57 @@ def get_dsa_module_spec_for_backend( return attention +def get_dsv4_hybrid_module_spec_for_backend( + config: TransformerConfig, backend: BackendSpecProvider = None +) -> ModuleSpec: + """Build the native SBHD DSv4 hybrid-attention module spec.""" + assert config.multi_latent_attention, "Currently only MLA supports sparse attention." + assert config.qk_l2_norm is False, "qk_l2_norm is not supported with MLA." + + rms_norm = config.normalization == "RMSNorm" + qk_norm = ( + backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp + ) + + compressor_spec = ModuleSpec( + module=Compressor, + submodules=CompressorSubmodules( + linear_wkv=backend.linear(), + linear_wgate=backend.linear(), + norm=backend.layer_norm(rms_norm=True, for_qk=False), + ), + ) + indexer_spec = ModuleSpec( + module=CSAIndexer, + submodules=CSAIndexerSubmodules( + linear_wq_b=backend.linear(), + linear_weights_proj=backend.linear(), + compressor=compressor_spec, + ), + ) + core_attention = ModuleSpec( + module=CompressedSparseAttention, + submodules=CompressedSparseAttentionSubmodules( + compressor=compressor_spec, indexer=indexer_spec + ), + ) + + return ModuleSpec( + module=DSv4HybridSelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=DSv4HybridSelfAttentionSubmodules( + linear_q_down_proj=backend.linear(), + linear_q_up_proj=backend.column_parallel_linear(), + linear_kv_proj=backend.column_parallel_linear(), + core_attention=core_attention, + linear_proj=backend.row_parallel_linear(), + q_layernorm=qk_norm, + kv_layernorm=qk_norm, + ), + metainfo={"fuse_input_layernorm": False}, + ) + + def get_experimental_attention_variant_module_spec( config: TransformerConfig, backend: BackendSpecProvider = None ) -> ModuleSpec: @@ -142,6 +205,8 @@ def get_experimental_attention_variant_module_spec( return get_gated_delta_net_module_spec(config=config, backend=backend) elif config.experimental_attention_variant == "dsa": return get_dsa_module_spec_for_backend(config=config, backend=backend) + elif config.experimental_attention_variant == "dsv4_hybrid": + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=backend) else: raise ValueError( f"Invalid experimental attention variant: {config.experimental_attention_variant}" diff --git a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py index c87ccd5ff31..45bf910f84d 100644 --- a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +++ b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py @@ -370,8 +370,11 @@ def _update_fp32_params_by_new_state(self): if not self.param_update_in_fp32: return for param, v in self.state.items(): - fp32_param = self.param_to_fp32_param[param] - fp32_param.data.copy_(v["master_param"]) + # Native FP32 params do not need a separate master parameter and are + # intentionally absent from param_to_fp32_param. + fp32_param = self.param_to_fp32_param.get(param) + if fp32_param is not None: + fp32_param.data.copy_(v["master_param"]) def update_fp32_param_by_new_param(self): """ diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index dac16f4a2ee..fd8c1864902 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -7,7 +7,6 @@ import math import warnings from abc import ABC, abstractmethod -from itertools import chain from logging import getLogger from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -1147,8 +1146,25 @@ def sharded_state_dict( state_dict = self.state_dict() + # Optimizer state ids enumerate the inner optimizer params: the fp32 main + # copies of float16 params, native fp32 params, and any frozen params, + # interleaved in the original param-group order. Map each fp32 main copy + # back to its model-side param; all other params already are model params. + main_param_id_to_model_param = { + id(main_param): model_param + for model_group, main_group in zip( + self.float16_groups, self.fp32_from_float16_groups, strict=True + ) + for model_param, main_param in zip(model_group, main_group, strict=True) + } + + def model_params_in_optimizer_order(): + for inner_group in self.optimizer.param_groups: + for param in inner_group['params']: + yield main_param_id_to_model_param.get(id(param), param) + id_to_sharded_param_map = get_param_id_to_sharded_param_map( - model_sharded_state_dict, chain.from_iterable(g for g in self.float16_groups) + model_sharded_state_dict, model_params_in_optimizer_order() ) _backfill_gtp_sharded_param_map( @@ -1159,6 +1175,20 @@ def sharded_state_dict( assert len(state_dict['fp32_from_fp16_params']) == len( state_dict['optimizer']['param_groups'] ) + # State ids of the fp32 main copies only, skipping native fp32 and frozen params. + float16_param_ids_per_group = [] + for state_group, inner_group in zip( + state_dict['optimizer']['param_groups'], self.optimizer.param_groups, strict=True + ): + float16_param_ids_per_group.append( + [ + param_id + for param_id, param in zip( + state_group['params'], inner_group['params'], strict=True + ) + if id(param) in main_param_id_to_model_param + ] + ) state_dict['fp32_from_fp16_params'] = [ [ make_sharded_optimizer_tensor( @@ -1166,10 +1196,10 @@ def sharded_state_dict( fp32_param, prefix=f'optimizer.state.fp32_param', ) - for param_id, fp32_param in zip(state_group['params'], fp32_group) + for param_id, fp32_param in zip(param_ids, fp32_group, strict=True) ] - for fp32_group, state_group in zip( - state_dict['fp32_from_fp16_params'], state_dict['optimizer']['param_groups'] + for fp32_group, param_ids in zip( + state_dict['fp32_from_fp16_params'], float16_param_ids_per_group, strict=True ) ] diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index 1daeacc9027..42b350ebd00 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -277,7 +277,7 @@ def _get_experimental_attention_variant_loss_scale_func(config): if loss_scale_func is not None: return loss_scale_func - if getattr(config, 'experimental_attention_variant', None) == 'dsa': + if getattr(config, 'experimental_attention_variant', None) in ('dsa', 'dsv4_hybrid'): from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, ) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 682b75fb701..1f94b853abe 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -303,6 +303,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -314,6 +315,7 @@ def __init__( self.config = config self.layer_number = layer_number self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.attn_mask_type = attn_mask_type self.attention_type = attention_type @@ -1490,8 +1492,16 @@ def forward( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv + rope_freqs_max_seqlen = ( + max(rope_max_seqlen_q, rope_max_seqlen_kv) + if rope_max_seqlen_q is not None and rope_max_seqlen_kv is not None + else None + ) else: cu_seqlens_q = cu_seqlens_kv = None + rope_freqs_max_seqlen = None if split_qkv: if q_pos_emb is not None: @@ -1504,6 +1514,7 @@ def forward( cu_seqlens=cu_seqlens_q, mscale=self._yarn_concentration_factor, cp_group=self.pg_collection.cp, + max_seqlen=rope_freqs_max_seqlen, ) else: query = inference_context.apply_rotary_emb_query( @@ -1522,6 +1533,7 @@ def forward( cu_seqlens=cu_seqlens_kv, mscale=self._yarn_concentration_factor, cp_group=self.pg_collection.cp, + max_seqlen=rope_freqs_max_seqlen, ) else: query, key, value = apply_fused_qkv_rotary_pos_emb( @@ -1658,6 +1670,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -1673,6 +1686,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) @@ -2075,6 +2089,7 @@ def __init__( attn_mask_type: AttnMaskType = AttnMaskType.padding, cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -2089,6 +2104,7 @@ def __init__( attention_type="cross", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, name=name, ) diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index e0b6af7aa7f..9991e6828d1 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -147,6 +147,7 @@ def __init__( pg_collection: ProcessGroupCollection = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -161,6 +162,7 @@ def __init__( pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, + is_mtp_layer=is_mtp_layer, ) # Resolve which classes to use for Q and KV linear up projections and norms, based on @@ -447,8 +449,16 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv + rope_freqs_max_seqlen = ( + max(rope_max_seqlen_q, rope_max_seqlen_kv) + if rope_max_seqlen_q is not None and rope_max_seqlen_kv is not None + else None + ) else: cu_seqlens_q = cu_seqlens_kv = None + rope_freqs_max_seqlen = None # ========================================= # Q down projection @@ -636,6 +646,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_freqs_max_seqlen, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -646,6 +657,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_freqs_max_seqlen, ) # query: [num_tokens, n, (kv_lora_rank + qk_pos_emb_head_dim)] diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py new file mode 100644 index 00000000000..9f32588a864 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -0,0 +1,909 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import copy +from dataclasses import dataclass +from functools import lru_cache +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn + +from megatron.core.fp8_utils import get_fp8_disabled_context +from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace +from megatron.core.models.common.embeddings import RotaryEmbedding, apply_rotary_pos_emb +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexerLossAutoScaler, + DSAIndexerLossLoggingHelper, + FusedDSAIndexerLoss, + fused_qk_topk_naive, + rotate_activation, +) +from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import nvtx_range_pop, nvtx_range_push + +# --------------------------------------------------------------------------- +# Helper functions for index computation +# --------------------------------------------------------------------------- + + +@lru_cache(maxsize=8) +def _get_window_topk_idxs_cached(window_size: int, seqlen: int, device_str: str) -> torch.Tensor: + """Compute sliding-window indices for a single sequence (cached). + + Returns: + indices: [seqlen, window_size] int tensor, -1 for invalid positions. + """ + base = torch.arange(seqlen, device=device_str).unsqueeze(1) + offsets = torch.arange(window_size, device=device_str) + matrix = (base - window_size + 1).clamp(min=0) + offsets + matrix = torch.where(matrix > base, -1, matrix) + return matrix + + +def get_window_topk_idxs( + window_size: int, batch_size: int, seqlen: int, device: torch.device +) -> torch.Tensor: + """Sliding-window indices [batch, seqlen, window_size].""" + matrix = _get_window_topk_idxs_cached(window_size, seqlen, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +@lru_cache(maxsize=8) +def _get_compress_topk_idxs_cached( + ratio: int, seqlen: int, offset: int, device_str: str +) -> torch.Tensor: + """Compute all-compressed-positions indices for a single sequence (cached). + + Returns: + indices: [seqlen, seqlen // ratio] int tensor, -1 for future positions. + """ + n_compressed = seqlen // ratio + matrix = torch.arange(n_compressed, device=device_str).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1, device=device_str).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix + + +def get_compress_topk_idxs( + ratio: int, batch_size: int, seqlen: int, offset: int, device: torch.device +) -> torch.Tensor: + """All-compressed-position indices [batch, seqlen, seqlen // ratio].""" + matrix = _get_compress_topk_idxs_cached(ratio, seqlen, offset, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +# --------------------------------------------------------------------------- +# Helper functions for RoPE +# --------------------------------------------------------------------------- + + +def _apply_rope( + x: torch.Tensor, + nope_dim: int, + pos_dim: int, + rotary_pos_emb_module: RotaryEmbedding, + config: TransformerConfig, + rotary_seq_len: int, + ratio: int = 1, + cp_group: torch.distributed.ProcessGroup = None, +) -> torch.Tensor: + """Apply RoPE to the last ``qk_pos_emb_head_dim`` dims, leaving the rest unchanged. + + Accepts both 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs. When the input is 3-D a temporary head dimension is inserted for + ``apply_rotary_pos_emb`` and removed before returning. + """ + if ratio == 1: + total_seq_len = rotary_seq_len + else: + total_seq_len = rotary_seq_len * ratio + # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's + # concentration factor (mscale) is NOT part of the DSv4 model contract -- + # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0 + # regardless of which rotary class is in use. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if config.apply_rope_fusion: + # ``mscale=1.0`` keeps the cached cos/sin free of yarn's + # concentration factor so the fused kernel sees the same + # rotation as the unfused split-rotate path (DSv4 "pure + # rotation" contract). + rotary_pos_cos, rotary_pos_sin = rotary_pos_emb_module.get_cached_cos_sin( + total_seq_len, dtype=x.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + else: + # Compressed-attention callers instantiate ``YarnRotaryEmbedding`` + # whenever ``compress_ratio > 1`` (regardless of ``config.rope_type``); + # its ``forward`` returns ``(emb, mscale)``. Base ``RotaryEmbedding`` + # returns a single tensor. Unpack either form uniformly; the + # caller-side ``mscale=1.0`` keeps the yarn concentration factor + # out of the rotation. + result = rotary_pos_emb_module(total_seq_len, packed_seq=False) + if isinstance(result, tuple): + rotary_pos_emb = result[0] + else: + rotary_pos_emb = result + if rotary_pos_emb is not None and ratio > 1: + rotary_pos_emb = rotary_pos_emb[:total_seq_len:ratio][:rotary_seq_len] + if rotary_pos_cos is not None and ratio > 1: + rotary_pos_cos = rotary_pos_cos[:total_seq_len:ratio][:rotary_seq_len] + if rotary_pos_sin is not None and ratio > 1: + rotary_pos_sin = rotary_pos_sin[:total_seq_len:ratio][:rotary_seq_len] + + squeeze_head = x.dim() == 3 + if squeeze_head: + x = x.unsqueeze(-2) + if config.apply_rope_fusion: + out = fused_mla_rope_inplace( + x, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + None, + cp_group.rank(), + cp_group.size(), + remove_interleaving=True, + ) + else: + x_nope, x_pe = torch.split(x, [nope_dim, pos_dim], dim=-1) + x_pe = apply_rotary_pos_emb( + x_pe, + rotary_pos_emb, + config=config, + cu_seqlens=None, + mscale=mscale, + cp_group=cp_group, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + out = torch.cat([x_nope, x_pe], dim=-1) + if squeeze_head: + out = out.squeeze(-2) + return out + + +# --------------------------------------------------------------------------- +# Sparse attention kernel (unfused, differentiable) +# --------------------------------------------------------------------------- + + +def unfused_compressed_sparse_attn( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Differentiable sparse attention with MQA and attention sink. + + Args: + query: [sq, b, np, hn] multi-head query. + kv_full: [n_kv, b, hn] single-head KV (original + compressed). + attn_sink: [np] per-head learnable bias. + topk_indices: [b, sq, topk] indices into kv_full (int32, -1 = invalid). + softmax_scale: float + + Returns: + output: [sq, b, np * hn] + """ + sq, b, np_, hn = query.size() + + # --- Gather KV at topk positions --- + # kv_full: [n_kv, b, hn] -> [b, n_kv, hn] + kv_t = kv_full.permute(1, 0, 2) + + safe_indices = topk_indices.clamp(min=0).long() # [b, sq, topk] + safe_indices_exp = safe_indices.unsqueeze(-1).expand(-1, -1, -1, hn) # [b, sq, topk, hn] + # [b, n_kv, hn] -> [b, 1, n_kv, hn] -> gather -> [b, sq, topk, hn] + kv_gathered = torch.gather( + kv_t.unsqueeze(1).expand(-1, sq, -1, -1), dim=2, index=safe_indices_exp + ) + + # --- Attention scores --- + # query: [sq, b, np, hn] -> [b, np, sq, hn] + q = query.permute(1, 2, 0, 3).float() + kv_g = kv_gathered.float() # [b, sq, topk, hn] + + # [b, np, sq, topk] + scores = torch.einsum("bnsh,bskh->bnsk", q, kv_g) * softmax_scale + + # Mask invalid + invalid_mask = (topk_indices < 0).unsqueeze(1) # [b, 1, sq, topk] + scores = scores.masked_fill(invalid_mask, float("-inf")) + + # --- Softmax with attention sink --- + sink = attn_sink.view(1, np_, 1, 1).float() + scores_max = scores.max(dim=-1, keepdim=True).values # [b, np, sq, 1] + scores_max = torch.max(scores_max, sink) + + exp_scores = torch.exp(scores - scores_max) # [b, np, sq, topk] + exp_sink = torch.exp(sink - scores_max) # [1, np, 1, 1] + + sum_exp = exp_scores.sum(dim=-1, keepdim=True) + exp_sink + attn_weights = exp_scores / sum_exp # [b, np, sq, topk] + + # --- Weighted sum --- + output = torch.einsum("bnsk,bskh->bnsh", attn_weights, kv_g) + output = output.to(query.dtype) + + # [b, np, sq, hn] -> [sq, b, np, hn] -> [sq, b, np * hn] + output = output.permute(2, 0, 1, 3).contiguous() + output = output.reshape(sq, b, np_ * hn) + return output + + +@torch.no_grad() +def _compute_unfused_csa_non_compressed_lse( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + window_indices: torch.Tensor, + softmax_scale: float, + chunk_size: int = 512, +) -> torch.Tensor: + """Return the detached sliding-window-plus-sink log mass for the CSA teacher. + + Args: + query: Query tensor in ``[sq, batch, heads, head_dim]`` layout. + kv_full: Original (non-compressed) KV in ``[sk, batch, head_dim]`` layout. + attn_sink: Per-head sink logits in ``[heads]`` layout. + window_indices: Local per-batch window indices in ``[batch, sq, window]`` layout. + softmax_scale: Scale applied to query-key logits. + chunk_size: Maximum number of flattened query rows processed at once. + + Returns: + Detached FP32 log-sum-exp values in ``[batch, heads, sq]`` layout. + """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if query.ndim != 4: + raise ValueError(f"query must have shape [sq, batch, heads, dim], got {query.shape}") + if attn_sink.ndim != 1: + raise ValueError(f"attn_sink must be 1D, got shape {tuple(attn_sink.shape)}") + + seqlen_q, batch_size, num_heads, head_dim = query.shape + if kv_full.ndim != 3 or kv_full.shape[1:] != (batch_size, head_dim): + raise ValueError( + "non-compressed KV must have shape " + f"[sk, {batch_size}, {head_dim}], got {tuple(kv_full.shape)}" + ) + if window_indices.ndim != 3 or window_indices.shape[:2] != (batch_size, seqlen_q): + raise ValueError( + "window_indices must have shape " + f"[{batch_size}, {seqlen_q}, window], got {tuple(window_indices.shape)}" + ) + if attn_sink.numel() != num_heads: + raise ValueError(f"attn_sink must contain {num_heads} values, got {attn_sink.numel()}") + if not (query.device == kv_full.device == attn_sink.device == window_indices.device): + raise ValueError("query, kv_full, attn_sink, and window_indices must share a device") + + n_kv = kv_full.shape[0] + q_flat = query.detach().permute(1, 0, 2, 3).reshape(-1, num_heads, head_dim) + kv_flat = kv_full.detach().permute(1, 0, 2).reshape(-1, head_dim) + batch_offsets = ( + torch.arange(batch_size, device=window_indices.device, dtype=torch.int64) * n_kv + ).view(batch_size, 1, 1) + window_indices_i64 = window_indices.to(dtype=torch.int64) + global_indices = torch.where( + window_indices_i64 >= 0, window_indices_i64 + batch_offsets, window_indices_i64 + ).reshape(batch_size * seqlen_q, -1) + + sink = attn_sink.detach().to(dtype=torch.float32).view(1, num_heads) + lse_chunks = [] + for start in range(0, q_flat.shape[0], chunk_size): + end = min(start + chunk_size, q_flat.shape[0]) + indices = global_indices[start:end] + gathered_kv = kv_flat.index_select(0, indices.clamp(min=0).reshape(-1)).reshape( + end - start, indices.shape[-1], head_dim + ) + window_logits = torch.einsum("rhd,rkd->rhk", q_flat[start:end].float(), gathered_kv.float()) + window_logits = (window_logits * softmax_scale).masked_fill( + (indices < 0).unsqueeze(1), float("-inf") + ) + lse_chunks.append(torch.logaddexp(torch.logsumexp(window_logits, dim=-1), sink)) + + if lse_chunks: + lse_flat = torch.cat(lse_chunks, dim=0) + else: + lse_flat = torch.empty((0, num_heads), dtype=torch.float32, device=query.device) + return lse_flat.reshape(batch_size, seqlen_q, num_heads).permute(0, 2, 1).contiguous() + + +# --------------------------------------------------------------------------- +# Compressor +# --------------------------------------------------------------------------- + + +@dataclass +class CompressorSubmodules: + """Submodule specs for CSA and HCA Compressor.""" + + linear_wkv: Union[ModuleSpec, type] = None + linear_wgate: Union[ModuleSpec, type] = None + norm: Union[ModuleSpec, type] = None + + +class Compressor(MegatronModule): + """Gated pooling compressor for CSA and HCA sparse attention. + + Compresses a sequence of tokens into a shorter sequence by pooling groups of + ``compress_ratio`` tokens using learned gated weights. + + For ``compress_ratio == 4``, overlapping compression is used (``coff = 2``). + For ``compress_ratio == 128``, non-overlapping compression is used (``coff = 1``). + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressorSubmodules, + compress_ratio: int, + head_dim: int, + rotate: bool = False, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + name: str | None = None, + ) -> None: + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError("Compressor requires an explicit ProcessGroupCollection") + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.overlap = compress_ratio == 4 + self.coff = 1 + int(self.overlap) + self.rotate = rotate + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + + self.rotary_pos_emb = rotary_pos_emb + + proj_out_dim = self.coff * head_dim + + with get_fp8_disabled_context(config, is_init=True): + self.linear_wkv = build_module( + submodules.linear_wkv, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wkv") if name is not None else None, + ) + + self.linear_wgate = build_module( + submodules.linear_wgate, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wgate") if name is not None else None, + ) + + # keep to high precision (FP32 in the reference DeepSeek V4 checkpoint) + _ape = torch.empty( + compress_ratio, proj_out_dim, device=torch.cuda.current_device(), dtype=torch.float32 + ) + config.init_method(_ape) + self.ape = mark_keep_in_fp32(nn.Parameter(_ape)) + + norm_config = copy.copy(config) + norm_config.normalization = "RMSNorm" + self.norm = build_module( + submodules.norm, config=norm_config, hidden_size=head_dim, eps=config.layernorm_epsilon + ) + + def backward_dw(self): + """Compute deferred weight gradients for the compressor projections.""" + self.linear_wkv.backward_dw() + self.linear_wgate.backward_dw() + + def _overlap_transform(self, tensor: torch.Tensor, fill_value: float = 0) -> torch.Tensor: + """Apply overlapping window transform for 4x compression. + + Input shape: [n_groups, ratio, b, coff * head_dim] + Output shape: [n_groups, 2 * ratio, b, head_dim] + """ + n_groups, ratio, b_dim, _ = tensor.size() + d = self.head_dim + new_tensor = tensor.new_full((n_groups, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + new_tensor[1:, :ratio] = tensor[:-1, :, :, :d] + return new_tensor + + def _project(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Project compressor values and gates outside any enclosing FP8 context.""" + with get_fp8_disabled_context(self.config): + kv, _ = self.linear_wkv(x) + score, _ = self.linear_wgate(x) + return kv, score + + def forward(self, x: torch.Tensor) -> Optional[torch.Tensor]: + """Compress hidden states into shorter KV sequence. + + Args: + x: [sq, b, hidden_size] + + Returns: + compressed_kv [sq // ratio, b, head_dim] or None if too short. + """ + nvtx_range_push("compressor") + + sq, b, _ = x.size() + ratio = self.compress_ratio + + if sq < ratio: + nvtx_range_pop("compressor") + return None + + kv, score = self._project(x) # [sq, b, coff * head_dim] + + cutoff = (sq // ratio) * ratio + if cutoff < sq: + kv = kv[:cutoff] + score = score[:cutoff] + + n_compressed = cutoff // ratio + + # Reshape: [n_compressed, ratio, b, coff * head_dim] + kv = kv.view(n_compressed, ratio, b, -1) + score = score.view(n_compressed, ratio, b, -1) + + # APE: [ratio, coff * head_dim] -> [1, ratio, 1, coff * head_dim] + score = score + self.ape.view(1, ratio, 1, -1) + + if self.overlap: + kv = self._overlap_transform(kv, fill_value=0) + score = self._overlap_transform(score, fill_value=float("-inf")) + + weights = torch.softmax(score, dim=1, dtype=torch.float32).to(kv.dtype) + kv = (kv * weights).sum(dim=1) # [n_compressed, b, head_dim] + + kv = self.norm(kv.to(x.dtype)) + + kv = _apply_rope( + kv, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + n_compressed, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + if self.rotate: + kv = rotate_activation(kv) + + nvtx_range_pop("compressor") + return kv # [n_compressed, b, head_dim] + + +# --------------------------------------------------------------------------- +# CSAIndexer +# --------------------------------------------------------------------------- + + +@dataclass +class CSAIndexerSubmodules: + """Submodule specs for CSAIndexer.""" + + linear_wq_b: Union[ModuleSpec, type] = None + linear_weights_proj: Union[ModuleSpec, type] = None + compressor: Union[ModuleSpec, type] = None + + +class CSAIndexer(MegatronModule): + """Learned top-k retrieval over compressed positions for CSA sparse attention. + + Computes index scores to select the most relevant compressed KV positions for each + query. Reuses the scoring logic from ``DSAIndexer`` (einsum -> relu -> weight -> sum + -> topk) and ``rotate_activation`` (Hadamard transform) from ``dsa.py``. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CSAIndexerSubmodules, + compress_ratio: int, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + name: str | None = None, + ) -> None: + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError("CSAIndexer requires an explicit ProcessGroupCollection") + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.hidden_size = config.hidden_size + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.q_lora_rank = ( + config.q_lora_rank if config.q_lora_rank is not None else config.hidden_size + ) + + self.index_n_heads = config.dsa_indexer_n_heads + self.index_head_dim = config.dsa_indexer_head_dim + self.index_topk = config.dsa_indexer_topk + + self.softmax_scale: float = self.index_head_dim**-0.5 + + self.rotary_pos_emb = rotary_pos_emb + + # Q projection + self.linear_wq_b = build_module( + submodules.linear_wq_b, + self.q_lora_rank, + self.index_n_heads * self.index_head_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_wq_b") if name is not None else None, + ) + + # The reference DeepSeek V4 checkpoint keeps this projection in BF16. + with get_fp8_disabled_context(config, is_init=True): + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + self.hidden_size, + self.index_n_heads, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + name=(name + ".linear_weights_proj") if name is not None else None, + ) + + # Own compressor (smaller head_dim, with Hadamard rotation) + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=compress_ratio, + head_dim=self.index_head_dim, + rotate=True, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".compressor") if name is not None else None, + ) + + def backward_dw(self): + """Compute deferred weight gradients for the indexer projections.""" + self.linear_wq_b.backward_dw() + self.linear_weights_proj.backward_dw() + self.compressor.backward_dw() + + def _project_weights(self, x: torch.Tensor) -> torch.Tensor: + """Project indexer weights outside any enclosing FP8 context.""" + with get_fp8_disabled_context(self.config): + weights, _ = self.linear_weights_proj(x) + return weights + + def forward_before_topk( + self, x: torch.Tensor, qr: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute Q, compressed K, and weights before top-k selection.""" + nvtx_range_push("indexer_before_topk") + + sq, bsz, _ = x.size() + + # Q path + q, _ = self.linear_wq_b(qr) # [sq, b, n_heads * head_dim] + q = q.reshape(sq, bsz, self.index_n_heads, self.index_head_dim) + q = _apply_rope( + q, + self.index_head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + sq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + q = rotate_activation(q) + + # K path: own compressor + k = self.compressor(x) # [sq//ratio, b, index_head_dim] + + weights = self._project_weights(x) # [sq, b, n_heads] + weights = weights * (self.index_n_heads**-0.5) + + nvtx_range_pop("indexer_before_topk") + return q, k, weights + + def forward( + self, x: torch.Tensor, qr: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Return (index_scores, topk_indices).""" + nvtx_range_push("indexer") + q, k, weights = self.forward_before_topk(x, qr) + nvtx_range_push("indexer_qk_topk") + effective_topk = min(self.index_topk, k.size(0)) + index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, effective_topk, mask) + nvtx_range_pop("indexer_qk_topk") + nvtx_range_pop("indexer") + return index_scores, topk_indices + + +# --------------------------------------------------------------------------- +# CompressedSparseAttention (core attention) +# --------------------------------------------------------------------------- + + +@dataclass +class CompressedSparseAttentionSubmodules: + """Submodule specs for CompressedSparseAttention.""" + + compressor: Union[ModuleSpec, type] = None + indexer: Union[ModuleSpec, type] = None + + +class CompressedSparseAttention(MegatronModule): + """Sparse core attention for CompressedSparseAttention. + + Combines sliding window attention with compressed KV attention. The spec always + provides compressor and indexer submodule specs; this ``__init__`` inspects + ``config.csa_compress_ratios[layer_idx]`` and conditionally builds them: + + * ``ratio == 0``: window-only (compressor and indexer NOT built) + * ``ratio == 4``: window + 4x compressed + learned Indexer (both built) + * ``ratio == 128``: window + 128x compressed, attend to all (compressor built only) + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressedSparseAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + attention_dropout: Optional[float] = None, + softmax_scale: Optional[float] = None, + k_channels: Optional[int] = None, + v_channels: Optional[int] = None, + cp_comm_type: str = "p2p", + pg_collection: Optional[ProcessGroupCollection] = None, + rotary_pos_emb: nn.Module = None, + compress_ratio: int = 0, + is_mtp_layer: bool = False, + name: str | None = None, + ): + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError( + "CompressedSparseAttention requires an explicit ProcessGroupCollection" + ) + self.pg_collection = pg_collection + + self.layer_number = layer_number + self.config.num_layers if is_mtp_layer else layer_number + self.compress_ratio = compress_ratio + self.window_size = config.csa_window_size + self.v_head_dim = config.v_head_dim + + self.n_local_heads = config.num_attention_heads + + if softmax_scale is None: + softmax_scale = config.v_head_dim**-0.5 + self.softmax_scale = softmax_scale + + # Learnable attention sink per head, kept in reference-checkpoint FP32. + self.attn_sink = mark_keep_in_fp32( + nn.Parameter(torch.zeros(self.n_local_heads, dtype=torch.float32)) + ) + + # Conditionally build Compressor (ratio > 1) + if self.compress_ratio > 1 and submodules.compressor is not None: + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=self.compress_ratio, + head_dim=config.v_head_dim, + rotate=False, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".compressor") if name is not None else None, + ) + else: + self.compressor = None + + # Conditionally build Indexer (ratio == 4) + if ( + self.compress_ratio == 4 + and not config.csa_dense_mode + and submodules.indexer is not None + ): + self.indexer = build_module( + submodules.indexer, + config=config, + compress_ratio=self.compress_ratio, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + name=(name + ".indexer") if name is not None else None, + ) + else: + self.indexer = None + + def backward_dw(self): + """Compute deferred gradients for the optional compressor and indexer projections.""" + if self.compressor is not None: + self.compressor.backward_dw() + if self.indexer is not None: + self.indexer.backward_dw() + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + x: torch.Tensor = None, + qr: torch.Tensor = None, + attn_mask_type: AttnMaskType = None, + attention_bias: torch.Tensor = None, + packed_seq_params=None, + ) -> torch.Tensor: + """Forward pass for CompressedSparseAttention. + + Args: + query: [sq, b, np, v_head_dim] + key: [sq, b, 1, v_head_dim] (single-head MQA; head dim squeezed internally) + value: unused (key == value in MQA) + attention_mask: attention mask (may be None for causal). + x: [sq, b, hidden_size] original hidden states. + qr: [sq, b, q_lora_rank] compressed query representation. + + Returns: + output: [sq, b, np * v_head_dim] + """ + nvtx_range_push("compressed_sparse_attn") + assert ( + packed_seq_params is None + ), "Packed sequence not supported for CompressedSparseAttention" + + sq, b, np, hn = query.size() + + # --- Step 1: Prepare single-head KV (squeeze singleton head dim) --- + kv = key.squeeze(-2) # [sq, b, 1, v_head_dim] -> [sq, b, v_head_dim] + + # --- Step 2: Compression --- + if self.compressor is not None and self.compress_ratio > 1: + compressed_kv = self.compressor(x) # [n_compressed, b, v_head_dim] + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + n_compressed = 0 + else: + kv_full = kv + n_compressed = 0 + + offset = sq # compressed indices start after original positions + + # --- Step 3: Window indices --- + window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) + + # --- Step 4: Compressed indices --- + indexer_loss = None + + if self.compress_ratio > 1 and n_compressed > 0: + nvtx_range_push("compressed_indices") + if self.indexer is not None: + x_det = x.detach() + qr_det = qr.detach() + + causal_mask = ( + torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) + ) + positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(b, -1, -1) + ) # [b, sq, n_compressed] + + if self.training and torch.is_grad_enabled(): + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det + ) + indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) + weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, kv, self.attn_sink, window_idxs, self.softmax_scale + ) + topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( + q_indexer, + weights_for_unfused, + k_indexer, + query.detach(), + key_for_loss.detach(), + self.softmax_scale, + min(self.indexer.index_topk, n_compressed), + indexer_loss_coeff, + causal_mask, + self.config.dsa_indexer_use_sparse_loss, + self.indexer.pg_collection, + None, + None, + None, + None, + self.config.calculate_per_token_loss, + True, + non_compressed_lse, + ) + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), + ) + else: + _, topk_indices_compressed = self.indexer(x_det, qr_det, mask=causal_mask) + + n_valid_per_pos = positions // self.compress_ratio # [sq, 1] + valid = (topk_indices_compressed >= 0) & (topk_indices_compressed < n_valid_per_pos) + compress_topk_idxs = torch.where( + valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) + ) + else: + compress_topk_idxs = get_compress_topk_idxs( + self.compress_ratio, b, sq, offset, query.device + ) + + topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1) + nvtx_range_pop("compressed_indices") + else: + topk_idxs = window_idxs + + topk_idxs = topk_idxs.int() + + # --- Step 5: Sparse attention --- + nvtx_range_push("sparse_attn_kernel") + output = unfused_compressed_sparse_attn( + query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale + ) + nvtx_range_pop("sparse_attn_kernel") + + # --- Step 6: Attach indexer loss --- + if indexer_loss is not None and self.training and torch.is_grad_enabled(): + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + nvtx_range_pop("compressed_sparse_attn") + return output diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py new file mode 100644 index 00000000000..df16ba09086 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -0,0 +1,696 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +from dataclasses import dataclass +from typing import NoReturn, Optional, Union + +import torch + +from megatron.core import tensor_parallel +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, + apply_rotary_pos_emb, +) +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.attention import Attention +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.torch_norm import LayerNormBuilder +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.core.typed_torch import apply_module +from megatron.core.utils import get_pg_size, is_te_min_version + +try: + from megatron.core.fusions.fused_mla_yarn_rope_apply import ( + fused_mla_rope_inplace, + fused_mla_rope_out_of_place, + ) +except Exception: + fused_mla_rope_inplace = None + fused_mla_rope_out_of_place = None + + +if HAVE_TE: + from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input +else: + TEColumnParallelLinear, TELinear, set_save_original_input = (None, None, None) + + +@torch.compile +def _q_rms_norm(q: torch.Tensor, eps: float) -> torch.Tensor: + """Fused RMS normalization for query tensor (no learnable weight).""" + return q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + + +@dataclass +class DSv4HybridSelfAttentionSubmodules: + """Submodules for the DSv4HybridAttention layer.""" + + q_layernorm: LayerNormBuilder + kv_layernorm: LayerNormBuilder + + linear_q_down_proj: Union[ModuleSpec, type] = None + linear_q_up_proj: Union[ModuleSpec, type] = None + linear_kv_proj: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + + +class DSv4HybridAttention(Attention): + """DeepSeek-v4 Hybrid Attention layer.""" + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + compress_ratio: Optional[int] = None, + name: str | None = None, + ) -> None: + + if pg_collection is None: + raise ValueError("DSv4 hybrid attention requires an explicit ProcessGroupCollection.") + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attention_type=attention_type, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + name=name, + ) + self.config: MLATransformerConfig + + assert ( + get_pg_size(self.pg_collection.tp) == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + + assert ( + not self.checkpoint_core_attention + ), "Checkpoint core attention is not supported in DSv4 Hybrid Attention." + assert ( + not self.offload_qkv_linear + ), "Offload qkv linear is not supported in DSv4 Hybrid Attention." + + self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads + + self.q_head_dim = self.config.v_head_dim + + self.key_hidden_size = self.q_head_dim + self.val_hidden_size = self.config.v_head_dim + + self.recompute_up_proj = ( + self.config.recompute_granularity == 'selective' + and "mla_up_proj" in self.config.recompute_modules + ) + self.qkv_up_checkpoint = None + + self.softmax_scale = None + + ratio_idx = self.config.num_layers + layer_number - 1 if is_mtp_layer else layer_number - 1 + if compress_ratio is None: + compress_ratio = self.config.csa_compress_ratios[ratio_idx] + use_compressed_yarn = compress_ratio > 1 + rope_base = ( + self.config.csa_compress_rotary_base if use_compressed_yarn else self.config.rotary_base + ) + self._dsv4_compress_ratio = compress_ratio + self._dsv4_rope_base = rope_base + self._dsv4_uses_yarn_rope = use_compressed_yarn + if not use_compressed_yarn: + self.rotary_pos_emb = RotaryEmbedding( + self.config.qk_pos_emb_head_dim, + rotary_percent=self.config.rotary_percent, + rotary_base=rope_base, + cp_group=self.pg_collection.cp, + ) + else: + self.rotary_pos_emb = YarnRotaryEmbedding( + self.config.qk_pos_emb_head_dim, + rotary_base=rope_base, + scaling_factor=self.config.rotary_scaling_factor, + original_max_position_embeddings=self.config.original_max_position_embeddings, + beta_fast=self.config.beta_fast, + beta_slow=self.config.beta_slow, + mscale=self.config.mscale, + mscale_all_dim=self.config.mscale_all_dim, + cp_group=self.pg_collection.cp, + ) + + core_attn_extra_kwargs = { + "rotary_pos_emb": self.rotary_pos_emb, + "compress_ratio": compress_ratio, + "is_mtp_layer": is_mtp_layer, + "name": (name + ".core_attention") if name is not None else None, + } + self.core_attention = build_module( + submodules.core_attention, + config=self.config, + layer_number=self.layer_number, + attn_mask_type=self.attn_mask_type, + attention_type=self.attention_type, + softmax_scale=self.softmax_scale, + k_channels=self.q_head_dim, + v_channels=self.config.v_head_dim, + cp_comm_type=cp_comm_type, + pg_collection=self.pg_collection, + **core_attn_extra_kwargs, + ) + + # Output. + self.o_local_groups = self.config.o_groups + assert ( + self.query_projection_size % self.config.o_groups == 0 + ), "num_attention_heads * v_head_dim must be divisible by o_groups" + group_proj_in_size = self.query_projection_size // self.config.o_groups + group_proj_out_size = self.config.o_groups * self.config.o_lora_rank + + _linear_o_group_proj = torch.empty( + group_proj_out_size, + group_proj_in_size, + device=torch.cuda.current_device(), + dtype=self.config.params_dtype, + ) + self.config.init_method(_linear_o_group_proj) + self.linear_o_group_proj = torch.nn.Parameter(_linear_o_group_proj) + + linear_proj_in_size = self.config.o_groups * self.config.o_lora_rank + + self.linear_proj = build_module( + submodules.linear_proj, + linear_proj_in_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=self.config.add_bias_linear, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name='proj', + tp_group=self.pg_collection.tp, + ) + + if ( + HAVE_TE + and isinstance(self.linear_proj, TELinear) + and ( + ( + self.config.fp8 + and self.config.fp8_recipe != 'delayed' + and is_te_min_version("2.6.0dev0") + ) + or (self.config.fp4 and is_te_min_version("2.7.0.dev0")) + ) + ): + # For fp8/fp4 training, the output of the fused core_attn is saved by itself, and + # linear_proj also saves the quantized tensor of this output. Here we set the + # linear_proj to save the original input tensors to avoid the extra memory usage of + # the quantized tensor. + set_save_original_input(self.linear_proj) + + def forward( + self, + hidden_states, + attention_mask, + key_value_states=None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + attention_bias=None, + packed_seq_params=None, + position_ids=None, + sequence_len_offset=None, + *, + inference_params=None, + ): + """Forward pass for DeepSeek-v4 Hybrid Attention""" + assert ( + rotary_pos_emb is None + ), "Rotary position embeddings should not be passed into DSv4HybridAttention." + assert ( + attention_bias is None + ), "Attention bias should not be passed into DSv4HybridAttention." + assert ( + rotary_pos_cos is None and rotary_pos_sin is None + ), "DSv4HybridAttention does not support Flash Decoding" + assert ( + not rotary_pos_cos_sin + ), "Flash-infer rope has not been tested with DSv4HybridAttention." + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridAttention." + assert ( + packed_seq_params is None + ), "Packed sequence is not supported for DSv4HybridAttention." + + # ===================== + # Query, Key, and Value + # ===================== + # Get the query, key and value tensors based on the type of attention - + # self or cross attn. + query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( + hidden_states, key_value_states, position_ids, None, inference_context=inference_context + ) + + # TODO: Currently, TE can only accept contiguous tensors for MLA + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + # ================================== + # core attention computation + # ================================== + # Need corresponding TE change + core_attn_manager = off_interface( + self.offload_core_attention and self.training, query, "core_attn" + ) + with core_attn_manager as query: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + packed_seq_params=None, + x=hidden_states, + qr=q_compressed, + ) + core_attn_out = core_attn_manager.group_offload( + core_attn_out, forced_released_tensors=[query, key, value] + ) + + if self.recompute_up_proj: + assert self.qkv_up_checkpoint is not None + self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) + self.qkv_up_checkpoint = None + + # inverse RoPE on last qk_pos_emb_head_dim of each head + seq_len = core_attn_out.size(0) + n_heads = self.num_attention_heads_per_partition + pos_dim = self.config.qk_pos_emb_head_dim + nope_dim = self.config.v_head_dim - pos_dim + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), n_heads, -1) + rope_seqlen = seq_len + # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's + # concentration factor (mscale) is NOT part of the DSv4 model contract -- + # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if self.config.apply_rope_fusion: + # ``mscale=1.0`` strips yarn's concentration factor from the + # cached cos/sin so the fused kernel matches the unfused + # path's forced ``mscale=1.0`` (DSv4 "pure rotation"). + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rope_seqlen, dtype=hidden_states.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + elif self._dsv4_uses_yarn_rope: + rotary_pos_emb, _ = self.rotary_pos_emb(rope_seqlen, packed_seq=False) + else: + rotary_pos_emb = self.rotary_pos_emb(rope_seqlen, packed_seq=False) + if self.config.apply_rope_fusion: + # Fused DSA backward retains the raw attention output O. Applying + # inverse RoPE to its view in-place corrupts the retained O used by + # the softmax backward, so this call needs private storage. + assert fused_mla_rope_out_of_place is not None + core_attn_out = fused_mla_rope_out_of_place( + core_attn_out, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + None, + self.pg_collection.cp.rank(), + self.pg_collection.cp.size(), + inverse=True, + remove_interleaving=True, + ) + else: + content_part, rot_part = torch.split( + core_attn_out, [core_attn_out.size(-1) - pos_dim, pos_dim], dim=-1 + ) + rot_part = apply_rotary_pos_emb( + rot_part, + rotary_pos_emb, + self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + inverse=True, + mla_output_remove_interleaving=True, + ) + core_attn_out = torch.cat([content_part, rot_part], dim=-1) + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), -1) + + # Grouped output + core_attn_out = core_attn_out.view( + core_attn_out.size(0), core_attn_out.size(1), self.o_local_groups, -1 + ) + wo_a_weight = self.linear_o_group_proj.view( + self.o_local_groups, self.config.o_lora_rank, -1 + ) + core_attn_out = torch.einsum("...gd,grd->...gr", core_attn_out, wo_a_weight) + core_attn_out = core_attn_out.reshape(*core_attn_out.shape[:-2], -1) + + # ================= + # Output. [sq, b, h] + # ================= + attn_proj_manager = off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") + with attn_proj_manager as core_attn_out: + output, bias = self.linear_proj(core_attn_out) + output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) + + return output, bias + + +class DSv4HybridSelfAttention(DSv4HybridAttention): + """DSv4Hybrid Self-attention layer class + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type=AttnMaskType.padding, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + compress_ratio: Optional[int] = None, + name: str | None = None, + ): + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + attention_type="self", + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + compress_ratio=compress_ratio, + name=name, + ) + + q_down_proj_kwargs = {} + if submodules.linear_q_down_proj in [TELinear]: + q_down_proj_kwargs['parallel_mode'] = 'duplicated' + else: + raise ValueError(f"Unsupported linear_q_down_proj: {submodules.linear_q_down_proj}") + + self.linear_q_down_proj = build_module( + submodules.linear_q_down_proj, + self.config.hidden_size, + self.config.q_lora_rank, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='q_down_proj', + skip_weight_param_allocation=False, + tp_group=None, + name=(name + ".linear_q_down_proj") if name is not None else None, + **q_down_proj_kwargs, + ) + + self.linear_q_up_proj = build_module( + submodules.linear_q_up_proj, + self.config.q_lora_rank, + self.config.num_attention_heads * self.q_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='q_up_proj', + tp_group=pg_collection.tp, + name=(name + ".linear_q_up_proj") if name is not None else None, + ) + + self.linear_kv_proj = build_module( + submodules.linear_kv_proj, + self.config.hidden_size, + self.config.v_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='kv_up_proj', + tp_group=pg_collection.tp, + name=(name + ".linear_kv_proj") if name is not None else None, + ) + self.kv_layernorm = submodules.kv_layernorm( + hidden_size=self.config.v_head_dim, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + self.q_layernorm = submodules.q_layernorm( + hidden_size=self.config.q_lora_rank, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + def get_query_key_value_tensors( + self, + hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + inference_context=None, + *, + inference_params=None, + ): + """ + Derives `query`, `key` and `value` tensors from `hidden_states`. + """ + # s = sequence length, b = batch size, h = hidden size, n = num attention heads + # Attention heads [s, b, n*h] + assert ( + hidden_states.ndim == 3 + ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" + assert ( + packed_seq_params is None + ), "Packed sequence is not supported for DSv4HybridAttention." + + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridSelfAttention." + + # ========================================= + # Prepare RoPE and seqlen related params + # ========================================= + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + inference_context, None, hidden_states, self.config, None + ) + + # rotary_pos_emb:[s, b, 1, 64] + # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's + # concentration factor (mscale) is NOT part of the DSv4 model contract -- + # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if self.config.apply_rope_fusion: + # ``mscale=1.0`` strips yarn's concentration factor from the + # cached cos/sin so the fused kernel matches the unfused + # path's forced ``mscale=1.0`` (DSv4 "pure rotation"). + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rotary_seq_len, dtype=hidden_states.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + elif self._dsv4_uses_yarn_rope: + rotary_pos_emb, _ = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + else: + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + + # ========================================= + # QKV down projection and layernorm + # ========================================= + # q_compressed: [s, b, q_lora_rank] + q_compressed, _ = self.linear_q_down_proj(hidden_states) + + kv_compressed = hidden_states + k_pos_emb = None + + # ========================================= + # Apply norm + # ========================================= + + if self.config.q_lora_rank is not None: + # q_compressed: [num_tokens, q_lora_rank] + q_compressed = apply_module(self.q_layernorm)(q_compressed) + + # ========================================= + # QKV up projection and RoPE apply + # ========================================= + + def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): + """Apply the up projection and RoPE to the SBHD query and key.""" + # q_compressed: [s, b, q_lora_rank] + # q: [s, b, n * (qk_head_dim + qk_pos_emb_head_dim)] + q, _ = self.linear_q_up_proj(q_compressed) + + # q: [num_tokens, n, q_head_dim] + q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) + q = _q_rms_norm(q, self.config.layernorm_epsilon) + + kv, _ = self.linear_kv_proj(kv_compressed) + kv = self.kv_layernorm(kv) + + # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] + if k_pos_emb is not None: + k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + + if self.config.apply_rope_fusion: + cp_rank = self.pg_collection.cp.rank() + cp_size = self.pg_collection.cp.size() + query = fused_mla_rope_inplace( + q, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + None, + cp_rank, + cp_size, + remove_interleaving=True, + ) + kv = kv.unsqueeze(-2) + kv = fused_mla_rope_inplace( + kv, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + None, + cp_rank, + cp_size, + remove_interleaving=True, + ) + key = kv + value = kv + else: + q_len = q.size()[0] + # Keep direct forward calls with shorter sequences aligned to their inputs. + rotary_pos_emb = rotary_pos_emb[0:q_len] + + # q_no_pe: [num_tokens, n, qk_head_dim] + # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] + q_no_pe, q_pos_emb = torch.split( + q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 + ) + + # RoPE and query (shared for wkv and latent) + # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] + q_pos_emb = apply_rotary_pos_emb( + q_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + # query: [num_tokens, n, (qk_head_dim + v_head_dim)] + query = torch.cat([q_no_pe, q_pos_emb], dim=-1) + + pos_dim = self.config.qk_pos_emb_head_dim + kv_no_pe, k_pos_emb = torch.split(kv, [kv.size(-1) - pos_dim, pos_dim], dim=-1) + + # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] + k_pos_emb = apply_rotary_pos_emb( + k_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + + # Single head: key = value = [num_tokens, 1, v_head_dim] + kv = torch.cat([kv_no_pe, k_pos_emb], dim=-1).unsqueeze(-2) + key = kv + value = kv + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + return query, key, value + + if self.recompute_up_proj: + quantization = self.config.fp8 or self.config.fp4 + self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) + query, key, value = self.qkv_up_checkpoint.checkpoint( + qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + else: + query, key, value = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + + return query, key, value, q_compressed, kv_compressed + + def backward_dw(self) -> NoReturn: + """Execute weight gradient computation""" + self._backward_kv_proj() + self._backward_q_proj() + self.core_attention.backward_dw() + self._backward_output_proj() + + def _backward_kv_proj(self): + """Computes weight gradients of KV projection layers""" + self.linear_kv_proj.backward_dw() + + def _backward_q_proj(self): + """Computes weight gradients of Q projection layers""" + self.linear_q_down_proj.backward_dw() + self.linear_q_up_proj.backward_dw() + + def _backward_output_proj(self): + """Computes weight gradients of output projection layer""" + self.linear_proj.backward_dw() + + def set_for_recompute_input_layernorm(self): + """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.""" + set_save_original_input(self.linear_q_down_proj) + set_save_original_input(self.linear_kv_proj) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index bfead2a25c5..fe8a1456208 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -7,7 +7,6 @@ import torch -from megatron.core import parallel_state from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -296,32 +295,72 @@ def save_loss_to_tracker( return tracker = DSAIndexerLossLoggingHelper.tracker + # Hybrid MTP layer numbers can exceed ``num_layers + mtp_num_layers`` + # because every prediction depth can contain multiple hybrid layers. + needed = max(num_layers, layer_number) if "values" not in tracker: - tracker["values"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + tracker["values"] = torch.zeros(needed, device=torch.cuda.current_device()) + elif tracker["values"].shape[0] < needed: + grown = torch.zeros( + needed, device=tracker["values"].device, dtype=tracker["values"].dtype + ) + grown[: tracker["values"].shape[0]] = tracker["values"] + tracker["values"] = grown tracker["values"][layer_number - 1] += loss.detach() tracker["reduce_group"] = reduce_group tracker["avg_group"] = avg_group @staticmethod - def clean_loss_in_tracker(): + def clean_loss_in_tracker(preserve_groups: bool = False): """Clear the indexer losses.""" tracker = DSAIndexerLossLoggingHelper.tracker + reduce_group = tracker.get("reduce_group") if preserve_groups else None + avg_group = tracker.get("avg_group") if preserve_groups else None if "values" in tracker: tracker["values"].zero_() - tracker["reduce_group"] = None - tracker["avg_group"] = None + tracker["reduce_group"] = reduce_group + tracker["avg_group"] = avg_group @staticmethod - def reduce_loss_in_tracker(): - """Collect and reduce the indexer losses across ranks.""" + def reduce_loss_in_tracker( + pg_collection: ProcessGroupCollection, num_layers: Optional[int] = None + ): + """Collect and reduce indexer losses across every pipeline rank. + + Args: + pg_collection: Process groups used for pipeline and data-parallel reductions. + num_layers: Total number of decoder and MTP layers. When provided, ranks without + local indexer losses contribute zeros to the pipeline-wide reduction. + """ tracker = DSAIndexerLossLoggingHelper.tracker - if "values" not in tracker: + pp_group = pg_collection.pp + + # Pipeline ranks can own different attention variants, so first agree on + # a common tracker size. Cache the result because layer allocation is + # static and the negotiation requires a device-to-host synchronization. + if tracker.get("agreed_size") is not None: + size = tracker["agreed_size"] + else: + local_size = tracker["values"].shape[0] if "values" in tracker else (num_layers or 0) + size_t = torch.tensor( + [local_size], device=torch.cuda.current_device(), dtype=torch.long + ) + torch.distributed.all_reduce(size_t, op=torch.distributed.ReduceOp.MAX, group=pp_group) + size = int(size_t.item()) + tracker["agreed_size"] = size + if size == 0: return + if "values" not in tracker: + tracker["values"] = torch.zeros(size, device=torch.cuda.current_device()) + elif tracker["values"].shape[0] < size: + grown = torch.zeros( + size, device=tracker["values"].device, dtype=tracker["values"].dtype + ) + grown[: tracker["values"].shape[0]] = tracker["values"] + tracker["values"] = grown values = tracker["values"] - torch.distributed.all_reduce( - values, group=parallel_state.get_pipeline_model_parallel_group() - ) + torch.distributed.all_reduce(values, group=pp_group) # Reduce indexer losses across ranks. if tracker.get('reduce_group') is not None: torch.distributed.all_reduce(values, group=tracker.get('reduce_group')) @@ -330,9 +369,7 @@ def reduce_loss_in_tracker(): values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG ) torch.distributed.all_reduce( - values, - group=parallel_state.get_data_parallel_group(with_context_parallel=False), - op=torch.distributed.ReduceOp.AVG, + values, group=pg_collection.dp, op=torch.distributed.ReduceOp.AVG ) @staticmethod @@ -340,9 +377,13 @@ def track_indexer_metrics( loss_scale: float, iteration: int, writer, + pg_collection: ProcessGroupCollection, wandb_writer=None, total_loss_dict=None, per_layer_logging: bool = False, + num_layers: Optional[int] = None, + num_indexer_layers: Optional[int] = None, + preserve_groups: bool = False, ): """Track the sparse attention indexer metrics for logging. @@ -350,20 +391,27 @@ def track_indexer_metrics( loss_scale: Scale factor for the loss. iteration: Current training iteration. writer: TensorBoard writer. + pg_collection: Process groups used for pipeline and data-parallel reductions. wandb_writer: Weights & Biases writer. total_loss_dict: Dictionary to accumulate total losses. per_layer_logging: Whether to log per-layer losses. + num_layers: Total number of decoder and MTP layers. Passing it makes ranks + without a local indexer participate in the pipeline reduction. + num_indexer_layers: Number of layers that own an indexer. Defaults to the + tracker size when every tracked layer owns one. + preserve_groups: Keep the saved reduction groups for CUDA Graph replays. """ - DSAIndexerLossLoggingHelper.reduce_loss_in_tracker() + DSAIndexerLossLoggingHelper.reduce_loss_in_tracker( + pg_collection=pg_collection, num_layers=num_layers + ) tracker = DSAIndexerLossLoggingHelper.tracker if "values" not in tracker: return indexer_loss_values = tracker["values"] * loss_scale - num_layers = indexer_loss_values.shape[0] - - # Average across all layers (assuming all layers have sparse attention) - avg_indexer_loss = indexer_loss_values.sum() / num_layers + if num_indexer_layers is None: + num_indexer_layers = indexer_loss_values.shape[0] + avg_indexer_loss = indexer_loss_values.sum() / max(num_indexer_layers, 1) # Log average loss if total_loss_dict is not None: @@ -378,7 +426,7 @@ def track_indexer_metrics( if wandb_writer is not None: wandb_writer.log({"indexer loss": avg_indexer_loss}, iteration) - DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + DSAIndexerLossLoggingHelper.clean_loss_in_tracker(preserve_groups=preserve_groups) def compute_dsa_indexer_loss( @@ -396,6 +444,7 @@ def compute_dsa_indexer_loss( key_positions: Optional[torch.Tensor] = None, query_valid_rows: Optional[torch.Tensor] = None, calculate_per_token_loss: bool = False, + non_compressed_lse: torch.Tensor | None = None, ) -> torch.Tensor: """ Compute KL divergence loss between index_scores and true attention_scores. @@ -421,6 +470,10 @@ def compute_dsa_indexer_loss( varlen_starts: Optional row-wise key start bounds [sq] for packed THD. varlen_ends: Optional row-wise key end bounds [sq] for packed THD. key_positions: Optional global key positions [sk] for packed THD. + non_compressed_lse: Optional detached FP32 log-sum-exp contribution + [batch, heads, seqlen_q] from teacher keys that are intentionally + omitted from ``key``. When provided, the selected ``key`` logits + are normalized with this external mass before heads are summed. Returns: index_loss: KL divergence loss (scalar). @@ -489,8 +542,8 @@ def compute_dsa_indexer_loss( attention_valid_mask = index_valid_mask if sparse_loss else base_valid_mask # [b, np, sq, sk] -> [b, np, sq, sk] - attention_scores = dsa_masking.masked_softmax( - attention_scores.float(), attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk), dim=-1 + attention_scores = _compute_indexer_teacher_probabilities( + attention_scores, attention_valid_mask, non_compressed_lse=non_compressed_lse ) # [b, sq, sk] -> [b, sq, sk] index_log_scores = dsa_masking.masked_log_softmax( @@ -504,7 +557,7 @@ def compute_dsa_indexer_loss( # attention scores are scattered to TP ranks in head dimension. torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) # The target is already non-negative because it is a sum of softmax probabilities. - attention_scores = dsa_indexer_loss.normalize_indexer_target(attention_scores) + attention_scores = _normalize_indexer_teacher_target(attention_scores, non_compressed_lse) return dsa_indexer_loss.indexer_loss_from_target( attention_scores, index_log_scores, @@ -514,6 +567,52 @@ def compute_dsa_indexer_loss( ) +def _compute_indexer_teacher_probabilities( + attention_scores: torch.Tensor, + attention_valid_mask: torch.Tensor, + non_compressed_lse: torch.Tensor | None = None, +) -> torch.Tensor: + """Normalize selected teacher logits, optionally with omitted attention mass. + + ``non_compressed_lse`` is a sufficient statistic for teacher logits that + must participate in the softmax denominator but must not appear in the + compressed-key target returned by this helper. + """ + b, np, sq, sk = attention_scores.shape + expanded_valid_mask = attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk) + if non_compressed_lse is None: + return dsa_masking.masked_softmax(attention_scores.float(), expanded_valid_mask, dim=-1) + + expected_shape = (b, np, sq) + if tuple(non_compressed_lse.shape) != expected_shape: + raise ValueError( + "non_compressed_lse must have shape [batch, heads, seqlen_q], " + f"got {tuple(non_compressed_lse.shape)}, expected {expected_shape}" + ) + if non_compressed_lse.device != attention_scores.device: + raise ValueError( + "non_compressed_lse and attention_scores must be on the same device, " + f"got {non_compressed_lse.device} and {attention_scores.device}" + ) + if non_compressed_lse.requires_grad: + raise ValueError("non_compressed_lse must be detached") + + masked_scores = attention_scores.float().masked_fill(~expanded_valid_mask, float("-inf")) + compressed_lse = torch.logsumexp(masked_scores, dim=-1) + full_lse = torch.logaddexp(non_compressed_lse.float(), compressed_lse) + probabilities = torch.exp(masked_scores - full_lse.unsqueeze(-1)) + return torch.where(expanded_valid_mask, probabilities, torch.zeros_like(probabilities)) + + +def _normalize_indexer_teacher_target( + target: torch.Tensor, non_compressed_lse: torch.Tensor | None +) -> torch.Tensor: + """L1-normalize teacher mass without changing the legacy DSA path.""" + if non_compressed_lse is None: + return dsa_indexer_loss.normalize_indexer_target(target) + return target / target.sum(dim=-1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny) + + def _compute_index_scores( q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor, use_relu: bool = True ) -> torch.Tensor: @@ -627,6 +726,7 @@ def fwd_fused_indexer_loss_naive( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """Naive implementation of forward pass for indexer loss.""" index_scores, topk_indices = fused_qk_topk_naive( @@ -656,6 +756,7 @@ def fwd_fused_indexer_loss_naive( key_positions=key_positions, query_valid_rows=query_valid_rows, calculate_per_token_loss=calculate_per_token_loss, + non_compressed_lse=non_compressed_lse, ) return topk_indices, indexer_loss @@ -680,6 +781,7 @@ def bwd_fused_indexer_loss_naive( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """Naive implementation of backward pass for indexer loss.""" query, _ = dsa_layout.ensure_sbhd(query, "query") @@ -752,8 +854,8 @@ def bwd_fused_indexer_loss_naive( else: index_valid_mask = base_valid_mask attention_valid_mask = index_valid_mask if sparse_loss else base_valid_mask - attention_scores_softmax = dsa_masking.masked_softmax( - attention_scores.float(), attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk), dim=-1 + attention_scores_softmax = _compute_indexer_teacher_probabilities( + attention_scores, attention_valid_mask, non_compressed_lse=non_compressed_lse ) # Free attention_scores immediately del attention_scores @@ -776,7 +878,9 @@ def bwd_fused_indexer_loss_naive( # L1 normalize. Fully masked packed/varlen rows can have zero summed # attention mass; clamp the denominator so those rows stay finite and are # later zeroed by the row-valid loss mask. - attention_scores_normalized = dsa_indexer_loss.normalize_indexer_target(attention_scores_sum) + attention_scores_normalized = _normalize_indexer_teacher_target( + attention_scores_sum, non_compressed_lse + ) # Free attention_scores_sum - no longer needed after normalization del attention_scores_sum @@ -890,6 +994,7 @@ def bwd_fused_indexer_loss_naive( "query_valid_rows", "calculate_per_token_loss", "use_relu", + "non_compressed_lse", ) @@ -916,6 +1021,7 @@ def forward( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """ Fused forward: index_scores never materialized in full. @@ -938,10 +1044,17 @@ def forward( query_valid_rows=query_valid_rows, calculate_per_token_loss=calculate_per_token_loss, use_relu=use_relu, + non_compressed_lse=non_compressed_lse, ) # Save for backward (recomputation strategy) - ctx.save_for_backward(q, weights, k, query, key, topk_indices) + saved_non_compressed_lse = ( + non_compressed_lse + if non_compressed_lse is not None + else q.new_empty(0, dtype=torch.float32) + ) + ctx.save_for_backward(q, weights, k, query, key, topk_indices, saved_non_compressed_lse) + ctx.has_non_compressed_lse = non_compressed_lse is not None ctx.softmax_scale = softmax_scale ctx.loss_coeff = loss_coeff ctx.sparse_loss = sparse_loss @@ -953,6 +1066,7 @@ def forward( ctx.query_valid_rows = query_valid_rows ctx.calculate_per_token_loss = calculate_per_token_loss ctx.use_relu = use_relu + ctx.num_inputs = len(ctx.needs_input_grad) return topk_indices, loss @@ -961,7 +1075,8 @@ def backward(ctx, grad_topk_indices, grad_loss): """ Backward: Recompute what we need. """ - q, weights, k, query, key, topk_indices = ctx.saved_tensors + q, weights, k, query, key, topk_indices, saved_non_compressed_lse = ctx.saved_tensors + non_compressed_lse = saved_non_compressed_lse if ctx.has_non_compressed_lse else None grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive( q, @@ -982,6 +1097,7 @@ def backward(ctx, grad_topk_indices, grad_loss): query_valid_rows=ctx.query_valid_rows, calculate_per_token_loss=ctx.calculate_per_token_loss, use_relu=ctx.use_relu, + non_compressed_lse=non_compressed_lse, ) grad_by_name = { @@ -991,8 +1107,10 @@ def backward(ctx, grad_topk_indices, grad_loss): # query and key are detached in forward, so return None for their gradients. "query": None, "key": None, + "non_compressed_lse": None, } - return tuple(grad_by_name.get(name) for name in _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES) + gradients = tuple(grad_by_name.get(name) for name in _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES) + return gradients[: ctx.num_inputs] class DSAIndexerLossAutoScaler(torch.autograd.Function): diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index 558b1b07a15..bf28600a1aa 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -433,6 +433,40 @@ def float_conversion(val): return conversion_helper(val, float_conversion) +def mark_keep_in_fp32(tensor: torch.Tensor) -> torch.Tensor: + """Mark a parameter or buffer so that ``Float16Module`` keeps it in FP32. + + Args: + tensor: The parameter or buffer to mark. + + Returns: + The same tensor, for call-site convenience. + """ + tensor.keep_in_fp32 = True + return tensor + + +def convert_module_to_dtype_except_fp32_marked( + module: torch.nn.Module, dtype: torch.dtype +) -> torch.nn.Module: + """Cast floating-point parameters and buffers except those marked to stay in FP32. + + Args: + module: The module to convert in place. + dtype: The target floating-point dtype. + + Returns: + The converted module. + """ + return module._apply( + lambda tensor: ( + tensor.to(dtype) + if tensor.is_floating_point() and not getattr(tensor, 'keep_in_fp32', False) + else tensor + ) + ) + + class Float16Module(MegatronModule): """Float 16 Module. @@ -455,13 +489,17 @@ def __init__(self, config: TransformerConfig, module: torch.nn.Module): self.pg_collection = getattr(module, 'pg_collection', None) if self.fp16: - self.add_module('module', module.half()) + self.add_module( + 'module', convert_module_to_dtype_except_fp32_marked(module, torch.half) + ) def float16_convertor(val): return val.half() elif self.bf16: - self.add_module('module', module.bfloat16()) + self.add_module( + 'module', convert_module_to_dtype_except_fp32_marked(module, torch.bfloat16) + ) def float16_convertor(val): return val.bfloat16() diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index aa21e78ce86..f42dacb9bb1 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -144,6 +144,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ) -> None: # TODO(nschank): Restructure so that the Attention initializer knows which specific @@ -156,6 +157,7 @@ def __init__( attn_mask_type=attn_mask_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) self.config: MLATransformerConfig @@ -484,6 +486,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): if pg_collection is None: @@ -498,6 +501,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) @@ -728,8 +732,16 @@ def get_query_key_value_tensors( cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded else: cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + rope_max_seqlen_q = packed_seq_params.max_seqlen_q + rope_max_seqlen_kv = packed_seq_params.max_seqlen_kv + rope_freqs_max_seqlen = ( + max(rope_max_seqlen_q, rope_max_seqlen_kv) + if rope_max_seqlen_q is not None and rope_max_seqlen_kv is not None + else None + ) else: cu_seqlens_q = cu_seqlens_kv = None + rope_freqs_max_seqlen = None # ========================================= # QKV down projection and layernorm @@ -941,6 +953,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_freqs_max_seqlen, ) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( @@ -951,6 +964,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po mscale=mscale, cp_group=self.pg_collection.cp, mla_rotary_interleaved=True, + max_seqlen=rope_freqs_max_seqlen, ) # query: [num_tokens, n, (qk_head_dim + v_head_dim)] @@ -1239,6 +1253,7 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, pp_layer_offset: Optional[int] = None, name: str | None = None, ): @@ -1254,6 +1269,7 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, pp_layer_offset=pp_layer_offset, name=name, ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 66cfaded213..327bc22fed6 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -291,8 +291,10 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa']] = None - """Type of attention variant to use. Currently support gated_delta_net and dsa.""" + experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = ( + None + ) + """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4_hybrid.""" experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None """Optional hook for experimental attention variants to receive the main loss scale.""" @@ -343,6 +345,22 @@ class TransformerConfig(ModelParallelConfig): dsa_indexer_k_norm_fp32: bool = False """Whether DSA indexer key LayerNorm should run on fp32 inputs.""" + #################### + # Compressed sparse attention + #################### + csa_window_size: int = 128 + """Sliding window size for compressed sparse attention.""" + + csa_compress_ratios: Optional[List[int]] = None + """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" + + csa_compress_rotary_base: float = 40000.0 + """RoPE base for compressed KV positions in compressed sparse attention.""" + + csa_dense_mode: bool = False + """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be + disabled.""" + #################### # linear attention #################### @@ -1396,6 +1414,31 @@ def __post_init__(self): "dsa_indexer_skip_topk_offset must be non-negative, got " f"{self.dsa_indexer_skip_topk_offset}." ) + elif self.experimental_attention_variant == "dsv4_hybrid": + assert self.multi_latent_attention, "DSv4 Hybrid requires multi_latent_attention." + assert self.csa_compress_ratios is not None, "csa_compress_ratios must be set" + mtp_layers = self.mtp_num_layers or 0 + expected_len = self.num_layers + mtp_layers + assert len(self.csa_compress_ratios) >= expected_len, ( + f"csa_compress_ratios length ({len(self.csa_compress_ratios)}) must be at least " + f"num_layers + mtp_num_layers ({self.num_layers} + {mtp_layers} = {expected_len})" + ) + assert all( + ratio in [0, 4, 128] for ratio in self.csa_compress_ratios + ), "csa_compress_ratios must be 0, 4, or 128" + assert ( + self.tensor_model_parallel_size == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + assert ( + self.context_parallel_size == 1 + ), "DSv4 Hybrid Attention does not support context parallelism yet." + assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." + if self.dsa_kernel_backend != "none": + raise ValueError( + "The native SBHD DSv4 slice requires dsa_kernel_backend='none'; " + "fused DSv4 backends are added by the follow-up kernel integration." + ) + self.hetereogenous_dist_checkpoint = True if self.fp8: # cannot support first last layer bf16 with delayed scaling @@ -2969,10 +3012,12 @@ class MLATransformerConfig(TransformerConfig): """Rank of Query tensor's low rank representation.""" kv_lora_rank: int = 512 - """Rank of Key and Value tensors' low rank representation.""" + """Rank of Key and Value tensors' low rank representation. + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_head_dim: int = 128 - """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim""" + """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_pos_emb_head_dim: int = 64 """Dimension of the position embedding in the QK projection.""" @@ -3010,6 +3055,12 @@ class MLATransformerConfig(TransformerConfig): mscale_all_dim: float = 0.0 """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + o_groups: int = 8 + """Number of groups for grouped low-rank output projection (wo_a).""" + + o_lora_rank: int = 1024 + """Low-rank dimension per group for grouped output (wo_a). Used when o_groups > 0.""" + cache_mla_latents: bool = False """Cache the low dimensional tensors for MLA rather than full KV cache. This is only for the dynamic inference backend and requires that @@ -3022,12 +3073,41 @@ class MLATransformerConfig(TransformerConfig): def __post_init__(self): super().__post_init__() - if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": + if ( + self.multi_latent_attention + and self.apply_rope_fusion + and self.rope_type != "yarn" + and self.experimental_attention_variant != "dsv4_hybrid" + ): raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") if self.attention_output_gate: raise NotImplementedError("Output gate is not supported for MLA yet.") + # DSv4 hybrid: derive qk_head_dim and kv_lora_rank from v_head_dim and qk_pos_emb_head_dim. + if self.experimental_attention_variant == "dsv4_hybrid": + assert ( + not self.mla_down_proj_fusion + ), "MLA down projection fusion must be disabled for DSv4 hybrid mode." + assert self.q_lora_rank is not None, "DSv4 hybrid mode requires q_lora_rank." + assert self.o_groups > 0, "DSv4 hybrid mode requires o_groups to be positive." + assert self.o_lora_rank > 0, "DSv4 hybrid mode requires o_lora_rank to be positive." + assert ( + self.num_attention_heads * self.v_head_dim + ) % self.o_groups == 0, ( + "num_attention_heads * v_head_dim must be divisible by o_groups." + ) + log_single_rank( + logger, + logging.WARNING, + "DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from " + "v_head_dim and qk_pos_emb_head_dim", + ) + derived = self.v_head_dim - self.qk_pos_emb_head_dim + assert derived > 0, "v_head_dim must be greater than qk_pos_emb_head_dim." + self.qk_head_dim = derived + self.kv_lora_rank = derived + if self.cache_mla_latents: assert ( self.apply_rope_fusion is False diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 5c55f2abe6c..de132b865a9 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -368,6 +368,8 @@ def __init__( attention_optional_kwargs["pg_collection"] = pg_collection if pp_layer_offset is not None: attention_optional_kwargs["pp_layer_offset"] = pp_layer_offset + if is_mtp_layer: + attention_optional_kwargs["is_mtp_layer"] = True # [Module 2: SelfAttention] self.self_attention = build_module( diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index d4f9eb9c0de..812dac1ce24 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -323,6 +323,15 @@ def no_rope_freq_type(x): # it's a single int but in str return int(x) + +def compress_ratios_type(x): + """Parse per-layer compression ratios for compressed sparse attention.""" + if isinstance(x, list): + return x + assert isinstance(x, str) + return _eval_pattern(x) + + def moe_freq_type(x): """Frequency between MoE layers and Dense layers. @@ -2213,6 +2222,7 @@ def _add_network_size_args(parser): "no_rope_freq", "moe_layer_freq", "linear_attention_freq", + "csa_compress_ratios", "moe_router_load_balancing_type", "moe_aux_loss_coeff", "cp_comm_type", @@ -3399,6 +3409,11 @@ def _add_mla_args(parser): help="Mscale for YaRN RoPE in multi-latent attention.") group.add_argument('--mscale-all-dim', type=float, default=0.0, help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") + group.add_argument('--o-groups', type=int, default=8, + help="Number of groups for grouped low-rank output projection (wo_a).") + group.add_argument('--o-lora-rank', type=int, default=1024, + help="Low-rank dimension per group for grouped output (wo_a). " + "Used when o-groups > 0.") group.add_argument('--cache-mla-latents', action='store_true', default=False, help="If set caches the mla down projected latents with mla flash decode.") group.add_argument( @@ -3423,6 +3438,15 @@ def _add_experimental_attention_variant_args(parser): 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.') + group.add_argument( + '--csa-compress-ratios', + type=compress_ratios_type, + default=None, + help='Per-layer compress ratios for compressed sparse attention. ' + 'Accepts a Python list expression such as "[0,0,4,128,4,128]" or ' + '"([0]+[4,128]*2)*3". Valid values are 0, 4, and 128, and the ' + 'list length must be at least num_layers plus mtp_num_layers.', + ) return parser def _add_heterogeneous_args(parser): diff --git a/megatron/training/training.py b/megatron/training/training.py index acdd727b82d..ec1954256a2 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2842,12 +2842,23 @@ def training_log( # Track sparse attention indexer loss. if args.dsa_indexer_loss_coeff is not None and args.dsa_indexer_loss_coeff > 0: indexer_loss_scale = 1 / get_num_microbatches() + assert isinstance( + pg_collection, ProcessGroupCollection + ), "DSA indexer logging requires a ProcessGroupCollection" DSAIndexerLossLoggingHelper.track_indexer_metrics( loss_scale=indexer_loss_scale, iteration=iteration, writer=writer, + pg_collection=pg_collection, wandb_writer=wandb_writer, total_loss_dict=total_loss_dict, + num_layers=args.num_layers + (args.mtp_num_layers or 0), + num_indexer_layers=( + sum(ratio == 4 for ratio in args.csa_compress_ratios) + if args.csa_compress_ratios is not None + else None + ), + preserve_groups=args.cuda_graph_impl != "none", ) # Dump memory snapshot and print metrics to stdout. diff --git a/tests/unit_tests/dist_checkpointing/test_optimizer.py b/tests/unit_tests/dist_checkpointing/test_optimizer.py index f93e09a43b7..7c319e0a14a 100644 --- a/tests/unit_tests/dist_checkpointing/test_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_optimizer.py @@ -79,6 +79,27 @@ def sharded_state_dict(self): return sharded_state_dict +class NativeFp32Model(torch.nn.Module): + """Parameters for an interleaved trainable/frozen BF16 and FP32 group.""" + + def __init__(self): + super().__init__() + self.pre = torch.nn.Linear(8, 8, bias=False) + self.frozen = torch.nn.Linear(8, 8, bias=False) + self.frozen.weight.requires_grad_(False) + self.gate = torch.nn.Parameter(torch.zeros(24, dtype=torch.float32)) + self.post = torch.nn.Linear(8, 8, bias=False) + self.config = TransformerConfig( + hidden_size=8, num_attention_heads=1, num_layers=1, bf16=True + ) + + def sharded_state_dict(self): + return { + key: ShardedTensor.from_rank_offsets(key, value) + for key, value in self.state_dict(keep_vars=True).items() + } + + class SwigluFactoryModel(torch.nn.Module): def __init__(self, pp_separate_model: bool = False): super().__init__() @@ -238,6 +259,65 @@ def test_optimizer_params(self, tmp_path_dist_ckpt): ] ) + def test_float16_optimizer_with_native_fp32_and_frozen_params(self): + """Native FP32 and frozen param ids must not shift BF16 checkpoint state.""" + from megatron.core.optimizer import OptimizerConfig + from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params + from megatron.core.transformer.module import ( + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, + ) + + Utils.initialize_model_parallel(1, 1) + model = NativeFp32Model().cuda() + model.gate = mark_keep_in_fp32(model.gate) + convert_module_to_dtype_except_fp32_marked(model, torch.bfloat16) + assert model.pre.weight.dtype == torch.bfloat16 + assert model.frozen.weight.dtype == torch.bfloat16 + assert not model.frozen.weight.requires_grad + assert model.gate.dtype == torch.float32 + assert model.post.weight.dtype == torch.bfloat16 + + # Use an explicit trainable BF16/frozen BF16/FP32/trainable BF16 order. + # Module.parameters() would yield the root gate before child parameters. + ordered_params = [model.pre.weight, model.frozen.weight, model.gate, model.post.weight] + for param in ordered_params: + if param.requires_grad: + param.grad = torch.zeros_like(param) + inner_optim = Adam(ordered_params) + inner_optim.step() + + optim = Float16OptimizerWithFloat16Params( + inner_optim, + OptimizerConfig(optimizer='adam', lr=1e-4, bf16=True), + None, + lambda opt, cfg: None, + ) + sharded_state_dict = optim.sharded_state_dict(model.sharded_state_dict()) + + # FP32 main copies pair with the BF16 params only, in optimizer order. + fp32_params = sharded_state_dict['fp32_from_fp16_params'][0] + assert [(sharded.key, tuple(sharded.data.shape)) for sharded in fp32_params] == [ + ('optimizer.state.fp32_param.pre.weight', (8, 8)), + ('optimizer.state.fp32_param.post.weight', (8, 8)), + ] + + # The frozen parameter has neither optimizer state nor an fp32 main copy. + state = sharded_state_dict['optimizer']['state'] + assert 1 not in state + + # Per-param state maps every trainable param, including native FP32, to the right key. + expected = {0: ('pre.weight', (8, 8)), 2: ('gate', (24,)), 3: ('post.weight', (8, 8))} + for param_id, (model_key, shape) in expected.items(): + for state_key in ('exp_avg', 'exp_avg_sq'): + sharded = state[param_id][state_key] + assert sharded.key == f'optimizer.state.{state_key}.{model_key}', sharded.key + assert tuple(sharded.data.shape) == shape, ( + param_id, + sharded.key, + sharded.data.shape, + ) + def initialize_pp_agnostic_model(pre_process=True, post_process=True, seed=0, **config_kwargs): torch.manual_seed(seed) diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 810f48092ee..a63edaba6a3 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -16,12 +16,16 @@ try: from megatron.core.fusions.fused_mla_yarn_rope_apply import ( - fused_apply_mla_rope_for_kv, fused_apply_mla_rope_for_q, + fused_mla_rope_inplace, + fused_mla_rope_kv_split, + fused_mla_rope_out_of_place, ) -except: - fused_apply_mla_rope_for_kv = None +except Exception: fused_apply_mla_rope_for_q = None + fused_mla_rope_inplace = None + fused_mla_rope_kv_split = None + fused_mla_rope_out_of_place = None def dtype_tols(dtype): @@ -54,7 +58,9 @@ def test_packed_freqs_returns_offset_mapped_output_for_context_parallel(self): t = torch.randn(4, 2, 8) freqs = torch.randn(8, 1, 1, 8) - out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) expected_freqs = torch.cat([freqs[0:1], freqs[3:4], freqs[4:5], freqs[7:8]], dim=0) expected = rope_utils_module._apply_rotary_pos_emb_bshd( @@ -69,7 +75,9 @@ def test_max_seqlen_freqs_returns_sequence_mapped_output_for_context_parallel(se t = torch.randn(4, 2, 8) freqs = torch.randn(4, 1, 1, 8) - out = rope_utils_module._apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, cp_group=cp_group) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) expected_freqs = torch.cat([freqs[1:2], freqs[2:3]], dim=0) expected_slices = [] @@ -83,9 +91,64 @@ def test_max_seqlen_freqs_returns_sequence_mapped_output_for_context_parallel(se torch.testing.assert_close(out, expected) + def test_missing_max_seqlen_preserves_legacy_packed_freq_mapping(self): + cp_group = FakeCPGroup(size=2, rank=0) + cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.int32) + t = torch.randn(4, 2, 8) + freqs = torch.randn(8, 1, 1, 8) -def _test_fused_apply_mla_rope_for_q(input_format): - assert fused_apply_mla_rope_for_q is not None + legacy_out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group + ) + explicit_out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens, freqs, cp_group=cp_group, max_seqlen=4 + ) + + torch.testing.assert_close(legacy_out, explicit_out) + + def test_shared_max_seqlen_maps_asymmetric_query_sequences_from_zero(self): + cp_group = FakeCPGroup(size=1, rank=0) + cu_seqlens_q = torch.tensor([0, 3, 6], dtype=torch.int32) + t = torch.randn(6, 2, 8) + freqs = torch.randn(4, 1, 1, 8) + + max_seqlen_q = 3 + max_seqlen_kv = freqs.size(0) + assert max_seqlen_q < max_seqlen_kv < t.size(0) + combined_max_seqlen = max(max_seqlen_q, max_seqlen_kv) + out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens_q, freqs, cp_group=cp_group, max_seqlen=combined_max_seqlen + ) + compatibility_out = rope_utils_module._apply_rotary_pos_emb_thd( + t, cu_seqlens_q, freqs, cp_group=cp_group + ) + + expected_freqs = freqs[torch.tensor([0, 1, 2, 0, 1, 2])] + expected = rope_utils_module._apply_rotary_pos_emb_bshd( + t.unsqueeze(1), expected_freqs + ).squeeze(1) + + torch.testing.assert_close(out, expected) + torch.testing.assert_close(out, compatibility_out) + + +class _SaveOutputForBackward(torch.autograd.Function): + """Minimal stand-in for a kernel whose backward consumes its output.""" + + @staticmethod + def forward(ctx, tensor): + output = tensor.clone() + ctx.save_for_backward(output) + return output + + @staticmethod + def backward(ctx, _grad_output): + (saved_output,) = ctx.saved_tensors + return saved_output + + +def _test_fused_mla_rope_inplace(input_format, inverse=False, remove_interleaving=False): + assert fused_mla_rope_inplace is not None num_heads = 32 q_dim = 128 emb_dim = 64 @@ -97,6 +160,7 @@ def _test_fused_apply_mla_rope_for_q(input_format): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -142,15 +206,25 @@ def _test_fused_apply_mla_rope_for_q(input_format): freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, + inverse=inverse, + mla_output_remove_interleaving=remove_interleaving, ) pytorch_output = torch.concat([no_pe, pe_output], dim=-1) pytorch_output.backward(pytorch_bwd_input, retain_graph=True) - fused_output = fused_apply_mla_rope_for_q( - fused_fwd_input, cos, sin, q_dim, emb_dim, cu_seqlens_q=cu_seqlens + fused_output = fused_mla_rope_inplace( + fused_fwd_input, + cos, + sin, + q_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=inverse, + remove_interleaving=remove_interleaving, ) fused_output.backward(fused_bwd_input, retain_graph=True) @@ -169,8 +243,8 @@ def _test_fused_apply_mla_rope_for_q(input_format): ) -def _test_fused_apply_mla_rope_for_kv(input_format): - assert fused_apply_mla_rope_for_kv is not None +def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False): + assert fused_mla_rope_kv_split is not None num_heads = 32 k_dim = 128 v_dim = 128 @@ -183,6 +257,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): multi_latent_attention=True, ) + max_seqlen = None if input_format == "sbhd": cu_seqlens = None seqlen = 1024 @@ -241,9 +316,11 @@ def _test_fused_apply_mla_rope_for_kv(input_format): freqs, transformer_config, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, + mla_output_remove_interleaving=remove_interleaving, ) if input_format == "sbhd": pe_output = pe_output.expand(-1, -1, num_heads, -1) @@ -255,7 +332,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): (pytorch_k_output, pytorch_v_output), (pytorch_bwd_k_input, pytorch_bwd_v_input) ) - fused_k_output, fused_v_output = fused_apply_mla_rope_for_kv( + fused_k_output, fused_v_output = fused_mla_rope_kv_split( fused_fwd_kv_input, fused_fwd_emb_input, cos, @@ -264,6 +341,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): k_dim, v_dim, cu_seqlens_kv=cu_seqlens, + remove_interleaving=remove_interleaving, ) torch.autograd.backward( (fused_k_output, fused_v_output), (fused_bwd_k_input, fused_bwd_v_input) @@ -301,13 +379,136 @@ def _test_fused_apply_mla_rope_for_kv(input_format): @pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("input_format", ["sbhd", "thd"]) -class TestFusedApplyMLARope: +class TestFusedMLARope: @pytest.mark.flaky_in_dev - def test_forward_backward_for_q(self, input_format): - _test_fused_apply_mla_rope_for_q(input_format) + @pytest.mark.parametrize("inverse", [False, True]) + @pytest.mark.parametrize("remove_interleaving", [False, True]) + def test_inplace_forward_backward(self, input_format, inverse, remove_interleaving): + _test_fused_mla_rope_inplace( + input_format, inverse=inverse, remove_interleaving=remove_interleaving + ) + + @pytest.mark.parametrize("remove_interleaving", [False, True]) + def test_kv_split_forward_backward(self, input_format, remove_interleaving): + _test_fused_mla_rope_kv_split(input_format, remove_interleaving=remove_interleaving) + + +@pytest.mark.experimental +@pytest.mark.internal +@pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("input_format", ["sbhd", "thd"]) +def test_out_of_place_inverse_rope_preserves_upstream_saved_output(input_format): + """Post-attention inverse RoPE must not overwrite an output saved for backward.""" + assert fused_mla_rope_out_of_place is not None + seqlen = 32 + batch_size = 1 + num_heads = 2 + nope_dim = 16 + emb_dim = 64 + dtype = torch.bfloat16 + + yarn_rope = YarnRotaryEmbedding(emb_dim, original_max_position_embeddings=seqlen) + freqs, mscale = yarn_rope(seqlen, 0) + cos = (torch.cos(freqs) * mscale).to(dtype) + sin = (torch.sin(freqs) * mscale).to(dtype) + + if input_format == "sbhd": + shape = (seqlen, batch_size, num_heads, nope_dim + emb_dim) + cu_seqlens = None + else: + shape = (2 * seqlen, num_heads, nope_dim + emb_dim) + cu_seqlens = torch.tensor([0, seqlen, 2 * seqlen], dtype=torch.int32, device="cuda") + + unsafe_source = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + unsafe_attention_output = _SaveOutputForBackward.apply(unsafe_source) + unsafe_reference = unsafe_attention_output.detach().clone() + unsafe_inverse_output = fused_mla_rope_inplace( + unsafe_attention_output, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + + assert unsafe_inverse_output.data_ptr() == unsafe_attention_output.data_ptr() + assert not torch.equal(unsafe_attention_output, unsafe_reference) + + source = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + attention_output = _SaveOutputForBackward.apply(source) + saved_reference = attention_output.detach().clone() + + inverse_output = fused_mla_rope_out_of_place( + attention_output, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + expected_inverse_output = fused_mla_rope_inplace( + saved_reference.clone(), + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=True, + remove_interleaving=True, + ) + + assert inverse_output.data_ptr() != attention_output.data_ptr() + torch.testing.assert_close(attention_output, saved_reference, rtol=0, atol=0) + torch.testing.assert_close(inverse_output, expected_inverse_output, rtol=0, atol=0) + + inverse_output.backward(torch.randn_like(inverse_output).contiguous()) + torch.testing.assert_close(source.grad, saved_reference, rtol=0, atol=0) + + +@pytest.mark.experimental +@pytest.mark.internal +@pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("input_format", ["sbhd", "thd"]) +def test_legacy_query_api_remains_in_place(input_format): + """The legacy API keeps its original mutation behavior and allocation profile.""" + assert fused_apply_mla_rope_for_q is not None + seqlen = 32 + batch_size = 1 + num_heads = 2 + nope_dim = 16 + emb_dim = 64 + dtype = torch.bfloat16 + + yarn_rope = YarnRotaryEmbedding(emb_dim, original_max_position_embeddings=seqlen) + freqs, mscale = yarn_rope(seqlen, 0) + cos = (torch.cos(freqs) * mscale).to(dtype) + sin = (torch.sin(freqs) * mscale).to(dtype) + + if input_format == "sbhd": + shape = (seqlen, batch_size, num_heads, nope_dim + emb_dim) + cu_seqlens = None + else: + shape = (2 * seqlen, num_heads, nope_dim + emb_dim) + cu_seqlens = torch.tensor([0, seqlen, 2 * seqlen], dtype=torch.int32, device="cuda") + + query = torch.randn(shape, dtype=dtype, device="cuda") + reference = query.clone() + expected = fused_mla_rope_inplace( + reference.clone(), cos, sin, nope_dim, emb_dim, cu_seqlens_q=cu_seqlens + ) + output = fused_apply_mla_rope_for_q( + query, cos, sin, qk_head_dim=nope_dim, emb_dim=emb_dim, cu_seqlens_q=cu_seqlens + ) - def test_forward_backward_for_kv(self, input_format): - _test_fused_apply_mla_rope_for_kv(input_format) + assert output.data_ptr() == query.data_ptr() + assert not torch.equal(query, reference) + torch.testing.assert_close(output, expected, rtol=0, atol=0) class TestApplyRotaryPosEmbMlaFusionConflict: diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index eb871568046..c5a1b426d01 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -68,6 +68,10 @@ "cpu_offloading_weights": False, "cross_entropy_fusion_impl": "native", "cross_entropy_loss_fusion": True, + "csa_compress_ratios": None, + "csa_compress_rotary_base": 40000.0, + "csa_dense_mode": False, + "csa_window_size": 128, "cuda_graph_impl": "none", "cuda_graph_retain_backward_graph": False, "cuda_graph_modules": [], diff --git a/tests/unit_tests/pipeline_parallel/test_schedules.py b/tests/unit_tests/pipeline_parallel/test_schedules.py index 92db675d193..3d50763ea61 100644 --- a/tests/unit_tests/pipeline_parallel/test_schedules.py +++ b/tests/unit_tests/pipeline_parallel/test_schedules.py @@ -377,7 +377,8 @@ def test_dsa_indexer_loss_scale_accepts_dict_output_tensor(): ) -def test_dsa_indexer_loss_scale_defaults_from_variant_without_mutating_config(): +@pytest.mark.parametrize("variant", ["dsa", "dsv4_hybrid"]) +def test_indexer_loss_scale_defaults_from_variant_without_mutating_config(variant): from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, ) @@ -385,7 +386,7 @@ def test_dsa_indexer_loss_scale_defaults_from_variant_without_mutating_config(): config = SimpleNamespace( calculate_per_token_loss=True, experimental_attention_variant_loss_scale_func=None, - experimental_attention_variant='dsa', + experimental_attention_variant=variant, grad_scale_func=lambda tensor: tensor * 7.0, num_moe_experts=None, mtp_num_layers=None, diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index dc65d541455..2fc55962ae8 100644 --- a/tests/unit_tests/test_fp8_utils.py +++ b/tests/unit_tests/test_fp8_utils.py @@ -23,6 +23,31 @@ reason_for_no_mxfp8 = "MXFP8 requires Transformer Engine and device arch >= 10" +@pytest.mark.skipif(not fp8_utils.HAVE_TE, reason="Transformer Engine is not installed") +@pytest.mark.parametrize( + ("is_init", "config_values", "te_helper"), + [ + ( + False, + {"fp8": "hybrid", "fp4": None, "fp8_param": False, "fp4_param": False}, + "fp8_autocast", + ), + (True, {"fp8": None, "fp4": None, "fp8_param": True, "fp4_param": False}, "fp8_model_init"), + ], +) +def test_get_fp8_disabled_context_uses_disabled_te_context(is_init, config_values, te_helper): + config = Mock(**config_values) + disabled_context = Mock() + + with patch.object( + fp8_utils.transformer_engine.pytorch, te_helper, return_value=disabled_context + ) as te_context: + result = fp8_utils.get_fp8_disabled_context(config, is_init=is_init) + + assert result is disabled_context + te_context.assert_called_once_with(enabled=False) + + class MockTELinear(nn.Module): """Mock TE Linear module for testing.""" diff --git a/tests/unit_tests/test_optimizer_cpu_offloading.py b/tests/unit_tests/test_optimizer_cpu_offloading.py index 33febbb3eb0..379acc9dbda 100644 --- a/tests/unit_tests/test_optimizer_cpu_offloading.py +++ b/tests/unit_tests/test_optimizer_cpu_offloading.py @@ -17,6 +17,20 @@ from torch.optim import Adam as GPUAdam from megatron.core.optimizer.cpu_offloading import HybridDeviceOptimizer +from megatron.core.transformer.module import ( + convert_module_to_dtype_except_fp32_marked, + mark_keep_in_fp32, +) + + +class Fp32MarkedToyNet(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4, bias=False) + self.scale = mark_keep_in_fp32(nn.Parameter(torch.ones(4))) + + def forward(self, x): + return self.proj(x) * self.scale class Net(nn.Module): @@ -71,6 +85,52 @@ def setup_seed(seed): torch.backends.cudnn.benchmark = False # Disable auto-tuner for reproducibility +def test_load_state_dict_with_native_fp32_param(): + """Round-trip state for a BF16 toy net with a parameter marked to stay in FP32.""" + model = Fp32MarkedToyNet().cuda() + convert_module_to_dtype_except_fp32_marked(model, torch.bfloat16) + assert model.proj.weight.dtype == torch.bfloat16 + assert model.scale.dtype == torch.float32 + + optimizer = HybridDeviceOptimizer( + model.parameters(), + offload_fraction=1.0, + cpu_optimizer_cls=Adam, + gpu_optimizer_cls=GPUAdam, + param_update_in_fp32=True, + overlap_cpu_optimizer_d2h_h2d=False, + lr=1e-3, + ) + inputs = torch.ones(2, 4, device="cuda", dtype=torch.bfloat16) + model(inputs).sum().backward() + optimizer.step() + + restored_model = Fp32MarkedToyNet().cuda() + convert_module_to_dtype_except_fp32_marked(restored_model, torch.bfloat16) + restored_model.load_state_dict(model.state_dict()) + restored_optimizer = HybridDeviceOptimizer( + restored_model.parameters(), + offload_fraction=1.0, + cpu_optimizer_cls=Adam, + gpu_optimizer_cls=GPUAdam, + param_update_in_fp32=True, + overlap_cpu_optimizer_d2h_h2d=False, + lr=1e-3, + ) + restored_optimizer.load_state_dict(optimizer.state_dict()) + + assert set(restored_optimizer.state) == set(restored_model.parameters()) + assert restored_model.proj.weight in restored_optimizer.param_to_fp32_param + assert restored_model.scale not in restored_optimizer.param_to_fp32_param + assert torch.equal( + restored_optimizer.param_to_fp32_param[restored_model.proj.weight], + optimizer.param_to_fp32_param[model.proj.weight], + ) + + restored_model(inputs).sum().backward() + restored_optimizer.step() + + @pytest.mark.skipif( torch.__version__ < '2.3.0', reason=( diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py new file mode 100644 index 00000000000..e7f68369431 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -0,0 +1,1418 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, + _apply_rope, + _compute_unfused_csa_non_compressed_lse, + get_compress_topk_idxs, + get_window_topk_idxs, + unfused_compressed_sparse_attn, +) +from megatron.core.transformer.experimental_attention_variant.dsa import ( + FusedDSAIndexerLoss, + compute_dsa_indexer_loss, + fused_qk_topk_naive, +) +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + + +def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + """Mock implementation of hadamard_transform for testing without the library installed.""" + return x * scale + + +class _DisabledContextTracker: + """Track whether a projection runs inside the FP8-disabled context.""" + + def __init__(self): + self.depth = 0 + self.entries = 0 + + def __call__(self, _config, is_init=False): + assert not is_init + return self + + def __enter__(self): + self.depth += 1 + self.entries += 1 + return self + + def __exit__(self, _exc_type, _exc_value, _traceback): + self.depth -= 1 + return False + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Automatically patch hadamard_transform in both dsa and csa modules if not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# =========================================================================== +# Helper function tests +# =========================================================================== + + +class _SingleRankTP: + @staticmethod + def size(): + return 1 + + +class _SingleRankPG: + tp = _SingleRankTP() + + +def test_unfused_csa_non_compressed_lse_matches_window_and_sink_oracle(): + torch.manual_seed(17) + seqlen_q, batch_size, n_kv = 3, 2, 5 + num_heads, head_dim = 2, 4 + query = torch.randn(seqlen_q, batch_size, num_heads, head_dim, requires_grad=True) + kv_full = torch.randn(n_kv, batch_size, head_dim, requires_grad=True) + sink = torch.randn(num_heads, requires_grad=True) + window_indices = torch.tensor([[[-1, 0], [0, 1], [1, 2]], [[-1, 0], [0, 2], [2, 4]]]) + + expected = torch.empty(batch_size, num_heads, seqlen_q) + with torch.no_grad(): + for batch in range(batch_size): + for row in range(seqlen_q): + for head in range(num_heads): + logits = [sink[head]] + for key_index in window_indices[batch, row]: + if key_index >= 0: + logits.append( + torch.dot(query[row, batch, head], kv_full[key_index, batch]) + ) + expected[batch, head, row] = torch.logsumexp(torch.stack(logits), dim=0) + + actual = _compute_unfused_csa_non_compressed_lse( + query, kv_full, sink, window_indices, softmax_scale=1.0 + ) + + assert actual.shape == (batch_size, num_heads, seqlen_q) + assert actual.dtype == torch.float32 + assert not actual.requires_grad + torch.testing.assert_close(actual, expected) + for teacher_tensor in (query, kv_full, sink): + assert teacher_tensor.grad is None + + +def _independent_csa_indexer_loss( + index_scores, + topk_indices, + query, + compressed_kv, + window_kv, + window_indices, + sink, + *, + sparse_loss, + loss_coeff, +): + """Compute a small-loop CSA teacher oracle with the complete denominator.""" + batch_size, seqlen_q, n_compressed = index_scores.shape + num_heads = query.shape[2] + losses = [] + for batch in range(batch_size): + for row in range(seqlen_q): + selected = ( + topk_indices[batch, row].tolist() if sparse_loss else list(range(n_compressed)) + ) + target = [] + for compressed_index in selected: + head_mass = 0.0 + for head in range(num_heads): + non_compressed_logits = [sink[head]] + for window_index in window_indices[batch, row]: + if window_index >= 0: + non_compressed_logits.append( + torch.dot(query[row, batch, head], window_kv[window_index, batch]) + ) + compressed_logits = [ + torch.dot(query[row, batch, head], compressed_kv[key_index, batch]) + for key_index in selected + ] + denominator = torch.logsumexp( + torch.stack(non_compressed_logits + compressed_logits), dim=0 + ) + selected_position = selected.index(compressed_index) + head_mass = head_mass + torch.exp( + compressed_logits[selected_position] - denominator + ) + target.append(head_mass) + target = torch.stack(target) + target = target / target.sum() + predict_log = torch.log_softmax(index_scores[batch, row, selected], dim=-1) + losses.append((target * (torch.log(target) - predict_log)).sum()) + return torch.stack(losses).mean() * loss_coeff + + +@pytest.mark.parametrize("sparse_loss", [False, True], ids=["dense", "sparse"]) +def test_csa_indexer_loss_uses_full_attention_denominator(sparse_loss): + torch.manual_seed(29) + seqlen_q, batch_size, num_heads, head_dim = 4, 1, 2, 3 + n_compressed, index_heads, index_dim = 3, 2, 2 + index_topk, loss_coeff = 2, 0.7 + + q = torch.randn(seqlen_q, batch_size, index_heads, index_dim, requires_grad=True) + weights = torch.randn(seqlen_q, batch_size, index_heads, requires_grad=True) + k = torch.randn(n_compressed, batch_size, index_dim, requires_grad=True) + query = torch.randn(seqlen_q, batch_size, num_heads, head_dim, requires_grad=True) + window_kv = torch.randn(seqlen_q, batch_size, head_dim, requires_grad=True) + compressed_kv = torch.randn(n_compressed, batch_size, head_dim, requires_grad=True) + sink = torch.randn(num_heads, requires_grad=True) + window_indices = torch.tensor([[[-1, 0], [0, 1], [1, 2], [2, 3]]]) + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, window_kv, sink, window_indices, softmax_scale=1.0 + ) + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, num_heads, -1) + compressed_mask = torch.zeros(seqlen_q, n_compressed) + + q_reference = q.detach().clone().requires_grad_(True) + weights_reference = weights.detach().clone().requires_grad_(True) + k_reference = k.detach().clone().requires_grad_(True) + index_scores_reference, topk_reference = fused_qk_topk_naive( + q_reference, k_reference, weights_reference, index_topk + ) + loss_reference = compute_dsa_indexer_loss( + index_scores_reference, + topk_reference, + query.detach(), + key_for_loss.detach(), + 1.0, + loss_coeff, + sparse_loss, + _SingleRankPG(), + mask=compressed_mask, + non_compressed_lse=non_compressed_lse, + ) + loss_reference.backward() + + topk_actual, loss_actual = FusedDSAIndexerLoss.apply( + q, + weights, + k, + query, + key_for_loss, + 1.0, + index_topk, + loss_coeff, + compressed_mask, + sparse_loss, + _SingleRankPG(), + None, + None, + None, + None, + False, + True, + non_compressed_lse, + ) + loss_actual.backward() + + independent_loss = _independent_csa_indexer_loss( + index_scores_reference.detach(), + topk_reference, + query.detach(), + compressed_kv.detach(), + window_kv.detach(), + window_indices, + sink.detach(), + sparse_loss=sparse_loss, + loss_coeff=loss_coeff, + ) + + torch.testing.assert_close(loss_actual, independent_loss) + torch.testing.assert_close(loss_actual, loss_reference) + torch.testing.assert_close(topk_actual, topk_reference) + torch.testing.assert_close(q.grad, q_reference.grad) + torch.testing.assert_close(weights.grad, weights_reference.grad) + torch.testing.assert_close(k.grad, k_reference.grad) + for teacher_tensor in (query, window_kv, compressed_kv, sink): + assert teacher_tensor.grad is None + + +class TestGetWindowTopkIdxs: + """Test get_window_topk_idxs helper.""" + + def test_basic_shape(self): + batch_size, seqlen, window_size = 2, 16, 4 + idxs = get_window_topk_idxs(window_size, batch_size, seqlen, torch.device("cpu")) + assert idxs.shape == (batch_size, seqlen, window_size) + + def test_causal_no_future(self): + """Indices should never exceed the query position.""" + seqlen, window_size = 32, 8 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + for i in range(seqlen): + valid = idxs[0, i][idxs[0, i] >= 0] + assert torch.all(valid <= i), f"Position {i} has future indices" + + def test_invalid_marked_minus_one(self): + """Early positions that cannot fill the window should use -1.""" + seqlen, window_size = 8, 4 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs[0, 0, 0] == -1 or idxs[0, 0, 0] == 0 + for pos in range(window_size, seqlen): + assert torch.all(idxs[0, pos] >= 0), f"Position {pos} has invalid -1" + + def test_window_larger_than_seqlen(self): + """Window larger than sequence length should still work.""" + seqlen, window_size = 4, 16 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs.shape == (1, seqlen, window_size) + + +class TestGetCompressTopkIdxs: + """Test get_compress_topk_idxs helper.""" + + def test_basic_shape(self): + ratio, batch_size, seqlen, offset = 4, 2, 32, 32 + idxs = get_compress_topk_idxs(ratio, batch_size, seqlen, offset, torch.device("cpu")) + n_compressed = seqlen // ratio + assert idxs.shape == (batch_size, seqlen, n_compressed) + + def test_offset_applied(self): + """Valid indices should be >= offset.""" + ratio, seqlen, offset = 4, 32, 100 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + valid = idxs[idxs >= 0] + if valid.numel() > 0: + assert torch.all(valid >= offset), "Valid indices should be offset" + + def test_causal_no_future(self): + """Compressed indices should respect causality.""" + ratio, seqlen, offset = 4, 32, 32 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + for i in range(seqlen): + n_valid = (i + 1) // ratio + valid = idxs[0, i][idxs[0, i] >= 0] + assert valid.numel() <= n_valid, f"Position {i} has too many valid compressed indices" + + def test_ratio_128(self): + """Test with large compression ratio.""" + ratio, seqlen, offset = 128, 256, 256 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + assert idxs.shape == (1, seqlen, seqlen // ratio) + + +# =========================================================================== +# unfused_compressed_sparse_attn tests +# =========================================================================== + + +class TestUnfusedCompressedSparseAttn: + """Test the unfused compressed sparse attention kernel.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_output_shape(self): + """Test output shape of unfused compressed sparse attention.""" + sq, b, np_, hn = 16, 2, 4, 64 + n_kv = sq + sq // 4 + topk = 8 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + + assert output.shape == (sq, b, np_ * hn) + assert output.dtype == query.dtype + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_invalid_indices_masked(self): + """Test that -1 indices are properly masked.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + + topk_indices = torch.full((b, sq, topk), -1, dtype=torch.int32).cuda() + topk_indices[:, :, 0] = 0 + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + assert not torch.isnan(output).any(), "Output should not contain NaN" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_gradient_flow(self): + """Test that gradients flow through sparse attention.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) + kv_full = torch.randn(n_kv, b, hn, dtype=torch.float32).cuda().requires_grad_(True) + attn_sink = torch.nn.Parameter(torch.zeros(np_, dtype=torch.float32).cuda()) + + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert kv_full.grad is not None + assert attn_sink.grad is not None + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +def _make_mla_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + csa_compress_ratios=None, + csa_window_size=8, + csa_dense_mode=False, + tensor_model_parallel_size=1, + sequence_parallel=False, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + dsa_indexer_use_sparse_loss=False, +): + """Helper to create MLATransformerConfig for CSA tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [0] * num_layers + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + csa_dense_mode=csa_dense_mode, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + + +def _make_compressor_submodules(): + """Create Compressor submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressorSubmodules( + linear_wkv=ModuleSpec(module=TELinear), + linear_wgate=ModuleSpec(module=TELinear), + norm=ModuleSpec(module=TENorm), + ) + + +def _make_csa_indexer_submodules(): + """Create CSAIndexer submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_weights_proj=ModuleSpec(module=TELinear), + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + ) + + +def _make_csa_submodules(): + """Create CompressedSparseAttention submodules spec.""" + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressedSparseAttentionSubmodules( + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + indexer=ModuleSpec(module=CSAIndexer, submodules=_make_csa_indexer_submodules()), + ) + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressor: + """Test Compressor module.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[4, 128, 4, 128]) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_output_shape(self, compress_ratio): + """Test that compressor produces correct output shape.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + + expected_len = seq_len // compress_ratio + assert output is not None + assert output.shape == (expected_len, batch_size, head_dim) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_too_short_input(self, compress_ratio): + """Test that compressor returns None when input is shorter than compress_ratio.""" + short_len = compress_ratio - 1 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(short_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + assert output is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_gradient_flow(self, compress_ratio): + """Test that gradients flow through the compressor.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + output = compressor(x) + loss = output.sum() + loss.backward() + + assert x.grad is not None + for name, param in compressor.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_projection_disables_fp8(self, compress_ratio, monkeypatch): + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=self.config.v_head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + tracker = _DisabledContextTracker() + calls = [] + + for name, projection in ( + ('linear_wkv', compressor.linear_wkv), + ('linear_wgate', compressor.linear_wgate), + ): + original_forward = projection.forward + + def checked_forward(*args, _name=name, _forward=original_forward, **kwargs): + assert tracker.depth > 0, f"{_name} ran outside the FP8-disabled context" + calls.append(_name) + return _forward(*args, **kwargs) + + monkeypatch.setattr(projection, 'forward', checked_forward) + + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + x = torch.randn( + compress_ratio * 2, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + compressor(x) + + assert calls == ['linear_wkv', 'linear_wgate'] + assert tracker.entries == 1 + + +# =========================================================================== +# CSAIndexer tests +# =========================================================================== + + +@pytest.mark.parametrize("seqlen", [32, 128]) +class TestCSAIndexer: + """Test CSAIndexer module basic functionality.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.compress_ratio = 4 + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4], dsa_indexer_topk=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.indexer = CSAIndexer( + config=cls.config, + submodules=_make_csa_indexer_submodules(), + compress_ratio=cls.compress_ratio, + rotary_pos_emb=cls.rotary_pos_emb, + pg_collection=cls.pg_collection, + ) + + yield + Utils.destroy_model_parallel() + + def test_csa_indexer_constructor(self, seqlen): + """Test CSAIndexer initialization.""" + assert isinstance(self.indexer, CSAIndexer) + assert self.indexer.compress_ratio == self.compress_ratio + assert self.indexer.index_n_heads == self.config.dsa_indexer_n_heads + assert self.indexer.index_head_dim == self.config.dsa_indexer_head_dim + assert self.indexer.index_topk == self.config.dsa_indexer_topk + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward(self, seqlen): + """Test CSAIndexer forward pass.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + index_scores, topk_indices = self.indexer(x, qr) + n_compressed = seqlen // self.compress_ratio + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + assert index_scores.dtype == torch.float32 + assert topk_indices.dtype == torch.long + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward_before_topk(self, seqlen): + """Test CSAIndexer forward_before_topk.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + q, k, weights = self.indexer.forward_before_topk(x, qr) + + assert q.shape == ( + seqlen, + batch_size, + self.config.dsa_indexer_n_heads, + self.config.dsa_indexer_head_dim, + ) + n_compressed = seqlen // self.compress_ratio + assert k.shape == (n_compressed, batch_size, self.config.dsa_indexer_head_dim) + assert weights.shape == (seqlen, batch_size, self.config.dsa_indexer_n_heads) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_weights_projection_disables_fp8(self, seqlen, monkeypatch): + tracker = _DisabledContextTracker() + self.indexer.cuda() + original_forward = self.indexer.linear_weights_proj.forward + + def checked_forward(*args, **kwargs): + assert tracker.depth > 0, "indexer weights projection ran under FP8" + return original_forward(*args, **kwargs) + + monkeypatch.setattr(self.indexer.linear_weights_proj, 'forward', checked_forward) + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + x = torch.randn(seqlen, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + weights = self.indexer._project_weights(x) + + assert weights.shape == (seqlen, 1, self.config.dsa_indexer_n_heads) + assert tracker.entries == 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_with_mask(self, seqlen): + """Test CSAIndexer with causal mask.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + n_compressed = seqlen // self.compress_ratio + causal_mask = torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(seqlen, -1) + positions = torch.arange(1, seqlen + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(batch_size, -1, -1) + ) + + index_scores, topk_indices = self.indexer(x, qr, mask=causal_mask) + + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + + +# =========================================================================== +# CompressedSparseAttention tests +# =========================================================================== + + +class TestCompressedSparseAttentionRatio1: + """Test CompressedSparseAttention with compress_ratio=1 (window-only).""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[0, 0, 0, 0], csa_window_size=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.csa = CompressedSparseAttention( + config=cls.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=cls.pg_collection, + rotary_pos_emb=rotary_pos_emb, + compress_ratio=0, + ) + + yield + Utils.destroy_model_parallel() + + def test_ratio1_no_compressor(self): + """With ratio=1, compressor and indexer should not be built.""" + assert self.csa.compressor is None + assert self.csa.indexer is None + + def test_mtp_layer_number_is_offset(self): + """MTP attention layers are numbered after all decoder layers.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + compress_ratio=0, + is_mtp_layer=True, + ) + + assert csa.layer_number == self.config.num_layers + 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_forward(self): + """Test forward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_backward(self): + """Test backward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.train() + self.csa.cuda() + + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressedSparseAttentionCompressed: + """Test CompressedSparseAttention with compress_ratio > 1.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], + csa_window_size=8, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=1.0, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + def _get_layer_number(self, compress_ratio): + """Return a layer_number (1-indexed) whose compress_ratio matches.""" + for i, r in enumerate(self.config.csa_compress_ratios): + if r == compress_ratio: + return i + 1 + raise ValueError(f"No layer with compress_ratio={compress_ratio}") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_constructor(self, compress_ratio): + """Test that compressor/indexer are conditionally built.""" + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + assert csa.compressor is not None + if compress_ratio == 4: + assert csa.indexer is not None + elif compress_ratio == 128: + assert csa.indexer is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_forward(self, compress_ratio): + """Test forward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_backward(self, compress_ratio): + """Test backward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + csa.train() + + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + for name, param in csa.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_eval_mode(self, compress_ratio): + """Test forward pass in eval mode.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + csa.eval() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + with torch.no_grad(): + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + +# =========================================================================== +# _apply_rope tests +# =========================================================================== + + +class TestApplyRope: + """Test ``_apply_rope`` — the layout-aware RoPE wrapper used by + Compressor / CSAIndexer / hybrid-attention callers. + + Behaviours covered: + + * 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs both work (3-D gets a temporary head-dim unsqueeze). + * Only the trailing ``pos_dim`` components are rotated; the leading + ``nope_dim`` slice is bit-exact unchanged. + * Both ``RotaryEmbedding`` (returns ``Tensor``) and + ``YarnRotaryEmbedding`` (returns ``(emb, mscale)`` tuple) — DSv4 + hybrid silently swaps the class based on ``compress_ratio``. + * Both unfused and fused (``config.apply_rope_fusion=True``) paths + produce the same output (within bf16 precision). + * For ``ratio > 1`` the rotary table is built at + ``rotary_seq_len * ratio`` and strided by ``ratio``. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(0) + model_parallel_cuda_manual_seed(0) + cls = request.cls + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + # head_dim 32 = nope 24 + pos 8 + cls.config = _make_mla_config(v_head_dim=32, qk_pos_emb_head_dim=8) + yield + Utils.destroy_model_parallel() + + def _make_rotary(self, kind: str): + from megatron.core.models.common.embeddings import RotaryEmbedding, YarnRotaryEmbedding + + pos_dim = self.config.qk_pos_emb_head_dim + if kind == 'rope': + return RotaryEmbedding( + pos_dim, rotary_percent=1.0, rotary_base=10000, cp_group=self.pg_collection.cp + ) + if kind == 'yarn': + return YarnRotaryEmbedding( + pos_dim, + rotary_base=40000, + scaling_factor=40, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1.0, + mscale_all_dim=0.0, + cp_group=self.pg_collection.cp, + ) + raise ValueError(kind) + + def _config_with(self, *, apply_rope_fusion: bool): + # Reuse the class-level config; only flip the fusion flag. + cfg = self.config + cfg.apply_rope_fusion = apply_rope_fusion + return cfg + + _ROTARY_FUSION_COMBOS = [ + pytest.param('rope', False, id='rope-unfused'), + pytest.param('rope', True, id='rope-fused'), + pytest.param('yarn', False, id='yarn-unfused'), + pytest.param('yarn', True, id='yarn-fused'), + ] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize(("rotary_kind", "apply_rope_fusion"), _ROTARY_FUSION_COMBOS) + @pytest.mark.parametrize("input_ndim", [3, 4], ids=['3d', '4d']) + @pytest.mark.parametrize("ratio", [1, 4], ids=['ratio_1', 'ratio_4']) + def test_apply_rope(self, rotary_kind, apply_rope_fusion, input_ndim, ratio): + """Output shape == input shape; no NaN; nope-dim slice is + bit-exact unchanged. Sweeps the valid combinations of rotary + class × apply_rope_fusion × input rank × ratio. Yarn's + tuple-return is covered by the ``'yarn-*'`` combos. + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch, heads = 8, 2, 4 + cfg = self._config_with(apply_rope_fusion=apply_rope_fusion) + + shape = (seq, batch, head_dim) if input_ndim == 3 else (seq, batch, heads, head_dim) + x = torch.randn(*shape, dtype=torch.bfloat16, device='cuda') + # ``fused_mla_rope_inplace`` mutates the input — give it a copy so + # the nope-dim equality check below still has the original. + out = _apply_rope( + x.clone(), + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + assert out.shape == x.shape + assert out.dtype == x.dtype + assert not torch.isnan(out).any() + # The leading nope_dim slice is the identity portion of RoPE. + assert torch.equal( + out[..., :nope], x[..., :nope] + ), "RoPE must not touch the first nope_dim components" + # Trailing pos_dim should rotate at non-zero positions. + pe_changed = (out[..., nope:] != x[..., nope:]).any(dim=-1).flatten() + assert pe_changed[ + 1: + ].any(), "RoPE should rotate the trailing pos_dim components for seq > 0" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("rotary_kind", ['rope', 'yarn']) + def test_3d_input_matches_4d_with_single_head(self, rotary_kind): + """For a single-head input, the 3-D ``(s, b, d)`` and 4-D + ``(s, b, 1, d)`` invocations must produce numerically identical + output (3-D path just inserts a temporary head dim). + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch = 8, 2 + cfg = self._config_with(apply_rope_fusion=False) + + x_3d = torch.randn(seq, batch, head_dim, dtype=torch.bfloat16, device='cuda') + x_4d = x_3d.unsqueeze(-2) + + out_3d = _apply_rope( + x_3d, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + out_4d = _apply_rope( + x_4d, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + + assert out_3d.shape == x_3d.shape + assert out_4d.shape == x_4d.shape + assert torch.equal(out_3d, out_4d.squeeze(-2)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("rotary_kind", ['rope', 'yarn']) + def test_ratio_strides_rotary_table(self, rotary_kind): + """For ``ratio > 1``, the rotary table is built at + ``rotary_seq_len * ratio`` and strided by ``ratio``. The result + with ``ratio=k`` must equal an ``apply_rope`` call on the same + positions of a length-``rotary_seq_len * k`` table. + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch, heads, ratio = 4, 1, 2, 4 + cfg = self._config_with(apply_rope_fusion=False) + + x_comp = torch.randn(seq, batch, heads, head_dim, dtype=torch.bfloat16, device='cuda') + out_comp = _apply_rope( + x_comp.clone(), + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + x_full = torch.zeros( + seq * ratio, batch, heads, head_dim, dtype=torch.bfloat16, device='cuda' + ) + x_full[::ratio][:seq] = x_comp + out_full = _apply_rope( + x_full, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq * ratio, + ratio=1, + cp_group=self.pg_collection.cp, + ) + out_ref = out_full[::ratio][:seq] + + assert torch.allclose(out_comp, out_ref, rtol=1e-3, atol=1e-3), ( + f"ratio={ratio} stride mismatch: " + f"max abs diff = {(out_comp - out_ref).abs().max().item():.3e}" + ) + + +# =========================================================================== +# csa_dense_mode tests +# =========================================================================== + + +class TestCompressedSparseAttentionDenseMode: + """Test that csa_dense_mode=True disables the indexer for ratio=4 layers.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], csa_window_size=8, csa_dense_mode=True + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_disables_indexer_for_ratio4(self): + """With csa_dense_mode=True, ratio=4 layers should NOT build an indexer.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + ).cuda() + + assert csa.compress_ratio == 4 + assert csa.compressor is not None, "Compressor should still be built" + assert csa.indexer is None, "Indexer should be disabled in dense mode" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_forward_ratio4(self): + """Forward pass should work for ratio=4 in dense mode (uses all compressed positions).""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + +class TestCSAHighPrecisionParams: + """Reference-checkpoint FP32 parameters survive BF16 model conversion.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4]) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ape_and_attn_sink_stay_fp32_after_bf16_conversion(self): + from megatron.core.transformer.module import Float16Module + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + name="decoder.layers.0.self_attention.core_attention", + ) + + assert csa.attn_sink.dtype == torch.float32 + assert csa.compressor.ape.dtype == torch.float32 + assert csa.indexer.compressor.ape.dtype == torch.float32 + + bf16_module = Float16Module(config=self.config, module=csa) + + assert bf16_module.module.attn_sink.dtype == torch.float32 + assert bf16_module.module.compressor.ape.dtype == torch.float32 + assert bf16_module.module.indexer.compressor.ape.dtype == torch.float32 + assert bf16_module.module.compressor.linear_wkv.weight.dtype == torch.bfloat16 + assert bf16_module.module.compressor.linear_wgate.weight.dtype == torch.bfloat16 diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py index 135c4802dd3..1881ffaaa40 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py @@ -28,6 +28,8 @@ DSAttention, DSAttentionSubmodules, FusedDSAIndexerLoss, + _compute_indexer_teacher_probabilities, + _normalize_indexer_teacher_target, _run_sparse_attention, _validate_nonpacked_cp_uniform_length, compute_dsa_indexer_loss, @@ -1658,6 +1660,25 @@ def test_rotate_activation_dtype_check(self): rotate_activation(x) +def test_indexer_teacher_probability_accepts_detached_external_mass(): + """An omitted-key LSE participates in the denominator without entering the target support.""" + attention_scores = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) + valid_mask = torch.ones((1, 2, 2), dtype=torch.bool) + non_compressed_lse = torch.tensor([[[0.5, 1.5]]]) + + actual = _compute_indexer_teacher_probabilities( + attention_scores, valid_mask, non_compressed_lse + ) + expected_denominator = torch.logaddexp( + torch.logsumexp(attention_scores, dim=-1), non_compressed_lse + ) + expected = torch.exp(attention_scores - expected_denominator.unsqueeze(-1)) + torch.testing.assert_close(actual, expected) + + normalized = _normalize_indexer_teacher_target(actual.sum(dim=1), non_compressed_lse) + torch.testing.assert_close(normalized.sum(dim=-1), torch.ones((1, 2))) + + @pytest.mark.parametrize("seqlen_and_topk", [[16, 32], [64, 32]]) class TestComputeDSAIndexerLoss: """Test compute_dsa_indexer_loss function.""" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py index 8d63a9bee11..749fb43d3fb 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_backend_tp_sp_parity.py @@ -681,3 +681,101 @@ def test_packed_cp_tp2_sequence_parallel_shared_skip_backend_matches_unfused_ref finally: DSAIndexerLossLoggingHelper.clean_loss_in_tracker() Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_indexer_loss_tracker_grows_for_hybrid_mtp_layer_numbers(): + """Hybrid MTP layers can have a layer number beyond the nominal layer count.""" + DSAIndexerLossLoggingHelper.tracker = {} + try: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(2.0, device="cuda"), layer_number=7, num_layers=5 + ) + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(3.0, device="cuda"), layer_number=9, num_layers=5 + ) + + values = DSAIndexerLossLoggingHelper.tracker["values"] + assert values.shape == (9,) + torch.testing.assert_close(values[6], torch.tensor(2.0, device="cuda")) + torch.testing.assert_close(values[8], torch.tensor(3.0, device="cuda")) + finally: + DSAIndexerLossLoggingHelper.tracker = {} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_indexer_metrics_average_only_ratio4_layers(monkeypatch: pytest.MonkeyPatch): + """Window and compressed-only layers must not dilute the indexer-loss average.""" + DSAIndexerLossLoggingHelper.tracker = { + "values": torch.tensor([0.0, 3.0, 0.0, 0.0, 0.0], device="cuda") + } + monkeypatch.setattr(DSAIndexerLossLoggingHelper, "reduce_loss_in_tracker", lambda **_: None) + total_loss_dict = {} + try: + DSAIndexerLossLoggingHelper.track_indexer_metrics( + loss_scale=0.5, + iteration=1, + writer=None, + pg_collection=object(), + total_loss_dict=total_loss_dict, + num_layers=5, + num_indexer_layers=1, + ) + + torch.testing.assert_close( + total_loss_dict["indexer loss"], torch.tensor(1.5, device="cuda") + ) + finally: + DSAIndexerLossLoggingHelper.tracker = {} + + +def test_indexer_loss_cleanup_can_preserve_graph_reduction_groups(): + """CUDA Graph replays keep stable group objects while clearing accumulated loss.""" + reduce_group = object() + avg_group = object() + DSAIndexerLossLoggingHelper.tracker = { + "values": torch.ones(2), + "reduce_group": reduce_group, + "avg_group": avg_group, + } + try: + DSAIndexerLossLoggingHelper.clean_loss_in_tracker(preserve_groups=True) + assert torch.count_nonzero(DSAIndexerLossLoggingHelper.tracker["values"]) == 0 + assert DSAIndexerLossLoggingHelper.tracker["reduce_group"] is reduce_group + assert DSAIndexerLossLoggingHelper.tracker["avg_group"] is avg_group + finally: + DSAIndexerLossLoggingHelper.tracker = {} + + +def test_indexer_metrics_reduce_across_pipeline_rank_without_indexer(): + """Every pipeline rank must join indexer loss reduction, even without a local indexer.""" + if Utils.world_size < 2: + pytest.skip("Cross-pipeline indexer reduction requires at least two distributed ranks") + + Utils.initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=2) + DSAIndexerLossLoggingHelper.tracker = {} + try: + num_layers = 5 + if parallel_state.get_pipeline_model_parallel_rank() == 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=torch.tensor(3.0, device="cuda"), layer_number=2, num_layers=num_layers + ) + + total_loss_dict = {} + DSAIndexerLossLoggingHelper.track_indexer_metrics( + loss_scale=0.5, + iteration=1, + writer=None, + pg_collection=ProcessGroupCollection.use_mpu_process_groups(required_pgs=['pp', 'dp']), + total_loss_dict=total_loss_dict, + num_layers=num_layers, + num_indexer_layers=1, + ) + + torch.testing.assert_close( + total_loss_dict["indexer loss"], torch.tensor(1.5, device="cuda") + ) + assert torch.count_nonzero(DSAIndexerLossLoggingHelper.tracker["values"]) == 0 + finally: + DSAIndexerLossLoggingHelper.tracker = {} + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py new file mode 100644 index 00000000000..6fd8b54f1cf --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -0,0 +1,599 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + +_SEED = 42 + + +def _mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + return x * scale + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Patch hadamard_transform in dsa/csa modules if the library is not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + _mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# --------------------------------------------------------------------------- +# Config / spec helpers +# --------------------------------------------------------------------------- + + +def _make_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + q_lora_rank=64, + o_groups=8, + o_lora_rank=64, + csa_compress_ratios=None, + csa_window_size=8, + tensor_model_parallel_size=1, + sequence_parallel=False, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + **extra_config_kwargs, +): + """Create an MLATransformerConfig for DSv4 hybrid attention tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [0, 4, 128, 4] + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=q_lora_rank, + kv_lora_rank=v_head_dim - qk_pos_emb_head_dim, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + o_groups=o_groups, + o_lora_rank=o_lora_rank, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + experimental_attention_variant='dsv4_hybrid', + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + **extra_config_kwargs, + ) + + +def _make_attention_spec(config): + """Build the full DSv4HybridSelfAttention ModuleSpec using the canonical spec builder.""" + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + +def test_module_spec_is_built_from_explicit_backend(): + """The neutral spec builder should use only its explicitly supplied backend.""" + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + Compressor, + CSAIndexer, + ) + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + + class Linear: + pass + + class ColumnParallelLinear: + pass + + class RowParallelLinear: + pass + + class Norm: + pass + + class Backend: + def linear(self): + return Linear + + def column_parallel_linear(self): + return ColumnParallelLinear + + def row_parallel_linear(self): + return RowParallelLinear + + def layer_norm(self, rms_norm=False, for_qk=False, has_residual=False): + return Norm + + spec = get_dsv4_hybrid_module_spec_for_backend(_make_config(), Backend()) + + assert spec.module is DSv4HybridSelfAttention + assert spec.submodules.linear_q_down_proj is Linear + assert spec.submodules.linear_q_up_proj is ColumnParallelLinear + assert spec.submodules.linear_kv_proj is ColumnParallelLinear + assert spec.submodules.linear_proj is RowParallelLinear + assert spec.submodules.core_attention.module is CompressedSparseAttention + assert spec.submodules.core_attention.submodules.compressor.module is Compressor + assert spec.submodules.core_attention.submodules.indexer.module is CSAIndexer + + +def test_config_includes_mtp_ratio_and_derives_dimensions(): + """DSv4 config should account for MTP and derive its shared Q/KV content width.""" + config = _make_config(num_layers=2, mtp_num_layers=1, csa_compress_ratios=[0, 4, 128]) + + expected_content_dim = config.v_head_dim - config.qk_pos_emb_head_dim + assert config.qk_head_dim == expected_content_dim + assert config.kv_lora_rank == expected_content_dim + assert config.hetereogenous_dist_checkpoint is True + + +def test_config_rejects_context_parallelism(): + """The SBHD slice should fail early instead of silently accepting unsupported CP.""" + with pytest.raises(AssertionError, match="does not support context parallelism"): + _make_config(context_parallel_size=2) + + +def test_config_rejects_fused_backend_in_native_slice(): + """Fused DSv4 backends belong to the follow-up kernel-integration slice.""" + with pytest.raises(ValueError, match="requires dsa_kernel_backend='none'"): + _make_config(dsa_kernel_backend="cudnn") + + +def test_config_accepts_hybrid_model_ratio_tail(): + """HybridModel may expand each MTP depth into multiple attention layers.""" + config = _make_config(num_layers=2, mtp_num_layers=1, csa_compress_ratios=[0, 4, 128, 4]) + assert config.csa_compress_ratios == [0, 4, 128, 4] + + +def test_constructor_requires_explicit_process_groups(): + """Production DSv4 construction must not read process groups from global MPU state.""" + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + + with pytest.raises(ValueError, match="explicit ProcessGroupCollection"): + DSv4HybridSelfAttention(config=None, submodules=None, layer_number=1) + + +def _build_attention(config, layer_number, pg_collection, **kwargs): + """Instantiate a DSv4HybridSelfAttention from config.""" + from megatron.core.transformer.spec_utils import build_module + + spec = _make_attention_spec(config) + return build_module( + spec, config=config, layer_number=layer_number, pg_collection=pg_collection, **kwargs + ) + + +# =========================================================================== +# Constructor tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionConstructor: + """Test construction of DSv4HybridSelfAttention in the supported TP=1 configuration.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_basic_construction(self): + """Verify the layer builds and has the expected sub-modules.""" + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert isinstance(attn, DSv4HybridSelfAttention) + assert hasattr(attn, 'linear_q_down_proj') + assert hasattr(attn, 'linear_q_up_proj') + assert hasattr(attn, 'linear_kv_proj') + assert hasattr(attn, 'linear_proj') + assert hasattr(attn, 'linear_o_group_proj') + assert hasattr(attn, 'core_attention') + assert hasattr(attn, 'q_layernorm') + assert hasattr(attn, 'kv_layernorm') + + def test_q_head_dim_equals_v_head_dim(self): + """q_head_dim must equal v_head_dim for DSv4 hybrid.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert attn.q_head_dim == config.v_head_dim + + def test_current_main_constructor_kwargs(self): + """Current TransformerLayer forwards module names and pipeline offsets.""" + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention( + config, + layer_number=1, + pg_collection=pg, + pp_layer_offset=0, + name="decoder.layers.0.self_attention", + ) + + assert attn._pp_layer_offset == 0 + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_rope_base_varies_with_compress_ratio(self, layer_number): + """Layers with compress_ratio > 1 should use csa_compress_rotary_base.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + ratios = [0, 4, 128, 4] + config = _make_config(csa_compress_ratios=ratios) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=layer_number, pg_collection=pg) + + ratio = ratios[layer_number - 1] + if ratio > 1: + expected_base = config.csa_compress_rotary_base + else: + expected_base = config.rotary_base + + # inv_freq is derived from rotary_base; verify the correct base was used + dim = config.qk_pos_emb_head_dim + recomputed_inv_freq = 1.0 / ( + expected_base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + assert torch.allclose( + attn.rotary_pos_emb.inv_freq.cpu(), recomputed_inv_freq, rtol=1e-5, atol=1e-5 + ) + + +# =========================================================================== +# Forward / backward tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionForwardBackward: + """Test forward and backward passes of DSv4HybridSelfAttention.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config(dsa_indexer_loss_coeff=1.0) + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_forward_output_shape(self, layer_number): + """Forward should produce [sq, b, hidden_size] output.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert output.dtype == torch.bfloat16 + assert not torch.isnan(output).any() + + @pytest.mark.parametrize("layer_number", [1, 2]) + def test_backward_gradient_flow(self, layer_number): + """Backward should produce gradients for all trainable parameters.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + attn.train() + + hidden = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + output, bias = attn(hidden_states=hidden, attention_mask=None) + loss = output.sum() + loss.backward() + + assert hidden.grad is not None, "No gradient on hidden_states" + for name, param in attn.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}" + + def test_eval_mode(self): + """Forward should work in eval mode.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + attn.eval() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + with torch.no_grad(): + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert not torch.isnan(output).any() + + def test_different_seq_lengths(self): + """Forward should handle various sequence lengths.""" + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=2, pg_collection=self.pg).cuda() + + for seq_len in [64, 128, 256]: + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + output, bias = attn(hidden_states=hidden, attention_mask=None) + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + + +# =========================================================================== +# get_query_key_value_tensors tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridQKV: + """Test get_query_key_value_tensors internals.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config() + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_qkv_shapes(self): + """Query, key, value should have correct shapes.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, q_compressed, kv_compressed = attn.get_query_key_value_tensors(hidden) + + n_heads = self.config.num_attention_heads + v_dim = self.config.v_head_dim + + assert q.shape == (seq_len, batch_size, n_heads, v_dim) + # key and value are single-head (MQA-style) with an extra head dim + assert k.shape[-1] == v_dim + assert v.shape[-1] == v_dim + + def test_key_equals_value(self): + """In the wkv path, key and value should be the same tensor.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) + assert torch.equal(k, v), "key and value should be identical in wkv path" + + +# =========================================================================== +# Grouped output projection tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridGroupedOutput: + """Test that grouped output projection (wo_a) parameters are created.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_o_group_proj_shape(self): + """linear_o_group_proj should have the correct shape.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + o_groups = 8 + o_lora_rank = 64 + config = _make_config(o_groups=o_groups, o_lora_rank=o_lora_rank) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + expected_out = o_groups * o_lora_rank + expected_in = (config.v_head_dim * config.num_attention_heads) // o_groups + assert attn.linear_o_group_proj.shape == (expected_out, expected_in) + assert attn.linear_o_group_proj.requires_grad + + +# =========================================================================== +# apply_rope_fusion tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridRopeFusion: + """Test that apply_rope_fusion=True works for both yarn and non-yarn layers. + + DSv4 Hybrid uses YarnRotaryEmbedding for layers with compress_ratio > 1 + and standard RotaryEmbedding for layers with compress_ratio <= 1. The + fused RoPE path must obtain cos/sin from both embedding classes via + get_cached_cos_sin. + + compress_ratios=[0, 4, 128, 4]: layer 1 has ratio 0 (standard + RotaryEmbedding), layers 2-4 have ratio > 1 (YarnRotaryEmbedding). + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_rope_fusion_forward_backward_parity(self): + """Fused RoPE forward/backward succeeds and matches the unfused path.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + fused_config = _make_config(apply_rope_fusion=True) + attn_fused = _build_attention(fused_config, layer_number=4, pg_collection=self.pg).cuda() + attn_fused.train() + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + unfused_config = _make_config(apply_rope_fusion=False) + attn_unfused = _build_attention( + unfused_config, layer_number=4, pg_collection=self.pg + ).cuda() + attn_unfused.train() + + hidden = torch.randn( + seq_len, batch_size, fused_config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + out_fused, _ = attn_fused(hidden_states=hidden, attention_mask=None) + out_unfused, _ = attn_unfused(hidden_states=hidden, attention_mask=None) + + assert out_fused.shape == (seq_len, batch_size, fused_config.hidden_size) + assert torch.isfinite(out_fused).all() + # The remaining difference is bf16 accumulation order between the fused + # Triton kernel and eager PyTorch operations. + torch.testing.assert_close(out_fused, out_unfused, atol=3e-2, rtol=3e-2) + + hidden_fused = hidden.detach().clone().requires_grad_(True) + hidden_unfused = hidden.detach().clone().requires_grad_(True) + + attn_fused(hidden_states=hidden_fused, attention_mask=None)[0].sum().backward() + attn_unfused(hidden_states=hidden_unfused, attention_mask=None)[0].sum().backward() + + assert hidden_fused.grad is not None + for name, param in attn_fused.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}" diff --git a/tests/unit_tests/transformer/test_module.py b/tests/unit_tests/transformer/test_module.py index 92f15b2f46d..5faf6c81ef1 100644 --- a/tests/unit_tests/transformer/test_module.py +++ b/tests/unit_tests/transformer/test_module.py @@ -4,7 +4,7 @@ import torch from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.module import Float16Module, MegatronModule +from megatron.core.transformer.module import Float16Module, MegatronModule, mark_keep_in_fp32 from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils @@ -163,3 +163,18 @@ def test_bf16_module(self): x = torch.ones((2, 2)).cuda() # inputs are converted to bf16 then outputs are converted to fp32 assert bf16_module(x).dtype == torch.float32 + + @pytest.mark.parametrize( + ('precision', 'dtype'), [('fp16', torch.float16), ('bf16', torch.bfloat16)] + ) + def test_keep_in_fp32_params(self, precision, dtype): + transformer_config = self.transformer_config + megatron_module = self.megatron_module + megatron_module.fp32_param = mark_keep_in_fp32( + torch.nn.Parameter(torch.zeros(4, dtype=torch.float32, device='cuda')) + ) + setattr(transformer_config, precision, True) + float16_module = Float16Module(config=transformer_config, module=megatron_module) + + assert float16_module.module.linear.weight.dtype == dtype + assert float16_module.module.fp32_param.dtype == torch.float32