Add D-PACE loss implementation for D-Flash training by weifanjiang · Pull Request #736 · vllm-project/speculators

weifanjiang · 2026-07-07T15:14:02Z

Purpose

This PR implements D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting. D-PACE proposes an alternative per-position loss weighting for training block-parallel draft models. Instead of applying a fixed exponential decay over block positions, D-PACE dynamically adjusts the weights as training progresses. This is based on the intuition that although earlier positions are naturally more important (i.e., an early rejection forces all subsequent draft positions to be discarded), when they are sufficiently stabilized, the training priority should shift to later positions to improve overall acceptance rate.

Two arguments are added to train.py to support D-PACE loss:

--per-position-loss-weight ["fixed-exp-decay", "dpace"] (default fixed-exp-decay). This flag enables D-PACE loss; while missing, it defaults to the original fixed exponential decay.
--dpace-alpha (default 0.5). This flag is only effective when D-PACE is enabled. It smoothes per-position confidence according to equation 7 of the paper. The paper uses 0.5 as the default value for their evaluation and justifies this choice in sensitivity analysis.

Note that D-PACE requires cross-entropy as the training loss (--loss-fn ce), allowing per-position confidence score computation to be piggybacked on the loss computation. Applying D-PACE to other losses is not supported and will result in ValueError.

D-PACE is currently enabled for D-Flash and D-Spark draft models. Theoretically, it is compatible with all block-parallel drafters.

Tests

D-Flash + D-PACE evaluation. Multiple D-Flash checkpoints are trained with different loss function configurations:

--loss-fn kld: KL Divergence (current default)
--loss-fn ce: Cross Entropy with fixed exponential decay
--per-position-loss-weight dpace --dpace-alpha ALPHA: D-PACE with ALPHA = 0.3, 0.5, 0.7.

The training set is 100K random samples from Magpie+Ultrachat. The verifier model is Qwen3-8B. Each checkpoint uses 1 A100 for hidden state generation and 7 A100s for training. Training configurations are identical for all checkpoints:

--target-layer-ids 1 9 17 25 34  --num-layers 5 --draft-vocab-size 32000 --max-anchors 3072 --block-size 8 --sliding-window 2048 --sliding-window-indices 0 1 2 3 4 --scheduler-type cosine --lr 0.0006 --epochs 1 --noise-std 0.0

Metric: since D-PACE is a pure training-time optimization and does not change draft model's forward pass complexity, its gains are captured entirely by acceptance rates.

Evaluation results: Mean-Accepted-Length (MAL) on Speculator Benchmarks:

Dataset	KL Divergence	CE	D-PACE $\alpha=0.3$	D-PACE $\alpha=0.5$	D-PACE $\alpha=0.7$
HumanEval	2.977	2.956	3.047	3.169	3.139
math_reasoning	3.275	3.322	3.541	3.532	3.531
qa	2.279	2.343	2.455	2.495	2.431
question	2.585	2.567	2.685	2.680	2.708
rag	2.302	2.269	2.393	2.372	2.336
summarization	1.906	1.918	1.981	1.980	1.991
tool_call	2.372	2.327	2.480	2.454	2.467
translation	2.096	2.072	2.176	2.183	2.191
writing	2.565	2.586	2.683	2.684	2.718
Mean	2.484	2.485	2.605	2.617	2.612

D-Spark + D-PACE evaluation. D-Spark checkpoints use Markov rank 256 and vanilla Markov head type. The confidence head is disabled. Remaining training configurations match above D-Flash experiments.

--loss-fn ce: Vanilla fixed-decay cross-entropy
--position-loss-weight dpace --dpace-alpha 0.5: D-PACE with $\alpha=0.5$

MALs:

Dataset	CE	D-PACE $\alpha=0.5$
HumanEval	3.145	3.302
math_reasoning	3.457	3.615
qa	2.352	2.492
question	2.671	2.770
rag	2.276	2.384
summarization	1.916	1.996
tool_call	2.451	2.538
translation	2.082	2.202
writing	2.686	2.776
Mean	2.560	2.675

