feat(specdec): warm-start DFlash fine-tuning from exported drafter via dflash_init_checkpoint

examples/speculative_decoding/doc/dflash.md

@@ -281,6 +281,30 @@ DFlash supports checkpoint resume transparently. Rotary embeddings are lazily
 initialized on first forward (matching EAGLE3's `_maybe_init_rope` pattern),
 avoiding meta-tensor issues during `from_pretrained` model construction.
 
+### Warm Start / Fine-Tuning (`dflash_init_checkpoint`)
+
+To fine-tune an already-trained drafter instead of training from scratch, set
+`dflash.dflash_init_checkpoint` to a local directory containing an exported
+DFlash checkpoint (the z-lab-compatible `model.safetensors` layout produced by
+`export_hf_checkpoint.py`; for a public drafter, download it first, e.g.
+`hf download z-lab/Qwen3-8B-DFlash-b16`):
+
+```yaml
+dflash:
+  dflash_init_checkpoint: /path/to/Qwen3-8B-DFlash-b16
+```
+
+The weights are loaded into `model.dflash_module` right after conversion.
+`dflash_architecture_config` must describe the same draft architecture as the
+checkpoint — note that draft head/MLP dims default to `Qwen3Config` values, not
+the base model's, so set `num_attention_heads` / `num_key_value_heads` /
+`head_dim` / `intermediate_size` explicitly. `dflash_block_size` may differ from
+the checkpoint's (weights are block-size agnostic); a `mask_token_id` or
+`target_layer_ids` mismatch loads but logs a degradation warning. Restore of a
+saved training checkpoint ignores the field. See
+`tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml` for a
+full fine-tuning pipeline.
+
 ### Export
 
 ```bash

modelopt/torch/speculative/config.py

@@ -144,6 +144,19 @@ class DFlashConfig(ModeloptBaseConfig):
         default={}, description="Config for the DFlash draft module architecture."
     )
 
+    dflash_init_checkpoint: str | None = ModeloptField(
+        default=None,
+        description=(
+            "Warm-start the draft module from an exported DFlash checkpoint before "
+            "fine-tuning. Must be a local directory containing ``model.safetensors`` in the "
+            "DFlashExporter / z-lab layout (e.g. a downloaded copy of "
+            "``z-lab/Qwen3-8B-DFlash-b16``). ``dflash_architecture_config`` must describe "
+            "the same draft architecture as the checkpoint (weights are block-size "
+            "agnostic, so ``dflash_block_size`` may differ). Ignored on restore — restored "
+            "models load their own trained weights."
+        ),
+    )
+
     dflash_use_torch_compile: bool = ModeloptField(
         default=True,
         description="Whether to use torch.compile on DFlash forward/loss methods.",

modelopt/torch/speculative/dflash/conversion.py

@@ -80,4 +80,7 @@ def restore_dflash_model(
 ) -> nn.Module:
     """Function for restoring a previously converted model to a DFlash model."""
     assert not metadata, "No metadata expected!"
+    # Never warm-start on restore: the restored model loads its own trained weights,
+    # and the init checkpoint may no longer exist where the model was trained.
+    config.dflash_init_checkpoint = None
     return convert_to_dflash_model(model, config)[0]

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -71,7 +71,9 @@
            lazy rope pattern needed for MLA models.
 """
 
+import json
 import logging
+import os
 
 import torch
 import torch.nn.functional as F
@@ -306,6 +308,12 @@ def modify(self, config):
         if base_device.type != "meta":
             self.dflash_module.to(self._base_model.dtype).to(base_device)
 
+        # Warm-start the draft module from an exported DFlash checkpoint (fine-tuning).
+        # Restore paths never reach here with a checkpoint set: restore_dflash_model
+        # clears the field since restored weights come from the restored model itself.
+        if config.dflash_init_checkpoint:
+            self._load_draft_init_weights(config.dflash_init_checkpoint)
+
         # Delete base model layers for offline training (save memory)
         if self.dflash_offline:
             self._base_model._modules.pop("layers")
@@ -317,6 +325,91 @@ def _build_draft_module(self, dflash_config):
         """Build the draft module. Subclasses override to use an augmented module."""
         return DFlashModule(dflash_config)
 
+    def _load_draft_init_weights(self, checkpoint: str):
+        """Warm-start ``self.dflash_module`` from an exported DFlash draft checkpoint.
+
+        ``checkpoint`` is a local directory holding ``model.safetensors`` (DFlashExporter /
+        z-lab layout). Keys map 1:1 onto ``dflash_module``. Missing keys are tolerated
+        only for submodules entirely absent from the checkpoint (e.g. the Domino/DSpark
+        heads when warm-starting from a plain DFlash backbone); partial overlap within a
+        submodule means an architecture mismatch and raises.
+        """
+        from safetensors.torch import load_file
+
+        weights_path = os.path.join(checkpoint, "model.safetensors")
+        if not os.path.isfile(weights_path):
+            raise ValueError(
+                f"dflash_init_checkpoint must be a local directory containing "
+                f"model.safetensors; not found at {weights_path}."
+            )
+
+        self._check_init_checkpoint_config(checkpoint, weights_path)
+
+        param = next(self.dflash_module.parameters())
+        state_dict = {
+            key: value.to(param.dtype)
+            for key, value in load_file(weights_path, device=str(param.device)).items()
+        }
+        arch_hint = (
+            "Set dflash_architecture_config to the checkpoint's architecture "
+            "(num_hidden_layers, heads, intermediate_size, ...)."
+        )
+        try:
+            # strict=False tolerates missing/unexpected keys but still raises on shape
+            # mismatches (e.g. a different draft depth changing the fc fusion width).
+            missing, unexpected = self.dflash_module.load_state_dict(state_dict, strict=False)
+        except RuntimeError as e:
+            raise ValueError(
+                f"DFlash init checkpoint {checkpoint} does not match the draft architecture "
+                f"built from dflash_architecture_config: {e} {arch_hint}"
+            ) from e
+        ckpt_submodules = {key.split(".", 1)[0] for key in state_dict}
+        hard_missing = [key for key in missing if key.split(".", 1)[0] in ckpt_submodules]
+        if unexpected or hard_missing:
+            raise ValueError(
+                f"DFlash init checkpoint {checkpoint} does not match the draft architecture "
+                f"built from dflash_architecture_config (unexpected keys: {unexpected}, "
+                f"missing keys: {hard_missing}). {arch_hint}"
+            )
+        if missing:
+            logger.warning(
+                "DFlash: submodules not in init checkpoint %s train from scratch: %s",
+                checkpoint,
+                sorted({key.split(".", 1)[0] for key in missing}),
+            )
+        logger.info(
+            "DFlash: warm-started draft module from %s (%d tensors).",
+            checkpoint,
+            len(state_dict),
+        )
+
+    def _check_init_checkpoint_config(self, checkpoint: str, weights_path: str):
+        """Warn when the init checkpoint's config disagrees with this conversion.
+
+        A ``mask_token_id`` or ``target_layer_ids`` mismatch loads cleanly (shapes match)
+        but silently degrades the warm start: the mask embedding and the fc fusion columns
+        were trained for the checkpoint's values.
+        """
+        config_path = os.path.join(os.path.dirname(weights_path), "config.json")
+        if not os.path.isfile(config_path):
+            return
+        with open(config_path) as f:
+            ckpt_dflash_config = json.load(f).get("dflash_config", {})
+        for name, ours in (
+            ("mask_token_id", self.mask_token_id),
+            ("target_layer_ids", list(self.target_layer_ids)),
+        ):
+            ckpt_value = ckpt_dflash_config.get(name)
+            if ckpt_value is not None and ckpt_value != ours:
+                logger.warning(
+                    "DFlash: init checkpoint %s was trained with %s=%s but this conversion "
+                    "uses %s — the warm start will be degraded.",
+                    checkpoint,
+                    name,
+                    ckpt_value,
+                    ours,
+                )
+
     def get_exporter(self):
         """Get the exporter for the DFlash draft model."""
         from modelopt.torch.export.plugins.hf_spec_export import DFlashExporter

tests/unit/torch/speculative/plugins/test_hf_dflash.py

@@ -279,6 +279,69 @@ def test_save_and_restore(self, tmp_path):
         tf_modelopt_state_and_output_tester(model_ref, model_test)
 
 
+class TestDFlashWarmStart:
+    """Test warm-starting the draft module from an exported DFlash checkpoint."""
+
+    def _export_drafter(self, tmp_path, **config_overrides):
+        """Convert a tiny model and export its drafter in z-lab layout."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dflash_config()
+        config.update(config_overrides)
+        mtsp.convert(model, [("dflash", config)])
+        export_dir = tmp_path / "exported_drafter"
+        model.get_exporter().export(export_dir)
+        return model, export_dir
+
+    def test_warm_start_loads_exported_weights(self, tmp_path):
+        """A fresh conversion with dflash_init_checkpoint matches the exported drafter.
+
+        Fine-tunes at a different block size: weights are block-size agnostic.
+        """
+        model_ref, export_dir = self._export_drafter(tmp_path)
+
+        model = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dflash_config(block_size=2 * BLOCK_SIZE)
+        config["dflash_init_checkpoint"] = str(export_dir)
+        mtsp.convert(model, [("dflash", config)])
+
+        ref_sd = model_ref.dflash_module.state_dict()
+        for key, value in model.dflash_module.state_dict().items():
+            assert torch.equal(value, ref_sd[key]), f"Mismatch after warm start: {key}"
+
+    def test_warm_start_bad_checkpoint_raises(self, tmp_path):
+        """A missing directory or a different draft architecture fails loudly."""
+        _, export_dir = self._export_drafter(tmp_path)
+
+        config = _get_dflash_config()
+        config["dflash_init_checkpoint"] = str(tmp_path / "no_such_dir")
+        with pytest.raises(ValueError, match="must be a local directory"):
+            mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)])
+
+        config = _get_dflash_config(num_layers=NUM_DRAFT_LAYERS + 1)
+        config["dflash_init_checkpoint"] = str(export_dir)
+        with pytest.raises(ValueError, match="does not match the draft architecture"):
+            mtsp.convert(get_tiny_llama(num_hidden_layers=4), [("dflash", config)])
+
+    def test_restore_ignores_init_checkpoint(self, tmp_path):
+        """Save/restore of a warm-started model must not re-read the init checkpoint."""
+        import shutil
+
+        _, export_dir = self._export_drafter(tmp_path)
+
+        mto.enable_huggingface_checkpointing()
+        model_ref = get_tiny_llama(num_hidden_layers=4)
+        config = _get_dflash_config()
+        config["dflash_init_checkpoint"] = str(export_dir)
+        mtsp.convert(model_ref, [("dflash", config)])
+        model_ref.save_pretrained(tmp_path / "warm_started_model")
+
+        # The init checkpoint is gone; restore must still work.
+        shutil.rmtree(export_dir)
+        model_test = AutoModelForCausalLM.from_pretrained(tmp_path / "warm_started_model")
+        assert isinstance(model_test, HFDFlashModel)
+        tf_modelopt_state_and_output_tester(model_ref, model_test)
+
+
 class TestDFlashLazyRotaryEmb:
     """Test lazy rotary embedding initialization (matching EAGLE3 pattern).
 

tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml

@@ -0,0 +1,107 @@
+# DFlash online FINE-TUNING for Qwen3-8B, warm-started from an exported DFlash checkpoint.
+#
+# Instead of training the drafter from scratch, dflash.dflash_init_checkpoint loads an
+# already-trained draft module (DFlashExporter / z-lab layout: model.safetensors +
+# config.json) into model.dflash_module before training. It must be a LOCAL directory —
+# stage the public drafter first (hf download z-lab/Qwen3-8B-DFlash-b16) or point
+# init_checkpoint at your own export directory.
+#
+# NOTE: dflash_architecture_config must describe the SAME draft architecture as the
+# checkpoint. modify() takes num_attention_heads/num_key_value_heads/head_dim/
+# intermediate_size from Qwen3Config defaults (NOT the base model), so set them
+# explicitly below; hidden_size/vocab/rope_theta are forced to the base model and need
+# not be set. A mismatch fails at conversion with a shape error. dflash_block_size may
+# differ from the checkpoint's (weights are block-size agnostic), but mask_token_id and
+# target_layer_ids should match — the loader warns if they don't.
+#
+# 3-step pipeline:
+#   task_0: Online DFlash fine-tuning (warm start)
+#   task_1: vLLM smoke test with DFlash speculative decoding
+#   task_2: MT-Bench per-category HF AR evaluation (1 GPU)
+#
+# The warm start is visible in the very first training logs: step 100 accuracy should
+# already be near the converged from-scratch level (~0.2, see hf_online_dflash.yaml)
+# instead of ~0.03.
+#
+# Regression criteria (set via environment):
+#   MAX_FINAL_LOSS: final loss must be below this (default: 4.5)
+#   MIN_FINAL_ACC:  final accuracy must be above this (default: 0.18)
+#
+# Usage:
+#   uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml --yes
+#   uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash_finetune.yaml --yes
+
+job_name: Qwen3-8B_DFlash_online_finetune
+pipeline:
+  global_vars:
+    hf_model: /hf-local/Qwen/Qwen3-8B
+    # Local dir with model.safetensors (a staged z-lab/Qwen3-8B-DFlash-b16 or your own export).
+    init_checkpoint: /hf-local/z-lab/Qwen3-8B-DFlash-b16
+
+  # Step 1: Online DFlash fine-tuning from the warm-start checkpoint
+  task_0:
+    script: common/specdec/dflash_online_training.sh
+    args:
+      - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml
+      - model.model_name_or_path=<<global_vars.hf_model>>
+      - data.data_path=/hf-local/modelopt/Speculative-Decoding-Dataset-v1-Qwen3-8B/sample-100K-openai.jsonl
+      - data.chat_template=examples/Qwen/Qwen3-8B/chat_template_train.jinja
+      - training.output_dir=/scratchspace/dflash_bs16_finetune
+      - training.per_device_train_batch_size=1
+      - training.num_train_epochs=1
+      - training.training_seq_len=4096
+      - training.save_steps=5000
+      - training.logging_steps=100
+      - training.disable_tqdm=true
+      - training.answer_only_loss=true
+      # Fine-tuning: lower LR than the 6e-4 from-scratch default.
+      - training.learning_rate=1.0e-4
+      - dflash.dflash_init_checkpoint=<<global_vars.init_checkpoint>>
+      # Match z-lab/Qwen3-8B-DFlash-b16: block_size 16, mask token 151669,
+      # 5 draft layers with Qwen3-8B dims.
+      - dflash.dflash_block_size=16
+      - dflash.dflash_num_anchors=512
+      - dflash.dflash_loss_decay_factor=7
+      - dflash.dflash_mask_token_id=151669
+      - dflash.dflash_architecture_config.num_hidden_layers=5
+      - dflash.dflash_architecture_config.num_attention_heads=32
+      - dflash.dflash_architecture_config.num_key_value_heads=8
+      - dflash.dflash_architecture_config.head_dim=128
+      - dflash.dflash_architecture_config.intermediate_size=12288
+    environment:
+      - MAX_FINAL_LOSS: "4.5"
+      - MIN_FINAL_ACC: "0.18"
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 1
+      ntasks_per_node: 1
+      gpus_per_node: 8
+
+  # Step 2: vLLM smoke test (uses exported checkpoint from training)
+  task_1:
+    script: common/specdec/vllm_smoke_test.sh
+    environment:
+      - HF_MODEL_CKPT: <<global_vars.hf_model>>
+      - DRAFT_CKPT_DIR: /scratchspace/dflash_bs16_finetune
+      - SPEC_METHOD: "dflash"
+      - NUM_SPEC_TOKENS: "7"
+      - MIN_ACCEPTANCE_LENGTH: "1.4"
+    slurm_config:
+      _factory_: "slurm_factory"
+      container: "vllm/vllm-openai:nightly"
+      nodes: 1
+      ntasks_per_node: 1
+      gpus_per_node: 1
+
+  # Step 3: HF AR evaluation
+  task_2:
+    script: common/specdec/ar_eval_mtbench.sh
+    args:
+      - --ckpt_dir /scratchspace/dflash_bs16_finetune
+      - --osl 512
+      - --steps 15
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 1
+      ntasks_per_node: 1
+      gpus_per_node: 1