Add DFlash model converter

scripts/train.py

@@ -406,7 +406,12 @@ def build_draft_model(
                 d2t=d2t,
                 verifier_name_or_path=args.verifier_name_or_path,
             )
-        return model_class.from_pretrained(args.from_pretrained, t2d=t2d, d2t=d2t)
+        return model_class.from_pretrained(
+            args.from_pretrained,
+            t2d=t2d,
+            d2t=d2t,
+            verifier=args.verifier_name_or_path,
+        )
 
     if args.speculator_type == "mtp":
         # MTP uses the verifier's own decoder config as the draft

src/speculators/convert/dflash/__init__.py

@@ -0,0 +1,5 @@
+"""DFlash checkpoint conversion utilities."""
+
+from speculators.convert.dflash.converter import DFlashConverter
+
+__all__ = ["DFlashConverter"]

src/speculators/convert/dflash/converter.py

@@ -0,0 +1,175 @@
+"""DFlash checkpoint converter.
+
+Converts an external DFlash checkpoint (e.g. ``z-lab/*-DFlash``) to a Speculators
+checkpoint that loads with ``DFlashDraftModel.from_pretrained(path)``.
+
+The draft transformer body (``layers.*``, ``fc``, ``hidden_norm``, ``norm``) already
+matches ``DFlashDraftModel`` so weights are copied as-is. The external checkpoint
+borrows the verifier's embedding and LM head at runtime, so ``embed_tokens`` /
+``lm_head`` / ``verifier_lm_head`` / ``verifier_norm`` are loaded from the verifier
+before saving.
+"""
+
+from pathlib import Path
+
+import torch
+from loguru import logger
+from transformers import PretrainedConfig
+
+from speculators.config import SpeculatorsConfig, VerifierConfig
+from speculators.convert.utils import (
+    ensure_checkpoint_is_local,
+    load_checkpoint_config,
+    load_checkpoint_weights,
+)
+from speculators.models.dflash import DFlashDraftModel, DFlashSpeculatorConfig
+from speculators.proposals.greedy import GreedyTokenProposalConfig
+
+__all__ = ["DFlashConverter"]
+
+# config.json keys that are not part of the draft transformer (Qwen3) config
+_NON_TRANSFORMER_KEYS = frozenset(
+    {"architectures", "auto_map", "block_size", "dflash_config", "num_target_layers"}
+)
+
+# state dict keys that are filled from the verifier (not the source checkpoint), so
+# their absence from the source weights is expected, not a conversion error
+_VERIFIER_FILLED_KEYS = frozenset(
+    {
+        "embed_tokens.weight",
+        "lm_head.weight",
+        "verifier_lm_head.weight",
+        "verifier_norm.weight",
+        "t2d",
+        "d2t",
+    }
+)
+
+
+class DFlashConverter:
+    """Convert an external DFlash checkpoint to speculators format.
+
+    Copies the draft transformer body as-is and fills the embedding, LM head, and
+    verifier norm from the verifier model so the saved checkpoint is self-contained.
+    """
+
+    def convert(
+        self,
+        input_path: str | Path,
+        output_path: str | Path,
+        base_model: str,
+        validate: bool = True,
+        aux_hidden_state_layer_ids: list[int] | None = None,
+        cache_dir: str | Path | None = None,
+    ) -> None:
+        logger.info(f"Converting DFlash checkpoint: {input_path}")
+
+        local_checkpoint_path = ensure_checkpoint_is_local(input_path, cache_dir)
+        source_config = load_checkpoint_config(local_checkpoint_path)
+        weights = load_checkpoint_weights(local_checkpoint_path)
+        logger.info(f"Loaded {len(weights)} weights")
+
+        config = self._build_config(
+            source_config, base_model, aux_hidden_state_layer_ids
+        )
+        saved_path = self._save(config, weights, output_path)
+        logger.success(f"Saved to: {saved_path}")
+
+        if validate:
+            self._validate(saved_path)
+
+    def _build_config(
+        self,
+        source_config: dict,
+        base_model: str,
+        aux_hidden_state_layer_ids: list[int] | None,
+    ) -> DFlashSpeculatorConfig:
+        dflash = source_config.get("dflash_config", {})
+        transformer_config = {
+            k: v for k, v in source_config.items() if k not in _NON_TRANSFORMER_KEYS
+        }
+
+        verifier_config_dict, _ = PretrainedConfig.get_config_dict(base_model)
+        source_hidden = transformer_config.get("hidden_size")
+        target_hidden = verifier_config_dict.get("hidden_size")
+        if source_hidden and target_hidden and source_hidden != target_hidden:
+            raise ValueError(
+                f"Architecture mismatch: source DFlash checkpoint has "
+                f"hidden_size={source_hidden} but base_model '{base_model}' has "
+                f"hidden_size={target_hidden}. Dimensions must match."
+            )
+
+        if aux_hidden_state_layer_ids is None:
+            target_layer_ids = dflash.get("target_layer_ids")
+            if target_layer_ids is None:
+                raise ValueError(
+                    "Checkpoint config has no `dflash_config.target_layer_ids`; "
+                    "pass `aux_hidden_state_layer_ids` explicitly."
+                )
+            # z-lab reads hidden_states[layer_id + 1] (index 0 is the embedding
+            # output) while speculators uses the layer id directly.
+            # Source: z-lab utils.extract_context_feature.
+            aux_hidden_state_layer_ids = [i + 1 for i in target_layer_ids]
+
+        speculators_config = SpeculatorsConfig(
+            algorithm="dflash",
+            proposal_methods=[
+                GreedyTokenProposalConfig(
+                    speculative_tokens=source_config["block_size"] - 1,
+                )
+            ],
+            default_proposal_method="greedy",
+            verifier=VerifierConfig(
+                name_or_path=base_model,
+                architectures=verifier_config_dict.get("architectures", []),
+            ),
+        )
+
+        return DFlashSpeculatorConfig(
+            transformer_layer_config=transformer_config,  # type: ignore[arg-type]
+            draft_vocab_size=transformer_config["vocab_size"],
+            block_size=source_config["block_size"],
+            aux_hidden_state_layer_ids=aux_hidden_state_layer_ids,
+            mask_token_id=dflash.get("mask_token_id"),
+            speculators_config=speculators_config,
+        )
+
+    def _save(
+        self,
+        config: DFlashSpeculatorConfig,
+        weights: dict[str, torch.Tensor],
+        output_path: str | Path,
+    ) -> Path:
+        model = DFlashDraftModel(config=config)
+
+        body = {k: v for k, v in weights.items() if k not in ("t2d", "d2t")}
+        missing, unexpected = model.load_state_dict(body, strict=False)
+        if unexpected:
+            raise ValueError(
+                "Unexpected keys in checkpoint -- the structure does not match "
+                f"DFlashDraftModel. Unexpected keys: {unexpected}"
+            )
+        critical_missing = [k for k in missing if k not in _VERIFIER_FILLED_KEYS]
+        if critical_missing:
+            raise ValueError(f"Draft weights missing after load: {critical_missing}")
+        logger.debug(f"Keys loaded from verifier at save time: {missing}")
+
+        # embed_tokens / lm_head / verifier_lm_head / verifier_norm come from the
+        # verifier; without this they would be saved as NaN.
+        model.load_verifier_weights()
+
+        model.to(dtype=next(iter(body.values())).dtype)  # type: ignore[call-arg]
+        model.save_pretrained(str(output_path))
+        return Path(output_path)
+
+    def _validate(self, output_path: Path) -> None:
+        logger.info("Validating converted DFlash checkpoint...")
+        try:
+            model = DFlashDraftModel.from_pretrained(str(output_path))
+        except (OSError, ValueError, RuntimeError) as exc:
+            logger.error(f"Validation failed: {exc}")
+            raise
+        for name in ("fc.weight", "lm_head.weight", "embed_tokens.weight"):
+            if torch.isnan(model.state_dict()[name]).any():
+                raise ValueError(f"Converted checkpoint has NaN in {name}")
+        logger.success("Validation succeeded")

