fix: support pretrained initialization for DFlash variants and EAGLE3.1

tests/test_initial_dflash_weights.py

@@ -0,0 +1,104 @@
+"""Published DFlash backbone initialization for DFlash2 finetuning."""
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from torchspec.models.draft.dflash2 import DFlash2Config, DFlash2DraftModel
+from torchspec.models.draft.keymap import to_export_keys
+from torchspec.training.checkpoint import load_initial_dflash_weights
+
+
+def model():
+    return DFlash2DraftModel(
+        DFlash2Config(
+            hidden_size=16,
+            intermediate_size=32,
+            num_hidden_layers=1,
+            num_attention_heads=2,
+            num_key_value_heads=1,
+            vocab_size=32,
+            max_position_embeddings=64,
+            num_target_layers=4,
+            target_layer_ids=[1, 3],
+            block_size=16,
+            conv_group_size=4,
+            selector_rank=4,
+            selector_top_k=4,
+            mask_token_id=31,
+        )
+    )
+
+
+def plain_weights(draft):
+    return to_export_keys(
+        {
+            k: v.clone() + 0.125
+            for k, v in draft.state_dict().items()
+            if k != "embed_tokens.weight"
+            and ".attention_conv." not in k
+            and ".mlp_conv." not in k
+            and not k.startswith("candidate_selector.")
+        }
+    )
+
+
+def test_shared_weights_loaded_new_modules_retained(tmp_path):
+    draft = model()
+    before = {k: v.clone() for k, v in draft.state_dict().items()}
+    weights = plain_weights(draft)
+    path = tmp_path / "model.safetensors"
+    save_file(weights, path)
+    load_initial_dflash_weights(draft, str(path))
+    after = draft.state_dict()
+    for key in before:
+        if (
+            key == "embed_tokens.weight"
+            or ".attention_conv." in key
+            or ".mlp_conv." in key
+            or key.startswith("candidate_selector.")
+        ):
+            assert torch.equal(after[key], before[key])
+        else:
+            assert torch.equal(after[key], before[key] + 0.125)
+    conv = draft.layers[0].attention_conv
+    x = torch.randn(1, 16, 16)
+    prepared, kernel = conv.prepare(x)
+    assert torch.equal(prepared, x)
+    assert torch.equal(conv.finish(x, kernel), x)
+
+
+@pytest.mark.parametrize(
+    "fault", ["missing_backbone", "bad_shape", "unexpected", "partial_dflash2"]
+)
+def test_reject_incomplete_or_incompatible_weights(tmp_path, fault):
+    draft = model()
+    weights = plain_weights(draft)
+    if fault == "missing_backbone":
+        del weights["fc.weight"]
+    elif fault == "bad_shape":
+        weights["fc.weight"] = weights["fc.weight"][:1]
+    elif fault == "unexpected":
+        weights["unknown.weight"] = torch.ones(1)
+    else:
+        weights["candidate_selector.hidden_projection.weight"] = (
+            draft.candidate_selector.hidden_projection.weight.detach().clone()
+        )
+    path = tmp_path / "model.safetensors"
+    save_file(weights, path)
+    with pytest.raises(ValueError):
+        load_initial_dflash_weights(draft, str(path))
+
+
+def test_complete_dflash2_checkpoint_loads_selector(tmp_path):
+    draft = model()
+    weights = to_export_keys(
+        {k: v.clone() for k, v in draft.state_dict().items() if k != "embed_tokens.weight"}
+    )
+    path = tmp_path / "model.safetensors"
+    save_file(weights, path)
+    other = model()
+    load_initial_dflash_weights(other, str(path))
+    assert torch.equal(
+        other.candidate_selector.predecessor_codebook, draft.candidate_selector.predecessor_codebook
+    )

tests/test_initial_draft_upgrades.py

