h-guo18 9048d13b86 [Feat]:Support DPace (#1724) ...

### What does this PR do?

Type of change: New feature

Adds the **D-PACE** (Dynamic Position-Aware Cross-Entropy) loss
objective for DFlash speculative-decoding training
([arXiv:2605.18810](https://arxiv.org/abs/2605.18810)). It replaces the
static exponential position decay with per-position CE weights derived
from the draft's own confidence `q_i = exp(-CE_i)`: smoothed `q̃_i =
(1-α)q_i + α` (Eq.7) and weighted by the suffix-sum of prefix products
`w_j = Σ_{m≥j} ∏_{i≤m} q̃_i` (Eq.8), which directly targets expected
accepted block length and shifts signal toward whichever positions
currently limit acceptance.

Selected via `dflash_loss_objective` — **D-PACE is now the default**
(`dpace`); set `dflash_loss_objective: decay` to restore the previous
static schedule. Smoothing via `dflash_dpace_alpha` (default 0.5).
Weights are detached from the gradient — training-only, ~2.3% overhead,
no architecture or inference change. Mutually exclusive with
`dflash_loss_decay_factor`.

### Usage

```yaml
# DFlash recipe / training config
dflash:
  dflash_loss_objective: dpace   # default: decay
  dflash_dpace_alpha: 0.5        # smoothing in (0, 1]; stable in [0.3, 0.7]
```

### Testing

CPU unit tests in
`tests/unit/torch/speculative/plugins/test_hf_dflash.py`: weights match
the paper closed form, are detached and non-increasing, the α smoothing
floor keeps later weights non-zero, and convert wires/validates the new
fields (rejects bad objective and degenerate α). Training validated on
Qwen3-8B (curve below).

<img width="1803" height="809" alt="image"
src="https://github.com/user-attachments/assets/d34dcd76-9e46-4051-94d4-c880b1987965"
/>

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ⚠️ Behavior change — D-PACE is
now the **default** objective, so DFlash training loss weighting changes
unless you set `dflash_loss_objective=decay` (which reproduces the
previous static-decay behavior).
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A (no new
dependency)
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅
- Did you get Claude approval on this PR?: ❌

### Additional Information

Reference: D-PACE, [arXiv:2605.18810](https://arxiv.org/abs/2605.18810).
See `examples/speculative_decoding/doc/dflash.md` for the math and
tuning notes.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added a **D-PACE** training loss objective for DFlash speculative
decoding (`dflash_loss_objective: dpace`), configurable via
`dflash_dpace_alpha` (default `0.5`).
* **Documentation**
* Documented D-PACE’s confidence-derived, dynamically weighted
per-position loss behavior (training-only) and noted that
`dflash_loss_decay_factor` is ignored with D-PACE.
* **Bug Fixes**
* Updated DFlash loss to reuse the precomputed per-token cross-entropy
in the non-KD path.
* **Tests**
* Added unit tests for D-PACE weight correctness, masking, gradient
detachment, monotonicity, smoothing, and config/validation behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

2026-06-20 02:14:31 +00:00

..

dflash.md

[Feat]:Support DPace (#1724)

2026-06-20 02:14:31 +00:00