D-PACE improves upon fixed-decay loss for both D-Flash and D-Spark drafters.

Checklist

I have filled in:

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan/results, such as providing test command and pasting the results.
(Optional) The necessary documentation update.
I (a human) have written or reviewed the code in this pr to the best of my ability.

coderabbitai · 2026-07-07T15:16:03Z

📝 Walkthrough

Walkthrough

Adds D-PACE (per-position confidence-based) loss weighting as an alternative to fixed-exponential decay. Introduces CLI flags, extends DFlashDraftModel's forward and get_trainer_kwargs to propagate the new configuration, adds branching logic in compute_metrics, and adds a new dpace_loss_weight helper plus optional precomputed CE support in compound_loss and loss_function.

Changes

D-PACE loss weighting

Layer / File(s)	Summary
dpace_loss_weight helper and loss function plumbing `src/speculators/models/metrics.py`	Adds `dpace_loss_weight` to compute masked per-position weights from per-token NLL, and extends `compound_loss`/`loss_function` with an optional `dpace_precomputed_ce` parameter to reuse precomputed elementwise CE.
dflash compute_metrics D-PACE branch `src/speculators/models/dflash/metrics.py`	Adds `per_position_loss_weight` and `dpace_alpha` parameters to `compute_metrics`, imports `ce_loss`/`dpace_loss_weight`, and branches `decay_fn` selection between the new D-PACE weighting and the existing fixed-exp-decay behavior.
DFlashDraftModel forward/trainer kwargs propagation `src/speculators/models/dflash/core.py`	Extends `forward` with `per_position_loss_weight` and `dpace_alpha` parameters forwarded into `compute_metrics`, and updates `get_trainer_kwargs` to propagate these values into train/validation kwargs.
CLI flags for D-PACE configuration `scripts/train.py`	Adds `--per-position-loss-weight` and `--dpace-alpha` CLI arguments to `parse_args()`.

Possibly related PRs

vllm-project/speculators#542: Also extends DFlashDraftModel.forward and dflash/metrics.compute_metrics signatures for configurable loss computation, overlapping the same call sites modified here.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: adding D-PACE loss support for D-Flash training.
Description check	✅ Passed	The description is directly related to the changeset and explains the D-PACE training update and tests.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (3)

src/speculators/models/metrics.py (2)
281-320: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Validate alpha range for numerical stability.

Confirmed against the D-PACE paper (arXiv:2605.18810) that this implementation matches the published smoothing/prefix-product/suffix-sum weighting scheme. However, the paper notes the smoothing parameter is only meaningful for a lower bound preventing the cumulative product in the weight from vanishing, and specifically that at α = 0, where the cumulative product vanishes and τ stalls near 1.0. There's no validation here that alpha stays within [0, 1]; an out-of-range value (e.g., negative or >1, or the degenerate alpha=0) would silently produce meaningless/unstable weights instead of failing fast.
🛡️ Proposed fix
     with torch.no_grad():
+        if not 0.0 < alpha <= 1.0:
+            raise ValueError(f"alpha must be in (0, 1], got {alpha}")
         # convert CE to per-position confidence
         q = torch.exp(-neg_log_q).float()
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/speculators/models/metrics.py` around lines 281 - 320, Validate the
smoothing parameter in dpace_loss_weight before computing the prefix product:
add an explicit check that alpha is within the supported range, rejecting
out-of-range values and the degenerate zero case with a clear failure. Keep the
existing weighting logic unchanged, but ensure the function fails fast for
invalid alpha so the numerical stability of the D-PACE weighting is preserved.
300-302: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value

Reshape relies on exact block-size divisibility.

num_anchors = q.shape[1] // block_size uses floor division; if q.shape[1] isn't an exact multiple of block_size, the subsequent .reshape(num_anchors, block_size) calls will raise a runtime error rather than a clear, actionable message. In the current call path (compute_metrics in dflash/metrics.py), seq_len is always num_anchors * block_size by construction, so this shouldn't trigger today, but a defensive assert would make future misuse fail with a clearer error.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/speculators/models/metrics.py` around lines 300 - 302, The reshape logic
in the metrics path assumes q.shape[1] is exactly divisible by block_size, so
add a defensive validation before computing num_anchors in the relevant reshape
block. Update the code around the q.reshape and loss_mask.reshape operations to
assert or explicitly check divisibility and raise a clear error message if the
sequence length is not compatible with block_size, using the existing
compute_metrics / metrics reshape flow as the location to safeguard.
scripts/train.py (1)
964-976: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Consider validating --dpace-alpha at parse time.

