File size: 6,534 Bytes
e65937c | 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 190 191 192 193 | """
Provides the entry points for converting non-speculators model checkpoints to
Speculators model format with the `convert_model` function.
It supports the following algorithms and conversion from their associated
research repositories:
- EAGLE3
- MTP
- DFlash
- Domino
Functions:
convert_model: Converts a model checkpoint to the Speculators format.
"""
import os
import tempfile
from typing import Literal
from loguru import logger
from transformers import PretrainedConfig
from speculators.convert.dflash.converter import DFlashConverter
from speculators.convert.domino.converter import DominoConverter
from speculators.convert.eagle.eagle3_converter import Eagle3Converter
from speculators.convert.mtp.converter import MTPConverter
__all__ = ["convert_model", "maybe_convert_external_checkpoint"]
def convert_model(
model: str,
verifier: str,
algorithm: Literal["eagle3", "mtp", "dflash", "domino"],
output_path: str = "converted",
validate_device: str | None = None,
**kwargs,
):
"""
Convert a non speculator's model checkpoint to a speculator's model checkpoint
for use within the Speculators library, Hugging Face Hub, or vLLM.
algorithm=="eagle3":
Eagle v3: https://github.com/SafeAILab/EAGLE
::
# general
convert_model(
model="./eagle/checkpoint",
verifier="meta-llama/Llama-3.1-8B-Instruct",
algorithm="eagle3",
)
# with normalization before the residual
convert_model(
model="./eagle/checkpoint",
verifier="meta-llama/Llama-3.1-8B-Instruct",
algorithm="eagle3",
norm_before_residual=True,
)
algorithm=="mtp":
MTP (Multi-Token Prediction): models with native MTP layers
(e.g. Qwen3-Next, Qwen3.5, Qwen3.5-MoE)
::
convert_model(
model="Qwen/Qwen3-Next-80B-A3B-Instruct",
verifier="Qwen/Qwen3-Next-80B-A3B-Instruct",
algorithm="mtp",
num_speculative_steps=3,
)
algorithm=="dflash":
DFlash: https://z-lab.ai/projects/dflash/
::
convert_model(
model="z-lab/Qwen3-8B-DFlash-b16",
verifier="Qwen/Qwen3-8B",
algorithm="dflash",
)
algorithm=="domino":
Domino: https://github.com/sgl-project/SpecForge
::
convert_model(
model="./domino/checkpoint",
verifier="Qwen/Qwen3-8B",
algorithm="domino",
)
:param model: Path to the input model checkpoint or Hugging Face model ID.
:param verifier: Verifier model checkpoint or Hugging Face model ID
to attach as the verification/base model for speculative decoding
:param algorithm: The conversion algorithm to use:
"eagle3", "mtp", "dflash", or "domino".
:param output_path: Directory path where the converted model will be saved.
:param kwargs: Additional keyword arguments for the conversion algorithm.
Options for Eagle3: {"norm_before_residual": true,
"eagle_aux_hidden_state_layer_ids": [1,23,44]}.
Options for MTP: {"num_speculative_steps": 3}.
Options for DFlash: {"aux_hidden_state_layer_ids": [2,10,18,26,34]}.
Options for Domino: {"aux_hidden_state_layer_ids": [2,10,18,26,34]};
the recurrent head's dimensions are read from the checkpoint's
dflash_config. Requires draft_vocab_size == vocab_size.
"""
if algorithm == "eagle3":
Eagle3Converter().convert(
model,
output_path,
verifier,
validate=validate_device is not None,
**kwargs,
)
elif algorithm == "mtp":
MTPConverter().convert(
model,
output_path,
verifier,
validate=validate_device is not None,
**kwargs,
)
elif algorithm == "dflash":
DFlashConverter().convert(
model,
output_path,
verifier,
validate=validate_device is not None,
**kwargs,
)
elif algorithm == "domino":
DominoConverter().convert(
model,
output_path,
verifier,
validate=validate_device is not None,
**kwargs,
)
else:
raise ValueError(f"Unsupported algorithm: {algorithm}")
def maybe_convert_external_checkpoint(
model: str | os.PathLike,
verifier: str | None = None,
cache_dir: str | os.PathLike | None = None,
output_path: str | None = None,
config_dict: dict | None = None,
) -> str:
"""Convert an external (non-speculators) checkpoint to speculators format.
A speculators checkpoint (config has ``speculators_model_type``) is returned
unchanged; otherwise the external format is detected and converted (which
requires ``verifier``) to ``output_path``, defaulting to a temp dir. Powers
the unified ``from_pretrained`` finetuning pathway. Pass ``config_dict`` to
reuse an already-loaded config and skip re-reading it.
"""
if config_dict is None:
config_dict, _ = PretrainedConfig.get_config_dict(model, cache_dir=cache_dir)
if "speculators_model_type" in config_dict:
return str(model)
architectures = config_dict.get("architectures") or []
dflash_config = config_dict.get("dflash_config") or {}
if dflash_config.get("projector_type") == "domino" or any(
"Domino" in architecture for architecture in architectures
):
algorithm: Literal["dflash", "domino"] = "domino"
elif "dflash_config" in config_dict or any(
"DFlash" in architecture for architecture in architectures
):
algorithm = "dflash"
else:
raise NotImplementedError(
f"Cannot auto-convert checkpoint '{model}': unrecognized external "
"format. Supported auto-conversion: DFlash and Domino."
)
if verifier is None:
raise ValueError(
f"Converting an external {algorithm} checkpoint requires a verifier. "
"Pass `verifier=<model id or path>`."
)
output_path = output_path or tempfile.mkdtemp(prefix="speculators_converted_")
logger.info(f"Auto-converting external {algorithm} checkpoint to {output_path}")
convert_model(
model=str(model),
verifier=verifier,
algorithm=algorithm,
output_path=output_path,
cache_dir=cache_dir,
)
return output_path
|