feat(dflash2): add selector stop-gradient control

docs/basic_usage/training.md

@@ -201,6 +201,11 @@ ramps it linearly to the configured value. Both schedule ratios default to
 zero. A newly initialized selector starts as a unary no-op, so enabling DFlash
 2 does not perturb the initial DFlash proposal scores.
 
+Set `training.dflash2_selector_stop_gradient: true` to keep the selector CE
+from updating the unary logits and draft hidden states. The selector parameters
+still train, and the primary DFlash/D-PACE/LK objective keeps its normal draft
+gradient path. The option defaults to `false`, preserving coupled training.
+
 The exported computation and parameter names match the public SGLang DFlash 2
 contract, including optional `output_multiplier` and
 `final_logit_softcapping` transforms from `dflash_config`.

examples/configs/README.md

@@ -280,7 +280,7 @@ Strategy-specific fields should be written only when tuning that objective:
 | Strategy | Fields and defaults |
 | --- | --- |
 | EAGLE3 | `training.ttt_length` (`7`), `training.lk_loss_type` (`null`; `lambda`, `alpha`, or `tv`), `training.kl_scale` (`1.0`), `training.kl_decay` (`1.0`) |
-| DFlash / DFlash 2 / Domino / D-PACE | `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`), `training.objective_chunk_blocks` (`128`; `0` materializes all objective logits), `training.loss_type` (`dflash`; fixed decay, or `dpace`; dynamic weighting), DFlash/DFlash 2's `training.lk_loss_type` (`null`; CE, `lambda`, `alpha`, or `tv`), `training.kl_scale` (`1.0`), `training.kl_decay` (`1.0`), DFlash 2's CE selector objective controls `training.dflash2_selector_loss_alpha` (`1.0`), `training.dflash2_selector_warmup_ratio` (`0.0`), and `training.dflash2_selector_ramp_ratio` (`0.0`), `training.dpace_alpha` (`0.5`), `training.lambda_base_start` (`1.0`), `training.lambda_base_decay_ratio` (`0.5`) |
+| DFlash / DFlash 2 / Domino / D-PACE | `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`), `training.objective_chunk_blocks` (`128`; `0` materializes all objective logits), `training.loss_type` (`dflash`; fixed decay, or `dpace`; dynamic weighting), DFlash/DFlash 2's `training.lk_loss_type` (`null`; CE, `lambda`, `alpha`, or `tv`), `training.kl_scale` (`1.0`), `training.kl_decay` (`1.0`), DFlash 2's CE selector objective controls `training.dflash2_selector_loss_alpha` (`1.0`), `training.dflash2_selector_warmup_ratio` (`0.0`), `training.dflash2_selector_ramp_ratio` (`0.0`), and `training.dflash2_selector_stop_gradient` (`false`), `training.dpace_alpha` (`0.5`), `training.lambda_base_start` (`1.0`), `training.lambda_base_decay_ratio` (`0.5`) |
 | DSpark | Token-pooled objective with valid-first-target anchors and distributed ratio telemetry. Configure the shared `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`; production recipes use `4.0`), and `training.objective_chunk_blocks` (`128`; `0` materializes all objective logits), plus `training.dspark_ce_loss_alpha` (`0.1`), `training.dspark_l1_loss_alpha` (`0.9`), and `training.dspark_confidence_head_alpha` (`1.0`). |
 | P-EAGLE | `training.num_depths` (`8`), `training.down_sample_ratio` (`0.8`), `training.down_sample_ratio_min` (`0.2`), `training.norm_before_residual` (`null`) |
 

specforge/algorithms/common/dflash_family_model.py