Per the D-PACE paper, alpha is only meaningful in [0, 1], with alpha=0 causing the weighting to degenerate. Argparse currently accepts any float unchecked. Adding a lightweight validator here (or in dpace_loss_weight, see comment in metrics.py) would catch a misconfigured training run early rather than after silently training with a degenerate loss weighting.

As per path instructions, "Check that scripts handle argument parsing robustly."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/train.py` around lines 964 - 976, Add parse-time validation for the
D-PACE smoothing argument so invalid values are rejected early. In the argument
setup around parser.add_argument for --dpace-alpha, enforce that the value is
within the valid [0, 1] range and disallow degenerate settings such as 0, either
via an argparse type/validator or a small reusable check used by
dpace_loss_weight. Keep the validation close to the existing D-Pace-specific
arguments so misconfigured runs fail immediately with a clear message.
Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/speculators/models/dflash/metrics.py`:
- Around line 55-64: Replace the inline lambda assigned to decay_fn in the
dflash metrics logic with a named nested function to satisfy Ruff E731 and
improve readability. In the per_position_loss_weight == "dpace" branch, define a
def that returns decay_mult instead of using lambda, and keep the existing
partial(dflash_loss_decay, gamma=gamma) path unchanged. Use a distinct,
descriptive function name near ce_loss, dpace_loss_weight, and decay_fn so the
intent is clear and no parameter name conflicts with the outer pos_idx.

---

Nitpick comments:
In `@scripts/train.py`:
- Around line 964-976: Add parse-time validation for the D-PACE smoothing
argument so invalid values are rejected early. In the argument setup around
parser.add_argument for --dpace-alpha, enforce that the value is within the
valid [0, 1] range and disallow degenerate settings such as 0, either via an
argparse type/validator or a small reusable check used by dpace_loss_weight.
Keep the validation close to the existing D-Pace-specific arguments so
misconfigured runs fail immediately with a clear message.

In `@src/speculators/models/metrics.py`:
- Around line 281-320: Validate the smoothing parameter in dpace_loss_weight
before computing the prefix product: add an explicit check that alpha is within
the supported range, rejecting out-of-range values and the degenerate zero case
with a clear failure. Keep the existing weighting logic unchanged, but ensure
the function fails fast for invalid alpha so the numerical stability of the
D-PACE weighting is preserved.
- Around line 300-302: The reshape logic in the metrics path assumes q.shape[1]
is exactly divisible by block_size, so add a defensive validation before
computing num_anchors in the relevant reshape block. Update the code around the
q.reshape and loss_mask.reshape operations to assert or explicitly check
divisibility and raise a clear error message if the sequence length is not
compatible with block_size, using the existing compute_metrics / metrics reshape
flow as the location to safeguard.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: d76c259c-9738-4941-87eb-0fba2af483e2

📥 Commits

Reviewing files that changed from the base of the PR and between 2d97b17 and 7b26944.

📒 Files selected for processing (4)

scripts/train.py
src/speculators/models/dflash/core.py
src/speculators/models/dflash/metrics.py
src/speculators/models/metrics.py

mergify · 2026-07-07T15:51:18Z

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

Signed-off-by: Weifan Jiang <weijiang@redhat.com>

fynnsu

Hi @weifanjiang thanks for the pr! I think we need to change how we're piping some things through speculators so that it fits within the existing abstractions.