@@ -0,0 +1,177 @@
+"""Strict weight migration from DFlash to DSpark and EAGLE3 to EAGLE3.1."""
+
+import pytest
+import torch
+from safetensors.torch import save_file
+from transformers import LlamaConfig
+from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config
+
+from torchspec.models.draft.deepseek_eagle import Eagle3DeepseekV2ForCausalLM
+from torchspec.models.draft.dflash import DFlashConfig, DFlashDraftModel
+from torchspec.models.draft.dspark import DSparkConfig, DSparkDraftModel
+from torchspec.models.draft.keymap import to_export_keys
+from torchspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3
+from torchspec.training.checkpoint import load_initial_dflash_weights, load_initial_draft_weights
+
+FAMILIES = ["dspark", "eagle_gqa", "eagle_mla"]
+
+
+def build(family, upgraded, seed=0, *, markov=True, confidence=True):
+    torch.manual_seed(seed)
+    kwargs = dict(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        vocab_size=64,
+        max_position_embeddings=128,
+        rms_norm_eps=1e-6,
+        pad_token_id=0,
+    )
+    if family == "dspark":
+        kwargs.update(
+            num_target_layers=2, target_hidden_size=32, target_num_hidden_layers=8, mask_token_id=63
+        )
+        if not upgraded:
+            return DFlashDraftModel(DFlashConfig(**kwargs))
+        return DSparkDraftModel(
+            DSparkConfig(
+                **kwargs,
+                markov_rank=8 if markov else 0,
+                enable_confidence_head=confidence,
+                confidence_head_with_markov=markov,
+            )
+        )
+    kwargs.update(
+        draft_vocab_size=32,
+        fc_norm=upgraded,
+        norm_output=upgraded,
+        num_aux_hidden_states=3,
+        target_hidden_size=32,
+    )
+    if family == "eagle_gqa":
+        return LlamaForCausalLMEagle3(LlamaConfig(**kwargs))
+    kwargs.update(
+        q_lora_rank=16,
+        kv_lora_rank=16,
+        qk_nope_head_dim=8,
+        qk_rope_head_dim=4,
+        v_head_dim=8,
+        n_routed_experts=1,
+        n_shared_experts=0,
+        first_k_dense_replace=0,
+        num_experts_per_tok=1,
+    )
+    return Eagle3DeepseekV2ForCausalLM(DeepseekV3Config(**kwargs))
+
+
+def upgrade_key(key):
+    return key.startswith(("markov_head.", "confidence_head.", "fc_norm."))
+
+
+def load(model, path, family):
+    loader = load_initial_dflash_weights if family == "dspark" else load_initial_draft_weights
+    return loader(model, str(path))
+
+
+def publish(state, path, exported=True):
+    save_file(to_export_keys(state) if exported else state, path)
+
+
+@pytest.mark.parametrize("family", FAMILIES)
+@pytest.mark.parametrize("exported", [False, True])
+@pytest.mark.parametrize("include_embedding", [False, True])
+def test_base_checkpoint_preserves_shared_weights_and_new_initialization(
+    tmp_path, family, exported, include_embedding
+):
+    source = build(family, False)
+    # Give pruned-vocabulary buffers meaningful values to catch accidental resets.
+    if family != "dspark":
+        source.d2t.copy_(torch.arange(32))
+        source.t2d[:32] = False
+    target = build(family, True, seed=1)
+    before = {k: v.clone() for k, v in target.state_dict().items()}
+    state = {
+        k: v
+        for k, v in source.state_dict().items()
+        if include_embedding or k != "embed_tokens.weight"
+    }
+    path = tmp_path / "model.safetensors"
+    publish(state, path, exported)
+    assert load(target, path, family) == path
+    for k, v in target.state_dict().items():
+        assert torch.equal(v, state[k] if k in state else before[k]), k
+    if family != "dspark":
+        assert target.norm_output
+        assert all(torch.equal(n.weight, torch.ones_like(n.weight)) for n in target.fc_norm)
+
+
+@pytest.mark.parametrize("family", FAMILIES)
+def test_complete_upgraded_checkpoint_restores_all_new_weights(tmp_path, family):
+    source = build(family, True)
+    with torch.no_grad():
+        for k, v in source.state_dict().items():
+            if upgrade_key(k):
+                v.fill_(0.37)
+    target = build(family, True, seed=1)
+    state = {k: v for k, v in source.state_dict().items() if k != "embed_tokens.weight"}
+    path = tmp_path / "model.safetensors"
+    publish(state, path)
+    load(target, path, family)
+    for k, v in state.items():
+        assert torch.equal(target.state_dict()[k], v), k
+
+
+@pytest.mark.parametrize("family", FAMILIES)
+@pytest.mark.parametrize("fault", ["partial_upgrade", "missing_backbone", "shape", "unexpected"])
+def test_bad_checkpoint_is_rejected_without_modifying_model(tmp_path, family, fault):
+    source = build(family, False)
+    target = build(family, True, seed=1)
+    before = {k: v.clone() for k, v in target.state_dict().items()}
+    state = dict(source.state_dict())
+    backbone_key = next(k for k in state if k.endswith("weight") and k != "embed_tokens.weight")
+    if fault == "partial_upgrade":
+        k = next(k for k in before if upgrade_key(k))
+        state[k] = before[k]
+    elif fault == "missing_backbone":
+        del state[backbone_key]
+    elif fault == "shape":
+        state[backbone_key] = state[backbone_key][:1]
+    else:
+        state["unknown.weight"] = torch.ones(1)
+    path = tmp_path / "model.safetensors"
+    publish(state, path)
+    with pytest.raises((RuntimeError, ValueError)):
+        load(target, path, family)
+    assert all(torch.equal(target.state_dict()[k], v) for k, v in before.items())
+
+
+@pytest.mark.parametrize("markov,confidence", [(True, False), (False, True), (False, False)])
+def test_dspark_respects_disabled_heads(tmp_path, markov, confidence):
+    source = build("dspark", False)
+    target = build("dspark", True, seed=1, markov=markov, confidence=confidence)
+    path = tmp_path / "model.safetensors"
+    publish(source.state_dict(), path)
+    load(target, path, "dspark")
+    for k, v in source.state_dict().items():
+        assert torch.equal(target.state_dict()[k], v)
+
+
+@pytest.mark.parametrize("family", ["eagle_gqa", "eagle_mla"])
+@pytest.mark.parametrize("missing", ["lm_head.weight", "d2t", "t2d"])
+def test_eagle_upgrade_still_requires_head_and_vocabulary_mapping(tmp_path, family, missing):
+    state = dict(build(family, False).state_dict())
+    del state[missing]
+    path = tmp_path / "model.safetensors"
+    publish(state, path)
+    with pytest.raises(RuntimeError, match="missing"):
+        load(build(family, True), path, family)
+
+
+@pytest.mark.parametrize("family", FAMILIES)
+def test_upgraded_checkpoint_cannot_silently_downgrade(tmp_path, family):
+    path = tmp_path / "model.safetensors"
+    publish(build(family, True).state_dict(), path)
+    with pytest.raises((RuntimeError, ValueError), match="unexpected"):
+        load(build(family, False), path, family)

