trending/Model-Optimizer
1
0
Fork 0
mirror of https://github.com/NVIDIA/Model-Optimizer.git synced 2026-10-02 03:14:52 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
Model-Optimizer/.pre-commit-config.yaml
T
h-guo18andClaude Opus 5 c2aaa44f60 [Speculative Decoding] DFlash2 draft variant (grouped sublayer convolution + candidate selector) (#2216) 
### What does this PR do?

Type of change: new feature

Adds **DFlash2** ([blog](https://inco.ai/blog/dflash2/)) as a draft
variant of the existing DFlash mode, selected with
`dflash_architecture_config.projector_type="dflash2"` alongside
`domino`, `dspark` and `lilicorr`.

DFlash2 keeps DFlash's one-pass parallel backbone and adds two
components that recover the acceptance a purely parallel draft loses:

- **Grouped dynamic depthwise convolution** around every attention and
MLP sublayer, giving each block position a view of its predecessors
*inside* the block. Taps do not cross the block boundary, so the draft
stays one forward pass.
- **Low-rank candidate selector** scoring transitions between adjacent
positions' top-k candidates, so serving walks one coherent path instead
of taking an independent argmax per position.

Both start as exact no-ops — the convolution's `base_kernel` is an
identity and `kernel_projection` is zeroed; the selector's
`successor_codebook` is zeroed — so a freshly built DFlash2 draft *is*
its DFlash backbone, and enabling the variant is an extension rather
than a perturbation. This matches the reference implementation
([SpecForge#772](https://github.com/sgl-project/SpecForge/pull/772),
merged) and the way `modeling_lilicorr` already installs this same
convolution class.

**This unblocks a recipe already shipped on `main`.**
`modeling_lilicorr._install_sublayer_convs` imports `DFlashGroupedConv`
from `modeling_dflash2`, so
`modelopt_recipes/general/speculative_decoding/lilicorr_conv.yaml`
raises at model build today and its CHANGELOG entry documents a feature
that cannot run. Landing this makes it runnable.

Module and parameter names match the SGLang/vLLM `DFlash2DraftModel`
loaders. Verified against the released `z-lab/Qwen3.8-27B-DFlash2`
checkpoint: **81 tensors, 21 name patterns, zero difference in either
direction**. The serving side,
[vllm-project/vllm#52816](https://github.com/vllm-project/vllm/pull/52816),
has since merged (`b389ac294`) with no change to the checkpoint
contract.

### Usage

```bash
python examples/speculative_decoding/main.py \
  --config modelopt_recipes/general/speculative_decoding/dflash2.yaml \
  model.model_name_or_path=Qwen/Qwen3-8B \
  data.data_path=<corpus>.jsonl \
  training.output_dir=<out>
```

```yaml
# modelopt_recipes/general/speculative_decoding/dflash2.yaml
dflash:
  dflash_selector_loss_alpha: 1.0      # weight of the candidate-selector CE term
  dflash_architecture_config:
    projector_type: dflash2
    conv_kernel_size: 2                # taps; must not exceed the block size
    conv_group_size: 16                # must divide hidden_size
    selector_rank: 256
    selector_top_k: 16
```

### Testing
<img width="2000" height="1320" alt="image"
src="https://github.com/user-attachments/assets/869004d1-92c6-41a4-a03d-fd824a06255c"
/>

**Unit** — 25 CPU tests in
`tests/unit/torch/speculative/plugins/test_hf_dflash2.py`; the full
`tests/unit/torch/speculative/` suite passes with no regressions. The
ones worth keeping pin invariants that a decreasing loss does not catch:

- the convolution is an exact identity on the **default** construction,
and its taps stay inside the block while a position still sees its
predecessors;
- the block-offset contract shared by the training objective and
`CandidateSelector.greedy_path` — a misaligned objective still
converges;
- the export fields the vLLM loader requires, including the top-level
`block_size` that `DFlash2Exporter` derives the nested copy from;
- which selector factors receive gradient on the first step.
`successor_codebook` starts at zero, so `predecessor_codebook` and
`hidden_projection` take one step to begin moving. That is a warm start,
not a dead branch, and both sides are asserted.

**End-to-end** — trained on Qwen3-8B against a plain DFlash control with
every other argument identical (plot above). Monotonic convergence, no
NaN/divergence, no DDP unused-parameter issues. Note the losses are
**not comparable** across arms: DFlash2's includes the selector CE term.

**Serving (vLLM)** — the exported drafter loads and drafts under the
merged DFlash2 path (`RESOLVED draft architectures:
['DFlash2DraftModel']`). Two notes for anyone reproducing: vLLM sizes
the convolution from `1 + num_speculative_tokens` at runtime rather than
from the checkpoint, so a `block_size=16` drafter is only correct at
`num_speculative_tokens=15`; and at that value the upstream path
currently hits an illegal memory access in `_cache_draft_logits`
([vllm#55279](https://github.com/vllm-project/vllm/issues/55279)),
independent of which checkpoint is used.

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

- Is this change backward compatible?: ✅ — additive. New
`projector_type`, its own registry and exporter, one new config field;
DFlash / Domino / DSpark / LiLiCorr numerics and `state_dict` contents
are untouched.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ —
`modeling_dflash2.py` is adapted from
[SpecForge#772](https://github.com/sgl-project/SpecForge/pull/772) and
carries its MIT notice. No new dependencies.
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ — under `0.48.0`.
- Did you get Claude approval on this PR?: ✅ — run on 2026-08-20; all
review threads addressed and resolved.

### Additional Information

Rebased onto current `main`. Two commits from the original branch were
dropped because
[#2342](https://github.com/NVIDIA/Model-Optimizer/pull/2342) landed them
first, with authorship preserved: the no-op sublayer seam in
`modeling_dflash.py`, and the `rope_theta`/`rope_parameters` fix —
`main`'s version of the latter is stricter, so this PR no longer touches
`hf_dflash.py` at all.


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

* **New Features**
* Added DFlash2 speculative decoding with grouped dynamic convolutions
and low-rank candidate selection.
* Added configurable selector-loss weighting, including an option to
disable it.
  * Added DFlash2 model conversion and export support.
  * Added checkpoints compatible with SGLang and vLLM DFlash2 serving.

* **Documentation**
* Added training recipes and a Qwen3-8B online DFlash2 training
configuration.

* **Tests**
* Added coverage for conversion, training, metrics, gradients, and
export compatibility.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-29 11:43:02 +08:00

204 lines
7.8 KiB
YAML
Raw Permalink Blame History

# NOTE: Make sure to update version in dev requirements (pyproject.toml) as well!
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v6.0.0
hooks:
- id: check-added-large-files
args: [--maxkb=500, --enforce-all]
exclude: >
(?x)^(
uv.lock|
examples/diffusers/quantization/assets/.*.png|
examples/diffusers/cache_diffusion/assets/.*.png|
)$
- id: check-json
exclude: ^.vscode/.*.json # vscode files can take comments
- id: check-merge-conflict
- id: check-symlinks
- id: check-toml
- id: mixed-line-ending
args: [--fix=lf]
- id: requirements-txt-fixer
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.20
hooks:
- id: ruff-check
args: [--fix, --exit-non-zero-on-fix]
- id: ruff-format
- repo: https://github.com/pre-commit/mirrors-mypy
rev: v2.1.0
hooks:
- id: mypy
- repo: https://github.com/pre-commit/mirrors-clang-format
rev: v21.1.0
hooks:
- id: clang-format
types_or: [c++, c, c#, cuda, java, javascript, objective-c, proto] # no json!
args: ["--style={ColumnLimit: 100}"]
- repo: https://github.com/pre-commit/pygrep-hooks
rev: v1.10.0
hooks:
- id: rst-backticks
- id: rst-directive-colons
- id: rst-inline-touching-normal
- repo: https://github.com/jumanjihouse/pre-commit-hook-yamlfmt
rev: 0.2.3
hooks:
- id: yamlfmt
args: [--mapping=2, --sequence=4, --offset=2, --implicit_start, --implicit_end, --preserve-quotes]
exclude: ^.github/workflows/
- repo: local
hooks:
- id: normalize-yaml-ext
name: normalize .yml to .yaml in required places, right now only yaml files in modelopt_recipes
entry: uv run --frozen --extra dev python tools/precommit/normalize_yaml_ext.py
language: system
files: ^modelopt_recipes/.*\.yml$
- id: check-modelopt-recipes
name: validate modelopt recipes
entry: uv run --frozen --extra dev python tools/precommit/check_modelopt_recipes.py
language: system
files: ^modelopt_recipes/
# configs/ contains reusable snippets (not full recipes) — skip recipe validation
exclude: ^modelopt_recipes/configs/
- id: sync-claude-skills
name: sync .claude/skills/ symlinks from plugin skills
entry: bash tools/precommit/sync_claude_skills.sh
language: system
files: ^plugins/modelopt/skills/
pass_filenames: false
- id: check-launcher-yaml
name: validate launcher YAML references to recipes and templates
entry: uv run --frozen --extra dev python tools/precommit/check_launcher_yaml.py
language: system
files: ^(tools/launcher/examples/.*\.yaml|tools/precommit/check_launcher_yaml\.py)$
# Instructions to change license file if ever needed:
# https://github.com/Lucas-C/pre-commit-hooks#removing-old-license-and-replacing-it-with-a-new-one
- repo: https://github.com/Lucas-C/pre-commit-hooks
rev: v1.5.5
hooks:
# Default hook for Apache 2.0 in python and shell files
- id: insert-license
alias: insert-license-py
args:
- --license-filepath
- ./LICENSE_HEADER
- --comment-style
- "#"
- --allow-past-years
types_or: [python, shell]
# NOTE: Exclude files that have copyright or license headers from another company or individual
# since we want to keep those above the license header added by this hook.
# Instead, we should manually add the license header to those files *after* the original header.
exclude: >
(?x)^(
modelopt/torch/quantization/utils/calib_utils.py|
modelopt/torch/quantization/ggml/codebooks.py|
modelopt/onnx/quantization/operators.py|
modelopt/onnx/quantization/ort_patching.py|
modelopt/torch/_deploy/utils/onnx_utils.py|
modelopt/torch/export/transformer_engine.py|
modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py|
modelopt/torch/quantization/export_onnx.py|
modelopt/torch/quantization/plugins/attention.py|
modelopt/torch/sparsity/attention_sparsity/methods/vsa_utils.py|
modelopt/torch/speculative/eagle/utils.py|
modelopt/torch/speculative/plugins/hf_domino.py|
modelopt/torch/speculative/plugins/modeling_domino.py|
modelopt/torch/speculative/plugins/hf_dflash.py|
modelopt/torch/speculative/plugins/modeling_dflash.py|
modelopt/torch/speculative/plugins/hf_dflash2.py|
modelopt/torch/speculative/plugins/modeling_dflash2.py|
modelopt/torch/speculative/plugins/hf_dspark.py|
modelopt/torch/speculative/plugins/modeling_dspark.py|
modelopt/torch/speculative/plugins/hf_lilicorr.py|
modelopt/torch/speculative/plugins/modeling_lilicorr.py|
modelopt/torch/speculative/plugins/hf_medusa.py|
modelopt/torch/utils/plugins/megatron_mmlu.py|
examples/deepseek/deepseek_v3/quantize_to_nvfp4.py|
examples/deepseek/deepseek_v3/ptq.py|
examples/diffusers/quantization/onnx_utils/export.py|
examples/llm_eval/lm_eval_hf.py|
examples/llm_eval/mmlu.py|
examples/llm_eval/modeling.py|
examples/onnx_ptq/bevformer/prepare_calibration.py|
examples/onnx_ptq/far3d/evaluate.py|
examples/onnx_ptq/petr/evaluate.py|
examples/onnx_ptq/trt_runner.py|
examples/llm_qat/train.py|
examples/llm_sparsity/weight_sparsity/finetune.py|
examples/specdec_bench/specdec_bench/models/specbench_medusa.py|
examples/speculative_decoding/main.py|
examples/speculative_decoding/medusa_utils.py|
examples/speculative_decoding/scripts/server_generate.py|
experimental/dms/models/qwen3/configuration_qwen3_dms.py|
experimental/dms/models/qwen3/modeling_qwen3_dms.py|
)$
# Default hook for Apache 2.0 in c/c++/cuda files
- id: insert-license
alias: insert-license-c
args:
- --license-filepath
- ./LICENSE_HEADER
- --comment-style
- "/*| *| */"
- --allow-past-years
types_or: [c++, cuda, c]
- repo: https://github.com/PyCQA/bandit
rev: 1.7.9
hooks:
- id: bandit
args: ["-c", "pyproject.toml", "-q"]
additional_dependencies: ["bandit[toml]"]
- repo: local
hooks:
- id: generate-arguments-md
name: Regenerate examples/llm_qat/ARGUMENTS.md
entry: uv run --frozen --extra dev python examples/llm_qat/arguments.py --generate_docs examples/llm_qat/ARGUMENTS.md
language: system
files: >-
(?x)^(
examples/llm_qat/arguments\.py|
modelopt/torch/distill/plugins/huggingface\.py|
modelopt/torch/opt/plugins/transformers\.py|
modelopt/torch/quantization/plugins/transformers_trainer\.py
)$
pass_filenames: false
- repo: https://github.com/DavidAnson/markdownlint-cli2
rev: v0.18.1
hooks:
- id: markdownlint-cli2
args: ["--fix"]
##### Manual hooks (Expect many false positives)
# These hooks are only run with `pre-commit run --all-files --hook-stage manual <hook_id>`
# Spell checker
- repo: https://github.com/crate-ci/typos
rev: v1.35.8
hooks:
- id: typos
stages: [manual]
# Link checker
- repo: https://github.com/lycheeverse/lychee.git
rev: v0.15.1
hooks:
- id: lychee
args: ["--no-progress", "--exclude-loopback"]
stages: [manual]
Reference in New Issue View Git Blame Copy Permalink