Signed-off-by: Weifan Jiang <weijiang@redhat.com>

…c schema Make scripts/train.py config-file-first without changing what any run does and without touching the model layer. - New src/speculators/train/config.py is the single source of truth: ~75 tunables as typed fields on cohesive pydantic-settings groups (verifier/draft/data/ generation/loss/optimizer/scheduler/trainer/logging + per-algo dflash/dspark/ peagle/mtp). The argparse CLI is GENERATED from these fields, so adding an argument is one typed field -- train.py has no per-flag add_argument. - --config run.yaml supplies any argument; precedence is CLI > YAML > coded defaults; --config is optional (no file == the old pure-CLI behaviour). - --dump-config prints the fully-resolved grouped YAML (a valid --config file); the resolved run.yaml is also saved next to each checkpoint for reproducibility. - Provenance = the argparse.SUPPRESS'd namespace (not model_fields_set, which the nested partial-update machinery pollutes). A YAML leaf counts as provided only under its owning group block; a misplaced key is reported as unrecognised and kept out of provenance so it can't mis-fire the draft-init conflict check. - ValidationError is routed through parser.error (argparse usage, exit code 2). Phase 1: the resolved config is flattened back to the vars(args)-shaped dict the five SpeculatorModel classes consume via **kwargs, so the model layer and Trainer are untouched; only main()'s TrainerConfig reads the typed groups. Pushing typed configs into the model layer is deferred to a follow-up PR. Includes the D-PACE loss knobs from #736 (--per-position-loss-weight, --dpace-alpha) as typed DFlash fields, with the "dpace requires ce loss / alpha in (0,1]" check ported to a cross-group model_validator. Equivalence is guarded by tests/unit/train/test_golden_cli_equivalence.py: the refactored flat working-dict is asserted byte-value-identical to the pre-refactor parser across a 15-invocation matrix. Expected values live in one golden_flat_dict.json (a _baseline all-defaults dict + per-invocation deltas). Removes src/speculators/utils/argparse_utils.py and its test: train.py was the only caller of explicitly_provided_dests, which the SUPPRESS'd-namespace provenance replaces. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Rahul-Tuli <rtuli@redhat.com>

…c schema Make scripts/train.py config-file-first without changing what any run does and without touching the model layer. - New src/speculators/train/config.py is the single source of truth: ~75 tunables as typed fields on cohesive pydantic-settings groups (verifier/draft/data/ generation/loss/optimizer/scheduler/trainer/logging + per-algo dflash/dspark/ peagle/mtp). The argparse CLI is GENERATED from these fields, so adding an argument is one typed field -- train.py has no per-flag add_argument. - --config run.yaml supplies any argument; precedence is CLI > YAML > coded defaults; --config is optional (no file == the old pure-CLI behaviour). - --dump-config prints the fully-resolved grouped YAML (a valid --config file); the resolved run.yaml is also saved next to each checkpoint for reproducibility. It runs the draft-init contract first, so a scaffold is never emitted for a config that would fail to load. - Provenance = the argparse.SUPPRESS'd namespace (not model_fields_set, which the nested partial-update machinery pollutes). A YAML leaf counts as provided only under its owning group block; a misplaced key is reported as unrecognised and kept out of provenance so it can't mis-fire the draft-init conflict check. - Every bad input is a clean argparse error (exit code 2, no traceback): schema ValidationError and a broken --config file (missing/unreadable/malformed/ non-mapping) alike; a missing required field names the flag to set. Phase 1: the resolved config is flattened back to the vars(args)-shaped dict the five SpeculatorModel classes consume via **kwargs, so the model layer and Trainer are untouched; only main()'s TrainerConfig reads the typed groups. Pushing typed configs into the model layer is deferred to a follow-up PR. Rebased onto current main. Absorbs the flags added since the branch point and ports them into the schema so behaviour is unchanged: - #736 D-PACE: --per-position-loss-weight, --dpace-alpha as typed DFlash fields, with the "dpace requires ce loss / alpha in (0,1]" check as a cross-group model_validator. - #651 scheduler: --scheduler-warmup-ratio, and --scheduler-type as a Literal (linear/cosine/none), threaded through the typed TrainerConfig build. - #749 default sliding-window attention: --sliding-window-indices renamed to --full-attention-indices (inverted semantics), matching upstream. The decoder-shaping flag set (used by the draft-init conflict check) is derived from the schema via a field tag, so it is not a second place to edit when a shaping knob is added. Equivalence is guarded by tests/unit/train/test_golden_cli_equivalence.py: the refactored flat working-dict is asserted byte-value-identical to the pre-refactor parser across a 17-invocation matrix (expected values in one golden_flat_dict.json -- a _baseline all-defaults dict + per-invocation deltas), and the generated CLI's option-string set is frozen against the pre-refactor parser (golden_cli_flags.json) so no flag can be silently renamed or dropped. Removes src/speculators/utils/argparse_utils.py and its test: train.py was the only caller of explicitly_provided_dests, which the SUPPRESS'd-namespace provenance replaces. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Rahul-Tuli <rtuli@redhat.com>

