File size: 6,071 Bytes
b0486d1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | from typing import Any, Literal
from pydantic import Field, field_serializer, field_validator
from transformers import AutoConfig, PretrainedConfig
from transformers.models.qwen3.modeling_qwen3 import (
Qwen3Config,
)
from speculators import SpeculatorModelConfig
__all__ = [
"DFlashSpeculatorConfig",
]
@SpeculatorModelConfig.register("dflash")
class DFlashSpeculatorConfig(SpeculatorModelConfig):
"""
Configuration for DFlash speculator with vocabulary mapping.
DFlash features vocabulary mapping between draft (64K) and target (128K)
vocabularies, enabling cross-tokenizer speculation.
:param transformer_layer_config: Configuration for the transformer decoder layer
:param draft_vocab_size: Size of draft model vocabulary for speculation
"""
speculators_model_type: Literal["dflash"] = "dflash"
architectures: list[str] = Field(
default_factory=lambda: ["DFlashSpeculator"],
description="Model architectures that can load these weights",
)
transformer_layer_config: PretrainedConfig = Field(
default_factory=Qwen3Config,
description="Configuration for the transformer decoder layer",
)
draft_vocab_size: int = Field(
default=32000,
description="Size of draft model vocabulary for speculation",
)
block_size: int = Field(
default=8,
description=(
"Default size of the draft block predicted with a forward pass of the model"
),
)
max_anchors: int = Field(
default=256,
description=(
"Maximum number of anchor positions to sample during training "
"(controls memory usage and training efficiency)"
),
)
target_hidden_size: int | None = Field(
default=None,
description="Hidden size of the target model (if different from draft model)",
)
aux_hidden_state_layer_ids: list[int] | None = Field(
default=None,
description="Layer IDs of the DFlash auxiliary hidden state layers",
)
decoder_layer_type: Literal["qwen3", "laguna_xs"] = Field(
default="qwen3",
description="Decoder layer implementation used by the DFlash drafter.",
)
mask_token_id: int | None = Field(
default=None,
description="Token ID used for masking",
)
sliding_window_non_causal: bool = Field(
default=False,
description="Use non-causal synthetic block attention for sliding-window layers.",
)
sliding_window_base: Literal["fixed_anchor", "moving_query"] = Field(
default="moving_query",
description=(
"Base-token sliding-window lower-bound policy. 'moving_query' matches "
"FlashAttention-style SWA during inference; 'fixed_anchor' preserves "
"the legacy DFlash training mask."
),
)
loss_type: Literal["distill", "dflash", "lk", "tv"] = Field(
default="distill",
description="DFlash objective. 'lk' uses hard-label LK loss.",
)
ce_weight: float | None = Field(
default=None,
description="Additive weight for hard-label DFlash CE.",
)
tv_weight: float | None = Field(
default=None,
description="Additive weight for full-distribution TV loss.",
)
kl_weight: float | None = Field(
default=None,
description="Additive weight for full-distribution KL distillation.",
)
lk_lambda: float = Field(
default=0.5,
description="Blend coefficient for hard-label LK loss.",
)
tv_temperature: float = Field(
default=1.0,
description="Teacher softmax temperature for TV/KL terms.",
)
cumacc_weight: bool = Field(
default=False,
description="Weight hard-label DFlash CE by draft cumulative acceptance.",
)
veri_cum_acc: bool = Field(
default=False,
description="Weight DFlash loss by verifier cumulative acceptance.",
)
veri_acc_temperature: float = Field(
default=1.0,
description="Temperature for verifier cumulative acceptance weighting.",
)
static_decay_weight: bool = Field(
default=True,
description="Apply DFlash position decay to hard-label CE.",
)
kl_distill_weight: float = Field(
default=0.0,
description="Back-compatible alias for kl_weight when kl_weight is unset.",
)
compile_decoder_layers: bool = Field(
default=True,
description=(
"If True, torch.compile each decoder layer forward during training. "
"The DFlash loss remains eager."
),
)
@field_serializer("transformer_layer_config")
def serialize_transformer_config(self, value: PretrainedConfig) -> dict:
"""Serialize transformer config to dict."""
return value.to_diff_dict()
@field_validator("transformer_layer_config", mode="before")
@classmethod
def validate_transformer_config(cls, value: Any) -> PretrainedConfig:
"""Validate and convert transformer config."""
if isinstance(value, dict):
config_class: type[PretrainedConfig] = Qwen3Config
if "model_type" in value:
config_class = AutoConfig.for_model(
model_type=value["model_type"]
).__class__
return config_class(**value)
return value
@property
def target_vocab_size(self) -> int:
"""Get target vocabulary size from transformer config."""
return self.transformer_layer_config.vocab_size
def resolve_loss_weights(self) -> tuple[float, float, float]:
if self.loss_type == "tv":
ce_default, tv_default = 0.0, 1.0
else:
ce_default, tv_default = 1.0, 0.0
ce = ce_default if self.ce_weight is None else self.ce_weight
tv = tv_default if self.tv_weight is None else self.tv_weight
kl = self.kl_distill_weight if self.kl_weight is None else self.kl_weight
return float(ce), float(tv), float(kl)
|