From 1f472cf5cd35cdc1d87d39c08ed5b5f78953c0d9 Mon Sep 17 00:00:00 2001
From: "Guan-Ming (Wesley) Chiu"
 <105915352+guan404ming@users.noreply.github.com>
Date: Thu, 18 Jun 2026 15:33:09 +0800
Subject: [PATCH 1/5] Add DFlash model converter

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
---
 src/speculators/convert/dflash/__init__.py  |   5 +
 src/speculators/convert/dflash/converter.py | 164 ++++++++++++++++++++
 src/speculators/convert/entrypoints.py      |  24 ++-
 tests/unit/convert/test_dflash_converter.py |  75 +++++++++
 4 files changed, 266 insertions(+), 2 deletions(-)
 create mode 100644 src/speculators/convert/dflash/__init__.py
 create mode 100644 src/speculators/convert/dflash/converter.py
 create mode 100644 tests/unit/convert/test_dflash_converter.py

diff --git a/src/speculators/convert/dflash/__init__.py b/src/speculators/convert/dflash/__init__.py
new file mode 100644
index 000000000..482646820
--- /dev/null
+++ b/src/speculators/convert/dflash/__init__.py
@@ -0,0 +1,5 @@
+"""DFlash checkpoint conversion utilities."""
+
+from speculators.convert.dflash.converter import DFlashConverter
+
+__all__ = ["DFlashConverter"]
diff --git a/src/speculators/convert/dflash/converter.py b/src/speculators/convert/dflash/converter.py
new file mode 100644
index 000000000..f424c0bd7
--- /dev/null
+++ b/src/speculators/convert/dflash/converter.py
@@ -0,0 +1,164 @@
+"""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"}
+)
+
+
+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.startswith("layers.")]
+        if critical_missing:
+            raise ValueError(
+                f"Draft layer 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")
diff --git a/src/speculators/convert/entrypoints.py b/src/speculators/convert/entrypoints.py
index 80c5357fc..3a961705a 100644
--- a/src/speculators/convert/entrypoints.py
+++ b/src/speculators/convert/entrypoints.py
@@ -9,6 +9,7 @@
 - EAGLE3
 - HASS
 - MTP
+- DFlash
 
 Functions:
     convert_model: Converts a model checkpoint to the Speculators format.
@@ -16,6 +17,7 @@
 
 from typing import Literal
 
+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
@@ -26,7 +28,7 @@
 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 +84,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 +131,13 @@ 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}")
diff --git a/tests/unit/convert/test_dflash_converter.py b/tests/unit/convert/test_dflash_converter.py
new file mode 100644
index 000000000..7992870dd
--- /dev/null
+++ b/tests/unit/convert/test_dflash_converter.py
@@ -0,0 +1,75 @@
+"""Unit tests for DFlashConverter config building."""
+
+from unittest.mock import patch
+
+import pytest
+
+from speculators.convert.dflash.converter import DFlashConverter
+
+
+def _source_config(**overrides):
+    config = {
+        "model_type": "qwen3",
+        "architectures": ["DFlashDraftModel"],
+        "auto_map": {"AutoModel": "dflash.DFlashDraftModel"},
+        "block_size": 16,
+        "num_target_layers": 36,
+        "dflash_config": {
+            "mask_token_id": 151669,
+            "target_layer_ids": [1, 9, 17, 25, 33],
+        },
+        "vocab_size": 151936,
+        "hidden_size": 4096,
+        "intermediate_size": 12288,
+        "num_hidden_layers": 5,
+        "num_attention_heads": 32,
+        "num_key_value_heads": 8,
+        "head_dim": 128,
+    }
+    config.update(overrides)
+    return config
+
+
+class TestBuildConfig:
+    @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict")
+    def test_happy_path(self, mock_get_config):
+        mock_get_config.return_value = (
+            {"hidden_size": 4096, "architectures": ["Qwen3ForCausalLM"]},
+            None,
+        )
+        config = DFlashConverter()._build_config(
+            _source_config(), "Qwen/Qwen3-8B", None
+        )
+
+        assert config.speculators_config.algorithm == "dflash"
+        assert config.speculators_config.verifier.name_or_path == "Qwen/Qwen3-8B"
+        assert config.block_size == 16
+        assert config.draft_vocab_size == 151936
+        assert config.mask_token_id == 151669
+        assert config.speculators_config.proposal_methods[0].speculative_tokens == 15
+        # z-lab target_layer_ids are offset by +1 to speculators layer ids
+        assert config.aux_hidden_state_layer_ids == [2, 10, 18, 26, 34]
+        # non-transformer keys are stripped from transformer_layer_config
+        assert config.transformer_layer_config.num_hidden_layers == 5
+        assert not hasattr(config.transformer_layer_config, "dflash_config")
+
+    @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict")
+    def test_explicit_aux_layer_ids_override(self, mock_get_config):
+        mock_get_config.return_value = ({"hidden_size": 4096}, None)
+        config = DFlashConverter()._build_config(
+            _source_config(), "Qwen/Qwen3-8B", [3, 11, 19]
+        )
+        assert config.aux_hidden_state_layer_ids == [3, 11, 19]
+
+    @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict")
+    def test_hidden_size_mismatch_raises(self, mock_get_config):
+        mock_get_config.return_value = ({"hidden_size": 2048}, None)
+        with pytest.raises(ValueError, match="Architecture mismatch"):
+            DFlashConverter()._build_config(_source_config(), "some/model", None)
+
+    @patch("speculators.convert.dflash.converter.PretrainedConfig.get_config_dict")
+    def test_missing_target_layer_ids_raises(self, mock_get_config):
+        mock_get_config.return_value = ({"hidden_size": 4096}, None)
+        source = _source_config(dflash_config={"mask_token_id": 151669})
+        with pytest.raises(ValueError, match="target_layer_ids"):
+            DFlashConverter()._build_config(source, "Qwen/Qwen3-8B", None)

From 0e3707339348248c4c47688afde2e05c54b72a05 Mon Sep 17 00:00:00 2001
From: "Guan-Ming (Wesley) Chiu"
 <105915352+guan404ming@users.noreply.github.com>
Date: Thu, 18 Jun 2026 17:11:41 +0800
Subject: [PATCH 2/5] Flag any missing draft weight in DFlash converter

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
---
 src/speculators/convert/dflash/converter.py | 19 ++++++++---
 tests/unit/convert/test_dflash_converter.py | 38 +++++++++++++++++++++
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/src/speculators/convert/dflash/converter.py b/src/speculators/convert/dflash/converter.py
index f424c0bd7..9bca42ea7 100644
--- a/src/speculators/convert/dflash/converter.py
+++ b/src/speculators/convert/dflash/converter.py
@@ -32,6 +32,19 @@
     {"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.
@@ -136,11 +149,9 @@ def _save(
                 "Unexpected keys in checkpoint -- the structure does not match "
                 f"DFlashDraftModel. Unexpected keys: {unexpected}"
             )
-        critical_missing = [k for k in missing if k.startswith("layers.")]
+        critical_missing = [k for k in missing if k not in _VERIFIER_FILLED_KEYS]
         if critical_missing:
-            raise ValueError(
-                f"Draft layer weights missing after load: {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
diff --git a/tests/unit/convert/test_dflash_converter.py b/tests/unit/convert/test_dflash_converter.py
index 7992870dd..b898d7660 100644
--- a/tests/unit/convert/test_dflash_converter.py
+++ b/tests/unit/convert/test_dflash_converter.py
@@ -3,8 +3,37 @@
 from unittest.mock import patch
 
 import pytest
+from transformers import Qwen3Config
 
+from speculators.config import SpeculatorsConfig, VerifierConfig
 from speculators.convert.dflash.converter import DFlashConverter
+from speculators.models.dflash import DFlashSpeculatorConfig
+from speculators.proposals.greedy import GreedyTokenProposalConfig
+
+
+def _tiny_dflash_config():
+    return DFlashSpeculatorConfig(
+        transformer_layer_config=Qwen3Config(
+            vocab_size=32,
+            hidden_size=16,
+            intermediate_size=32,
+            num_hidden_layers=1,
+            num_attention_heads=2,
+            num_key_value_heads=1,
+            head_dim=8,
+            max_position_embeddings=32,
+        ),
+        draft_vocab_size=32,
+        block_size=4,
+        aux_hidden_state_layer_ids=[0],
+        mask_token_id=1,
+        speculators_config=SpeculatorsConfig(
+            algorithm="dflash",
+            proposal_methods=[GreedyTokenProposalConfig(speculative_tokens=3)],
+            default_proposal_method="greedy",
+            verifier=VerifierConfig(name_or_path="dummy", architectures=[]),
+        ),
+    )
 
 
 def _source_config(**overrides):
@@ -73,3 +102,12 @@ def test_missing_target_layer_ids_raises(self, mock_get_config):
         source = _source_config(dflash_config={"mask_token_id": 151669})
         with pytest.raises(ValueError, match="target_layer_ids"):
             DFlashConverter()._build_config(source, "Qwen/Qwen3-8B", None)
+
+
+class TestSave:
+    def test_missing_draft_weights_raise(self, tmp_path):
+        # No source weights: every draft-body weight (fc, norm, hidden_norm,
+        # layers.*) is missing and must be flagged, not silently kept as NaN.
+        # Raises before load_verifier_weights, so no verifier download.
+        with pytest.raises(ValueError, match="Draft weights missing"):
+            DFlashConverter()._save(_tiny_dflash_config(), {}, tmp_path)

From e2dbd647d8354a50304b4e615f70b875d23faf43 Mon Sep 17 00:00:00 2001
From: "Guan-Ming (Wesley) Chiu"
 <105915352+guan404ming@users.noreply.github.com>
Date: Thu, 18 Jun 2026 22:46:52 +0800
Subject: [PATCH 3/5] Auto-convert external checkpoints in from_pretrained

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
---
 scripts/train.py                       |  7 +++-
 src/speculators/convert/entrypoints.py | 50 +++++++++++++++++++++++++-
 src/speculators/model.py               | 14 ++++++--
 tests/unit/convert/test_entrypoints.py | 42 ++++++++++++++++++++++
 4 files changed, 109 insertions(+), 4 deletions(-)
 create mode 100644 tests/unit/convert/test_entrypoints.py

diff --git a/scripts/train.py b/scripts/train.py
index c6dfbf1b1..e6a7a40e2 100644
--- a/scripts/train.py
+++ b/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
diff --git a/src/speculators/convert/entrypoints.py b/src/speculators/convert/entrypoints.py
index 3a961705a..97350b81d 100644
--- a/src/speculators/convert/entrypoints.py
+++ b/src/speculators/convert/entrypoints.py
@@ -15,14 +15,18 @@
     convert_model: Converts a model checkpoint to the Speculators format.
 """
 