## Purpose This PR implements [D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting](https://arxiv.org/abs/2605.18810). D-PACE proposes an alternative per-position loss weighting for training block-parallel draft models. Instead of applying a fixed exponential decay over block positions, D-PACE dynamically adjusts the weights as training progresses. This is based on the intuition that although earlier positions are naturally more important (i.e., an early rejection forces all subsequent draft positions to be discarded), when they are sufficiently stabilized, the training priority should shift to later positions to improve overall acceptance rate. Two arguments are added to `train.py` to support D-PACE loss: - `--per-position-loss-weight ["fixed-exp-decay", "dpace"]` (default `fixed-exp-decay`). This flag enables D-PACE loss; while missing, it defaults to the original fixed exponential decay. - `--dpace-alpha` (default `0.5`). This flag is only effective when D-PACE is enabled. It smoothes per-position confidence according to equation 7 of the paper. The paper uses 0.5 as the default value for their evaluation and justifies this choice in sensitivity analysis. Note that D-PACE requires cross-entropy as the training loss (`--loss-fn ce`), allowing per-position confidence score computation to be piggybacked on the loss computation. Applying D-PACE to other losses are not supported and will result in `ValueError`. D-PACE is currently enabled for D-Flash and D-Spark draft models. Theoretically, it is compatible with all block-parallel drafters. ## Tests **D-Flash + D-PACE evaluation.** Multiple D-Flash checkpoints are trained with different loss function configurations: - `--loss-fn kld`: KL Divergence (current default) - `--loss-fn ce`: Cross Entropy with fixed exponential decay - `--per-position-loss-weight dpace --dpace-alpha ALPHA`: D-PACE with `ALPHA` = 0.3, 0.5, 0.7. The training set is 100K random samples from Magpie+Ultrachat. The verifier model is Qwen3-8B. Each checkpoint uses 1 A100 for hidden state generation and 7 A100s for training. Training configurations are identical for all checkpoints: ``` --target-layer-ids 1 9 17 25 34 --num-layers 5 --draft-vocab-size 32000 --max-anchors 3072 --block-size 8 --sliding-window 2048 --sliding-window-indices 0 1 2 3 4 --scheduler-type cosine --lr 0.0006 --epochs 1 --noise-std 0.0 ``` Metric: since D-PACE is a pure training-time optimization and does not change draft model's forward pass complexity, its gains are captured entirely by acceptance rates. Evaluation results: Mean-Accepted-Length (MAL) on [Speculator Benchmarks](https://huggingface.co/datasets/RedHatAI/speculator_benchmarks): | Dataset | KL Divergence | CE | D-PACE $\alpha=0.3$ | D-PACE $\alpha=0.5$| D-PACE $\alpha=0.7$ | |---|---:|---:|---:|---:|---:| | HumanEval | 2.977 | 2.956 | 3.047 | **3.169** | 3.139 | | math_reasoning | 3.275 | 3.322 | **3.541** | 3.532 | 3.531 | | qa | 2.279 | 2.343 | 2.455 | **2.495** | 2.431 | | question | 2.585 | 2.567 | 2.685 | 2.680 | **2.708** | | rag | 2.302 | 2.269 | **2.393** | 2.372 | 2.336 | | summarization | 1.906 | 1.918 | 1.981 | 1.980 | **1.991** | | tool_call | 2.372 | 2.327 | **2.480** | 2.454 | 2.467 | | translation | 2.096 | 2.072 | 2.176 | 2.183 | **2.191** | | writing | 2.565 | 2.586 | 2.683 | 2.684 | **2.718** | | Mean | 2.484 | 2.485 | 2.605 | **2.617** | 2.612 | **D-Spark + D-PACE evaluation.** D-Spark checkpoints use Markov rank 256 and vanilla Markov head type. The confidence head is disabled. Remaining training configurations match above D-Flash experiments. - `--loss-fn ce`: Vanilla fixed-decay cross-entropy - `--position-loss-weight dpace --dpace-alpha 0.5`: D-PACE with $\alpha=0.5$ MALs: | Dataset | CE | D-PACE $\alpha=0.5$ | |---|---:|---:| | HumanEval | 3.145 | **3.302** | | math_reasoning | 3.457 | **3.615** | | qa | 2.352 | **2.492** | | question | 2.671 | **2.770** | | rag | 2.276 | **2.384** | | summarization | 1.916 | **1.996** | | tool_call | 2.451 | **2.538** | | translation | 2.082 | **2.202** | | writing | 2.686 | **2.776** | | Mean | 2.560 | **2.675** | D-PACE improves upon fixed-decay loss for both D-Flash and D-Spark drafters. ## Checklist I have filled in: - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [x] The test plan/results, such as providing test command and pasting the results. - [ ] (Optional) The necessary documentation update. - [x] I (a human) have written or reviewed the code in this pr to the best of my ability. --------- Signed-off-by: Weifan Jiang <weijiang@redhat.com> Co-authored-by: Weifan Jiang <weifanjiang@a100-08.nemg-001.lab.rdu2.dc.redhat.com> Co-authored-by: Weifan Jiang <weijiang@redhat.com> Signed-off-by: Roderick Wu <Roderick-Wu@a100-04.nemg-001.lab.rdu2.dc.redhat.com>

