diff --git a/docs/basic_usage/training.md b/docs/basic_usage/training.md index c67b394c0..0f474ad51 100644 --- a/docs/basic_usage/training.md +++ b/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`. diff --git a/examples/configs/README.md b/examples/configs/README.md index 402fd3e70..74dde9e05 100644 --- a/examples/configs/README.md +++ b/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`) | diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index ca65c6cb9..4d9ab1810 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/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. diff --git a/specforge/algorithms/dflash/providers.py b/specforge/algorithms/dflash/providers.py index 7dd3b216f..7185f0751 100644 --- a/specforge/algorithms/dflash/providers.py +++ b/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 diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index b50c07262..1d5855132 100644 --- a/specforge/algorithms/model_providers.py +++ b/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, diff --git a/specforge/config/schema.py b/specforge/config/schema.py index d1173273e..85bf54748 100644 --- a/specforge/config/schema.py +++ b/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 diff --git a/tests/test_config/test_schema.py b/tests/test_config/test_schema.py index 9e5aa1233..26c2d5801 100644 --- a/tests/test_config/test_schema.py +++ b/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), diff --git a/tests/test_modeling/test_dflash2.py b/tests/test_modeling/test_dflash2.py index 6140886f1..6969e5bba 100644 --- a/tests/test_modeling/test_dflash2.py +++ b/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):