diff --git a/src/speculators/convert/dflash/converter.py b/src/speculators/convert/dflash/converter.py index 5d76f1514..1147ec2c5 100644 --- a/src/speculators/convert/dflash/converter.py +++ b/src/speculators/convert/dflash/converter.py @@ -117,7 +117,13 @@ def _build_config( # z-lab reads hidden_states[layer_id + 1] (index 0 is the embedding # output) while speculators uses the layer id directly. # Source: z-lab utils.extract_context_feature. - aux_hidden_state_layer_ids = [i + 1 for i in target_layer_ids] + # Exclude the last verifier layer — it is always included + # implicitly by launch_vllm.py (--include-last-layer) and split + # off as verifier_last_hidden_states during training. + num_verifier_layers = verifier_config_dict["num_hidden_layers"] + aux_hidden_state_layer_ids = [ + i + 1 for i in target_layer_ids if i + 1 != num_verifier_layers + ] speculators_config = SpeculatorsConfig( algorithm="dflash", @@ -142,6 +148,67 @@ def _build_config( speculators_config=speculators_config, ) + def _remap_weights( + self, + weights: dict[str, torch.Tensor], + config: DFlashSpeculatorConfig, + model: DFlashDraftModel, + ) -> dict[str, torch.Tensor]: + """Remap checkpoint weights to match DFlashDraftModel's state dict. + + Handles Laguna-style fused ``qkv_proj`` → separate ``q/k/v_proj``, + drops ``g_proj`` and ``aux_hidden_norms`` (no DFlash equivalent), and + slices ``fc.weight`` when the source has more target layers than needed. + """ + has_fused_qkv = any("qkv_proj" in k for k in weights) + if not has_fused_qkv: + return weights + + tl = config.transformer_layer_config + q_dim = tl.num_attention_heads * tl.head_dim + kv_dim = tl.num_key_value_heads * tl.head_dim + + remapped: dict[str, torch.Tensor] = {} + dropped: list[str] = [] + + for key, tensor in weights.items(): + if "qkv_proj" in key: + projection_keys = [ + key.replace("qkv_proj", proj) + for proj in ("q_proj", "k_proj", "v_proj") + ] + conflicts = [pk for pk in projection_keys if pk in weights] + if conflicts: + raise ValueError( + f"Checkpoint contains both fused qkv_proj and separate " + f"projection keys: {conflicts}" + ) + q, k, v = tensor.split([q_dim, kv_dim, kv_dim], dim=0) + remapped[projection_keys[0]] = q + remapped[projection_keys[1]] = k + remapped[projection_keys[2]] = v + elif ".g_proj." in key or key.startswith("aux_hidden_norms."): + dropped.append(key) + elif key == "fc.weight": + model_fc_dim = model.fc.in_features + if tensor.shape[1] > model_fc_dim: + logger.info( + f"Slicing fc.weight from {tensor.shape[1]} to {model_fc_dim}" + ) + remapped[key] = tensor[:, :model_fc_dim] + else: + remapped[key] = tensor + else: + remapped[key] = tensor + + if dropped: + logger.info(f"Dropped {len(dropped)} incompatible keys: {dropped}") + logger.info( + f"Remapped {sum(1 for k in weights if 'qkv_proj' in k)} fused qkv_proj " + f"→ separate q/k/v_proj" + ) + return remapped + def _save( self, config: DFlashSpeculatorConfig, @@ -151,6 +218,7 @@ def _save( model = DFlashDraftModel(config=config) body = {k: v for k, v in weights.items() if k not in ("t2d", "d2t")} + body = self._remap_weights(body, config, model) missing, unexpected = model.load_state_dict(body, strict=False) if unexpected: raise ValueError( diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index 71bf32c23..defd53d06 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -1,4 +1,5 @@ import logging +from copy import deepcopy from typing import ClassVar import torch @@ -95,7 +96,19 @@ def __init__( config.transformer_layer_config.hidden_size, eps=config.transformer_layer_config.rms_norm_eps, # type: ignore[arg-type] ) - self.rotary_emb = Qwen3RotaryEmbedding(config.transformer_layer_config) # type: ignore[arg-type] + rotary_config = config.transformer_layer_config + rope_params = getattr(rotary_config, "rope_parameters", None) + if rope_params and "sliding_attention" in rope_params: + if self.uses_full_attn: + logger.warning( + "Flattening nested rope_parameters to the sliding_attention " + "variant, but this model has %d full-attention layer(s). " + "Full-attention layers may use incorrect rope scaling.", + num_draft_layers - len(self.sliding_window_indices), + ) + rotary_config = deepcopy(rotary_config) + rotary_config.rope_parameters = rope_params["sliding_attention"] + self.rotary_emb = Qwen3RotaryEmbedding(rotary_config) # type: ignore[arg-type] self.fc = nn.Linear( len(self.target_layer_ids) * config.transformer_layer_config.hidden_size, diff --git a/tests/unit/convert/test_dflash_converter.py b/tests/unit/convert/test_dflash_converter.py index 959a56bb9..4a8b98f25 100644 --- a/tests/unit/convert/test_dflash_converter.py +++ b/tests/unit/convert/test_dflash_converter.py @@ -1,31 +1,40 @@ -"""Unit tests for DFlashConverter config building.""" +"""Unit tests for DFlashConverter config building and weight remapping.""" +import logging from unittest.mock import patch import pytest +import torch from transformers import Qwen3Config from speculators.config import SpeculatorsConfig, VerifierConfig from speculators.convert.dflash.converter import DFlashConverter -from speculators.models.dflash import DFlashSpeculatorConfig +from speculators.models.dflash import DFlashDraftModel, DFlashSpeculatorConfig from speculators.proposals.greedy import GreedyTokenProposalConfig +_HIDDEN = 16 +_NUM_HEADS = 2 +_NUM_KV_HEADS = 1 +_HEAD_DIM = 8 +_Q_DIM = _NUM_HEADS * _HEAD_DIM # 16 +_KV_DIM = _NUM_KV_HEADS * _HEAD_DIM # 8 -def _tiny_dflash_config(): + +def _tiny_dflash_config(num_aux_layers=1): return DFlashSpeculatorConfig( transformer_layer_config=Qwen3Config( vocab_size=32, - hidden_size=16, + hidden_size=_HIDDEN, intermediate_size=32, num_hidden_layers=1, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=8, + num_attention_heads=_NUM_HEADS, + num_key_value_heads=_NUM_KV_HEADS, + head_dim=_HEAD_DIM, max_position_embeddings=32, ), draft_vocab_size=32, block_size=4, - aux_hidden_state_layer_ids=[0], + aux_hidden_state_layer_ids=list(range(num_aux_layers)), mask_token_id=1, speculators_config=SpeculatorsConfig( algorithm="dflash", @@ -63,7 +72,11 @@ class TestBuildConfig: @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict") def test_happy_path(self, mock_get_config): mock_get_config.return_value = ( - {"hidden_size": 4096, "architectures": ["Qwen3ForCausalLM"]}, + { + "hidden_size": 4096, + "num_hidden_layers": 36, + "architectures": ["Qwen3ForCausalLM"], + }, None, ) config = DFlashConverter()._build_config( @@ -82,6 +95,23 @@ def test_happy_path(self, mock_get_config): assert config.transformer_layer_config.num_hidden_layers == 5 assert not hasattr(config.transformer_layer_config, "dflash_config") + @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict") + def test_excludes_last_verifier_layer(self, mock_get_config): + mock_get_config.return_value = ( + { + "hidden_size": 4096, + "num_hidden_layers": 34, + "architectures": ["Qwen3ForCausalLM"], + }, + None, + ) + # target_layer_ids [1, 9, 17, 25, 33] → +1 → [2, 10, 18, 26, 34] + # but 34 == num_hidden_layers, so it should be excluded + config = DFlashConverter()._build_config( + _source_config(), "Qwen/Qwen3-8B", None + ) + assert config.aux_hidden_state_layer_ids == [2, 10, 18, 26] + @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict") def test_explicit_aux_layer_ids_override(self, mock_get_config): mock_get_config.return_value = ({"hidden_size": 4096}, None) @@ -98,12 +128,122 @@ def test_hidden_size_mismatch_raises(self, mock_get_config): @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict") def test_missing_target_layer_ids_raises(self, mock_get_config): - mock_get_config.return_value = ({"hidden_size": 4096}, None) + mock_get_config.return_value = ( + {"hidden_size": 4096, "num_hidden_layers": 36}, + None, + ) source = _source_config(dflash_config={"mask_token_id": 151669}) with pytest.raises(ValueError, match="target_layer_ids"): DFlashConverter()._build_config(source, "Qwen/Qwen3-8B", None) +class TestRemapWeights: + def _make_fused_weights(self): + qkv = torch.randn(_Q_DIM + 2 * _KV_DIM, _HIDDEN) + return { + "layers.0.self_attn.qkv_proj.weight": qkv, + "layers.0.self_attn.g_proj.weight": torch.randn(2, _HIDDEN), + "layers.0.self_attn.o_proj.weight": torch.randn(_HIDDEN, _Q_DIM), + "aux_hidden_norms.0.weight": torch.randn(_HIDDEN), + "fc.weight": torch.randn(_HIDDEN, _HIDDEN * 2), + "norm.weight": torch.randn(_HIDDEN), + } + + def test_splits_fused_qkv(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + weights = self._make_fused_weights() + qkv = weights["layers.0.self_attn.qkv_proj.weight"] + + remapped = DFlashConverter()._remap_weights(weights, config, model) + + assert "layers.0.self_attn.qkv_proj.weight" not in remapped + assert torch.equal(remapped["layers.0.self_attn.q_proj.weight"], qkv[:_Q_DIM]) + assert torch.equal( + remapped["layers.0.self_attn.k_proj.weight"], + qkv[_Q_DIM : _Q_DIM + _KV_DIM], + ) + assert torch.equal( + remapped["layers.0.self_attn.v_proj.weight"], + qkv[_Q_DIM + _KV_DIM :], + ) + + def test_drops_g_proj_and_aux_hidden_norms(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + remapped = DFlashConverter()._remap_weights( + self._make_fused_weights(), config, model + ) + assert not any("g_proj" in k for k in remapped) + assert not any("aux_hidden_norms" in k for k in remapped) + + def test_slices_fc_weight(self): + config = _tiny_dflash_config(num_aux_layers=1) + model = DFlashDraftModel(config=config) + weights = self._make_fused_weights() + # fc from checkpoint is wider than model expects + wide_fc = torch.randn(_HIDDEN, _HIDDEN * 3) + weights["fc.weight"] = wide_fc + + remapped = DFlashConverter()._remap_weights(weights, config, model) + assert remapped["fc.weight"].shape[1] == model.fc.in_features + assert torch.equal(remapped["fc.weight"], wide_fc[:, : model.fc.in_features]) + + def test_passthrough_when_no_fused_qkv(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + weights = {"norm.weight": torch.randn(_HIDDEN)} + remapped = DFlashConverter()._remap_weights(weights, config, model) + assert remapped is weights + + def test_preserves_other_keys(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + weights = self._make_fused_weights() + remapped = DFlashConverter()._remap_weights(weights, config, model) + assert "layers.0.self_attn.o_proj.weight" in remapped + assert "norm.weight" in remapped + + def test_mixed_fused_and_separate_qkv_raises(self): + config = _tiny_dflash_config() + model = DFlashDraftModel(config=config) + weights = self._make_fused_weights() + weights["layers.0.self_attn.q_proj.weight"] = torch.randn(_Q_DIM, _HIDDEN) + with pytest.raises(ValueError, match="both fused qkv_proj and separate"): + DFlashConverter()._remap_weights(weights, config, model) + + +class TestRopeParameters: + _ROPE_SLIDING = { + "rope_type": "default", + "rope_theta": 10000.0, + } + + def test_nested_rope_params_warns_with_full_attn(self, caplog): + config = _tiny_dflash_config() + tl = config.transformer_layer_config + tl.layer_types = ["full_attention"] + tl.rope_parameters = { + "sliding_attention": dict(self._ROPE_SLIDING), + "full_attention": {"rope_type": "default", "rope_theta": 1000000.0}, + } + with caplog.at_level(logging.WARNING, logger="speculators.models.dflash.core"): + DFlashDraftModel(config=config) + assert any("full-attention layer" in r.message for r in caplog.records) + + def test_nested_rope_params_no_warn_all_sliding(self, caplog): + config = _tiny_dflash_config() + tl = config.transformer_layer_config + tl.layer_types = ["sliding_attention"] + tl.sliding_window = 512 + tl.rope_parameters = { + "sliding_attention": dict(self._ROPE_SLIDING), + } + with caplog.at_level(logging.WARNING, logger="speculators.models.dflash.core"): + DFlashDraftModel(config=config) + assert not any("full-attention layer" in r.message for r in caplog.records) + + class TestSave: def test_missing_draft_weights_raise(self, tmp_path): # No source weights: every draft-body weight (fc, norm, hidden_norm,