…c schema Make scripts/train.py config-file-first without changing what any run does and without touching the model layer. - New src/speculators/train/config.py is the single source of truth: ~75 tunables as typed fields on cohesive pydantic-settings groups (verifier/draft/data/ generation/loss/optimizer/scheduler/trainer/logging + per-algo dflash/dspark/ peagle/mtp). The argparse CLI is GENERATED from these fields, so adding an argument is one typed field -- train.py has no per-flag add_argument. - --config run.yaml supplies any argument; precedence is CLI > YAML > coded defaults; --config is optional (no file == the old pure-CLI behaviour). - --dump-config prints the fully-resolved grouped YAML (a valid --config file); the resolved run.yaml is also saved next to each checkpoint for reproducibility. It runs the draft-init contract first, so a scaffold is never emitted for a config that would fail to load. - Provenance = the argparse.SUPPRESS'd namespace (not model_fields_set, which the nested partial-update machinery pollutes). A YAML leaf counts as provided only under its owning group block; a misplaced key is reported as unrecognised and kept out of provenance so it can't mis-fire the draft-init conflict check. - Every bad input is a clean argparse error (exit code 2, no traceback): schema ValidationError and a broken --config file (missing/unreadable/malformed/ non-mapping) alike; a missing required field names the flag to set. Phase 1: the resolved config is flattened back to the vars(args)-shaped dict the five SpeculatorModel classes consume via **kwargs, so the model layer and Trainer are untouched; only main()'s TrainerConfig reads the typed groups. Pushing typed configs into the model layer is deferred to a follow-up PR. Rebased onto current main. Absorbs the flags added since the branch point and ports them into the schema so behaviour is unchanged: - #736 D-PACE: --per-position-loss-weight, --dpace-alpha as typed DFlash fields, with the "dpace requires ce loss / alpha in (0,1]" check as a cross-group model_validator. - #651 scheduler: --scheduler-warmup-ratio, and --scheduler-type as a Literal (linear/cosine/none), threaded through the typed TrainerConfig build. - #749 default sliding-window attention: --sliding-window-indices renamed to --full-attention-indices (inverted semantics), matching upstream. The decoder-shaping flag set (used by the draft-init conflict check) is derived from the schema via a field tag, so it is not a second place to edit when a shaping knob is added. Equivalence is guarded by tests/unit/train/test_golden_cli_equivalence.py: the refactored flat working-dict is asserted byte-value-identical to the pre-refactor parser across a 17-invocation matrix (expected values in one golden_flat_dict.json -- a _baseline all-defaults dict + per-invocation deltas), and the generated CLI's option-string set is frozen against the pre-refactor parser (golden_cli_flags.json) so no flag can be silently renamed or dropped. Removes src/speculators/utils/argparse_utils.py and its test: train.py was the only caller of explicitly_provided_dests, which the SUPPRESS'd-namespace provenance replaces. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Rahul-Tuli <rtuli@redhat.com>

