From 7b12023d94568ad9b3135e157b2290b465d1492f Mon Sep 17 00:00:00 2001 From: khazic Date: Tue, 1 Sep 2026 21:15:49 +0800 Subject: [PATCH 1/6] feat(dflash2): co-train the DFlash2 drafter against vLLM rollout 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 --- README.md | 16 +- ci/run_example_test.sh | 15 +- examples/run_qwen3-8b_drafter_dflash2_vllm.sh | 110 +++++++ pyproject.toml | 1 + .../test_dflash2_backend_contract.py | 44 ++- .../test_vllm_dflash2_runtime_contract.py | 277 ++++++++++++++++++ .../unit/test_convert_speculators_dflash2.py | 215 ++++++++++++++ verl_speco/convert_speculators_dflash2.py | 275 +++++++++++++++++ verl_speco/integration/rollout_publish.py | 2 + verl_speco/integration/sglang_runtime.py | 9 +- verl_speco/integration/vllm_runtime.py | 191 ++++++++++-- .../models/dflash2/configuration_dflash2.py | 18 ++ 12 files changed, 1145 insertions(+), 28 deletions(-) create mode 100644 examples/run_qwen3-8b_drafter_dflash2_vllm.sh create mode 100644 tests/integration/test_vllm_dflash2_runtime_contract.py create mode 100644 tests/unit/test_convert_speculators_dflash2.py create mode 100644 verl_speco/convert_speculators_dflash2.py diff --git a/README.md b/README.md index 044aacf8..6193382d 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ training, and hot-update logic through `verl_speco`. - **Multiple drafter backends**: includes EAGLE-1, EAGLE-2, EAGLE3, DFlash, DSpark, Domino, and P-EAGLE trainer backends under `verl_speco.backends`. - **vLLM and SGLang integration**: supports EAGLE-1, EAGLE-2, EAGLE3, DFlash, - and DSpark speculative decoding on vLLM, plus EAGLE3 and DFlash on SGLang, + DFlash2, and DSpark speculative decoding on vLLM, plus EAGLE3 and DFlash on SGLang, with drafter collection and hot-update logic integrated through the rollout engine. - **GPU and NPU examples**: provides example scripts for vLLM, SGLang, and @@ -63,6 +63,7 @@ faster end-to-end training without accuracy regression. | EAGLE-2 | vLLM | FSDP | Available | | EAGLE3 | vLLM, SGLang | FSDP | Available | | DFlash | vLLM, SGLang | FSDP | Available | +| DFlash2 | vLLM via DFlash | FSDP | Available | | DSpark | vLLM | FSDP | Available | | Domino | vLLM, SGLang via DFlash | FSDP | Available | | P-EAGLE | Not wired in this overlay | FSDP | Training only | @@ -95,6 +96,7 @@ drafter backend you use. | EAGLE-1 / EAGLE-2 | Engine version with native EAGLE support | Runtime-specific | - | | EAGLE3 | >= 0.18.0 | >= 0.18.0 | >= 0.5.10 | | DFlash | >= 0.20.2 | >= 0.20.2 | >= 0.5.12 | +| DFlash2 | >= 0.28.0 (served as DFlash) | - | - | | DSpark | GPU: [main](https://github.com/vllm-project/vllm/tree/main)
NPU: [`dc68bd8`](https://github.com/vllm-project/vllm/tree/dc68bd8c4199b00631fe71eb37313f406cc66ac1) | NPU: [`8214d19`](https://github.com/vllm-project/vllm-ascend/tree/8214d19f8b505484b839469444887b404db2e3a8) | - | | Domino | DFlash-compatible runtime with Domino projector support | Runtime-specific | Runtime-specific | | P-EAGLE | Not wired | Not wired | Not wired | @@ -102,6 +104,18 @@ drafter backend you use. For vLLM DFlash, the drafter checkpoint must use the DFlash draft model config expected by the runtime. +For vLLM DFlash2, keep `speculative_algorithm=DFLASH2`: the overlay maps it onto +vLLM's DFlash method and the engine picks the DFlash2 draft (dynamic +convolutions plus candidate selector) from the checkpoint's `DFlash2DraftModel` +architecture, so both the drafter training loop and the rollout drafter run +DFlash2. The checkpoint must use the z-lab layout with the DFlash2 knobs under +`dflash_config`; `python -m verl_speco.convert_speculators_dflash2` rewrites a +speculators-format drafter (for example `mgoin/Qwen3-4B-speculator.dflash2`) into +it. vLLM sizes the convolution block as the bonus token plus +`rollout.spec_verify_tokens`, so set `spec_verify_tokens = dflash2_block_size - 1` +(see `examples/run_qwen3-8b_drafter_dflash2_vllm.sh`). SGLang co-training of +DFlash2 is not wired yet. + For vLLM DSpark on GPU, use vLLM main. For vLLM DSpark on NPU, follow the version pairing documented by [vLLM-Ascend PR #11153](https://github.com/vllm-project/vllm-ascend/pull/11153): diff --git a/ci/run_example_test.sh b/ci/run_example_test.sh index 531562a8..23557a39 100644 --- a/ci/run_example_test.sh +++ b/ci/run_example_test.sh @@ -15,6 +15,9 @@ case "${platform}/${backend}/${drafter}" in gpu/vllm/dspark) example="examples/run_qwen3-8b_drafter_dspark_vllm.sh" ;; + gpu/vllm/dflash2) + example="examples/run_qwen3-8b_drafter_dflash2_vllm.sh" + ;; gpu/vllm/peagle|gpu/vllm/domino) example="examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh" ;; @@ -43,7 +46,7 @@ case "${platform}/${backend}/${drafter}" in example="examples/run_qwen3-8b_drafter_dflash_sglang.sh" ;; *) - echo "usage: $0 {gpu|npu} {vllm|sglang} {eagle3|megatron-eagle3|dflash|dspark|peagle|domino}" >&2 + echo "usage: $0 {gpu|npu} {vllm|sglang} {eagle3|megatron-eagle3|dflash|dflash2|dspark|peagle|domino}" >&2 exit 2 ;; esac @@ -69,6 +72,7 @@ done # The example owns the collect-to-train algorithm pairing; the runner only needs # the collect side of it to pick the matching cached draft model. separate_training="false" +spec_verify_tokens_default="4" case "${drafter}" in eagle3|megatron-eagle3) @@ -83,6 +87,13 @@ case "${drafter}" in draft_model="${SPECO_DSPARK_DRAFT_MODEL:-}" draft_algorithm="DSPARK" ;; + dflash2) + draft_model="${SPECO_DFLASH2_DRAFT_MODEL:-}" + draft_algorithm="DFLASH2" + # vLLM sizes the DFlash2 convolution block as 1 + spec_verify_tokens, and + # the trainer folds by dflash2_block_size (default 8), so the two must agree. + spec_verify_tokens_default="7" + ;; peagle) draft_model="${SPECO_EAGLE3_DRAFT_MODEL:-}" draft_algorithm="EAGLE3" @@ -186,7 +197,7 @@ overrides=( "actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${SPECO_REF_LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-1}" "actor_rollout_ref.rollout.drafter.rollout.spec_steps=${SPECO_SPEC_STEPS:-3}" "actor_rollout_ref.rollout.drafter.rollout.spec_topk=${SPECO_SPEC_TOPK:-1}" - "actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=${SPECO_SPEC_VERIFY_TOKENS:-4}" + "actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=${SPECO_SPEC_VERIFY_TOKENS:-${spec_verify_tokens_default}}" "actor_rollout_ref.rollout.drafter.training.step=${SPECO_DRAFTER_TRAINING_STEPS:-1}" "actor_rollout_ref.rollout.drafter.training.batch_size_per_gpu=${SPECO_DRAFTER_BATCH_SIZE_PER_GPU:-1}" "actor_rollout_ref.rollout.drafter.training.collect_interval_steps=${SPECO_COLLECT_INTERVAL_STEPS:-1}" diff --git a/examples/run_qwen3-8b_drafter_dflash2_vllm.sh b/examples/run_qwen3-8b_drafter_dflash2_vllm.sh new file mode 100644 index 00000000..36d87db8 --- /dev/null +++ b/examples/run_qwen3-8b_drafter_dflash2_vllm.sh @@ -0,0 +1,110 @@ +set -x + +# GPU example for a DFlash2 drafter co-trained against vLLM's DFlash proposer. +# +# DFlash2 (dynamic convolutions + candidate selector) is served by vLLM >= 0.28.0 +# through the DFlash speculative method; the engine picks the DFlash2 draft from +# the checkpoint's "DFlash2DraftModel" architecture. The drafter checkpoint must +# use the z-lab layout (DFlash2 knobs under dflash_config); convert a +# speculators-format drafter with +# python -m verl_speco.convert_speculators_dflash2 --input ... --target ... --output ... +# The rollout block is the bonus token plus spec_verify_tokens mask tokens, so +# spec_verify_tokens must equal dflash2_block_size - 1 (8 -> 7 below). +project_name='verl_grpo_example_dflash2_drafter' +exp_name='qwen3_8b_dflash2_drafter_vllm_gpu' + +gen_tp=2 +train_sp=1 +ray_num_cpus=${SPECO_RAY_NUM_CPUS:-64} +ray_worker_soft_limit=${SPECO_RAY_WORKER_SOFT_LIMIT:-8} + +MODEL_PATH=/path/to/model +CKPTS_DIR=/path/to/checkpoint +TRAIN_FILE=/path/to/train_file +TEST_FILE=/path/to/test_file +DRAFTER_PATH=/path/to/vllm-compatible-dflash2-drafter + + +PYTHONUNBUFFERED=1 python3 -m verl_speco.main \ + algorithm.adv_estimator=grpo \ + ray_kwargs.ray_init.num_cpus=${ray_num_cpus} \ + +ray_kwargs.ray_init._system_config.prestart_worker_first_driver=false \ + +ray_kwargs.ray_init._system_config.num_workers_soft_limit=${ray_worker_soft_limit} \ + data.train_files=${TRAIN_FILE} \ + data.val_files=${TEST_FILE} \ + data.train_batch_size=16 \ + data.max_prompt_length=512 \ + data.max_response_length=8192 \ + data.filter_overlong_prompts=True \ + data.filter_overlong_prompts_workers=256 \ + data.truncation='error' \ + actor_rollout_ref.rollout.temperature=0.6 \ + actor_rollout_ref.model.path=${MODEL_PATH} \ + actor_rollout_ref.actor.optim.lr=1e-6 \ + actor_rollout_ref.model.use_remove_padding=True \ + actor_rollout_ref.actor.ppo_mini_batch_size=16 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.actor.use_kl_loss=True \ + actor_rollout_ref.actor.kl_loss_coef=0.001 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.entropy_coeff=0 \ + actor_rollout_ref.actor.calculate_entropy=False \ + actor_rollout_ref.model.enable_gradient_checkpointing=True \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \ + actor_rollout_ref.actor.ulysses_sequence_parallel_size=${train_sp} \ + actor_rollout_ref.ref.ulysses_sequence_parallel_size=${train_sp} \ + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.actor.use_dynamic_bsz=True \ + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \ + actor_rollout_ref.rollout.name=vllm \ + actor_rollout_ref.rollout.enforce_eager=False \ + actor_rollout_ref.rollout.enable_chunked_prefill=True \ + actor_rollout_ref.rollout.enable_prefix_caching=True \ + actor_rollout_ref.rollout.max_num_seqs=256 \ + actor_rollout_ref.rollout.max_num_batched_tokens=12288 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \ + actor_rollout_ref.rollout.n=5 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \ + actor_rollout_ref.ref.fsdp_config.param_offload=True \ + actor_rollout_ref.rollout.drafter.enable=True \ + actor_rollout_ref.rollout.drafter.enable_drafter_training=True \ + actor_rollout_ref.rollout.drafter.model_path=${DRAFTER_PATH} \ + actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH2 \ + actor_rollout_ref.rollout.drafter.training.collect_hidden_states_from_sgl=False \ + actor_rollout_ref.rollout.drafter.training.collect_hidden_states_from_old_logprob=True \ + actor_rollout_ref.rollout.drafter.training.old_logprob_hidden_capture_impl=forward_hook \ + actor_rollout_ref.rollout.drafter.training.dflash2_block_size=8 \ + actor_rollout_ref.rollout.drafter.training.dflash2_num_anchors=64 \ + actor_rollout_ref.rollout.drafter.training.dflash2_loss_decay_gamma=7 \ + actor_rollout_ref.rollout.drafter.training.dflash2_selector_loss_weight=1.0 \ + actor_rollout_ref.rollout.drafter.training.dflash_max_window=512 \ + actor_rollout_ref.rollout.drafter.rollout.spec_steps=1 \ + actor_rollout_ref.rollout.drafter.rollout.spec_topk=1 \ + actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=7 \ + actor_rollout_ref.rollout.drafter.training.step=10 \ + actor_rollout_ref.rollout.drafter.training.collect_interval_steps=5 \ + actor_rollout_ref.rollout.drafter.training.training_interval_steps=5 \ + actor_rollout_ref.rollout.drafter.training.publish_async=True \ + actor_rollout_ref.rollout.drafter.training.publish_dtype=bf16 \ + actor_rollout_ref.rollout.drafter.training.draft_update_weights_bucket_megabytes=512 \ + actor_rollout_ref.rollout.drafter.training.draft_update_pause_generation=True \ + actor_rollout_ref.rollout.drafter.training.draft_update_flush_before=False \ + actor_rollout_ref.rollout.drafter.training.draft_update_flush_after=True \ + actor_rollout_ref.rollout.load_format="auto" \ + actor_rollout_ref.actor.strategy=fsdp2 \ + algorithm.use_kl_in_reward=False \ + trainer.val_before_train=False \ + trainer.critic_warmup=0 \ + trainer.logger='["console", "wandb"]' \ + trainer.project_name=${project_name} \ + trainer.experiment_name=${exp_name} \ + trainer.n_gpus_per_node=16 \ + trainer.nnodes=1 \ + trainer.default_local_dir=${CKPTS_DIR} \ + trainer.total_training_steps=100 \ + trainer.save_freq=20 \ + trainer.test_freq=5 \ + trainer.total_epochs=6 $@ diff --git a/pyproject.toml b/pyproject.toml index ca8a9c7b..e4755cb2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ Repository = "https://github.com/verl-project/verl-SpeCo" verl-speco = "verl_speco.main:main" verl-speco-draft-train = "verl_speco.draft_train_launcher:main" verl-speco-inspect-features = "verl_speco.inspect_feature_store:main" +verl-speco-convert-speculators-dflash2 = "verl_speco.convert_speculators_dflash2:main" [tool.setuptools.dynamic] version = { attr = "verl_speco.__version__" } diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index 576f7f05..82a93d38 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -22,6 +22,8 @@ from __future__ import annotations +import json + import pytest @@ -775,11 +777,47 @@ def test_plain_dflash_checkpoint_warm_starts_without_the_dflash2_modules() -> No backend._load_draft_checkpoint(model, "", normalized_state=backbone_only) -def test_dflash2_rejected_by_vllm_config_builder() -> None: +def test_dflash2_maps_to_the_vllm_dflash_method() -> None: + """vLLM serves DFlash2 through its DFlash proposer, dispatching on the architecture.""" from verl_speco.integration.vllm_runtime import _speculative_method_from_drafter - with pytest.raises(ValueError, match="not an engine-level speculative algorithm"): - _speculative_method_from_drafter({"speculative_algorithm": "DFLASH2"}) + assert _speculative_method_from_drafter({"speculative_algorithm": "DFLASH2"}) == "dflash" + + +def test_dflash2_config_serializes_the_runtime_dflash_config_block() -> None: + """A trainer-saved DFlash2 config must stay servable by vLLM (nested knobs).""" + pytest.importorskip("torch") + pytest.importorskip("transformers") + from verl_speco.models.dflash2 import DFlash2Config + + config = DFlash2Config( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=256, + target_layer_ids=[1, 3], + mask_token_id=255, + block_size=4, + selector_top_k=8, + ) + + serialized = json.loads(config.to_json_string()) + nested = serialized["dflash_config"] + assert nested["block_size"] == 4 + assert nested["conv_kernel_size"] == 2 + assert nested["conv_group_size"] == 16 + assert nested["selector_rank"] == 256 + assert nested["selector_top_k"] == 8 + assert nested["mask_token_id"] == 255 + assert nested["target_layer_ids"] == [1, 3] + assert serialized["architectures"] == ["DFlash2DraftModel"] + + # The nested block must round-trip through the overlay's own loader too. + reloaded = DFlash2Config.from_dict(serialized) + assert reloaded.selector_top_k == 8 + assert reloaded.block_size == 4 def test_dflash2_rejected_by_sglang_config_builder() -> None: diff --git a/tests/integration/test_vllm_dflash2_runtime_contract.py b/tests/integration/test_vllm_dflash2_runtime_contract.py new file mode 100644 index 00000000..fe69f2ab --- /dev/null +++ b/tests/integration/test_vllm_dflash2_runtime_contract.py @@ -0,0 +1,277 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""vLLM co-training contract for the DFlash2 drafter (served through DFlash).""" + +from __future__ import annotations + +import json + +import pytest + +from verl_speco.integration import vllm_runtime +from verl_speco.integration.vllm_runtime import ( + _assert_vllm_supports_dflash2, + _draft_param_name_candidates, + _normalize_dflash2_runtime_aliases, + _validate_vllm_dflash2_block_size, + _validate_vllm_dflash_drafter_config, + build_vllm_speculative_config_from_drafter, +) + +_DFLASH2_CONFIG = { + "architectures": ["DFlash2DraftModel"], + "model_type": "qwen3", + "dflash_config": { + "block_size": 8, + "conv_kernel_size": 2, + "conv_group_size": 16, + "selector_rank": 256, + "selector_top_k": 16, + "mask_token_id": 151669, + "target_layer_ids": [0, 8, 16, 24, 32], + }, +} + + +def _write_drafter(tmp_path, config, name="dflash2-drafter"): + model_path = tmp_path / name + model_path.mkdir() + (model_path / "config.json").write_text(json.dumps(config), encoding="utf-8") + return model_path + + +def _drafter(model_path, **overrides): + config = { + "enable": True, + "enable_drafter_training": True, + "speculative_algorithm": "DFLASH2", + "model_path": str(model_path), + "rollout": {"spec_steps": 1, "spec_verify_tokens": 7}, + "training": {"dflash2_block_size": 8}, + "vllm": {}, + } + config.update(overrides) + return config + + +@pytest.fixture +def vllm_has_dflash2(monkeypatch): + monkeypatch.setattr(vllm_runtime, "_vllm_supports_dflash2", lambda: True) + + +def test_dflash2_speculative_config_is_the_dflash_contract( + tmp_path, vllm_has_dflash2 +) -> None: + model_path = _write_drafter(tmp_path, _DFLASH2_CONFIG) + + config = build_vllm_speculative_config_from_drafter(_drafter(model_path)) + + assert config == { + "draft_sample_method": "greedy", + "method": "dflash", + "model": str(model_path), + "num_speculative_tokens": 7, + } + + +def test_dflash2_speculative_config_needs_spec_verify_tokens( + tmp_path, vllm_has_dflash2 +) -> None: + model_path = _write_drafter(tmp_path, _DFLASH2_CONFIG) + + with pytest.raises(ValueError, match="spec_verify_tokens"): + build_vllm_speculative_config_from_drafter( + _drafter(model_path, rollout={"spec_steps": 3}) + ) + + +def test_dflash2_speculative_config_refuses_a_vllm_without_dflash2( + tmp_path, monkeypatch +) -> None: + model_path = _write_drafter(tmp_path, _DFLASH2_CONFIG) + monkeypatch.setattr(vllm_runtime, "_vllm_supports_dflash2", lambda: False) + + with pytest.raises(ValueError, match="vLLM 0.28.0"): + build_vllm_speculative_config_from_drafter(_drafter(model_path)) + + +def test_dflash2_capability_probe_is_skipped_without_vllm(monkeypatch) -> None: + monkeypatch.setattr(vllm_runtime, "_vllm_supports_dflash2", lambda: None) + _assert_vllm_supports_dflash2() + + +def test_dflash2_capability_probe_reads_the_installed_vllm(monkeypatch) -> None: + import importlib.util + + real_find_spec = importlib.util.find_spec + + def fake_find_spec(name, *args, **kwargs): + if name == "vllm": + return object() + if name == vllm_runtime._VLLM_DFLASH2_MODULE: + return None + return real_find_spec(name, *args, **kwargs) + + monkeypatch.setattr(importlib.util, "find_spec", fake_find_spec) + assert vllm_runtime._vllm_supports_dflash2() is False + + monkeypatch.setattr(importlib.util, "find_spec", lambda name, *a, **k: None) + assert vllm_runtime._vllm_supports_dflash2() is None + + +def test_dflash2_validator_rejects_a_plain_dflash_checkpoint(tmp_path) -> None: + """A DFlash checkpoint would load and silently serve without the DFlash2 modules.""" + model_path = _write_drafter( + tmp_path, + {"architectures": ["DFlashDraftModel"], **_DFLASH2_CONFIG["dflash_config"]}, + ) + + with pytest.raises(ValueError, match="DFlash2 drafter checkpoint"): + _validate_vllm_dflash_drafter_config(model_path, algorithm="DFLASH2") + + +def test_dflash2_validator_requires_the_runtime_hyperparameters(tmp_path) -> None: + config = json.loads(json.dumps(_DFLASH2_CONFIG)) + del config["dflash_config"]["selector_top_k"] + del config["dflash_config"]["conv_group_size"] + model_path = _write_drafter(tmp_path, config) + + with pytest.raises( + ValueError, match=r"missing \['conv_group_size', 'selector_top_k'\]" + ): + _validate_vllm_dflash_drafter_config(model_path, algorithm="DFLASH2") + + +def test_dflash2_validator_accepts_top_level_hyperparameters(tmp_path) -> None: + """Configs saved by the trainer keep the knobs flat; the alias patch nests them.""" + config = { + "architectures": ["Qwen3DFlash2Model"], + **_DFLASH2_CONFIG["dflash_config"], + } + model_path = _write_drafter(tmp_path, config) + + _validate_vllm_dflash_drafter_config(model_path, algorithm="DFLASH2") + + +def test_dflash2_validator_skips_a_missing_checkpoint(tmp_path) -> None: + _validate_vllm_dflash_drafter_config(tmp_path / "absent", algorithm="DFLASH2") + _validate_vllm_dflash_drafter_config(None, algorithm="DFLASH2") + + +def test_dflash2_block_size_must_match_the_engine_block( + tmp_path, vllm_has_dflash2 +) -> None: + model_path = _write_drafter(tmp_path, _DFLASH2_CONFIG) + + with pytest.raises(ValueError, match="spec_verify_tokens=16 but block_size=8"): + build_vllm_speculative_config_from_drafter( + _drafter(model_path, rollout={"spec_verify_tokens": 16}) + ) + + +def test_dflash2_block_size_falls_back_to_the_checkpoint() -> None: + config = json.loads(json.dumps(_DFLASH2_CONFIG)) + _validate_vllm_dflash2_block_size(config, {"training": {}}, 7) + with pytest.raises(ValueError, match="block_size=8"): + _validate_vllm_dflash2_block_size(config, {"training": {}}, 3) + # Training config wins over the checkpoint. + _validate_vllm_dflash2_block_size( + config, {"training": {"dflash2_block_size": 4}}, 3 + ) + # Nothing to compare against: accept. + _validate_vllm_dflash2_block_size(None, {"training": {}}, 3) + _validate_vllm_dflash2_block_size({"architectures": ["DFlash2DraftModel"]}, {}, 3) + + +def test_dflash2_codebooks_publish_under_the_engine_parameter_names() -> None: + """Trainer nn.Embedding ``.weight`` -> vLLM bare ``nn.Parameter``.""" + candidates = _draft_param_name_candidates( + "draft_model.candidate_selector.predecessor_codebook.weight" + ) + assert "model.candidate_selector.predecessor_codebook" in candidates + assert "candidate_selector.predecessor_codebook" in candidates + # The spelled-out name is still tried first for engines that keep the module. + assert candidates.index( + "candidate_selector.predecessor_codebook.weight" + ) < candidates.index("candidate_selector.predecessor_codebook") + + successor = _draft_param_name_candidates( + "module.draft_model.candidate_selector.successor_codebook.weight" + ) + assert "model.candidate_selector.successor_codebook" in successor + + # Every other DFlash2 parameter already matches the engine spelling. + projection = _draft_param_name_candidates( + "draft_model.candidate_selector.hidden_projection.weight" + ) + assert "model.candidate_selector.hidden_projection.weight" in projection + assert not any(name.endswith("hidden_projection") for name in projection) + conv = _draft_param_name_candidates( + "draft_model.layers.0.attention_conv.base_kernel" + ) + assert "model.layers.0.attention_conv.base_kernel" in conv + + +def test_dflash2_runtime_aliases_nest_flat_hyperparameters() -> None: + config = { + "architectures": ["DFlash2DraftModel"], + "block_size": 8, + "conv_kernel_size": 2, + "conv_group_size": 16, + "selector_rank": 256, + "selector_top_k": 16, + "dflash_config": {"target_layer_ids": [0, 8]}, + } + + assert _normalize_dflash2_runtime_aliases(config) is True + assert config["dflash_config"] == { + "target_layer_ids": [0, 8], + "block_size": 8, + "conv_kernel_size": 2, + "conv_group_size": 16, + "selector_rank": 256, + "selector_top_k": 16, + } + # Idempotent once nested. + assert _normalize_dflash2_runtime_aliases(config) is False + + +def test_dflash2_runtime_aliases_create_the_nested_block_and_ignore_others() -> None: + config = {"architectures": ["Qwen3DFlash2Model"], "selector_top_k": 4} + assert _normalize_dflash2_runtime_aliases(config) is True + assert config["dflash_config"] == {"selector_top_k": 4} + + dflash = {"architectures": ["DFlashDraftModel"], "selector_top_k": 4} + assert _normalize_dflash2_runtime_aliases(dflash) is False + assert "dflash_config" not in dflash + + assert ( + _normalize_dflash2_runtime_aliases({"architectures": "DFlash2DraftModel"}) + is False + ) + + +def test_dflash2_runtime_aliases_refuse_conflicts() -> None: + config = { + "architectures": ["DFlash2DraftModel"], + "selector_top_k": 16, + "dflash_config": {"selector_top_k": 8}, + } + with pytest.raises(ValueError, match="selector_top_k conflicts"): + _normalize_dflash2_runtime_aliases(config) + + with pytest.raises(TypeError, match="must be a mapping"): + _normalize_dflash2_runtime_aliases( + {"architectures": ["DFlash2DraftModel"], "dflash_config": "bad"} + ) diff --git a/tests/unit/test_convert_speculators_dflash2.py b/tests/unit/test_convert_speculators_dflash2.py new file mode 100644 index 00000000..c0e235c8 --- /dev/null +++ b/tests/unit/test_convert_speculators_dflash2.py @@ -0,0 +1,215 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import json +import os + +import pytest + +from verl_speco.convert_speculators_dflash2 import ( + convert_speculators_dflash2_checkpoint, + convert_speculators_dflash2_config, + main, +) + + +def _speculators_config(**overrides): + config = { + "architectures": ["DFlash2DraftModel"], + "aux_hidden_state_layer_ids": [1, 9, 17, 25, 33], + "block_size": 8, + "conv_group_size": 16, + "conv_kernel_size": 2, + "draft_vocab_size": 151936, + "dtype": "bfloat16", + "mask_token_id": 151669, + "selector_rank": 256, + "selector_top_k": 16, + "speculators_config": {"algorithm": "dflash2"}, + "speculators_model_type": "dflash2", + "transformer_layer_config": { + "attention_bias": False, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "intermediate_size": 9728, + "layer_types": ["sliding_attention"] * 5, + "max_position_embeddings": 40960, + "max_window_layers": 28, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-6, + "rope_parameters": {"rope_theta": 1000000, "rope_type": "default"}, + "sliding_window": 2048, + "use_sliding_window": True, + "vocab_size": 151936, + }, + } + config.update(overrides) + return config + + +_TARGET_CONFIG = { + "num_hidden_layers": 36, + "bos_token_id": 151643, + "eos_token_id": 151645, + "vocab_size": 151936, +} + + +def test_convert_emits_the_zlab_layout() -> None: + converted = convert_speculators_dflash2_config( + _speculators_config(), _TARGET_CONFIG + ) + + assert converted["architectures"] == ["DFlash2DraftModel"] + assert converted["model_type"] == "qwen3" + assert converted["is_causal"] is False + assert converted["num_target_layers"] == 36 + assert converted["hidden_size"] == 2560 + assert converted["num_hidden_layers"] == 5 + assert converted["rope_theta"] == 1000000 + assert converted["rope_parameters"] == { + "rope_theta": 1000000, + "rope_type": "default", + } + assert converted["eos_token_id"] == 151645 + assert converted["bos_token_id"] == 151643 + assert converted["dflash_config"] == { + "block_size": 8, + "conv_group_size": 16, + "conv_kernel_size": 2, + "mask_token_id": 151669, + "selector_rank": 256, + "selector_top_k": 16, + # hidden_states indices minus one: decoder-layer indices. + "target_layer_ids": [0, 8, 16, 24, 32], + } + # Speculators keys must not leak into the runtime contract. + assert "speculators_config" not in converted + assert "transformer_layer_config" not in converted + assert "aux_hidden_state_layer_ids" not in converted + + +def test_convert_defaults_to_full_attention_but_can_keep_the_window() -> None: + full = convert_speculators_dflash2_config(_speculators_config(), _TARGET_CONFIG) + assert full["use_sliding_window"] is False + assert full["sliding_window"] is None + assert full["layer_types"] == ["full_attention"] * 5 + + windowed = convert_speculators_dflash2_config( + _speculators_config(), _TARGET_CONFIG, keep_sliding_window=True + ) + assert windowed["use_sliding_window"] is True + assert windowed["sliding_window"] == 2048 + assert windowed["layer_types"] == ["sliding_attention"] * 5 + + +def test_convert_prefers_explicit_target_layer_ids() -> None: + converted = convert_speculators_dflash2_config( + _speculators_config(target_layer_ids=[2, 4, 6, 8, 10]), _TARGET_CONFIG + ) + assert converted["dflash_config"]["target_layer_ids"] == [2, 4, 6, 8, 10] + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ( + {"speculators_model_type": "eagle3", "speculators_config": {}}, + "not a speculators DFlash2", + ), + ({"transformer_layer_config": None}, "transformer_layer_config"), + ({"aux_hidden_state_layer_ids": [1, 9, 17, 25, 40]}, "exceed the target"), + ({"aux_hidden_state_layer_ids": [0, 9, 17, 25, 33]}, "must be >= 1"), + ({"draft_vocab_size": 32000}, "share the target vocabulary"), + ({"mask_token_id": None}, "lacks mask_token_id"), + ({"selector_rank": None}, "lacks \\['selector_rank'\\]"), + ], +) +def test_convert_rejects_malformed_configs(overrides, match) -> None: + with pytest.raises(ValueError, match=match): + convert_speculators_dflash2_config( + _speculators_config(**overrides), _TARGET_CONFIG + ) + + +def test_convert_rejects_missing_layer_ids() -> None: + config = _speculators_config() + del config["aux_hidden_state_layer_ids"] + with pytest.raises(ValueError, match="neither target_layer_ids"): + convert_speculators_dflash2_config(config, _TARGET_CONFIG) + + +def _write_dirs(tmp_path): + source = tmp_path / "speculators" + target = tmp_path / "target" + source.mkdir() + target.mkdir() + (source / "config.json").write_text( + json.dumps(_speculators_config()), encoding="utf-8" + ) + (source / "model.safetensors").write_bytes(b"weights") + (target / "config.json").write_text(json.dumps(_TARGET_CONFIG), encoding="utf-8") + return source, target + + +def test_convert_checkpoint_writes_config_and_copies_weights(tmp_path) -> None: + source, target = _write_dirs(tmp_path) + output = tmp_path / "converted" + + converted = convert_speculators_dflash2_checkpoint( + str(source), str(target), str(output) + ) + + written = json.loads((output / "config.json").read_text(encoding="utf-8")) + assert written == converted + assert (output / "model.safetensors").read_bytes() == b"weights" + assert not os.path.islink(output / "model.safetensors") + + +def test_convert_checkpoint_can_symlink_weights_and_overwrites(tmp_path) -> None: + source, target = _write_dirs(tmp_path) + output = tmp_path / "converted" + output.mkdir() + (output / "model.safetensors").write_bytes(b"stale") + + convert_speculators_dflash2_checkpoint( + str(source), str(target), str(output), link_weights=True + ) + + assert os.path.islink(output / "model.safetensors") + assert (output / "model.safetensors").read_bytes() == b"weights" + + +def test_convert_checkpoint_requires_weights(tmp_path) -> None: + source, target = _write_dirs(tmp_path) + (source / "model.safetensors").unlink() + with pytest.raises(FileNotFoundError, match="no safetensors"): + convert_speculators_dflash2_checkpoint( + str(source), str(target), str(tmp_path / "out") + ) + + +def test_main_converts_and_reports(tmp_path, capsys) -> None: + source, target = _write_dirs(tmp_path) + output = tmp_path / "converted" + + main(["--input", str(source), "--target", str(target), "--output", str(output)]) + + assert (output / "config.json").exists() + assert "DFlash2DraftModel" in capsys.readouterr().out diff --git a/verl_speco/convert_speculators_dflash2.py b/verl_speco/convert_speculators_dflash2.py new file mode 100644 index 00000000..6d56cec9 --- /dev/null +++ b/verl_speco/convert_speculators_dflash2.py @@ -0,0 +1,275 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Convert a speculators-format DFlash2 drafter into the z-lab checkpoint layout. + +Public DFlash2 drafters trained with the ``speculators`` library (for example +``mgoin/Qwen3-4B-speculator.dflash2``) ship a ``config.json`` that nests the +transformer hyperparameters under ``transformer_layer_config`` and describes the +target through ``speculators_config``. Both this overlay's trainer +(``DFlash2Config.from_dflash2_pretrained``) and vLLM releases up to 0.28.0 read +the released z-lab layout instead: a flat Qwen3-style config with the DFlash2 +knobs under ``dflash_config``. The weights use the same parameter names in both +layouts, so only ``config.json`` is rewritten; the safetensors are copied as-is. + +Run:: + + python -m verl_speco.convert_speculators_dflash2 \ + --input /path/to/speculators-dflash2 \ + --target /path/to/target-model \ + --output /path/to/dflash2-drafter +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import shutil +from typing import Any + +# Transformer hyperparameters copied verbatim from ``transformer_layer_config``. +_TRANSFORMER_KEYS = ( + "attention_bias", + "attention_dropout", + "head_dim", + "hidden_act", + "hidden_size", + "initializer_range", + "intermediate_size", + "max_position_embeddings", + "num_attention_heads", + "num_hidden_layers", + "num_key_value_heads", + "rms_norm_eps", + "rope_scaling", + "rope_theta", + "rope_parameters", + "vocab_size", +) +# Sliding-window attention knobs. The trainer's DFlash backbone attends over the +# full anchored context, so the converted drafter defaults to full attention to +# keep training and rollout on the same model. +_SLIDING_WINDOW_KEYS = ( + "layer_types", + "max_window_layers", + "sliding_window", + "use_sliding_window", +) +_DFLASH2_KEYS = ( + "block_size", + "conv_kernel_size", + "conv_group_size", + "selector_rank", + "selector_top_k", +) +_WEIGHT_PATTERNS = ("*.safetensors", "*.safetensors.index.json") + + +def _load_json(path: str) -> dict[str, Any]: + with open(path, "r", encoding="utf-8") as f: + loaded = json.load(f) + if not isinstance(loaded, dict): + raise ValueError(f"{path} does not hold a JSON object") + return loaded + + +def _target_layer_ids(speculators_config: dict[str, Any]) -> list[int]: + """The DFlash context layers, as decoder-layer indices. + + speculators records ``aux_hidden_state_layer_ids`` as indices into + ``output_hidden_states`` (entry 0 is the embedding output), while the z-lab + layout's ``target_layer_ids`` count decoder layers, hence the ``- 1``; this + matches the ``eagle_aux_hidden_state_layer_ids = target_layer_ids + 1`` alias + vLLM applies in the other direction. + """ + if "target_layer_ids" in speculators_config: + return [int(layer_id) for layer_id in speculators_config["target_layer_ids"]] + aux_layer_ids = speculators_config.get("aux_hidden_state_layer_ids") + if not aux_layer_ids: + raise ValueError( + "speculators DFlash2 config has neither target_layer_ids nor " + "aux_hidden_state_layer_ids" + ) + layer_ids = [int(layer_id) - 1 for layer_id in aux_layer_ids] + if min(layer_ids) < 0: + raise ValueError( + f"aux_hidden_state_layer_ids={aux_layer_ids!r} must be >= 1 (hidden_states indices)" + ) + return layer_ids + + +def convert_speculators_dflash2_config( + speculators_config: dict[str, Any], + target_config: dict[str, Any], + *, + keep_sliding_window: bool = False, +) -> dict[str, Any]: + """Build the z-lab layout ``config.json`` for a speculators DFlash2 drafter.""" + model_type = str(speculators_config.get("speculators_model_type") or "").lower() + algorithm = str( + (speculators_config.get("speculators_config") or {}).get("algorithm") or "" + ).lower() + if "dflash2" not in (model_type, algorithm): + raise ValueError( + "not a speculators DFlash2 config: speculators_model_type=" + f"{model_type!r} algorithm={algorithm!r}" + ) + transformer = speculators_config.get("transformer_layer_config") + if not isinstance(transformer, dict): + raise ValueError("speculators DFlash2 config lacks transformer_layer_config") + + target_text = target_config.get("text_config") or target_config + target_num_hidden_layers = int(target_text["num_hidden_layers"]) + target_layer_ids = _target_layer_ids(speculators_config) + if max(target_layer_ids) >= target_num_hidden_layers: + raise ValueError( + f"target_layer_ids={target_layer_ids} exceed the target's " + f"num_hidden_layers={target_num_hidden_layers}" + ) + + converted: dict[str, Any] = { + "architectures": ["DFlash2DraftModel"], + "model_type": str(transformer.get("model_type") or "qwen3"), + "is_causal": False, + "tie_word_embeddings": False, + "use_cache": True, + "dtype": str(speculators_config.get("dtype") or "bfloat16"), + } + for key in _TRANSFORMER_KEYS: + if transformer.get(key) is not None: + converted[key] = transformer[key] + rope_parameters = transformer.get("rope_parameters") + if isinstance(rope_parameters, dict) and "rope_theta" not in converted: + # transformers < 5 reads rope_theta at the top level. + if rope_parameters.get("rope_theta") is not None: + converted["rope_theta"] = rope_parameters["rope_theta"] + if keep_sliding_window: + for key in _SLIDING_WINDOW_KEYS: + if key in transformer: + converted[key] = transformer[key] + else: + converted["use_sliding_window"] = False + converted["sliding_window"] = None + converted["layer_types"] = ["full_attention"] * int( + converted["num_hidden_layers"] + ) + for key in ("bos_token_id", "eos_token_id", "pad_token_id"): + value = transformer.get(key) + if value is None: + value = target_text.get(key) + converted[key] = value + + draft_vocab_size = speculators_config.get("draft_vocab_size") + if draft_vocab_size is not None and int(draft_vocab_size) != int( + converted["vocab_size"] + ): + raise ValueError( + "DFlash2 drafts share the target vocabulary; " + f"draft_vocab_size={draft_vocab_size} != vocab_size={converted['vocab_size']}" + ) + + mask_token_id = speculators_config.get("mask_token_id") + if mask_token_id is None: + raise ValueError("speculators DFlash2 config lacks mask_token_id") + dflash_config = { + key: speculators_config[key] + for key in _DFLASH2_KEYS + if speculators_config.get(key) is not None + } + missing = sorted(set(_DFLASH2_KEYS) - set(dflash_config)) + if missing: + raise ValueError(f"speculators DFlash2 config lacks {missing}") + dflash_config["mask_token_id"] = int(mask_token_id) + dflash_config["target_layer_ids"] = target_layer_ids + converted["dflash_config"] = dflash_config + converted["num_target_layers"] = target_num_hidden_layers + return converted + + +def convert_speculators_dflash2_checkpoint( + input_dir: str, + target_dir: str, + output_dir: str, + *, + keep_sliding_window: bool = False, + link_weights: bool = False, +) -> dict[str, Any]: + """Write the converted drafter to ``output_dir`` and return its config.""" + speculators_config = _load_json(os.path.join(input_dir, "config.json")) + target_config = _load_json(os.path.join(target_dir, "config.json")) + converted = convert_speculators_dflash2_config( + speculators_config, target_config, keep_sliding_window=keep_sliding_window + ) + + weight_files = sorted( + path + for pattern in _WEIGHT_PATTERNS + for path in glob.glob(os.path.join(input_dir, pattern)) + ) + if not weight_files: + raise FileNotFoundError(f"no safetensors weights under {input_dir}") + + os.makedirs(output_dir, exist_ok=True) + with open(os.path.join(output_dir, "config.json"), "w", encoding="utf-8") as f: + json.dump(converted, f, indent=2, sort_keys=True) + f.write("\n") + for source in weight_files: + destination = os.path.join(output_dir, os.path.basename(source)) + if os.path.lexists(destination): + os.remove(destination) + if link_weights: + os.symlink(os.path.abspath(source), destination) + else: + shutil.copy2(source, destination) + return converted + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument( + "--input", required=True, help="speculators-format DFlash2 drafter directory" + ) + parser.add_argument( + "--target", + required=True, + help="target model directory (its config.json supplies the layer count and special tokens)", + ) + parser.add_argument("--output", required=True, help="converted drafter directory") + parser.add_argument( + "--keep-sliding-window", + action="store_true", + help="keep the speculators sliding-window attention settings instead of full attention", + ) + parser.add_argument( + "--link-weights", + action="store_true", + help="symlink the safetensors into the output instead of copying them", + ) + args = parser.parse_args(argv) + converted = convert_speculators_dflash2_checkpoint( + args.input, + args.target, + args.output, + keep_sliding_window=args.keep_sliding_window, + link_weights=args.link_weights, + ) + print( + f"wrote {args.output}: architectures={converted['architectures']} " + f"dflash_config={json.dumps(converted['dflash_config'], sort_keys=True)}" + ) + + +if __name__ == "__main__": + main() diff --git a/verl_speco/integration/rollout_publish.py b/verl_speco/integration/rollout_publish.py index fdedb93b..104f0d98 100644 --- a/verl_speco/integration/rollout_publish.py +++ b/verl_speco/integration/rollout_publish.py @@ -745,8 +745,10 @@ def export_actor_lm_head_weight( ) normalized_row_indices = _normalize_lm_head_row_indices(row_indices) + # Block drafters that train against the target's own lm_head rows. is_dflash = drafter_speculative_algorithm(getattr(worker, "config", None)) in { "DFLASH", + "DFLASH2", "DSPARK", } if ( diff --git a/verl_speco/integration/sglang_runtime.py b/verl_speco/integration/sglang_runtime.py index 6a2edffa..3b629120 100644 --- a/verl_speco/integration/sglang_runtime.py +++ b/verl_speco/integration/sglang_runtime.py @@ -641,11 +641,14 @@ def _server_args_overrides_from_drafter( # (dynamic convolutions + candidate selector) ride in the checkpoint's # dflash_config, not a distinct engine-level method. DFLASH2 is never a # valid SGLang ServerArgs algorithm, so fail loud instead of forwarding - # the raw string. + # the raw string. The vLLM overlay maps DFLASH2 onto its DFlash path + # (vllm_runtime._speculative_method_from_drafter); the SGLang co-train + # path is not validated for DFlash2 yet. raise ValueError( "DFLASH2 is not an engine-level speculative algorithm; DFlash2 is served as a DFlash " - "checkpoint. Keep DFLASH2 for drafter training (which this overlay runs offline) and " - "set actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH to serve a trained " + "checkpoint. SGLang co-training of a DFlash2 drafter is not wired in this overlay: use " + "actor_rollout_ref.rollout.name=vllm (vLLM >= 0.28.0) for DFLASH2 co-training, or set " + "actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH to serve a trained " "DFlash2 checkpoint as a frozen rollout drafter; its dflash_config carries the DFlash2 " "convolution and selector hyperparameters." ) diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py index e47273b6..2aeba1b5 100644 --- a/verl_speco/integration/vllm_runtime.py +++ b/verl_speco/integration/vllm_runtime.py @@ -27,6 +27,7 @@ import json import logging import os +import re import sys import threading import time @@ -814,31 +815,173 @@ def _drafter_algorithm(drafter_cfg: dict[str, Any]) -> str: # Draft architectures vLLM can serve through its DFlash speculative path. # Keep in sync with the alias sets in verl_speco/models/auto.py. +_DFLASH2_SERVABLE_ARCHITECTURES = frozenset({"DFlash2DraftModel", "Qwen3DFlash2Model"}) _DFLASH_SERVABLE_ARCHITECTURES = frozenset( - {"DFlashDraftModel", "DFlash2DraftModel", "Qwen3DFlash2Model"} + {"DFlashDraftModel", *_DFLASH2_SERVABLE_ARCHITECTURES} +) +# Hyperparameters vLLM's DFlash2 draft (qwen3_dflash2.py) indexes straight out of +# ``dflash_config`` when it builds the model; a missing key is a KeyError deep in +# engine startup, so the checkpoint contract is checked up front instead. +_DFLASH2_RUNTIME_KEYS = ( + "conv_kernel_size", + "conv_group_size", + "selector_rank", + "selector_top_k", +) +# vLLM module that carries the DFlash2 draft class; its presence is the +# capability probe for DFlash2 rollout (vllm-project/vllm#52816, v0.28.0+). +_VLLM_DFLASH2_MODULE = "vllm.model_executor.models.qwen3_dflash2" +# Trainer-side spelling of the DFlash2 selector codebooks (see +# ``_draft_param_name_candidates`` for the engine-side rename). +_DFLASH2_CODEBOOK_SUFFIX = re.compile( + r"candidate_selector\.(predecessor|successor)_codebook\.weight$" ) -def _validate_vllm_dflash_drafter_config( - spec_model_path: Any, algorithm: str = "DFLASH" +def _normalize_dflash2_runtime_aliases(config: Any) -> bool: + """Mirror DFlash2 hyperparameters into ``dflash_config`` for vLLM. + + vLLM's DFlash2 draft reads the convolution and selector knobs strictly from + ``dflash_config`` (the z-lab checkpoint layout), while a checkpoint this + overlay saved carries them at the top level. Copy missing keys down so both + layouts serve; a conflicting pair is a corrupt checkpoint and fails loud. + """ + architectures = _get_nested(config, ("architectures",), None) or [] + if isinstance(architectures, str): + architectures = [architectures] + if _DFLASH2_SERVABLE_ARCHITECTURES.isdisjoint(str(name) for name in architectures): + return False + + dflash_config = _get_nested(config, ("dflash_config",), None) + if dflash_config is not None and not hasattr(dflash_config, "get"): + raise TypeError("DFlash2 dflash_config must be a mapping when provided") + changed = False + for key in (*_DFLASH2_RUNTIME_KEYS, "block_size"): + top_level = _get_nested(config, (key,), None) + nested = _get_nested(dflash_config, (key,), None) + if top_level is None or top_level == nested: + continue + if nested is not None: + raise ValueError( + f"DFlash2 {key} conflicts with dflash_config.{key}: {top_level!r} != {nested!r}" + ) + if dflash_config is None: + dflash_config = {} + _set_child(config, "dflash_config", dflash_config) + _set_child(dflash_config, key, top_level) + changed = True + return changed + + +def _vllm_supports_dflash2() -> bool | None: + """Whether the installed vLLM ships the DFlash2 draft; ``None`` without vLLM.""" + import importlib.util + + try: + if importlib.util.find_spec("vllm") is None: + return None + return importlib.util.find_spec(_VLLM_DFLASH2_MODULE) is not None + except (ImportError, ValueError): + return None + + +def _assert_vllm_supports_dflash2() -> None: + if _vllm_supports_dflash2() is False: + raise ValueError( + "DFLASH2 rollout needs a vLLM that ships the DFlash2 draft model " + f"({_VLLM_DFLASH2_MODULE}, vllm-project/vllm#52816, released in vLLM 0.28.0); " + "the installed vLLM does not have it. Upgrade vLLM, or keep " + "actor_rollout_ref.rollout.drafter.enable=false and train the DFlash2 drafter offline." + ) + + +def _dflash2_config_value(config: dict[str, Any], key: str) -> Any: + """Read a DFlash2 knob from the checkpoint config, top level first.""" + if key in config: + return config[key] + nested = config.get("dflash_config") + if isinstance(nested, dict): + return nested.get(key) + return None + + +def _validate_vllm_dflash2_block_size( + config: dict[str, Any] | None, + drafter_cfg: dict[str, Any], + num_speculative_tokens: int, ) -> None: - if not spec_model_path: + """Pin the DFlash2 block to the engine's ``1 + num_speculative_tokens``. + + The dynamic convolutions are causal *within* a block: vLLM sizes that block + as the bonus token plus ``num_speculative_tokens`` mask tokens, while the + trainer folds the drafted sequence by ``dflash2_block_size``. If the two + disagree, the served conv reads across positions the trained conv never saw + (or vice versa), which shows up as a silently weak drafter rather than an + error, so refuse the mismatch here. + """ + training_cfg = drafter_cfg.get("training") or {} + block_size = _positive_int_or_none(training_cfg.get("dflash2_block_size")) + if block_size is None and config is not None: + block_size = _positive_int_or_none(_dflash2_config_value(config, "block_size")) + if block_size is None: return + if int(num_speculative_tokens) + 1 != int(block_size): + raise ValueError( + "DFLASH2 rollout requires actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens " + "== block_size - 1 so the served convolution block matches the trained one: got " + f"spec_verify_tokens={num_speculative_tokens} but block_size={block_size} " + "(from drafter.training.dflash2_block_size or the checkpoint's dflash_config)." + ) + +def _load_vllm_dflash_drafter_config(spec_model_path: Any) -> dict[str, Any] | None: + if not spec_model_path: + return None config_path = os.path.join(os.fspath(spec_model_path), "config.json") if not os.path.exists(config_path): - return - + return None try: with open(config_path, "r", encoding="utf-8") as f: - config = json.load(f) + return json.load(f) except json.JSONDecodeError as exc: raise ValueError( f"Invalid DFlash drafter config.json at {config_path}: {exc}" ) from exc + +def _validate_vllm_dflash_drafter_config( + spec_model_path: Any, algorithm: str = "DFLASH" +) -> None: + config = _load_vllm_dflash_drafter_config(spec_model_path) + if config is None: + return + config_path = os.path.join(os.fspath(spec_model_path), "config.json") + architectures = config.get("architectures") or [] algorithm = str(algorithm or "DFLASH").strip().upper() + if algorithm == "DFLASH2": + # vLLM dispatches on the architecture: only the DFlash2 names reach the + # draft class with the convolutions and the selector. A plain DFlash + # checkpoint would load fine and silently serve without them. + if _DFLASH2_SERVABLE_ARCHITECTURES.isdisjoint(architectures): + raise ValueError( + "vLLM DFLASH2 requires actor_rollout_ref.rollout.drafter.model_path to point " + "to a DFlash2 drafter checkpoint with architectures in " + f"{sorted(_DFLASH2_SERVABLE_ARCHITECTURES)}; got architectures={architectures!r} " + f"from {config_path}. Use speculative_algorithm=DFLASH for a plain DFlash drafter." + ) + missing = [ + key + for key in _DFLASH2_RUNTIME_KEYS + if _dflash2_config_value(config, key) is None + ] + if missing: + raise ValueError( + "vLLM DFLASH2 requires the drafter config.json to carry the DFlash2 " + f"hyperparameters {list(_DFLASH2_RUNTIME_KEYS)} (top level or under dflash_config); " + f"missing {missing} in {config_path}." + ) + return if algorithm == "DSPARK": if not _is_dspark_config(config): raise ValueError( @@ -922,18 +1065,15 @@ def _speculative_method_from_drafter(drafter_cfg: dict[str, Any]) -> str: "drafter training." ) if algorithm == "DFLASH2": - # Same story as Domino: DFlash2 is a DFlash variant whose extra modules - # (dynamic convolutions + candidate selector) ride in the checkpoint's - # dflash_config, not a distinct engine-level method. DFLASH2 is never a - # valid vLLM method, so fail loud instead of forwarding the raw string, - # mirroring sglang_runtime._server_args_overrides_from_drafter. - raise ValueError( - "DFLASH2 is not an engine-level speculative algorithm; DFlash2 is served as a DFlash " - "checkpoint. Keep DFLASH2 for drafter training (which this overlay runs offline) and " - "set actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH to serve a trained " - "DFlash2 checkpoint as a frozen rollout drafter; its dflash_config carries the DFlash2 " - "convolution and selector hyperparameters." - ) + # DFlash2 is a DFlash variant, not an engine-level method: vLLM runs it + # through the DFlash proposer (method="dflash") and picks the DFlash2 draft + # class (dynamic convolutions + candidate selector, vllm-project/vllm#52816, + # first released in v0.28.0) from the checkpoint's ``DFlash2DraftModel`` + # architecture. DFLASH2 itself is never a valid vLLM method string, so map + # it here; the checkpoint contract is enforced by + # ``_validate_vllm_dflash_drafter_config`` and the engine capability by + # ``_assert_vllm_supports_dflash2``. + return "dflash" if algorithm == "DSPARK": return "dflash" if _is_vllm_ascend_runtime_hint() else "dspark" @@ -1070,6 +1210,13 @@ def build_vllm_speculative_config_from_drafter( "actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens " f"must be positive for vLLM {method.upper()} speculative decoding" ) + if algorithm == "DFLASH2": + _assert_vllm_supports_dflash2() + _validate_vllm_dflash2_block_size( + _load_vllm_dflash_drafter_config(spec_model_path), + drafter_cfg, + num_speculative_tokens, + ) else: num_speculative_tokens = _positive_int_or_none( _first_present( @@ -1728,6 +1875,7 @@ def patched_eagle_config_init(self, *args, **kwargs): current(self, *args, **kwargs) if str(method or "eagle").strip().lower() == "dflash": _normalize_dflash_target_layer_aliases(self) + _normalize_dflash2_runtime_aliases(self) if _is_dspark_hf_config(self): _set_child(self, "architectures", ["DFlashDraftModel"]) @@ -2222,6 +2370,11 @@ def _draft_param_name_candidates(name: str) -> list[str]: candidates.append(candidate) if "midlayer." in candidate: candidates.append(candidate.replace("midlayer.", "model.layers.0.")) + # The trainer keeps the DFlash2 selector codebooks in nn.Embedding + # modules (``..._codebook.weight``); vLLM's CandidateSelector holds them as + # bare parameters (``..._codebook``), like the released z-lab checkpoints. + if _DFLASH2_CODEBOOK_SUFFIX.search(candidate): + candidates.append(candidate[: -len(".weight")]) for candidate in list(candidates): if not candidate.startswith("model."): candidates.append(f"model.{candidate}") diff --git a/verl_speco/models/dflash2/configuration_dflash2.py b/verl_speco/models/dflash2/configuration_dflash2.py index 6dd95b3f..2c0bce9c 100644 --- a/verl_speco/models/dflash2/configuration_dflash2.py +++ b/verl_speco/models/dflash2/configuration_dflash2.py @@ -75,6 +75,24 @@ def __init__( self.selector_top_k = int(selector_top_k) self.selector_loss_weight = float(selector_loss_weight) + def to_dict(self): + """Serialize with the z-lab ``dflash_config`` block alongside the flat keys. + + vLLM's DFlash2 draft reads the convolution and selector knobs strictly + from ``dflash_config``, so a checkpoint saved by the trainer has to carry + that block to be servable as a rollout drafter. The flat keys stay for + this overlay's own loaders; ``from_dflash2_pretrained`` lets a top-level + value win over the nested one, so the two never disagree. + """ + output = super().to_dict() + nested = dict(output.get("dflash_config") or {}) + for key in _NESTED_DFLASH_KEYS: + if output.get(key) is not None: + nested[key] = output[key] + if nested: + output["dflash_config"] = nested + return output + @classmethod def from_dflash2_pretrained(cls, model_path: str): config_path = os.path.join(model_path, "config.json") From c629f9222663cc3ca9f7de28c8faa22bd78b63d5 Mon Sep 17 00:00:00 2001 From: khazic Date: Tue, 1 Sep 2026 21:20:07 +0800 Subject: [PATCH 2/6] ci(dflash2): pass the DFlash2 training knobs in the example runner Signed-off-by: khazic --- ci/run_example_test.sh | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/ci/run_example_test.sh b/ci/run_example_test.sh index 23557a39..a0c4e2a0 100644 --- a/ci/run_example_test.sh +++ b/ci/run_example_test.sh @@ -239,6 +239,15 @@ if [[ "${drafter}" == "dflash" ]]; then ) fi +if [[ "${drafter}" == "dflash2" ]]; then + overrides+=( + "actor_rollout_ref.rollout.drafter.training.dflash2_block_size=${SPECO_DFLASH2_BLOCK_SIZE:-8}" + "actor_rollout_ref.rollout.drafter.training.dflash2_num_anchors=${SPECO_DFLASH2_NUM_ANCHORS:-8}" + "actor_rollout_ref.rollout.drafter.training.dflash2_loss_decay_gamma=${SPECO_DFLASH2_LOSS_DECAY_GAMMA:-7}" + "actor_rollout_ref.rollout.drafter.training.dflash_max_window=${SPECO_DFLASH_MAX_WINDOW:-64}" + ) +fi + if [[ "${drafter}" == "dspark" ]]; then overrides+=( "actor_rollout_ref.rollout.drafter.rollout.spec_steps=${SPECO_DSPARK_SPEC_STEPS:-1}" From 72fa480e8b3096132e9c1030173a99b31237e1b1 Mon Sep 17 00:00:00 2001 From: khazic Date: Tue, 1 Sep 2026 22:51:45 +0800 Subject: [PATCH 3/6] fix(vllm): stage the draft publish IPC buckets in non-expandable segments 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 --- ...st_vllm_draft_update_allocator_contract.py | 64 +++++++++++++++++++ verl_speco/integration/vllm_runtime.py | 42 ++++++++++-- 2 files changed, 100 insertions(+), 6 deletions(-) create mode 100644 tests/integration/test_vllm_draft_update_allocator_contract.py diff --git a/tests/integration/test_vllm_draft_update_allocator_contract.py b/tests/integration/test_vllm_draft_update_allocator_contract.py new file mode 100644 index 00000000..1124c0b6 --- /dev/null +++ b/tests/integration/test_vllm_draft_update_allocator_contract.py @@ -0,0 +1,64 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The vLLM draft publish stages its CUDA-IPC buckets in non-expandable segments.""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from verl_speco.integration.vllm_runtime import _ipc_safe_allocator + + +@pytest.fixture +def fake_verl_device(monkeypatch): + calls: list[bool] = [] + # The CPU CI runs without verl installed; stub the parents so the import + # machinery reaches the fake leaf module. + for parent in ("verl", "verl.utils"): + if parent not in sys.modules: + monkeypatch.setitem(sys.modules, parent, types.ModuleType(parent)) + module = types.ModuleType("verl.utils.device") + module.set_expandable_segments = calls.append # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "verl.utils.device", module) + return calls + + +def test_ipc_send_disables_expandable_segments_and_restores_them( + fake_verl_device, +) -> None: + with _ipc_safe_allocator(True): + assert fake_verl_device == [False] + assert fake_verl_device == [False, True] + + +def test_ipc_send_restores_expandable_segments_on_failure(fake_verl_device) -> None: + with pytest.raises(RuntimeError, match="send failed"): + with _ipc_safe_allocator(True): + raise RuntimeError("send failed") + assert fake_verl_device == [False, True] + + +def test_shm_transport_leaves_the_allocator_alone(fake_verl_device) -> None: + with _ipc_safe_allocator(False): + pass + assert fake_verl_device == [] + + +def test_verl_without_the_helper_is_a_no_op(monkeypatch) -> None: + monkeypatch.setitem(sys.modules, "verl.utils.device", None) + with _ipc_safe_allocator(True): + pass diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py index 2aeba1b5..7099d3eb 100644 --- a/verl_speco/integration/vllm_runtime.py +++ b/verl_speco/integration/vllm_runtime.py @@ -2438,6 +2438,33 @@ async def _maybe_call_vllm_server_method( return await method.remote(*args, **kwargs) +@contextmanager +def _ipc_safe_allocator(enabled: bool): + """Stage the IPC buckets in non-expandable CUDA segments. + + CUDA tensors shared over IPC out of an expandable segment carry an fd-based + handle that the receiver can only import through ``pidfd_getfd`` (Linux >= + 5.6); on older kernels the vLLM worker fails the whole draft update with + "does not support the pidfd_getfd syscall". verl's own actor->rollout sync + flips expandable segments off around its send for the same reason and turns + them back on afterwards, so the draft publish mirrors that. A verl without + the helper never enabled expandable segments, so there is nothing to do. + """ + if not enabled: + yield + return + try: + from verl.utils.device import set_expandable_segments + except ImportError: + yield + return + set_expandable_segments(False) + try: + yield + finally: + set_expandable_segments(True) + + async def speco_vllm_update_draft_weights( self, weights: Any, *args, global_steps: int | None = None, **kwargs ): @@ -2494,12 +2521,15 @@ async def speco_vllm_update_draft_weights( kwargs={**kwargs, "use_shm": use_shm}, ) - sender = BucketedWeightSender( - zmq_handle=_draft_zmq_handle_from_base(self.zmq_handle), - bucket_size_mb=int(bucket_mb), - use_shm=use_shm, - ) - await sender.async_send_weights(_named_weight_iter(weights)) + # Only the CUDA-IPC transport shares device allocations; shm stages + # through host memory and is unaffected by the allocator mode. + with _ipc_safe_allocator(not use_shm): + sender = BucketedWeightSender( + zmq_handle=_draft_zmq_handle_from_base(self.zmq_handle), + bucket_size_mb=int(bucket_mb), + use_shm=use_shm, + ) + await sender.async_send_weights(_named_weight_iter(weights)) if future is not None: await future From 8a997da1cbc84e1c7092ad1a51437a195c528682 Mon Sep 17 00:00:00 2001 From: khazic Date: Tue, 1 Sep 2026 23:19:13 +0800 Subject: [PATCH 4/6] fix(dflash2): publish the selector codebooks under vLLM's parameter names 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 --- .../test_vllm_dflash2_runtime_contract.py | 21 ++++++++++++++ verl_speco/integration/vllm_runtime.py | 28 +++++++++++++++---- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/tests/integration/test_vllm_dflash2_runtime_contract.py b/tests/integration/test_vllm_dflash2_runtime_contract.py index fe69f2ab..9bc94024 100644 --- a/tests/integration/test_vllm_dflash2_runtime_contract.py +++ b/tests/integration/test_vllm_dflash2_runtime_contract.py @@ -22,6 +22,7 @@ from verl_speco.integration import vllm_runtime from verl_speco.integration.vllm_runtime import ( _assert_vllm_supports_dflash2, + _dflash2_engine_param_name, _draft_param_name_candidates, _normalize_dflash2_runtime_aliases, _validate_vllm_dflash2_block_size, @@ -275,3 +276,23 @@ def test_dflash2_runtime_aliases_refuse_conflicts() -> None: _normalize_dflash2_runtime_aliases( {"architectures": ["DFlash2DraftModel"], "dflash_config": "bad"} ) + + +def test_dflash2_engine_param_name_renames_only_the_codebooks() -> None: + """The IPC receiver hands vLLM's load_weights the engine spelling.""" + assert ( + _dflash2_engine_param_name("candidate_selector.predecessor_codebook.weight") + == "candidate_selector.predecessor_codebook" + ) + assert ( + _dflash2_engine_param_name("candidate_selector.successor_codebook.weight") + == "candidate_selector.successor_codebook" + ) + for untouched in ( + "candidate_selector.hidden_projection.weight", + "candidate_selector.predecessor_codebook", + "layers.0.attention_conv.base_kernel", + "layers.0.mlp_conv.kernel_projection.weight", + "fc.weight", + ): + assert _dflash2_engine_param_name(untouched) == untouched diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py index 7099d3eb..bc622dc8 100644 --- a/verl_speco/integration/vllm_runtime.py +++ b/verl_speco/integration/vllm_runtime.py @@ -838,6 +838,22 @@ def _drafter_algorithm(drafter_cfg: dict[str, Any]) -> str: ) +def _dflash2_engine_param_name(name: str) -> str: + """Spell a published DFlash2 parameter the way vLLM's draft names it. + + The trainer keeps the selector codebooks in ``nn.Embedding`` modules + (``..._codebook.weight``) while vLLM's ``CandidateSelector`` holds them as + bare parameters (``..._codebook``), like the released z-lab checkpoints; + every other DFlash2 parameter already matches. vLLM's ``load_weights`` + refuses the trainer spelling ("Attempted to load nested weight ... into a + single parameter"), so the rename has to happen before the weights reach + the engine. + """ + if _DFLASH2_CODEBOOK_SUFFIX.search(name): + return name[: -len(".weight")] + return name + + def _normalize_dflash2_runtime_aliases(config: Any) -> bool: """Mirror DFlash2 hyperparameters into ``dflash_config`` for vLLM. @@ -2370,11 +2386,9 @@ def _draft_param_name_candidates(name: str) -> list[str]: candidates.append(candidate) if "midlayer." in candidate: candidates.append(candidate.replace("midlayer.", "model.layers.0.")) - # The trainer keeps the DFlash2 selector codebooks in nn.Embedding - # modules (``..._codebook.weight``); vLLM's CandidateSelector holds them as - # bare parameters (``..._codebook``), like the released z-lab checkpoints. - if _DFLASH2_CODEBOOK_SUFFIX.search(candidate): - candidates.append(candidate[: -len(".weight")]) + engine_name = _dflash2_engine_param_name(candidate) + if engine_name != candidate: + candidates.append(engine_name) for candidate in list(candidates): if not candidate.startswith("model."): candidates.append(f"model.{candidate}") @@ -2993,6 +3007,10 @@ def on_bucket_received(bucket_weights): changed = True if "midlayer." in n: n = n.replace("midlayer.", "layers.0.") + if is_dflash: + # DFlash2 selector codebooks: trainer ``.weight`` -> engine bare + # parameter; vLLM's load_weights rejects the trainer spelling. + n = _dflash2_engine_param_name(n) if ( is_eagle3 and n != "lm_head.weight" From 63e75299fc8dc5f5859e2b101683a12f85dd2b81 Mon Sep 17 00:00:00 2001 From: khazic Date: Wed, 2 Sep 2026 10:02:09 +0800 Subject: [PATCH 5/6] refactor(dflash2): apply review cleanups to the vLLM co-train path 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 --- .../backends/dflash2_trainer_backend.py | 2 + verl_speco/convert_speculators_dflash2.py | 11 +++-- verl_speco/integration/vllm_runtime.py | 40 +++++++++++-------- 3 files changed, 33 insertions(+), 20 deletions(-) diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py index 7434ed24..8382568a 100644 --- a/verl_speco/backends/dflash2_trainer_backend.py +++ b/verl_speco/backends/dflash2_trainer_backend.py @@ -161,6 +161,8 @@ class DFlash2TrainerBackend(DFlashTrainerBackend): # ``nn.Parameter`` tensors, while this overlay holds them in ``nn.Embedding`` # modules, whose state dict spells the same tensor with a trailing # ``.weight``. Every other DFlash2 parameter name matches upstream exactly. + # Publish-side inverse: vllm_runtime._dflash2_engine_param_name — keep the + # two in sync when adding aliases here. _CHECKPOINT_KEY_ALIASES = { "candidate_selector.predecessor_codebook": ( "candidate_selector.predecessor_codebook.weight" diff --git a/verl_speco/convert_speculators_dflash2.py b/verl_speco/convert_speculators_dflash2.py index 6d56cec9..ed88bed9 100644 --- a/verl_speco/convert_speculators_dflash2.py +++ b/verl_speco/convert_speculators_dflash2.py @@ -151,10 +151,13 @@ def convert_speculators_dflash2_config( if transformer.get(key) is not None: converted[key] = transformer[key] rope_parameters = transformer.get("rope_parameters") - if isinstance(rope_parameters, dict) and "rope_theta" not in converted: - # transformers < 5 reads rope_theta at the top level. - if rope_parameters.get("rope_theta") is not None: - converted["rope_theta"] = rope_parameters["rope_theta"] + # transformers < 5 reads rope_theta at the top level. + if ( + "rope_theta" not in converted + and isinstance(rope_parameters, dict) + and rope_parameters.get("rope_theta") is not None + ): + converted["rope_theta"] = rope_parameters["rope_theta"] if keep_sliding_window: for key in _SLIDING_WINDOW_KEYS: if key in transformer: diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py index bc622dc8..a3d2cd39 100644 --- a/verl_speco/integration/vllm_runtime.py +++ b/verl_speco/integration/vllm_runtime.py @@ -27,7 +27,6 @@ import json import logging import os -import re import sys import threading import time @@ -822,6 +821,10 @@ def _drafter_algorithm(drafter_cfg: dict[str, Any]) -> str: # Hyperparameters vLLM's DFlash2 draft (qwen3_dflash2.py) indexes straight out of # ``dflash_config`` when it builds the model; a missing key is a KeyError deep in # engine startup, so the checkpoint contract is checked up front instead. +# Superset lists of the nested-key contract live in +# verl_speco/models/dflash2/configuration_dflash2.py (_NESTED_DFLASH_KEYS) and +# verl_speco/convert_speculators_dflash2.py (_DFLASH2_KEYS); they cannot be +# imported here because this module must stay importable without transformers. _DFLASH2_RUNTIME_KEYS = ( "conv_kernel_size", "conv_group_size", @@ -831,10 +834,11 @@ def _drafter_algorithm(drafter_cfg: dict[str, Any]) -> str: # vLLM module that carries the DFlash2 draft class; its presence is the # capability probe for DFlash2 rollout (vllm-project/vllm#52816, v0.28.0+). _VLLM_DFLASH2_MODULE = "vllm.model_executor.models.qwen3_dflash2" -# Trainer-side spelling of the DFlash2 selector codebooks (see -# ``_draft_param_name_candidates`` for the engine-side rename). -_DFLASH2_CODEBOOK_SUFFIX = re.compile( - r"candidate_selector\.(predecessor|successor)_codebook\.weight$" +# Trainer-side spelling of the DFlash2 selector codebooks. Load-side inverse: +# DFlash2TrainerBackend._CHECKPOINT_KEY_ALIASES — keep the two in sync. +_DFLASH2_CODEBOOK_WEIGHT_SUFFIXES = ( + "candidate_selector.predecessor_codebook.weight", + "candidate_selector.successor_codebook.weight", ) @@ -849,7 +853,7 @@ def _dflash2_engine_param_name(name: str) -> str: single parameter"), so the rename has to happen before the weights reach the engine. """ - if _DFLASH2_CODEBOOK_SUFFIX.search(name): + if name.endswith(_DFLASH2_CODEBOOK_WEIGHT_SUFFIXES): return name[: -len(".weight")] return name @@ -913,12 +917,10 @@ def _assert_vllm_supports_dflash2() -> None: def _dflash2_config_value(config: dict[str, Any], key: str) -> Any: """Read a DFlash2 knob from the checkpoint config, top level first.""" - if key in config: - return config[key] - nested = config.get("dflash_config") - if isinstance(nested, dict): - return nested.get(key) - return None + return _first_present( + _get_nested(config, (key,), None), + _get_nested(config, ("dflash_config", key), None), + ) def _validate_vllm_dflash2_block_size( @@ -966,9 +968,12 @@ def _load_vllm_dflash_drafter_config(spec_model_path: Any) -> dict[str, Any] | N def _validate_vllm_dflash_drafter_config( - spec_model_path: Any, algorithm: str = "DFLASH" + spec_model_path: Any, + algorithm: str = "DFLASH", + config: dict[str, Any] | None = None, ) -> None: - config = _load_vllm_dflash_drafter_config(spec_model_path) + if config is None: + config = _load_vllm_dflash_drafter_config(spec_model_path) if config is None: return config_path = os.path.join(os.fspath(spec_model_path), "config.json") @@ -1216,8 +1221,11 @@ def build_vllm_speculative_config_from_drafter( rollout_drafter_cfg = drafter_cfg.get("rollout") or {} if method in ("dflash", "dspark"): + drafter_checkpoint_config = _load_vllm_dflash_drafter_config(spec_model_path) if method == "dflash" or algorithm == "DSPARK": - _validate_vllm_dflash_drafter_config(spec_model_path, algorithm=algorithm) + _validate_vllm_dflash_drafter_config( + spec_model_path, algorithm=algorithm, config=drafter_checkpoint_config + ) num_speculative_tokens = _positive_int_or_none( rollout_drafter_cfg.get("spec_verify_tokens") ) @@ -1229,7 +1237,7 @@ def build_vllm_speculative_config_from_drafter( if algorithm == "DFLASH2": _assert_vllm_supports_dflash2() _validate_vllm_dflash2_block_size( - _load_vllm_dflash_drafter_config(spec_model_path), + drafter_checkpoint_config, drafter_cfg, num_speculative_tokens, ) From 95258386387b6e474417ac1bcaee020d4f2e36f5 Mon Sep 17 00:00:00 2001 From: khazic Date: Wed, 2 Sep 2026 10:03:10 +0800 Subject: [PATCH 6/6] refactor(ci,dflash2): keep the DFlash2 spec_verify_tokens default in 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 --- ci/run_example_test.sh | 9 ++++----- verl_speco/integration/vllm_runtime.py | 7 +++++-- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/ci/run_example_test.sh b/ci/run_example_test.sh index a0c4e2a0..76b9c4f0 100644 --- a/ci/run_example_test.sh +++ b/ci/run_example_test.sh @@ -72,7 +72,6 @@ done # The example owns the collect-to-train algorithm pairing; the runner only needs # the collect side of it to pick the matching cached draft model. separate_training="false" -spec_verify_tokens_default="4" case "${drafter}" in eagle3|megatron-eagle3) @@ -90,9 +89,6 @@ case "${drafter}" in dflash2) draft_model="${SPECO_DFLASH2_DRAFT_MODEL:-}" draft_algorithm="DFLASH2" - # vLLM sizes the DFlash2 convolution block as 1 + spec_verify_tokens, and - # the trainer folds by dflash2_block_size (default 8), so the two must agree. - spec_verify_tokens_default="7" ;; peagle) draft_model="${SPECO_EAGLE3_DRAFT_MODEL:-}" @@ -197,7 +193,7 @@ overrides=( "actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${SPECO_REF_LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-1}" "actor_rollout_ref.rollout.drafter.rollout.spec_steps=${SPECO_SPEC_STEPS:-3}" "actor_rollout_ref.rollout.drafter.rollout.spec_topk=${SPECO_SPEC_TOPK:-1}" - "actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=${SPECO_SPEC_VERIFY_TOKENS:-${spec_verify_tokens_default}}" + "actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=${SPECO_SPEC_VERIFY_TOKENS:-4}" "actor_rollout_ref.rollout.drafter.training.step=${SPECO_DRAFTER_TRAINING_STEPS:-1}" "actor_rollout_ref.rollout.drafter.training.batch_size_per_gpu=${SPECO_DRAFTER_BATCH_SIZE_PER_GPU:-1}" "actor_rollout_ref.rollout.drafter.training.collect_interval_steps=${SPECO_COLLECT_INTERVAL_STEPS:-1}" @@ -241,6 +237,9 @@ fi if [[ "${drafter}" == "dflash2" ]]; then overrides+=( + # vLLM sizes the DFlash2 convolution block as 1 + spec_verify_tokens, and + # the trainer folds by dflash2_block_size (default 8), so the two must agree. + "actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=${SPECO_DFLASH2_SPEC_VERIFY_TOKENS:-7}" "actor_rollout_ref.rollout.drafter.training.dflash2_block_size=${SPECO_DFLASH2_BLOCK_SIZE:-8}" "actor_rollout_ref.rollout.drafter.training.dflash2_num_anchors=${SPECO_DFLASH2_NUM_ANCHORS:-8}" "actor_rollout_ref.rollout.drafter.training.dflash2_loss_decay_gamma=${SPECO_DFLASH2_LOSS_DECAY_GAMMA:-7}" diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py index a3d2cd39..60a9f310 100644 --- a/verl_speco/integration/vllm_runtime.py +++ b/verl_speco/integration/vllm_runtime.py @@ -2469,8 +2469,11 @@ def _ipc_safe_allocator(enabled: bool): 5.6); on older kernels the vLLM worker fails the whole draft update with "does not support the pidfd_getfd syscall". verl's own actor->rollout sync flips expandable segments off around its send for the same reason and turns - them back on afterwards, so the draft publish mirrors that. A verl without - the helper never enabled expandable segments, so there is nothing to do. + them back on afterwards, so the draft publish mirrors that. Restoring + ``True`` unconditionally (torch has no public getter for the prior state) + matches the state verl's own per-step sync leaves behind on every verl that + ships the helper; a verl without it never enabled expandable segments, and + the ImportError guard then leaves the allocator untouched. """ if not enabled: yield