src/speculators/convert/entrypoints.py

@@ -9,24 +9,31 @@
 - EAGLE3
 - HASS
 - MTP
+- DFlash
 
 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.eagle.eagle3_converter import Eagle3Converter
 from speculators.convert.eagle.eagle_converter import EagleConverter
 from speculators.convert.mtp.converter import MTPConverter
 
-__all__ = ["convert_model"]
+__all__ = ["convert_model", "maybe_convert_external_checkpoint"]
 
 
 def convert_model(
     model: str,
     verifier: str,
-    algorithm: Literal["eagle", "eagle3", "mtp"],
+    algorithm: Literal["eagle", "eagle3", "mtp", "dflash"],
     output_path: str = "converted",
     validate_device: str | None = None,
     **kwargs,
@@ -82,17 +89,27 @@ def convert_model(
             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",
+        )
+
     :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:
-        "eagle", "eagle3", or "mtp".
+        "eagle", "eagle3", "mtp", or "dflash".
     :param output_path: Directory path where the converted model will be saved.
     :param kwargs: Additional keyword arguments for the conversion algorithm.
         Options for Eagle: {"layernorms": true, "fusion_bias": true}.
         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]}.
     """
 
     if algorithm == "eagle":
@@ -119,5 +136,61 @@ def convert_model(
             validate=validate_device is not None,
             **kwargs,
         )
+    elif algorithm == "dflash":
+        DFlashConverter().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 []
+    algorithm: Literal["dflash"] = "dflash"
+    if not (
+        "dflash_config" in config_dict or any("DFlash" in a for a in architectures)
+    ):
+        raise NotImplementedError(
+            f"Cannot auto-convert checkpoint '{model}': unrecognized external "
+            "format. Supported auto-conversion: DFlash."
+        )
+
+    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

src/speculators/model.py

@@ -293,6 +293,8 @@ def from_pretrained(
             If None, automatically detects the available format.
         :param weights_only: Whether to only load model weights without optimizer
             states or other training artifacts.
+        :param verifier: Verifier model id/path used to auto-convert an external
+            (non-speculators) checkpoint; ignored for speculators checkpoints.
         :param kwargs: Additional keyword arguments passed to the model constructor
             and loading process.
         :return: A SpeculatorModel instance of the appropriate subclass, loaded with
@@ -304,6 +306,23 @@ def from_pretrained(
                     "Either `config` or `pretrained_model_name_or_path` must be "
                     "provided to load a SpeculatorModel."
                 )
+            # Auto-convert external (non-speculators) checkpoints so one
+            # `from_pretrained` pathway finetunes both formats. Detect format
+            # once here and only invoke the converter when needed.
+            config_dict, _ = PretrainedConfig.get_config_dict(
+                pretrained_model_name_or_path, cache_dir=cache_dir
+            )
+            if "speculators_model_type" not in config_dict:
+                from speculators.convert.entrypoints import (  # noqa: PLC0415
+                    maybe_convert_external_checkpoint,
+                )
+
+                pretrained_model_name_or_path = maybe_convert_external_checkpoint(
+                    pretrained_model_name_or_path,
+                    verifier=kwargs.get("verifier"),
+                    cache_dir=cache_dir,
+                    config_dict=config_dict,
+                )
             config = cls.config_class.from_pretrained(
                 pretrained_model_name_or_path,
                 cache_dir=cache_dir,
@@ -314,8 +333,6 @@ def from_pretrained(
             )
 
         if not isinstance(config, SpeculatorModelConfig):
-            # once conversion is added, need to handle the case where a non speculator
-            # config is passed in as a kwarg and auto convert
             raise TypeError(
                 f"Expected config to be an instance of SpeculatorModelConfig, "
                 f"got {type(config)}."