coderabbitai Bot reviewed Jul 7, 2026

View reviewed changes

Comment thread src/speculators/models/dflash/metrics.py Outdated

mergify Bot added the quality-failed label Jul 7, 2026

Weifan Jiang and others added 3 commits July 7, 2026 12:50

add DPACE loss support for DFlash training

89e6f99

Signed-off-by: Weifan Jiang <weijiang@redhat.com>

remove redundant argument path

ccf3465

Signed-off-by: Weifan Jiang <weijiang@redhat.com>

fix make style / make quality issues

b76c359

Signed-off-by: Weifan Jiang <weijiang@redhat.com>

weifanjiang force-pushed the dpace branch from 7b26944 to b76c359 Compare July 7, 2026 16:51

mergify Bot removed the quality-failed label Jul 7, 2026

fynnsu requested changes Jul 7, 2026

View reviewed changes

Weifan Jiang added 2 commits July 7, 2026 17:54

refactor loss computation; enable d-pace for d-spark

ce287cd

Signed-off-by: Weifan Jiang <weijiang@redhat.com>

fix make style/quality issues

fb2e959

Signed-off-by: Weifan Jiang <weijiang@redhat.com>

fynnsu approved these changes Jul 8, 2026

View reviewed changes

fynnsu added the ready This PR is ready for review label Jul 8, 2026

fynnsu merged commit 21033a7 into vllm-project:main Jul 8, 2026
9 of 10 checks passed

rahul-tuli mentioned this pull request Jul 8, 2026

refactor(train): config-file-first CLI generated from a typed pydantic schema #744

Draft

4 tasks

orestis-z mentioned this pull request Jul 9, 2026

RFC: Ablation Study — Domino vs DSpark Performance Comparison #757

Open

dsikka mentioned this pull request Jul 13, 2026

Q3 Roadmap (WIP) #781

Open

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add D-PACE loss implementation for D-Flash training#736

Add D-PACE loss implementation for D-Flash training#736
fynnsu merged 5 commits into
vllm-project:mainfrom
weifanjiang:dpace

weifanjiang commented Jul 7, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 7, 2026 •

edited

Loading

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

mergify Bot commented Jul 7, 2026

Uh oh!

fynnsu left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

weifanjiang commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Tests

Checklist

Uh oh!

coderabbitai Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mergify Bot commented Jul 7, 2026

Uh oh!

fynnsu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

weifanjiang commented Jul 7, 2026 •

edited

Loading

coderabbitai Bot commented Jul 7, 2026 •

edited

Loading