+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(
@@ -141,3 +145,47 @@ def convert_model(
         )
     else:
         raise ValueError(f"Unsupported algorithm: {algorithm}")
+
+
+def maybe_convert_external_checkpoint(
+    model: str,
+    verifier: str | None = None,
+    cache_dir: str | None = None,
+    output_path: str | 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.
+    """
+    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 []
+    if "dflash_config" in config_dict or any("DFlash" in a for a in architectures):
+        algorithm = "dflash"
+    else:
+        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
diff --git a/src/speculators/model.py b/src/speculators/model.py
index d711ca3c3..20a728317 100644
--- a/src/speculators/model.py
+++ b/src/speculators/model.py
@@ -242,6 +242,7 @@ def from_pretrained(
         revision: str = "main",
         use_safetensors: bool | None = None,
         weights_only: bool = True,
+        verifier: str | None = None,
         t2d: torch.Tensor | None = None,
         d2t: torch.Tensor | None = None,
         **kwargs,
@@ -293,6 +294,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 +307,15 @@ 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.
+            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=verifier, cache_dir=cache_dir
+            )
             config = cls.config_class.from_pretrained(
                 pretrained_model_name_or_path,
                 cache_dir=cache_dir,
@@ -314,8 +326,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)}."
diff --git a/tests/unit/convert/test_entrypoints.py b/tests/unit/convert/test_entrypoints.py
new file mode 100644
index 000000000..f044b5b5e
--- /dev/null
+++ b/tests/unit/convert/test_entrypoints.py
@@ -0,0 +1,42 @@
+"""Unit tests for the external-checkpoint auto-conversion dispatch."""
+
+from unittest.mock import patch
+
+import pytest
+
+from speculators.convert.entrypoints import maybe_convert_external_checkpoint
+
+
+class TestMaybeConvertExternalCheckpoint:
+    @patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
+    def test_speculators_checkpoint_passes_through(self, mock_cfg):
+        mock_cfg.return_value = ({"speculators_model_type": "dflash"}, None)
+        assert maybe_convert_external_checkpoint("some/speculators-model") == (
+            "some/speculators-model"
+        )
+
+    @patch("speculators.convert.entrypoints.convert_model")
+    @patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
+    def test_external_dflash_converts(self, mock_cfg, mock_convert):
+        mock_cfg.return_value = ({"dflash_config": {}, "model_type": "qwen3"}, None)
+        out = maybe_convert_external_checkpoint(
+            "z-lab/Qwen3-8B-DFlash-b16",
+            verifier="Qwen/Qwen3-8B",
+            output_path="/tmp/out",
+        )
+        assert out == "/tmp/out"
+        mock_convert.assert_called_once()
+        assert mock_convert.call_args.kwargs["algorithm"] == "dflash"
+        assert mock_convert.call_args.kwargs["verifier"] == "Qwen/Qwen3-8B"
+
+    @patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
+    def test_external_without_verifier_raises(self, mock_cfg):
+        mock_cfg.return_value = ({"architectures": ["DFlashDraftModel"]}, None)
+        with pytest.raises(ValueError, match="requires a verifier"):
+            maybe_convert_external_checkpoint("z-lab/Qwen3-8B-DFlash-b16")
+
+    @patch("speculators.convert.entrypoints.PretrainedConfig.get_config_dict")
+    def test_unrecognized_format_raises(self, mock_cfg):
+        mock_cfg.return_value = ({"model_type": "qwen3"}, None)
+        with pytest.raises(NotImplementedError, match="unrecognized external"):
+            maybe_convert_external_checkpoint("some/model", verifier="Qwen/Qwen3-8B")

From 4676344146a08deedf2a28252fa8218494bd6574 Mon Sep 17 00:00:00 2001
From: "Guan-Ming (Wesley) Chiu"
 <105915352+guan404ming@users.noreply.github.com>
Date: Fri, 19 Jun 2026 08:07:29 +0800
Subject: [PATCH 4/5] Fix ci error

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
---
 src/speculators/convert/entrypoints.py      | 12 +++++++-----
 src/speculators/model.py                    |  5 +++--
 tests/unit/convert/test_dflash_converter.py |  2 +-
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/speculators/convert/entrypoints.py b/src/speculators/convert/entrypoints.py
index 97350b81d..cf459e5bb 100644
--- a/src/speculators/convert/entrypoints.py
+++ b/src/speculators/convert/entrypoints.py
@@ -15,6 +15,7 @@
     convert_model: Converts a model checkpoint to the Speculators format.
 """
 
+import os
 import tempfile
 from typing import Literal
 
@@ -148,9 +149,9 @@ def convert_model(
 
 
 def maybe_convert_external_checkpoint(
-    model: str,
+    model: str | os.PathLike,
     verifier: str | None = None,
-    cache_dir: str | None = None,
+    cache_dir: str | os.PathLike | None = None,
     output_path: str | None = None,
 ) -> str:
     """Convert an external (non-speculators) checkpoint to speculators format.
@@ -165,9 +166,10 @@ def maybe_convert_external_checkpoint(
         return str(model)
 
     architectures = config_dict.get("architectures") or []
-    if "dflash_config" in config_dict or any("DFlash" in a for a in architectures):
-        algorithm = "dflash"
-    else:
+    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."
diff --git a/src/speculators/model.py b/src/speculators/model.py
index 20a728317..c90e4608c 100644
--- a/src/speculators/model.py
+++ b/src/speculators/model.py
@@ -242,7 +242,6 @@ def from_pretrained(
         revision: str = "main",
         use_safetensors: bool | None = None,
         weights_only: bool = True,
-        verifier: str | None = None,
         t2d: torch.Tensor | None = None,
         d2t: torch.Tensor | None = None,
         **kwargs,
@@ -314,7 +313,9 @@ def from_pretrained(
             )
 
             pretrained_model_name_or_path = maybe_convert_external_checkpoint(
-                pretrained_model_name_or_path, verifier=verifier, cache_dir=cache_dir
+                pretrained_model_name_or_path,
+                verifier=kwargs.get("verifier"),
+                cache_dir=cache_dir,
             )
             config = cls.config_class.from_pretrained(
                 pretrained_model_name_or_path,
diff --git a/tests/unit/convert/test_dflash_converter.py b/tests/unit/convert/test_dflash_converter.py
index b898d7660..959a56bb9 100644
--- a/tests/unit/convert/test_dflash_converter.py
+++ b/tests/unit/convert/test_dflash_converter.py
@@ -75,7 +75,7 @@ def test_happy_path(self, mock_get_config):
         assert config.block_size == 16
         assert config.draft_vocab_size == 151936
         assert config.mask_token_id == 151669
-        assert config.speculators_config.proposal_methods[0].speculative_tokens == 15
+        assert config.speculators_config.proposal_methods[0].speculative_tokens == 15  # type: ignore[attr-defined]
         # z-lab target_layer_ids are offset by +1 to speculators layer ids
         assert config.aux_hidden_state_layer_ids == [2, 10, 18, 26, 34]
         # non-transformer keys are stripped from transformer_layer_config

From 3496fd22ce801f49c0f78fce97c60516366942dd Mon Sep 17 00:00:00 2001
From: "Guan-Ming (Wesley) Chiu"
 <105915352+guan404ming@users.noreply.github.com>
Date: Sat, 20 Jun 2026 13:48:55 +0800
Subject: [PATCH 5/5] Refactor external checkpoint conversion logic in
 SpeculatorModel

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
---
 src/speculators/convert/entrypoints.py |  7 +++++--
 src/speculators/model.py               | 22 ++++++++++++++--------
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/src/speculators/convert/entrypoints.py b/src/speculators/convert/entrypoints.py
index cf459e5bb..3cefcf7ed 100644
--- a/src/speculators/convert/entrypoints.py
+++ b/src/speculators/convert/entrypoints.py
@@ -153,15 +153,18 @@ def maybe_convert_external_checkpoint(
     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.
+    the unified ``from_pretrained`` finetuning pathway. Pass ``config_dict`` to
+    reuse an already-loaded config and skip re-reading it.
     """
-    config_dict, _ = PretrainedConfig.get_config_dict(model, cache_dir=cache_dir)
+    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)
 
diff --git a/src/speculators/model.py b/src/speculators/model.py
index c90e4608c..fb837b1c0 100644
--- a/src/speculators/model.py
+++ b/src/speculators/model.py
@@ -307,16 +307,22 @@ def from_pretrained(
                     "provided to load a SpeculatorModel."
                 )
             # Auto-convert external (non-speculators) checkpoints so one
-            # `from_pretrained` pathway finetunes both formats.
-            from speculators.convert.entrypoints import (  # noqa: PLC0415
-                maybe_convert_external_checkpoint,
+            # `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,
-            )
+                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,