Integrate DSA training into Megatron-LM

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).

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.

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(),
+        )

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
         # ==================================

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.")

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