diff --git a/docs/README_megatron_bridge_dsa.md b/docs/README_megatron_bridge_dsa.md new file mode 100644 index 00000000000..a8febfe2484 --- /dev/null +++ b/docs/README_megatron_bridge_dsa.md @@ -0,0 +1,22 @@ +# Megatron-Bridge DSA Enablement Plan + +## Goals +- Surface DeepSeek Sparse Attention configuration through Megatron-Bridge so downstream Hugging Face models can opt into DSA-aware training runs without manual patching. +- Provide conversion utilities that translate Hugging Face checkpoint metadata into Megatron DSA settings (e.g., top-k, indexer ranks). +- Guarantee backwards compatibility for existing dense-only pipelines. + +## Implementation Outline +- **Configuration Plumbing** + - Extend bridge configuration schemas with a `dsa` stanza that mirrors Megatron’s `MLATransformerConfig` additions. + - Add CLI flags/env overrides so pipelines can toggle `warmup` or `sparse` modes at launch. +- **Checkpoint & State Mapping** + - Introduce helpers that initialize the DSA indexer weights when importing dense checkpoints, including an optional dense warmup phase to populate the lightning indexer. + - Update save/load paths to persist DSA-specific buffers (indexer KL stats, top-k history) alongside existing MLA caches. +- **Workflow Integration** + - Provide recipe templates demonstrating dense warmup scheduling followed by sparse fine-tuning within the bridge orchestration layer. + - Add validation hooks that report DSA sparsity statistics back to orchestrators for monitoring. + +## Testing & Rollout +- Unit-test the new configuration serialization/deserialization paths with both dense-only and DSA-enabled configs. +- Add smoke tests that round-trip a Hugging Face model through Megatron-Bridge with `dsa_training_mode=sparse` and confirm reproducible outputs. +- Coordinate documentation updates so downstream teams understand the staged training process (dense warmup → sparse adaptation). diff --git a/docs/README_transformer_engine_dsa.md b/docs/README_transformer_engine_dsa.md new file mode 100644 index 00000000000..411a242fb1b --- /dev/null +++ b/docs/README_transformer_engine_dsa.md @@ -0,0 +1,22 @@ +# Transformer Engine DSA Integration Plan + +## Objectives +- Add sparse token selection support to the MLA attention kernels so they can consume DSA top-k masks from Megatron-LM during both forward and backward passes. +- Provide a lightweight indexer training API that exposes lightning indexer activations and gradients without duplicating projections in Python. +- Maintain feature parity across Hopper and Blackwell kernels while keeping existing dense-path regression tests green. + +## Proposed Workstreams +- **Kernel Extensions** + - Extend MLA fused attention kernels to accept an additional per-query sparse mask (block-structured) and fallback to dense behavior when the mask is absent. + - Add optional FP8 blockwise quantization utilities shared with the inference repo to avoid reimplementing quantization logic. +- **Indexer Autograd Support** + - Export custom autograd functions that surface the KL-divergence loss hooks currently implemented in Python so Megatron can attach DSA losses without manual bookkeeping. + - Ensure gradients can be accumulated independently for the indexer parameters while preventing cross-talk with the dense attention path. +- **Configuration & Testing** + - Introduce a `dsa_mode` flag mirroring the Megatron configuration to make staged roll-out (disabled → warmup → sparse) explicit in TE. + - Add unit tests for warmup and sparse modes including mask broadcasting edge cases and distributed TP=1/2 validation. + +## Dependencies & Coordination +- Align Tensor shapes and scaling factors with Megatron’s `DeepSeekSparseIndexer` to avoid redundant transpose/cast operations. +- Coordinate release timelines with Megatron-LM so the kernel API lands before the training-side default flips to sparse mode. +- Work with QA to extend existing MLA CI coverage (nightly Hopper + Blackwell) to include sparse attention regression suites. diff --git a/megatron/core/transformer/deepseek_sparse_attention.py b/megatron/core/transformer/deepseek_sparse_attention.py new file mode 100644 index 00000000000..45a92019a0c --- /dev/null +++ b/megatron/core/transformer/deepseek_sparse_attention.py @@ -0,0 +1,208 @@ +import math +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from megatron.core import parallel_state +from megatron.core.transformer.moe.moe_utils import ( + MoEAuxLossAutoScaler, + save_to_aux_losses_tracker, +) +from megatron.core.transformer.transformer_config import MLATransformerConfig + + +_EPS = 1e-6 + + +def _hadamard_transform(x: Tensor) -> Tensor: + """Apply an in-place Hadamard transform along the last dimension.""" + hidden = x.size(-1) + if hidden & (hidden - 1) != 0: + raise ValueError( + f"DSA requires index head dimension to be a power of two, but got {hidden}." + ) + stride = 1 + out = x + while stride < hidden: + out = out.reshape(*out.shape[:-1], -1, 2, stride) + first = out[..., 0, :] + second = out[..., 1, :] + out = torch.cat((first + second, first - second), dim=-1) + stride <<= 1 + return out / math.sqrt(hidden) + + +def _apply_rotary(x: Tensor, freqs_cis: Tensor) -> Tensor: + """Apply rotary positional embedding to the provided tensor.""" + dtype = x.dtype + complex_x = torch.view_as_complex( + x.float().view(*x.shape[:-1], -1, 2) + ) + freqs = freqs_cis.view(1, x.size(1), 1, -1) + rotated = torch.view_as_real(complex_x * freqs).flatten(-2) + return rotated.to(dtype) + + +@dataclass +class DSAOutput: + topk_indices: Tensor + additive_mask: Optional[Tensor] + kl_loss: Tensor + + +class DeepSeekSparseIndexer(nn.Module): + """Implements the DeepSeek Sparse Attention indexer for training.""" + + def __init__(self, config: MLATransformerConfig, layer_number: int) -> None: + super().__init__() + if config.tensor_model_parallel_size != 1: + raise ValueError("DSA training currently requires tensor_model_parallel_size == 1.") + + self.config = config + self.layer_number = layer_number + + self.training_mode = getattr(config, "dsa_training_mode", "disabled") + self.index_heads = getattr(config, "dsa_index_n_heads", 64) + self.head_dim = getattr(config, "dsa_index_head_dim", 128) + self.topk = getattr(config, "dsa_index_topk", 2048) + self.loss_weight = getattr(config, "dsa_kl_loss_weight", 1.0) + self.detach_inputs = getattr(config, "dsa_detach_indexer_inputs", True) + self.log_metrics = getattr(config, "dsa_log_aux_metrics", False) + + self.rope_dim = min(config.qk_pos_emb_head_dim, self.head_dim) + + self.q_proj = nn.Linear(config.q_lora_rank, self.index_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False) + self.k_norm = nn.LayerNorm(self.head_dim, eps=config.layernorm_epsilon) + self.weight_proj = nn.Linear(config.hidden_size, self.index_heads, bias=False) + + def _maybe_detach(self, tensor: Tensor) -> Tensor: + return tensor.detach() if self.detach_inputs else tensor + + def _compute_index_scores( + self, + hidden_states: Tensor, + qr_states: Tensor, + freqs_cis: Optional[Tensor], + ) -> Tensor: + bsz, seqlen, _ = hidden_states.shape + + q_latent = self.q_proj(qr_states).view(bsz, seqlen, self.index_heads, self.head_dim) + q_pe, q_nope = torch.split( + q_latent, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + if self.rope_dim > 0 and freqs_cis is not None: + q_pe = _apply_rotary(q_pe, freqs_cis) + q_latent = torch.cat((q_pe, q_nope), dim=-1) + + k_latent = self.k_proj(hidden_states) + k_latent = self.k_norm(k_latent) + if self.rope_dim > 0 and freqs_cis is not None: + k_pe, k_nope = torch.split( + k_latent, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + k_pe = _apply_rotary(k_pe.unsqueeze(2), freqs_cis).squeeze(2) + k_latent = torch.cat((k_pe, k_nope), dim=-1) + + q_latent = _hadamard_transform(q_latent) + k_latent = _hadamard_transform(k_latent) + + weights = self.weight_proj(hidden_states) * (self.index_heads ** -0.5) + + raw = torch.einsum("bqhd,bkd->bqhk", q_latent, k_latent) + raw = F.relu(raw) + index_scores = (raw * weights.unsqueeze(-1)).sum(dim=2) + return index_scores + + def forward( + self, + hidden_states: Tensor, + qr_states: Tensor, + freqs_cis: Optional[Tensor], + causal_mask: Optional[Tensor], + *, + attention_scores: Optional[Tensor] = None, + query: Optional[Tensor] = None, + key: Optional[Tensor] = None, + attn_scale: Optional[float] = None, + training: bool, + ) -> Optional[DSAOutput]: + if self.training_mode == "disabled": + return None + + hidden_states = self._maybe_detach(hidden_states) + qr_states = self._maybe_detach(qr_states) + + bsz, seqlen, _ = hidden_states.shape + index_scores = self._compute_index_scores(hidden_states, qr_states, freqs_cis) + + if causal_mask is not None: + index_scores = index_scores + causal_mask.unsqueeze(0) + + if query is not None and key is not None: + q = query.permute(1, 0, 2, 3).float() + k = key.permute(1, 0, 2, 3).float() + head_dim = min(q.size(-1), k.size(-1), self.config.qk_head_dim) + q = q[..., :head_dim] + k = k[..., :head_dim] + dense_logits = torch.einsum("bshd,bthd->bsht", q, k) + if attn_scale is not None: + dense_logits = dense_logits * attn_scale + if causal_mask is not None: + dense_logits = dense_logits + causal_mask.unsqueeze(0).unsqueeze(2) + dense_probs = torch.softmax(dense_logits, dim=-1).detach() + dense_probs = dense_probs.mean(dim=2) + elif attention_scores is not None: + dense_logits = attention_scores.float() + if causal_mask is not None: + dense_logits = dense_logits + causal_mask.unsqueeze(0).unsqueeze(2) + dense_probs = torch.softmax(dense_logits, dim=-1).detach() + dense_probs = dense_probs.mean(dim=2) + else: + dense_probs = torch.softmax(index_scores, dim=-1).detach() + + topk = min(self.topk, seqlen) + topk_indices = torch.topk(index_scores, k=topk, dim=-1).indices + + if self.training_mode == "warmup": + pred_probs = torch.softmax(index_scores, dim=-1) + target_probs = dense_probs + kl = F.kl_div( + torch.log(pred_probs + _EPS), target_probs + _EPS, reduction="batchmean" + ) + additive_mask = None + else: + gathered_pred = torch.gather(index_scores, dim=-1, index=topk_indices) + pred_probs = torch.softmax(gathered_pred, dim=-1) + target_selected = torch.gather(dense_probs, dim=-1, index=topk_indices) + target_selected = target_selected / (target_selected.sum(dim=-1, keepdim=True) + _EPS) + kl = F.kl_div( + torch.log(pred_probs + _EPS), target_selected + _EPS, reduction="batchmean" + ) + additive_mask = index_scores.new_full((bsz, seqlen, seqlen), float("-inf")) + additive_mask.scatter_(-1, topk_indices, 0.0) + + if training and self.loss_weight > 0.0: + scaled_loss = self.loss_weight * kl + index_scores = MoEAuxLossAutoScaler.apply(index_scores, scaled_loss) + + if training and self.log_metrics: + num_layers = self.config.num_layers + if self.config.mtp_num_layers is not None: + num_layers += self.config.mtp_num_layers + save_to_aux_losses_tracker( + "dsa_indexer_loss", + kl.detach(), + self.layer_number, + num_layers, + reduce_group=parallel_state.get_data_parallel_group(with_context_parallel=True), + ) + + return DSAOutput( + topk_indices=topk_indices, + additive_mask=additive_mask, + kl_loss=kl.detach(), + ) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index a8893ebec36..2fb0dd5204a 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -37,6 +37,7 @@ from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import MLATransformerConfig from megatron.core.utils import deprecate_inference_params, is_te_min_version +from megatron.core.transformer.deepseek_sparse_attention import DeepSeekSparseIndexer try: from megatron.core.fusions.fused_mla_yarn_rope_apply import ( @@ -188,6 +189,13 @@ def __init__( # the quantized tensor. set_save_original_input(self.linear_proj) + self.dsa_indexer = ( + DeepSeekSparseIndexer(self.config, self.layer_number) + if self.config.dsa_training_mode != "disabled" + else None + ) + self.dsa_last_kl_loss = None + def forward( self, hidden_states, @@ -257,6 +265,58 @@ def forward( if value is not None: value = value.contiguous() + def _mask_to_bss(mask: torch.Tensor, batch: int) -> torch.Tensor: + if mask.dim() == 2: + mask = mask.unsqueeze(0) + if mask.dim() == 4: + mask = mask.squeeze(1) + if mask.dim() == 3 and mask.size(0) == 1 and batch > 1: + mask = mask.expand(batch, -1, -1) + return mask + + def _ensure_4d(mask: torch.Tensor, batch: int) -> torch.Tensor: + if mask.dim() == 2: + mask = mask.unsqueeze(0).unsqueeze(0) + elif mask.dim() == 3: + if mask.size(0) == 1 and batch > 1: + mask = mask.expand(batch, -1, -1) + mask = mask.unsqueeze(1) + elif mask.dim() == 4 and mask.size(0) == 1 and batch > 1: + mask = mask.expand(batch, -1, -1, -1) + return mask + + if ( + self.dsa_indexer is not None + and inference_context is None + and packed_seq_params is None + and attention_mask is not None + ): + batch = hidden_states.size(1) + dsa_hidden = hidden_states.permute(1, 0, 2).contiguous() + if self.config.q_lora_rank is not None: + qr_states = q_compressed.permute(1, 0, 2).contiguous() + else: + qr_states = dsa_hidden + causal_mask = _mask_to_bss(attention_mask, batch) + dsa_output = self.dsa_indexer( + dsa_hidden, + qr_states, + freqs_cis=None, + causal_mask=causal_mask, + query=query, + key=key, + attn_scale=self.softmax_scale, + training=self.training, + ) + if dsa_output is not None: + self.dsa_last_kl_loss = dsa_output.kl_loss + if dsa_output.additive_mask is not None: + dsa_mask = dsa_output.additive_mask.unsqueeze(1) + attention_mask = _ensure_4d(attention_mask, batch) + attention_mask = attention_mask + dsa_mask + else: + attention_mask = _ensure_4d(attention_mask, batch) + # ================================== # core attention computation # ================================== diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 68bdfbcf021..01371949648 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1626,6 +1626,28 @@ class MLATransformerConfig(TransformerConfig): This is only for the dynamic inference backend and requires that Flash MLA is installed.""" + dsa_training_mode: str = "disabled" + """Training mode for DeepSeek Sparse Attention. Supported values: + 'disabled', 'warmup', and 'sparse'.""" + + dsa_index_n_heads: int = 64 + """Number of indexer heads used by DSA.""" + + dsa_index_head_dim: int = 128 + """Head dimension for the DSA indexer.""" + + dsa_index_topk: int = 2048 + """Number of top tokens selected per query in sparse mode.""" + + dsa_kl_loss_weight: float = 1.0 + """Weight applied to the DSA KL loss.""" + + dsa_detach_indexer_inputs: bool = True + """Whether to detach inputs feeding the DSA indexer (matches DeepSeek warm-up protocol).""" + + dsa_log_aux_metrics: bool = False + """If enabled, logs DSA auxiliary losses through the MoE logging utility.""" + def __post_init__(self): super().__post_init__() if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": @@ -1635,3 +1657,16 @@ def __post_init__(self): assert ( self.apply_rope_fusion is False ), "Rope Fusion is not compatible with caching latents" + + if self.dsa_training_mode not in {"disabled", "warmup", "sparse"}: + raise ValueError( + f"Unsupported dsa_training_mode '{self.dsa_training_mode}'. " + "Choose from 'disabled', 'warmup', or 'sparse'." + ) + if self.dsa_training_mode != "disabled": + if not self.multi_latent_attention: + raise ValueError("DSA requires multi_latent_attention to be enabled.") + if self.dsa_index_topk <= 0: + raise ValueError("dsa_index_topk must be positive when DSA is enabled.") + if self.dsa_index_head_dim & (self.dsa_index_head_dim - 1) != 0: + raise ValueError("dsa_index_head_dim must be a power of two for DSA.") diff --git a/megatron/training/yaml_arguments.py b/megatron/training/yaml_arguments.py index 405d7b70fad..3ab4f18e99c 100644 --- a/megatron/training/yaml_arguments.py +++ b/megatron/training/yaml_arguments.py @@ -393,6 +393,10 @@ def core_config_from_args(args, dataclass=TransformerConfig): for f in dataclasses.fields(dataclass): if hasattr(args, f.name): kw_args[f.name] = getattr(args, f.name) + elif f.default is not dataclasses.MISSING: + kw_args[f.name] = f.default + elif getattr(f, "default_factory", dataclasses.MISSING) is not dataclasses.MISSING: + kw_args[f.name] = f.default_factory() else: raise Exception(f"Missing argument {f.name} for {str(dataclass)} config") return kw_args diff --git a/tests/unit_tests/transformer/test_deepseek_sparse_attention.py b/tests/unit_tests/transformer/test_deepseek_sparse_attention.py new file mode 100644 index 00000000000..6bb0042a2cf --- /dev/null +++ b/tests/unit_tests/transformer/test_deepseek_sparse_attention.py @@ -0,0 +1,88 @@ +import torch + +from megatron.core.transformer.deepseek_sparse_attention import DeepSeekSparseIndexer +from megatron.core.transformer.transformer_config import MLATransformerConfig + + +def _make_config(mode: str) -> MLATransformerConfig: + return MLATransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + q_lora_rank=32, + kv_lora_rank=64, + qk_head_dim=32, + qk_pos_emb_head_dim=16, + v_head_dim=32, + dsa_training_mode=mode, + dsa_index_n_heads=4, + dsa_index_head_dim=32, + dsa_index_topk=3, + dsa_kl_loss_weight=0.5, + ) + + +def _dummy_inputs(batch: int, seq: int, config: MLATransformerConfig): + hidden = torch.randn(batch, seq, config.hidden_size) + qr = torch.randn(batch, seq, config.q_lora_rank or config.hidden_size) + query = torch.randn(seq, batch, config.num_attention_heads, config.qk_head_dim) + key = torch.randn_like(query) + causal = torch.full((batch, seq, seq), float("-inf")) + causal = torch.triu(causal, diagonal=1) + return hidden, qr, query, key, causal + + +def test_dsa_warmup_outputs_mask_is_none(): + torch.manual_seed(1) + config = _make_config("warmup") + indexer = DeepSeekSparseIndexer(config, layer_number=1) + hidden, qr, query, key, causal = _dummy_inputs(batch=2, seq=5, config=config) + + result = indexer( + hidden_states=hidden, + qr_states=qr, + freqs_cis=None, + causal_mask=causal, + query=query, + key=key, + attn_scale=1.0, + training=True, + ) + + assert result is not None + assert result.additive_mask is None + assert result.topk_indices.shape == (2, 5, config.dsa_index_topk) + assert torch.isfinite(result.kl_loss) + + +def test_dsa_sparse_outputs_mask_applies_topk(): + torch.manual_seed(2) + config = _make_config("sparse") + indexer = DeepSeekSparseIndexer(config, layer_number=1) + hidden, qr, query, key, causal = _dummy_inputs(batch=1, seq=4, config=config) + + result = indexer( + hidden_states=hidden, + qr_states=qr, + freqs_cis=None, + causal_mask=causal, + query=query, + key=key, + attn_scale=1.0, + training=True, + ) + + assert result is not None + assert result.additive_mask is not None + mask = result.additive_mask + assert mask.shape == (1, 4, 4) + + topk = result.topk_indices[0] + for row in range(topk.size(0)): + selected = topk[row] + assert torch.all(mask[0, row, selected] == 0) + dropped_selector = torch.ones(mask.size(-1), dtype=bool) + dropped_selector[selected] = False + dropped_values = mask[0, row, dropped_selector] + if dropped_values.numel() > 0: + assert torch.all(dropped_values <= 0)