diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py new file mode 100644 index 00000000..576f7f05 --- /dev/null +++ b/tests/integration/test_dflash2_backend_contract.py @@ -0,0 +1,868 @@ +# 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 +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 +``tests/special_standalone/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 _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") + 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") + pytest.importorskip("transformers") + 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") + pytest.importorskip("transformers") + 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") + pytest.importorskip("transformers") + 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") + pytest.importorskip("transformers") + 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_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") + 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_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 verl_speco.models.dflash2 import DFlash2Config + + # 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 + + +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 + + +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) + + +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 DFlashConfig + + assert DFlashConfig().rope_theta == pytest.approx(10000.0) + + +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 + + 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.""" + 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: + """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 = 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.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 _is_dflash2_module_key(key) + } + + 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_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 + + +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``. + """ + pytest.importorskip("transformers") + 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: + pytest.importorskip("transformers") + 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/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"), 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/tests/special_standalone/dflash2_gpu_smoke.py b/tests/special_standalone/dflash2_gpu_smoke.py new file mode 100644 index 00000000..718731d1 --- /dev/null +++ b/tests/special_standalone/dflash2_gpu_smoke.py @@ -0,0 +1,202 @@ +# 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() + # 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 + 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): + 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() + 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 + # 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 + ) + 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"unary_only_acc={base_acc:.4f} lift={sel_acc - base_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() diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py new file mode 100644 index 00000000..7434ed24 --- /dev/null +++ b/verl_speco/backends/dflash2_trainer_backend.py @@ -0,0 +1,385 @@ +# 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.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.") + + +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. + + 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]) + # 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 + # 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( + selector_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. + # 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) + # 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 + # 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": ( + (torch.argmax(scores, dim=-1) == row_targets) & scored + ) + .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( + float(active_targets.numel()), dtype=torch.float32, device=device + ), + } + return self.selector_loss_weight * selector_loss, metrics + + +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" + + 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 _validate_normalized_state( + self, draft_model, normalized_state, model_path: str + ) -> None: + """Fail loud when a checkpoint carries DFlash2 modules under other names. + + 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) 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 DFlash2TrainerBackend._CHECKPOINT_KEY_ALIASES." + ) + + 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 + return int( + training_cfg.get( + "dflash2_block_size", getattr(drafter_config, "block_size", 8) + ) + ) + + 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=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), + 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 + ) + # 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 ( + 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=block_size, + 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..1a4dd902 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 @@ -50,6 +51,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, @@ -153,6 +186,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 +519,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) @@ -575,6 +657,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(), @@ -589,6 +686,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"} @@ -605,6 +704,7 @@ def forward( device=device, ), } + diagnostics.update(auxiliary_metrics) return ( loss, @@ -617,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 @@ -703,7 +808,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), @@ -762,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: @@ -908,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] = {} @@ -931,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/backends/domino_trainer_backend.py b/verl_speco/backends/domino_trainer_backend.py index 87367e57..55c2e724 100644 --- a/verl_speco/backends/domino_trainer_backend.py +++ b/verl_speco/backends/domino_trainer_backend.py @@ -50,9 +50,11 @@ from verl_speco.backends.dflash_trainer_backend import ( DFlashTrainerBackend, DFlashTrainingModel, + _block_acceptance_counts, _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 @@ -374,6 +376,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 +411,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(), @@ -514,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 61d58267..f8e42e02 100644 --- a/verl_speco/backends/dspark_trainer_backend.py +++ b/verl_speco/backends/dspark_trainer_backend.py @@ -24,9 +24,11 @@ from verl_speco.backends.dflash_trainer_backend import ( DFlashTrainerBackend, DFlashTrainingModel, + _block_acceptance_counts, _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 @@ -592,6 +594,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 +644,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 ( @@ -764,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/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..57bb5c4f 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): @@ -101,6 +102,8 @@ def _is_dflash_config(drafter_cfg: Any, model_configs: tuple[Any, ...]) -> bool: architecture in { "DFlashDraftModel", + "DFlash2DraftModel", + "Qwen3DFlash2Model", "DSparkDraftModel", "Qwen3DSparkModel", "DominoDraftModel", @@ -179,7 +182,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 = [] @@ -187,6 +194,13 @@ 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.extend( ( @@ -288,7 +302,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..11c409ef 100644 --- a/verl_speco/integration/sglang_runtime.py +++ b/verl_speco/integration/sglang_runtime.py @@ -174,6 +174,10 @@ 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( + # 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") @@ -603,6 +607,19 @@ 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. 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 {} training_cfg = drafter_cfg.get("training") or {} diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py index d179ddbf..e47273b6 100644 --- a/verl_speco/integration/vllm_runtime.py +++ b/verl_speco/integration/vllm_runtime.py @@ -812,6 +812,13 @@ def _drafter_algorithm(drafter_cfg: dict[str, Any]) -> str: ) +# Draft architectures vLLM can serve through its DFlash speculative path. +# Keep in sync with the alias sets in verl_speco/models/auto.py. +_DFLASH_SERVABLE_ARCHITECTURES = frozenset( + {"DFlashDraftModel", "DFlash2DraftModel", "Qwen3DFlash2Model"} +) + + def _validate_vllm_dflash_drafter_config( spec_model_path: Any, algorithm: str = "DFLASH" ) -> None: @@ -842,10 +849,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 _DFLASH_SERVABLE_ARCHITECTURES.isdisjoint(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 +921,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" 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/__init__.py b/verl_speco/models/dflash/__init__.py index 07344d68..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 DFlashConfig +from .configuration_dflash import DFlashConfig, resolve_rope_theta from .modeling_dflash import ( DFlashAttention, DFlashDecoderLayer, @@ -31,4 +31,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..4c2f0a46 100644 --- a/verl_speco/models/dflash/configuration_dflash.py +++ b/verl_speco/models/dflash/configuration_dflash.py @@ -12,11 +12,89 @@ # See the License for the specific language governing permissions and # limitations under the License. import json +import logging import os -from typing import Optional +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.""" + if source is None: + return None + if isinstance(source, Mapping): + return source.get(key) + 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``. + + 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. + """ + 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(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 +117,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 +136,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 diff --git a/verl_speco/models/dflash/modeling_dflash.py b/verl_speco/models/dflash/modeling_dflash.py index 62aa8ff4..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( @@ -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..a31b2178 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" @@ -852,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"] = ( @@ -876,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", @@ -938,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", @@ -947,6 +962,13 @@ 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_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", "sanitized_rows": f"{prefix}/sanitized_rows", "masked_rows": f"{prefix}/masked_rows", "sampled_vocab_size": f"{prefix}/sampled_vocab_size", @@ -1030,7 +1052,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: