feat: DSA conversion (dense-MLA -> DeepSeek Sparse Attention), with a Kimi K2 example

docs/advanced.md

@@ -6,6 +6,7 @@ This page covers the specialized features layered on top of the core training st
 
 - [Custom Modeling](#custom-modeling)
   - [Expert Parallelism Backends](#expert-parallelism-backends)
+  - [DSA Conversion](#dsa-conversion)
 - [Multimodal Training](#multimodal-training)
   - [Supported Families](#supported-families)
   - [Enabling VLM Mode](#enabling-vlm-mode)
@@ -67,6 +68,76 @@ DeepEP requires some careful tuning to achieve optimal performance, tuning param
 
 With DeepEP, gradient clipping is currently not supported. (`optim.max_norm` is set to `None` automatically.)
 
+### DSA Conversion
+
+DSA-capable families (`glm_moe_dsa`, `kimi_k2_dsa`) share one attention module (`src/prime_rl/trainer/models/layers/dsa.py`) that runs in two modes, toggled by the model's `use_sparse_attn` field (default `True`): the normal DeepSeek Sparse Attention (DSA) top-k path, or an ordinary dense-causal fallback over the same MLA projections. Dense mode exists to convert an existing dense-MLA checkpoint to DSA via continued pretraining, mirroring the two-stage recipe DeepSeek (arXiv:2512.02556) and GLM (arXiv:2602.15763) used. Any dense-MLA family gets this for free by building its DSA sibling on the same module — see `kimi_k2_dsa` alongside `kimi_k2` for the pattern.
+
+**0. Bootstrap.** A dense checkpoint has no `indexer.*` weights at all. Point `model.name` at the dense checkpoint, force the DSA config/model class (`impl = "custom"`), and set `[model.debug] random_init_indexer = true` for exactly this one run: the loader strips the (absent) indexer keys from the strict checkpoint load instead of erroring, then randomly initializes them. Turn it back off from the very next run onward — by then the checkpoint has real indexer weights.
+
+**1. Indexer warm-up.** On `[model]`: `use_sparse_attn = false`, `freeze_all_except_indexer = true` (freezes everything *but* the indexer — plain `freeze_sparse_indexer = false` alone does not freeze the rest of the model) and `indexer_kl_coeff = <coeff>` (also flips the HF config's `train_indexer` on). Only the indexer trains, against a KL loss to the real dense attention distribution.
+
+**2. Sparse adaptation.** `use_sparse_attn = true`, `indexer_kl_coeff = <coeff>` (still on); `freeze_all_except_indexer = false` (the default) so everything trains jointly with the ordinary LM loss; the indexer's KL target now comes from the actual sparse-selected keys instead of the full distribution.
+
+Either data format works, since neither the indexer-KL loss nor the dense-mode forward pass consults `loss_mask` — ordinary SFT masking doesn't interfere with training the indexer over every token. DeepSeek-V3.2 and GLM-5's own from-scratch conversions used continued-pretraining-style raw text (`[data] type = "raw_text"`, `RawTextDataConfig`); IndexCache's from-scratch conversion (arXiv:2603.12201, applied to GLM-4.7-Flash) used plain SFT data instead (`type = "sft"`) — 1,000 warm-up steps + 4,000 sparse-training steps at a 200K context length, no raw-text corpus needed. SFT data is the more practical default unless you're specifically trying to reproduce DeepSeek's/GLM-5's own recipe. Run each stage as a separate `uv run sft` invocation, chained via checkpoint resume (`ckpt.resume_step`/`ckpt.resume_path`).
+
+Once converted, recover any quality lost to sparsification with the existing `opd` algorithm: point `algorithm.teacher` at the *original* dense checkpoint (served independently — never wired into the DSA student's weight-broadcast group, since it lacks the `indexer.*` parameters) and train the converted checkpoint as the student.
+
+**Example: converting Kimi K2 to DSA.**
+
+First prepare a local checkpoint directory: copy `moonshotai/Kimi-K2-Instruct`'s safetensors files as-is, but replace its `config.json` with one whose `model_type` is `"kimi_k2_dsa"` (not `"kimi_k2"`) and which adds the indexer fields (`index_topk`, `index_n_heads`, `index_head_dim`, ...) — this is what tells prime-rl to load the checkpoint's (otherwise-identical) weights into the DSA-capable model class instead of the plain dense one.
+
+```toml
+# bootstrap.toml — one-off, produces the first real (randomly-initialized-indexer) checkpoint
+[model]
+name = "/path/to/kimi-k2-dsa-prepared"   # the copy with the edited config.json above
+impl = "custom"
+
+[model.debug]
+random_init_indexer = true
+
+[data]
+type = "fake"   # any single step works; this run exists only to materialize+save the bootstrap checkpoint
+```
+
+```toml
+# stage_a.toml — indexer warm-up (SFT data, per IndexCache's real recipe — swap
+# [data] for type = "raw_text" + RawTextDataConfig fields to follow DeepSeek/GLM-5's
+# continued-pretraining recipe instead)
+max_steps = 1000   # IndexCache's own figure for this stage
+
+[model]
+name = "<bootstrap.toml's output checkpoint>"
+impl = "custom"
+use_sparse_attn = false
+freeze_all_except_indexer = true
+indexer_kl_coeff = 1.0
+
+[data]
+type = "sft"
+name = "<a long-context SFT dataset>"
+seq_len = 200000   # IndexCache's own figure for this recipe; adjust to your budget
+```
+
+```toml
+# stage_b.toml — sparse adaptation, resumes from stage_a's output
+max_steps = 4000   # IndexCache's own figure for this stage
+
+[model]
+name = "<stage_a.toml's output checkpoint>"
+impl = "custom"
+use_sparse_attn = true
+indexer_kl_coeff = 1.0
+
+[data]
+type = "sft"
+name = "<same dataset, or your real SFT mix>"
+seq_len = 200000
+```
+
+Recovery afterward is `algorithm.teacher` pointed at an independently-served `moonshotai/Kimi-K2-Instruct` (the original dense checkpoint) under the `opd` algorithm, training the converted checkpoint as the student.
+
+Kimi K2 is a 1T-parameter model — all four runs above need real multi-node compute, not a single-box smoke test.
+
 ## Multimodal Training
 
 ### Supported Families

packages/prime-rl-configs/src/prime_rl/configs/sft.py

@@ -73,10 +73,12 @@ class LossMaskConfig(BaseConfig):
     """Tool messages contribute to the loss."""
 
 
-class SFTDataConfig(BaseDataConfig):
-    type: Literal["sft"] = "sft"
+class HfInterleavedDataConfig(BaseDataConfig):
+    """Shared shape for datasets loaded (and optionally interleaved) straight from the HF
+    Hub — subsets/splits/probabilities/stopping strategy, plus epoch shuffling. Subclasses
+    add whatever's specific to how each row gets turned into a training sample."""
 
-    name: str = "PrimeIntellect/Reverse-Text-SFT"
+    name: str
     """HF dataset name or path."""
 
     subsets: list[str] | None = None
@@ -97,10 +99,6 @@ class SFTDataConfig(BaseDataConfig):
     seed: int = 0
     """Random seed for shuffling. Re-shuffled per epoch by adding the epoch count to the seed."""
 
-    # Configuring
-    loss_mask: LossMaskConfig = LossMaskConfig()
-    """Which message types contribute to the loss."""
-
     @model_validator(mode="after")
     def validate_subsets_and_splits(self):
         if self.subsets is not None or self.splits is not None:
@@ -122,6 +120,23 @@ def validate_subsets_and_splits(self):
         return self
 
 
+class SFTDataConfig(HfInterleavedDataConfig):
+    type: Literal["sft"] = "sft"
+
+    name: str = "PrimeIntellect/Reverse-Text-SFT"
+    """HF dataset name or path."""
+
+    loss_mask: LossMaskConfig = LossMaskConfig()
+    """Which message types contribute to the loss."""
+
+
+class RawTextDataConfig(HfInterleavedDataConfig):
+    type: Literal["raw_text"] = "raw_text"
+
+    text_column: str = "text"
+    """Column holding the raw text to tokenize."""
+
+
 class SFTValConfig(BaseConfig):
     interval: int = Field(50, ge=1)
     """Run validation every N training steps."""
@@ -132,7 +147,7 @@ class SFTValConfig(BaseConfig):
     data: SFTDataConfig
 
 
-DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig, Field(discriminator="type")]
+DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig | RawTextDataConfig, Field(discriminator="type")]
 
 
 class BaseDeploymentConfig(BaseConfig):

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -106,6 +106,12 @@ class DebugModelConfig(BaseConfig):
     random_init: bool = False
     """Randomly initialize the model instead of loading weights."""
 
+    random_init_indexer: bool = False
+    """Bootstrap a DSA conversion: strip `indexer.*` keys before loading a dense predecessor
+    checkpoint (which has none) and randomly initialize them afterward, instead of erroring
+    on the missing keys. One-time — turn off again once the checkpoint has real indexer
+    weights (i.e. after the first checkpoint save of the warm-up stage)."""
+
     force_balanced_routing: bool = False
     """Replace MoE token-choice routing with a round-robin assignment so every expert sees an equal share. Intended for fake-data smoke tests where untrained routing would otherwise OOM under severe imbalance. Gating scores are still gathered from the override indices so the forward pass stays consistent."""
 
@@ -218,6 +224,27 @@ class ModelConfig(BaseModelConfig):
     freeze_moe_router: bool = False
     """Freeze MoE router parameters during training."""
 
+    freeze_sparse_indexer: bool = True
+    """Freeze DSA sparse-attention indexer parameters during training (default: True, the
+    common case of loading an already-converted DSA checkpoint). Set False during DSA
+    conversion's indexer warm-up stage, together with the model's `train_indexer=True`
+    and `indexer_kl_coeff`, to train the indexer instead."""
+
+    indexer_kl_coeff: float | None = None
+    """Weight of the DSA indexer-vs-attention KL loss (see `compute_indexer_kl_loss`),
+    added to the training loss. `None` (default) disables it. Only meaningful when the
+    model config sets `train_indexer=True`; a no-op otherwise."""
+
+    use_sparse_attn: bool | None = None
+    """Override a DSA-capable model's `use_sparse_attn` HF config field (dense-causal
+    fallback vs. the real DSA top-k path — see the "DSA Conversion" doc section). `None`
+    (default) leaves whatever the checkpoint's own `config.json` says untouched."""
+
+    freeze_all_except_indexer: bool = False
+    """DSA conversion's indexer warm-up stage: freeze every parameter except the sparse
+    indexer's (overrides `freeze_sparse_indexer`/`freeze_moe_router`, which would otherwise
+    also freeze it). Off by default — this is a training-recipe stage, not a normal mode."""
+
     lora: LoRAConfig | None = None
     """LoRA configuration. If None, LoRA is disabled."""
 

src/prime_rl/trainer/model.py

@@ -44,12 +44,12 @@
     get_custom_vlm_cls,
     supports_custom_impl,
 )
-from prime_rl.trainer.models.glm_moe_dsa.sparse_mla_attention import Indexer
 from prime_rl.trainer.models.layers.checkpointing import (
     get_supported_targets,
     set_selective_activation_checkpointing,
     supports_selective_activation_checkpointing,
 )
+from prime_rl.trainer.models.layers.dsa import Indexer, strip_indexer_from_state_dict
 from prime_rl.trainer.models.layers.fp8_linear import replace_linear_with_fp8_blockwise_linear
 from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head
 from prime_rl.trainer.models.layers.moe import LatentMoE, MoE, TokenChoiceTopKRouter
@@ -386,12 +386,17 @@ def apply_fp32_moe_router(model: nn.Module) -> None:
 def freeze_sparse_indexer(model: nn.Module) -> None:
     """Freeze DSA sparse-attention indexer parameters.
 
-    The indexer's `compute_sparse_indices` forward runs under `torch.no_grad()`, so its
-    params never receive a gradient and cannot be trained. Left with requires_grad=True
-    they stay stateless in the optimizer, which breaks strict checkpoint resume: DCP
-    materializes optimizer state for every requires_grad param at load time, but the
-    stateless params were never saved -> "Missing key in checkpoint state_dict". Freezing
-    them keeps the saved and loaded optimizer state symmetric.
+    Default (config.freeze_sparse_indexer=True): the indexer is a fixed, pretrained
+    component (the common case — loading an already-converted DSA checkpoint like GLM-5).
+    Left with requires_grad=True with nothing feeding it a gradient, it would stay
+    stateless in the optimizer, which breaks strict checkpoint resume: DCP materializes
+    optimizer state for every requires_grad param at load time, but a stateless param was
+    never saved -> "Missing key in checkpoint state_dict". Freezing keeps the saved and
+    loaded optimizer state symmetric.
+
+    Set config.freeze_sparse_indexer=False (and model.train_indexer=True) during DSA
+    conversion's indexer warm-up stage, where `SparseMlaAttention`'s differentiable
+    `Indexer.score()` path and `compute_indexer_kl_loss` do supply a gradient.
     """
     logger = get_logger()
     num_frozen = 0
@@ -406,6 +411,27 @@ def freeze_sparse_indexer(model: nn.Module) -> None:
         logger.info(f"Froze {num_frozen} sparse indexer parameters")
 
 
+def freeze_all_except_indexer(model: nn.Module) -> None:
+    """DSA conversion's indexer warm-up stage: freeze every parameter except the indexer's,
+    mirroring `freeze_all_except_lora_and_specified`'s shape (`trainer/lora.py`). The paper
+    recipe (arXiv:2512.02556, arXiv:2602.15763) trains only the indexer against a frozen
+    base model before turning sparsity on — this is the mechanism for that, an explicit
+    override of `freeze_sparse_indexer`/`freeze_moe_router` (which would otherwise freeze
+    the one thing this stage needs trainable)."""
+    logger = get_logger()
+    num_trainable = 0
+    for name, param in model.named_parameters():
+        if "indexer." in name:
+            param.requires_grad = True
+            num_trainable += 1
+        else:
+            param.requires_grad = False
+
+    if num_trainable == 0:
+        raise ValueError("No indexer parameters found to leave trainable. Is this a DSA-capable model?")
+    logger.info(f"Froze all parameters except {num_trainable} sparse indexer parameters")
+
+
 def apply_force_balanced_routing(model: nn.Module) -> None:
     """Force MoE token-choice routers into round-robin assignment for fake-data smoke tests."""
     logger = get_logger()
@@ -556,6 +582,12 @@ def get_model(
                 f"({sum(t == 'full' for t in indexer_types)}/{len(indexer_types)} full layers)"
             )
 
+    if config.indexer_kl_coeff is not None:
+        model_config.train_indexer = True
+
+    if config.use_sparse_attn is not None:
+        model_config.use_sparse_attn = config.use_sparse_attn
+
     # Ensure pad_token_id is set (some models like Qwen3MoE don't have it).
     # In transformers v5, token IDs moved from PretrainedConfig to GenerationConfig.
     if not hasattr(model_config, "pad_token_id") or model_config.pad_token_id is None:
@@ -884,6 +916,8 @@ def _init_buffers_post_meta():
     load_dcp_start_time = time.perf_counter()
     state_dict = model.state_dict()
     state_dict = strip_lora_from_state_dict(state_dict)
+    if config.debug.random_init_indexer:
+        state_dict = strip_indexer_from_state_dict(state_dict)
     if model.config.tie_word_embeddings:
         del state_dict["lm_head.weight"]
     dcp_load(
@@ -897,11 +931,15 @@ def _init_buffers_post_meta():
 
     _move_buffers_to_cuda(model, config)
 
+    indexer_modules = [m for m in model.modules() if isinstance(m, Indexer)] if config.debug.random_init_indexer else []
+    if config.debug.random_init_indexer and not indexer_modules:
+        raise ValueError("debug.random_init_indexer=True but no Indexer modules found. Is this a DSA-capable model?")
     lora_modules = [m for m in model.modules() if hasattr(m, "_init_lora_parameters")]
-    if lora_modules:
+
+    if lora_modules or indexer_modules:
         generator: torch.Generator | None = None
         if parallel_dims.dp_replicate_enabled:
-            # Synchronize LoRA initialization across dp_replicate ranks by broadcasting a seed
+            # Synchronize random init across dp_replicate ranks by broadcasting a seed
             dp_replicate_mesh = parallel_dims.world_mesh["dp_replicate"]
             seed_tensor = torch.empty(1, dtype=torch.long, device="cuda")
             if dp_replicate_mesh.get_local_rank() == 0:
@@ -910,6 +948,8 @@ def _init_buffers_post_meta():
             generator = torch.Generator(device="cuda").manual_seed(seed_tensor.item())
         for module in lora_modules:
             module._init_lora_parameters(generator)
+        for module in indexer_modules:
+            module._init_indexer_parameters(generator)
     logger.debug(f"Loaded weights using HF DCP in {time.perf_counter() - load_dcp_start_time:.2f} seconds")
 
 
@@ -1178,10 +1218,15 @@ def setup_model(
     if config.moe_router_dtype == "float32":
         apply_fp32_moe_router(model)
 
-    # The DSA sparse-attention indexer runs its forward under torch.no_grad(), so it is
-    # never trainable. Freeze it so optimizer state stays symmetric across checkpoint
-    # save/resume. No-op for models without a sparse indexer.
-    freeze_sparse_indexer(model)
+    # No-op for models without a sparse indexer. See freeze_sparse_indexer's docstring for
+    # when to disable this (DSA conversion's indexer warm-up stage).
+    if config.freeze_sparse_indexer:
+        freeze_sparse_indexer(model)
+
+    # Overrides freeze_moe_router/freeze_sparse_indexer above: DSA conversion's indexer
+    # warm-up stage needs everything BUT the indexer frozen.
+    if config.freeze_all_except_indexer:
+        freeze_all_except_indexer(model)
 
     if config.debug.force_balanced_routing:
         apply_force_balanced_routing(model)

src/prime_rl/trainer/models/__init__.py

@@ -16,6 +16,7 @@
 from prime_rl.trainer.models.glm_moe_dsa import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
 from prime_rl.trainer.models.gpt_oss import GptOssConfig, GptOssForCausalLM
 from prime_rl.trainer.models.kimi_k2 import KimiK2Config, KimiK2ForCausalLM
+from prime_rl.trainer.models.kimi_k2_dsa import KimiK2DsaConfig, KimiK2DsaForCausalLM
 from prime_rl.trainer.models.laguna import LagunaConfig, LagunaForCausalLM
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, cast_float_and_contiguous
 from prime_rl.trainer.models.llama import LlamaForCausalLM
@@ -31,6 +32,7 @@
 AutoConfig.register("glm4_moe", Glm4MoeConfig, exist_ok=True)
 AutoConfig.register("glm_moe_dsa", GlmMoeDsaConfig, exist_ok=True)
 AutoConfig.register("kimi_k2", KimiK2Config, exist_ok=True)
+AutoConfig.register("kimi_k2_dsa", KimiK2DsaConfig, exist_ok=True)
 AutoConfig.register("laguna", LagunaConfig, exist_ok=True)
 AutoConfig.register("minimax_m2", MiniMaxM2Config, exist_ok=True)
 AutoConfig.register("nemotron_h", NemotronHConfig, exist_ok=True)
@@ -46,6 +48,7 @@
 _CUSTOM_CAUSAL_LM_MAPPING.register(Glm4MoeConfig, Glm4MoeForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(GlmMoeDsaConfig, GlmMoeDsaForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(KimiK2Config, KimiK2ForCausalLM, exist_ok=True)
+_CUSTOM_CAUSAL_LM_MAPPING.register(KimiK2DsaConfig, KimiK2DsaForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(LagunaConfig, LagunaForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(MiniMaxM2Config, MiniMaxM2ForCausalLM, exist_ok=True)
 _CUSTOM_CAUSAL_LM_MAPPING.register(NemotronHConfig, NemotronHForCausalLM, exist_ok=True)