diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2907f626226..fc109c985e5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -37,6 +37,7 @@ Changelog - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. - The DFlash2, LiLiCorr, DSpark and Domino training forwards now return their metrics in one ``draft_metrics`` dict of detached 0-dim tensors (call ``.item()`` for a Python float), replacing ``selector_metrics``, ``lilicorr_metrics``, ``dspark_metrics`` and ``domino_metrics``. DSpark's ``ce_loss``/``l1_loss``/``base_accuracy`` are now ``loss_ce``/``loss_tvd``/``backbone_accuracy``, and Domino's ``final_loss``/``base_loss``/``base_accuracy`` are now ``loss_final``/``loss_backbone``/``backbone_accuracy``. +- DSpark and Domino training now honor ``dflash_loss_objective`` instead of always applying the static ``dflash_loss_decay_factor`` decay, so they train with D-PACE by default. Set ``dflash_loss_objective: decay`` to keep the previous weighting, as the shipped ``dspark.yaml`` and ``domino.yaml`` recipes now do, and finish runs started on an earlier version before upgrading: resuming one restores the objective saved in its checkpoint (``dpace`` unless the run set it) instead of the recipe's. **Deprecations** diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index a0c2cc34bb1..fa5c3e7f780 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -895,11 +895,6 @@ def _draft_logits(self, batch): logits = logits.view(*batch.anchors.shape, self.dflash_block_size, -1) return DraftLogits(final=logits, backbone=logits) - @property - def _position_weighting(self): - """How block positions are weighted in the loss: ``"dpace"`` or ``"decay"``.""" - return self.dflash_loss_objective - def _loss_context(self, batch): """The step's labels, the draft's logits and the position weights.""" targets = _block_targets( @@ -915,7 +910,7 @@ def _loss_context(self, batch): self._draft_logits(batch), targets, weighting={ - "scheme": self._position_weighting, + "scheme": self.dflash_loss_objective, "first": 1 - self._label_shift, "decay_factor": self.dflash_loss_decay_factor, "dpace_alpha": self.dflash_dpace_alpha, @@ -975,7 +970,7 @@ def forward( - Position weighting by D-PACE or exponential decay A variant changes the step only through its hooks: ``_online_base_outputs``, - ``_draft_logits`` (its head), ``_position_weighting`` and ``_compute_loss``. + ``_draft_logits`` (its head) and ``_compute_loss``. """ if not self.training: if self.dflash_offline: diff --git a/modelopt/torch/speculative/plugins/hf_domino.py b/modelopt/torch/speculative/plugins/hf_domino.py index 0a4bc9fdbd4..658ca266c37 100644 --- a/modelopt/torch/speculative/plugins/hf_domino.py +++ b/modelopt/torch/speculative/plugins/hf_domino.py @@ -138,11 +138,6 @@ def _build_objective(self, config): """CE on the labels, for the corrected logits and for the backbone's.""" return [objective_term("ce", 1.0)] - @property - def _position_weighting(self): - # dflash_loss_objective is not wired for Domino yet: it always decays. - return "decay" - def _online_base_outputs(self, input_ids, attention_mask, **kwargs): """The target hidden states; Domino does not use target-logit KD.""" with torch.no_grad(): @@ -181,7 +176,10 @@ def _draft_logits(self, batch): ) def _compute_loss(self, ctx): - """``(1 - lambda_base) * objective(final) + lambda_base * objective(backbone)``.""" + """``(1 - lambda_base) * objective(final) + lambda_base * objective(backbone)``. + + Both share the step's position weights; D-PACE reads them off the served logits. + """ out = super()._compute_loss(ctx) lambda_base = self._current_lambda_base() loss = backbone_loss = out.loss diff --git a/modelopt/torch/speculative/plugins/hf_dspark.py b/modelopt/torch/speculative/plugins/hf_dspark.py index 69b5fdc3f1f..34a7391162f 100644 --- a/modelopt/torch/speculative/plugins/hf_dspark.py +++ b/modelopt/torch/speculative/plugins/hf_dspark.py @@ -119,11 +119,6 @@ def _build_objective(self, config): objective_term("tvd", config.dflash_l1_loss_alpha), ] - @property - def _position_weighting(self): - # dflash_loss_objective is not wired for DSpark yet: it always decays. - return "decay" - def _online_base_outputs(self, input_ids, attention_mask, **kwargs): """The target hidden states, and the base's final hidden for the loss to project.""" # The inner model, not the CausalLM: the loss projects only the rows it reads. diff --git a/modelopt_recipes/general/speculative_decoding/domino.yaml b/modelopt_recipes/general/speculative_decoding/domino.yaml index e7237b11b8f..055f78508ad 100644 --- a/modelopt_recipes/general/speculative_decoding/domino.yaml +++ b/modelopt_recipes/general/speculative_decoding/domino.yaml @@ -64,7 +64,9 @@ dflash: dflash_use_torch_compile: false # Domino does not use target-logit KD; it trains its own base/final CE losses. dflash_self_logit_distillation: false - # gamma for exponential loss decay (block_size=16 -> 7). + # Static exponential position decay, gamma=7 for block_size=16: the weighting this recipe + # was tuned with. The default objective, 'dpace', would ignore the decay factor. + dflash_loss_objective: decay dflash_loss_decay_factor: 7.0 # Qwen3 has no native mask token; 151669 is an unused id used by the reference. dflash_mask_token_id: 151669 diff --git a/modelopt_recipes/general/speculative_decoding/dspark.yaml b/modelopt_recipes/general/speculative_decoding/dspark.yaml index db69f3c4b9e..084f7145b24 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark.yaml +++ b/modelopt_recipes/general/speculative_decoding/dspark.yaml @@ -68,7 +68,9 @@ dflash: # DSpark computes the target distribution internally for its TVD/confidence # terms; the DFlash KD path is unused. dflash_self_logit_distillation: false - # gamma for exponential loss decay (block_size=16 -> 7). + # Static exponential position decay, gamma=7 for block_size=16: the weighting this recipe + # was tuned with. The default objective, 'dpace', would ignore the decay factor. + dflash_loss_objective: decay dflash_loss_decay_factor: 7.0 # Qwen3 has no native mask token; 151669 is an unused id used by the reference. dflash_mask_token_id: 151669 diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md index 8494f6f4b76..074b2420d4b 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/domino.md @@ -52,7 +52,8 @@ Everything in `dflash.md` applies. Domino adds: `dflash_self_logit_distillation` is **false** — Domino trains its own base/final CE losses rather than distilling target logits. Recipe defaults also differ from DFlash's: `block_size` 16, `num_anchors` 256, `num_train_epochs` 6, -`training_seq_len` 3072, `warmup_ratio` 0.04, `dflash_loss_decay_factor` 7.0. +`training_seq_len` 3072, `warmup_ratio` 0.04, and the static position decay +(`dflash_loss_objective: decay`, `dflash_loss_decay_factor` 7.0) instead of D-PACE. `max_grad_norm: 1.0` is stated explicitly in the recipe but is *not* a delta — it is the `transformers.TrainingArguments` default that `dflash.yaml` inherits by not setting it. @@ -101,5 +102,6 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i | `shift_label=False is not supported: projector_type 'domino' trains only next-token alignment` | `shift_label: false` | Leave it at `true` | | `pure_draft_prefix_len must be in [0, N] (block_size=B), got X` | Prefix length >= block size | Lower it below `block_size` | | `DominoLambdaCallback: state.max_steps unset (<=0); lambda_base curriculum disabled` (warning) | `training.max_steps` not set | Set `training.max_steps`, else the curriculum never runs | +| `dflash_loss_decay_factor=7.0 is ignored when dflash_loss_objective='dpace'` (warning) when resuming a run started on an earlier ModelOpt version | That version saved the unset objective as `dpace` but trained with the static decay; resuming restores the saved `dpace`, not the recipe's `decay` pin | Finish the run on the version that started it | | DDP error about unused parameters | `ddp_find_unused_parameters` turned off | Keep it `true` | | Draft trains but quality is poor | Draft dims left at `Qwen3Config` defaults instead of matching the base | Set the GQA/FFN dims explicitly | diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md index bb2c6992ac5..36f2e7849d9 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dspark.md @@ -91,7 +91,8 @@ Total loss is `ce_alpha*CE + l1_alpha*TVD + conf_alpha*confidence_BCE`. `dflash_self_logit_distillation` is **false** for DSpark — it computes the target distribution internally for the TVD and confidence terms, so the DFlash KD path is unused. Recipe defaults also differ from DFlash's: `block_size` 16, `num_anchors` 256, -`num_train_epochs` 6, `training_seq_len` 3072, `warmup_ratio` 0.04. +`num_train_epochs` 6, `training_seq_len` 3072, `warmup_ratio` 0.04, and the static position +decay (`dflash_loss_objective: decay`, `dflash_loss_decay_factor` 7.0) instead of D-PACE. ## Per-model adjustments @@ -149,6 +150,7 @@ failures (`seq_len` divisibility, offline eval, mask token, chat template) are i | `shift_label=False is not supported: projector_type 'dspark' trains only next-token alignment` | `dflash_architecture_config.shift_label: false` | Remove it; DSpark only trains next-token alignment | | `dflash_confidence_head_alpha > 0 but the confidence head was not built` | Loss term enabled without the head | Set `dflash_architecture_config.use_confidence_head=true`, or set the alpha to 0 | | `An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate` | Offline/streaming conversion deleted base layers | Keep `estimate_ar=false` and `ar_validate_steps=0`; evaluate after export | +| `dflash_loss_decay_factor=7.0 is ignored when dflash_loss_objective='dpace'` (warning) when resuming a run started on an earlier ModelOpt version | That version saved the unset objective as `dpace` but trained with the static decay; resuming restores the saved `dpace`, not the recipe's `decay` pin | Finish the run on the version that started it | | Draft trains but acceptance length is poor | Draft dims left at defaults instead of matching the base | Set the GQA/FFN dims explicitly (see *Per-model adjustments*) | | Acceptance length capped despite clean training (streaming) | `EAGLE_CAPTURE_IDS` final layer wrong, or the vLLM aux-capture fix (vllm#46788) missing | Correct the ids; use a container with the fix | | Trainer init crash on a serve container | tensorboard absent | `training.report_to=none` | diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index ec86aa115c7..84793ae4442 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -1615,6 +1615,51 @@ def test_next_token_variants_refuse_shift_label_false_at_convert(head_fields): mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)]) +@pytest.mark.parametrize("objective", ["dpace", "decay"]) +@pytest.mark.parametrize( + ("head_fields", "metric", "source"), + [ + ({"projector_type": "dspark", "markov_rank": 4}, "loss_ce", "final"), + ( + {"projector_type": "domino", "emb_dim": 8, "gru_hidden_dim": 8}, + "loss_backbone", + "backbone", + ), + ], + ids=["dspark", "domino"], +) +def test_next_token_variants_weight_positions_by_the_objective( + objective, head_fields, metric, source +): + """Weights follow dflash_loss_objective from block position 0; D-PACE reads the served logits. + + Every objective term is reduced under them, Domino's backbone CE included. + """ + config = get_dflash_config() + config["dflash_architecture_config"].update(head_fields) + config.update(dflash_loss_objective=objective, dflash_loss_decay_factor=4.0) + model = get_tiny_llama(num_hidden_layers=4) + mtsp.convert(model, [("dflash", config)]) + model.float().train() + with patch.object(model, "_compute_loss", wraps=model._compute_loss) as compute_loss: + out = model(**_dflash_batch(model.config.vocab_size)) + ctx = compute_loss.call_args.args[0] + assert not torch.equal(ctx.logits.final, ctx.logits.backbone) + + def ce(logits): + return F.cross_entropy( + logits.flatten(0, 2), ctx.targets.token_ids.flatten(), reduction="none" + ) + + mask = ctx.targets.mask + weights = _block_position_weights( + mask, objective, 0, 4.0, 0.5, target_ce=lambda: ce(ctx.logits.final).view_as(mask) + ) + assert torch.equal(ctx.weights, weights) + want = (ce(getattr(ctx.logits, source)) * weights.flatten()).sum() / (weights.sum() + 1e-6) + assert torch.equal(out.draft_metrics[metric], want) + + class TestTeacherLogits: """The base distribution is projected only at the rows, and vocab entries, asked for. diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml index 7adf3400cd7..8760c2047af 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_domino.yaml @@ -63,6 +63,7 @@ pipeline: # Domino knobs (also set in the recipe; repeated here for visibility). - dflash.dflash_block_size=16 - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_objective=decay - dflash.dflash_loss_decay_factor=7 - dflash.dflash_mask_token_id=151669 - dflash.dflash_self_logit_distillation=false