dspark.py
6.2 KB · 161 lines · python Raw
1 # coding=utf-8
2 """DSpark draft model: DFlash backbone + EAGLE-style Markov and confidence heads.
3
4 DSpark shares SpecForge's DFlash block-diffusion drafter (dual-source KV
5 injection via :class:`DFlashDraftModel`, anchor sampling, MASK-token noise
6 stream) and adds two heads on top:
7
8 - Markov head: a low-rank learned bigram bias added to the draft logits,
9 conditioned on the (teacher-forced) previous token. Improves the per-token
10 distribution without touching the backbone.
11 - Confidence head (AcceptRatePredictor): predicts a per-draft-position
12 acceptance probability, trained against the empirical draft-vs-target
13 accept rate (used at inference time for adaptive block length).
14
15 Ported from TorchSpec PR #129 (``torchspec/models/draft/dspark.py``). The Markov
16 / confidence modeling code is adapted from DeepSeek's DeepSpec
17 (``deepspec/modeling/dspark/{markov_head,common}.py``, MIT License).
18
19 SpecForge differences vs TorchSpec (load-bearing):
20 - There is no ``DFlashConfig``; SpecForge's :class:`DFlashDraftModel` uses a
21 plain ``Qwen3Config`` plus a ``config.dflash_config`` dict. So
22 :class:`DSparkConfig` subclasses ``Qwen3Config`` and declares the DSpark
23 fields as top-level attributes; DFlash-carried fields (``block_size``,
24 ``num_target_layers``, ``dflash_config``) stay as before.
25 - The draft model has no ``embed_tokens`` of its own (the embedding lives on
26 the target and is passed into the online wrapper), and the context
27 projection is ``self.fc`` (not ``context_proj``). The heads only depend on
28 ``config.hidden_size`` / ``config.vocab_size``, so this does not matter for
29 construction.
30 """
31
32 from typing import Optional
33
34 import torch
35 import torch.nn as nn
36 from transformers.models.qwen3.modeling_qwen3 import Qwen3Config
37
38 from specforge.modeling.draft.dflash import DFlashDraftModel
39
40
41 class DSparkConfig(Qwen3Config):
42 """Configuration for the DSpark draft model.
43
44 Extends ``Qwen3Config`` (SpecForge's DFlash draft is config-light and reads a
45 plain ``Qwen3Config``). DSpark-specific fields are declared here; the
46 DFlash-carried fields (``block_size``, ``num_target_layers``, and the nested
47 ``dflash_config`` dict holding ``target_layer_ids`` / ``mask_token_id``) are
48 consumed by the :class:`DFlashDraftModel` base ``__init__`` and must be
49 present on the config object before constructing the model.
50 """
51
52 model_type = "dspark"
53
54 def __init__(
55 self,
56 markov_rank: int = 256,
57 markov_head_type: str = "vanilla",
58 enable_confidence_head: bool = True,
59 confidence_head_with_markov: bool = True,
60 **kwargs,
61 ):
62 super().__init__(**kwargs)
63 self.markov_rank = markov_rank
64 self.markov_head_type = markov_head_type
65 self.enable_confidence_head = enable_confidence_head
66 self.confidence_head_with_markov = confidence_head_with_markov
67
68
69 class VanillaMarkov(nn.Module):
70 """Low-rank learned bigram bias added to the draft logits.
71
72 Adapted from DeepSpec's ``deepspec/modeling/dspark/markov_head.py``.
73 """
74
75 def __init__(self, *, vocab_size: int, markov_rank: int):
76 super().__init__()
77 self.vocab_size = int(vocab_size)
78 self.markov_rank = int(markov_rank)
79 self.markov_head_type = "vanilla"
80 assert (
81 self.markov_rank > 0
82 ), f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}."
83 self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank)
84 self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False)
85
86 def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor:
87 return self.markov_w1(token_ids.long())
88
89 def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor:
90 return self.markov_w2(latent_states)
91
92 def compute_step_bias(self, token_ids: torch.Tensor) -> torch.Tensor:
93 return self.project_bias(self.get_prev_embeddings(token_ids))
94
95 def apply_block_logits(
96 self,
97 base_logits: torch.Tensor,
98 *,
99 token_ids: torch.Tensor,
100 ) -> torch.Tensor:
101 if base_logits.size(2) == 0:
102 return base_logits
103 return base_logits + self.compute_step_bias(token_ids)
104
105
106 class AcceptRatePredictor(nn.Module):
107 """Per-position acceptance-probability predictor (a single linear head).
108
109 Adapted from DeepSpec's ``deepspec/modeling/dspark/common.py``.
110 """
111
112 def __init__(self, input_dim: int):
113 super().__init__()
114 self.proj = nn.Linear(int(input_dim), 1)
115
116 def forward(self, features: torch.Tensor) -> torch.Tensor:
117 return self.proj(features).squeeze(-1)
118
119
120 def build_markov_head(config) -> Optional[nn.Module]:
121 markov_rank = int(getattr(config, "markov_rank", 0))
122 assert markov_rank >= 0, f"markov_rank must be >= 0, got {markov_rank}"
123 if markov_rank == 0:
124 return None
125
126 markov_head_type = str(getattr(config, "markov_head_type", "vanilla")).lower()
127 if markov_head_type == "vanilla":
128 return VanillaMarkov(vocab_size=config.vocab_size, markov_rank=markov_rank)
129 raise NotImplementedError(
130 f"markov_head_type={markov_head_type!r} is not supported yet; only 'vanilla' "
131 "is implemented as it is recommended by the authors."
132 )
133
134
135 class DSparkDraftModel(DFlashDraftModel):
136 """DSpark draft network: DFlash backbone + Markov / confidence heads."""
137
138 config_class = DSparkConfig
139
140 def __init__(self, config) -> None:
141 super().__init__(config)
142
143 self.markov_rank = int(getattr(config, "markov_rank", 0))
144 self.confidence_head_with_markov = bool(
145 getattr(config, "confidence_head_with_markov", True)
146 )
147
148 self.markov_head = build_markov_head(config)
149
150 self.confidence_head: Optional[nn.Module] = None
151 if getattr(config, "enable_confidence_head", False):
152 conf_input_dim = config.hidden_size
153 if self.confidence_head_with_markov:
154 if self.markov_head is None:
155 raise ValueError(
156 "confidence_head_with_markov=True requires a Markov head "
157 "(markov_rank > 0)."
158 )
159 conf_input_dim += self.markov_rank
160 self.confidence_head = AcceptRatePredictor(conf_input_dim)
161