diff --git a/tests/test_initial_dflash_weights.py b/tests/test_initial_dflash_weights.py new file mode 100644 index 00000000..b340b6ae --- /dev/null +++ b/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 + ) diff --git a/tests/test_initial_draft_upgrades.py b/tests/test_initial_draft_upgrades.py new file mode 100644 index 00000000..74d168d7 --- /dev/null +++ b/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) diff --git a/torchspec/training/checkpoint.py b/torchspec/training/checkpoint.py index 76c16de9..38f8b804 100644 --- a/torchspec/training/checkpoint.py +++ b/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. diff --git a/torchspec/training/dflash_trainer.py b/torchspec/training/dflash_trainer.py index 82b5125c..4982a6bd 100644 --- a/torchspec/training/dflash_trainer.py +++ b/torchspec/training/dflash_trainer.py @@ -134,6 +134,9 @@ def init_model( draft_model = self._build_draft_model(config) if dist.get_rank() == 0: + initial_path = getattr(self.args, "initial_draft_model_path", None) + if initial_path: + checkpoint.load_initial_dflash_weights(draft_model, initial_path) draft_model.load_embedding( target_model_path, embedding_key=getattr(self.args, "embedding_key", "model.embed_tokens.weight"),