@@ -212,6 +212,7 @@ def __init__(
         selector_loss_alpha: float = 1.0,
         selector_warmup_ratio: float = 0.0,
         selector_ramp_ratio: float = 0.0,
+        selector_stop_gradient: bool = False,
         lk_loss_type: Optional[str] = None,
         kl_scale: float = 1.0,
         kl_decay: float = 1.0,
@@ -250,6 +251,7 @@ def __init__(
         self.selector_loss_alpha = float(selector_loss_alpha)
         self.selector_warmup_ratio = float(selector_warmup_ratio)
         self.selector_ramp_ratio = float(selector_ramp_ratio)
+        self.selector_stop_gradient = bool(selector_stop_gradient)
         self.lk_loss_type = lk_loss_type
         self.kl_scale = float(kl_scale)
         self.kl_decay = float(kl_decay)
@@ -484,6 +486,12 @@ def _selector_chunk_terms(
         ``checkpointed_chunk_reduce``'s flat tuple contract.
         """
 
+        if self.selector_stop_gradient:
+            # Isolate only the selector objective. The caller still uses the
+            # original tensors for the primary DFlash/D-PACE/LK objective.
+            objective_logits = objective_logits.detach()
+            hidden = hidden.detach()
+
         # Match serving exactly: train only against the strict unary top-k.
         # Candidate misses are a backbone/recall failure, not a selector
         # classification example, so they carry no selector gradient.

specforge/algorithms/dflash/providers.py

@@ -87,6 +87,9 @@ def resume_contract(_config, draft_model, training_model):
                 "dflash2_selector_ramp_ratio": float(
                     training_model.selector_ramp_ratio
                 ),
+                "dflash2_selector_stop_gradient": bool(
+                    training_model.selector_stop_gradient
+                ),
             }
         )
     return contract

specforge/algorithms/model_providers.py

@@ -394,6 +394,7 @@ def build_dflash_model(
             selector_loss_alpha=cfg.training.dflash2_selector_loss_alpha,
             selector_warmup_ratio=cfg.training.dflash2_selector_warmup_ratio,
             selector_ramp_ratio=cfg.training.dflash2_selector_ramp_ratio,
+            selector_stop_gradient=cfg.training.dflash2_selector_stop_gradient,
             lk_loss_type=cfg.training.lk_loss_type,
             kl_scale=cfg.training.kl_scale,
             kl_decay=cfg.training.kl_decay,

specforge/config/schema.py

@@ -575,6 +575,9 @@ class TrainingConfig(StrictConfigModel):
     dflash2_selector_warmup_ratio: float = Field(default=0.0, ge=0.0, le=1.0)
     #: Fraction of optimizer steps used to ramp the selector weight to its target.
     dflash2_selector_ramp_ratio: float = Field(default=0.0, ge=0.0, le=1.0)
+    #: Stop selector gradients at the unary/backbone boundary while preserving
+    #: the primary DFlash/D-PACE/LK gradient path.
+    dflash2_selector_stop_gradient: bool = False
     lambda_base_start: float = 1.0
     lambda_base_decay_ratio: float = 0.5
     dspark_ce_loss_alpha: float = 0.1

tests/test_config/test_schema.py

@@ -340,10 +340,15 @@ def test_dflash2_selector_objective_settings_are_bounded(self):
         payload["training"]["dflash2_selector_loss_alpha"] = 0.25
         payload["training"]["dflash2_selector_warmup_ratio"] = 0.1
         payload["training"]["dflash2_selector_ramp_ratio"] = 0.2
+        payload["training"]["dflash2_selector_stop_gradient"] = True
         config = Config.model_validate(payload)
         self.assertEqual(config.training.dflash2_selector_loss_alpha, 0.25)
         self.assertEqual(config.training.dflash2_selector_warmup_ratio, 0.1)
         self.assertEqual(config.training.dflash2_selector_ramp_ratio, 0.2)
+        self.assertTrue(config.training.dflash2_selector_stop_gradient)
+
+        default_config = Config.model_validate(_online_payload("dflash"))
+        self.assertFalse(default_config.training.dflash2_selector_stop_gradient)
 
         for field, invalid in (
             ("dflash2_selector_loss_alpha", -0.1),

tests/test_modeling/test_dflash2.py

@@ -176,6 +176,7 @@ def test_resume_contract_tracks_dflash2_specific_semantics(self):
             num_anchors=8,
             selector_loss_alpha=0.75,
             selector_ramp_ratio=0.2,
+            selector_stop_gradient=True,
             selector_warmup_ratio=0.1,
         )
 
@@ -188,6 +189,7 @@ def test_resume_contract_tracks_dflash2_specific_semantics(self):
         self.assertEqual(contract["dflash2_selector_loss_alpha"], 0.75)
         self.assertEqual(contract["dflash2_selector_warmup_ratio"], 0.1)
         self.assertEqual(contract["dflash2_selector_ramp_ratio"], 0.2)
+        self.assertTrue(contract["dflash2_selector_stop_gradient"])
         self.assertEqual(contract["dflash_lk_loss_type"], "lambda")
 
     def test_backward_reaches_convolution_parameters(self):
@@ -641,6 +643,57 @@ def transform_unary_logits(logits):
             self.assertIsNotNone(parameter.grad)
             self.assertGreater(parameter.grad.abs().sum().item(), 0.0)
 
+    def test_selector_stop_gradient_isolates_backbone_inputs(self):
+        class Draft(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.candidate_selector = CandidateSelector(
+                    hidden_size=4,
+                    vocab_size=4,
+                    state_rank=2,
+                    top_k=2,
+                    initializer_range=0.2,
+                )
+
+            @staticmethod
+            def transform_unary_logits(logits):
+                return logits.float()
+
+        draft = Draft()
+        with torch.no_grad():
+            draft.candidate_selector.successor_codebook.normal_(std=0.2)
+        model = OnlineDFlashModel(
+            draft_model=draft,
+            target_lm_head=nn.Identity(),
+            target_embed_tokens=nn.Embedding(4, 4),
+            mask_token_id=3,
+            block_size=2,
+            attention_backend="eager",
+            selector_stop_gradient=True,
+        )
+        hidden = torch.tensor(
+            [[[[0.0, 0.0, 0.0, 0.0], [0.0, 3.0, 2.0, 1.0]]]],
+            requires_grad=True,
+        )
+        terms = model._dflash_objective_chunk_terms(
+            hidden,
+            torch.tensor([[[0, 2]]]),
+            torch.tensor([[[0.0, 1.0]]]),
+            torch.tensor([[[0, 1]]]),
+        )
+        (terms.ce_loss_num + terms.selector_ce_num).backward()
+
+        reference = hidden.detach().clone().requires_grad_(True)
+        base_ce = torch.nn.functional.cross_entropy(
+            reference[0, 0, 1].unsqueeze(0),
+            torch.tensor([2]),
+        )
+        base_ce.backward()
+        torch.testing.assert_close(hidden.grad, reference.grad)
+        for parameter in draft.candidate_selector.parameters():
+            self.assertIsNotNone(parameter.grad)
+            self.assertGreater(parameter.grad.abs().sum().item(), 0.0)
+
     def test_zero_effective_selector_alpha_keeps_selector_in_autograd_graph(self):
         class Draft(nn.Module):
             def __init__(self):