diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index dedef4a..c3373a8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -71,6 +71,9 @@ jobs: tests/test_eagle4_validation.py \ tests/test_fa4_import_compat.py \ tests/test_flex_attention_backend.py \ + tests/test_dmtp.py \ + tests/test_training_resume_rng.py \ + tests/test_evaluation_contract.py \ tests/test_open_source_entrypoints.py \ tests/test_release_audit.py \ tests/test_shipped_config_defaults.py \ diff --git a/aspec/commands/_recipe_smoke.py b/aspec/commands/_recipe_smoke.py index b8d012b..9e0f0c7 100644 --- a/aspec/commands/_recipe_smoke.py +++ b/aspec/commands/_recipe_smoke.py @@ -31,6 +31,19 @@ def audit_state(trainer): result = _base_audit_state(trainer) draft = getattr(trainer, "draft_model", None) + if getattr(getattr(draft, "config", None), "mtp_conditioning", None): + from aspec.commands._checkpoint_audit import tensor_hash + + result["mtp_conditioning"] = { + "policy": draft.config.mtp_conditioning, + "frozen": tensor_hash((n, p) for n, p in draft.named_parameters() + if n.startswith(("mtp_source.", "embed_tokens."))), + "fusion": tensor_hash(draft.mtp_fusion.named_parameters()) + if draft.mtp_fusion is not None else None, + } + expected = [p for p in draft.parameters() if p.requires_grad] + if [id(p) for p in trainer.optimizer.model_params] != [id(p) for p in expected]: + raise AssertionError("MTP conditioning optimizer ownership differs from the Draft") if getattr(getattr(draft, "config", None), "ffn_type", "dense") == "ling_moe": result["ling_expert_bias"] = [layer.mlp.gate.expert_bias.detach().cpu().tolist() for layer in draft.layers] @@ -117,6 +130,10 @@ def train_from_queue(self, step, num_batches): def save_model(self, step, force_sync=False): super().save_model(step, force_sync) current = audit_state(self._trainer) + if current.get("mtp_conditioning"): + initial = json.loads((self._audit_root / f"recipe-audit-initial{self._audit_suffix}.json").read_text()) + if current["mtp_conditioning"]["frozen"] != initial["mtp_conditioning"]["frozen"]: + raise AssertionError("MTP conditioning changed its frozen source/embedding") if current.get("eagle4", {}).get("policy", {}).get("stage") == "candidate_head": initial = json.loads((self._audit_root / f"recipe-audit-initial{self._audit_suffix}.json").read_text()) if initial["eagle4"]["body"] != current["eagle4"]["body"]: diff --git a/aspec/commands/convert_to_hf.py b/aspec/commands/convert_to_hf.py index a469064..585ec97 100644 --- a/aspec/commands/convert_to_hf.py +++ b/aspec/commands/convert_to_hf.py @@ -163,6 +163,9 @@ def _serving_config(raw_config: dict) -> dict: from aspec.config.draft_attention import resolve_draft_attention from aspec.config.draft_structure import resolve_draft_structure + if raw_config.get("mtp_conditioning"): + raise NotImplementedError("MTP-conditioned Drafts require native --format aspec export") + if raw_config.get("architectures") == ["NativeMTPDraftModel"]: native = json.loads(json.dumps(raw_config["native_config"])) native["architectures"] = ["Qwen3_5ForCausalLMMTP"] @@ -461,6 +464,10 @@ def convert_checkpoint( } if eagle4_asset is not None: manifest["eagle4"] = eagle4_asset + if raw_config.get("mtp_conditioning"): + from aspec.config.mtp_conditioning import mtp_conditioning_semantics + + manifest["mtp_conditioning"] = mtp_conditioning_semantics(raw_config) if artifact_format == "aspec": manifest["compute_dtype"] = str(model.embed_tokens.weight.dtype) manifest["tensor_dtypes"] = dict(Counter(str(t.dtype) for t in tensors.values())) @@ -473,7 +480,7 @@ def convert_checkpoint( if policy.get("dsa") is not None: manifest["draft_structure"]["attention"]["dsa"] = policy["dsa"] algorithm = training_contract["algorithm"] - block_proposal = algorithm in {"dspark", "dflash"} + block_proposal = algorithm in {"dspark", "dflash"} and not raw_config.get("mtp_conditioning") mla_proposal = raw_config.get("architectures") == ["Eagle3DeepseekV2ForCausalLM"] proposal = ("aspec.draft_artifact.MLAEagleConsumer" if mla_proposal else "aspec.draft_artifact.BlockDraftConsumer" if block_proposal else None) @@ -483,7 +490,7 @@ def convert_checkpoint( "validation": "required_per_artifact", "scope": ("asset and proposal; caller owns verifier features/cache/verification" if proposal else "asset loading and weights-only warmstart; no native proposal component"), - "borrowed_tensors": ["Target output head"] if block_proposal else [], + "borrowed_tensors": ["Target output head"] if algorithm in {"dspark", "dflash"} else [], } (output_dir / "export_manifest.json").write_text( json.dumps(manifest, indent=2, sort_keys=True) + "\n", diff --git a/aspec/config/dflash_qwen3_5_9b_dmtp_config.json b/aspec/config/dflash_qwen3_5_9b_dmtp_config.json new file mode 100644 index 0000000..3059932 --- /dev/null +++ b/aspec/config/dflash_qwen3_5_9b_dmtp_config.json @@ -0,0 +1,106 @@ +{ + "architectures": [ + "DFlashDraftModel" + ], + "model_type": "dflash", + "hidden_size": 4096, + "intermediate_size": 12288, + "head_dim": 128, + "num_hidden_layers": 5, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "vocab_size": 248320, + "rms_norm_eps": 1e-06, + "max_position_embeddings": 262144, + "rope_theta": 10000000.0, + "num_target_layers": 5, + "target_hidden_size": 4096, + "target_num_hidden_layers": 32, + "target_layer_ids": [ + 1, + 8, + 15, + 22, + 29 + ], + "mask_token_id": 248070, + "tie_word_embeddings": false, + "mtp_conditioning": { + "mode": "hidden", + "native_config": { + "attention_bias": false, + "attention_dropout": 0, + "attn_output_gate": true, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_value_head_dim": 128, + "max_position_embeddings": 262144, + "mlp_only_layers": [], + "model_type": "qwen3_5_text", + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 16, + "num_hidden_layers": 32, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "use_cache": true, + "vocab_size": 248320, + "mamba_ssm_dtype": "float32", + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25 + } + } + } +} diff --git a/aspec/config/draft_algorithm.py b/aspec/config/draft_algorithm.py index 0abfc12..134c03f 100644 --- a/aspec/config/draft_algorithm.py +++ b/aspec/config/draft_algorithm.py @@ -73,4 +73,9 @@ def draft_model_semantics(config) -> dict: "enable_post_mlp_norm", "input_embedding_scale", "output_multiplier", "final_logit_softcapping", "is_causal", ) - return {key: raw[key] for key in fields if key in raw} + semantics = {key: raw[key] for key in fields if key in raw} + if raw.get("mtp_conditioning") is not None: + from aspec.config.mtp_conditioning import resolve_mtp_conditioning + + semantics["mtp_conditioning"] = resolve_mtp_conditioning(raw) + return semantics diff --git a/aspec/config/dspark_qwen3_5_9b_dmtp_config.json b/aspec/config/dspark_qwen3_5_9b_dmtp_config.json new file mode 100644 index 0000000..ea37245 --- /dev/null +++ b/aspec/config/dspark_qwen3_5_9b_dmtp_config.json @@ -0,0 +1,110 @@ +{ + "architectures": [ + "DSparkDraftModel" + ], + "model_type": "dspark", + "hidden_size": 4096, + "intermediate_size": 12288, + "head_dim": 128, + "num_hidden_layers": 5, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "vocab_size": 248320, + "rms_norm_eps": 1e-06, + "max_position_embeddings": 262144, + "rope_theta": 10000000.0, + "num_target_layers": 5, + "target_hidden_size": 4096, + "target_num_hidden_layers": 32, + "target_layer_ids": [ + 1, + 8, + 15, + 22, + 29 + ], + "mask_token_id": 248070, + "markov_rank": 256, + "markov_head_type": "vanilla", + "enable_confidence_head": true, + "confidence_head_with_markov": true, + "tie_word_embeddings": false, + "mtp_conditioning": { + "mode": "hidden", + "native_config": { + "attention_bias": false, + "attention_dropout": 0, + "attn_output_gate": true, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_value_head_dim": 128, + "max_position_embeddings": 262144, + "mlp_only_layers": [], + "model_type": "qwen3_5_text", + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "num_attention_heads": 16, + "num_hidden_layers": 32, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "use_cache": true, + "vocab_size": 248320, + "mamba_ssm_dtype": "float32", + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25 + } + } + } +} diff --git a/aspec/config/mtp_conditioning.py b/aspec/config/mtp_conditioning.py new file mode 100644 index 0000000..bce8d93 --- /dev/null +++ b/aspec/config/mtp_conditioning.py @@ -0,0 +1,103 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Explicit MTP-prefix semantics, independent of execution and Target payloads.""" + +import copy +from collections.abc import Mapping + + +def resolve_mtp_conditioning(raw): + policy = raw.get("mtp_conditioning") + if policy is None: + return None + if not isinstance(policy, Mapping) or set(policy) != {"mode", "native_config"}: + raise ValueError("mtp_conditioning requires exactly mode and native_config") + if policy["mode"] not in {"token", "hidden"}: + raise ValueError("MTP conditioning mode must be token or hidden") + architectures = raw.get("architectures") + if ((architectures and architectures[0] not in { + "DFlashDraftModel", "DSparkDraftModel", "Qwen3DSparkModel"}) + or raw.get("model_type") in {"dflash2", "k3_dspark"}): + raise ValueError("MTP conditioning currently requires DFlash or DSpark full GQA") + if not isinstance(policy["native_config"], Mapping): + raise ValueError("MTP conditioning native_config must be a native Qwen3.5 config mapping") + from aspec.models.draft.mtp import NativeMTPConfig + + native = NativeMTPConfig(native_config=policy["native_config"]).native_config + if (raw.get("ffn_type", "dense") != "dense" + or raw.get("attention_mode", "gqa") != "gqa" + or raw.get("is_causal", False) + or any(kind != "full_attention" for kind in raw.get("layer_types", []))): + raise ValueError("MTP conditioning currently requires Dense bidirectional full GQA Drafts") + if native.get("attention_dropout", 0): + raise ValueError("Frozen MTP conditioning requires zero attention dropout") + for name, expected in ( + ("hidden_size", raw.get("target_hidden_size")), + ("vocab_size", raw.get("vocab_size")), + ("num_hidden_layers", raw.get("target_num_hidden_layers")), + ): + if native.get(name) != expected: + raise ValueError(f"MTP native {name} disagrees with its Target/Draft binding") + if raw.get("hidden_size") != native["hidden_size"]: + raise ValueError("MTP conditioning requires the same frozen Target embedding width") + if "full_attention" not in native.get("layer_types", []): + raise ValueError("MTP native config must identify its full-attention predictor") + return {"mode": policy["mode"], "native_config": copy.deepcopy(native)} + + +def mtp_conditioning_semantics(raw): + policy = resolve_mtp_conditioning(raw) + if policy is None: + return None + return { + **policy, + "prefix_layout": "dual_anchor", + "first_candidate": "teacher_forced_in_training_native_mtp_in_proposals", + "target_context": "strictly_before_original_anchor", + "prediction_state": "target_h[a-1]_and_token_x[a]_predict_x[a+1]", + "feature_computation": "frozen_draft_side_native_mtp", + "feature_stage": "native_predictor_post_norm", + "fusion": "zero_initialized_embedding_residual" if policy["mode"] == "hidden" else None, + "frozen": ["native_mtp", "target_embedding", "target_lm_head"], + } + + +def validate_mtp_target_config(native, target): + """Check the native predictor's mathematical binding, beyond tensor shapes.""" + target = target.get("text_config", target) + if target.get("model_type") != "qwen3_5_text": + raise ValueError("MTP conditioning requires the native Qwen3.5 Dense Target") + from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig + + target = Qwen3_5TextConfig(**target).to_dict() + keys = ("model_type", "hidden_size", "intermediate_size", "num_hidden_layers", + "num_attention_heads", "num_key_value_heads", "head_dim", "vocab_size", + "rms_norm_eps", "hidden_act", "attention_bias", "attention_dropout", + "layer_types", "rope_parameters", "attn_output_gate", "max_position_embeddings") + for key in keys: + if key in native and native[key] != target.get(key): + raise ValueError("MTP conditioning native config differs from its Target: " + key) + + +def mtp_conditioning_memory(policy): + """Additional BF16 model/optimizer bytes; neither total fit nor activation peak.""" + if policy is None: + return None + from aspec.models.dmtp import FrozenMTPFeatures + + source = FrozenMTPFeatures(policy["native_config"]) + frozen = sum(p.numel() for p in source.parameters()) + width = policy["native_config"]["hidden_size"] + fusion = 2 * width * width + 2 * width if policy["mode"] == "hidden" else 0 + return { + "scope": "additional Draft-side components; no duplicated vocabulary embedding/head", + "frozen_weight_bytes": 2 * frozen, + "trainable_weight_bytes": 2 * fusion, + "gradient_bytes": 2 * fusion, + "fp32_master_bytes": 4 * fusion, + "fp32_gradient_buffer_bytes": 4 * fusion, + "adam_moment_bytes": 8 * fusion, + "derived_feature_bytes_per_token": 2 * width if fusion else 0, + "activation_peak_bytes": None, + } diff --git a/aspec/controller/eval.py b/aspec/controller/eval.py index a5fbd9b..667e141 100644 --- a/aspec/controller/eval.py +++ b/aspec/controller/eval.py @@ -120,7 +120,10 @@ def _check_idle_timeout( def eval_selection_metric(metrics): """Retrieval recall is not decoder acceptance; each stage selects its own proxy.""" - return "eval/candidate_top1_recall" if "eval/candidate_top1_recall" in metrics else "eval/simulated_acc_len" + for key in ("eval/candidate_top1_recall", "eval/suffix_simulated_acc_len"): + if key in metrics: + return key + return "eval/simulated_acc_len" def update_checkpoint_eval_meta( @@ -142,13 +145,17 @@ def update_checkpoint_eval_meta( logger.warning(f"Checkpoint meta.json not found at {meta_path}, skipping eval meta update") return current_best - for key in ("eval/avg_loss", "eval/avg_acc", "eval/simulated_acc_len", + selection_metric = eval_selection_metric(eval_metrics) + for key in ("eval/avg_loss", "eval/avg_acc", "eval/simulated_acc_len", "eval/suffix_simulated_acc_len", "eval/candidate_top1_recall", "eval/candidate_teacher_mass"): if key in eval_metrics: metadata[key] = eval_metrics[key] + if selection_metric in eval_metrics: + metadata["selection_metric"] = selection_metric + if selection_metric == "eval/suffix_simulated_acc_len": + metadata.pop("eval/simulated_acc_len", None) _write_checkpoint_metadata(meta_path, metadata) - selection_metric = eval_selection_metric(eval_metrics) score = eval_metrics.get(selection_metric) if score is not None and math.isfinite(score) and score > current_best: current_best = score diff --git a/aspec/controller/loop.py b/aspec/controller/loop.py index a193bc4..3013dbc 100644 --- a/aspec/controller/loop.py +++ b/aspec/controller/loop.py @@ -1995,8 +1995,10 @@ def check_progress(): } if "train/avg_acc" in metrics: postfix["acc"] = f"{metrics['train/avg_acc']:.3f}" - if "train/simulated_acc_len" in metrics: - postfix["acc_len"] = f"{metrics['train/simulated_acc_len']:.2f}" + for key, label in (("train/simulated_acc_len", "acc_len"), + ("train/suffix_simulated_acc_len", "suffix_acc_len")): + if key in metrics: + postfix[label] = f"{metrics[key]:.2f}" if enable_perf: postfix["P~"] = ( f"{metrics.get('perf/producer_batch_rate_estimate', 0):.1f}" diff --git a/aspec/draft_artifact.py b/aspec/draft_artifact.py index a27f00c..10e5f17 100644 --- a/aspec/draft_artifact.py +++ b/aspec/draft_artifact.py @@ -68,6 +68,10 @@ def validate_native_asset(path, *, expected_config=None, weights_only=False): config = AutoDraftModelConfig.from_dict(raw) if resolve_draft_algorithm(config.to_dict()) != manifest.get("algorithm"): raise ValueError("Native artifact algorithm identity mismatch") + from aspec.config.mtp_conditioning import mtp_conditioning_semantics + + if manifest.get("mtp_conditioning") != mtp_conditioning_semantics(config.to_dict()): + raise ValueError("Native artifact MTP conditioning semantics mismatch") if expected_config is not None and draft_model_semantics(config) != draft_model_semantics(expected_config): raise ValueError("Native artifact semantics disagree with the requested Draft") from aspec.config.draft_attention import resolve_draft_attention @@ -141,6 +145,8 @@ def __init__(self, draft_model, *, block_size): from aspec.models.dflash import DFlashModel self.model = draft_model.eval() + if getattr(draft_model.config, "mtp_conditioning", None): + raise ValueError("MTP-conditioned assets require an MTP-aware proposal consumer") self.algorithm = resolve_draft_algorithm(draft_model.config.to_dict()) if self.algorithm not in {"dspark", "dflash"}: raise ValueError("BlockDraftConsumer requires a DSpark or DFlash asset") diff --git a/aspec/models/dflash.py b/aspec/models/dflash.py index a2178b4..1afb268 100644 --- a/aspec/models/dflash.py +++ b/aspec/models/dflash.py @@ -142,6 +142,8 @@ class DFlashModel(nn.Module): - Per-position loss_mask application """ + prediction_start = 1 + def __init__( self, draft_model, @@ -169,6 +171,8 @@ def __init__( raise ValueError(f"dflash_dpace_alpha must be in [0, 1], got {dpace_alpha}") self.draft_model = draft_model + self.mtp_conditioning = getattr(getattr(draft_model, "config", None), "mtp_conditioning", None) + self.prediction_start = type(self).prediction_start + int(self.mtp_conditioning is not None) self.block_size = block_size self.num_anchors = num_anchors if anchor_sampling_scope not in {"row", "sample"}: @@ -181,11 +185,17 @@ def __init__( self.loss_decay_gamma = loss_decay_gamma self.ce_loss_alpha = float(ce_loss_alpha) self.l1_loss_alpha = float(l1_loss_alpha) + if self.mtp_conditioning and (block_size <= self.prediction_start + or self.loss_objective != "decay"): + raise ValueError("MTP conditioning requires decay loss and an unknown suffix slot") execution = resolve_draft_execution_capability( attention_backend, draft_context_parallel_size ) self.attention_backend = execution.attention_backend self.draft_context_parallel_size = execution.context_parallel_size + if self.mtp_conditioning and (self.draft_context_parallel_size != 1 + or anchor_sampling_scope != "row"): + raise ValueError("MTP conditioning currently requires CP1 row anchors") self.direct_cp_backend = execution.direct_cp_backend from aspec.config.draft_attention import ( draft_context_shard_alignment, @@ -234,12 +244,34 @@ def __init__( self.direct_cp_group = None @property - def uses_target_hidden_states(self) -> bool: + def uses_target_distribution(self) -> bool: config = getattr(self.draft_model, "config", None) if getattr(config, "dsa_training_stage", None) == "indexer_warmup": return False return self.l1_loss_alpha > 0 + @property + def uses_target_hidden_states(self) -> bool: + return self.uses_target_distribution or bool( + self.mtp_conditioning and self.mtp_conditioning["mode"] == "hidden" + ) + + def _prepare_draft_prefix(self, input_ids, last_hidden_states, attention_mask, draft_prefix, + *packed_metadata): + if not self.mtp_conditioning: + if draft_prefix is not None: + raise ValueError("Unconditioned Draft does not accept an MTP prefix") + return None + if any(value is not None for value in packed_metadata): + raise ValueError("MTP conditioning does not accept packed or sharded documents") + if input_ids.shape[0] != 1: + raise ValueError("MTP conditioning currently requires microbatch size one") + if draft_prefix is None: + from aspec.models.dmtp import training_mtp_prefix + + return training_mtp_prefix(self.draft_model, input_ids, last_hidden_states, attention_mask) + return draft_prefix + def _sample_anchor_positions( self, seq_len: int, @@ -257,6 +289,11 @@ def _sample_anchor_positions( raise ValueError("Sample anchor sampling requires packed_document_lengths from the collator") if (injected_anchors is None) != (injected_keep_mask is None): raise ValueError("Anchor replay requires both injected_anchors and injected_keep_mask") + if self.mtp_conditioning: + # There is no verifier h[-1] for the first document token. Keep + # this exclusion in sampling/replay, not a stochastic runtime error. + loss_mask = loss_mask.clone() + loss_mask[:, 0] = 0 if injected_anchors is not None: return validate_injected_anchors( injected_anchors, injected_keep_mask, @@ -354,6 +391,16 @@ def _create_noise_embed( return self.draft_model.embed_tokens(noise_ids) + def _condition_noise_embed(self, embeddings, input_ids, anchors, keep, draft_prefix): + if self.mtp_conditioning: + from aspec.models.dmtp import condition_noise_embeddings + + return condition_noise_embeddings(self.draft_model, embeddings, input_ids, + anchors, keep, draft_prefix, self.block_size) + if draft_prefix is not None: + raise ValueError("This training wrapper does not accept a conditioned draft prefix") + return embeddings + def _block_mask_options(self) -> dict: return {} @@ -471,6 +518,7 @@ def _draft_backbone( dsa_losses_out: Optional[list] = None, packed_document_lengths: Optional[tuple[int, ...]] = None, attention_mask: Optional[torch.Tensor] = None, + draft_prefix=None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: """ Shared DFlash backbone (context features → anchor sampling → noise @@ -630,6 +678,9 @@ def _draft_backbone( # 3. Create noise embeddings (anchor token + MASK tokens) noise_embedding = self._create_noise_embed(input_ids, anchor_positions, block_keep_mask) + noise_embedding = self._condition_noise_embed( + noise_embedding, input_ids, anchor_positions, block_keep_mask, draft_prefix + ) # 4. Create position IDs context_position_ids, draft_position_ids = self._create_position_ids( @@ -963,6 +1014,7 @@ def forward( anchor_seed: Optional[int] = None, packed_document_lengths: Optional[tuple[int, ...]] = None, attention_mask: Optional[torch.Tensor] = None, + draft_prefix=None, ) -> Tuple[ torch.Tensor, torch.Tensor, @@ -979,8 +1031,7 @@ def forward( loss: scalar training loss (objective-weighted) accuracy: scalar accuracy (binary mask, no decay) loss_per_position: [block_size] mean loss at each within-block position - (index 0 is the anchor slot and always 0; indices 1..B-1 are the - predicted tokens at 1..B-1 steps past the anchor) + (slots before prediction_start are known inputs and always 0) acc_per_position: [block_size] mean accuracy at each within-block position count_per_position: [block_size] valid label count at each within-block position before loss decay is applied @@ -998,6 +1049,10 @@ def forward( qsa_losses = [] dsa_losses = [] # 1-6. Shared backbone → draft hidden states + anchor bookkeeping. + draft_prefix = self._prepare_draft_prefix( + input_ids, last_hidden_states, attention_mask, draft_prefix, + ctx_doc_ids, base_position_ids, packed_context_positions, packed_document_lengths, + ) draft_hidden, anchor_positions, block_keep_mask, n_blocks = self._draft_backbone( input_ids, hidden_states_list, @@ -1014,6 +1069,7 @@ def forward( dsa_losses_out=dsa_losses, packed_document_lengths=packed_document_lengths, attention_mask=attention_mask, + draft_prefix=draft_prefix, ) if getattr(self.draft_model.config, "dsa_training_stage", None) == "indexer_warmup": @@ -1054,7 +1110,7 @@ def forward( weight_mask = weight_mask * valid_label_mask.float() pos_in_block = torch.arange(self.block_size, device=device).view(1, 1, -1) - weight_mask = weight_mask * (pos_in_block > 0).float() + weight_mask = weight_mask * (pos_in_block >= self.prediction_start).float() # Gather original loss_mask at label positions original_loss_mask_gathered = torch.gather( @@ -1077,7 +1133,7 @@ def forward( ce_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") loss_per_token = self.ce_loss_alpha * ce_per_token - if self.uses_target_hidden_states: + if self.uses_target_distribution: if last_hidden_states is None: raise ValueError( "DFlash L1 distillation (l1_loss_alpha > 0) requires target " @@ -1107,18 +1163,22 @@ def forward( ): # Loss decay: exp(-(k-1)/γ) so k=1 (1st prediction) gets weight 1.0 k = torch.arange(self.block_size, device=device).view(1, 1, -1) - decay_weights = torch.exp(-(k - 1).clamp(min=0).float() / self.loss_decay_gamma) + decay_weights = torch.exp( + -(k - self.prediction_start).clamp(min=0).float() / self.loss_decay_gamma + ) objective_weights = weight_mask * decay_weights elif self.loss_objective == "dpace": dpace_weights = torch.ones_like(weight_mask) - if self.block_size > 1: + if self.block_size > self.prediction_start: with torch.no_grad(): - target_confidences = torch.exp(-loss_per_token_by_position[..., 1:].float()) + target_confidences = torch.exp( + -loss_per_token_by_position[..., self.prediction_start:].float() + ) dpace_pred_weights = _dpace_position_weights( target_confidences, self.dpace_alpha, ).to(dtype=weight_mask.dtype) - dpace_weights[..., 1:] = dpace_pred_weights + dpace_weights[..., self.prediction_start:] = dpace_pred_weights objective_weights = weight_mask * dpace_weights flat_weights = objective_weights.view(-1) diff --git a/aspec/models/dmtp.py b/aspec/models/dmtp.py new file mode 100644 index 0000000..5b75b8d --- /dev/null +++ b/aspec/models/dmtp.py @@ -0,0 +1,181 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Native MTP features and dual-anchor conditioning, inspired by DMTP (Shen, 2026). + +Method: https://sylvan820.github.io/D-MTP/ (November 9 residual update). +Independent implementation reusing ASpec's retained DFlash backbone/objective +and native Qwen3.5 predictor. Configuration, scheduling, optimization and +recovery stay in the existing ASpec workflow. This is not a claim of +reproducing the author's training recipe or downstream serving runtime. +""" + +from dataclasses import dataclass + +import torch +from torch import nn + + +@dataclass(frozen=True) +class MTPPrefix: + # Row i is the state predicting x[i+2], after consuming x[i+1]. + prediction_hidden: torch.Tensor | None = None + # One actual MTP proposal per anchor; None means teacher-forced training. + first_candidate_ids: torch.Tensor | None = None + + +class MTPAnchorFusion(nn.Module): + def __init__(self, embedding_size, mtp_hidden_size, eps=1e-6): + super().__init__() + from aspec.models.draft.dflash import DFlashRMSNorm + + self.embedding_norm = DFlashRMSNorm(embedding_size, eps=eps) + self.hidden_norm = DFlashRMSNorm(mtp_hidden_size, eps=eps) + self.proj = nn.Linear(embedding_size + mtp_hidden_size, embedding_size, bias=False) + nn.init.zeros_(self.proj.weight) + + def forward(self, embeddings, hidden): + features = torch.cat((self.embedding_norm(embeddings), + self.hidden_norm(hidden.detach())), dim=-1) + return embeddings + self.proj(features) + + +class FrozenMTPFeatures(nn.Module): + """The native predictor body, borrowing the Draft's frozen Target embedding. + + It owns no embedding or output head, and therefore does not duplicate the + vocabulary-sized Target bindings in memory/checkpoints. Native predictor + mathematics are shared with NativeMTPDraftModel, rather than reimplemented. + """ + + def __init__(self, native_config): + super().__init__() + from aspec.models.draft.mtp import NativeMTPConfig, NativeMTPDraftModel + + with torch.device("meta"): + source = NativeMTPDraftModel(NativeMTPConfig(native_config=native_config)) + self.native = source.native + for name in ("pre_fc_norm_embedding", "pre_fc_norm_hidden", "fc", "midlayer", "norm"): + setattr(self, name, getattr(source, name)) + self.rotary_emb = type(source.rotary_emb)(self.native, device="cpu") + self._rotary_ready = False + self.requires_grad_(False) + self.eval() + + def train(self, mode=True): + return super().train(False) + + def _apply(self, fn, recurse=True): + self._rotary_ready = False + return super()._apply(fn, recurse=recurse) + + def initialize_native_weights(self, target_model_path): + from aspec.models.draft.base import load_tensor_from_pretrained + from aspec.models.draft.mtp import NativeMTPDraftModel + + state = {name: load_tensor_from_pretrained(target_model_path, + NativeMTPDraftModel.native_key(name)) + for name in self.state_dict()} + self.load_state_dict(state, strict=True, assign=True) + self.requires_grad_(False) + + @torch.no_grad() + def validate_native_weights(self, target_model_path): + from aspec.models.draft.base import load_tensor_from_pretrained + from aspec.models.draft.mtp import NativeMTPDraftModel + + for name, actual in self.state_dict().items(): + if hasattr(actual, "full_tensor"): + actual = actual.full_tensor() + expected = load_tensor_from_pretrained(target_model_path, + NativeMTPDraftModel.native_key(name)) + if actual.shape != expected.shape or not torch.equal( + actual.cpu(), expected.to(dtype=actual.dtype) + ): + raise ValueError("Frozen MTP differs from its Target: " + name) + + @torch.no_grad() + def forward(self, hidden_states, input_embeds, position_ids): + from aspec.models.draft.mtp import NativeMTPDraftModel + + if not self._rotary_ready: + # Device moves, to_empty and checkpoint restoration do not carry + # the native non-persistent frequencies. Regenerate them once. + self.rotary_emb = type(self.rotary_emb)(self.native, device=hidden_states.device) + self._rotary_ready = True + return NativeMTPDraftModel.predict_embeddings( + self, hidden_states.detach(), input_embeds.detach(), position_ids + ).detach() + + +def training_mtp_prefix(draft, input_ids, last_hidden_states, attention_mask=None): + policy = draft.config.mtp_conditioning + if policy["mode"] == "token": + return MTPPrefix() + if last_hidden_states is None or last_hidden_states.shape[:2] != input_ids.shape: + raise ValueError("MTP hidden conditioning requires aligned post-norm Target final states") + if input_ids.shape[0] != 1: + raise ValueError("MTP conditioning currently requires microbatch size one") + length = input_ids.shape[1] + if attention_mask is not None: + if attention_mask.shape != input_ids.shape: + raise ValueError("MTP attention mask must align with input token IDs") + length = int(attention_mask.sum()) + if (not attention_mask[:, :length].bool().all() + or attention_mask[:, length:].bool().any()): + raise ValueError("MTP conditioning requires right-padded complete documents") + if length < 2: + return MTPPrefix(last_hidden_states.new_zeros( + (1, max(input_ids.shape[1] - 1, 0), draft.config.target_hidden_size))) + positions = torch.arange(length - 1, device=input_ids.device).unsqueeze(0) + hidden = draft.mtp_source(last_hidden_states[:, :length - 1], + draft.embed_tokens(input_ids[:, 1:length]), positions) + hidden = torch.nn.functional.pad(hidden, (0, 0, 0, input_ids.shape[1] - length)) + return MTPPrefix(hidden.detach()) + + +def condition_noise_embeddings(draft, embeddings, input_ids, anchors, keep, prefix, block_size): + """Add T1 while retaining the original A cutoff and native state timing.""" + if not isinstance(prefix, MTPPrefix): + raise ValueError("MTP conditioning requires an explicit derived prefix") + if ((anchors < 1) | (anchors + 1 >= input_ids.shape[1]))[keep].any(): + raise ValueError("MTP anchors need a previous Target state and an in-bounds T1") + if prefix.first_candidate_ids is None: + ids = input_ids.gather(1, (anchors + 1).clamp(max=input_ids.shape[1] - 1)) + else: + ids = prefix.first_candidate_ids + if ids.shape != anchors.shape or ids.dtype != torch.long: + raise ValueError("MTP proposal IDs must be int64 with the anchor shape") + ids = torch.where(keep, ids, torch.full_like(ids, draft.mask_token_id)) + first = draft.embed_tokens(ids) + if draft.config.mtp_conditioning["mode"] == "hidden": + hidden = prefix.prediction_hidden + expected = (input_ids.shape[0], input_ids.shape[1] - 1, draft.config.target_hidden_size) + if hidden is None or hidden.shape != expected: + raise ValueError(f"MTP prediction hidden must have shape {expected}") + indices = (anchors - 1).clamp_min(0).unsqueeze(-1).expand(-1, -1, hidden.shape[-1]) + selected = hidden.gather(1, indices) + # Inactive features must be finite before trainable normalization/projection; + # masking their loss afterwards cannot prevent NaN parameter gradients. + selected = torch.where(keep.unsqueeze(-1), selected, 0.0) + first = draft.mtp_fusion(first, selected) + elif prefix.prediction_hidden is not None: + raise ValueError("Token-only conditioning does not consume MTP hidden features") + blocks = embeddings.reshape(input_ids.shape[0], anchors.shape[1], block_size, -1) + return torch.cat((blocks[:, :, :1], first.unsqueeze(2), blocks[:, :, 2:]), dim=2).flatten(1, 2) + + +@torch.no_grad() +def frozen_mtp_predictions(mtp, input_ids, final_hidden_states): + """Causal document-local predictor pass; no future token is consumed. + + x[a] is the known pending token A. Its prediction state is output row a-1, + formed from Target h[a-1] and x[a], and predicts the unverified x[a+1]. + """ + if mtp.training or any(p.requires_grad for p in mtp.parameters()): + raise ValueError("The native MTP feature producer must be frozen and in eval mode") + if (input_ids.ndim != 2 or input_ids.shape[0] != 1 or input_ids.shape[1] < 2 + or final_hidden_states.shape[:2] != input_ids.shape): + raise ValueError("MTP feature extraction requires one complete unpadded document") + positions = torch.arange(input_ids.shape[1] - 1, device=input_ids.device).unsqueeze(0) + return mtp.predict(final_hidden_states[:, :-1].detach(), input_ids[:, 1:], positions).detach() diff --git a/aspec/models/draft/dflash.py b/aspec/models/draft/dflash.py index 0876e20..83c682b 100644 --- a/aspec/models/draft/dflash.py +++ b/aspec/models/draft/dflash.py @@ -82,6 +82,7 @@ def __init__( norm_topk_prob: bool = True, router_aux_loss_coef: Optional[float] = None, cudnn_gqa_mode: str = "auto", + mtp_conditioning: Optional[dict] = None, **kwargs, ): # Transformers 5 validates YaRN parameters inside PretrainedConfig's @@ -164,6 +165,11 @@ def __init__( setattr(self, key, value) if self.model_type != "k3_dspark" and getattr(self, "attention_mode", "gqa") != "mla": resolve_dflash_rope_parameters(self) + from aspec.config.mtp_conditioning import resolve_mtp_conditioning + + self.mtp_conditioning = resolve_mtp_conditioning({ + **self.to_dict(), "mtp_conditioning": mtp_conditioning, + }) class DFlashRMSNorm(nn.Module): @@ -1140,6 +1146,15 @@ def __init__(self, config: PretrainedConfig): # Final norm self.final_norm = DFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.mtp_fusion = self.mtp_source = None + if getattr(config, "mtp_conditioning", None): + from aspec.models.dmtp import FrozenMTPFeatures, MTPAnchorFusion + + self.freeze_embedding() + self.mtp_source = FrozenMTPFeatures(config.mtp_conditioning["native_config"]) + if config.mtp_conditioning["mode"] == "hidden": + self.mtp_fusion = MTPAnchorFusion(config.hidden_size, target_hidden_size, + eps=config.rms_norm_eps) self.last_csa_indexer_losses = () self.last_qsa_indexer_losses = () self.last_dsa_indexer_losses = () @@ -1148,6 +1163,36 @@ def __init__(self, config: PretrainedConfig): self.configure_qsa_trainable_parameters() self.configure_dsa_trainable_parameters() + def initialize_mtp_conditioning(self, target_model_path): + if self.mtp_source is not None: + self.mtp_source.initialize_native_weights(target_model_path) + + @torch.no_grad() + def validate_mtp_binding(self, target_model_path, embedding_key): + if self.mtp_source is None: + return + from transformers import AutoConfig + + from aspec.config.mtp_conditioning import validate_mtp_target_config + from aspec.models.draft.base import load_tensor_from_pretrained + + validate_mtp_target_config( + self.config.mtp_conditioning["native_config"], + AutoConfig.from_pretrained(target_model_path).to_dict(), + ) + + actual = self.embed_tokens.weight + if hasattr(actual, "full_tensor"): + actual = actual.full_tensor() + expected = load_tensor_from_pretrained(target_model_path, embedding_key) + if actual.shape != expected.shape: + raise ValueError("MTP conditioning embedding shape differs from its Target") + for start in range(0, expected.shape[0], 1024): + if not torch.equal(actual[start:start + 1024].cpu(), + expected[start:start + 1024].to(dtype=actual.dtype)): + raise ValueError("MTP conditioning frozen embedding differs from its Target") + self.mtp_source.validate_native_weights(target_model_path) + def configure_dsa_trainable_parameters(self) -> None: layers = [layer.self_attn.dsa_layer for layer in self.layers if getattr(layer.self_attn, "dsa_layer", None) is not None] diff --git a/aspec/models/draft/mtp.py b/aspec/models/draft/mtp.py index 72779d4..590ecb7 100644 --- a/aspec/models/draft/mtp.py +++ b/aspec/models/draft/mtp.py @@ -143,7 +143,11 @@ def validate_target_binding(self, target_model_path, *, embedding_key, lm_head_k raise ValueError("Native MTP frozen tensor differs from its Target: " + name) def predict(self, hidden_states, input_ids, position_ids, layout=None): - embeds = self.pre_fc_norm_embedding(self.embed_tokens(input_ids)) + return self.predict_embeddings(hidden_states, self.embed_tokens(input_ids), position_ids, layout) + + def predict_embeddings(self, hidden_states, input_embeds, position_ids, layout=None): + """Shared native mathematics for trainable MTP and frozen feature extraction.""" + embeds = self.pre_fc_norm_embedding(input_embeds) hidden = self.pre_fc_norm_hidden(hidden_states) projected = self.fc(torch.cat((embeds, hidden), dim=-1)) rotary_positions = position_ids.unsqueeze(0).expand(3, -1, -1) diff --git a/aspec/models/dspark.py b/aspec/models/dspark.py index 8a32aef..678c357 100644 --- a/aspec/models/dspark.py +++ b/aspec/models/dspark.py @@ -71,6 +71,8 @@ def _cached_projection_chunks(counts, row_budget, block_size, *, copies, limit=N class DSparkModel(DFlashModel): """DSpark training wrapper (DFlash backbone + Markov/L1/confidence heads).""" + prediction_start = 0 + def __init__( self, draft_model, @@ -146,7 +148,7 @@ def _indexer_warmup_kind(self) -> str | None: return None @property - def uses_target_hidden_states(self) -> bool: + def uses_target_distribution(self) -> bool: if self._indexer_warmup_kind() is not None: # The warm-up phase optimizes only the detached indexer KL. Loading # verifier hidden states would retain and transfer target activations @@ -154,19 +156,19 @@ def uses_target_hidden_states(self) -> bool: return False # The confidence head regresses onto an accept rate derived from the # verifier distribution, so the L1 weight alone does not cover it. - return super().uses_target_hidden_states or ( + return super().uses_target_distribution or ( self.draft_model.confidence_head is not None and self.confidence_head_alpha > 0 ) def _decay_weights(self, device: torch.device) -> torch.Tensor: - """exp(-k/gamma) over within-block position k (DeepSpec convention). + """Weight one at the first unknown prediction, then decay within the block. - Every slot 0..B-1 is a real prediction here, so slot 0 (the first - predicted token) gets weight 1.0 and later slots decay. + Plain DSpark starts at slot zero; conditioning makes T1 visible there + and starts the suffix weights at slot one. """ k = torch.arange(self.block_size, device=device).view(1, 1, -1) if self.loss_decay_gamma is not None and self.loss_decay_gamma > 0: - return torch.exp(-k.float() / self.loss_decay_gamma) + return torch.exp(-(k - self.prediction_start).clamp_min(0).float() / self.loss_decay_gamma) return torch.ones_like(k, dtype=torch.float32) def forward( @@ -184,6 +186,7 @@ def forward( anchor_seed: Optional[int] = None, packed_document_lengths: Optional[tuple[int, ...]] = None, attention_mask: Optional[torch.Tensor] = None, + draft_prefix=None, ) -> Tuple[ torch.Tensor, torch.Tensor, @@ -211,6 +214,10 @@ def forward( qsa_losses = [] dsa_losses = [] # ---- Shared DFlash backbone (steps 1-6) → draft hidden states ---- + draft_prefix = self._prepare_draft_prefix( + input_ids, last_hidden_states, attention_mask, draft_prefix, + ctx_doc_ids, base_position_ids, packed_context_positions, packed_document_lengths, + ) draft_hidden, anchor_positions, block_keep_mask, n_blocks = self._draft_backbone( input_ids, hidden_states_list, @@ -227,11 +234,11 @@ def forward( dsa_losses_out=dsa_losses, packed_document_lengths=packed_document_lengths, attention_mask=attention_mask, + draft_prefix=draft_prefix, ) # ---- Labels + eval mask (DeepSpec convention) ---- - # Slot j predicts the token at anchor+j+1 (the real anchor token seeds - # slot 0). All block_size slots are supervised — there is no masked - # anchor slot, unlike DFlash. + # Slot j predicts anchor+j+1. Plain DSpark supervises every slot; + # conditioning excludes visible T1 using prediction_start below. label_offsets = torch.arange(1, self.block_size + 1, device=device).view(1, 1, -1) label_indices = anchor_positions.unsqueeze(-1) + label_offsets # [B, nb, bs] valid_label_mask = label_indices < seq_len @@ -257,6 +264,9 @@ def forward( eval_bool = block_keep_mask.unsqueeze(-1) & valid_label_mask & (target_loss_mask > 0.5) eval_bool = eval_bool.to(torch.int32).cumprod(dim=-1).bool() eval_mask = eval_bool.float() # [B, nb, bs] + eval_mask = eval_mask * ( + torch.arange(self.block_size, device=device) >= self.prediction_start + ) decay_weight_mask = eval_mask * self._decay_weights(device) local_den = decay_weight_mask.sum() @@ -304,7 +314,7 @@ def forward( anchor_token_ids = torch.gather(input_ids, 1, anchor_positions) # [B, nb] prev_token_ids = torch.cat([anchor_token_ids.unsqueeze(-1), target_ids[:, :, :-1]], dim=-1) aligned_hidden = None - if self.uses_target_hidden_states: + if self.uses_target_distribution: if last_hidden_states is None: raise ValueError( "DSpark L1/confidence losses require target last_hidden_states; set " @@ -512,9 +522,17 @@ def forward( **dsa_components, } if aligned_hidden is not None: - loss_components.update(acceptance_components(acceptance_totals)) - - # ---- Metrics (cross-entropy based; all block_size slots are productive) ---- + observations = acceptance_components(acceptance_totals) + if self.mtp_conditioning: + observations = { + ("suffix_tau_probabilistic" if key == "tau_probabilistic" else + f"suffix_accept_rate@{int(key.split('@')[1]) - self.prediction_start}"): value + for key, value in observations.items() + if key == "tau_probabilistic" or int(key.split('@')[1]) >= self.prediction_start + } + loss_components.update(observations) + + # ---- Metrics over the unknown prediction slots ---- with torch.no_grad(): accuracy = correct_position_sum.sum() / eval_mask.sum().clamp(min=1e-6) count_per_position = eval_mask.sum(dim=(0, 1)) @@ -608,7 +626,9 @@ def _prediction_terms(self, hidden, target_ids, prev_ids, head_weight, position_loss = (ce * eval_mask).sum((0, 1)) position_correct = correct.float().sum((0, 1)) acceptance_totals = ( - acceptance_statistics(accept_rate, eval_mask) + F.pad(acceptance_statistics( + accept_rate[..., self.prediction_start:], eval_mask[..., self.prediction_start:] + ), (self.prediction_start, 0)) if accept_rate is not None else logits.new_zeros( (2, self.block_size + 1), dtype=torch.float32, ) diff --git a/aspec/planning.py b/aspec/planning.py index 7515a0e..7576ee9 100644 --- a/aspec/planning.py +++ b/aspec/planning.py @@ -558,7 +558,46 @@ def _load_draft_descriptor(config: DictConfig) -> dict[str, Any]: if structure["ffn_type"] != "dense" and _select(config, "training.compile_model", False): raise PlanCompileError("MoE Draft torch.compile is not yet supported; disable compile_model") + from aspec.config.mtp_conditioning import resolve_mtp_conditioning + + try: + conditioning = resolve_mtp_conditioning({**raw, "layer_types": attention["layer_types"]}) + except (ValueError, TypeError) as exc: + raise PlanCompileError(str(exc)) from exc + if conditioning: + from aspec.config.mtp_conditioning import validate_mtp_target_config + + target_raw = _read_local_target_config(_select(config, "model.target_model_path")) + if target_raw is not None: + try: + validate_mtp_target_config(conditioning["native_config"], target_raw) + except ValueError as exc: + raise PlanCompileError(str(exc)) from exc + if family not in {"dflash", "dspark"}: + raise PlanCompileError("MTP conditioning requires DFlash or DSpark") + for key, required in ( + ("draft_context_parallel_size", 1), ("micro_batch_size", 1), + ("dflash_packing", False), ("compile_model", False), + ("anchor_sampling_scope", "row"), + ("fsdp_strategy", "REPLICATE"), + ): + if _select(config, "training." + key) != required: + raise PlanCompileError(f"MTP conditioning requires training.{key}={required}") + if _select(config, "dataset.eval_micro_batch_size") not in {None, 1}: + raise PlanCompileError("MTP conditioning requires dataset.eval_micro_batch_size=1") + if conditioning["mode"] == "hidden" and not _select(config, "inference.store_last_hidden_states"): + raise PlanCompileError("MTP hidden conditioning requires inference.store_last_hidden_states=true") + start = 2 if family == "dflash" else 1 + if _select(config, "training.dflash_block_size") <= start: + raise PlanCompileError("MTP conditioning requires at least one unknown suffix prediction") + if family == "dflash" and _select(config, "training.dflash_loss_objective") != "decay": + raise PlanCompileError("MTP-conditioned DFlash currently requires the decay objective") + if (_select(config, "training.training_num_nodes", 1) != 1 + or _select(config, "cluster.mode", "single_node") != "single_node"): + raise PlanCompileError("MTP conditioning currently supports a single node") + return { + "mtp_conditioning": conditioning, "structure": structure, "eagle4": eagle4, "attention": attention, @@ -1307,6 +1346,18 @@ def _workload( ["embedding", "candidate_head"] + (["lm_head"] if _select(config, "model.freeze_lm_head", False) else [])) if family in {"dspark", "dflash", "dflash2"}: objective["anchor_sampling_scope"] = anchor_scope + if draft.get("mtp_conditioning"): + from aspec.config.mtp_conditioning import mtp_conditioning_semantics + + raw = json.loads(Path(draft["path"]).read_text()) + objective["mtp_conditioning"] = mtp_conditioning_semantics(raw) + start = 2 if family == "dflash" else 1 + objective.update( + prediction_start=start, + candidate_budget=objective["block_size"] - (1 if family == "dflash" else 0), + suffix_prediction_budget=objective["block_size"] - start, + acceptance_scope="teacher_forced_T1_conditional_suffix_only", + ) objective.update( { "learning_rate": _select(config, "training.learning_rate"), @@ -1392,6 +1443,10 @@ def _policy( runtime: Mapping[str, Any], context_contract: Mapping[str, Any], ) -> ExecutionPolicy: + from aspec.config.mtp_conditioning import mtp_conditioning_memory, resolve_mtp_conditioning + + raw = json.loads(Path(_select(config, "model.draft_model_config")).read_text()) + conditioning_memory = mtp_conditioning_memory(resolve_mtp_conditioning(raw)) cluster_mode = _select(config, "cluster.mode", "single_node") data_plane = _select(config, "mooncake.data_plane", "store") transport_protocol = _select(config, "mooncake.protocol", "tcp") @@ -1647,6 +1702,7 @@ def _policy( }, capacity={ "device_memory": { + **({"mtp_conditioning": conditioning_memory} if conditioning_memory else {}), "sglang_static_fraction": (None if runtime["inference_engine_type"] == "offline" else _select(config, "inference.sglang.mem_fraction_static")) }, diff --git a/aspec/reporting.py b/aspec/reporting.py index b0b282c..52fcdb0 100644 --- a/aspec/reporting.py +++ b/aspec/reporting.py @@ -339,6 +339,8 @@ def render_run_report(report: Mapping[str, Any]) -> str: evaluation = training.get("evaluation") if evaluation and evaluation.get("enabled"): lines.append(f"held_out_eval: {len(evaluation.get('history', []))} points; proxy_only=true") + if evaluation.get("selection_metric"): + lines.append(f"selection_metric: {evaluation['selection_metric']}") best = evaluation.get("best_checkpoint") lines.append(f"best_checkpoint: {best.get('checkpoint') if best else None}") if performance: diff --git a/aspec/training/checkpoint.py b/aspec/training/checkpoint.py index 63c4fac..26a59bc 100644 --- a/aspec/training/checkpoint.py +++ b/aspec/training/checkpoint.py @@ -802,6 +802,19 @@ def load_initial_draft_weights( manifest_path = checkpoint_path.parent / "export_manifest.json" native = manifest_path.is_file() and json.loads(manifest_path.read_text()).get("storage_format") == "aspec_native" + conditioning = getattr(draft_model.config, "mtp_conditioning", None) + transition = False + if conditioning: + saved_config_path = checkpoint_path.parent / "config.json" + if not saved_config_path.is_file(): + raise ValueError("MTP conditioning warmstart requires the source Draft config") + saved_raw = json.loads(saved_config_path.read_text()) + saved_conditioning = saved_raw.get("mtp_conditioning") + transition = saved_conditioning is None + if saved_conditioning is not None and saved_conditioning != conditioning: + raise ValueError("MTP conditioning warmstart policy differs from the requested Draft") + if transition and any(key.startswith(("mtp_source.", "mtp_fusion.")) for key in tensors): + raise ValueError("Plain Draft warmstart contains undeclared MTP tensors") if native: from aspec.draft_artifact import validate_native_asset @@ -812,7 +825,10 @@ def load_initial_draft_weights( # This entry is weights-only initialization, not exact recovery. # Preserve model bindings and strict tensors while allowing the new # run to choose its indexer stage/objective policy. - validate_native_asset(checkpoint_path.parent, expected_config=draft_model.config, weights_only=True) + expected_config = draft_model.config + if transition: + expected_config = {**draft_model.config.to_dict(), "mtp_conditioning": None} + validate_native_asset(checkpoint_path.parent, expected_config=expected_config, weights_only=True) if hasattr(draft_model, "keep_csa_projection_precision"): draft_model.keep_csa_projection_precision() state = tensors @@ -824,6 +840,15 @@ def load_initial_draft_weights( # DFlash-family serving assets normally borrow this frozen tensor from # the Target. The caller loads it from that Target immediately afterward. state["embed_tokens.weight"] = draft_model.state_dict()["embed_tokens.weight"] + if transition: + # Explicit weights-only phase: only the new, initialized frozen source + # and zero residual are absent from the predecessor. Everything else + # still uses the ordinary strict loader. + for key, value in draft_model.state_dict().items(): + if key.startswith(("mtp_source.", "mtp_fusion.")): + if value.is_meta: + raise ValueError("Initialize the frozen MTP source before weights-only warmstart") + state[key] = value draft_model.load_state_dict(state, strict=True) if hasattr(draft_model, "t2d"): from aspec.models.vocab import validate_vocab_mapping diff --git a/aspec/training/dflash_trainer.py b/aspec/training/dflash_trainer.py index bbd85be..4c37223 100644 --- a/aspec/training/dflash_trainer.py +++ b/aspec/training/dflash_trainer.py @@ -92,6 +92,7 @@ def __init__(self, args: Namespace): self.l1_loss_alpha = getattr(args, "dflash_l1_loss_alpha", 0.0) self._last_hs_prenorm = getattr(args, "last_hidden_states_prenorm", False) self.verifier_norm: Optional[torch.nn.Module] = None + self._mtp_binding_validated = False def init_model( self, @@ -143,7 +144,10 @@ def init_model( draft_model = self._build_draft_model(config) + conditioning = getattr(config, "mtp_conditioning", None) if dist.get_rank() == 0: + if conditioning: + draft_model.initialize_mtp_conditioning(target_model_path) initial_path = getattr(self.args, "initial_draft_model_path", None) if initial_path: checkpoint.load_initial_draft_weights( @@ -197,6 +201,15 @@ def init_model( _unwrapped = getattr(self.model, "_orig_mod", self.model) # torch.compile self.dflash = getattr(_unwrapped, "module", _unwrapped) # DDP/replicate self.draft_model = self.dflash.draft_model + if conditioning: + self._anchor_slot_offset = self.dflash.prediction_start + if hasattr(self, "confidence_head_alpha"): + self._extra_loss_component_keys = [ + "ce_loss", "l1_loss", "confidence_loss", + *(f"suffix_accept_rate@{index}" for index in range( + self.block_size - self.dflash.prediction_start)), + "suffix_tau_probabilistic", + ] total_steps = self.args.lr_total_steps decay_style = getattr(self.args, "lr_decay_style", "cosine") @@ -234,6 +247,11 @@ def init_model( self.lr_scheduler = self.optimizer.lr_scheduler checkpoint_payload = checkpoint.load(self) + if conditioning: + self.draft_model.validate_mtp_binding( + target_model_path, + embedding_key=getattr(self.args, "embedding_key", "model.embed_tokens.weight"), + ) self._init_target_lm_head(target_model_path) @@ -248,6 +266,7 @@ def init_model( self.target_lm_head_weight = ( None if self.target_lm_head is None else self.target_lm_head.lm_head.weight ) + self._mtp_binding_validated = bool(conditioning) # Frozen Target helpers can initialize temporary random parameters. # Restore the training RNG only after all model initialization ends. @@ -296,6 +315,27 @@ def get_feature_contract_status(self) -> dict: from aspec.models.draft.moe import QwenDraftMoE draft = getattr(self, "draft_model", None) + if getattr(getattr(draft, "config", None), "mtp_conditioning", None): + import inspect + import sys + + from aspec.config.mtp_conditioning import mtp_conditioning_semantics + + status = dict(status) + status["mtp_conditioning"] = { + **mtp_conditioning_semantics(draft.config.to_dict()), + "prediction_start": self.dflash.prediction_start, + "frozen_source_parameters": sum(p.numel() for p in draft.mtp_source.parameters()), + "fusion_parameters": (sum(p.numel() for p in draft.mtp_fusion.parameters()) + if draft.mtp_fusion is not None else 0), + "target_binding_validated": self._mtp_binding_validated, + "runtime_origin": { + "python": sys.executable, + "trainer": inspect.getfile(type(self)), + "draft": inspect.getfile(type(draft)), + "native_features": inspect.getfile(type(draft.mtp_source)), + }, + } layers = list(getattr(draft, "layers", [])) if layers and any(isinstance(layer.mlp, (QwenDraftMoE, LingDraftMoE)) for layer in layers): status = dict(status) @@ -337,6 +377,8 @@ def get_feature_contract_status(self) -> dict: return status def save_draft_model_for_serving(self, output_dir: str) -> None: + if getattr(self.draft_model.config, "mtp_conditioning", None): + raise NotImplementedError("MTP-conditioned Draft requires native --format aspec export") if getattr(self.draft_model.config, "ffn_type", "dense") != "dense": raise NotImplementedError( "MoE Draft serving export is not yet supported; retain the complete training checkpoint" @@ -633,6 +675,10 @@ def eval_from_cache(self) -> dict: return self._aggregate_eval_metrics(all_metrics) + def _prediction_metric_prefix(self) -> str: + """MTP observations cover the unknown suffix conditional on visible T1.""" + return "suffix_" if getattr(getattr(self, "dflash", None), "mtp_conditioning", None) else "" + def _aggregate_eval_metrics(self, all_step_metrics: list[dict]) -> dict: if not all_step_metrics: return {} @@ -659,7 +705,7 @@ def _aggregate_eval_metrics(self, all_step_metrics: list[dict]) -> dict: count_key="count_pp", ) - # Drop anchor slot (index 0) — see _aggregate_metrics for rationale. + # Observe only unknown predictions, excluding visible T1 when conditioned. pred_loss_pp = avg_loss_pp[self._anchor_slot_offset :] pred_acc_pp = avg_acc_pp[self._anchor_slot_offset :] pred_count_pp = count_pp[self._anchor_slot_offset :] @@ -673,22 +719,23 @@ def _aggregate_eval_metrics(self, all_step_metrics: list[dict]) -> dict: weighted_avg_loss, avg_acc = self._compute_scalar_metrics( pred_loss_pp, pred_acc_pp, pred_count_pp ) + prefix = self._prediction_metric_prefix() metrics: dict = { "eval/avg_loss": weighted_avg_loss, "eval/avg_acc": avg_acc, - "eval/simulated_acc_len": simulated_acc_len, + f"eval/{prefix}simulated_acc_len": simulated_acc_len, } for i in range(pred_loss_pp.shape[0]): - metrics[f"eval/ploss_{i}"] = pred_loss_pp[i].item() - metrics[f"eval/acc_{i}"] = pred_acc_pp[i].item() + metrics[f"eval/{prefix}ploss_{i}"] = pred_loss_pp[i].item() + metrics[f"eval/{prefix}acc_{i}"] = pred_acc_pp[i].item() metrics.update(loss_components) if dist.get_rank() == 0: logger.info( f"eval: loss={weighted_avg_loss:.4f}, acc={avg_acc:.4f}, " - f"sim_acc_len={simulated_acc_len:.2f}" + f"{prefix}sim_acc_len={simulated_acc_len:.2f}" ) return metrics @@ -844,9 +891,7 @@ def _aggregate_metrics( count_key="count_per_position", ) - # Skip index 0 (anchor slot, always zero); indices 1..B-1 are the - # predicted tokens at 1..B-1 steps past the anchor. Re-index to 0..B-2 - # so the naming matches Eagle3 (acc_0 = first predicted token). + # Re-index unknown predictions from zero; conditioned suffix_0 denotes T2. pred_loss_pp = avg_loss_pp[self._anchor_slot_offset :] pred_acc_pp = avg_acc_pp[self._anchor_slot_offset :] pred_count_pp = count_pp[self._anchor_slot_offset :] @@ -860,11 +905,12 @@ def _aggregate_metrics( simulated_acc_len += cumulative avg_loss, avg_acc = self._compute_scalar_metrics(pred_loss_pp, pred_acc_pp, pred_count_pp) + prefix = self._prediction_metric_prefix() metrics = { "train/avg_loss": avg_loss, "train/avg_acc": avg_acc, - "train/simulated_acc_len": simulated_acc_len, + f"train/{prefix}simulated_acc_len": simulated_acc_len, "train/grad_norm": grad_norm.item() if grad_norm is not None else 0.0, "train/global_step": self.global_step, "train/lr": self.optimizer.get_learning_rate(), @@ -872,8 +918,8 @@ def _aggregate_metrics( } for i in range(pred_loss_pp.shape[0]): - metrics[f"train/ploss_{i}"] = pred_loss_pp[i].item() - metrics[f"train/acc_{i}"] = pred_acc_pp[i].item() + metrics[f"train/{prefix}ploss_{i}"] = pred_loss_pp[i].item() + metrics[f"train/{prefix}acc_{i}"] = pred_acc_pp[i].item() metrics.update(loss_components) diff --git a/configs/recipes/README.md b/configs/recipes/README.md index e9c3f3d..56e6628 100644 --- a/configs/recipes/README.md +++ b/configs/recipes/README.md @@ -13,3 +13,5 @@ The [shared Target walkthrough](../../examples/qwen3-8b-shared-target/README.md) the recipe for experiment comparison, separate curves and recovery from different checkpoints. Supplemental variants belong under `configs/examples/`; benchmark comparisons remain under `configs/benchmarks/`. + +MTP-conditioned DSpark uses `qwen3_5_9b_dmtp.yaml`; the [workflow](../../examples/qwen3.5-9b-dmtp/README.md) describes hidden/token modes, DFlash selection and native-only assets. diff --git a/configs/recipes/qwen3_5_9b_dmtp.yaml b/configs/recipes/qwen3_5_9b_dmtp.yaml new file mode 100644 index 0000000..e30cd74 --- /dev/null +++ b/configs/recipes/qwen3_5_9b_dmtp.yaml @@ -0,0 +1,106 @@ +# Release training starting point: three passes, held-out eval and checkpoints. +# Hyperparameters are not a convergence/quality claim; see docs/user-guide/recipes.md. +# MTP-hidden-conditioned DSpark, native Qwen3.5 predictor frozen on each Draft. +# CP1, MBS1, Dense full GQA, REPLICATE, no packing/whole-model compile. +# Conditional-suffix observations do not certify a serving implementation. + +model: + target_model_path: ${oc.env:ASPEC_TARGET_MODEL,Qwen/Qwen3.5-9B} + native_context_length: 262144 + trust_remote_code: false + embedding_key: model.language_model.embed_tokens.weight + norm_key: model.language_model.norm.weight + draft_model_config: aspec/config/dspark_qwen3_5_9b_dmtp_config.json + +dataset: + eval_interval: 100 + eval_micro_batch_size: 1 + train_data_path: ${oc.env:ASPEC_TRAIN_DATA} + eval_data_path: ${oc.env:ASPEC_EVAL_DATA} + chat_template: null + min_loss_tokens: 32 + shuffle_dataset: true + +training: + attention_backend: flex_attention + draft_context_parallel_size: 1 + micro_batch_size: 1 + draft_accumulation_steps: 128 + learning_rate: 6e-4 + min_lr: 6e-5 + weight_decay: 0.0 + max_concurrent_batches: 1 + streaming_dispatch_with_training: true + pipeline_dispatch_with_training: true + max_grad_norm: 1.0 + max_seq_length: 4096 + num_epochs: 3 + seed: 42 + training_num_gpus_per_node: 4 + training_num_nodes: 1 + placement_strategy: inference_first + ttt_length: 7 + fsdp_strategy: REPLICATE + fsdp_reduce_dtype: bfloat16 + prefetch_depth: 1 + save_interval: 100 + save_per_epoch: false + save_final_checkpoint: true + warmup_ratio: 0.04 + dflash_block_size: 7 + dflash_packing: false + dspark_num_anchors: 512 + dspark_num_target_layers: 5 + dspark_loss_decay_gamma: 4.0 + dspark_ce_loss_alpha: 0.1 + dspark_l1_loss_alpha: 0.9 + dspark_confidence_head_alpha: 1.0 + +inference: + inference_engine_type: sgl + store_last_hidden_states: true + target_prefill_memory_optimized: false + inference_num_gpus: 4 + inference_num_gpus_per_engine: 1 + inference_num_gpus_per_node: 8 + inference_fetch_batch: 8 + inference_buffer_threshold: 2 + max_sample_pool_size: 12 + inference_batch_size: 2 + sglang: + tp_size: 1 + mem_fraction_static: 0.5 + +mooncake: + master_server_address: null + metadata_server: null + data_plane: transfer_engine + protocol: nvlink_intra + require_nvlink: true + global_segment_size: 64GB + local_buffer_size: 8GB + host_buffer_size: 14GB + max_seq_len: 4096 + hidden_dim: 4096 + async_put_pool_size: 1 + kv_lease_ttl_s: 60 + get_retry_max_wait_seconds: 120 + gpu_ring_slots: 3 + source_gpu_ring_slots: 3 + destination_gpu_ring_slots: 1 + enable_gpu_direct: true + allow_host_fallback: false + enable_hard_pin: false + +logging: + report_to: none + +debug: + enable_perf_metrics: true + performance_measurement_required: false + performance_warmup_steps: 2 + performance_drain_steps: 1 + timeline_enabled: false + +output_dir: outputs/qwen3_5_9b_dmtp +cache_dir: cache/qwen3_5_9b_dmtp diff --git a/docs/contributing/validation.md b/docs/contributing/validation.md index 8757e7e..f8041ed 100644 --- a/docs/contributing/validation.md +++ b/docs/contributing/validation.md @@ -15,7 +15,7 @@ python -m tools.check_docs python -m ruff check aspec tests tools ``` -`contracts` checks all 13 configs/12 workflows, resource budgets and training defaults, negative +`contracts` checks all 14 configs/13 workflows, resource budgets and training defaults, negative acceptance cases, and existing algorithm/packing/CP/sharing/multi-node contracts. Model names, data paths and node IPs in CPU tests are fixtures, never real GPU inputs. No models are downloaded by this command. CPU contracts do not establish GPU kernel parity or available capacity. @@ -80,6 +80,7 @@ bash tools/run_in_env.sh python -m tools.recipes export \ | Scope | Required additional evidence | | --- | --- | | DFlash / EAGLE3 / Native MTP | Algorithm contract tests plus the actual recipe's train/resume/assets; Native MTP initialization must use the native predictor | +| MTP-conditioned DSpark / DFlash | Native timing/loss/frozen-state tests, real train/resume/native assets and conditional-suffix observations; no full serving acceptance claim | | EAGLE4 | Both stage reports and candidate frozen-body audit; stage-local resume; no official candidate-recipe claim | | Ling | Correct model tokenizer/template and feature layers; use Ling-prepared data; MoE routing/bias state is retained | | 27B CP2 | Genuine long supervised inputs, actual long work in runtime metrics, GPU CP loss/gradient parity and measured memory fit | diff --git a/docs/framework/REFERENCES.md b/docs/framework/REFERENCES.md index 11b3207..0f33c80 100644 --- a/docs/framework/REFERENCES.md +++ b/docs/framework/REFERENCES.md @@ -18,6 +18,7 @@ the retained source identities; applicable license and notice texts are distribu | [Ling-V2](https://github.com/inclusionAI/Ling-V2) and [DeepSpec](https://github.com/deepseek-ai/DeepSpec) | Declared model, routing and DSpark implementation derivations | | [DFlash](https://github.com/z-lab/dflash) `07ebd93db9f472af339b644bb70221ad8428328a` | DFlash2 model adaptation inherited through TorchSpec; Z Lab MIT notices retained | | [Transformers](https://github.com/huggingface/transformers) v4.57.1 `8cb5963cc22174954e7dca2c0a3320b7dc2f4edc` | Adapted `Qwen3MoeSparseMoeBlock` and load-balancing objective in `aspec/models/draft/moe.py` | +| [DMTP](https://sylvan820.github.io/D-MTP/) | Inspiration for dual-anchor MTP hidden conditioning and zero-initialized embedding residual; independent ASpec integration reuses the retained native predictor, without copying DMTP code | | [Speculators](https://github.com/vllm-project/speculators) `7ec7aba9450b533334e8f7992baf4e93be3c1a57` | Native MTP recurrent objective and checkpoint bindings; ASpec supplies its own training/layout integration | | [Ling-3.0-flash](https://huggingface.co/inclusionAI/Ling-3.0-flash) `fcc60a52b34c69c6271d0ac8a8cb0815020b6a43` | Adapted MoE routing and independently retained reference class bodies; Apache-2.0 source headers preserved | | [EAGLE](https://github.com/SafeAILab/EAGLE), [GPT-NeoX](https://github.com/EleutherAI/gpt-neox) and Transformers through TorchSpec | Model/data helpers in `aspec/models/draft/base.py` and `aspec/data/preprocessing.py`; original source headers retain the transitive attribution | diff --git a/docs/support.md b/docs/support.md index 4558455..58bd775 100644 --- a/docs/support.md +++ b/docs/support.md @@ -13,6 +13,7 @@ and tests do not imply CPU training support. | EAGLE3 | Recurrent TTT, full/pruned vocabulary, GQA/MHA or materialized MLA; FA4 packing/CP | No EAGLE3 MoE, anchored mode or whole-model compile | | EAGLE4 | Dense and selected upstream MoE; Draft/candidate stages; packing/CP and native assets | Candidate Top-K forward-KL is explicitly implemented from public semantics, not an unpublished official recipe; no serving integration | | DFlash2 | Experimental full/causal/SWA training with Flex, cuDNN THD or FA4 packing/CP | Experimental objective; no claim to reproduce an official training recipe | +| MTP-conditioned DSpark / DFlash | Frozen native Qwen3.5 predictor, token/hidden dual-anchor training, exact recovery and strict native assets | Dense full GQA, CP1, MBS1, REPLICATE, one node, non-packing, decay objective; no generic serving consumer | | Native MTP | Qwen3.5-9B native single predictor, same-Target frozen embedding/head, multi-step training, recovery and export/stitch | Text-only Dense MTP; SDPA CP1 or independently validated FA4 packing/CP; no generic multi-layer/MoE MTP | | Dense / Qwen MoE / Ling MoE | Structure-specific FFNs, routing state, checkpointing and strict assets | No promise of every algorithm × structure combination | | MHA/GQA / materialized MLA | Explicit head geometry and tensor bindings | MLA materializes K/V; no absorbed-MLA claim | diff --git a/docs/user-guide/algorithms.md b/docs/user-guide/algorithms.md index 69fe89c..917f01b 100644 --- a/docs/user-guide/algorithms.md +++ b/docs/user-guide/algorithms.md @@ -96,6 +96,18 @@ anchor-shared window. FA4 preserves convolution, selector, document boundaries and CP ownership. Functional results do not establish official-recipe reproduction or long-term quality. +## Native MTP conditioning for DSpark/DFlash + +`configs/recipes/qwen3_5_9b_dmtp.yaml` uses the native Qwen3.5 predictor as a frozen +condition for a DSpark Draft. DFlash uses the corresponding DMTP Draft JSON. +The ordinary trainer owns fusion parameters, optimizer, recovery and native assets; +Target selected-layer/final-state transport stays unchanged. Dense full GQA, CP1, +MBS1, REPLICATE, one node and non-packing are the initial scope. Other combinations +fail Plan validation. Follow the [workflow and timing contract](../../examples/qwen3.5-9b-dmtp/README.md). +Given teacher-forced T1, only unknown suffix positions receive supervision. Suffix +acceptance observations are conditional; native export requires an MTP-aware downstream +consumer and does not certify generic SGLang serving. + ## Native Qwen3.5 MTP `configs/recipes/qwen3_5_9b_mtp.yaml` initializes the actual native predictor and binds frozen embeddings/head diff --git a/docs/user-guide/recipes.md b/docs/user-guide/recipes.md index d7585dc..d16df31 100644 --- a/docs/user-guide/recipes.md +++ b/docs/user-guide/recipes.md @@ -19,6 +19,7 @@ connections on your hardware before training. | `dflash` | Qwen3-8B / DFlash | 1 node, 3 GPUs: 1T + 2D; 4K | Train a block-parallel DFlash Draft | | `eagle3` | Qwen3-8B / EAGLE3 | 1 node, 3 GPUs: 1T + 2D; 4K | Train a recurrent EAGLE3 Draft | | `eagle4` | Qwen3-8B / EAGLE4 | 1 node, 3 GPUs: 1T + 2D; 4K; sequential stages | Train a Draft, then candidate retrieval projections | +| `dmtp` | Qwen3.5-9B / MTP-conditioned DSpark | 1 node, 8 GPUs: 4T + 4D; 4K, CP1 | Train a Draft conditioned on the frozen native predictor; native assets | | `native-mtp` | Qwen3.5-9B / native predictor | 1 node, 3 GPUs: 1T + 2D; 4K | Fine-tune the Target's native predictor | | `ling-tiny` | Ling-3.0-tiny / Dense DSpark | 1 node, 2 GPUs: 1T + 1D; 4K | Train a small Ling Draft | | `ling-flash` | Ling-3.0-flash / Dense DSpark | 1 node, 3 GPUs: 2T (TP2) + 1D; 4K | Train against a larger Ling Target | @@ -41,6 +42,8 @@ For a short single-node DSpark run with cuDNN attention, CP1 and no packing, use Both start with four TP1 Target engines and four Draft GPUs and require a matching local Target checkpoint plus prepared token IDs/masks. The 9B DSpark example is separate from native MTP training. +MTP-conditioned DSpark is available as `aspec recipes plan dmtp`. See the [Qwen3.5-9B workflow](../../examples/qwen3.5-9b-dmtp/README.md) for supervision timing, recovery and native export. + ## Inputs and normal training Install using [the environment setup](installation.md). Download the exact Target checkpoint or use diff --git a/examples/qwen3.5-9b-dmtp/README.md b/examples/qwen3.5-9b-dmtp/README.md new file mode 100644 index 0000000..304d0f3 --- /dev/null +++ b/examples/qwen3.5-9b-dmtp/README.md @@ -0,0 +1,93 @@ +# Native MTP conditions a DSpark or DFlash Draft + +This example adds a frozen Qwen3.5-9B native MTP predictor to the ordinary ASpec +training workflow. It trains the Draft and, in hidden mode, a zero-initialized +residual fusion. The Target, its embedding/output head and the native predictor +stay frozen. It does not define a separate training engine or a serving runtime. + +## Configuration and training + +Activate the [isolated Linux environment](../../docs/user-guide/installation.md). +Set disjoint training/evaluation inputs using the [supported data formats](../../docs/user-guide/data.md): + +```bash +export ASPEC_TARGET_MODEL=/path/to/Qwen3.5-9B +export ASPEC_TRAIN_DATA=/path/to/train.parquet +export ASPEC_EVAL_DATA=/path/to/heldout.parquet +export RAY_ADDRESS=local + +aspec --config configs/recipes/qwen3_5_9b_dmtp.yaml --dry-run +aspec --config configs/recipes/qwen3_5_9b_dmtp.yaml +aspec --report outputs/qwen3_5_9b_dmtp +``` + +The starting recipe uses DSpark, five full-attention GQA layers, four Draft GPUs +and four TP1 Target GPUs. `mtp_conditioning` in the Draft JSON has exactly two +fields: `mode` (`token` or `hidden`) and `native_config` (the native Qwen3.5 text +configuration). Hidden mode requires `inference.store_last_hidden_states=true`. +Ordinary selected-layer and final Target features are transported unchanged; +the Draft derives the native MTP state locally. No second Target transformer or +vocabulary-sized embedding/head is allocated for the MTP body. + +To use DFlash, select `aspec/config/dflash_qwen3_5_9b_dmtp_config.json`, set +`training.dflash_num_target_layers=5`, and choose the DFlash block/anchor/loss +settings. The current objective is `decay`. To compare token-only conditioning, +copy the JSON to a new file and change `mtp_conditioning.mode` to `token`. +A plain Draft omits `mtp_conditioning` entirely. These are separate workloads. + +The first supported scope is Dense full bidirectional GQA, one complete document +per microbatch, CP1, REPLICATE, one node, no packing or whole-model compilation. +Plan rejects unsupported combinations. The native config and all frozen tensors +must match the same Target; loading and recovery validate that binding. + +## Timing and supervision + +For original anchor `a`, Target context ends strictly before `a`. Native MTP takes +post-norm Target `h[a-1]` and known token `x[a]`, producing the state that predicts +`T1=x[a+1]`. The Draft input is `[A,T1,MASK,...]`. Training teacher-forces T1; +hidden mode adds a normalized, zero-initialized embedding residual at its input +slot. With identical Draft weights, this initialization has the same forward output +as token-only conditioning. + +DFlash supervises input slots `2..B-1`; DSpark supervises prediction slots `1..B-1`. +The visible T1 is excluded from loss, accuracy and decay position zero. DSpark +reports `suffix_accept_rate@k` and `suffix_tau_probabilistic` conditional on T1. +Both algorithms report `train/suffix_simulated_acc_len` and +`eval/suffix_simulated_acc_len`; their per-position `suffix_acc_0` and +`suffix_ploss_0` start at T2. Checkpoint selection and evaluation metadata use the +suffix proxy explicitly. Plain DSpark/DFlash runs retain their ordinary metric names. +These observations omit acceptance of the native MTP proposal and cannot establish +full-block accepted length or serving speed. + +## Recovery and assets + +Use a fresh output directory for exact resume, preserving the model policy, +source identity and full learning-rate horizon: + +```bash +aspec --config configs/recipes/qwen3_5_9b_dmtp.yaml \ + training.load_path=outputs/qwen3_5_9b_dmtp/checkpoints \ + output_dir=outputs/qwen3_5_9b_dmtp-resumed +aspec export --format aspec --input-dir /path/to/complete/iter_checkpoint \ + --output-dir outputs/qwen3_5_9b_dmtp-asset +``` + +Native assets retain conditioning config, fusion and frozen predictor weights. +`aspec.draft_artifact.load_native_draft` strictly reloads them. Generic SGLang +export and the ordinary block-proposal consumer reject conditioned assets; +a downstream MTP-aware consumer must implement the actual proposed T1 and its +verification/cache timing. Export does not certify that consumer. + +`model.initial_draft_model_path` is weights-only initialization. A plain Draft +can start a new conditioned phase with initialized native tensors and zero +fusion; undeclared or partial MTP tensors fail strict loading. Exact resume +cannot change the conditioning mode or native policy. + +## Framework validation + +Feature regressions cover configuration rejection, native predictor timing, +future-token isolation, visible-token loss masking, frozen parameter ownership, +weights-only initialization, exact recovery and strict native asset loading. +The recipe checks exercise the ordinary training, recovery and export entry points. +See [contributor validation](../../docs/contributing/validation.md) for the +additional GPU and recipe checks required for this path. diff --git a/tests/test_dmtp.py b/tests/test_dmtp.py new file mode 100644 index 0000000..d06247e --- /dev/null +++ b/tests/test_dmtp.py @@ -0,0 +1,548 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""DMTP timing, known-token loss exclusion, and frozen-feature boundaries.""" + +import copy +import json +from types import MethodType, SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.nn import functional as F + +from aspec.models.dflash import DFlashModel, _create_dflash_mask_mod +from aspec.models.dmtp import ( + MTPAnchorFusion, + MTPPrefix, + frozen_mtp_predictions, +) +from aspec.models.draft.dflash import DFlashConfig, DFlashDraftModel + + +def native_config(): + return {"model_type": "qwen3_5_text", "hidden_size": 16, "intermediate_size": 32, + "num_attention_heads": 2, "num_key_value_heads": 1, "head_dim": 8, + "vocab_size": 32, "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["full_attention"], + "rms_norm_eps": 1e-6, "mtp_num_hidden_layers": 1, + "rope_parameters": {"rope_type": "default", "rope_theta": 10000, + "partial_rotary_factor": 1.0, "mrope_section": [1, 1, 2], + "mrope_interleaved": True}} + + +def wrapper(use_hidden=True, algorithm="dflash"): + from aspec.models.draft.dspark import DSparkConfig, DSparkDraftModel + from aspec.models.dspark import DSparkModel + + cfg_cls, model_cls, root_cls = ((DFlashConfig, DFlashDraftModel, DFlashModel) + if algorithm == "dflash" else + (DSparkConfig, DSparkDraftModel, DSparkModel)) + config = cfg_cls(architectures=["DFlashDraftModel" if algorithm == "dflash" else "DSparkDraftModel"], + hidden_size=16, intermediate_size=32, num_hidden_layers=1, + num_attention_heads=2, num_key_value_heads=1, head_dim=8, + vocab_size=32, num_target_layers=1, target_hidden_size=16, + target_num_hidden_layers=4, target_layer_ids=[1], + mask_token_id=31, max_position_embeddings=128, + mtp_conditioning={"mode": "hidden" if use_hidden else "token", + "native_config": native_config()}) + draft = model_cls(config) + draft.embed_tokens.requires_grad_(False) + return root_cls(draft, block_size=5, num_anchors=2, attention_backend="sdpa") + + +def _nonzero_model(algorithm): + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(17) + model = wrapper(algorithm=algorithm) + nn.init.normal_(model.draft_model.mtp_fusion.proj.weight, std=0.1) + return model + + +def _fusion_inputs(): + generator = torch.Generator().manual_seed(31) + ids = torch.arange(12)[None] + return dict(input_ids=ids, hidden_states_list=torch.randn(1, 12, 16, generator=generator), + loss_mask=torch.ones_like(ids), lm_head_weight=torch.randn(32, 16, generator=generator), + last_hidden_states=torch.randn(1, 12, 16, generator=generator), + injected_anchors=torch.tensor([[3, 5]]), + injected_keep_mask=torch.ones(1, 2, dtype=torch.bool), + draft_prefix=MTPPrefix(torch.randn(1, 11, 16, generator=generator, requires_grad=True))) + + +def _reference_noise(self, embeddings, input_ids, anchors, keep, prefix): + """Scalar anchor indexing and explicit RMS/linear equations, independent of fusion helpers.""" + draft = self.draft_model + fusion = draft.mtp_fusion + + def normalize(value, norm): + value = value.float() + return value * torch.rsqrt(value.square().mean(-1, keepdim=True) + norm.variance_epsilon) * norm.weight + + blocks = embeddings.reshape(1, anchors.shape[1], self.block_size, -1) + result = [] + for index, anchor in enumerate(anchors[0].tolist()): + token = input_ids[:, anchor + 1] if keep[0, index] else input_ids.new_tensor([draft.mask_token_id]) + first = F.embedding(token, draft.embed_tokens.weight) + hidden = (prefix.prediction_hidden[:, anchor - 1].detach() if keep[0, index] + else torch.zeros_like(first)) + features = torch.cat((normalize(first, fusion.embedding_norm), + normalize(hidden, fusion.hidden_norm)), dim=-1) + first = first + F.linear(features, fusion.proj.weight) + result.append(torch.cat((blocks[:, index, :1], first[:, None], blocks[:, index, 2:]), dim=1)) + return torch.cat(result, dim=1) + + +def _reference_logits_and_loss(model, hidden, inputs): + """Dense equations for the conditioned CE/L1/confidence objective on complete blocks.""" + draft = model.draft_model + hidden = hidden.reshape(1, 2, model.block_size, -1) + offsets = torch.arange(model.block_size)[None, None] + positions = inputs["injected_anchors"][..., None] + offsets + is_dspark = hasattr(draft, "markov_head") + positions = positions + int(is_dspark) + labels = inputs["input_ids"][:, positions[0]] + logits = F.linear(hidden, inputs["lm_head_weight"]) + previous = None + if is_dspark: + previous = F.embedding(inputs["input_ids"][:, positions[0] - 1], + draft.markov_head.markov_w1.weight) + logits = logits + F.linear(previous, draft.markov_head.markov_w2.weight) + ce = F.cross_entropy(logits.flatten(0, 2), labels.flatten(), reduction="none").view_as(labels) + terms = model.ce_loss_alpha * ce + if is_dspark: + teacher = inputs["last_hidden_states"][:, positions[0] - 1] + target = F.linear(teacher, inputs["lm_head_weight"]).float().softmax(-1) + l1 = (logits.float().softmax(-1) - target).abs().sum(-1) + confidence = F.linear(torch.cat((hidden, previous), dim=-1), + draft.confidence_head.proj.weight, draft.confidence_head.proj.bias).squeeze(-1) + accept = (1 - 0.5 * l1).clamp(0, 1).detach() + terms = terms + model.l1_loss_alpha * l1 + model.confidence_head_alpha * F.binary_cross_entropy_with_logits( + confidence, accept, reduction="none") + weights = (offsets >= model.prediction_start) * torch.exp( + -(offsets - model.prediction_start).clamp_min(0).float() / model.loss_decay_gamma) + weights = weights.expand_as(terms) + return logits, (terms * weights).sum() / weights.sum() + + +def _assert_parameter_gradients(actual, expected, *, exact=False): + actual_parameters = dict(actual.named_parameters()) + expected_parameters = dict(expected.named_parameters()) + assert actual_parameters.keys() == expected_parameters.keys() + for name, parameter in actual_parameters.items(): + reference = expected_parameters[name] + assert (parameter.grad is None) == (reference.grad is None), name + if parameter.grad is not None: + assert torch.isfinite(parameter.grad).all(), name + torch.testing.assert_close(parameter.grad, reference.grad, + rtol=0 if exact else 2e-5, atol=0 if exact else 2e-6, msg=name) + if not parameter.requires_grad: + assert parameter.grad is None, name + + +@pytest.mark.parametrize("algorithm", ["dflash", "dspark"]) +def test_nonzero_fusion_logits_loss_and_gradients_match_independent_reference(algorithm): + actual = _nonzero_model(algorithm) + reference = copy.deepcopy(actual) + reference._condition_noise_embed = MethodType(_reference_noise, reference) + inputs, reference_inputs = _fusion_inputs(), _fusion_inputs() + captured = [] + handle = actual.draft_model.register_forward_hook(lambda _, args, output: captured.append(output)) + result = actual(**inputs) + handle.remove() + expected_hidden, _, _, _ = reference._draft_backbone( + reference_inputs["input_ids"], reference_inputs["hidden_states_list"], reference_inputs["loss_mask"], + injected_anchors=reference_inputs["injected_anchors"], + injected_keep_mask=reference_inputs["injected_keep_mask"], draft_prefix=reference_inputs["draft_prefix"]) + logits, expected_loss = _reference_logits_and_loss(reference, expected_hidden, reference_inputs) + actual_logits, _ = _reference_logits_and_loss(actual, captured[0], inputs) + torch.testing.assert_close(actual_logits, logits, rtol=2e-5, atol=2e-6) + torch.testing.assert_close(result[0], expected_loss, rtol=2e-5, atol=2e-6) + result[0].backward() + expected_loss.backward() + _assert_parameter_gradients(actual, reference) + assert inputs["draft_prefix"].prediction_hidden.grad is None + for parameter in actual.draft_model.mtp_fusion.parameters(): + assert parameter.grad is not None and parameter.grad.abs().sum() > 0 + + +@pytest.mark.parametrize("algorithm", ["dflash", "dspark"]) +def test_nonzero_fusion_gradient_checkpointing_preserves_outputs_and_gradients(algorithm): + actual = _nonzero_model(algorithm) + reference = copy.deepcopy(actual) + actual.draft_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + assert actual.training and actual.draft_model.gradient_checkpointing + result, expected = actual(**_fusion_inputs()), reference(**_fusion_inputs()) + torch.testing.assert_close(result, expected, rtol=2e-5, atol=2e-6) + result[0].backward() + expected[0].backward() + _assert_parameter_gradients(actual, reference) + + +@pytest.mark.parametrize("algorithm", ["dflash", "dspark"]) +@pytest.mark.parametrize("all_inactive", [False, True]) +def test_inactive_nan_mtp_features_preserve_finite_outputs_and_valid_gradients(algorithm, all_inactive): + actual = _nonzero_model(algorithm) + reference = copy.deepcopy(actual) + inputs, reference_inputs = _fusion_inputs(), _fusion_inputs() + for values in [inputs, reference_inputs]: + values["injected_anchors"] = torch.tensor([[3, 0]]) + values["injected_keep_mask"] = torch.tensor([[not all_inactive, False]]) + values["attention_mask"] = torch.tensor([[1] * 8 + [0] * 4]) + values["loss_mask"] = values["attention_mask"].clone() + with torch.no_grad(): + # Inactive anchor slots use row zero; padded feature rows are also irrelevant. + inputs["draft_prefix"].prediction_hidden[:, 0] = float("nan") + inputs["draft_prefix"].prediction_hidden[:, 7:] = float("nan") + result, expected = actual(**inputs), reference(**reference_inputs) + assert torch.isfinite(result[0]) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + result[0].backward() + expected[0].backward() + _assert_parameter_gradients(actual, reference, exact=True) + + +def test_residual_starts_at_token_only_and_only_changes_t1(): + model = wrapper() + ids = torch.arange(12).unsqueeze(0) + anchors, keep = torch.tensor([[3, 5]]), torch.tensor([[True, False]]) + base = model._create_noise_embed(ids, anchors, keep) + hidden = torch.randn(1, 11, 16, requires_grad=True) + actual = model._condition_noise_embed(base, ids, anchors, keep, MTPPrefix(hidden)) + expected = base.clone() + expected[:, 1] = model.draft_model.embed_tokens(ids[:, 4]) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + actual[:, 1].sum().backward() + assert model.draft_model.mtp_fusion.proj.weight.grad.abs().sum() > 0 + assert hidden.grad is None + assert model.draft_model.embed_tokens.weight.grad is None + + +def test_fusion_selects_state_predicting_t1_and_accepts_actual_proposals(): + model = wrapper() + ids = torch.arange(12).unsqueeze(0) + anchors, keep = torch.tensor([[3, 5]]), torch.ones(1, 2, dtype=torch.bool) + hidden = torch.arange(11 * 16).reshape(1, 11, 16).float() + seen = [] + model.draft_model.mtp_fusion.register_forward_pre_hook(lambda _, values: seen.append(values)) + proposed = torch.tensor([[20, 21]]) + base = model._create_noise_embed(ids, anchors, keep) + actual = model._condition_noise_embed(base, ids, anchors, keep, MTPPrefix(hidden, proposed)) + torch.testing.assert_close(seen[0][1], hidden[:, [2, 4]]) + torch.testing.assert_close(actual[:, [1, 6]], model.draft_model.embed_tokens(proposed)) + + +def test_mtp_features_are_detached_even_after_residual_becomes_nonzero(): + fusion = MTPAnchorFusion(16, 16) + nn.init.normal_(fusion.proj.weight) + hidden = torch.randn(1, 16, requires_grad=True) + fusion(torch.randn(1, 16), hidden).square().sum().backward() + assert hidden.grad is None + assert fusion.hidden_norm.weight.grad.abs().sum() > 0 + + +def test_only_suffix_positions_receive_loss_and_decay_restarts_at_t2(monkeypatch): + model = wrapper() + features = torch.randn(1, 10, 16, requires_grad=True) + anchors = torch.tensor([[3, 5]]) + keep = torch.ones_like(anchors, dtype=torch.bool) + monkeypatch.setattr(model, "_draft_backbone", lambda *a, **k: (features, anchors, keep, 2)) + ids = torch.arange(12).unsqueeze(0) + head = torch.randn(32, 16) + result = model(input_ids=ids, hidden_states_list=None, loss_mask=torch.ones_like(ids), + lm_head_weight=head, draft_prefix=MTPPrefix(torch.randn(1, 11, 16))) + logits = F.linear(features, head).reshape(1, 2, 5, 32) + labels = ids[:, anchors[0, :, None] + torch.arange(5)] + ce = F.cross_entropy(logits[:, :, 2:].flatten(0, 2), labels[:, :, 2:].flatten(), + reduction="none").reshape(1, 2, 3) + weights = torch.exp(-torch.arange(3).float() / model.loss_decay_gamma) + expected = (ce * weights).sum() / (2 * weights.sum()) + torch.testing.assert_close(result[0], expected) + assert result[4].tolist() == [0, 0, 2, 2, 2] + result[0].backward() + assert features.grad.reshape(1, 2, 5, 16)[:, :, :2].count_nonzero() == 0 + + +def test_context_cutoff_stays_before_a_even_with_t1_visible(): + mask = _create_dflash_mask_mod(torch.tensor([[3]]), torch.tensor([[True]]), 12, 5) + for query in range(5): + assert mask(0, 0, query, 2) + assert not mask(0, 0, query, 3) + assert not mask(0, 0, query, 4) + assert mask(0, 0, query, 13) # T1 is visible within the draft block. + + +def test_frozen_mtp_uses_previous_target_state_and_known_pending_token(): + class Predictor(nn.Module): + def predict(self, hidden, ids, positions): + self.seen = hidden, ids, positions + return hidden + ids.unsqueeze(-1) + mtp = Predictor().eval() + ids = torch.arange(6).unsqueeze(0) + final = torch.randn(1, 6, 16, requires_grad=True) + output = frozen_mtp_predictions(mtp, ids, final) + torch.testing.assert_close(mtp.seen[0], final[:, :-1]) + assert mtp.seen[1].tolist() == [[1, 2, 3, 4, 5]] + assert mtp.seen[2].tolist() == [[0, 1, 2, 3, 4]] + assert not output.requires_grad + mtp.train() + with pytest.raises(ValueError, match="frozen"): + frozen_mtp_predictions(mtp, ids, final) + + +def test_invalid_timing_and_packed_input_fail_closed(): + model = wrapper() + ids = torch.arange(12).unsqueeze(0) + anchors, keep = torch.tensor([[0, 3]]), torch.ones(1, 2, dtype=torch.bool) + base = model._create_noise_embed(ids, anchors, keep) + with pytest.raises(ValueError, match="previous Target"): + model._condition_noise_embed(base, ids, anchors, keep, MTPPrefix(torch.randn(1, 11, 16))) + with pytest.raises(ValueError, match="packed"): + model(input_ids=ids, hidden_states_list=None, loss_mask=torch.ones_like(ids), + lm_head_weight=torch.randn(32, 16), ctx_doc_ids=torch.zeros_like(ids)) + with pytest.raises(ValueError, match="complete unpadded"): + frozen_mtp_predictions(SimpleNamespace(training=False, parameters=lambda: []), + ids.expand(2, -1), torch.randn(2, 12, 16)) + + +def test_frozen_provider_matches_native_and_materializes_meta_buffers(monkeypatch): + from aspec.models.dmtp import FrozenMTPFeatures, training_mtp_prefix + from aspec.models.draft.mtp import NativeMTPConfig, NativeMTPDraftModel + + expected = NativeMTPDraftModel(NativeMTPConfig(native_config=native_config())) + expected.requires_grad_(False).eval() + tensors = {expected.native_key(k): v for k, v in expected.state_dict().items()} + monkeypatch.setattr("aspec.models.draft.base.load_tensor_from_pretrained", lambda _, k: tensors[k]) + provider = FrozenMTPFeatures(native_config()) + provider.initialize_native_weights("unused") + provider.to_empty(device="cpu") + provider.initialize_native_weights("unused") + provider.train() + assert not provider.training and not any(p.requires_grad for p in provider.parameters()) + ids, final = torch.randint(0, 32, (1, 7)), torch.randn(1, 7, 16) + actual = provider(final[:, :-1], expected.embed_tokens(ids[:, 1:]), torch.arange(6)[None]) + torch.testing.assert_close(actual, frozen_mtp_predictions(expected, ids, final), rtol=0, atol=0) + provider.validate_native_weights("unused") + draft = wrapper().draft_model + draft.mtp_source = provider + draft.embed_tokens.load_state_dict(expected.embed_tokens.state_dict()) + prefix = training_mtp_prefix(draft, ids, final) + torch.testing.assert_close(prefix.prediction_hidden, actual, rtol=0, atol=0) + shorter = training_mtp_prefix(draft, ids, final, torch.tensor([[1, 1, 1, 1, 0, 0, 0]])) + torch.testing.assert_close(shorter.prediction_hidden[:, :3], actual[:, :3]) + assert shorter.prediction_hidden[:, 3:].count_nonzero() == 0 + assert not any(k.startswith(("embed_tokens", "lm_head")) for k in provider.state_dict()) + + +@pytest.mark.parametrize("algorithm,start", [("dflash", 2), ("dspark", 1)]) +def test_native_policy_is_checkpoint_identity_and_suffix_supervision(algorithm, start, monkeypatch): + from aspec.config.draft_algorithm import draft_model_semantics + from aspec.models.draft.auto import AutoDraftModelConfig + + model = wrapper(algorithm=algorithm) + cfg = model.draft_model.config + semantics = draft_model_semantics(cfg) + assert semantics["mtp_conditioning"] == cfg.mtp_conditioning + assert draft_model_semantics(semantics) == semantics + restored = AutoDraftModelConfig.from_dict(cfg.to_dict()) + assert restored.mtp_conditioning == cfg.mtp_conditioning + features = torch.randn(1, 10, 16, requires_grad=True) + anchors, keep = torch.tensor([[3, 5]]), torch.ones(1, 2, dtype=torch.bool) + monkeypatch.setattr(model, "_draft_backbone", lambda *a, **k: (features, anchors, keep, 2)) + ids = torch.arange(12)[None] + output = model(ids, None, torch.ones_like(ids), torch.randn(32, 16), + last_hidden_states=torch.randn(1, 12, 16), + draft_prefix=MTPPrefix(torch.randn(1, 11, 16))) + assert model.prediction_start == start + assert output[4][:start].count_nonzero() == 0 + assert output[4][start:].tolist() == [2] * (5 - start) + output[0].backward() + assert features.grad.view(1, 2, 5, 16)[:, :, :start].count_nonzero() == 0 + if algorithm == "dspark": + assert "tau_probabilistic" not in output[5] + assert "suffix_tau_probabilistic" in output[5] + assert len([k for k in output[5] if k.startswith("suffix_accept_rate")]) == 4 + + +@pytest.mark.parametrize("algorithm", ["dflash", "dspark"]) +@pytest.mark.parametrize("mode", [None, "token", "hidden"]) +@pytest.mark.parametrize("namespace", ["train", "eval"]) +def test_prediction_metrics_identify_suffix_scope(algorithm, mode, namespace, monkeypatch): + from aspec.training.dflash_trainer import DFlashTrainer + from aspec.training.dspark_trainer import DSparkTrainer + + cls = DFlashTrainer if algorithm == "dflash" else DSparkTrainer + trainer = object.__new__(cls) + conditioned = mode is not None + trainer.dflash = SimpleNamespace(mtp_conditioning={"mode": mode} if conditioned else None) + trainer._anchor_slot_offset = int(algorithm == "dflash") + int(conditioned) + trainer._extra_loss_component_keys = [] + trainer.loss_decay_gamma, trainer.global_step = 1.0, 7 + trainer.optimizer = SimpleNamespace(get_learning_rate=lambda: 0.001) + monkeypatch.setattr("aspec.training.dflash_trainer.dist.get_rank", lambda: 0) + monkeypatch.setattr("aspec.training.dflash_trainer.dist.all_reduce", lambda *a, **k: None) + offset = trainer._anchor_slot_offset + losses = torch.tensor([0.0] * offset + [1.0, 2.0, 3.0]) + accuracies = torch.tensor([0.0] * offset + [0.5, 0.5, 0.5]) + counts = torch.tensor([0.0] * offset + [1.0, 1.0, 1.0]) + if namespace == "train": + result = trainer._aggregate_metrics([dict(loss_per_position=losses, acc_per_position=accuracies, + count_per_position=counts)], 7, grad_norm=torch.tensor(0.0)) + else: + result = trainer._aggregate_eval_metrics([dict(loss_pp=losses, acc_pp=accuracies, count_pp=counts)]) + prefix = "suffix_" if conditioned else "" + assert result[f"{namespace}/{prefix}simulated_acc_len"] == pytest.approx(0.875) + assert result[f"{namespace}/{prefix}ploss_0"] == 1.0 + assert result[f"{namespace}/{prefix}acc_0"] == 0.5 + opposite = "" if conditioned else "suffix_" + for name in ["simulated_acc_len", "ploss_0", "acc_0"]: + assert f"{namespace}/{opposite}{name}" not in result + + +@pytest.mark.parametrize("metric", ["eval/simulated_acc_len", "eval/suffix_simulated_acc_len"]) +def test_checkpoint_selection_persists_the_observed_metric_scope(tmp_path, metric): + from aspec.controller.eval import eval_selection_metric, update_checkpoint_eval_meta + + checkpoint = tmp_path / "iter_0000002" + checkpoint.mkdir() + (checkpoint / "meta.json").write_text(json.dumps({"global_step": 1, "eval/simulated_acc_len": 0.1})) + metrics = {"eval/avg_loss": 1.0, "eval/avg_acc": 0.5, metric: 0.875} + assert eval_selection_metric(metrics) == metric + assert update_checkpoint_eval_meta(str(tmp_path), 1, metrics, -float("inf")) == 0.875 + assert (tmp_path / "best_checkpointed_iteration.txt").read_text() == "2" + for path in [checkpoint / "meta.json", tmp_path / "best_meta.json"]: + metadata = json.loads(path.read_text()) + assert metadata[metric] == 0.875 + assert metadata["selection_metric"] == metric + if metric == "eval/suffix_simulated_acc_len": + assert "eval/simulated_acc_len" not in metadata + + +@pytest.mark.parametrize("algorithm", ["dflash", "dspark"]) +def test_strict_native_asset_and_weights_only_transition(tmp_path, monkeypatch, algorithm): + from safetensors.torch import save_file + + from aspec.commands import convert_to_hf + from aspec.draft_artifact import BlockDraftConsumer, load_native_draft + from aspec.training.checkpoint import load_initial_draft_weights + + model = _nonzero_model(algorithm).eval() + draft = model.draft_model + # A real source initialized from the native predictor, without a Target head copy. + from aspec.models.draft.mtp import NativeMTPConfig, NativeMTPDraftModel + + native = NativeMTPDraftModel(NativeMTPConfig(native_config=native_config())) + state = {native.native_key(k): v for k, v in native.state_dict().items()} + monkeypatch.setattr("aspec.models.draft.base.load_tensor_from_pretrained", lambda _, k: state[k]) + draft.initialize_mtp_conditioning("unused") + inputs = _fusion_inputs() + inputs.pop("draft_prefix") + backbone_outputs = [] + handle = draft.register_forward_hook(lambda _, args, output: backbone_outputs.append(output)) + with torch.no_grad(): + before = model(**inputs) + handle.remove() + root = tmp_path / "iter_0000003" + (root / "model").mkdir(parents=True) + (root / "model/.metadata").touch() + (root / "draft_config.json").write_text(json.dumps(draft.config.to_dict())) + monkeypatch.setattr(convert_to_hf, "_load_fsdp_state_dict", lambda _: { + "draft_model." + key: value for key, value in draft.state_dict().items()}) + output = tmp_path / "native" + manifest = convert_to_hf.convert_checkpoint(input_dir=root, output_dir=output, artifact_format="aspec") + assert manifest["draft_structure"]["loader_contract"]["proposal_entrypoint"] is None + restored = load_native_draft(output) + assert not restored.embed_tokens.weight.requires_grad + assert not any(p.requires_grad for p in restored.mtp_source.parameters()) + for key, value in draft.state_dict().items(): + torch.testing.assert_close(restored.state_dict()[key], value, rtol=0, atol=0) + restored_model = type(model)(restored, block_size=5, num_anchors=2, attention_backend="sdpa").eval() + handle = restored.register_forward_hook(lambda _, args, output: backbone_outputs.append(output)) + with torch.no_grad(): + torch.testing.assert_close(restored_model(**inputs), before, rtol=0, atol=0) + handle.remove() + torch.testing.assert_close(backbone_outputs[1], backbone_outputs[0], rtol=0, atol=0) + logits, _ = _reference_logits_and_loss(model, backbone_outputs[0], inputs) + restored_logits, _ = _reference_logits_and_loss(restored_model, backbone_outputs[1], inputs) + torch.testing.assert_close(restored_logits, logits, rtol=0, atol=0) + with pytest.raises(ValueError, match="MTP-aware"): + BlockDraftConsumer(restored, block_size=5) + with pytest.raises(NotImplementedError, match="native"): + convert_to_hf._serving_config(draft.config.to_dict()) + plain = tmp_path / "plain" + plain.mkdir() + cfg = {**draft.config.to_dict(), "mtp_conditioning": None} + (plain / "config.json").write_text(json.dumps(cfg)) + tensors = {k: v for k, v in draft.state_dict().items() + if not k.startswith(("mtp_source.", "mtp_fusion."))} + save_file(tensors, str(plain / "model.safetensors")) + load_initial_draft_weights(draft, str(plain)) + tensors["mtp_fusion.proj.weight"] = draft.mtp_fusion.proj.weight + save_file(tensors, str(plain / "model.safetensors")) + with pytest.raises(ValueError, match="undeclared"): + load_initial_draft_weights(draft, str(plain)) + + +@pytest.mark.parametrize("key,value", [ + ("micro_batch_size", 2), ("draft_context_parallel_size", 2), + ("dflash_packing", True), ("compile_model", True), ("fsdp_strategy", "FULL_SHARD"), +]) +def test_conditioned_plan_rejects_unimplemented_combinations(key, value, monkeypatch): + from pathlib import Path + + from aspec.config.train_config import load_config + from aspec.planning import compile_execution_plan + + monkeypatch.setenv("ASPEC_TRAIN_DATA", "examples/data/sample_conversations.jsonl") + monkeypatch.setenv("ASPEC_EVAL_DATA", "examples/data/sample_conversations.jsonl") + path = Path(__file__).parents[1] / "configs/recipes/qwen3_5_9b_dmtp.yaml" + config = load_config(str(path), save_snapshot=False) + plan = compile_execution_plan(config) + assert plan.workload.objective["prediction_start"] == 1 + assert plan.workload.objective["suffix_prediction_budget"] == 6 + config.training[key] = value + with pytest.raises(ValueError): + compile_execution_plan(config) + + +def test_native_condition_never_reads_future_tokens_or_target_states(monkeypatch): + from aspec.models.dmtp import FrozenMTPFeatures + from aspec.models.draft.mtp import NativeMTPConfig, NativeMTPDraftModel + + native = NativeMTPDraftModel(NativeMTPConfig(native_config=native_config())) + tensors = {native.native_key(k): v for k, v in native.state_dict().items()} + monkeypatch.setattr("aspec.models.draft.base.load_tensor_from_pretrained", lambda _, k: tensors[k]) + provider = FrozenMTPFeatures(native_config()) + provider.initialize_native_weights("unused") + ids = torch.randint(0, 32, (1, 8)) + hidden = torch.randn(1, 8, 16) + positions = torch.arange(8)[None] + before = provider(hidden, native.embed_tokens(ids), positions) + changed_ids, changed_hidden = ids.clone(), hidden.clone() + changed_ids[:, 4:] = (changed_ids[:, 4:] + 1) % 32 + changed_hidden[:, 4:] += 10 + after = provider(changed_hidden, native.embed_tokens(changed_ids), positions) + torch.testing.assert_close(before[:, :4], after[:, :4], rtol=0, atol=0) + + +@pytest.mark.parametrize("algorithm", ["dflash", "dspark"]) +def test_zero_hidden_residual_matches_token_conditioned_workload(algorithm): + token = wrapper(use_hidden=False, algorithm=algorithm) + hidden = wrapper(use_hidden=True, algorithm=algorithm) + with torch.no_grad(): + for name, parameter in token.draft_model.named_parameters(): + if not parameter.is_meta: + hidden.draft_model.get_parameter(name).copy_(parameter) + ids = torch.arange(12)[None] + kwargs = dict(input_ids=ids, hidden_states_list=torch.randn(1, 12, 16), + loss_mask=torch.ones_like(ids), lm_head_weight=torch.randn(32, 16), + last_hidden_states=torch.randn(1, 12, 16), + injected_anchors=torch.tensor([[3, 5]]), + injected_keep_mask=torch.ones(1, 2, dtype=torch.bool)) + expected = token(**kwargs, draft_prefix=MTPPrefix()) + actual = hidden(**kwargs, draft_prefix=MTPPrefix(torch.randn(1, 11, 16))) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + actual[0].backward() + assert hidden.draft_model.mtp_fusion.proj.weight.grad.abs().sum() > 0 + assert all(p.grad is None for p in hidden.draft_model.mtp_source.parameters()) diff --git a/tests/test_evaluation_contract.py b/tests/test_evaluation_contract.py index 4ac3db8..f02d147 100644 --- a/tests/test_evaluation_contract.py +++ b/tests/test_evaluation_contract.py @@ -65,7 +65,8 @@ def test_eval_rebatch_preserves_document_length_and_local_cp_ownership(cp_size): torch.testing.assert_close(teacher[:, :hi - lo], sample["last_hidden_states"]) -def test_exact_resume_retains_eval_history_only_through_selected_frontier(tmp_path, monkeypatch): +@pytest.mark.parametrize("metric_key", ["eval/simulated_acc_len", "eval/suffix_simulated_acc_len"]) +def test_exact_resume_retains_eval_history_only_through_selected_frontier(tmp_path, monkeypatch, metric_key): from aspec.controller import eval as evaluation target = tmp_path / "target" @@ -82,8 +83,8 @@ def test_exact_resume_retains_eval_history_only_through_selected_frontier(tmp_pa load_path=str(checkpoint)) identity = evaluation.eval_cache_identity(args, "sha256:eval", 2) (previous / "evaluation.json").write_text(json.dumps({"cache_identity": identity, "history": [ - {"step": 8, "checkpoint": str(checkpoint), "metrics": {"eval/simulated_acc_len": 0.3}}, - {"step": 16, "checkpoint": str(later), "metrics": {"eval/simulated_acc_len": 0.9}}, + {"step": 8, "checkpoint": str(checkpoint), "metrics": {metric_key: 0.3}}, + {"step": 16, "checkpoint": str(later), "metrics": {metric_key: 0.9}}, ]})) controller, group = Mock(), Mock() controller.get_dataset_admission_evidence.remote.return_value = {"eval": {"dataset_fingerprint": "sha256:eval"}} diff --git a/tests/test_release_recipes.py b/tests/test_release_recipes.py index b96960f..98a1a4c 100644 --- a/tests/test_release_recipes.py +++ b/tests/test_release_recipes.py @@ -83,7 +83,7 @@ def test_qwen8b_recipes_plan_without_target_override(name, monkeypatch): def test_catalog_exactly_covers_formal_files(): entries = recipes.catalog() - assert len(entries) == 12 + assert len(entries) == 13 declared = {p for entry in entries.values() for p in entry["configs"]} assert declared == { str(p.relative_to(recipes.ROOT)) for p in (recipes.ROOT / "configs/recipes").glob("*.yaml") diff --git a/tests/test_training_resume_rng.py b/tests/test_training_resume_rng.py index cb6a223..6eb390b 100644 --- a/tests/test_training_resume_rng.py +++ b/tests/test_training_resume_rng.py @@ -16,12 +16,20 @@ from types import SimpleNamespace from unittest.mock import Mock +import pytest import torch from aspec.training import checkpoint, dflash_trainer +from aspec.training.dspark_trainer import DSparkTrainer -def test_training_rng_is_restored_after_target_helper_initialization(monkeypatch): +@pytest.mark.parametrize("trainer_class,anchor_offset", [ + (dflash_trainer.DFlashTrainer, 1), (DSparkTrainer, 0), +]) +@pytest.mark.parametrize("explicit_none", [False, True]) +def test_training_rng_is_restored_after_target_helper_initialization( + monkeypatch, trainer_class, anchor_offset, explicit_none, +): saved_rng = torch.get_rng_state().clone() class Draft(torch.nn.Linear): @@ -36,16 +44,20 @@ def freeze_embedding(self): def keep_csa_projection_precision(self): pass - trainer = object.__new__(dflash_trainer.DFlashTrainer) + # Legacy Draft/wrapper extensions deliberately have no MTP-specific methods + # or attributes. Disabled optional features must not require those hooks. + trainer = object.__new__(trainer_class) + trainer._anchor_slot_offset = anchor_offset trainer.args = SimpleNamespace( lr_total_steps=32, learning_rate=0.001, max_grad_norm=1, continual_training=False ) trainer.dp_rank, trainer.global_step = 0, 0 trainer.dp_mesh, trainer.fsdp_cpu_offload = None, False trainer.prof = SimpleNamespace(on_init_end=Mock()) - trainer._draft_config_class = lambda **_: SimpleNamespace( - num_target_layers=1, target_hidden_size=2, target_num_hidden_layers=2 - ) + config = dict(num_target_layers=1, target_hidden_size=2, target_num_hidden_layers=2) + if explicit_none: + config["mtp_conditioning"] = None + trainer._draft_config_class = lambda **_: SimpleNamespace(**config) trainer._get_init_weight_context_manager = lambda: nullcontext trainer._build_draft_model = lambda _: Draft(2, 2) @@ -83,3 +95,5 @@ def initialize_target(_path): trainer.init_model({}, "unused") assert trainer.global_step == 16 assert torch.equal(torch.get_rng_state(), saved_rng) + assert trainer._anchor_slot_offset == anchor_offset + assert not trainer._mtp_binding_validated diff --git a/tools/recipe_catalog.json b/tools/recipe_catalog.json index 98e3385..6c0873f 100644 --- a/tools/recipe_catalog.json +++ b/tools/recipe_catalog.json @@ -135,5 +135,16 @@ "target_gpus": 8, "draft_gpus": 8, "native_only_export": false + }, + "dmtp": { + "configs": [ + "configs/recipes/qwen3_5_9b_dmtp.yaml" + ], + "model": "Qwen3.5-9B", + "nodes": 1, + "gpus": 8, + "target_gpus": 4, + "draft_gpus": 4, + "native_only_export": true } } diff --git a/tools/recipes.py b/tools/recipes.py index b20616e..787114a 100644 --- a/tools/recipes.py +++ b/tools/recipes.py @@ -204,6 +204,9 @@ def main(): "tests/test_eagle3_trainer.py", "tests/test_eagle4.py", "tests/test_native_mtp.py", + "tests/test_dmtp.py", + "tests/test_training_resume_rng.py", + "tests/test_evaluation_contract.py", "tests/test_ling_moe.py", "tests/test_multi_draft_plan.py", "tests/test_multinode_tcp_reference.py",