torchspec/training/checkpoint.py

@@ -155,24 +155,101 @@ def resolve_initial_draft_checkpoint(path: str) -> Path:
     return checkpoint_path
 
 
-def load_initial_draft_weights(draft_model: torch.nn.Module, path: str) -> Path:
-    """Load a published draft checkpoint into a freshly constructed draft model.
-
-    Unrelated to ``load()`` below, which resumes a training checkpoint with its optimizer and
-    scheduler state. Published drafts carry the serving key names ``tools/convert_to_hf.py``
-    writes, so keys are mapped back first, and the load is strict in both directions: a mismatch
-    would otherwise leave part of the model randomly initialised for a whole run.
-    """
+def _load_initial_weights(
+    draft_model: torch.nn.Module,
+    path: str,
+    *,
+    upgrade_keys: set[str],
+    allow_missing_embedding: bool,
+    error_type: type[Exception] = RuntimeError,
+) -> Path:
+    """Load shared weights strictly; new modules must be wholly absent or complete."""
     checkpoint_path = resolve_initial_draft_checkpoint(path)
+    expected = draft_model.state_dict()
     with safe_open(checkpoint_path, framework="pt", device="cpu") as f:
-        tensors = {key: f.get_tensor(key) for key in f.keys()}
+        tensors = to_internal_keys({key: f.get_tensor(key) for key in f.keys()}, expected)
 
