feat(specdec): warm-start DFlash fine-tuning from exported drafter via dflash_init_checkpoint#1978
feat(specdec): warm-start DFlash fine-tuning from exported drafter via dflash_init_checkpoint#1978h-guo18 wants to merge 3 commits into
Conversation
…a dflash_init_checkpoint Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
|
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. |
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
Codecov Report❌ Patch coverage is
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
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
What does this PR do?
Type of change: new feature
Adds warm-start fine-tuning for DFlash: a new
dflash_init_checkpointfield onDFlashConfigloads an already-trained, exported drafter (DFlashExporter / z-lab layout) intomodel.dflash_moduleat conversion time, so training continues from it instead of random init. Must be a local directory containingmodel.safetensors(stage public drafters withhf downloadfirst).dflash_architecture_configfails with a clear error;mask_token_id/target_layer_idsmismatch vs the checkpoint'sconfig.jsonlogs a degradation warning.dflash_block_sizemay differ (weights are block-size agnostic).restore_dflash_modelclears 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 inexamples/speculative_decoding/doc/dflash.md.Usage
Testing
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.Qwen/Qwen3-8Bwithdflash_init_checkpointpointing at a local copy ofz-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.).None; existing configs/checkpoints unaffected)CONTRIBUTING.md: N/A