from __future__ import annotations
from typing import Any, Iterable, TYPE_CHECKING
from .base import ModelBase, TextModel, gguf, logger
if TYPE_CHECKING:
from torch import Tensor
@ModelBase.register("Qwen3DSparkModel", "DSparkForCausalLM", "DsparkSpeculator", "DSparkModel")
class DSparkModel(TextModel):
"""Convert a DSpark checkpoint with its target vocabulary and correction weights."""
model_arch = gguf.MODEL_ARCH.DSPARK
def set_vocab(self):
# dspark drafter ships no tokenizer; it ties to the TARGET model's
# vocab and always operates on token IDs the target's own tokenizer
# already produced (never on strings). tokenizer.ggml.model=none means
# llama.cpp loads zero real vocab entries by default, which then makes
# every token id fail the generic "token < n_vocab" batch-validation
# check in llama-batch.cpp -- not just detokenization, ANY decode()
# call. add_vocab_size() tells the "none" tokenizer loader (see
# llama_vocab::impl::load in src/llama-vocab.cpp) to fill in that many
# placeholder/dummy entries, purely so vocab.n_tokens() reports the
# real (target) vocab width and batch validation passes. These
# placeholder entries carry no real strings and this GGUF cannot be
# used standalone for text I/O -- don't try to "fix" them into
# something meaningful.
self._set_vocab_none()
vocab_size = self.hparams.get("vocab_size")
if vocab_size is not None:
self.gguf_writer.add_vocab_size(int(vocab_size))
# explicit head/structural remap; per-layer decoder tensors fall through to
# the standard tensor_map (self.map_tensor_name) used by TextModel.
_name_map = {
"fc": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_FC],
"hidden_norm": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_HIDDEN_NORM],
"log_snr_embed.fc1": "dspark.log_snr_fc1",
"log_snr_embed.fc2": "dspark.log_snr_fc2",
"hidden_correction.hidden_norm": "dspark.correction_hidden_norm",
"hidden_correction.embed_norm": "dspark.correction_embed_norm",
"hidden_correction.gate_proj": "dspark.correction_gate",
"hidden_correction.up_proj": "dspark.correction_up",
"hidden_correction.down_proj": "dspark.correction_down",
# HF Qwen3DSparkModel names: markov_w1 = prev-token Embed [vocab, rank],
# markov_w2 = Linear [vocab, rank]; confidence_head is a proj with a bias.
"markov_head.markov_w1": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_MARKOV_HEAD_A],
"markov_head.prev_embed": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_MARKOV_HEAD_A],
"markov_head.markov_w2": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_MARKOV_HEAD_B],
"confidence_head.proj": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_CONFIDENCE_HEAD],
# legacy aliases (kept so older exports still convert):
"markov_head.down": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_MARKOV_HEAD_A],
"markov_head.up": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_MARKOV_HEAD_B],
"confidence_head": gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_CONFIDENCE_HEAD],
}
def set_gguf_parameters(self):
super().set_gguf_parameters()
hp = self.hparams
correction = bool(hp.get("enable_hidden_correction", False))
exposure = bool(hp.get("markov_deploy_exposure", False))
if correction and (not exposure or hp.get("markov_head_type", "vanilla") != "vanilla"):
raise ValueError("Production DSpark requires vanilla deploy-exposure Markov")
if exposure and not correction:
raise ValueError("Deploy exposure without hidden correction is unsupported")
self.gguf_writer.add_bool("dspark.dspark.hidden_correction", correction)
self.gguf_writer.add_bool("dspark.dspark.markov_deploy_exposure", exposure)
if correction:
self.gguf_writer.add_uint32("dspark.dspark.correction_size", int(hp.get("hidden_correction_intermediate_size") or hp["hidden_size"]))
snr = bool(hp.get("log_snr_conditioning", False))
self.gguf_writer.add_bool("dspark.dspark.log_snr_conditioning", snr)
if snr:
self.gguf_writer.add_float32("dspark.dspark.min_log_snr", float(hp["min_log_snr"]))
self.gguf_writer.add_float32("dspark.dspark.max_log_snr", float(hp["max_log_snr"]))
# dspark drafter trunk is a small transformer; block_count is its depth.
# (exports carry this as "num_hidden_layers"; TextModel already wired
# block_count from that, so nothing extra needed here.)
block_size = int(hp.get("block_size", 7))
self.gguf_writer.add_dspark_block_size(block_size)
mask_token_id = hp.get("mask_token_id")
if mask_token_id is not None:
self.gguf_writer.add_dspark_mask_token_id(int(mask_token_id))
target_layers = hp.get("target_layer_ids")
if target_layers is not None:
self.gguf_writer.add_dspark_target_layers([int(x) for x in target_layers])
markov_rank = hp.get("markov_rank")
if markov_rank is not None:
self.gguf_writer.add_dspark_markov_rank(int(markov_rank))
enable_conf = hp.get("enable_confidence_head")
if enable_conf is not None:
self.gguf_writer.add_dspark_confidence_head(bool(enable_conf))
conf_with_markov = hp.get("confidence_head_with_markov")
if conf_with_markov is not None:
self.gguf_writer.add_dspark_confidence_head_with_markov(bool(conf_with_markov))
logger.info("dspark: block_size=%d target_layers=%s hidden_correction=%s", block_size, target_layers, correction)
def modify_tensors(self, data_torch: "Tensor", name: str, bid: int | None) -> Iterable[tuple[str, "Tensor"]]:
n = name
# strip a leading model. / drafter. wrapper if present
for prefix in ("model.", "drafter.", "dspark."):
if n.startswith(prefix):
n = n[len(prefix):]
break
# structural / head tensors with a direct mapping
if n == "hidden_correction.gate_up_proj.weight":
width = int(self.hparams.get("hidden_correction_intermediate_size") or self.hparams["hidden_size"])
if tuple(data_torch.shape) != (2 * width, 2 * int(self.hparams["hidden_size"])):
raise ValueError("Unexpected fused hidden-correction projection shape")
if int(self.hparams.get("fused_param_tp", 1)) != 1:
raise ValueError("Unfuse TP-interleaved correction weights before GGUF conversion")
return [("dspark.correction_gate.weight", data_torch[:width]),
("dspark.correction_up.weight", data_torch[width:])]
for src, dst in self._name_map.items():
if n == f"{src}.weight":
return [(dst + ".weight", data_torch)]
if n == f"{src}.bias":
return [(dst + ".bias", data_torch)]
if n in ("norm.weight", "final_norm.weight"):
return [(gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.OUTPUT_NORM] + ".weight", data_torch)]
if n in ("lm_head.weight",):
return [(gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.OUTPUT] + ".weight", data_torch)]
if n in ("embed_tokens.weight", "tok_embeddings.weight"):
return [(gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.TOKEN_EMBD] + ".weight", data_torch)]
# per-layer decoder tensors fall through to the standard mapping. The base
# class resolves layers.{bid}.<attn/mlp...> to blk.{bid}.<gguf name>.
return [(self.map_tensor_name(n), data_torch)]