-    draft_model.load_state_dict(
-        to_internal_keys(tensors, draft_model.state_dict().keys()), strict=True
+    present_upgrade = upgrade_keys.intersection(tensors)
+    if present_upgrade and present_upgrade != upgrade_keys:
+        raise error_type(
+            f"Incomplete {type(draft_model).__name__} upgrade weights: "
+            f"missing={sorted(upgrade_keys - present_upgrade)}"
+        )
+    allowed_missing = {"embed_tokens.weight"} if allow_missing_embedding else set()
+    if not present_upgrade:
+        allowed_missing |= upgrade_keys
+    missing = set(expected) - set(tensors) - allowed_missing
+    unexpected = set(tensors) - set(expected)
+    mismatched = {
+        key for key in tensors.keys() & expected.keys() if tensors[key].shape != expected[key].shape
+    }
+    if missing or unexpected or mismatched:
+        raise error_type(
+            f"Incompatible initial {type(draft_model).__name__} checkpoint: "
+            f"missing={sorted(missing)}, unexpected={sorted(unexpected)}, "
+            f"shape_mismatch={sorted(mismatched)}"
+        )
+    # Validate before modifying the model. Retain only explicitly allowed initialization.
+    retained = set(expected) - set(tensors)
+    draft_model.load_state_dict({**expected, **tensors}, strict=True)
+    logger.info(
+        "Loaded %d initial draft tensors from %s; retained initialization for %s",
+        len(tensors),
+        checkpoint_path,
+        sorted(retained),
     )
     return checkpoint_path
 
 
+def load_initial_dflash_weights(draft_model: torch.nn.Module, path: str) -> Path:
+    """Warm-start DFlash, DFlash2 or DSpark from a compatible published backbone.
+
+    DFlash2 convolutions/selector or DSpark Markov/confidence heads may be wholly
+    absent, retaining their constructor initialization. Partial additions and missing
+    backbone weights are errors. The embedding is loaded separately from the target.
+    """
+    from torchspec.models.draft.dflash2 import DFlash2DraftModel
+    from torchspec.models.draft.dspark import DSparkDraftModel
+
+    extra = set()
+    if isinstance(draft_model, DFlash2DraftModel):
+        extra = {
+            key
+            for key in draft_model.state_dict()
+            if ".attention_conv." in key
+            or ".mlp_conv." in key
+            or key.startswith("candidate_selector.")
+        }
+    elif isinstance(draft_model, DSparkDraftModel):
+        extra = {
+            key
+            for key in draft_model.state_dict()
+            if key.startswith(("markov_head.", "confidence_head."))
+        }
+    return _load_initial_weights(
+        draft_model, path, upgrade_keys=extra, allow_missing_embedding=True, error_type=ValueError
+    )
+
+
+def load_initial_draft_weights(draft_model: torch.nn.Module, path: str) -> Path:
+    """Load published draft weights, including a compatible EAGLE3 -> 3.1 warm start.
+
+    EAGLE's new per-input FC norms may be wholly absent; their unit-scale constructor
+    initialization is retained. Post-norm feedback is selected by config.norm_output
+    and adds no weights. This changes computation, so it is a training initialization,
+    not an output-preserving conversion. Head and vocabulary mappings remain strict.
+    The frozen EAGLE embedding may be absent because the trainer loads it from target.
+    This path does not restore optimizer or scheduler state.
+    """
+    from torchspec.models.draft.base import Eagle3DraftModel
+
+    is_eagle = isinstance(draft_model, Eagle3DraftModel)
+    extra = {key for key in draft_model.state_dict() if is_eagle and key.startswith("fc_norm.")}
+    return _load_initial_weights(
+        draft_model, path, upgrade_keys=extra, allow_missing_embedding=is_eagle
+    )
+
+
 def resolve_resume_model_dir(args: Any) -> Path | None:
     """Model directory ``training.load_path`` resolves to, or None if there is nothing to resume.