feat(dflash2): co-train the DFlash2 drafter against vLLM rollout - #77
Open
khazic wants to merge 6 commits into
Open
feat(dflash2): co-train the DFlash2 drafter against vLLM rollout#77khazic wants to merge 6 commits into
khazic wants to merge 6 commits into
Conversation
Map DFLASH2 onto vLLM's dflash method (the engine picks the DFlash2 draft from the DFlash2DraftModel architecture, vLLM >= 0.28.0), validate the checkpoint contract and the block size against spec_verify_tokens up front, publish the selector codebooks under the engine's bare parameter names, mirror the DFlash2 knobs into dflash_config for trainer-saved configs, and add a converter for speculators-format DFlash2 drafters plus the example, CI arm and docs. Signed-off-by: khazic <khazzz1c@gmail.com>
Signed-off-by: khazic <khazzz1c@gmail.com>
…ents CUDA tensors shared over IPC out of an expandable segment need pidfd_getfd on the receiving vLLM worker (Linux >= 5.6); on older kernels every draft publish failed with 'does not support the pidfd_getfd syscall'. Mirror verl's own actor->rollout sync and flip expandable segments off around the bucketed send, restoring them afterwards. Signed-off-by: khazic <khazzz1c@gmail.com>
…ames on the IPC path The bucketed IPC receiver feeds the translated names to the engine's load_weights, which rejects the trainer's nn.Embedding spelling (candidate_selector.*_codebook.weight) for vLLM's bare parameters. Route both publish paths through one rename helper. Signed-off-by: khazic <khazzz1c@gmail.com>
Load the drafter config.json once per speculative-config build, replace the codebook-suffix regex with str.endswith (drops the re import), reuse _get_nested for the top-level-then-nested knob lookup, flatten the converter's rope_theta branch, and cross-reference the trainer/engine codebook alias tables. Signed-off-by: khazic <khazzz1c@gmail.com>
…the per-drafter block Match the dspark pattern (SPECO_DFLASH2_SPEC_VERIFY_TOKENS) instead of threading a drafter-aware default through the shared overrides array, and document why the IPC allocator guard restores expandable segments unconditionally. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
Author
|
Sanitized logs of the three GPU validation runs (internal paths and hosts rewritten). The excerpts below keep the launch command, every per-step metrics line (acceptance length, drafter train/publish flags, reward), and the drafter/publish lifecycle messages; the complete files are on the End-to-end smoke, 8 GRPO steps (4 train+publish cycles) (excerpt)Frozen-drafter baseline, 100 GRPO steps (excerpt) |
Contributor
Author
|
Co-training arm, continued from the previous comment: Co-training run, 100 GRPO steps (50 train+publish cycles) (excerpt) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds online co-training for the DFlash2 drafter against vLLM rollout. Since #62, DFlash2 could only be trained offline:
_speculative_method_from_drafterrefusedDFLASH2for both engines. vLLM ships the DFlash2 draft model since v0.28.0 (qwen3_dflash2.py, dynamic convolutions plus candidate selector, vllm-project/vllm#52816), served through the DFlash speculative method and dispatched on the checkpoint'sDFlash2DraftModelarchitecture, so the full collect / train / publish loop now runs with DFlash2 the same way it does for DFlash and DSpark.Changes
DFLASH2maps to vLLM methoddflash; the SGLang path keeps its fail-loud (not wired) and now points users at vLLM.vllm.model_executor.models.qwen3_dflash2(>= 0.28.0);dflash_configconv/selector hyperparameters vLLM indexes unconditionally;rollout.spec_verify_tokens == block_size - 1, because vLLM sizes the convolution block as the bonus token plus the speculative tokens while the trainer folds bydflash2_block_size; a mismatch would silently serve a conv trained on a different block layout.nn.Embeddingmodules on the trainer (..._codebook.weight) but as bare parameters in vLLM (z-lab checkpoint spelling). A shared_dflash2_engine_param_namerenames them on both publish paths (graph-safe per-param update and the bucketed-IPCload_weightspath, which rejects the trainer spelling with "Attempted to load nested weight ... into a single parameter").pidfd_getfd(< Linux 5.6), sharing tensors allocated in expandable segments makes the receiving vLLM worker fail the whole update._ipc_safe_allocatornow flips expandable segments off around the bucketed send and restores them, mirroring what verl's own actor-to-rollout sync does.DFlash2Config.to_dictemits the nesteddflash_configblock (z-lab layout) so a checkpoint saved by the trainer stays loadable by vLLM; the EAGLEConfig alias patch also mirrors flat DFlash2 keys intodflash_configfor older flat checkpoints.python -m verl_speco.convert_speculators_dflash2rewrites a speculators-format DFlash2 drafter (e.g.mgoin/Qwen3-4B-speculator.dflash2) into the z-lab layout; vLLM 0.28.0 predates speculators-format DFlash2 loading (Add support for loading dflash2 model in speculators format vllm-project/vllm#53797), and the weights use identical parameter names in both layouts, so onlyconfig.jsonis rewritten.examples/run_qwen3-8b_drafter_dflash2_vllm.sh, adflash2arm inci/run_example_test.sh, README rows and the DFlash2 serving note.Test
CPU:
tests/integration/test_vllm_dflash2_runtime_contract.py,tests/integration/test_vllm_draft_update_allocator_contract.py,tests/unit/test_convert_speculators_dflash2.py(new), plus the updatedtest_dflash2_backend_contract.pyand the existingtest_vllm_runtime_contract.py: 42 passed against vLLM 0.28.0, and the full CPU suite shows no regressions relative to the base commit in the same environment.GPU (single node, vLLM 0.28.0, torch 2.13, flash-attn 2.8.3): GRPO on GSM8K with Qwen3-4B as the target and the converted
mgoin/Qwen3-4B-speculator.dflash2drafter, viaci/run_example_test.sh gpu vllm dflash2.DFlash2Qwen3Modelper publish (convolutions and selector codebooks included), acceptance length stable at 4.25 to 4.58 across publishes, per-cycle drafter loss 4.31 -> 3.21, 3.86 -> 2.89, 3.88 -> 3.05, 3.46 -> 2.65.drafter.enable_drafter_training=false): acceptance 4.470 +- 0.113, no drift-induced decay on this task.The complete sanitized logs of all three runs (key phases inline, full files as gist links) are attached in the PR comments.
Notes
hidden_state_window_tokens_per_sample) exceeds typical response lengths, most samples are dropped at collection and the drafter silently never trains (schedule_reason=inconsistent_data_versionevery cycle, because some drafter workers end up with no samples). Documenting or fail-louding that interaction is worth a separate issue; it is not DFlash2-specific.