From 95f4274b065ee92aed6100b6f3d1b4c511fb8ea5 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 16:47:52 +0800 Subject: [PATCH 01/16] feat: add DFlash2 drafter training backend DFlash2 (https://inco.ai/blog/dflash2/) keeps DFlash's block-diffusion drafting and target-context backbone and adds two modules on top: - a two-tap dynamic depthwise convolution wrapped around every attention and MLP sublayer, which counteracts the accuracy decay DFlash shows toward the end of a block; - a candidate selector that keeps selector_top_k candidates per block position and traces one coherent path with a low-rank bilinear score over adjacent candidates: S_t(a,b) = U_t(b) + . Implemented as a DFlash variant, mirroring how Domino and DSpark extend the same backbone: DFlash2DraftModel subclasses DFlashDraftModel and DFlash2TrainerBackend subclasses DFlashTrainerBackend, so the block-drafter plumbing, anchor sampling, packing and FSDP2 wrapping are reused rather than duplicated. Defaults match the released z-lab/Qwen3.8-27B-DFlash2 checkpoint (block_size 8, conv 2-tap / group 16, selector rank 256 / top-k 16). Two deltas over the upstream inference implementation are load-bearing here: - The convolution is causal along the sequence axis. Upstream only ever drafts a single block, but training packs n_blocks blocks into one flat draft sequence, so the conv is applied block-locally; otherwise tap 1 at a block's first position would read the last position of the previous, unrelated block. - Upstream allocates base_kernel with torch.empty and relies on from_pretrained to fill it. Training from scratch needs a real init, so the conv starts as an identity passthrough (tap 0 = 1, later taps = 0, zeroed projection): a freshly built DFlash2 is numerically identical to DFlash and learns the correction. The selector's training objective teacher-forces the predecessor to the ground truth, so it scores all positions in one shot instead of the sequential trace inference uses; it is a cross-entropy over the drafter's own top-k restricted to rows whose ground truth survived that top-k. It requires loss_mode=full_vocab and fails loud otherwise, since a restricted vocabulary would make the candidate ids meaningless. To avoid copying DFlashTrainingModel.forward wholesale (as Domino does), DFlashTrainingModel gains a documented _auxiliary_loss extension point that returns None for plain DFlash, leaving every existing backend's behaviour and numerics unchanged. Signed-off-by: khazic --- .../test_dflash2_backend_contract.py | 346 ++++++++++++++++++ .../backends/dflash2_trainer_backend.py | 304 +++++++++++++++ verl_speco/backends/dflash_trainer_backend.py | 50 +++ verl_speco/backends/factory.py | 5 + verl_speco/config/speco_base.yaml | 17 + .../integration/oldlogprob_layer_ids.py | 6 +- verl_speco/models/auto.py | 17 + verl_speco/models/dflash/modeling_dflash.py | 15 + verl_speco/models/dflash2/__init__.py | 28 ++ .../models/dflash2/configuration_dflash2.py | 94 +++++ verl_speco/models/dflash2/modeling_dflash2.py | 262 +++++++++++++ verl_speco/trainer/base_trainer.py | 5 +- verl_speco/trainer/draft_training_loop.py | 21 +- 13 files changed, 1165 insertions(+), 5 deletions(-) create mode 100644 tests/integration/test_dflash2_backend_contract.py create mode 100644 verl_speco/backends/dflash2_trainer_backend.py create mode 100644 verl_speco/models/dflash2/__init__.py create mode 100644 verl_speco/models/dflash2/configuration_dflash2.py create mode 100644 verl_speco/models/dflash2/modeling_dflash2.py diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py new file mode 100644 index 00000000..345e7a55 --- /dev/null +++ b/tests/integration/test_dflash2_backend_contract.py @@ -0,0 +1,346 @@ +"""Contract tests for the DFlash2 drafter backend. + +CPU-light: they exercise the two DFlash2 modules (the two-tap dynamic +convolution and the candidate selector), the block-locality invariant the +training layout requires, the algorithm routing, and the block-drafter +classification. The full training forward is validated on GPU by +``ci/dflash2_gpu_smoke.py``. +""" + +from __future__ import annotations + +import pytest + + +def _tiny_dflash2_config(**overrides): + from verl_speco.models.dflash2 import DFlash2Config + + kwargs = dict( + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=2, + num_hidden_layers=1, + vocab_size=32, + num_target_layers=4, + num_context_layers=2, + target_hidden_size=8, + target_num_hidden_layers=4, + target_layer_ids=[1, 3], + mask_token_id=31, + block_size=4, + num_anchors=8, + conv_kernel_size=2, + conv_group_size=4, + selector_rank=6, + selector_top_k=5, + rms_norm_eps=1e-6, + max_position_embeddings=64, + ) + kwargs.update(overrides) + return DFlash2Config(**kwargs) + + +def test_dflash2_model_builds_conv_and_selector() -> None: + pytest.importorskip("torch") + pytest.importorskip("transformers") + from verl_speco.models.dflash2 import DFlash2DraftModel + + config = _tiny_dflash2_config() + model = DFlash2DraftModel(config) + + # Every layer gets a conv around attention and around the MLP. + for layer in model.layers: + assert layer.attention_conv is not None + assert layer.mlp_conv is not None + assert layer.attention_conv.kernel_size == config.conv_kernel_size + assert layer.attention_conv.group_size == config.conv_group_size + groups = config.hidden_size // config.conv_group_size + assert layer.attention_conv.kernel_projection.out_features == ( + 2 * config.conv_kernel_size * groups + ) + + selector = model.candidate_selector + assert selector.predecessor_codebook.num_embeddings == config.vocab_size + assert selector.successor_codebook.num_embeddings == config.vocab_size + assert selector.hidden_projection.in_features == config.hidden_size + assert selector.hidden_projection.out_features == config.selector_rank + + +def test_plain_dflash_layers_keep_conv_hooks_inert() -> None: + """The conv hooks live on the shared DFlash layer; DFlash must not use them.""" + pytest.importorskip("torch") + pytest.importorskip("transformers") + from verl_speco.models.dflash import DFlashConfig, DFlashDraftModel + + config = DFlashConfig( + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=2, + num_hidden_layers=1, + vocab_size=32, + num_target_layers=4, + num_context_layers=2, + target_hidden_size=8, + target_num_hidden_layers=4, + target_layer_ids=[1, 3], + mask_token_id=31, + ) + model = DFlashDraftModel(config) + for layer in model.layers: + assert layer.attention_conv is None + assert layer.mlp_conv is None + + +def test_conv_is_identity_at_init() -> None: + """A freshly built DFlash2 conv must be a passthrough. + + The correction is learned on top of DFlash, so an untrained DFlash2 has to + start numerically equal to DFlash rather than perturbing the backbone. + """ + pytest.importorskip("torch") + import torch + + from verl_speco.models.dflash2 import GroupedDynamicCausalConv + + conv = GroupedDynamicCausalConv( + hidden_size=8, kernel_size=2, group_size=4, block_size=4 + ) + hidden = torch.randn(2, 8, 8) + prepared, dynamic = conv.prepare(hidden) + torch.testing.assert_close(prepared, hidden) + torch.testing.assert_close(conv.finish(hidden, dynamic), hidden) + + +def test_conv_does_not_leak_across_block_boundaries() -> None: + """Tap 1 is causal, so it must never read the previous block's last row. + + Training packs ``n_blocks`` blocks into one flat draft sequence. If the + convolution ran over that flat axis, position 0 of block i would mix in the + final position of block i-1, which belongs to an unrelated anchor. + """ + pytest.importorskip("torch") + import torch + + from verl_speco.models.dflash2 import GroupedDynamicCausalConv + + block_size = 4 + conv = GroupedDynamicCausalConv( + hidden_size=8, kernel_size=2, group_size=4, block_size=block_size + ) + # Make tap 1 the only contributor so any cross-block read is visible. + with torch.no_grad(): + conv.base_kernel[0, 0, :] = 0.0 + conv.base_kernel[0, 1, :] = 1.0 + + hidden = torch.randn(1, 2 * block_size, 8) + out, _ = conv.prepare(hidden) + + # First row of each block has no in-block predecessor, so it must be zero. + torch.testing.assert_close(out[:, 0], torch.zeros_like(out[:, 0])) + torch.testing.assert_close( + out[:, block_size], torch.zeros_like(out[:, block_size]) + ) + # Interior rows read their own block's previous row. + torch.testing.assert_close(out[:, 1], hidden[:, 0]) + torch.testing.assert_close(out[:, block_size + 1], hidden[:, block_size]) + + +def test_conv_rejects_length_that_is_not_a_block_multiple() -> None: + pytest.importorskip("torch") + import torch + + from verl_speco.models.dflash2 import GroupedDynamicCausalConv + + conv = GroupedDynamicCausalConv( + hidden_size=8, kernel_size=2, group_size=4, block_size=4 + ) + with pytest.raises(ValueError, match="multiple of block_size"): + conv.prepare(torch.randn(1, 6, 8)) + + +def test_selector_pair_scores_match_the_sequential_selector() -> None: + """The vectorized training path must agree with the inference-time trace. + + ``pair_scores`` teacher-forces the predecessor; feeding it the path the + sequential ``select`` actually took has to reproduce the same scores. + """ + pytest.importorskip("torch") + import torch + + from verl_speco.models.dflash2 import CandidateSelector + + torch.manual_seed(0) + config = _tiny_dflash2_config() + selector = CandidateSelector(config) + + batch, block = 2, config.block_size + hidden = torch.randn(batch, block, config.hidden_size) + logits = torch.randn(batch, block, config.vocab_size) + anchor_ids = torch.randint(0, config.vocab_size, (batch,)) + + path, candidates = selector.select(hidden, logits, anchor_ids) + assert path.shape == (batch, block) + + # Rebuild the predecessor sequence the trace used: anchor, then its choices. + predecessors = torch.cat([anchor_ids.unsqueeze(1), path[:, :-1]], dim=1) + top_k = candidates.shape[-1] + unary = torch.gather(logits, 2, candidates) + with torch.no_grad(): + flat_scores = selector.pair_scores( + hidden.reshape(-1, config.hidden_size), + unary.reshape(-1, top_k), + candidates.reshape(-1, top_k), + predecessors.reshape(-1), + ).reshape(batch, block, top_k) + replayed = torch.gather( + candidates, 2, flat_scores.argmax(dim=-1, keepdim=True) + ).squeeze(-1) + torch.testing.assert_close(replayed, path) + + +def test_dflash2_training_model_adds_selector_loss() -> None: + """The selector objective must actually contribute to the total loss.""" + pytest.importorskip("torch") + pytest.importorskip("transformers") + import torch + + from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainingModel + from verl_speco.models.dflash2 import DFlash2DraftModel + + config = _tiny_dflash2_config() + bsz, seq_len = 2, 16 + input_ids = torch.randint(0, config.vocab_size, (bsz, seq_len)) + hidden_states_list = [ + torch.randn(bsz, seq_len, config.target_hidden_size) + for _ in config.target_layer_ids + ] + loss_mask = torch.ones(bsz, seq_len, dtype=torch.long) + lm_head_weight = torch.randn(config.vocab_size, config.hidden_size) + + def run(selector_loss_weight): + torch.manual_seed(0) + model = DFlash2TrainingModel( + draft_model=DFlash2DraftModel(config), + block_size=config.block_size, + num_anchors=config.num_anchors, + selector_loss_weight=selector_loss_weight, + ) + return model(input_ids, hidden_states_list, loss_mask, lm_head_weight) + + loss_off, _, _, _, _, diagnostics_off = run(0.0) + loss_on, _, _, _, _, diagnostics_on = run(1.0) + + assert "selector_loss" not in diagnostics_off + assert "selector_loss" in diagnostics_on + assert float(diagnostics_on["selector_active_count"]) > 0 + # With the weight on, the total loss carries the extra selector term. + assert float(loss_on) > float(loss_off) + assert float(loss_on) == pytest.approx( + float(loss_off) + float(diagnostics_on["selector_loss"]), rel=1e-4 + ) + + +def test_dflash2_training_model_rejects_restricted_vocab() -> None: + """Selector candidates are real token ids, so a restricted vocab is invalid.""" + pytest.importorskip("torch") + pytest.importorskip("transformers") + + from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainingModel + from verl_speco.models.dflash2 import DFlash2DraftModel + + config = _tiny_dflash2_config() + with pytest.raises(ValueError, match="full_vocab"): + DFlash2TrainingModel( + draft_model=DFlash2DraftModel(config), + block_size=config.block_size, + num_anchors=config.num_anchors, + loss_mode="restricted_ce", + ) + + +def test_dflash2_backend_is_registered_in_the_factory() -> None: + from verl_speco.backends.factory import SUPPORTED_DRAFTER_ALGORITHMS + + assert "DFLASH2" in SUPPORTED_DRAFTER_ALGORITHMS + + +def test_dflash2_uses_dflash_aux_layers() -> None: + from verl_speco.integration.oldlogprob_layer_ids import DFLASH_FAMILY_ALGORITHMS + + assert "DFLASH2" in DFLASH_FAMILY_ALGORITHMS + + +def test_dflash2_config_lifts_nested_dflash_config(tmp_path) -> None: + """Upstream z-lab checkpoints nest the DFlash2 knobs under dflash_config.""" + pytest.importorskip("transformers") + import json + + from verl_speco.models.dflash2 import DFlash2Config + + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "architectures": ["DFlash2DraftModel"], + "hidden_size": 8, + "intermediate_size": 16, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "num_hidden_layers": 1, + "vocab_size": 32, + "dflash_config": { + "block_size": 8, + "conv_kernel_size": 2, + "conv_group_size": 16, + "selector_rank": 256, + "selector_top_k": 16, + "mask_token_id": 31, + "target_layer_ids": [1, 3], + }, + } + ) + ) + + config = DFlash2Config.from_dflash2_pretrained(str(tmp_path)) + assert config.model_type == "dflash2" + assert config.block_size == 8 + assert config.conv_kernel_size == 2 + assert config.conv_group_size == 16 + assert config.selector_rank == 256 + assert config.selector_top_k == 16 + assert config.mask_token_id == 31 + assert config.target_layer_ids == [1, 3] + + +def test_dflash2_config_routes_through_auto(tmp_path) -> None: + pytest.importorskip("transformers") + import json + + from verl_speco.models.auto import AutoDraftModelConfig + from verl_speco.models.dflash2 import DFlash2Config + + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "architectures": ["DFlash2DraftModel"], + "hidden_size": 8, + "intermediate_size": 16, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "num_hidden_layers": 1, + "vocab_size": 32, + "dflash_config": {"block_size": 8, "selector_top_k": 16}, + } + ), + encoding="utf-8", + ) + + loaded = AutoDraftModelConfig.from_file(str(config_path)) + assert isinstance(loaded, DFlash2Config) + assert loaded.architectures == ["DFlash2DraftModel"] + # The nested z-lab block must survive routing through AutoDraftModelConfig. + assert loaded.block_size == 8 + assert loaded.selector_top_k == 16 diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py new file mode 100644 index 00000000..dcb3cc0f --- /dev/null +++ b/verl_speco/backends/dflash2_trainer_backend.py @@ -0,0 +1,304 @@ +# 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 logging +import os +from copy import deepcopy + +import torch +import torch.nn.functional as F + +from verl_speco.backends.dflash_trainer_backend import ( + DFlashTrainerBackend, + DFlashTrainingModel, +) +from verl_speco.models.dflash2 import DFlash2Config, DFlash2DraftModel +from verl_speco.trainer.checkpoint import log_drafter_checkpoint_step + +logger = logging.getLogger(__name__) + + +class DFlash2TrainingModel(DFlashTrainingModel): + """DFlash training wrapper plus the DFlash2 candidate-selector objective. + + The dynamic convolutions live inside the draft model, so they are already + exercised by the inherited CE path. Only the selector needs its own loss: it + learns to re-rank the drafter's own top-k candidates using the previous + token, which at training time is teacher-forced to the ground truth. + """ + + _no_split_modules = ["DFlashDecoderLayer"] + + def __init__(self, *args, selector_loss_weight: float = 1.0, **kwargs): + super().__init__(*args, **kwargs) + self.selector_loss_weight = float(selector_loss_weight) + if self.loss_mode != "full_vocab": + raise ValueError( + "DFlash2's candidate selector scores real vocabulary ids, so it " + f"requires loss_mode='full_vocab'; got {self.loss_mode!r}. A " + "restricted/sampled vocabulary would make the selector's " + "candidate ids meaningless." + ) + + def _auxiliary_loss( + self, + *, + input_ids, + safe_label_indices, + active_mask, + active_hidden, + active_logits, + active_targets, + active_weights, + ): + selector = self.draft_model.candidate_selector + if self.selector_loss_weight <= 0: + return None, {} + + device = active_hidden.device + top_k = min(selector.top_k, active_logits.shape[-1]) + unary, candidates = torch.topk(active_logits, top_k, dim=-1, sorted=False) + + # Teacher-forced predecessor: the token immediately before the slot this + # row predicts. Block-relative position 0 is the anchor and never carries + # loss weight, so index - 1 stays inside the sequence for every active row. + predecessor_indices = (safe_label_indices - 1).clamp(min=0) + bsz, n_blocks, block_size = safe_label_indices.shape + flat_predecessor_indices = predecessor_indices.reshape(bsz, -1) + predecessor_ids = torch.gather(input_ids, 1, flat_predecessor_indices) + predecessor_ids = predecessor_ids.reshape(-1)[active_mask] + + scores = selector.pair_scores( + active_hidden, unary.float(), candidates, predecessor_ids + ) + + # The selector can only learn on rows whose ground truth survived the + # drafter's own top-k; elsewhere there is no correct choice to make. + target_hits = candidates == active_targets.unsqueeze(-1) + learnable = target_hits.any(dim=-1) + selector_loss = torch.zeros((), dtype=torch.float32, device=device) + selector_correct = torch.zeros((), dtype=torch.float32, device=device) + selector_tokens = torch.zeros((), dtype=torch.float32, device=device) + if bool(learnable.any()): + learn_scores = scores[learnable] + learn_weights = active_weights[learnable].float() + learn_targets = torch.argmax(target_hits[learnable].int(), dim=-1) + per_row = F.cross_entropy(learn_scores, learn_targets, reduction="none") + finite = torch.isfinite(per_row) + per_row = torch.where(finite, per_row, torch.zeros_like(per_row)) + learn_weights = learn_weights * finite.to(learn_weights.dtype) + selector_loss = (per_row * learn_weights).sum() / learn_weights.sum().clamp( + min=1e-6 + ) + with torch.no_grad(): + selector_tokens = finite.float().sum() + selector_correct = ( + ((torch.argmax(learn_scores, dim=-1) == learn_targets) & finite) + .float() + .sum() + ) + + with torch.no_grad(): + metrics = { + "selector_loss": selector_loss.detach().float(), + "selector_correct_count": selector_correct, + "selector_token_count": selector_tokens, + "selector_coverage_count": learnable.float().sum(), + "selector_active_count": torch.tensor( + float(active_targets.numel()), dtype=torch.float32, device=device + ), + } + return self.selector_loss_weight * selector_loss, metrics + + +class DFlash2TrainerBackend(DFlashTrainerBackend): + @property + def model_type(self): + return "dflash2" + + def _training_value(self, training_cfg, dflash2_key: str, dflash_key: str, default): + value = training_cfg.get(dflash2_key, None) + if value is not None: + return value + return training_cfg.get(dflash_key, default) + + def _normalize_dflash_config( + self, drafter_config, target_hf_config, normalized_state, spec_model_path + ): + training_cfg = self.config.rollout.drafter.training + if training_cfg.get("dflash2_num_target_layers", None) is not None: + if getattr(drafter_config, "num_context_layers", None) is None: + drafter_config.num_context_layers = int( + training_cfg["dflash2_num_target_layers"] + ) + return super()._normalize_dflash_config( + drafter_config, target_hf_config, normalized_state, spec_model_path + ) + + def _build_fallback_config(self, target_hf_config): + training_cfg = self.config.rollout.drafter.training + target_text_config = getattr(target_hf_config, "text_config", target_hf_config) + hidden_size_cfg = self._training_value( + training_cfg, "dflash2_hidden_size", "dflash_hidden_size", None + ) + hidden_size = int( + hidden_size_cfg + if hidden_size_cfg is not None + else target_text_config.hidden_size + ) + num_context_layers = int( + self._training_value( + training_cfg, "dflash2_num_target_layers", "dflash_num_target_layers", 5 + ) + ) + target_num_hidden_layers = int( + getattr(target_text_config, "num_hidden_layers", 36) + ) + mask_token_id_cfg = self._training_value( + training_cfg, "dflash2_mask_token_id", "dflash_mask_token_id", None + ) + mask_token_id = int( + mask_token_id_cfg + if mask_token_id_cfg is not None + else target_text_config.vocab_size - 1 + ) + target_layer_ids = self._training_value( + training_cfg, "dflash2_target_layer_ids", "dflash_target_layer_ids", None + ) + if target_layer_ids is None: + from verl_speco.models.dflash import build_target_layer_ids + + target_layer_ids = build_target_layer_ids( + num_context_layers, target_num_hidden_layers + ) + return DFlash2Config( + hidden_size=hidden_size, + intermediate_size=int( + getattr(target_text_config, "intermediate_size", hidden_size * 4) + ), + num_hidden_layers=int( + self._training_value( + training_cfg, + "dflash2_num_hidden_layers", + "dflash_num_hidden_layers", + 5, + ) + ), + num_attention_heads=int(getattr(target_text_config, "num_attention_heads")), + num_key_value_heads=int( + getattr( + target_text_config, + "num_key_value_heads", + getattr(target_text_config, "num_attention_heads"), + ) + ), + vocab_size=int(target_text_config.vocab_size), + rms_norm_eps=float(getattr(target_text_config, "rms_norm_eps", 1e-6)), + max_position_embeddings=int( + getattr(target_text_config, "max_position_embeddings", 32768) + ), + rope_theta=float(getattr(target_text_config, "rope_theta", 10000.0)), + num_target_layers=target_num_hidden_layers, + num_context_layers=num_context_layers, + target_hidden_size=int(target_text_config.hidden_size), + target_num_hidden_layers=target_num_hidden_layers, + target_layer_ids=target_layer_ids, + mask_token_id=mask_token_id, + block_size=int(training_cfg.get("dflash2_block_size", 8)), + num_anchors=int(training_cfg.get("dflash2_num_anchors", 512)), + loss_decay_gamma=float(training_cfg.get("dflash2_loss_decay_gamma", 7.0)), + conv_kernel_size=int(training_cfg.get("dflash2_conv_kernel_size", 2)), + conv_group_size=int(training_cfg.get("dflash2_conv_group_size", 16)), + selector_rank=int(training_cfg.get("dflash2_selector_rank", 256)), + selector_top_k=int(training_cfg.get("dflash2_selector_top_k", 16)), + selector_loss_weight=float( + training_cfg.get("dflash2_selector_loss_weight", 1.0) + ), + architectures=["DFlash2DraftModel"], + ) + + def build_model(self): + target_model_path = self.config.model.path + spec_model_path = self.config.rollout.drafter.model_path + config_path = ( + os.path.join(spec_model_path, "config.json") if spec_model_path else None + ) + target_hf_config = self._get_target_hf_config() + normalized_state = None + + if config_path and os.path.exists(config_path): + drafter_config = DFlash2Config.from_dflash2_pretrained(spec_model_path) + if spec_model_path and os.path.exists(spec_model_path): + log_drafter_checkpoint_step( + logger, spec_model_path, action="Loading DFlash2 drafter weights" + ) + normalized_state = self._normalize_draft_state_dict( + self._load_draft_state_dict(spec_model_path) + ) + else: + drafter_config = self._build_fallback_config(target_hf_config) + + if not isinstance(drafter_config, DFlash2Config): + raise TypeError( + f"DFlash2 config is not a DFlash2Config: {type(drafter_config)}" + ) + drafter_config = self._normalize_dflash_config( + drafter_config, target_hf_config, normalized_state, spec_model_path + ) + + draft_model = DFlash2DraftModel(deepcopy(drafter_config)) + if ( + spec_model_path + and os.path.exists(spec_model_path) + and os.path.exists(config_path) + ): + self._load_draft_checkpoint( + draft_model, spec_model_path, normalized_state=normalized_state + ) + draft_model.load_embedding(target_model_path) + draft_model.freeze_embedding() + + self.target_lm_head = self._build_target_lm_head( + target_model_path, target_hf_config + ) + training_cfg = self.config.rollout.drafter.training + return DFlash2TrainingModel( + draft_model=draft_model, + block_size=int( + training_cfg.get( + "dflash2_block_size", getattr(drafter_config, "block_size", 8) + ) + ), + num_anchors=int( + training_cfg.get( + "dflash2_num_anchors", getattr(drafter_config, "num_anchors", 512) + ) + ), + loss_decay_gamma=float( + training_cfg.get( + "dflash2_loss_decay_gamma", + getattr(drafter_config, "loss_decay_gamma", 7.0), + ) + ), + selector_loss_weight=float( + training_cfg.get( + "dflash2_selector_loss_weight", + getattr(drafter_config, "selector_loss_weight", 1.0), + ) + ), + ), drafter_config + + +__all__ = ["DFlash2TrainerBackend", "DFlash2TrainingModel"] diff --git a/verl_speco/backends/dflash_trainer_backend.py b/verl_speco/backends/dflash_trainer_backend.py index 838f2552..03ed4196 100644 --- a/verl_speco/backends/dflash_trainer_backend.py +++ b/verl_speco/backends/dflash_trainer_backend.py @@ -153,6 +153,41 @@ def __init__( self.sampled_ce_negatives = max(int(sampled_ce_negatives), 0) self._tensor_template_cache: dict[tuple, torch.Tensor] = {} + def _auxiliary_loss( + self, + *, + input_ids: torch.Tensor, + safe_label_indices: torch.Tensor, + active_mask: torch.Tensor, + active_hidden: torch.Tensor, + active_logits: torch.Tensor, + active_targets: torch.Tensor, + active_weights: torch.Tensor, + ) -> tuple[torch.Tensor | None, dict[str, torch.Tensor]]: + """Extra loss term contributed by a DFlash variant's own head. + + Called once per step with the intermediates of the main CE path, at the + point where they all exist. Plain DFlash has no auxiliary head and + returns ``None``, leaving the loss untouched. + + Args: + input_ids: ``[bsz, seq_len]`` full sequence, for variants that need + neighbouring tokens (for example a predecessor token). + safe_label_indices: ``[bsz, n_blocks, block_size]`` clamped absolute + positions each block slot predicts. + active_mask: ``[bsz * n_blocks * block_size]`` bool mask of the rows + that carry loss weight. + active_hidden: ``[num_active, hidden]`` backbone states of those rows. + active_logits: ``[num_active, vocab]`` drafter logits of those rows. + active_targets: ``[num_active]`` ground-truth token ids. + active_weights: ``[num_active]`` per-row loss weights. + + Returns: + tuple: ``(loss_or_None, metrics)`` to add to the total loss and the + diagnostics dict. + """ + return None, {} + def _cached_arange( self, name: str, @@ -451,6 +486,20 @@ def forward( local_ploss_sum = (active_loss * active_loss_weights).sum() loss = local_ploss_sum / valid_token_count + auxiliary_metrics: dict[str, torch.Tensor] = {} + if active_targets.numel() > 0: + auxiliary_loss, auxiliary_metrics = self._auxiliary_loss( + input_ids=input_ids, + safe_label_indices=safe_label_indices, + active_mask=active_mask, + active_hidden=active_hidden, + active_logits=active_logits, + active_targets=active_targets, + active_weights=active_weights, + ) + if auxiliary_loss is not None: + loss = loss + auxiliary_loss + with torch.no_grad(): correct = torch.zeros_like(binary_eval_mask, dtype=torch.bool) top1_correct_count = torch.zeros((), dtype=torch.float32, device=device) @@ -605,6 +654,7 @@ def forward( device=device, ), } + diagnostics.update(auxiliary_metrics) return ( loss, diff --git a/verl_speco/backends/factory.py b/verl_speco/backends/factory.py index 7cc72aa6..5042689f 100644 --- a/verl_speco/backends/factory.py +++ b/verl_speco/backends/factory.py @@ -31,6 +31,7 @@ "EAGLE2", "EAGLE3", "DFLASH", + "DFLASH2", "DSPARK", "DOMINO", "PEAGLE", @@ -61,6 +62,10 @@ def build_trainer_backend(config, model_config) -> Any: from verl_speco.backends.dflash_trainer_backend import DFlashTrainerBackend return DFlashTrainerBackend(config, model_config) + if algorithm == "DFLASH2": + from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainerBackend + + return DFlash2TrainerBackend(config, model_config) if algorithm == "DSPARK": from verl_speco.backends.dspark_trainer_backend import DSparkTrainerBackend diff --git a/verl_speco/config/speco_base.yaml b/verl_speco/config/speco_base.yaml index f4f17481..bf9f88e8 100644 --- a/verl_speco/config/speco_base.yaml +++ b/verl_speco/config/speco_base.yaml @@ -163,6 +163,23 @@ actor_rollout_ref: domino_shift_label: true domino_lambda_base_start: 1.0 domino_lambda_base_decay_steps: 2000 + # DFlash2 (DFlash variant: two-tap dynamic depthwise convolutions around + # each sublayer + a candidate selector that re-ranks the drafter's own + # top-k with a low-rank bilinear score over adjacent candidates). + # Defaults match the released z-lab/Qwen3.8-27B-DFlash2 checkpoint. + dflash2_block_size: 8 + dflash2_num_anchors: 512 + dflash2_loss_decay_gamma: 7.0 + dflash2_hidden_size: null + dflash2_num_target_layers: 5 + dflash2_num_hidden_layers: 5 + dflash2_mask_token_id: null + dflash2_target_layer_ids: null + dflash2_conv_kernel_size: 2 + dflash2_conv_group_size: 16 + dflash2_selector_rank: 256 + dflash2_selector_top_k: 16 + dflash2_selector_loss_weight: 1.0 # EAGLE-1 / EAGLE-2 draft training (single-step feature regression + # full-vocab soft-CE distillation against the frozen target head). eagle1_num_hidden_layers: 1 diff --git a/verl_speco/integration/oldlogprob_layer_ids.py b/verl_speco/integration/oldlogprob_layer_ids.py index 73351e2a..3100023f 100644 --- a/verl_speco/integration/oldlogprob_layer_ids.py +++ b/verl_speco/integration/oldlogprob_layer_ids.py @@ -18,8 +18,9 @@ from typing import Any # Drafters that consume the DFlash aux context layers instead of the EAGLE -# aux-plus-final layout. Domino is a DFlash variant, so it shares the layout. -DFLASH_FAMILY_ALGORITHMS = frozenset({"DFLASH", "DSPARK", "DOMINO"}) +# aux-plus-final layout. Domino and DFlash2 are DFlash variants, so they share +# the layout. +DFLASH_FAMILY_ALGORITHMS = frozenset({"DFLASH", "DFLASH2", "DSPARK", "DOMINO"}) def _get_nested(config: Any, path: tuple[str, ...], default=None): @@ -188,6 +189,7 @@ def _dflash_num_context_layers( _get_nested(training_cfg, ("dspark_num_target_layers",), None) ) candidates.append(_get_nested(training_cfg, ("domino_num_target_layers",), None)) + candidates.append(_get_nested(training_cfg, ("dflash2_num_target_layers",), None)) candidates.extend( ( _get_nested(training_cfg, ("dflash_num_target_layers",), None), diff --git a/verl_speco/models/auto.py b/verl_speco/models/auto.py index 985ecb36..6a7acaf1 100644 --- a/verl_speco/models/auto.py +++ b/verl_speco/models/auto.py @@ -41,6 +41,11 @@ "Qwen3DominoModel", } +_DFLASH2_ARCHITECTURE_ALIASES = { + "DFlash2DraftModel", + "Qwen3DFlash2Model", +} + def _normalize_int_list(value): if value is None: @@ -201,6 +206,7 @@ def from_file(cls, config_path: str): architecture not in cls._config_mapping and architecture not in _DSPARK_ARCHITECTURE_ALIASES and architecture not in _DOMINO_ARCHITECTURE_ALIASES + and architecture not in _DFLASH2_ARCHITECTURE_ALIASES ): raise ValueError(f"Architecture {architecture} not supported") @@ -219,6 +225,17 @@ def from_file(cls, config_path: str): config["model_type"] = DominoConfig.model_type config["architectures"] = ["DominoDraftModel"] config.setdefault("projector_type", "domino") + elif architecture in _DFLASH2_ARCHITECTURE_ALIASES: + from .dflash2 import DFlash2Config + + config_class = DFlash2Config + # Upstream z-lab checkpoints nest the DFlash2 knobs under + # ``dflash_config``; lift them so they survive ``from_dict``. + nested = config.get("dflash_config") or {} + for key, value in nested.items(): + config.setdefault(key, value) + config["model_type"] = DFlash2Config.model_type + config["architectures"] = ["DFlash2DraftModel"] elif architecture in _EAGLE3_ARCHITECTURE_ALIASES: config = _normalize_eagle3_config_dict(config) diff --git a/verl_speco/models/dflash/modeling_dflash.py b/verl_speco/models/dflash/modeling_dflash.py index 62aa8ff4..d2713f52 100644 --- a/verl_speco/models/dflash/modeling_dflash.py +++ b/verl_speco/models/dflash/modeling_dflash.py @@ -247,6 +247,11 @@ def __init__(self, config: PretrainedConfig): self.post_attention_layernorm = DFlashRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) + # Optional dynamic-convolution hooks, populated only by DFlash2. Every + # other DFlash-family drafter leaves them None, so their forward is + # unchanged. + self.attention_conv = None + self.mlp_conv = None def forward( self, @@ -259,6 +264,9 @@ def forward( ) -> torch.Tensor: residual = draft_hidden draft_hidden = self.input_layernorm(draft_hidden) + attention_kernel = None + if self.attention_conv is not None: + draft_hidden, attention_kernel = self.attention_conv.prepare(draft_hidden) draft_hidden = self.self_attn( draft_hidden=draft_hidden, context_hidden=context_hidden, @@ -267,11 +275,18 @@ def forward( block_mask=block_mask, dense_attention_mask=dense_attention_mask, ) + if attention_kernel is not None: + draft_hidden = self.attention_conv.finish(draft_hidden, attention_kernel) draft_hidden = residual + draft_hidden residual = draft_hidden draft_hidden = self.post_attention_layernorm(draft_hidden) + mlp_kernel = None + if self.mlp_conv is not None: + draft_hidden, mlp_kernel = self.mlp_conv.prepare(draft_hidden) draft_hidden = self.mlp(draft_hidden) + if mlp_kernel is not None: + draft_hidden = self.mlp_conv.finish(draft_hidden, mlp_kernel) return residual + draft_hidden diff --git a/verl_speco/models/dflash2/__init__.py b/verl_speco/models/dflash2/__init__.py new file mode 100644 index 00000000..2993d363 --- /dev/null +++ b/verl_speco/models/dflash2/__init__.py @@ -0,0 +1,28 @@ +# 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 verl_speco.models.dflash2.configuration_dflash2 import DFlash2Config +from verl_speco.models.dflash2.modeling_dflash2 import ( + CandidateSelector, + DFlash2DraftModel, + GroupedDynamicCausalConv, + grouped_dynamic_convolve, +) + +__all__ = [ + "CandidateSelector", + "DFlash2Config", + "DFlash2DraftModel", + "GroupedDynamicCausalConv", + "grouped_dynamic_convolve", +] diff --git a/verl_speco/models/dflash2/configuration_dflash2.py b/verl_speco/models/dflash2/configuration_dflash2.py new file mode 100644 index 00000000..6dd95b3f --- /dev/null +++ b/verl_speco/models/dflash2/configuration_dflash2.py @@ -0,0 +1,94 @@ +# 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 + +from verl_speco.models.dflash import DFlashConfig + +# Keys the upstream z-lab DFlash2 checkpoints nest under ``dflash_config`` +# instead of putting at the config top level. ``from_dflash2_pretrained`` +# lifts them so the rest of this overlay can read them as plain attributes. +_NESTED_DFLASH_KEYS = ( + "block_size", + "conv_kernel_size", + "conv_group_size", + "selector_rank", + "selector_top_k", + "mask_token_id", + "target_layer_ids", +) + + +class DFlash2Config(DFlashConfig): + """Configuration for the DFlash2 draft model. + + DFlash2 keeps the DFlash target-context backbone and block-diffusion drafting + and adds two modules on top (see https://inco.ai/blog/dflash2/): + + - a two-tap dynamic depthwise convolution wrapped around every attention and + MLP sublayer, which counteracts the accuracy decay DFlash shows toward the + end of a block; + - a candidate selector that keeps ``selector_top_k`` candidates per block + position and traces one coherent path through them with a low-rank + bilinear score over adjacent candidates. + + Defaults match the released ``z-lab/Qwen3.8-27B-DFlash2`` checkpoint. + """ + + model_type = "dflash2" + + def __init__( + self, + *args, + block_size: int = 8, + num_anchors: int = 512, + loss_decay_gamma: float = 7.0, + conv_kernel_size: int = 2, + conv_group_size: int = 16, + selector_rank: int = 256, + selector_top_k: int = 16, + selector_loss_weight: float = 1.0, + **kwargs, + ): + architectures = kwargs.pop("architectures", None) + super().__init__(*args, **kwargs) + self.architectures = architectures or ["DFlash2DraftModel"] + self.block_size = int(block_size) + self.num_anchors = int(num_anchors) + self.loss_decay_gamma = float(loss_decay_gamma) + self.conv_kernel_size = int(conv_kernel_size) + self.conv_group_size = int(conv_group_size) + self.selector_rank = int(selector_rank) + self.selector_top_k = int(selector_top_k) + self.selector_loss_weight = float(selector_loss_weight) + + @classmethod + def from_dflash2_pretrained(cls, model_path: str): + config_path = os.path.join(model_path, "config.json") + with open(config_path, "r", encoding="utf-8") as f: + config = json.load(f) + + # Upstream checkpoints keep the DFlash2-specific knobs in a nested + # ``dflash_config`` block; a top-level value (from a checkpoint this + # overlay wrote itself) always wins. + nested = config.get("dflash_config") or {} + for key in _NESTED_DFLASH_KEYS: + if key not in config and key in nested: + config[key] = nested[key] + + config["model_type"] = cls.model_type + config["architectures"] = ["DFlash2DraftModel"] + return cls.from_dict(config) diff --git a/verl_speco/models/dflash2/modeling_dflash2.py b/verl_speco/models/dflash2/modeling_dflash2.py new file mode 100644 index 00000000..9889c6fa --- /dev/null +++ b/verl_speco/models/dflash2/modeling_dflash2.py @@ -0,0 +1,262 @@ +# 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 torch +import torch.nn as nn +import torch.nn.functional as F + +from verl_speco.models.dflash import DFlashDraftModel + +from .configuration_dflash2 import DFlash2Config + + +def grouped_dynamic_convolve( + hidden: torch.Tensor, + dynamic: torch.Tensor, + base: torch.Tensor, + group_size: int, +) -> torch.Tensor: + """Depthwise causal convolution with a content-adaptive per-group correction. + + ``base`` is a static per-channel kernel of shape ``[kernel_size, hidden]``; + ``dynamic`` carries one extra coefficient per (position, tap, group) so every + ``group_size`` channels share a correction. Faithful to the upstream + ``_grouped_dynamic_convolve`` in z-lab/dflash. + + Args: + hidden: ``[batch, length, hidden]`` activations. ``length`` must already + be block-local, since tap ``offset`` reads position ``t - offset``. + dynamic: ``[batch, length, kernel_size, groups]`` correction coefficients. + base: ``[kernel_size, hidden]`` static kernel. + group_size: Number of channels sharing one dynamic coefficient. + + Returns: + torch.Tensor: Same shape as ``hidden``. + """ + batch, length, hidden_size = hidden.shape + groups = hidden_size // group_size + blocks = hidden.view(batch, length, groups, group_size) + dynamic = dynamic.reshape(batch, length, base.shape[0], groups, 1) + output = torch.zeros_like(blocks) + for offset in range(base.shape[0]): + if offset == 0: + values = blocks + else: + values = F.pad(blocks[:, :-offset], (0, 0, 0, 0, offset, 0)) + kernel = base[offset].view(1, 1, groups, group_size).to(hidden.dtype) + output = output + kernel * values + output = torch.addcmul(output, dynamic[:, :, offset].to(hidden.dtype), values) + return output.view_as(hidden) + + +class GroupedDynamicCausalConv(nn.Module): + """Two-tap dynamic depthwise convolution wrapped around a transformer sublayer. + + ``prepare`` runs before the sublayer and also emits the dynamic kernel that + ``finish`` applies after it, so one projection feeds both taps. + + Unlike the upstream inference implementation, which only ever sees a single + draft block, the training path packs ``n_blocks`` blocks into one flat + sequence. The convolution is causal along the sequence axis, so it is applied + block-locally here; otherwise tap 1 at a block's first position would read the + last position of the previous, unrelated block. + """ + + def __init__( + self, hidden_size: int, kernel_size: int, group_size: int, block_size: int + ): + super().__init__() + if hidden_size % group_size != 0: + raise ValueError( + f"DFlash2 conv_group_size={group_size} must divide hidden_size={hidden_size}" + ) + self.hidden_size = int(hidden_size) + self.kernel_size = int(kernel_size) + self.group_size = int(group_size) + self.block_size = int(block_size) + groups = hidden_size // group_size + # Start as an identity passthrough (tap 0 = 1, later taps = 0) with a + # zeroed projection, so a freshly built DFlash2 is numerically identical + # to DFlash and learns the correction from there. + base_kernel = torch.zeros(2, self.kernel_size, hidden_size) + base_kernel[:, 0, :] = 1.0 + self.base_kernel = nn.Parameter(base_kernel) + self.kernel_projection = nn.Linear( + hidden_size, 2 * self.kernel_size * groups, bias=False + ) + nn.init.zeros_(self.kernel_projection.weight) + + def _to_block_local( + self, hidden: torch.Tensor + ) -> tuple[torch.Tensor, tuple[int, ...]]: + batch, length, hidden_size = hidden.shape + if length % self.block_size != 0: + raise ValueError( + f"DFlash2 conv expects the draft sequence length ({length}) to be a " + f"multiple of block_size ({self.block_size}); the causal taps must not " + "cross a block boundary." + ) + n_blocks = length // self.block_size + return ( + hidden.reshape(batch * n_blocks, self.block_size, hidden_size), + (batch, length, hidden_size), + ) + + def prepare(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + block_hidden, shape = self._to_block_local(hidden) + groups = self.hidden_size // self.group_size + dynamic = self.kernel_projection(block_hidden).view( + *block_hidden.shape[:-1], 2, self.kernel_size, groups + ) + convolved = grouped_dynamic_convolve( + block_hidden, dynamic[..., 0, :, :], self.base_kernel[0], self.group_size + ) + return convolved.reshape(*shape), dynamic[..., 1, :, :] + + def finish(self, hidden: torch.Tensor, dynamic: torch.Tensor) -> torch.Tensor: + block_hidden, shape = self._to_block_local(hidden) + convolved = grouped_dynamic_convolve( + block_hidden, dynamic, self.base_kernel[1], self.group_size + ) + return convolved.reshape(*shape) + + +class CandidateSelector(nn.Module): + """Scores adjacent (predecessor, candidate) pairs to pick one coherent path. + + For block position ``t``, candidate ``b`` and the token ``a`` chosen at + ``t - 1``:: + + S_t(a, b) = U_t(b) + + + where ``U_t(b)`` is the drafter's own logit for ``b``, ``A``/``B`` are the + predecessor/successor codebooks and ``H`` projects the backbone hidden state. + """ + + def __init__(self, config: DFlash2Config): + super().__init__() + self.rank = int(config.selector_rank) + self.top_k = int(config.selector_top_k) + self.predecessor_codebook = nn.Embedding(config.vocab_size, self.rank) + self.successor_codebook = nn.Embedding(config.vocab_size, self.rank) + self.hidden_projection = nn.Linear(config.hidden_size, self.rank, bias=False) + + def pair_scores( + self, + hidden: torch.Tensor, + unary: torch.Tensor, + candidates: torch.Tensor, + predecessor_ids: torch.Tensor, + ) -> torch.Tensor: + """Score every candidate against a known predecessor, all positions at once. + + Training teacher-forces the predecessor (it is the ground-truth previous + token), so unlike :meth:`select` this needs no sequential loop. + + Args: + hidden: ``[n, hidden]`` backbone states for the scored positions. + unary: ``[n, k]`` drafter logits for the candidates. + candidates: ``[n, k]`` candidate token ids. + predecessor_ids: ``[n]`` ground-truth previous token per position. + + Returns: + torch.Tensor: ``[n, k]`` pair scores. + """ + projected = self.hidden_projection(hidden) + gated = self.predecessor_codebook(predecessor_ids) * projected + successor = self.successor_codebook(candidates) + return unary + torch.einsum("nr,nkr->nk", gated, successor) + + @torch.no_grad() + def select( + self, + hidden: torch.Tensor, + logits: torch.Tensor, + anchor_ids: torch.Tensor, + temperature: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Greedy/sampled path trace through the per-position candidate sets. + + Mirrors the upstream inference-time selector: the predecessor at position + ``t`` is whatever was chosen at ``t - 1``, seeded by ``anchor_ids``. + + Args: + hidden: ``[batch, block, hidden]`` backbone states. + logits: ``[batch, block, vocab]`` drafter logits. + anchor_ids: ``[batch]`` token preceding the block. + temperature: 0 for argmax, otherwise softmax sampling. + + Returns: + tuple: ``(path [batch, block], candidates [batch, block, k])``. + """ + top_k = min(self.top_k, logits.shape[-1]) + unary, candidates = torch.topk(logits, top_k, dim=-1, sorted=False) + projected = self.hidden_projection(hidden) + predecessor = anchor_ids + path = [] + for position in range(hidden.shape[1]): + gated = self.predecessor_codebook(predecessor) * projected[:, position] + scores = unary[:, position] + torch.einsum( + "br,bkr->bk", gated, self.successor_codebook(candidates[:, position]) + ) + if temperature > 0: + probs = torch.softmax(scores.float() / temperature, dim=-1) + index = torch.multinomial(probs, num_samples=1)[:, 0] + else: + index = torch.argmax(scores, dim=-1) + predecessor = candidates[:, position].gather(-1, index[:, None])[:, 0] + path.append(predecessor) + return torch.stack(path, dim=1), candidates + + +class DFlash2DraftModel(DFlashDraftModel): + """DFlash block drafter plus the DFlash2 dynamic convolutions and selector. + + The convolutions are attached to the existing :class:`DFlashDecoderLayer` + hooks, which are inert (``None``) for every other DFlash-family drafter, so + DFlash / Domino / DSpark / JetSpec keep their exact behaviour. + """ + + config_class = DFlash2Config + + def __init__(self, config: DFlash2Config): + super().__init__(config) + self.block_size = int(getattr(config, "block_size", 8)) + self.conv_kernel_size = int(getattr(config, "conv_kernel_size", 2)) + self.conv_group_size = int(getattr(config, "conv_group_size", 16)) + for layer in self.layers: + layer.attention_conv = GroupedDynamicCausalConv( + config.hidden_size, + self.conv_kernel_size, + self.conv_group_size, + self.block_size, + ) + layer.mlp_conv = GroupedDynamicCausalConv( + config.hidden_size, + self.conv_kernel_size, + self.conv_group_size, + self.block_size, + ) + self.candidate_selector = CandidateSelector(config) + + +__all__ = [ + "CandidateSelector", + "DFlash2Config", + "DFlash2DraftModel", + "GroupedDynamicCausalConv", + "grouped_dynamic_convolve", +] diff --git a/verl_speco/trainer/base_trainer.py b/verl_speco/trainer/base_trainer.py index 0975145a..f82e6855 100644 --- a/verl_speco/trainer/base_trainer.py +++ b/verl_speco/trainer/base_trainer.py @@ -818,13 +818,14 @@ def _has_mesh_dim(self, dim_name: str) -> bool: def _is_block_drafter_backend(self) -> bool: return getattr(self.backend, "model_type", None) in { "dflash", + "dflash2", "dspark", "domino", } def _block_drafter_metric_prefix(self) -> str: model_type = str(getattr(self.backend, "model_type", "dflash") or "dflash") - if model_type in {"dspark", "domino"}: + if model_type in {"dspark", "domino", "dflash2"}: return model_type return "dflash" @@ -1030,7 +1031,7 @@ def _build_draft_model(self): pending_target_weight = self._pending_target_lm_head_weight if ( getattr(self.backend, "model_type", None) - in {"eagle3", "dflash", "dspark", "domino"} + in {"eagle3", "dflash", "dflash2", "dspark", "domino"} and torch.is_tensor(pending_target_weight) and pending_target_weight.dim() == 2 ): diff --git a/verl_speco/trainer/draft_training_loop.py b/verl_speco/trainer/draft_training_loop.py index 6ea7e277..096d9b5d 100644 --- a/verl_speco/trainer/draft_training_loop.py +++ b/verl_speco/trainer/draft_training_loop.py @@ -305,6 +305,25 @@ def _fill_if_missing( "target_num_hidden_layers", ), ), + # DFlash2 is served as a DFlash checkpoint whose dflash_config carries the + # selector/conv hyperparameters, matching the released z-lab layout. + "dflash2": ( + "dflash_config", + ( + "block_size", + "num_anchors", + "loss_decay_gamma", + "conv_kernel_size", + "conv_group_size", + "selector_rank", + "selector_top_k", + "target_layer_ids", + "num_context_layers", + "num_target_layers", + "target_num_hidden_layers", + "mask_token_id", + ), + ), "dspark": ( "dspark_config", ( @@ -339,7 +358,7 @@ def _rewrite_standalone_block_runtime_config( contract and only merge the alias fields needed by vLLM/SGLang. """ backend_type = getattr(getattr(trainer, "backend", None), "model_type", None) - if backend_type not in {"dflash", "dspark", "domino"}: + if backend_type not in {"dflash", "dflash2", "dspark", "domino"}: return if completed_future is not None: From f9467647f88fd3e0237a38be467b143b3c8ec775 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 16:59:00 +0800 Subject: [PATCH 02/16] test: add DFlash2 GPU training smoke Mirrors tests/special_standalone/domino_gpu_smoke.py: drives the real training path on a real target, reporting the DFlash CE signals plus the DFlash2-specific selector_loss / selector_acc / selector_coverage. Signed-off-by: khazic --- tests/special_standalone/dflash2_gpu_smoke.py | 191 ++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 tests/special_standalone/dflash2_gpu_smoke.py diff --git a/tests/special_standalone/dflash2_gpu_smoke.py b/tests/special_standalone/dflash2_gpu_smoke.py new file mode 100644 index 00000000..794ead68 --- /dev/null +++ b/tests/special_standalone/dflash2_gpu_smoke.py @@ -0,0 +1,191 @@ +# 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. +"""Hardware smoke test for the DFlash2 drafter training backend. + +Drives the real training path on GPU with a real target model (default +Qwen3-4B): it runs the frozen target forward to collect the DFlash-style +multi-layer context hidden states, builds the DFlash2 draft via +``DFlash2TrainerBackend.build_model``, and runs several optimizer steps through +``compute_loss`` (which invokes the block-drafter forward with the two-tap +dynamic convolutions and the candidate-selector objective). + +The draft is cold-started, so the useful signals are: + * ``loss`` trending down and ``accuracy`` rising (the DFlash CE path, now + routed through the dynamic convolutions), + * ``selector_loss`` trending down and ``selector_acc`` rising (the DFlash2 + candidate selector learning to re-rank the drafter's own top-k), + * ``selector_coverage`` (fraction of scored rows whose ground truth survived + the drafter's top-k) rising as the backbone improves. + +Run: + python tests/special_standalone/dflash2_gpu_smoke.py \ + --target /path/to/target-model --steps 120 +""" + +from __future__ import annotations + +import argparse + +import torch +from omegaconf import OmegaConf +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +PROMPTS = [ + "Explain why the sky appears blue during the day, in a few sentences.", + "Write a short Python function that returns the nth Fibonacci number and explain it.", + "Summarize the water cycle and its main stages in a short paragraph.", + "Describe the main differences between TCP and UDP for a networking student.", +] + + +def _build_batch(target, tokenizer, target_layer_ids, device): + """One packed batch: input_ids, loss_mask, and concatenated context hidden states.""" + id_chunks, mask_chunks, hidden_chunks = [], [], [] + for text in PROMPTS: + messages = [{"role": "user", "content": text}] + prompt = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + enc = tokenizer(prompt, return_tensors="pt").to(device) + with torch.no_grad(): + out = target(input_ids=enc["input_ids"], output_hidden_states=True) + layers = [out.hidden_states[i][0] for i in target_layer_ids] # each [S, H] + id_chunks.append(enc["input_ids"][0]) + mask_chunks.append(torch.ones(enc["input_ids"].size(1), device=device)) + hidden_chunks.append(torch.cat(layers, dim=-1)) # [S, num_ctx*H] + + input_ids = torch.cat(id_chunks).unsqueeze(0) + loss_mask = torch.cat(mask_chunks).unsqueeze(0) + hidden = torch.cat(hidden_chunks).unsqueeze(0).to(torch.bfloat16) + return { + "input_ids": input_ids, + "loss_mask": loss_mask, + "hidden_states": hidden, + "attention_mask": torch.ones_like(input_ids), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--target", required=True, help="path or HF id of the target causal LM" + ) + parser.add_argument("--steps", type=int, default=120) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--num-context-layers", type=int, default=5) + parser.add_argument("--block-size", type=int, default=8) + parser.add_argument("--selector-top-k", type=int, default=16) + parser.add_argument("--selector-rank", type=int, default=256) + parser.add_argument("--selector-loss-weight", type=float, default=1.0) + args = parser.parse_args() + + device = "cuda" + torch.manual_seed(0) + + print(f"[smoke] loading target {args.target}") + tokenizer = AutoTokenizer.from_pretrained(args.target) + target = ( + AutoModelForCausalLM.from_pretrained(args.target, torch_dtype=torch.bfloat16) + .to(device) + .eval() + ) + target_cfg = AutoConfig.from_pretrained(args.target) + + from verl_speco.models.dflash import build_target_layer_ids + + target_layers = int(getattr(target_cfg, "num_hidden_layers")) + target_layer_ids = build_target_layer_ids(args.num_context_layers, target_layers) + print(f"[smoke] context layers={target_layer_ids} (of {target_layers})") + + batch = _build_batch(target, tokenizer, target_layer_ids, device) + print( + f"[smoke] batch seq_len={batch['input_ids'].size(1)} " + f"hidden={batch['hidden_states'].size(-1)}" + ) + + cfg = OmegaConf.create( + { + "rollout": { + "drafter": { + "speculative_algorithm": "DFLASH2", + "model_path": "/dev/null/does-not-exist", + "training": { + "dflash2_block_size": args.block_size, + "dflash2_num_anchors": 128, + "dflash2_num_target_layers": args.num_context_layers, + "dflash2_num_hidden_layers": 1, + "dflash2_selector_top_k": args.selector_top_k, + "dflash2_selector_rank": args.selector_rank, + "dflash2_selector_loss_weight": args.selector_loss_weight, + "lr": args.lr, + }, + } + }, + "model": {"path": args.target}, + } + ) + + from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainerBackend + + backend = DFlash2TrainerBackend(cfg, target_cfg) + model, drafter_cfg = backend.build_model() + model = model.to(device).to(torch.bfloat16).train() + backend.target_lm_head = backend.target_lm_head.to(device).to(torch.bfloat16) + optimizer = backend.setup_optimizer(model, cfg.rollout.drafter.training) + n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + conv = model.draft_model.layers[0].attention_conv + print( + f"[smoke] block_size={model.block_size} conv_kernel={conv.kernel_size} " + f"conv_group={conv.group_size} selector_rank={drafter_cfg.selector_rank} " + f"selector_top_k={drafter_cfg.selector_top_k} trainable_params={n_params:,}" + ) + + first = None + for step in range(args.steps): + out = backend.compute_loss(model, batch, 0) + num_tokens = out["local_num_tokens"].clamp_min(1) + loss = out["total_local_ploss"] / num_tokens + optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + + if step % 10 == 0 or step == args.steps - 1: + d = out["diagnostics"] + total = float(loss) + acc = float(out["accuracy"]) + sel_loss = float(d["selector_loss"]) + sel_tokens = max(float(d["selector_token_count"]), 1.0) + sel_acc = float(d["selector_correct_count"]) / sel_tokens + coverage = float(d["selector_coverage_count"]) / max( + float(d["selector_active_count"]), 1.0 + ) + if first is None: + first = (total, acc, sel_loss, sel_acc, coverage) + print( + f"[smoke] step {step:3d} loss={total:.4f} acc={acc:.4f} " + f"selector_loss={sel_loss:.4f} selector_acc={sel_acc:.4f} " + f"selector_coverage={coverage:.4f}" + ) + + print( + f"[smoke] DONE loss {first[0]:.4f}->{total:.4f} acc {first[1]:.4f}->{acc:.4f} " + f"selector_loss {first[2]:.4f}->{sel_loss:.4f} " + f"selector_acc {first[3]:.4f}->{sel_acc:.4f} " + f"selector_coverage {first[4]:.4f}->{coverage:.4f}" + ) + + +if __name__ == "__main__": + main() From 8bc16f965558ef854cdd316a3e14fb8788553775 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 17:08:37 +0800 Subject: [PATCH 03/16] test: drive the DFlash2 smoke under autocast like the real trainer The smoke hard-cast the module to bf16 (copied from the Domino smoke), which makes the shared DFlash forward assign a bf16 cross_entropy result into its fp32 loss_per_token buffer. The real training path keeps fp32 parameters and runs the forward under torch.amp.autocast(bfloat16), which leaves cross_entropy in fp32. Match that instead; the Domino smoke only survives the hard cast because its own forward explicitly floats the logits. Signed-off-by: khazic --- tests/special_standalone/dflash2_gpu_smoke.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/special_standalone/dflash2_gpu_smoke.py b/tests/special_standalone/dflash2_gpu_smoke.py index 794ead68..df8359c0 100644 --- a/tests/special_standalone/dflash2_gpu_smoke.py +++ b/tests/special_standalone/dflash2_gpu_smoke.py @@ -140,8 +140,11 @@ def main() -> None: backend = DFlash2TrainerBackend(cfg, target_cfg) model, drafter_cfg = backend.build_model() - model = model.to(device).to(torch.bfloat16).train() - backend.target_lm_head = backend.target_lm_head.to(device).to(torch.bfloat16) + # Match the real training path (base_trainer): fp32 parameters driven under + # bf16 autocast, rather than hard-casting the module to bf16. Autocast keeps + # cross_entropy in fp32, which is what the shared DFlash forward expects. + model = model.to(device).train() + backend.target_lm_head = backend.target_lm_head.to(device) optimizer = backend.setup_optimizer(model, cfg.rollout.drafter.training) n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) conv = model.draft_model.layers[0].attention_conv @@ -153,7 +156,8 @@ def main() -> None: first = None for step in range(args.steps): - out = backend.compute_loss(model, batch, 0) + with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): + out = backend.compute_loss(model, batch, 0) num_tokens = out["local_num_tokens"].clamp_min(1) loss = out["total_local_ploss"] / num_tokens optimizer.zero_grad() From dd86f507cf1023b2cdb4acdacb78aa89d10a8138 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 17:10:57 +0800 Subject: [PATCH 04/16] test: give the DFlash2 smoke bf16 params AND autocast The production forward stacks both: FSDP MixedPrecision(param_dtype=bf16) makes the parameters bf16, and the surrounding autocast keeps cross_entropy in fp32. Emulating only one of the two breaks in a different place each time (bf16 params alone scatter a bf16 CE result into the fp32 loss_per_token buffer; autocast alone leaves the RMSNorm weights fp32, so attention q/k become fp32 while v stays bf16). Signed-off-by: khazic --- tests/special_standalone/dflash2_gpu_smoke.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/special_standalone/dflash2_gpu_smoke.py b/tests/special_standalone/dflash2_gpu_smoke.py index df8359c0..fe5ba641 100644 --- a/tests/special_standalone/dflash2_gpu_smoke.py +++ b/tests/special_standalone/dflash2_gpu_smoke.py @@ -140,11 +140,13 @@ def main() -> None: backend = DFlash2TrainerBackend(cfg, target_cfg) model, drafter_cfg = backend.build_model() - # Match the real training path (base_trainer): fp32 parameters driven under - # bf16 autocast, rather than hard-casting the module to bf16. Autocast keeps - # cross_entropy in fp32, which is what the shared DFlash forward expects. - model = model.to(device).train() - backend.target_lm_head = backend.target_lm_head.to(device) + # Match the real training path, which stacks two things: FSDP + # MixedPrecision(param_dtype=bf16) gives the forward bf16 parameters (so the + # RMSNorm weights stay bf16 and attention q/k/v agree), and the surrounding + # torch.amp.autocast keeps cross_entropy in fp32 (so its result can be + # scattered into the fp32 loss_per_token buffer). Emulate both here. + model = model.to(device).to(torch.bfloat16).train() + backend.target_lm_head = backend.target_lm_head.to(device).to(torch.bfloat16) optimizer = backend.setup_optimizer(model, cfg.rollout.drafter.training) n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) conv = model.draft_model.layers[0].attention_conv From d5a37db5cece034359ed38726352dda82e54ce9a Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 17:14:19 +0800 Subject: [PATCH 05/16] test: cover DFLASH2 in the drafter factory contract test_factory_lists_every_supported_algorithm asserts the exact supported set, so it has to learn about DFLASH2. Also parametrize the hidden-states layout test over it: DFlash2 consumes the same DFlash aux context layers, and tagging it eagle3_aux_plus_last would make DFlash preprocessing fail closed. Signed-off-by: khazic --- tests/integration/test_drafter_backend_factory_contract.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/integration/test_drafter_backend_factory_contract.py b/tests/integration/test_drafter_backend_factory_contract.py index 3d17c118..0f2bf14f 100644 --- a/tests/integration/test_drafter_backend_factory_contract.py +++ b/tests/integration/test_drafter_backend_factory_contract.py @@ -34,6 +34,7 @@ def test_factory_lists_every_supported_algorithm() -> None: "EAGLE2", "EAGLE3", "DFLASH", + "DFLASH2", "DSPARK", "DOMINO", "PEAGLE", @@ -61,6 +62,9 @@ def test_factory_rejects_unknown_algorithm_before_importing_a_backend() -> None: # Tagging it eagle3_aux_plus_last makes DFlash preprocessing fail closed. ("DOMINO", {}, "dflash_aux"), ("domino", {}, "dflash_aux"), + # DFlash2 is likewise a DFlash variant on the same context layers. + ("DFLASH2", {}, "dflash_aux"), + ("dflash2", {}, "dflash_aux"), ("DSPARK", {}, "dflash_aux_plus_last"), ("DSPARK", {"dspark_l1_loss_alpha": 0.0}, "dflash_aux"), ("DSPARK", {"dspark_l1_loss_alpha": None}, "dflash_aux"), From 0a3239913c0c48696313269af7b294b22f040218 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 17:28:58 +0800 Subject: [PATCH 06/16] fix(dflash2): address review of the DFlash2 backend - Pin the conv block size to the trainer's block size. They came from two independent sources (checkpoint config vs dflash2_block_size), and when the training value is a multiple of the config value the block-multiple guard still passes while each conv block spans several anchor blocks, so the causal tap reads across an anchor boundary silently. Resolve once in build_model. - Gate dflash2_num_target_layers on the running algorithm. It sat after domino_num_target_layers in a flat priority list, and speco_base.yaml defines the Domino key unconditionally, so it could never be reached; a DFLASH2 run with a non-default value would build a drafter expecting N context layers while the rollout collected 5. - Compute the selector loss over all active rows, zero-weighting the ones whose ground truth missed the drafter's top-k, instead of slicing them out. Slicing made the autograd graph coverage-dependent, so on a cold-started drafter the selector parameters could receive gradients on some ranks and not others and desync the FSDP2/DDP reduction. - Log the selector diagnostics: they were computed and returned but absent from _record_dflash_training_metrics' scalar_keys, so the only signals showing whether the selector learns were dropped outside the standalone smoke. - Wire DFLASH2 into both SGLang gates: it now requests the DFlash aux hidden states like the rest of the family, and rejects DFLASH2 as an engine-level algorithm with a Domino-style message instead of forwarding the raw string to ServerArgs. - Drop the dead _normalize_dflash_config override (num_context_layers defaults to 5, never None, so the guard never fired). - Tests: importorskip transformers in the four tests that reach it through DFlashConfig (they errored instead of skipping without it), fix the smoke path in the module docstring, and add a regression test for the block-size drift. Signed-off-by: khazic --- .../test_dflash2_backend_contract.py | 44 +++++++++- .../backends/dflash2_trainer_backend.py | 81 ++++++++++--------- .../integration/oldlogprob_layer_ids.py | 21 ++++- verl_speco/integration/sglang_runtime.py | 14 +++- verl_speco/trainer/base_trainer.py | 6 ++ 5 files changed, 123 insertions(+), 43 deletions(-) diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index 345e7a55..c6d31eb1 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -4,7 +4,7 @@ convolution and the candidate selector), the block-locality invariant the training layout requires, the algorithm routing, and the block-drafter classification. The full training forward is validated on GPU by -``ci/dflash2_gpu_smoke.py``. +``tests/special_standalone/dflash2_gpu_smoke.py``. """ from __future__ import annotations @@ -100,6 +100,7 @@ def test_conv_is_identity_at_init() -> None: start numerically equal to DFlash rather than perturbing the backbone. """ pytest.importorskip("torch") + pytest.importorskip("transformers") import torch from verl_speco.models.dflash2 import GroupedDynamicCausalConv @@ -121,6 +122,7 @@ def test_conv_does_not_leak_across_block_boundaries() -> None: final position of block i-1, which belongs to an unrelated anchor. """ pytest.importorskip("torch") + pytest.importorskip("transformers") import torch from verl_speco.models.dflash2 import GroupedDynamicCausalConv @@ -149,6 +151,7 @@ def test_conv_does_not_leak_across_block_boundaries() -> None: def test_conv_rejects_length_that_is_not_a_block_multiple() -> None: pytest.importorskip("torch") + pytest.importorskip("transformers") import torch from verl_speco.models.dflash2 import GroupedDynamicCausalConv @@ -167,6 +170,7 @@ def test_selector_pair_scores_match_the_sequential_selector() -> None: sequential ``select`` actually took has to reproduce the same scores. """ pytest.importorskip("torch") + pytest.importorskip("transformers") import torch from verl_speco.models.dflash2 import CandidateSelector @@ -260,6 +264,44 @@ def test_dflash2_training_model_rejects_restricted_vocab() -> None: ) +def test_backend_pins_conv_block_size_to_the_trainer_block_size() -> None: + """The convs and the trainer must not disagree about the block size. + + The convs are built from the drafter config while the training wrapper takes + its own ``dflash2_block_size``. If the training value is a multiple of the + config value, ``_to_block_local``'s guard still passes but each conv "block" + spans several anchor blocks, so the causal tap reads across an anchor + boundary: exactly the leak the module is supposed to prevent, silently. + """ + pytest.importorskip("torch") + pytest.importorskip("transformers") + from omegaconf import OmegaConf + + from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainerBackend + from verl_speco.models.dflash2 import DFlash2Config + + backend = DFlash2TrainerBackend( + OmegaConf.create( + { + "rollout": { + "drafter": { + "speculative_algorithm": "DFLASH2", + "model_path": "", + # Deliberately a multiple of the config's block_size=4, + # so the block-multiple guard alone would not catch it. + "training": {"dflash2_block_size": 8}, + } + }, + "model": {"path": ""}, + } + ), + None, + ) + config = _tiny_dflash2_config(block_size=4) + assert isinstance(config, DFlash2Config) + assert backend._resolved_block_size(config) == 8 + + def test_dflash2_backend_is_registered_in_the_factory() -> None: from verl_speco.backends.factory import SUPPORTED_DRAFTER_ALGORITHMS diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py index dcb3cc0f..faab3de0 100644 --- a/verl_speco/backends/dflash2_trainer_backend.py +++ b/verl_speco/backends/dflash2_trainer_backend.py @@ -86,35 +86,37 @@ def _auxiliary_loss( # The selector can only learn on rows whose ground truth survived the # drafter's own top-k; elsewhere there is no correct choice to make. + # Those rows are zero-weighted rather than sliced out, so the autograd + # graph covers the same parameters on every rank. Slicing here would + # make the selector parameters receive gradients only on the ranks that + # happened to have coverage this step, which desyncs the FSDP2/DDP + # gradient reduction. target_hits = candidates == active_targets.unsqueeze(-1) learnable = target_hits.any(dim=-1) - selector_loss = torch.zeros((), dtype=torch.float32, device=device) - selector_correct = torch.zeros((), dtype=torch.float32, device=device) - selector_tokens = torch.zeros((), dtype=torch.float32, device=device) - if bool(learnable.any()): - learn_scores = scores[learnable] - learn_weights = active_weights[learnable].float() - learn_targets = torch.argmax(target_hits[learnable].int(), dim=-1) - per_row = F.cross_entropy(learn_scores, learn_targets, reduction="none") - finite = torch.isfinite(per_row) - per_row = torch.where(finite, per_row, torch.zeros_like(per_row)) - learn_weights = learn_weights * finite.to(learn_weights.dtype) - selector_loss = (per_row * learn_weights).sum() / learn_weights.sum().clamp( - min=1e-6 - ) - with torch.no_grad(): - selector_tokens = finite.float().sum() - selector_correct = ( - ((torch.argmax(learn_scores, dim=-1) == learn_targets) & finite) - .float() - .sum() - ) + # argmax over an all-False row yields 0; the row is masked out below. + row_targets = torch.argmax(target_hits.int(), dim=-1) + per_row = F.cross_entropy(scores, row_targets, reduction="none") + finite = torch.isfinite(per_row) + per_row = torch.where(finite, per_row, torch.zeros_like(per_row)) + row_weights = ( + active_weights.float() + * learnable.to(torch.float32) + * finite.to(torch.float32) + ) + selector_loss = (per_row * row_weights).sum() / row_weights.sum().clamp( + min=1e-6 + ) with torch.no_grad(): + scored = learnable & finite metrics = { "selector_loss": selector_loss.detach().float(), - "selector_correct_count": selector_correct, - "selector_token_count": selector_tokens, + "selector_correct_count": ( + (torch.argmax(scores, dim=-1) == row_targets) & scored + ) + .float() + .sum(), + "selector_token_count": scored.float().sum(), "selector_coverage_count": learnable.float().sum(), "selector_active_count": torch.tensor( float(active_targets.numel()), dtype=torch.float32, device=device @@ -134,17 +136,20 @@ def _training_value(self, training_cfg, dflash2_key: str, dflash_key: str, defau return value return training_cfg.get(dflash_key, default) - def _normalize_dflash_config( - self, drafter_config, target_hf_config, normalized_state, spec_model_path - ): + def _resolved_block_size(self, drafter_config) -> int: + """Single source of truth for the block size. + + The convolutions are built from ``drafter_config.block_size`` while the + training wrapper takes its own ``dflash2_block_size``. If those disagree + the conv's notion of a block spans several anchor blocks, and its causal + tap silently reads across an anchor boundary without tripping the + block-multiple guard. Resolve once and apply to both. + """ training_cfg = self.config.rollout.drafter.training - if training_cfg.get("dflash2_num_target_layers", None) is not None: - if getattr(drafter_config, "num_context_layers", None) is None: - drafter_config.num_context_layers = int( - training_cfg["dflash2_num_target_layers"] - ) - return super()._normalize_dflash_config( - drafter_config, target_hf_config, normalized_state, spec_model_path + return int( + training_cfg.get( + "dflash2_block_size", getattr(drafter_config, "block_size", 8) + ) ) def _build_fallback_config(self, target_hf_config): @@ -257,6 +262,10 @@ def build_model(self): drafter_config = self._normalize_dflash_config( drafter_config, target_hf_config, normalized_state, spec_model_path ) + # Pin the config's block size to the one the trainer will use, so the + # convolutions and the block layout cannot drift apart. + block_size = self._resolved_block_size(drafter_config) + drafter_config.block_size = block_size draft_model = DFlash2DraftModel(deepcopy(drafter_config)) if ( @@ -276,11 +285,7 @@ def build_model(self): training_cfg = self.config.rollout.drafter.training return DFlash2TrainingModel( draft_model=draft_model, - block_size=int( - training_cfg.get( - "dflash2_block_size", getattr(drafter_config, "block_size", 8) - ) - ), + block_size=block_size, num_anchors=int( training_cfg.get( "dflash2_num_anchors", getattr(drafter_config, "num_anchors", 512) diff --git a/verl_speco/integration/oldlogprob_layer_ids.py b/verl_speco/integration/oldlogprob_layer_ids.py index 3100023f..b0405f1b 100644 --- a/verl_speco/integration/oldlogprob_layer_ids.py +++ b/verl_speco/integration/oldlogprob_layer_ids.py @@ -180,7 +180,11 @@ def _build_dflash_target_layer_ids( def _dflash_num_context_layers( - drafter_cfg: Any, model_configs: tuple[Any, ...], *, is_dspark: bool = False + drafter_cfg: Any, + model_configs: tuple[Any, ...], + *, + is_dspark: bool = False, + is_dflash2: bool = False, ) -> int: training_cfg = _get_nested(drafter_cfg, ("training",), {}) or {} candidates = [] @@ -188,8 +192,14 @@ def _dflash_num_context_layers( candidates.append( _get_nested(training_cfg, ("dspark_num_target_layers",), None) ) + # Variant-specific knobs must be gated on the running algorithm: the base + # YAML defines every variant's *_num_target_layers unconditionally, so an + # ungated entry would let one variant's default shadow another's. + if is_dflash2: + candidates.append( + _get_nested(training_cfg, ("dflash2_num_target_layers",), None) + ) candidates.append(_get_nested(training_cfg, ("domino_num_target_layers",), None)) - candidates.append(_get_nested(training_cfg, ("dflash2_num_target_layers",), None)) candidates.extend( ( _get_nested(training_cfg, ("dflash_num_target_layers",), None), @@ -290,7 +300,12 @@ def resolve_oldlogprob_aux_layer_ids( if target_num_hidden_layers is None: return None return _build_dflash_target_layer_ids( - _dflash_num_context_layers(drafter_cfg, model_configs, is_dspark=is_dspark), + _dflash_num_context_layers( + drafter_cfg, + model_configs, + is_dspark=is_dspark, + is_dflash2=algorithm == "DFLASH2", + ), int(target_num_hidden_layers), ) diff --git a/verl_speco/integration/sglang_runtime.py b/verl_speco/integration/sglang_runtime.py index d13bf59d..9afdfd9c 100644 --- a/verl_speco/integration/sglang_runtime.py +++ b/verl_speco/integration/sglang_runtime.py @@ -174,7 +174,7 @@ def _drafter_uses_dflash_aux_hidden(drafter_cfg: dict[str, Any]) -> bool: algorithm = str(drafter_cfg.get("speculative_algorithm", "") or "").upper() training_cfg = drafter_cfg.get("training") or {} return bool( - algorithm in {"DFLASH", "DSPARK"} + algorithm in {"DFLASH", "DFLASH2", "DSPARK"} and drafter_cfg.get("enable") and drafter_cfg.get("enable_drafter_training") and training_cfg.get("collect_hidden_states_from_sgl") @@ -603,6 +603,18 @@ def _server_args_overrides_from_drafter( "enables the Domino correction head on engines that support it, keeping DOMINO for " "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 SGLang ServerArgs algorithm, so fail loud instead of forwarding + # the raw string. + raise ValueError( + "DFLASH2 is not an engine-level speculative algorithm; DFlash2 is served as a DFlash " + "checkpoint. Set actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH for the " + "rollout/serve path; the trained checkpoint's dflash_config carries the DFlash2 " + "convolution and selector hyperparameters, keeping DFLASH2 for drafter training." + ) rollout_cfg = drafter_cfg.get("rollout") or {} training_cfg = drafter_cfg.get("training") or {} diff --git a/verl_speco/trainer/base_trainer.py b/verl_speco/trainer/base_trainer.py index f82e6855..f3178825 100644 --- a/verl_speco/trainer/base_trainer.py +++ b/verl_speco/trainer/base_trainer.py @@ -948,6 +948,12 @@ def _record_dflash_training_metrics(self, loss_dict: dict[str, Any]) -> None: "ce_weighted_token_count": f"{prefix}/ce_weighted_token_count", "l1_loss_sum": f"{prefix}/l1_loss_sum", "l1_weighted_token_count": f"{prefix}/l1_weighted_token_count", + # DFlash2 candidate selector. + "selector_loss": f"{prefix}/selector_loss", + "selector_correct_count": f"{prefix}/selector_correct_count", + "selector_token_count": f"{prefix}/selector_token_count", + "selector_coverage_count": f"{prefix}/selector_coverage_count", + "selector_active_count": f"{prefix}/selector_active_count", "sanitized_rows": f"{prefix}/sanitized_rows", "masked_rows": f"{prefix}/masked_rows", "sampled_vocab_size": f"{prefix}/sampled_vocab_size", From bfa50b8f4a5d447799ce1e59cf5119534d463e09 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 18:10:57 +0800 Subject: [PATCH 07/16] test(dflash2): measure the selector's lift over unary-only ranking selector_acc on its own is not evidence that the selector learned anything. S_t(a, b) starts from the drafter's own logit U_t(b), so once the backbone ranks the ground truth first by itself the selector scores perfectly whether or not the bilinear term contributes. Report the unary-only ranking accuracy on the same scored rows as the control, so the lift between them isolates the selector's actual contribution. Signed-off-by: khazic --- tests/special_standalone/dflash2_gpu_smoke.py | 5 +++++ verl_speco/backends/dflash2_trainer_backend.py | 11 +++++++++++ verl_speco/trainer/base_trainer.py | 1 + 3 files changed, 17 insertions(+) diff --git a/tests/special_standalone/dflash2_gpu_smoke.py b/tests/special_standalone/dflash2_gpu_smoke.py index fe5ba641..718731d1 100644 --- a/tests/special_standalone/dflash2_gpu_smoke.py +++ b/tests/special_standalone/dflash2_gpu_smoke.py @@ -174,6 +174,10 @@ def main() -> None: sel_loss = float(d["selector_loss"]) sel_tokens = max(float(d["selector_token_count"]), 1.0) sel_acc = float(d["selector_correct_count"]) / sel_tokens + # Unary-only ranking on the same rows. The selector's score starts + # from the drafter's own logit, so this is the baseline it has to + # beat; sel_acc alone says nothing once the backbone converges. + base_acc = float(d["selector_base_correct_count"]) / sel_tokens coverage = float(d["selector_coverage_count"]) / max( float(d["selector_active_count"]), 1.0 ) @@ -182,6 +186,7 @@ def main() -> None: print( f"[smoke] step {step:3d} loss={total:.4f} acc={acc:.4f} " f"selector_loss={sel_loss:.4f} selector_acc={sel_acc:.4f} " + f"unary_only_acc={base_acc:.4f} lift={sel_acc - base_acc:+.4f} " f"selector_coverage={coverage:.4f}" ) diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py index faab3de0..f944034c 100644 --- a/verl_speco/backends/dflash2_trainer_backend.py +++ b/verl_speco/backends/dflash2_trainer_backend.py @@ -109,6 +109,12 @@ def _auxiliary_loss( with torch.no_grad(): scored = learnable & finite + # Control for the selector's own accuracy. S_t(a, b) starts from the + # drafter's logit U_t(b), so once the backbone ranks the truth first + # on its own the selector scores perfectly without the bilinear term + # having learned anything. Measuring the unary-only ranking on the + # same rows is what isolates the selector's actual contribution: + # the lift is selector_correct_count - selector_base_correct_count. metrics = { "selector_loss": selector_loss.detach().float(), "selector_correct_count": ( @@ -116,6 +122,11 @@ def _auxiliary_loss( ) .float() .sum(), + "selector_base_correct_count": ( + (torch.argmax(unary, dim=-1) == row_targets) & scored + ) + .float() + .sum(), "selector_token_count": scored.float().sum(), "selector_coverage_count": learnable.float().sum(), "selector_active_count": torch.tensor( diff --git a/verl_speco/trainer/base_trainer.py b/verl_speco/trainer/base_trainer.py index f3178825..469e6984 100644 --- a/verl_speco/trainer/base_trainer.py +++ b/verl_speco/trainer/base_trainer.py @@ -951,6 +951,7 @@ def _record_dflash_training_metrics(self, loss_dict: dict[str, Any]) -> None: # DFlash2 candidate selector. "selector_loss": f"{prefix}/selector_loss", "selector_correct_count": f"{prefix}/selector_correct_count", + "selector_base_correct_count": f"{prefix}/selector_base_correct_count", "selector_token_count": f"{prefix}/selector_token_count", "selector_coverage_count": f"{prefix}/selector_coverage_count", "selector_active_count": f"{prefix}/selector_active_count", From 142eb1532aa9118b3a333b1a3445a2ddee63a7a4 Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 18:18:54 +0800 Subject: [PATCH 08/16] test: add the license header to the DFlash2 contract tests check_license.py covers tests/ as well as verl_speco/, so the new contract test file failed the pre-commit job. Every sibling test file already carries it. Signed-off-by: khazic --- tests/integration/test_dflash2_backend_contract.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index c6d31eb1..63706431 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -1,3 +1,16 @@ +# 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. """Contract tests for the DFlash2 drafter backend. CPU-light: they exercise the two DFlash2 modules (the two-tap dynamic From 9b61cc4035b8154f4223daa5a103530ec8add4ca Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 18:50:25 +0800 Subject: [PATCH 09/16] fix(dflash2): keep the selector loss out of the drafter backbone The selector is an auxiliary head that re-ranks candidates the drafter already produced, but its cross-entropy was reaching the backbone through two paths: torch.topk passes gradient through the selected logits, and the selector's hidden projection read the backbone state directly. So the selector objective was reshaping the drafter itself rather than only training the re-ranker. That is wrong on its own terms, and it also invalidates any DFlash/DFlash2 comparison: the two backbones would be trained under different effective objectives instead of differing only by architecture, so a measured delta could not be attributed to the architecture change. Detach both inputs and pin the invariant with a test. Signed-off-by: khazic --- .../test_dflash2_backend_contract.py | 61 +++++++++++++++++++ .../backends/dflash2_trainer_backend.py | 15 ++++- 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index 63706431..0abd31aa 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -259,6 +259,67 @@ def run(selector_loss_weight): ) +def test_selector_loss_does_not_backprop_into_the_backbone() -> None: + """The selector must train as a re-ranking head, not reshape the drafter. + + torch.topk passes gradient through the selected logits and the selector's + hidden projection reads the backbone state, so without detaching, the + selector's cross-entropy would flow back into the draft model. That would + both change the drafter's effective objective and make any DFlash/DFlash2 + comparison confounded, since the two backbones would no longer differ only + by architecture. + """ + pytest.importorskip("torch") + pytest.importorskip("transformers") + import torch + + from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainingModel + from verl_speco.models.dflash2 import DFlash2DraftModel + + torch.manual_seed(0) + config = _tiny_dflash2_config() + model = DFlash2TrainingModel( + draft_model=DFlash2DraftModel(config), + block_size=config.block_size, + num_anchors=config.num_anchors, + selector_loss_weight=1.0, + ) + + n_rows = 12 + active_hidden = torch.randn(n_rows, config.hidden_size, requires_grad=True) + active_logits = torch.randn(n_rows, config.vocab_size, requires_grad=True) + bsz, n_blocks = 1, 3 + safe_label_indices = torch.arange( + 1, 1 + n_blocks * config.block_size + ).reshape(bsz, n_blocks, config.block_size) + active_mask = torch.zeros( + bsz * n_blocks * config.block_size, dtype=torch.bool + ) + active_mask[:n_rows] = True + + loss, _ = model._auxiliary_loss( + input_ids=torch.randint( + 0, config.vocab_size, (bsz, n_blocks * config.block_size + 1) + ), + safe_label_indices=safe_label_indices, + active_mask=active_mask, + active_hidden=active_hidden, + active_logits=active_logits, + active_targets=torch.randint(0, config.vocab_size, (n_rows,)), + active_weights=torch.ones(n_rows), + ) + loss.backward() + + assert active_hidden.grad is None, ( + "selector loss leaked gradient into the backbone hidden states" + ) + assert active_logits.grad is None, ( + "selector loss leaked gradient into the drafter logits via topk" + ) + # The selector's own parameters must still be trained. + assert model.draft_model.candidate_selector.hidden_projection.weight.grad is not None + + def test_dflash2_training_model_rejects_restricted_vocab() -> None: """Selector candidates are real token ids, so a restricted vocab is invalid.""" pytest.importorskip("torch") diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py index f944034c..080fe08e 100644 --- a/verl_speco/backends/dflash2_trainer_backend.py +++ b/verl_speco/backends/dflash2_trainer_backend.py @@ -69,7 +69,18 @@ def _auxiliary_loss( device = active_hidden.device top_k = min(selector.top_k, active_logits.shape[-1]) - unary, candidates = torch.topk(active_logits, top_k, dim=-1, sorted=False) + # Detach the backbone inputs. The selector is an auxiliary head that + # re-ranks whatever the drafter already produced, so its objective must + # not reshape the drafter itself: torch.topk passes gradient through the + # selected logits, and hidden_projection would otherwise push gradient + # back through the backbone as well. Leaving them attached also breaks + # any DFlash/DFlash2 comparison, because the two backbones would then be + # trained under different effective objectives rather than differing + # only by architecture. + unary, candidates = torch.topk( + active_logits.detach(), top_k, dim=-1, sorted=False + ) + selector_hidden = active_hidden.detach() # Teacher-forced predecessor: the token immediately before the slot this # row predicts. Block-relative position 0 is the anchor and never carries @@ -81,7 +92,7 @@ def _auxiliary_loss( predecessor_ids = predecessor_ids.reshape(-1)[active_mask] scores = selector.pair_scores( - active_hidden, unary.float(), candidates, predecessor_ids + selector_hidden, unary.float(), candidates, predecessor_ids ) # The selector can only learn on rows whose ground truth survived the From 6e97e04d958c403ba1337fa348c9c242483138fb Mon Sep 17 00:00:00 2001 From: khazic Date: Mon, 24 Aug 2026 20:34:59 +0800 Subject: [PATCH 10/16] feat(dflash): track the acceptance length of drafted blocks The block drafters report per-position accuracies, but those are marginals: each position is counted independently. A speculative verifier stops at the first mismatch, so what governs the achievable speedup is the length of the correct prefix, which the marginals cannot be combined into. Add _block_acceptance_counts() and call it from the DFlash, Domino and DSpark forwards, so every block drafter emits the metric the base trainer already had prefix-generic plumbing for. The helper takes the drafted positions only and leaves the slicing to each caller, because the layouts differ: DFlash keeps an unscored anchor at column 0, while Domino and DSpark build shifted labels where every column is a real prediction and slicing would both drop a token and stop a wrong first token from truncating the block. DFlash intersects the scoring mask with the reweighted mask, since correctness is only recorded where the reweighted mask is positive; without that, a position zeroed by loss decay or front_position_weight would read as a mismatch and truncate every block. Reported as mean_acceptance_length under the same convention as the rollout-side drafter/spec_decode/mean_acceptance_length, whose leading 1.0 is the token the target emits itself at each verification step, so the two are directly comparable. The raw sum and block count are published alongside it. Signed-off-by: khazic --- .../test_dflash2_backend_contract.py | 133 ++++++++++++++++-- verl_speco/backends/dflash_trainer_backend.py | 49 +++++++ verl_speco/backends/domino_trainer_backend.py | 11 ++ verl_speco/backends/dspark_trainer_backend.py | 11 ++ verl_speco/trainer/base_trainer.py | 14 ++ 5 files changed, 209 insertions(+), 9 deletions(-) diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index 0abd31aa..9a118ff8 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -154,9 +154,7 @@ def test_conv_does_not_leak_across_block_boundaries() -> None: # First row of each block has no in-block predecessor, so it must be zero. torch.testing.assert_close(out[:, 0], torch.zeros_like(out[:, 0])) - torch.testing.assert_close( - out[:, block_size], torch.zeros_like(out[:, block_size]) - ) + torch.testing.assert_close(out[:, block_size], torch.zeros_like(out[:, block_size])) # Interior rows read their own block's previous row. torch.testing.assert_close(out[:, 1], hidden[:, 0]) torch.testing.assert_close(out[:, block_size + 1], hidden[:, block_size]) @@ -289,12 +287,10 @@ def test_selector_loss_does_not_backprop_into_the_backbone() -> None: active_hidden = torch.randn(n_rows, config.hidden_size, requires_grad=True) active_logits = torch.randn(n_rows, config.vocab_size, requires_grad=True) bsz, n_blocks = 1, 3 - safe_label_indices = torch.arange( - 1, 1 + n_blocks * config.block_size - ).reshape(bsz, n_blocks, config.block_size) - active_mask = torch.zeros( - bsz * n_blocks * config.block_size, dtype=torch.bool + safe_label_indices = torch.arange(1, 1 + n_blocks * config.block_size).reshape( + bsz, n_blocks, config.block_size ) + active_mask = torch.zeros(bsz * n_blocks * config.block_size, dtype=torch.bool) active_mask[:n_rows] = True loss, _ = model._auxiliary_loss( @@ -317,7 +313,9 @@ def test_selector_loss_does_not_backprop_into_the_backbone() -> None: "selector loss leaked gradient into the drafter logits via topk" ) # The selector's own parameters must still be trained. - assert model.draft_model.candidate_selector.hidden_projection.weight.grad is not None + assert ( + model.draft_model.candidate_selector.hidden_projection.weight.grad is not None + ) def test_dflash2_training_model_rejects_restricted_vocab() -> None: @@ -460,3 +458,120 @@ def test_dflash2_config_routes_through_auto(tmp_path) -> None: # The nested z-lab block must survive routing through AutoDraftModelConfig. assert loaded.block_size == 8 assert loaded.selector_top_k == 16 + + +def test_acceptance_counts_stop_at_the_first_mismatch() -> None: + """Acceptance is a prefix property, not a sum of per-position marginals.""" + pytest.importorskip("torch") + pytest.importorskip("transformers") + pytest.importorskip("safetensors") + import torch + + from verl_speco.backends.dflash_trainer_backend import _block_acceptance_counts + + # Drafted positions only; the caller has already dropped any anchor column. + correct = torch.tensor( + [ + [ + [1, 1, 1, 1, 1, 1, 1], # every drafted position correct + [1, 1, 0, 1, 1, 1, 1], # miss at index 2 truncates the tail + [0, 1, 1, 1, 1, 1, 1], # miss at index 0 accepts nothing + [0, 0, 0, 0, 0, 0, 0], # nothing correct + ] + ], + dtype=torch.bool, + ) + scored = torch.ones(1, 4, 7, dtype=torch.bool) + + accepted_sum, scored_blocks = _block_acceptance_counts(correct, scored) + # 7 + 2 + 0 + 0, over 4 scored blocks. + assert float(accepted_sum) == 9.0 + assert float(scored_blocks) == 4.0 + # The marginals would have credited the truncated blocks with 6 apiece. + assert correct.sum(dim=-1).tolist() == [[7, 6, 6, 0]] + + +def test_acceptance_counts_treat_unscored_positions_as_terminal() -> None: + """A block that runs past its supervised region cannot keep accepting.""" + pytest.importorskip("torch") + pytest.importorskip("transformers") + pytest.importorskip("safetensors") + import torch + + from verl_speco.backends.dflash_trainer_backend import _block_acceptance_counts + + correct = torch.ones(1, 2, 7, dtype=torch.bool) + scored = torch.tensor( + [ + [ + [1, 1, 1, 0, 0, 0, 0], # supervised through index 2 + [0, 0, 0, 0, 0, 0, 0], # nothing to score + ] + ], + dtype=torch.bool, + ) + + accepted_sum, scored_blocks = _block_acceptance_counts(correct, scored) + assert float(accepted_sum) == 3.0 + # A block with nothing to score must not dilute the mean. + assert float(scored_blocks) == 1.0 + + +def test_acceptance_counts_do_not_assume_an_anchor_column() -> None: + """Shifted-label drafters have no anchor, so position 0 must be able to truncate. + + DFlash keeps an unscored anchor at column 0 and slices it off before calling. + Domino and DSpark build labels with a ``+ 1`` offset, so their column 0 is a + real prediction. If this helper sliced internally it would both discard that + token and let a wrong first token pass unnoticed, over-reporting acceptance + on exactly the case the metric exists to catch. + """ + pytest.importorskip("torch") + pytest.importorskip("transformers") + pytest.importorskip("safetensors") + import torch + + from verl_speco.backends.dflash_trainer_backend import _block_acceptance_counts + + # Wrong at position 0, right everywhere after: a verifier accepts nothing. + correct = torch.tensor([[[0, 1, 1, 1, 1, 1, 1, 1]]], dtype=torch.bool) + scored = torch.ones(1, 1, 8, dtype=torch.bool) + + accepted_sum, scored_blocks = _block_acceptance_counts(correct, scored) + assert float(accepted_sum) == 0.0 + assert float(scored_blocks) == 1.0 + + # And an all-correct block of the same layout yields every position. + accepted_sum, _ = _block_acceptance_counts( + torch.ones(1, 1, 8, dtype=torch.bool), scored + ) + assert float(accepted_sum) == 8.0 + + +def test_mean_acceptance_length_counts_the_target_bonus_token() -> None: + """The reported metric adds the target's own token, matching the rollout metric.""" + pytest.importorskip("torch") + from types import SimpleNamespace + + from verl_speco.trainer.base_trainer import DrafterBaseTrainer + + trainer = DrafterBaseTrainer.__new__(DrafterBaseTrainer) + trainer.backend = SimpleNamespace(model_type="dflash2") + trainer.config = SimpleNamespace( + rollout=SimpleNamespace(drafter=SimpleNamespace(training={})) + ) + trainer._training_metric_steps = 1 + trainer._training_metric_sums = { + "dflash2/accepted_length_sum": 9.0, + "dflash2/scored_block_count": 4.0, + } + trainer.optimizer = None + trainer.optimizer_steps_total = 0 + + metrics = trainer.get_training_metrics() + + # 9 / 4 accepted drafts, plus the token the target emits itself. + assert metrics["dflash2/mean_acceptance_length"] == pytest.approx(3.25) + # The raw sum and count are published too, so the ratio can be debugged. + assert metrics["dflash2/accepted_length_sum"] == pytest.approx(9.0) + assert metrics["dflash2/scored_block_count"] == pytest.approx(4.0) diff --git a/verl_speco/backends/dflash_trainer_backend.py b/verl_speco/backends/dflash_trainer_backend.py index 03ed4196..f8a3f29a 100644 --- a/verl_speco/backends/dflash_trainer_backend.py +++ b/verl_speco/backends/dflash_trainer_backend.py @@ -50,6 +50,38 @@ def forward(self, hidden_states): return self.fc(hidden_states) +def _block_acceptance_counts( + block_correct: torch.Tensor, block_scored: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Accepted draft tokens summed over blocks, and the number of scored blocks. + + A greedy verifier takes a drafted block prefix and rejects everything from + the first mismatch onward, so position ``k`` only counts when ``0..k`` are + all correct. That is a prefix property, and it cannot be recovered from the + per-position marginal accuracies, which count each position independently. + + Both arguments cover the **drafted** positions only, shaped + ``[bsz, n_blocks, n_drafted]``. Callers own that slicing because the block + layouts differ: DFlash keeps an unscored anchor at column 0 and must drop + it, while Domino and DSpark build shifted labels where every column is a + real prediction. Slicing inside here would silently discard their first + drafted token and, worse, stop a wrong first token from truncating the + block, which is the exact failure the metric exists to catch. + + ``block_scored`` must be the same mask the correctness was computed under, + so that an unscored position reads as "no more prefix" rather than as a + mismatch. Unscored positions terminate the prefix, because nothing beyond + the supervised region can be counted as accepted. + + Returned as a sum and a count rather than a mean so the two reduce correctly + across microbatches and across ranks. ``cumsum`` rather than ``cumprod`` + because it is the more broadly supported primitive across backends. + """ + broken = (~(block_correct & block_scored)).cumsum(dim=-1) + accepted = (broken == 0).sum(dim=-1) + return accepted.sum().float(), block_scored.any(dim=-1).sum().float() + + def _create_dflash_mask_mod( anchor_positions: torch.Tensor, block_keep_mask: torch.Tensor, @@ -624,6 +656,21 @@ def forward( ) loss_per_position = loss_sum_per_position / count_per_pos acc_per_position = correct_per_position / count_per_pos + # Prefix acceptance per block; the per-position accuracies above are + # marginals and cannot be combined into it. Column 0 is this layout's + # anchor, so it is dropped. The scoring mask is intersected with the + # reweighted mask because `correct` was only written where + # `active_mask` held: a position zeroed by loss decay or by + # `front_position_weight` is unscored, not a mismatch, and reading it + # as a mismatch would truncate every block's prefix. + block_drafted = slice(1, None) + block_scored = (binary_weights > 0) & ( + flat_weights.view(bsz, n_blocks, self.block_size) > 0 + ) + accepted_length_sum, scored_block_count = _block_acceptance_counts( + correct.view(bsz, n_blocks, self.block_size)[:, :, block_drafted], + block_scored[:, :, block_drafted], + ) masked_rows = (binary_eval_mask <= 0.5).sum().to(dtype=torch.float32) diagnostics = { "correct_count": correct.sum().float(), @@ -638,6 +685,8 @@ def forward( "loss_sum_per_position": loss_sum_per_position, "correct_per_position": correct_per_position, "count_per_position": count_per_position, + "accepted_length_sum": accepted_length_sum, + "scored_block_count": scored_block_count, "sampled_vocab_size": torch.tensor( float(restricted_vocab.numel()) if self.loss_mode in {"restricted_ce", "sampled_ce"} diff --git a/verl_speco/backends/domino_trainer_backend.py b/verl_speco/backends/domino_trainer_backend.py index 87367e57..5b3f893d 100644 --- a/verl_speco/backends/domino_trainer_backend.py +++ b/verl_speco/backends/domino_trainer_backend.py @@ -50,6 +50,7 @@ from verl_speco.backends.dflash_trainer_backend import ( DFlashTrainerBackend, DFlashTrainingModel, + _block_acceptance_counts, _create_dflash_dense_attention_mask, _create_dflash_mask_mod, ) @@ -374,6 +375,14 @@ def forward(self, input_ids, hidden_states_list, loss_mask, lm_head_weight): min=1.0 ) acc_per_position = correct_per_position / count_per_position.clamp(min=1.0) + # Prefix acceptance per block; the per-position accuracies above are + # marginals and cannot be combined into it. Labels here are shifted + # by one, so unlike DFlash there is no anchor column: every position + # is a real prediction and none may be sliced off. + accepted_length_sum, scored_block_count = _block_acceptance_counts( + correct.view(bsz, n_blocks, self.block_size), + binary_weights > 0, + ) valid_token_count = active_weights.sum().to(dtype=torch.float32) weighted_token_count = flat_weights.sum().to(dtype=torch.float32) accuracy = correct.float().sum() / binary_eval_mask.float().sum().clamp( @@ -401,6 +410,8 @@ def forward(self, input_ids, hidden_states_list, loss_mask, lm_head_weight): "loss_sum_per_position": loss_sum_per_position.detach(), "correct_per_position": correct_per_position.detach(), "count_per_position": count_per_position.detach(), + "accepted_length_sum": accepted_length_sum.detach(), + "scored_block_count": scored_block_count.detach(), "local_ploss_sum": (loss_per_token * binary_eval_mask.float()) .sum() .detach(), diff --git a/verl_speco/backends/dspark_trainer_backend.py b/verl_speco/backends/dspark_trainer_backend.py index 61d58267..6b247da3 100644 --- a/verl_speco/backends/dspark_trainer_backend.py +++ b/verl_speco/backends/dspark_trainer_backend.py @@ -24,6 +24,7 @@ from verl_speco.backends.dflash_trainer_backend import ( DFlashTrainerBackend, DFlashTrainingModel, + _block_acceptance_counts, _create_dflash_dense_attention_mask, _create_dflash_mask_mod, ) @@ -592,6 +593,14 @@ def forward( min=1.0 ) acc_per_position = correct_per_position / count_per_position.clamp(min=1.0) + # Prefix acceptance per block; the per-position accuracies above are + # marginals and cannot be combined into it. Labels here are shifted + # by one, so unlike DFlash there is no anchor column: every position + # is a real prediction and none may be sliced off. + accepted_length_sum, scored_block_count = _block_acceptance_counts( + correct.view(bsz, n_blocks, self.block_size), + binary_weights > 0, + ) valid_token_count = active_weights.sum().to(dtype=torch.float32) weighted_token_count = flat_weights.sum().to(dtype=torch.float32) accuracy = correct.float().sum() / binary_eval_mask.float().sum().clamp( @@ -634,6 +643,8 @@ def forward( "loss_sum_per_position": loss_sum_per_position.detach(), "correct_per_position": correct_per_position.detach(), "count_per_position": count_per_position.detach(), + "accepted_length_sum": accepted_length_sum.detach(), + "scored_block_count": scored_block_count.detach(), "local_ploss_sum": local_ploss_sum.detach(), } return ( diff --git a/verl_speco/trainer/base_trainer.py b/verl_speco/trainer/base_trainer.py index 469e6984..a31b2178 100644 --- a/verl_speco/trainer/base_trainer.py +++ b/verl_speco/trainer/base_trainer.py @@ -853,6 +853,16 @@ def get_training_metrics(self) -> dict[str, float]: eval_tokens = sums.get(f"{prefix}/eval_token_count", 0.0) if eval_tokens > 0: metrics[f"{prefix}/accuracy"] = correct / eval_tokens + scored_blocks = sums.get(f"{prefix}/scored_block_count", 0.0) + if scored_blocks > 0: + # Same convention as the rollout-side drafter/spec_decode/ + # mean_acceptance_length: the leading 1.0 is the token the target + # emits itself at each verification step, not a drafted one, so the + # training metric predicts the served one instead of sitting a token + # below it on the same dashboard. + metrics[f"{prefix}/mean_acceptance_length"] = 1.0 + ( + sums.get(f"{prefix}/accepted_length_sum", 0.0) / scored_blocks + ) quality_tokens = sums.get(f"{prefix}/quality_token_count", 0.0) if quality_tokens > 0: metrics[f"{prefix}/top1_acc"] = ( @@ -877,6 +887,8 @@ def get_training_metrics(self) -> dict[str, float]: f"{prefix}/ce_weighted_token_count", f"{prefix}/l1_weighted_token_count", f"{prefix}/quality_token_count", + f"{prefix}/accepted_length_sum", + f"{prefix}/scored_block_count", f"{prefix}/sanitized_rows", f"{prefix}/masked_rows", f"{prefix}/sampled_vocab_size", @@ -939,6 +951,8 @@ def _record_dflash_training_metrics(self, loss_dict: dict[str, Any]) -> None: scalar_keys = { "correct_count": f"{prefix}/correct_count", "eval_token_count": f"{prefix}/eval_token_count", + "accepted_length_sum": f"{prefix}/accepted_length_sum", + "scored_block_count": f"{prefix}/scored_block_count", "top1_correct_count": f"{prefix}/top1_correct_count", "top5_correct_count": f"{prefix}/top5_correct_count", "quality_token_count": f"{prefix}/quality_token_count", From dcf452e3785d95a56a9a304bc661d6042faa1d57 Mon Sep 17 00:00:00 2001 From: khazic Date: Thu, 27 Aug 2026 21:04:57 +0800 Subject: [PATCH 11/16] fix(dflash): resolve the RoPE base from rope_parameters as well transformers 5 moved the RoPE base out of a top-level rope_theta and into a rope_parameters dict. The released DFlash-family drafter checkpoints and modern target configs carry only the nested spelling, so DFlashConfig fell back to its 10000.0 default: three orders of magnitude below the 1e7 those models were trained at, silently, on both the checkpoint path and the cold-start path that reads the base off the target config. Resolve both spellings in one helper and use it in DFlashConfig plus every DFlash-family fallback config, so a DFlash baseline and a DFlash2 arm cannot end up on different RoPE bases. Signed-off-by: khazic --- verl_speco/backends/dflash_trainer_backend.py | 3 +- verl_speco/backends/domino_trainer_backend.py | 3 +- verl_speco/backends/dspark_trainer_backend.py | 3 +- verl_speco/models/dflash/__init__.py | 4 +- .../models/dflash/configuration_dflash.py | 51 +++++++++++++++++-- 5 files changed, 57 insertions(+), 7 deletions(-) diff --git a/verl_speco/backends/dflash_trainer_backend.py b/verl_speco/backends/dflash_trainer_backend.py index f8a3f29a..3a4b6262 100644 --- a/verl_speco/backends/dflash_trainer_backend.py +++ b/verl_speco/backends/dflash_trainer_backend.py @@ -29,6 +29,7 @@ DFlashConfig, DFlashDraftModel, build_target_layer_ids, + resolve_rope_theta, ) from verl_speco.models.dflash.flex_attention import compile_friendly_create_block_mask from verl_speco.models.target.target_head import TargetHead @@ -802,7 +803,7 @@ def _build_fallback_config(self, target_hf_config): max_position_embeddings=int( getattr(target_text_config, "max_position_embeddings", 32768) ), - rope_theta=float(getattr(target_text_config, "rope_theta", 10000.0)), + rope_theta=resolve_rope_theta(target_text_config), num_target_layers=target_num_hidden_layers, num_context_layers=num_context_layers, target_hidden_size=int(target_text_config.hidden_size), diff --git a/verl_speco/backends/domino_trainer_backend.py b/verl_speco/backends/domino_trainer_backend.py index 5b3f893d..55c2e724 100644 --- a/verl_speco/backends/domino_trainer_backend.py +++ b/verl_speco/backends/domino_trainer_backend.py @@ -54,6 +54,7 @@ _create_dflash_dense_attention_mask, _create_dflash_mask_mod, ) +from verl_speco.models.dflash import resolve_rope_theta from verl_speco.models.dflash.flex_attention import compile_friendly_create_block_mask from verl_speco.models.domino import DominoConfig, DominoDraftModel from verl_speco.trainer.checkpoint import log_drafter_checkpoint_step @@ -525,7 +526,7 @@ def _build_fallback_config(self, target_hf_config): max_position_embeddings=int( getattr(target_text_config, "max_position_embeddings", 32768) ), - rope_theta=float(getattr(target_text_config, "rope_theta", 10000.0)), + rope_theta=resolve_rope_theta(target_text_config), num_target_layers=target_num_hidden_layers, num_context_layers=num_context_layers, target_hidden_size=int(target_text_config.hidden_size), diff --git a/verl_speco/backends/dspark_trainer_backend.py b/verl_speco/backends/dspark_trainer_backend.py index 6b247da3..f8e42e02 100644 --- a/verl_speco/backends/dspark_trainer_backend.py +++ b/verl_speco/backends/dspark_trainer_backend.py @@ -28,6 +28,7 @@ _create_dflash_dense_attention_mask, _create_dflash_mask_mod, ) +from verl_speco.models.dflash import resolve_rope_theta from verl_speco.models.dflash.flex_attention import compile_friendly_create_block_mask from verl_speco.models.dspark import DSparkConfig, DSparkDraftModel from verl_speco.trainer.checkpoint import log_drafter_checkpoint_step @@ -775,7 +776,7 @@ def _build_fallback_config(self, target_hf_config): max_position_embeddings=int( getattr(target_text_config, "max_position_embeddings", 32768) ), - rope_theta=float(getattr(target_text_config, "rope_theta", 10000.0)), + rope_theta=resolve_rope_theta(target_text_config), num_target_layers=target_num_hidden_layers, num_context_layers=num_context_layers, target_hidden_size=int(target_text_config.hidden_size), diff --git a/verl_speco/models/dflash/__init__.py b/verl_speco/models/dflash/__init__.py index 07344d68..0a442466 100644 --- a/verl_speco/models/dflash/__init__.py +++ b/verl_speco/models/dflash/__init__.py @@ -11,7 +11,7 @@ # 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 .configuration_dflash import DFlashConfig +from .configuration_dflash import DEFAULT_ROPE_THETA, DFlashConfig, resolve_rope_theta from .modeling_dflash import ( DFlashAttention, DFlashDecoderLayer, @@ -23,6 +23,7 @@ ) __all__ = [ + "DEFAULT_ROPE_THETA", "DFlashConfig", "DFlashDraftModel", "DFlashAttention", @@ -31,4 +32,5 @@ "DFlashRMSNorm", "DFlashRotaryEmbedding", "build_target_layer_ids", + "resolve_rope_theta", ] diff --git a/verl_speco/models/dflash/configuration_dflash.py b/verl_speco/models/dflash/configuration_dflash.py index f01e7198..ca2fa6c2 100644 --- a/verl_speco/models/dflash/configuration_dflash.py +++ b/verl_speco/models/dflash/configuration_dflash.py @@ -13,10 +13,50 @@ # limitations under the License. import json import os -from typing import Optional +from collections.abc import Mapping +from typing import Any, Optional from transformers import PretrainedConfig +DEFAULT_ROPE_THETA = 10000.0 + + +def _lookup(source: Any, key: str): + """Read ``key`` from either a mapping or a config object.""" + if source is None: + return None + if isinstance(source, Mapping): + return source.get(key) + return getattr(source, key, None) + + +def resolve_rope_theta(source: Any, default: float = DEFAULT_ROPE_THETA) -> float: + """Resolve the RoPE base from a config that may nest it under ``rope_parameters``. + + transformers 5 moved the RoPE base out of a top-level ``rope_theta`` and into + a ``rope_parameters`` dict. Released DFlash-family drafter checkpoints and + modern target configs carry only the nested spelling, so reading + ``rope_theta`` alone silently falls back to ``default``, which is three orders + of magnitude off for a model trained at 1e7. A top-level value still wins, so + a config this overlay wrote itself round-trips unchanged. + ``verl_speco.integration.sglang_patch`` bridges the same split on the serving + side. + + Args: + source: A config object or a raw config mapping. + default: RoPE base to use when neither spelling carries one. + + Returns: + float: The resolved RoPE base. + """ + top_level = _lookup(source, "rope_theta") + if top_level is not None: + return float(top_level) + nested = _lookup(_lookup(source, "rope_parameters"), "rope_theta") + if nested is not None: + return float(nested) + return float(default) + class DFlashConfig(PretrainedConfig): """Configuration for the DFlash draft model. @@ -39,7 +79,7 @@ def __init__( vocab_size: int = 152064, rms_norm_eps: float = 1e-6, max_position_embeddings: int = 32768, - rope_theta: float = 10000.0, + rope_theta: Optional[float] = None, num_target_layers: int = 36, num_context_layers: Optional[int] = 5, target_hidden_size: int = 4096, @@ -58,7 +98,12 @@ def __init__( self.vocab_size = vocab_size self.rms_norm_eps = rms_norm_eps self.max_position_embeddings = max_position_embeddings - self.rope_theta = rope_theta + # ``modeling_dflash`` reads ``rope_theta`` directly, so keep it a plain + # attribute, but accept the transformers-5 ``rope_parameters`` spelling + # the released checkpoints use. + self.rope_theta = resolve_rope_theta( + {"rope_theta": rope_theta, "rope_parameters": kwargs.get("rope_parameters")} + ) self.num_target_layers = num_target_layers self.num_context_layers = num_context_layers self.target_hidden_size = target_hidden_size From 3492314680b247f073c0f0d8abaee84a9a8fa5cf Mon Sep 17 00:00:00 2001 From: khazic Date: Thu, 27 Aug 2026 21:05:14 +0800 Subject: [PATCH 12/16] fix(dflash2): load the upstream selector codebooks instead of dropping them The released z-lab DFlash2 checkpoint stores the two selector codebooks as bare 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, so nothing failed: the base loader drops keys the model does not have with only a debug log, and its required-key gate covers the DFlash backbone alone, which left both codebooks at their random init while training carried on. Rename them on load, and fail loud when a checkpoint carries some DFlash2 modules under unrecognized names, so a future upstream rename cannot degrade into a silent cold start. A checkpoint carrying none of them is still accepted: warm-starting DFlash2 from a plain DFlash backbone stays a legitimate flow. Signed-off-by: khazic --- .../backends/dflash2_trainer_backend.py | 86 ++++++++++++++++++- 1 file changed, 85 insertions(+), 1 deletion(-) diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py index 080fe08e..803bad7b 100644 --- a/verl_speco/backends/dflash2_trainer_backend.py +++ b/verl_speco/backends/dflash2_trainer_backend.py @@ -24,11 +24,33 @@ DFlashTrainerBackend, DFlashTrainingModel, ) +from verl_speco.models.dflash import resolve_rope_theta from verl_speco.models.dflash2 import DFlash2Config, DFlash2DraftModel from verl_speco.trainer.checkpoint import log_drafter_checkpoint_step logger = logging.getLogger(__name__) +# Substrings that mark a parameter as belonging to one of the two modules DFlash2 +# adds on top of the DFlash backbone. +_DFLASH2_MODULE_KEY_MARKERS = ("attention_conv.", "mlp_conv.", "candidate_selector.") + +# Upstream z-lab DFlash2 checkpoints store the two selector codebooks as bare +# ``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 already matches upstream exactly. +_DFLASH2_CHECKPOINT_KEY_ALIASES = { + "candidate_selector.predecessor_codebook": ( + "candidate_selector.predecessor_codebook.weight" + ), + "candidate_selector.successor_codebook": ( + "candidate_selector.successor_codebook.weight" + ), +} + + +def _is_dflash2_module_key(key: str) -> bool: + return any(marker in key for marker in _DFLASH2_MODULE_KEY_MARKERS) + class DFlash2TrainingModel(DFlashTrainingModel): """DFlash training wrapper plus the DFlash2 candidate-selector objective. @@ -158,6 +180,68 @@ def _training_value(self, training_cfg, dflash2_key: str, dflash_key: str, defau return value return training_cfg.get(dflash_key, default) + def _normalize_draft_state_dict(self, state_dict): + """Rename the upstream selector codebooks onto their ``nn.Embedding`` keys. + + The base normalizer only strips wrapper prefixes, so without this the two + codebooks arrive under names the model does not have and are dropped by + ``_load_draft_checkpoint`` with nothing but a debug log. + """ + normalized = super()._normalize_draft_state_dict(state_dict) + for source, target in _DFLASH2_CHECKPOINT_KEY_ALIASES.items(): + if source not in normalized: + continue + value = normalized.pop(source) + normalized.setdefault(target, value) + return normalized + + def _assert_dflash2_modules_are_complete( + self, draft_model, normalized_state, model_path: str + ) -> None: + """Fail loud when a checkpoint carries DFlash2 modules under other names. + + ``_load_draft_checkpoint`` drops keys the model does not have with only a + debug log, and its required-key gate covers the DFlash backbone only. An + upstream rename would therefore leave the convolutions or the selector + silently at their cold-start values, which reads as a merely weak drafter + rather than as a failed load. A checkpoint carrying none of these keys is + still accepted: warm-starting DFlash2 from a plain DFlash backbone is a + legitimate flow, and the DFlash2 modules cold-start as an identity + passthrough by design. + """ + expected = { + key for key in draft_model.state_dict() if _is_dflash2_module_key(key) + } + present = expected.intersection(normalized_state) + stray = { + key for key in normalized_state if _is_dflash2_module_key(key) + }.difference(expected) + if not present and not stray: + return + if present == expected and not stray: + return + raise ValueError( + "DFlash2 checkpoint carries only part of the DFlash2 modules under the " + "expected parameter names, so the rest would silently stay at their " + f"cold-start values: missing={sorted(expected - present)} " + f"unrecognized={sorted(stray)} model_path={model_path}. Add the " + "renamed keys to _DFLASH2_CHECKPOINT_KEY_ALIASES." + ) + + def _load_draft_checkpoint( + self, draft_model, model_path: str, normalized_state=None + ) -> None: + if normalized_state is None: + normalized_state = self._normalize_draft_state_dict( + self._load_draft_state_dict(model_path) + ) + self._assert_dflash2_modules_are_complete( + draft_model, normalized_state, model_path + ) + super()._load_draft_checkpoint( + draft_model, model_path, normalized_state=normalized_state + ) + def _resolved_block_size(self, drafter_config) -> int: """Single source of truth for the block size. @@ -236,7 +320,7 @@ def _build_fallback_config(self, target_hf_config): max_position_embeddings=int( getattr(target_text_config, "max_position_embeddings", 32768) ), - rope_theta=float(getattr(target_text_config, "rope_theta", 10000.0)), + rope_theta=resolve_rope_theta(target_text_config), num_target_layers=target_num_hidden_layers, num_context_layers=num_context_layers, target_hidden_size=int(target_text_config.hidden_size), From 2ee2c997f7eb5e1c13fa89da8255d5e889c5f82f Mon Sep 17 00:00:00 2001 From: khazic Date: Thu, 27 Aug 2026 21:05:14 +0800 Subject: [PATCH 13/16] fix(dflash2): give vLLM and SGLang the same DFLASH2 guard DFLASH2 is not an engine-level algorithm, so SGLang failed loud with guidance while vLLM fell through to a bare 'Unsupported speculative_algorithm' string. Worse, the guidance both should give (serve a trained DFlash2 checkpoint as DFLASH) was unfollowable on vLLM, whose DFlash drafter validator accepted only architectures=['DFlashDraftModel'] and so rejected DFlash2DraftModel. Add the DFLASH2 branch to the vLLM method resolver with the same message, accept the DFlash2 architecture on the DFlash serve path, and classify it in the old-logprob aux-layer lookup so a serve-only run resolves the DFlash layout. Drop DFLASH2 from the SGLang aux-hidden set: that path needs the same drafter.enable the ServerArgs override rejects DFLASH2 under, so it was unreachable. Signed-off-by: khazic --- .../test_dflash2_backend_contract.py | 300 ++++++++++++++++++ .../integration/oldlogprob_layer_ids.py | 2 + verl_speco/integration/sglang_runtime.py | 13 +- verl_speco/integration/vllm_runtime.py | 26 +- 4 files changed, 335 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index 9a118ff8..14838800 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -575,3 +575,303 @@ def test_mean_acceptance_length_counts_the_target_bonus_token() -> None: # The raw sum and count are published too, so the ratio can be debugged. assert metrics["dflash2/accepted_length_sum"] == pytest.approx(9.0) assert metrics["dflash2/scored_block_count"] == pytest.approx(4.0) + + +def _dflash2_backend(**training_overrides): + from omegaconf import OmegaConf + + from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainerBackend + + training = {"dflash2_block_size": 4} + training.update(training_overrides) + return DFlash2TrainerBackend( + OmegaConf.create( + { + "rollout": { + "drafter": { + "speculative_algorithm": "DFLASH2", + "model_path": "", + "training": training, + } + }, + "model": {"path": ""}, + } + ), + None, + ) + + +def test_config_reads_rope_theta_from_rope_parameters(tmp_path) -> None: + """The released checkpoint carries the RoPE base only under rope_parameters. + + transformers 5 moved it there, so reading ``rope_theta`` alone falls back to + the 10000.0 default: three orders of magnitude below the 1e7 the checkpoint + was trained at, and silent, because nothing downstream can tell a defaulted + base from a real one. + """ + pytest.importorskip("transformers") + import json + + from verl_speco.models.dflash2 import DFlash2Config + + (tmp_path / "config.json").write_text( + json.dumps( + { + "architectures": ["DFlash2DraftModel"], + "hidden_size": 8, + "intermediate_size": 16, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "num_hidden_layers": 1, + "vocab_size": 32, + "rope_parameters": {"rope_theta": 10000000, "rope_type": "default"}, + "dflash_config": {"block_size": 8}, + } + ), + encoding="utf-8", + ) + + config = DFlash2Config.from_dflash2_pretrained(str(tmp_path)) + assert config.rope_theta == pytest.approx(1e7) + + +def test_top_level_rope_theta_wins_over_rope_parameters() -> None: + """A config this overlay wrote itself must round-trip unchanged.""" + pytest.importorskip("transformers") + from verl_speco.models.dflash import DFlashConfig + + config = DFlashConfig( + rope_theta=500000.0, + rope_parameters={"rope_theta": 10000000, "rope_type": "default"}, + ) + assert config.rope_theta == pytest.approx(500000.0) + + +def test_rope_theta_defaults_when_neither_spelling_is_present() -> None: + pytest.importorskip("transformers") + from verl_speco.models.dflash import DEFAULT_ROPE_THETA, DFlashConfig + + assert DFlashConfig().rope_theta == pytest.approx(DEFAULT_ROPE_THETA) + + +def test_resolve_rope_theta_reads_config_objects_and_mappings() -> None: + from types import SimpleNamespace + + from verl_speco.models.dflash import resolve_rope_theta + + assert resolve_rope_theta({"rope_parameters": {"rope_theta": 7.0}}) == 7.0 + assert ( + resolve_rope_theta( + SimpleNamespace(rope_parameters=SimpleNamespace(rope_theta=7.0)) + ) + == 7.0 + ) + assert resolve_rope_theta(None, default=3.0) == 3.0 + # An empty/unset rope_parameters must not shadow the default. + assert resolve_rope_theta({"rope_parameters": None}, default=3.0) == 3.0 + + +def test_fallback_config_reads_the_target_rope_parameters() -> None: + """Cold start hits the same split: modern target configs nest the base too. + + Without this the draft's RoPE base silently disagrees with the target it is + drafting for, and every DFlash-family arm of an A/B would sit at 10000.0 + regardless of the target. + """ + pytest.importorskip("torch") + pytest.importorskip("transformers") + from types import SimpleNamespace + + target_config = SimpleNamespace( + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=2, + num_hidden_layers=4, + vocab_size=32, + rms_norm_eps=1e-6, + max_position_embeddings=64, + rope_parameters={"rope_theta": 10000000, "rope_type": "default"}, + ) + + config = _dflash2_backend()._build_fallback_config(target_config) + assert config.rope_theta == pytest.approx(1e7) + + +def _upstream_spelling(state_dict): + """Rewrite a model state dict into the upstream z-lab key spelling.""" + renamed = {} + for key, value in state_dict.items(): + if key.endswith("_codebook.weight"): + key = key[: -len(".weight")] + renamed[key] = value + return renamed + + +def test_upstream_selector_codebooks_load_onto_the_embedding_keys() -> None: + """Upstream stores the codebooks as bare Parameters, this overlay as Embeddings. + + ``_load_draft_checkpoint`` drops keys the model does not have with only a + debug log, and its required-key gate covers the DFlash backbone only, so + without the rename the two codebooks stay randomly initialized and nothing + says so. + """ + pytest.importorskip("torch") + pytest.importorskip("transformers") + import torch + + from verl_speco.models.dflash2 import DFlash2DraftModel + + reference = DFlash2DraftModel(_tiny_dflash2_config()) + with torch.no_grad(): + reference.candidate_selector.predecessor_codebook.weight.normal_() + reference.candidate_selector.successor_codebook.weight.normal_() + for layer in reference.layers: + layer.attention_conv.base_kernel.normal_() + + backend = _dflash2_backend() + normalized = backend._normalize_draft_state_dict( + _upstream_spelling(reference.state_dict()) + ) + assert "candidate_selector.predecessor_codebook.weight" in normalized + assert "candidate_selector.predecessor_codebook" not in normalized + + loaded = DFlash2DraftModel(_tiny_dflash2_config()) + backend._load_draft_checkpoint(loaded, "", normalized_state=normalized) + + torch.testing.assert_close( + loaded.candidate_selector.predecessor_codebook.weight, + reference.candidate_selector.predecessor_codebook.weight, + ) + torch.testing.assert_close( + loaded.candidate_selector.successor_codebook.weight, + reference.candidate_selector.successor_codebook.weight, + ) + torch.testing.assert_close( + loaded.layers[0].attention_conv.base_kernel, + reference.layers[0].attention_conv.base_kernel, + ) + + +def test_partially_renamed_dflash2_modules_fail_loud() -> None: + """A future upstream rename must not degrade into a silent cold start.""" + pytest.importorskip("torch") + pytest.importorskip("transformers") + from verl_speco.models.dflash2 import DFlash2DraftModel + + model = DFlash2DraftModel(_tiny_dflash2_config()) + state = dict(model.state_dict()) + state["candidate_selector.prev_codebook"] = state.pop( + "candidate_selector.predecessor_codebook.weight" + ) + + backend = _dflash2_backend() + with pytest.raises(ValueError, match="only part of the DFlash2 modules"): + backend._load_draft_checkpoint(model, "", normalized_state=state) + + +def test_plain_dflash_checkpoint_warm_starts_without_the_dflash2_modules() -> None: + """Warm-starting DFlash2 from a DFlash backbone stays a legitimate flow. + + The DFlash2 modules cold-start as an identity passthrough by design, so a + checkpoint carrying none of them is a deliberate configuration, not a + mismatch. + """ + pytest.importorskip("torch") + pytest.importorskip("transformers") + from verl_speco.models.dflash2 import DFlash2DraftModel + + model = DFlash2DraftModel(_tiny_dflash2_config()) + backbone_only = { + key: value + for key, value in model.state_dict().items() + if not any( + marker in key + for marker in ("attention_conv.", "mlp_conv.", "candidate_selector.") + ) + } + + backend = _dflash2_backend() + backend._load_draft_checkpoint(model, "", normalized_state=backbone_only) + + +def test_dflash2_rejected_by_vllm_config_builder() -> None: + 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"}) + + +def test_dflash2_rejected_by_sglang_config_builder() -> None: + from verl_speco.integration.sglang_runtime import ( + _server_args_overrides_from_drafter, + ) + + with pytest.raises(ValueError, match="not an engine-level speculative algorithm"): + _server_args_overrides_from_drafter( + {"enable": True, "speculative_algorithm": "DFLASH2"}, + supported_fields={"speculative_algorithm"}, + ) + + +def test_dflash2_never_reaches_the_sglang_aux_hidden_path() -> None: + """The engine rejects DFLASH2 under the same ``enable`` this helper requires.""" + from verl_speco.integration.sglang_runtime import _drafter_uses_dflash_aux_hidden + + assert not _drafter_uses_dflash_aux_hidden( + { + "enable": True, + "enable_drafter_training": True, + "speculative_algorithm": "DFLASH2", + "training": {"collect_hidden_states_from_sgl": True}, + } + ) + + +def test_vllm_dflash_path_accepts_a_dflash2_checkpoint(tmp_path) -> None: + """The DFLASH2 error tells users to serve the checkpoint as DFLASH. + + That advice is only followable if the drafter validator accepts the DFlash2 + architecture on the DFlash path. + """ + import json + + from verl_speco.integration.vllm_runtime import ( + _validate_vllm_dflash_drafter_config, + ) + + (tmp_path / "config.json").write_text( + json.dumps({"architectures": ["DFlash2DraftModel"]}), encoding="utf-8" + ) + _validate_vllm_dflash_drafter_config(str(tmp_path), algorithm="DFLASH") + + +def test_vllm_dflash_path_still_rejects_an_eagle_checkpoint(tmp_path) -> None: + import json + + from verl_speco.integration.vllm_runtime import ( + _validate_vllm_dflash_drafter_config, + ) + + (tmp_path / "config.json").write_text( + json.dumps({"architectures": ["LlamaForCausalLMEagle3"]}), encoding="utf-8" + ) + with pytest.raises(ValueError, match="DFlash-family drafter checkpoint"): + _validate_vllm_dflash_drafter_config(str(tmp_path), algorithm="DFLASH") + + +def test_dflash2_architecture_classifies_as_a_dflash_config() -> None: + """A serve-only run sets algorithm=DFLASH, so the architecture must classify.""" + from types import SimpleNamespace + + from verl_speco.integration.oldlogprob_layer_ids import ( + resolve_oldlogprob_aux_layer_ids, + ) + + layer_ids = resolve_oldlogprob_aux_layer_ids( + {"training": {"dflash_num_target_layers": 5}}, + target_num_hidden_layers=36, + model_configs=[SimpleNamespace(architectures=["DFlash2DraftModel"])], + ) + assert layer_ids is not None + assert len(layer_ids) == 5 diff --git a/verl_speco/integration/oldlogprob_layer_ids.py b/verl_speco/integration/oldlogprob_layer_ids.py index b0405f1b..57bb5c4f 100644 --- a/verl_speco/integration/oldlogprob_layer_ids.py +++ b/verl_speco/integration/oldlogprob_layer_ids.py @@ -102,6 +102,8 @@ def _is_dflash_config(drafter_cfg: Any, model_configs: tuple[Any, ...]) -> bool: architecture in { "DFlashDraftModel", + "DFlash2DraftModel", + "Qwen3DFlash2Model", "DSparkDraftModel", "Qwen3DSparkModel", "DominoDraftModel", diff --git a/verl_speco/integration/sglang_runtime.py b/verl_speco/integration/sglang_runtime.py index 9afdfd9c..11c409ef 100644 --- a/verl_speco/integration/sglang_runtime.py +++ b/verl_speco/integration/sglang_runtime.py @@ -174,7 +174,11 @@ def _drafter_uses_dflash_aux_hidden(drafter_cfg: dict[str, Any]) -> bool: algorithm = str(drafter_cfg.get("speculative_algorithm", "") or "").upper() training_cfg = drafter_cfg.get("training") or {} return bool( - algorithm in {"DFLASH", "DFLASH2", "DSPARK"} + # DFLASH2, like DOMINO, never reaches this path: it is not an + # engine-level algorithm, so the ServerArgs override below fails loud + # for it under the same ``enable`` this check requires. A trained + # DFlash2 checkpoint is served as DFLASH. + algorithm in {"DFLASH", "DSPARK"} and drafter_cfg.get("enable") and drafter_cfg.get("enable_drafter_training") and training_cfg.get("collect_hidden_states_from_sgl") @@ -611,9 +615,10 @@ def _server_args_overrides_from_drafter( # the raw string. raise ValueError( "DFLASH2 is not an engine-level speculative algorithm; DFlash2 is served as a DFlash " - "checkpoint. Set actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH for the " - "rollout/serve path; the trained checkpoint's dflash_config carries the DFlash2 " - "convolution and selector hyperparameters, keeping DFLASH2 for drafter training." + "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." ) rollout_cfg = drafter_cfg.get("rollout") or {} diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py index d179ddbf..176865c3 100644 --- a/verl_speco/integration/vllm_runtime.py +++ b/verl_speco/integration/vllm_runtime.py @@ -812,6 +812,10 @@ def _drafter_algorithm(drafter_cfg: dict[str, Any]) -> str: ) +# Draft architectures vLLM can serve through its DFlash speculative path. +_DFLASH_SERVABLE_ARCHITECTURES = frozenset({"DFlashDraftModel", "DFlash2DraftModel"}) + + def _validate_vllm_dflash_drafter_config( spec_model_path: Any, algorithm: str = "DFLASH" ) -> None: @@ -842,10 +846,15 @@ def _validate_vllm_dflash_drafter_config( ) return - if architectures and "DFlashDraftModel" not in architectures: + # DFlash2 is served as a DFlash checkpoint (the engine reads its convolution + # and selector hyperparameters out of dflash_config), which is what the + # DFLASH2 fail-loud above tells users to do, so its architecture has to be + # accepted here or that advice would be unfollowable. + if architectures and not _DFLASH_SERVABLE_ARCHITECTURES.intersection(architectures): raise ValueError( "vLLM DFlash requires actor_rollout_ref.rollout.drafter.model_path " - "to point to a DFlash drafter checkpoint with architectures=['DFlashDraftModel']; " + "to point to a DFlash-family drafter checkpoint with architectures in " + f"{sorted(_DFLASH_SERVABLE_ARCHITECTURES)}; " f"got architectures={architectures!r} from {config_path}. " "Do not use an EAGLE/EAGLE3 drafter path with speculative_algorithm=DFLASH." ) @@ -909,6 +918,19 @@ def _speculative_method_from_drafter(drafter_cfg: dict[str, Any]) -> str: "enables the Domino correction head on engines that support it, keeping DOMINO for " "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." + ) if algorithm == "DSPARK": return "dflash" if _is_vllm_ascend_runtime_hint() else "dspark" From 458f178978a9f249218de942e3ef4dd29f9c34d2 Mon Sep 17 00:00:00 2001 From: khazic Date: Thu, 27 Aug 2026 21:22:02 +0800 Subject: [PATCH 14/16] refactor(dflash): move the checkpoint key aliases and the variant validation hook onto the base backend /simplify pass over the three fixes. The DFlash2 rename was a second, differently shaped pass wrapped around the base normalizer, and the completeness check was a full _load_draft_checkpoint override whose default-fill was copied verbatim from the base. Both become declarative: a _CHECKPOINT_KEY_ALIASES table applied by the base normalizer, and a _validate_normalized_state hook the base loader calls after its own backbone gate. Also promote the base loader's dropped-key report to a warning when keys were unexpected or shape-mismatched (missing keys stay at debug, the embedding is loaded separately), so the silent drop this fix works around is visible for DFlash / DSpark / Domino too. Signed-off-by: khazic --- .../test_dflash2_backend_contract.py | 125 +++++------------- .../integration/test_vllm_runtime_contract.py | 15 +++ .../backends/dflash2_trainer_backend.py | 84 ++++-------- verl_speco/backends/dflash_trainer_backend.py | 32 ++++- verl_speco/integration/vllm_runtime.py | 7 +- verl_speco/models/dflash/__init__.py | 3 +- verl_speco/models/dflash/modeling_dflash.py | 4 +- 7 files changed, 117 insertions(+), 153 deletions(-) diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index 14838800..58cc2ff0 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -54,6 +54,30 @@ def _tiny_dflash2_config(**overrides): return DFlash2Config(**kwargs) +def _dflash2_backend(**training_overrides): + from omegaconf import OmegaConf + + from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainerBackend + + training = {"dflash2_block_size": 4} + training.update(training_overrides) + return DFlash2TrainerBackend( + OmegaConf.create( + { + "rollout": { + "drafter": { + "speculative_algorithm": "DFLASH2", + "model_path": "", + "training": training, + } + }, + "model": {"path": ""}, + } + ), + None, + ) + + def test_dflash2_model_builds_conv_and_selector() -> None: pytest.importorskip("torch") pytest.importorskip("transformers") @@ -347,28 +371,11 @@ def test_backend_pins_conv_block_size_to_the_trainer_block_size() -> None: """ pytest.importorskip("torch") pytest.importorskip("transformers") - from omegaconf import OmegaConf - - from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainerBackend from verl_speco.models.dflash2 import DFlash2Config - backend = DFlash2TrainerBackend( - OmegaConf.create( - { - "rollout": { - "drafter": { - "speculative_algorithm": "DFLASH2", - "model_path": "", - # Deliberately a multiple of the config's block_size=4, - # so the block-multiple guard alone would not catch it. - "training": {"dflash2_block_size": 8}, - } - }, - "model": {"path": ""}, - } - ), - None, - ) + # Deliberately a multiple of the config's block_size=4, so the + # block-multiple guard alone would not catch it. + backend = _dflash2_backend(dflash2_block_size=8) config = _tiny_dflash2_config(block_size=4) assert isinstance(config, DFlash2Config) assert backend._resolved_block_size(config) == 8 @@ -577,30 +584,6 @@ def test_mean_acceptance_length_counts_the_target_bonus_token() -> None: assert metrics["dflash2/scored_block_count"] == pytest.approx(4.0) -def _dflash2_backend(**training_overrides): - from omegaconf import OmegaConf - - from verl_speco.backends.dflash2_trainer_backend import DFlash2TrainerBackend - - training = {"dflash2_block_size": 4} - training.update(training_overrides) - return DFlash2TrainerBackend( - OmegaConf.create( - { - "rollout": { - "drafter": { - "speculative_algorithm": "DFLASH2", - "model_path": "", - "training": training, - } - }, - "model": {"path": ""}, - } - ), - None, - ) - - def test_config_reads_rope_theta_from_rope_parameters(tmp_path) -> None: """The released checkpoint carries the RoPE base only under rope_parameters. @@ -649,9 +632,9 @@ def test_top_level_rope_theta_wins_over_rope_parameters() -> None: def test_rope_theta_defaults_when_neither_spelling_is_present() -> None: pytest.importorskip("transformers") - from verl_speco.models.dflash import DEFAULT_ROPE_THETA, DFlashConfig + from verl_speco.models.dflash import DFlashConfig - assert DFlashConfig().rope_theta == pytest.approx(DEFAULT_ROPE_THETA) + assert DFlashConfig().rope_theta == pytest.approx(10000.0) def test_resolve_rope_theta_reads_config_objects_and_mappings() -> None: @@ -700,12 +683,10 @@ def test_fallback_config_reads_the_target_rope_parameters() -> None: def _upstream_spelling(state_dict): """Rewrite a model state dict into the upstream z-lab key spelling.""" - renamed = {} - for key, value in state_dict.items(): - if key.endswith("_codebook.weight"): - key = key[: -len(".weight")] - renamed[key] = value - return renamed + return { + key.removesuffix(".weight") if key.endswith("_codebook.weight") else key: value + for key, value in state_dict.items() + } def test_upstream_selector_codebooks_load_onto_the_embedding_keys() -> None: @@ -760,7 +741,7 @@ def test_partially_renamed_dflash2_modules_fail_loud() -> None: from verl_speco.models.dflash2 import DFlash2DraftModel model = DFlash2DraftModel(_tiny_dflash2_config()) - state = dict(model.state_dict()) + state = model.state_dict() state["candidate_selector.prev_codebook"] = state.pop( "candidate_selector.predecessor_codebook.weight" ) @@ -779,16 +760,14 @@ def test_plain_dflash_checkpoint_warm_starts_without_the_dflash2_modules() -> No """ pytest.importorskip("torch") pytest.importorskip("transformers") + from verl_speco.backends.dflash2_trainer_backend import _is_dflash2_module_key from verl_speco.models.dflash2 import DFlash2DraftModel model = DFlash2DraftModel(_tiny_dflash2_config()) backbone_only = { key: value for key, value in model.state_dict().items() - if not any( - marker in key - for marker in ("attention_conv.", "mlp_conv.", "candidate_selector.") - ) + if not _is_dflash2_module_key(key) } backend = _dflash2_backend() @@ -828,38 +807,6 @@ def test_dflash2_never_reaches_the_sglang_aux_hidden_path() -> None: ) -def test_vllm_dflash_path_accepts_a_dflash2_checkpoint(tmp_path) -> None: - """The DFLASH2 error tells users to serve the checkpoint as DFLASH. - - That advice is only followable if the drafter validator accepts the DFlash2 - architecture on the DFlash path. - """ - import json - - from verl_speco.integration.vllm_runtime import ( - _validate_vllm_dflash_drafter_config, - ) - - (tmp_path / "config.json").write_text( - json.dumps({"architectures": ["DFlash2DraftModel"]}), encoding="utf-8" - ) - _validate_vllm_dflash_drafter_config(str(tmp_path), algorithm="DFLASH") - - -def test_vllm_dflash_path_still_rejects_an_eagle_checkpoint(tmp_path) -> None: - import json - - from verl_speco.integration.vllm_runtime import ( - _validate_vllm_dflash_drafter_config, - ) - - (tmp_path / "config.json").write_text( - json.dumps({"architectures": ["LlamaForCausalLMEagle3"]}), encoding="utf-8" - ) - with pytest.raises(ValueError, match="DFlash-family drafter checkpoint"): - _validate_vllm_dflash_drafter_config(str(tmp_path), algorithm="DFLASH") - - def test_dflash2_architecture_classifies_as_a_dflash_config() -> None: """A serve-only run sets algorithm=DFLASH, so the architecture must classify.""" from types import SimpleNamespace diff --git a/tests/integration/test_vllm_runtime_contract.py b/tests/integration/test_vllm_runtime_contract.py index 22c9a85b..2fcec378 100644 --- a/tests/integration/test_vllm_runtime_contract.py +++ b/tests/integration/test_vllm_runtime_contract.py @@ -409,6 +409,21 @@ def test_vllm_dflash_validator_rejects_dspark_when_algorithm_is_dflash( _validate_vllm_dflash_drafter_config(model_path, algorithm="DFLASH") +def test_vllm_dflash_validator_accepts_a_dflash2_checkpoint(tmp_path) -> None: + """DFLASH2 is rejected as an engine method and served as a DFlash checkpoint. + + That advice is only followable if the DFlash path accepts the DFlash2 + architecture. + """ + model_path = tmp_path / "dflash2-drafter" + model_path.mkdir() + (model_path / "config.json").write_text( + '{"architectures": ["DFlash2DraftModel"]}', encoding="utf-8" + ) + + _validate_vllm_dflash_drafter_config(model_path, algorithm="DFLASH") + + def test_vllm_dspark_validator_accepts_markov_head_config(tmp_path) -> None: model_path = tmp_path / "dspark-drafter" model_path.mkdir() diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py index 803bad7b..7434ed24 100644 --- a/verl_speco/backends/dflash2_trainer_backend.py +++ b/verl_speco/backends/dflash2_trainer_backend.py @@ -34,19 +34,6 @@ # adds on top of the DFlash backbone. _DFLASH2_MODULE_KEY_MARKERS = ("attention_conv.", "mlp_conv.", "candidate_selector.") -# Upstream z-lab DFlash2 checkpoints store the two selector codebooks as bare -# ``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 already matches upstream exactly. -_DFLASH2_CHECKPOINT_KEY_ALIASES = { - "candidate_selector.predecessor_codebook": ( - "candidate_selector.predecessor_codebook.weight" - ), - "candidate_selector.successor_codebook": ( - "candidate_selector.successor_codebook.weight" - ), -} - def _is_dflash2_module_key(key: str) -> bool: return any(marker in key for marker in _DFLASH2_MODULE_KEY_MARKERS) @@ -170,6 +157,19 @@ def _auxiliary_loss( class DFlash2TrainerBackend(DFlashTrainerBackend): + # Upstream z-lab DFlash2 checkpoints store the two selector codebooks as bare + # ``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. + _CHECKPOINT_KEY_ALIASES = { + "candidate_selector.predecessor_codebook": ( + "candidate_selector.predecessor_codebook.weight" + ), + "candidate_selector.successor_codebook": ( + "candidate_selector.successor_codebook.weight" + ), + } + @property def model_type(self): return "dflash2" @@ -180,66 +180,36 @@ def _training_value(self, training_cfg, dflash2_key: str, dflash_key: str, defau return value return training_cfg.get(dflash_key, default) - def _normalize_draft_state_dict(self, state_dict): - """Rename the upstream selector codebooks onto their ``nn.Embedding`` keys. - - The base normalizer only strips wrapper prefixes, so without this the two - codebooks arrive under names the model does not have and are dropped by - ``_load_draft_checkpoint`` with nothing but a debug log. - """ - normalized = super()._normalize_draft_state_dict(state_dict) - for source, target in _DFLASH2_CHECKPOINT_KEY_ALIASES.items(): - if source not in normalized: - continue - value = normalized.pop(source) - normalized.setdefault(target, value) - return normalized - - def _assert_dflash2_modules_are_complete( + def _validate_normalized_state( self, draft_model, normalized_state, model_path: str ) -> None: """Fail loud when a checkpoint carries DFlash2 modules under other names. - ``_load_draft_checkpoint`` drops keys the model does not have with only a - debug log, and its required-key gate covers the DFlash backbone only. An - upstream rename would therefore leave the convolutions or the selector - silently at their cold-start values, which reads as a merely weak drafter - rather than as a failed load. A checkpoint carrying none of these keys is - still accepted: warm-starting DFlash2 from a plain DFlash backbone is a - legitimate flow, and the DFlash2 modules cold-start as an identity - passthrough by design. + The base loader drops keys the model does not have, and its required-key + gate covers the DFlash backbone only. An upstream rename would therefore + leave the convolutions or the selector at their cold-start values, which + reads as a merely weak drafter rather than as a failed load. A checkpoint + carrying none of these keys is still accepted: warm-starting DFlash2 from + a plain DFlash backbone is a legitimate flow, and the DFlash2 modules + cold-start as an identity passthrough by design. """ expected = { key for key in draft_model.state_dict() if _is_dflash2_module_key(key) } present = expected.intersection(normalized_state) stray = { - key for key in normalized_state if _is_dflash2_module_key(key) - }.difference(expected) - if not present and not stray: - return - if present == expected and not stray: + key + for key in normalized_state + if _is_dflash2_module_key(key) and key not in expected + } + if not stray and (not present or present == expected): return raise ValueError( "DFlash2 checkpoint carries only part of the DFlash2 modules under the " "expected parameter names, so the rest would silently stay at their " f"cold-start values: missing={sorted(expected - present)} " f"unrecognized={sorted(stray)} model_path={model_path}. Add the " - "renamed keys to _DFLASH2_CHECKPOINT_KEY_ALIASES." - ) - - def _load_draft_checkpoint( - self, draft_model, model_path: str, normalized_state=None - ) -> None: - if normalized_state is None: - normalized_state = self._normalize_draft_state_dict( - self._load_draft_state_dict(model_path) - ) - self._assert_dflash2_modules_are_complete( - draft_model, normalized_state, model_path - ) - super()._load_draft_checkpoint( - draft_model, model_path, normalized_state=normalized_state + "renamed keys to DFlash2TrainerBackend._CHECKPOINT_KEY_ALIASES." ) def _resolved_block_size(self, drafter_config) -> int: diff --git a/verl_speco/backends/dflash_trainer_backend.py b/verl_speco/backends/dflash_trainer_backend.py index 3a4b6262..1a4dd902 100644 --- a/verl_speco/backends/dflash_trainer_backend.py +++ b/verl_speco/backends/dflash_trainer_backend.py @@ -717,6 +717,11 @@ def forward( class DFlashTrainerBackend: + # Checkpoint parameter names that differ from this overlay's own spelling, + # as ``{checkpoint key: model key}``. Variants fill this in when the released + # checkpoint holds a tensor in a different module type than the overlay does. + _CHECKPOINT_KEY_ALIASES: dict[str, str] = {} + def __init__(self, config, target_model_config): self.config = config self.target_model_config = target_model_config @@ -862,8 +867,27 @@ def _normalize_draft_state_dict( normalized_key = normalized_key[len(prefix) :] break normalized_state[normalized_key] = value + # Aliases run after prefix stripping, and a key already spelled the way + # the model spells it wins, so a checkpoint this overlay wrote itself is + # never rewritten. + for source, target in self._CHECKPOINT_KEY_ALIASES.items(): + if source in normalized_state: + normalized_state.setdefault(target, normalized_state.pop(source)) return normalized_state + def _validate_normalized_state( + self, + draft_model: DFlashDraftModel, + normalized_state: dict[str, torch.Tensor], + model_path: str, + ) -> None: + """Variant hook to reject a checkpoint the base gate cannot judge. + + The base gate only knows the DFlash backbone, and unrecognized keys are + dropped rather than raised on, so a variant whose extra modules arrive + under other names must say so here. + """ + def _infer_num_context_layers_from_state( self, normalized_state: dict[str, torch.Tensor], target_hidden_size: int ) -> int | None: @@ -1008,6 +1032,7 @@ def _load_draft_checkpoint( "DFlash/DSpark checkpoint does not use the canonical vLLM parameter names; " f"missing={missing_backbone_keys} model_path={model_path}" ) + self._validate_normalized_state(draft_model, normalized_state, model_path) model_state = draft_model.state_dict() filtered_state: dict[str, torch.Tensor] = {} @@ -1031,7 +1056,12 @@ def _load_draft_checkpoint( missing, _ = draft_model.load_state_dict(filtered_state, strict=False) if unexpected or missing or mismatched: - logger.debug( + # Dropped and shape-mismatched keys mean checkpoint tensors did not + # reach the model, which shows up later only as a weak drafter, so + # say so at warning level. Missing keys alone are routine (the + # embedding is loaded separately), so they stay at debug. + log = logger.warning if (unexpected or mismatched) else logger.debug + log( "DFlash draft checkpoint load report from %s: loaded=%s missing=%s unexpected=%s mismatched=%s", model_path, len(filtered_state), diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py index 176865c3..e47273b6 100644 --- a/verl_speco/integration/vllm_runtime.py +++ b/verl_speco/integration/vllm_runtime.py @@ -813,7 +813,10 @@ def _drafter_algorithm(drafter_cfg: dict[str, Any]) -> str: # Draft architectures vLLM can serve through its DFlash speculative path. -_DFLASH_SERVABLE_ARCHITECTURES = frozenset({"DFlashDraftModel", "DFlash2DraftModel"}) +# Keep in sync with the alias sets in verl_speco/models/auto.py. +_DFLASH_SERVABLE_ARCHITECTURES = frozenset( + {"DFlashDraftModel", "DFlash2DraftModel", "Qwen3DFlash2Model"} +) def _validate_vllm_dflash_drafter_config( @@ -850,7 +853,7 @@ def _validate_vllm_dflash_drafter_config( # and selector hyperparameters out of dflash_config), which is what the # DFLASH2 fail-loud above tells users to do, so its architecture has to be # accepted here or that advice would be unfollowable. - if architectures and not _DFLASH_SERVABLE_ARCHITECTURES.intersection(architectures): + if architectures and _DFLASH_SERVABLE_ARCHITECTURES.isdisjoint(architectures): raise ValueError( "vLLM DFlash requires actor_rollout_ref.rollout.drafter.model_path " "to point to a DFlash-family drafter checkpoint with architectures in " diff --git a/verl_speco/models/dflash/__init__.py b/verl_speco/models/dflash/__init__.py index 0a442466..1a1d8402 100644 --- a/verl_speco/models/dflash/__init__.py +++ b/verl_speco/models/dflash/__init__.py @@ -11,7 +11,7 @@ # 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 .configuration_dflash import DEFAULT_ROPE_THETA, DFlashConfig, resolve_rope_theta +from .configuration_dflash import DFlashConfig, resolve_rope_theta from .modeling_dflash import ( DFlashAttention, DFlashDecoderLayer, @@ -23,7 +23,6 @@ ) __all__ = [ - "DEFAULT_ROPE_THETA", "DFlashConfig", "DFlashDraftModel", "DFlashAttention", diff --git a/verl_speco/models/dflash/modeling_dflash.py b/verl_speco/models/dflash/modeling_dflash.py index d2713f52..38d46db8 100644 --- a/verl_speco/models/dflash/modeling_dflash.py +++ b/verl_speco/models/dflash/modeling_dflash.py @@ -22,7 +22,7 @@ from safetensors import safe_open from transformers import PretrainedConfig, PreTrainedModel -from .configuration_dflash import DFlashConfig +from .configuration_dflash import DFlashConfig, resolve_rope_theta from .flex_attention import compile_friendly_flex_attention @@ -143,7 +143,7 @@ def __init__(self, config: PretrainedConfig): self.rotary_emb = DFlashRotaryEmbedding( self.head_dim, max_position_embeddings=getattr(config, "max_position_embeddings", 32768), - base=getattr(config, "rope_theta", 10000.0), + base=resolve_rope_theta(config), ) def forward( From 27c680e6fb616f231b83f75ac1e0a90407655c3b Mon Sep 17 00:00:00 2001 From: khazic Date: Thu, 27 Aug 2026 21:48:48 +0800 Subject: [PATCH 15/16] fix(dflash): say so when the draft rotary cannot apply the target's RoPE scaling /code-review follow-up. resolve_rope_theta reads the base out of rope_parameters and discards the rest, but DFlashRotaryEmbedding takes only a base, so a target using yarn / linear / llama3 scaling gets a draft whose rotary phase diverges past the original context length. That was already true before this branch, and it is invisible for exactly the reason a defaulted base is. Warn once per distinct rope_type rather than per decoder layer. Applying the scaling itself is a larger change than this PR should carry. Signed-off-by: khazic --- .../test_dflash2_backend_contract.py | 41 +++++++++++++++++++ .../models/dflash/configuration_dflash.py | 40 +++++++++++++++++- 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index 58cc2ff0..e7ac3bd1 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -822,3 +822,44 @@ def test_dflash2_architecture_classifies_as_a_dflash_config() -> None: ) assert layer_ids is not None assert len(layer_ids) == 5 + + +def test_scaled_rope_is_warned_about_not_silently_dropped(caplog) -> None: + """The draft rotary applies the base only, so a scaled target diverges. + + That divergence is invisible for the same reason a defaulted base is, so it + has to be said out loud rather than discarded with the rest of + ``rope_parameters``. + """ + import logging + + from verl_speco.models.dflash import configuration_dflash + + configuration_dflash._WARNED_ROPE_TYPES.discard("yarn") + with caplog.at_level(logging.WARNING, logger=configuration_dflash.__name__): + theta = configuration_dflash.resolve_rope_theta( + {"rope_parameters": {"rope_theta": 5e5, "rope_type": "yarn", "factor": 4.0}} + ) + + assert theta == pytest.approx(5e5) + assert "rope_type='yarn'" in caplog.text + + # Warned once per distinct rope_type, not once per decoder layer. + caplog.clear() + with caplog.at_level(logging.WARNING, logger=configuration_dflash.__name__): + configuration_dflash.resolve_rope_theta( + {"rope_parameters": {"rope_theta": 5e5, "rope_type": "yarn"}} + ) + assert caplog.text == "" + + +def test_default_rope_type_is_not_warned_about(caplog) -> None: + import logging + + from verl_speco.models.dflash import configuration_dflash + + with caplog.at_level(logging.WARNING, logger=configuration_dflash.__name__): + configuration_dflash.resolve_rope_theta( + {"rope_parameters": {"rope_theta": 1e7, "rope_type": "default"}} + ) + assert caplog.text == "" diff --git a/verl_speco/models/dflash/configuration_dflash.py b/verl_speco/models/dflash/configuration_dflash.py index ca2fa6c2..4c2f0a46 100644 --- a/verl_speco/models/dflash/configuration_dflash.py +++ b/verl_speco/models/dflash/configuration_dflash.py @@ -12,14 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. import json +import logging import os from collections.abc import Mapping from typing import Any, Optional from transformers import PretrainedConfig +logger = logging.getLogger(__name__) + DEFAULT_ROPE_THETA = 10000.0 +# rope_type values that need no scaling beyond the base. +_UNSCALED_ROPE_TYPES = frozenset({"default", ""}) + +# rope_type values already warned about, so a per-layer rotary build does not +# repeat the same line once per decoder layer. +_WARNED_ROPE_TYPES: set[str] = set() + def _lookup(source: Any, key: str): """Read ``key`` from either a mapping or a config object.""" @@ -30,6 +40,32 @@ def _lookup(source: Any, key: str): return getattr(source, key, None) +def _warn_once_if_scaling_is_ignored(source: Any, rope_parameters: Any) -> None: + """Warn when the config asks for RoPE scaling the DFlash draft cannot apply. + + ``DFlashRotaryEmbedding`` takes only a base, so a target using yarn, linear + or llama3 scaling gets a draft whose rotary phase diverges from it past the + original context length. That is invisible for the same reason a defaulted + base is, so say it once per distinct ``rope_type``. + """ + rope_type = _lookup(rope_parameters, "rope_type") + if rope_type is None: + rope_scaling = _lookup(source, "rope_scaling") + rope_type = _lookup(rope_scaling, "rope_type") or _lookup(rope_scaling, "type") + if rope_type is None: + return + rope_type = str(rope_type) + if rope_type in _UNSCALED_ROPE_TYPES or rope_type in _WARNED_ROPE_TYPES: + return + _WARNED_ROPE_TYPES.add(rope_type) + logger.warning( + "DFlash draft RoPE ignores rope_type=%r: the draft rotary applies the base only, " + "so its phase diverges from a target using scaled RoPE beyond the original " + "context length.", + rope_type, + ) + + def resolve_rope_theta(source: Any, default: float = DEFAULT_ROPE_THETA) -> float: """Resolve the RoPE base from a config that may nest it under ``rope_parameters``. @@ -49,10 +85,12 @@ def resolve_rope_theta(source: Any, default: float = DEFAULT_ROPE_THETA) -> floa Returns: float: The resolved RoPE base. """ + rope_parameters = _lookup(source, "rope_parameters") + _warn_once_if_scaling_is_ignored(source, rope_parameters) top_level = _lookup(source, "rope_theta") if top_level is not None: return float(top_level) - nested = _lookup(_lookup(source, "rope_parameters"), "rope_theta") + nested = _lookup(rope_parameters, "rope_theta") if nested is not None: return float(nested) return float(default) From 6315c40f354e5155bbad269d406d9ef1d588aa14 Mon Sep 17 00:00:00 2001 From: khazic Date: Thu, 27 Aug 2026 21:53:08 +0800 Subject: [PATCH 16/16] test(dflash2): guard the new rope tests on transformers The CPU unit-test job runs without transformers and relies on every test that imports verl_speco.models skipping itself. The three rope-resolver tests import it transitively through verl_speco.models.dflash, so they errored there while passing locally. Signed-off-by: khazic --- tests/integration/test_dflash2_backend_contract.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py index e7ac3bd1..576f7f05 100644 --- a/tests/integration/test_dflash2_backend_contract.py +++ b/tests/integration/test_dflash2_backend_contract.py @@ -638,6 +638,7 @@ def test_rope_theta_defaults_when_neither_spelling_is_present() -> None: def test_resolve_rope_theta_reads_config_objects_and_mappings() -> None: + pytest.importorskip("transformers") from types import SimpleNamespace from verl_speco.models.dflash import resolve_rope_theta @@ -831,6 +832,7 @@ def test_scaled_rope_is_warned_about_not_silently_dropped(caplog) -> None: has to be said out loud rather than discarded with the rest of ``rope_parameters``. """ + pytest.importorskip("transformers") import logging from verl_speco.models.dflash import configuration_dflash @@ -854,6 +856,7 @@ def test_scaled_rope_is_warned_about_not_silently_dropped(caplog) -> None: def test_default_rope_type_is_not_warned_about(caplog) -> None: + pytest.importorskip("transformers") import logging from verl_speco.models.dflash import configuration_dflash