File size: 3,104 Bytes
3d55cc1 1afd905 3d55cc1 1afd905 | 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 | 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",
)
mask_token_id: int | None = Field(
default=None,
description="Token ID used for masking",
)
@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
|