Skip to content

feat(specdec): warm-start DFlash fine-tuning from exported drafter via dflash_init_checkpoint#1978

Draft
h-guo18 wants to merge 3 commits into
mainfrom
haoguo/dflash-finetune-example
Draft

feat(specdec): warm-start DFlash fine-tuning from exported drafter via dflash_init_checkpoint#1978
h-guo18 wants to merge 3 commits into
mainfrom
haoguo/dflash-finetune-example

Conversation

@h-guo18

@h-guo18 h-guo18 commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: new feature

Adds warm-start fine-tuning for DFlash: a new dflash_init_checkpoint field on DFlashConfig loads an already-trained, exported drafter (DFlashExporter / z-lab layout) into model.dflash_module at conversion time, so training continues from it instead of random init. Must be a local directory containing model.safetensors (stage public drafters with hf download first).

  • Shape/key mismatch vs dflash_architecture_config fails with a clear error; mask_token_id / target_layer_ids mismatch vs the checkpoint's config.json logs a degradation warning. dflash_block_size may differ (weights are block-size agnostic).
  • Domino/DSpark can warm-start their DFlash backbone from a plain DFlash checkpoint (missing head submodules train from scratch, with a warning).
  • Restore of a saved training checkpoint ignores the field (restore_dflash_model clears it), so resume works even if the init checkpoint is gone.

Also adds a launcher example (tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml) and a doc section in examples/speculative_decoding/doc/dflash.md.

Usage

# recipe YAML (or launcher dotlist: dflash.dflash_init_checkpoint=...)
dflash:
  dflash_init_checkpoint: /hf-local/z-lab/Qwen3-8B-DFlash-b16  # local dir with model.safetensors
  dflash_block_size: 16
  dflash_mask_token_id: 151669
  # must match the checkpoint's draft architecture (dims do NOT inherit from the base model):
  dflash_architecture_config:
    num_hidden_layers: 5
    num_attention_heads: 32
    num_key_value_heads: 8
    head_dim: 128
    intermediate_size: 12288
uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml --yes

Testing

  • New unit tests: TestDFlashWarmStart (export→load round-trip at a different block size, missing-dir / arch-mismatch errors, restore ignores the field). All dflash/domino/dspark unit tests pass.
  • End-to-end: converted Qwen/Qwen3-8B with dflash_init_checkpoint pointing at a local copy of z-lab/Qwen3-8B-DFlash-b16 — all 58 tensors load bit-exact, target_layer_ids=[1, 9, 17, 25, 33] matches the checkpoint, no warnings.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

  • Is this change backward compatible?: ✅ (new optional field, default None; existing configs/checkpoints unaffected)
  • If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
  • Did you write any new necessary tests?: ✅
  • Did you update Changelog?: ❌ (can add if desired)
  • Did you get Claude approval on this PR?: ❌ (pending)

…a dflash_init_checkpoint

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

copy-pr-bot Bot commented Jul 15, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 58ef60cd-eeba-4386-9686-a6013097b94d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch haoguo/dflash-finetune-example

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

@h-guo18 h-guo18 changed the title feat(specdec): warm-start DFlash fine-tuning from exported drafter vi… feat(specdec): warm-start DFlash fine-tuning from exported drafter via dflash_init_checkpoint Jul 15, 2026
@github-actions

github-actions Bot commented Jul 15, 2026

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1978/

Built to branch gh-pages at 2026-07-15 04:33 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

@codecov

codecov Bot commented Jul 15, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.93939% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.88%. Comparing base (f479e78) to head (52d8ca8).

Files with missing lines Patch % Lines
modelopt/torch/speculative/plugins/hf_dflash.py 93.54% 2 Missing ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##             main    #1978   +/-   ##
=======================================
  Coverage   77.87%   77.88%           
=======================================
  Files         522      522           
  Lines       58452    58485   +33     
=======================================
+ Hits        45522    45553   +31     
- Misses      12930    12932    +2     
Flag Coverage Δ
unit 55.41% <93.93%> (+0.02%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

h-guo18 added 2 commits July 15, 2026 04:25
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant