diff --git a/configs/qwen3.5-4b-mtp.json b/configs/qwen3.5-4b-mtp.json new file mode 100644 index 000000000..29a3858db --- /dev/null +++ b/configs/qwen3.5-4b-mtp.json @@ -0,0 +1,37 @@ +{ + "architectures": [ + "Qwen3_5MTPDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoModel": "mtp.Qwen3_5MTPDraftModel" + }, + "attn_output_gate": true, + "bos_token_id": 248043, + "eos_token_id": 248044, + "head_dim": 256, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9216, + "max_position_embeddings": 262144, + "model_type": "qwen3", + "mtp_config": { + "share_lm_head": true + }, + "num_attention_heads": 16, + "num_hidden_layers": 1, + "num_key_value_heads": 4, + "pad_token_id": 248044, + "partial_rotary_factor": 0.25, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 10000000, + "sliding_window": null, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 248320 +} diff --git a/docs/basic_usage/training.md b/docs/basic_usage/training.md index 39efc3e5a..e8c333356 100644 --- a/docs/basic_usage/training.md +++ b/docs/basic_usage/training.md @@ -234,6 +234,7 @@ The checked-in examples are the canonical starting points: | EAGLE3 offline disaggregated | [`qwen3-8b-eagle3-offline-disaggregated.yaml`](../../examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml) | | Ascend NPU DFlash online | [`qwen3.5-4b-dflash-online-npu.yaml`](../../examples/configs/qwen3.5-4b-dflash-online-npu.yaml) | | Ascend NPU Domino online | [`qwen3.5-4b-domino-online-npu.yaml`](../../examples/configs/qwen3.5-4b-domino-online-npu.yaml) | +| Ascend NPU MTP disaggregated | [`qwen3.5-4b-mtp-disaggregated-npu.yaml`](../../examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml) | ## Online and offline data @@ -262,6 +263,7 @@ The unified runtime supports text training in these combinations: | DFlash | Yes, consumer DP | Yes, DP | Yes, consumer DP | | Domino | Yes, consumer DP | Yes, DP | Yes, consumer DP | | DSpark | Yes, consumer DP | Yes, DP | Yes, consumer DP | +| MTP | Yes, consumer DP | Yes, DP | Yes, consumer DP | | P-EAGLE | Yes, consumer DP, batch size 1 | No | No | Unsupported combinations fail explicitly during config validation or run @@ -274,10 +276,10 @@ assembly. In particular: - attention backends are strategy-specific: EAGLE3 accepts `sdpa`, `flex_attention`, `fa`, or offline `usp`; P-EAGLE requires `flex_attention`; DFlash, Domino, and DSpark accept `eager`, `sdpa`, or - `flex_attention`; + `flex_attention`; MTP accepts `eager` or `sdpa`; - P-EAGLE requires `training.batch_size=1` and reuses EAGLE3's server capture schema; -- offline feature training supports EAGLE3, DFlash, Domino, and DSpark; +- offline feature training supports EAGLE3, DFlash, Domino, DSpark, and MTP; - every online run is disaggregated and uses `model.target_backend=sglang`; finite runs may omit both step fields so the producer can publish the exact optimizer horizon derived from the prepared prompt plan; @@ -509,6 +511,23 @@ specforge export --to hf \ Pass `--vocab-mapping /path/to/mapping.pt` when the checkpoint predates the mapping buffers or when you intentionally need to refresh them. +MTP is deployed by merging its trained native head back into the target model. +The merge command accepts the same runtime checkpoint shapes as the generic +exporter (`training_state.pt`, a step/latest directory, or the run output +directory): + +```bash +python scripts/merge_mtp_to_base.py \ + --base-model-path Qwen/Qwen3.5-4B \ + --mtp-checkpoint-path ./outputs/qwen3.5-4b-mtp/qwen3.5-4b-mtp-latest \ + --draft-config configs/qwen3.5-4b-mtp.json \ + --output-path ./exports/Qwen3.5-4B-MTP +``` + +An already-exported HF MTP draft directory can also be passed as +`--mtp-checkpoint-path`; in that case its own `config.json` is used and +`--draft-config` may be omitted. + ## Troubleshooting ### Late OOM or non-finite hidden states on online runs diff --git a/examples/README.md b/examples/README.md index be9b23903..4ed9bf6cc 100644 --- a/examples/README.md +++ b/examples/README.md @@ -38,6 +38,7 @@ NPU, offline, and managed/external-service variants, is in | `examples/configs/qwen3.5-4b-dflash-online-npu.yaml` | Disaggregated NPU SGLang capture | DFlash | | `examples/configs/qwen3.5-4b-dflash-disaggregated-npu.yaml` | Managed local NPU full stack | DFlash | | `examples/configs/qwen3.5-4b-domino-online-npu.yaml` | Disaggregated NPU SGLang capture | Domino | +| `examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml` | Managed local NPU full stack | MTP | Online configs point `data.train_data_path` at raw conversation data. Offline configs expect strategy-specific feature checkpoints in diff --git a/examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml b/examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml new file mode 100644 index 000000000..8b22ab424 --- /dev/null +++ b/examples/configs/qwen3.5-4b-mtp-disaggregated-npu.yaml @@ -0,0 +1,77 @@ +# Single-node managed full stack for Ascend NPU, launched with one +# `specforge train` command. Device ordinals are injected via +# ASCEND_RT_VISIBLE_DEVICES. +model: + target_model_path: "Qwen/Qwen3.5-4B" + draft_model_config: "configs/qwen3.5-4b-mtp.json" + target_backend: sglang + trust_remote_code: true + # Qwen3.5-4B nests its text decoder under model.language_model and ties + # lm_head to the embedding, so both keys point at the same tensor. + embedding_key: "model.language_model.embed_tokens.weight" + lm_head_key: "model.language_model.embed_tokens.weight" + torch_dtype: "bfloat16" + sglang_attention_backend: ascend +data: + train_data_path: "./cache/dataset/train_regen.jsonl" + max_length: 32768 + chat_template: "qwen3.5" + build_dataset_num_proc: 32 + cache_dir: "./cache" +training: + strategy: "mtp" + num_epochs: 10 + max_steps: 10000 + batch_size: 2 + accumulation_steps: 4 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1 + attention_backend: "sdpa" + save_interval: 10000 + log_interval: 50 + dist_timeout: 30 + seed: 42 +tracking: + report_to: "tensorboard" +run_id: "qwen3.5-4b-mtp-npu-managed" +output_dir: "./outputs/qwen3.5-4b-mtp-npu-managed" + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 10 + disaggregated: + control_dir: outputs/qwen3.5-4b-mtp-npu-managed/control + consumer_state_dir: outputs/qwen3.5-4b-mtp-npu-managed/consumer-state + backend: mooncake + managed_local: + # Validated 16-card A3 layout: 6 TP=1 capture servers own devices 0-5, + # the 10-rank trainer owns devices 6-15 (global batch = 2*4*10 = 80). + # The pipeline is capture-bound, so scale capture servers first. + trainer_cuda_visible_devices: + ["6", "7", "8", "9", "10", "11", "12", "13", "14", "15"] + mooncake: + # Host segment mounted by the capture-server sink for feature storage. + # 16 GiB fits a 64GB-card host next to the target model and KV cache. + global_segment_size_bytes: 17179869184 + capture_servers: + - port: 30000 + cuda_visible_devices: ["0"] + tp_size: 1 + - port: 30001 + cuda_visible_devices: ["1"] + tp_size: 1 + - port: 30002 + cuda_visible_devices: ["2"] + tp_size: 1 + - port: 30003 + cuda_visible_devices: ["3"] + tp_size: 1 + - port: 30004 + cuda_visible_devices: ["4"] + tp_size: 1 + - port: 30005 + cuda_visible_devices: ["5"] + tp_size: 1 diff --git a/scripts/merge_mtp_to_base.py b/scripts/merge_mtp_to_base.py new file mode 100755 index 000000000..ca695b9c2 --- /dev/null +++ b/scripts/merge_mtp_to_base.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +# coding=utf-8 +"""Merge a trained MTP draft checkpoint back into the base target model. + +Thin CLI wrapper around ``specforge.export.mtp.merge_mtp_into_base``; the merge +logic and the per-family key mapping live in the package. + +Example: + python scripts/merge_mtp_to_base.py \ + --base-model-path PATH/TO/Qwen3.5-4B \ + --mtp-checkpoint-path PATH/TO/outputs/qwen3.5-4b-mtp/RUN-latest \ + --draft-config configs/qwen3.5-4b-mtp.json \ + --output-path PATH/TO/Qwen3.5-4B-MTP \ + --key-format sglang +""" + +import argparse + +from specforge.export.mtp import merge_mtp_into_base + + +def main(): + parser = argparse.ArgumentParser( + description="Merge trained MTP weights back into the base Qwen3.5 model." + ) + parser.add_argument( + "--base-model-path", + type=str, + required=True, + help="Path to the original Qwen3.5 base model checkpoint.", + ) + parser.add_argument( + "--mtp-checkpoint-path", + type=str, + required=True, + help=( + "SpecForge runtime checkpoint/output path, or an already-exported " + "HF MTP draft directory." + ), + ) + parser.add_argument( + "--draft-config", + type=str, + default=None, + help=( + "Draft config JSON (required for a SpecForge runtime checkpoint; " + "the exported HF directory already contains config.json)." + ), + ) + parser.add_argument( + "--output-path", + type=str, + required=True, + help="Directory to write the merged checkpoint.", + ) + parser.add_argument( + "--key-format", + type=str, + default="sglang", + choices=["sglang", "hf"], + help=( + "MTP key layout. Both 'sglang' and 'hf' produce the flat native " + "layout (mtp.layers.0.* / mtp.norm.weight) that SGLang's flat " + "Qwen3_5ForCausalLMMTP and HF/vLLM MTP modules expect; the argument " + "is kept for backward compatibility." + ), + ) + args = parser.parse_args() + + merge_mtp_into_base( + args.base_model_path, + args.mtp_checkpoint_path, + args.output_path, + args.key_format, + draft_config_path=args.draft_config, + ) + + +if __name__ == "__main__": + main() diff --git a/specforge/algorithms/builtin.py b/specforge/algorithms/builtin.py index c2d8527bd..8e4ea7e9a 100644 --- a/specforge/algorithms/builtin.py +++ b/specforge/algorithms/builtin.py @@ -6,6 +6,7 @@ from specforge.algorithms.domino.providers import create_registration as domino from specforge.algorithms.dspark.providers import create_registration as dspark from specforge.algorithms.eagle3.providers import create_registration as eagle3 +from specforge.algorithms.mtp.providers import create_registration as mtp from specforge.algorithms.peagle.providers import create_registration as peagle from specforge.algorithms.registry import AlgorithmRegistry @@ -13,7 +14,7 @@ def builtin_algorithm_registry() -> AlgorithmRegistry: """Return a fresh immutable catalog without module-level mutation.""" - return AlgorithmRegistry((eagle3(), peagle(), dflash(), domino(), dspark())) + return AlgorithmRegistry((eagle3(), peagle(), dflash(), domino(), dspark(), mtp())) __all__ = ["builtin_algorithm_registry"] diff --git a/specforge/algorithms/common/dflash_family_data.py b/specforge/algorithms/common/hidden_states_data.py similarity index 68% rename from specforge/algorithms/common/dflash_family_data.py rename to specforge/algorithms/common/hidden_states_data.py index 99c0df9c8..770ed710c 100644 --- a/specforge/algorithms/common/dflash_family_data.py +++ b/specforge/algorithms/common/hidden_states_data.py @@ -1,4 +1,7 @@ -"""Shared DFlash-family normalization and padding adapters.""" +"""Shared hidden-states normalization and padding adapters. + +Used by the DFlash-family (DFlash/Domino/DSpark) and MTP algorithms. +""" from __future__ import annotations @@ -9,6 +12,7 @@ NORMALIZER_ID = "dflash_family_offline_v1" DSPARK_NORMALIZER_ID = "dspark_offline_v1" +MTP_NORMALIZER_ID = "mtp_offline_v1" def _normalize_hidden_states( @@ -185,15 +189,99 @@ def collate(features): return collate +def normalize_mtp_offline_sample(raw, max_len: int): + """Normalize MTP capture tensors (no aux-layer concat, final hidden only).""" + + input_ids = raw["input_ids"][:max_len].unsqueeze(0) + loss_mask = raw["loss_mask"][:max_len].unsqueeze(0) + target_last_hidden_states = _normalize_hidden_states( + raw, + "target_last_hidden_states", + max_len, + description="MTP target_last_hidden_states", + ) + lengths = { + input_ids.shape[1], + loss_mask.shape[1], + target_last_hidden_states.shape[1], + } + if len(lengths) != 1: + raise ValueError( + "offline MTP features have mismatched sequence lengths after " + f"truncation: input_ids={input_ids.shape[1]}, " + f"loss_mask={loss_mask.shape[1]}, " + f"target_last_hidden_states={target_last_hidden_states.shape[1]}" + ) + return { + "input_ids": input_ids, + "loss_mask": loss_mask, + "target_last_hidden_states": target_last_hidden_states, + } + + +def build_mtp_offline_reader( + strategy, + hidden_states_path, + *, + run_id, + ttt_length, + max_len, +): + # Transitional runtime import; the composition root will inject this port. + from specforge.runtime.data_plane.offline_reader import OfflineManifestReader + + return OfflineManifestReader( + hidden_states_path, + run_id=run_id, + strategy=strategy, + feature_keys=( + "input_ids", + "loss_mask", + "target_last_hidden_states", + ), + target_repr="hidden_state", + ttt_length=ttt_length, + max_len=max_len, + ) + + +def build_mtp_offline_normalizer(max_len, **_topology): + return partial(normalize_mtp_offline_sample, max_len=max_len) + + +def build_mtp_collator(): + def collate(features): + return pad_and_concatenate_features( + features, + sequence_axes={ + "input_ids": 1, + "loss_mask": 1, + "target_last_hidden_states": 1, + }, + required_keys=( + "input_ids", + "loss_mask", + "target_last_hidden_states", + ), + ) + + return collate + + __all__ = [ "DSPARK_NORMALIZER_ID", + "MTP_NORMALIZER_ID", "NORMALIZER_ID", "build_collator", "build_dspark_collator", "build_dspark_offline_normalizer", "build_dspark_offline_reader", + "build_mtp_collator", + "build_mtp_offline_normalizer", + "build_mtp_offline_reader", "build_offline_normalizer", "build_offline_reader", "normalize_dspark_offline_sample", + "normalize_mtp_offline_sample", "normalize_offline_sample", ] diff --git a/specforge/algorithms/dflash/providers.py b/specforge/algorithms/dflash/providers.py index 2cc119ac9..c5bc5837b 100644 --- a/specforge/algorithms/dflash/providers.py +++ b/specforge/algorithms/dflash/providers.py @@ -8,7 +8,7 @@ empty_options, no_missing_checkpoint_keys, ) -from specforge.algorithms.common.dflash_family_data import ( +from specforge.algorithms.common.hidden_states_data import ( NORMALIZER_ID, build_collator, build_offline_normalizer, diff --git a/specforge/algorithms/domino/providers.py b/specforge/algorithms/domino/providers.py index af6506c41..16dfcca1b 100644 --- a/specforge/algorithms/domino/providers.py +++ b/specforge/algorithms/domino/providers.py @@ -5,7 +5,7 @@ from functools import partial from specforge.algorithms.common.defaults import no_missing_checkpoint_keys -from specforge.algorithms.common.dflash_family_data import ( +from specforge.algorithms.common.hidden_states_data import ( NORMALIZER_ID, build_collator, build_offline_normalizer, diff --git a/specforge/algorithms/dspark/providers.py b/specforge/algorithms/dspark/providers.py index d3e37be3f..c48951cc0 100644 --- a/specforge/algorithms/dspark/providers.py +++ b/specforge/algorithms/dspark/providers.py @@ -8,7 +8,7 @@ empty_options, no_missing_checkpoint_keys, ) -from specforge.algorithms.common.dflash_family_data import ( +from specforge.algorithms.common.hidden_states_data import ( DSPARK_NORMALIZER_ID, build_dspark_collator, build_dspark_offline_normalizer, diff --git a/specforge/algorithms/mtp/__init__.py b/specforge/algorithms/mtp/__init__.py new file mode 100644 index 000000000..2b409daf0 --- /dev/null +++ b/specforge/algorithms/mtp/__init__.py @@ -0,0 +1,5 @@ +"""MTP algorithm registration.""" + +from specforge.algorithms.mtp.providers import create_registration + +__all__ = ["create_registration"] diff --git a/specforge/algorithms/mtp/providers.py b/specforge/algorithms/mtp/providers.py new file mode 100644 index 000000000..bde2aa7b3 --- /dev/null +++ b/specforge/algorithms/mtp/providers.py @@ -0,0 +1,318 @@ +"""Built-in MTP (Multi-Token Prediction) registration and executable providers. + +MTP fine-tunes the single-layer draft head shipped natively with Qwen3.5-style +checkpoints. The training signal is the target model's *final* (post-norm) +hidden state — unlike EAGLE3 it needs no aux-layer concat, so the offline +capture persists only ``input_ids`` / ``loss_mask`` / ``target_last_hidden_states`` +(the streaming layout still carries the aux tensor the server patch always +produces; the strategy ignores it). + +Draft construction initializes from the *native* ``mtp.*`` weights inside the +target checkpoint (fine-tuning) and shares + freezes the target embedding and +lm_head, mirroring the serving layout consumed by SGLang's +``Qwen3_5ForCausalLMMTP``. +""" + +from __future__ import annotations + +from functools import partial + +from specforge.algorithms.common.defaults import ( + empty_options, + no_missing_checkpoint_keys, + one_loss_token, + online_needs_input_tools, +) +from specforge.algorithms.common.hidden_states_data import ( + MTP_NORMALIZER_ID, + build_mtp_collator, + build_mtp_offline_normalizer, + build_mtp_offline_reader, +) +from specforge.algorithms.common.providers import ( + AlgorithmProviders, + DraftConfigProvider, + ModelProvider, + OfflineCaptureLayout, + OfflineDataProvider, + ServerCaptureLayout, + ServerStreamingProvider, + StepProvider, + make_registration, +) +from specforge.algorithms.contracts import ( + AlgorithmCapabilities, + AlgorithmSpec, + DraftRequirement, + FeatureContract, + FeatureMode, + OfflineStorageContract, +) + +ALGORITHM_NAME = "mtp" +DRAFT_ARCHITECTURE = "Qwen3_5MTPDraftModel" + +# MTP persists no aux-layer tensor, but the capture plan requires a non-empty +# layer list; a single layer keeps the (discarded) aux capture cheap. +_CAPTURE_LAYER_IDS = [1] + + +def build_step(wrapped_model, *, target_head=None, **_options): + del target_head + from specforge.training.strategies.base import MTPTrainStrategy + + return MTPTrainStrategy(wrapped_model) + + +def resume_contract(_config, draft_model, training_model): + """Persist resolved MTP model and objective semantics.""" + + mtp_config = getattr(draft_model.config, "mtp_config", None) or {} + return { + "mtp_draft_num_hidden_layers": int( + getattr(draft_model.config, "num_hidden_layers", 1) + ), + "mtp_draft_vocab_size": int(getattr(draft_model.config, "vocab_size", 0)), + "mtp_share_lm_head": bool(mtp_config.get("share_lm_head", True)), + "mtp_attention_backend": str( + getattr(draft_model.config, "_attn_implementation", "") + ), + } + + +def _init_from_native_mtp(cfg, draft_model) -> None: + """Initialize the draft's ``mtp.*`` weights from the target checkpoint. + + Qwen3.5-style target checkpoints ship a native MTP head whose keys match + the draft's flat ``mtp.*`` layout. The target lm_head may be shared rather + than duplicated under that prefix. Loading the required keys turns training + into fine-tuning of the native head — the only training mode for this + algorithm. Initialization is strict by default: missing or partial native + state fails rather than silently leaving trainable tensors randomized. + """ + + if cfg.model.draft_checkpoint_path: + print( + "[mtp] native target initialization skipped; weights come from the " + "warm-start draft checkpoint." + ) + return + + from specforge.modeling.target.checkpoint import ( + load_selected_tensors, + resolve_checkpoint_dir, + ) + + target_path = cfg.model.target_model_path + prefix = draft_model.NATIVE_KEY_PREFIX + try: + checkpoint_dir = resolve_checkpoint_dir( + target_path, cache_dir=cfg.model.cache_dir + ) + native_mtp = load_selected_tensors( + checkpoint_dir, lambda key: key.startswith(prefix) + ) + scan_error = None + except Exception as exc: # pragma: no cover - depends on target checkpoint + native_mtp = {} + scan_error = exc + + if native_mtp: + model_keys = set(draft_model.native_state_dict()) + required_keys = set(draft_model.required_native_state_keys()) + extra_keys = set(draft_model.allowed_extra_native_state_keys()) + loaded_keys = set(native_mtp) + missing = sorted(required_keys - loaded_keys) + unexpected = sorted(loaded_keys - model_keys - extra_keys) + if missing or unexpected: + details = [] + if missing: + details.append(f"missing required native keys: {missing}") + if unexpected: + details.append(f"unexpected native keys: {unexpected}") + raise RuntimeError( + f"[mtp] incompatible native {prefix}* state in {target_path}: " + + "; ".join(details) + ) + draft_model.load_state_dict(native_mtp, strict=False) + print( + f"[mtp] initialized {len(native_mtp)} native {prefix}* weights from " + f"{target_path} (native-MTP fine-tune)." + ) + return + + detail = f" (scan failed: {scan_error})" if scan_error is not None else "" + raise RuntimeError( + f"[mtp] no native {prefix}* weights found in {target_path}{detail}. MTP " + "training fine-tunes the native MTP head shipped with the target " + "checkpoint and does not start from random initialization by default. " + "Point model.target_model_path at a checkpoint that ships native MTP " + "weights (e.g. Qwen3.5), or set model.draft_checkpoint_path to resume " + "from a trained MTP draft." + ) + + +def _share_target_embeddings(cfg, draft_model, torch_dtype) -> None: + """Share (and freeze) the target checkpoint's embed_tokens and lm_head.""" + + from specforge.modeling.target.target_utils import TargetEmbeddingsAndHead + + target_components = TargetEmbeddingsAndHead.from_pretrained( + cfg.model.target_model_path, + embed_key=cfg.model.embedding_key, + lm_head_key=cfg.model.lm_head_key, + cache_dir=cfg.model.cache_dir, + device="cpu", + dtype=torch_dtype, + trust_remote_code=cfg.model.trust_remote_code, + ) + draft_model.share_target_embeddings( + target_components.embed_tokens.weight, + lm_head_weight=target_components.lm_head.weight, + ) + print("[mtp] shared target embed_tokens/lm_head with the draft (frozen).") + + +def build_draft(cfg, draft_config): + import torch + + from specforge.modeling.auto import AutoDraftModel + from specforge.training.model_loading import warm_start_draft_model + from specforge.utils import get_local_device + + torch_dtype = getattr(torch, cfg.model.torch_dtype) + draft_config._attn_implementation = cfg.training.attention_backend + draft_model = AutoDraftModel.from_config(draft_config, torch_dtype=torch_dtype) + + _init_from_native_mtp(cfg, draft_model) + _share_target_embeddings(cfg, draft_model, torch_dtype) + + if cfg.model.draft_checkpoint_path: + warm_start_draft_model( + draft_model, + cfg.model.draft_checkpoint_path, + draft_config=draft_config, + strategy=cfg.training.strategy, + cache_dir=cfg.model.cache_dir, + trust_remote_code=cfg.model.trust_remote_code, + ) + return draft_model.to(device=get_local_device(), dtype=torch_dtype) + + +def build_training_model(config, draft_model, draft_config, target_config, tokenizer): + from specforge.algorithms.model_providers import AlgorithmModelParts + from specforge.core.mtp import OnlineMTPModel + + return AlgorithmModelParts( + model=OnlineMTPModel(draft_model=draft_model), + capture_layers=None, + ) + + +def resolve_capture_layers(config, draft_config, target_config): + return list(_CAPTURE_LAYER_IDS) + + +def create_registration(): + return make_registration(algorithm_spec(), algorithm_providers()) + + +def algorithm_spec() -> AlgorithmSpec: + ready = { + "input_ids", + "loss_mask", + "target_last_hidden_states", + } + return AlgorithmSpec( + name=ALGORITHM_NAME, + draft=DraftRequirement( + compatible_architectures={DRAFT_ARCHITECTURE}, + default_architecture=DRAFT_ARCHITECTURE, + ), + feature_contracts=( + FeatureContract( + mode=FeatureMode.OFFLINE, + modality="text", + required_tensors=ready, + allowed_target_representations={"hidden_state"}, + default_target_representation="hidden_state", + storage=OfflineStorageContract( + format="specforge_hidden_states_v1", + required_tensors=ready, + normalizer=MTP_NORMALIZER_ID, + ), + ), + FeatureContract( + mode=FeatureMode.STREAMING, + modality="text", + required_tensors=ready, + allowed_target_representations={"hidden_state"}, + default_target_representation="hidden_state", + ), + ), + capabilities=AlgorithmCapabilities( + attention_backends={"eager", "sdpa"}, + ), + ) + + +def algorithm_providers() -> AlgorithmProviders: + return AlgorithmProviders( + algorithm_name=ALGORITHM_NAME, + step=StepProvider( + build=build_step, + options=empty_options, + resume_contract=resume_contract, + allowed_missing_checkpoint_keys=no_missing_checkpoint_keys, + uses_external_target_head=False, + ), + model=ModelProvider( + draft_config=DraftConfigProvider( + architecture=DRAFT_ARCHITECTURE, + expected_auto_map_model="mtp.Qwen3_5MTPDraftModel", + ), + build_draft=build_draft, + build_training_model=build_training_model, + resolve_capture_layers=resolve_capture_layers, + minimum_loss_tokens=one_loss_token, + needs_input_tools=online_needs_input_tools, + default_dataloader_num_workers=8, + ), + offline=( + OfflineDataProvider( + modality="text", + normalizer_id=MTP_NORMALIZER_ID, + capture_layout=OfflineCaptureLayout( + capture_method="dflash", + aux_feature=None, + last_hidden_feature="target_last_hidden_states", + passthrough=( + ("input_ids", "input_ids"), + ("loss_mask", "loss_mask"), + ), + ), + build_reader=partial(build_mtp_offline_reader, ALGORITHM_NAME), + build_normalizer=build_mtp_offline_normalizer, + build_collator=build_mtp_collator, + ), + ), + server_streaming=( + ServerStreamingProvider( + modality="text", + capture_method="dflash", + target_representation="hidden_state", + layout=ServerCaptureLayout( + aux_feature="hidden_states", + last_hidden_feature="target_last_hidden_states", + passthrough=( + ("input_ids", "input_ids", ()), + ("loss_mask", "loss_mask", ()), + ), + ), + build_collator=build_mtp_collator, + ), + ), + ) + + +__all__ = ["algorithm_providers", "algorithm_spec", "create_registration"] diff --git a/specforge/core/mtp.py b/specforge/core/mtp.py new file mode 100644 index 000000000..230a66edf --- /dev/null +++ b/specforge/core/mtp.py @@ -0,0 +1,144 @@ +# coding=utf-8 +"""Online training wrapper for single-layer MTP (architecture-independent). + +MTP predicts the next token from the current token's embedding plus the target +model's last hidden state. Shift is performed inside this wrapper; the target +backend is expected to return *raw* input_ids and last_hidden_states (DFlash +style), not the pre-shifted output of generate_eagle3_data. +""" + +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class OnlineMTPModel(nn.Module): + """ + Online MTP training wrapper. + + Architecture-agnostic: any registered MTP draft module exposing + ``forward(input_ids, hidden_states, attention_mask, position_ids)`` and a + ``config`` with ``pad_token_id`` can be plugged in (see + ``specforge/modeling/draft/mtp/``). + + Args: + draft_model: The MTP draft model (e.g. ``modeling/draft/mtp/qwen3_5.py``). + ploss_decay: Per-layer loss decay. For a single MTP layer this is + unused, but kept for multi-layer extension. + """ + + def __init__( + self, + draft_model: nn.Module, + ploss_decay: float = 1.0, + ) -> None: + super().__init__() + self.draft_model = draft_model + self.ploss_decay = ploss_decay + + def _shift_for_next_token( + self, + logits: torch.Tensor, + input_ids: torch.Tensor, + loss_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Shift logits/labels/mask to match vLLM speculative decoding. + + In serving, the draft model's input_ids are the target input_ids shifted + right by one (the draft fuses token x_{t+1} with target hidden state + h_t) and it predicts the token after that (x_{t+2}). Training therefore + uses: + - draft input: input_ids[:, 1:] (x_1..x_T, padded) + - label: x_2..x_T followed by a pad (length matches logits) + """ + shift_logits = logits[:, :-1, :].contiguous() + # x_2..x_T has length seq_len-2; pad one position so its length equals + # seq_len-1 (same as shift_logits). The padded position is ignored. + shift_labels = F.pad(input_ids[:, 2:], (0, 1), value=-100).contiguous() + shift_mask = F.pad(loss_mask[:, 2:], (0, 1), value=0).contiguous() + return shift_logits, shift_labels, shift_mask + + def forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + loss_mask: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: + """ + Args: + input_ids: raw token ids, [batch, seq_len]. + hidden_states: target model last hidden states, [batch, seq_len, hidden]. + loss_mask: [batch, seq_len]. + attention_mask: optional padding mask, [batch, seq_len]. + position_ids: optional position ids, [batch, seq_len]. + + Returns: + loss: scalar weighted loss. + acc_corrects: per-layer per-position correct tensors. + acc_denoms: per-layer per-position denominator tensors. + """ + # Draft input is the target sequence shifted right by one. The last + # position is padded because there is no x_{T+1}; the corresponding + # logit is dropped in _shift_for_next_token. + pad_token_id = getattr(self.draft_model.config, "pad_token_id", 0) + shifted_input_ids = F.pad(input_ids[:, 1:], (0, 1), value=pad_token_id) + + # The padding mask must follow the same shift so the synthetic pad token + # at the last position is not attended to. + if attention_mask is not None: + shifted_attention_mask = F.pad(attention_mask[:, 1:], (0, 1), value=0).to( + attention_mask.dtype + ) + else: + shifted_attention_mask = None + + # Serving evaluates the shifted draft token x[t+1] at its own position + # p[t+1], even though it is fused with the target hidden state h[t]. + # Preserve caller-supplied offsets (for packed/non-zero-based sequences) + # and give the synthetic final token the next monotonic position. + batch_size, seq_len = input_ids.shape + if position_ids is None: + position_ids = ( + torch.arange(seq_len, dtype=torch.long, device=input_ids.device) + .unsqueeze(0) + .expand(batch_size, -1) + ) + elif position_ids.shape != input_ids.shape: + raise ValueError( + "position_ids must have the same [batch, seq_len] shape as " + f"input_ids; got {tuple(position_ids.shape)} and " + f"{tuple(input_ids.shape)}" + ) + shifted_position_ids = torch.cat( + (position_ids[:, 1:], position_ids[:, -1:] + 1), dim=1 + ) + + outputs = self.draft_model( + input_ids=shifted_input_ids, + hidden_states=hidden_states, + attention_mask=shifted_attention_mask, + position_ids=shifted_position_ids, + ) + logits = outputs.logits + + shift_logits, shift_labels, shift_mask = self._shift_for_next_token( + logits, input_ids, loss_mask + ) + + flat_logits = shift_logits.view(-1, shift_logits.size(-1)) + flat_labels = shift_labels.view(-1) + losses = F.cross_entropy(flat_logits, flat_labels, reduction="none") + losses = losses * shift_mask.view(-1).float() + loss = losses.sum() / shift_mask.sum().clamp_min(1) + + with torch.no_grad(): + preds = shift_logits.argmax(dim=-1) + corrects = (preds == shift_labels).float() * shift_mask.float() + denoms = shift_mask.float() + + # Single-layer MTP: wrap in length-1 lists for E1 evaluator compatibility. + return loss, [corrects], [denoms] diff --git a/specforge/export/mtp.py b/specforge/export/mtp.py new file mode 100644 index 000000000..49a665066 --- /dev/null +++ b/specforge/export/mtp.py @@ -0,0 +1,381 @@ +# coding=utf-8 +"""Merge a trained MTP draft checkpoint back into the base target checkpoint. + +This module owns the MTP-specific merge policy (native key prefix handling, +shared-embedding backfill, config patching). The model-independent merge +machinery — copying non-weight files, replacing keys by prefix, shard/index +writing — lives in ``specforge/modeling/target/checkpoint.py``. +Architecture-specific knowledge (target-side embed/lm_head key candidates, +native key prefix) comes from the registered MTP draft class — see +``specforge/modeling/draft/mtp/``. +""" + +from __future__ import annotations + +import glob +import json +import os +import shutil +from typing import Dict, List, Optional, Tuple + +import torch + +from specforge.modeling.target.checkpoint import ( + load_selected_tensors, + load_tensors_by_keys, + merge_state_into_checkpoint, +) + + +def _default_key_candidates() -> Tuple[List[str], List[str], str]: + """Base-class defaults, imported lazily to keep this module import-light.""" + + from specforge.modeling.draft.mtp.base import MTPDraftModel + + return ( + list(MTPDraftModel.TARGET_EMBED_KEY_CANDIDATES), + list(MTPDraftModel.TARGET_HEAD_KEY_CANDIDATES), + MTPDraftModel.NATIVE_KEY_PREFIX, + ) + + +def _resolve_key_candidates( + draft_config_source: str, +) -> Tuple[List[str], List[str], str]: + """Return (embed candidates, head candidates, native prefix) for the draft. + + Reads the draft ``config.json`` and resolves its + ``architectures[0]`` through the draft registry, so each MTP family can + override its target-side key candidates on the draft class. + """ + + embed, head, prefix = _default_key_candidates() + config_path = ( + draft_config_source + if os.path.isfile(draft_config_source) + else os.path.join(draft_config_source, "config.json") + ) + if os.path.exists(config_path): + try: + with open(config_path, "r") as f: + architectures = json.load(f).get("architectures") or [] + if architectures: + from specforge.modeling.draft.registry import DRAFT_REGISTRY + + draft_cls = DRAFT_REGISTRY.get(architectures[0]) + if draft_cls is not None: + embed = list( + getattr(draft_cls, "TARGET_EMBED_KEY_CANDIDATES", embed) + ) + head = list(getattr(draft_cls, "TARGET_HEAD_KEY_CANDIDATES", head)) + prefix = getattr(draft_cls, "NATIVE_KEY_PREFIX", prefix) + except Exception as exc: # pragma: no cover - defensive + print(f" warning: could not resolve draft key candidates: {exc}") + return embed, head, prefix + + +def convert_mtp_keys( + state_dict: Dict[str, torch.Tensor], fmt: str, prefix: str = "mtp." +) -> Dict[str, torch.Tensor]: + """Convert MTP weight keys to the requested output format. + + Training already saves the flat native layout that both SGLang and + HF/vLLM MTP modules expect, so ``sglang`` and ``hf`` both return it + unchanged; ``fmt`` is kept for backward compatibility. A legacy nested + layout (``mtp.model.layers.0.*``) is normalized to flat. + """ + + converted = {} + for k, v in state_dict.items(): + # Normalize legacy nested keys (mtp.model.layers.* -> mtp.layers.*). + if k.startswith(f"{prefix}model.layers."): + new_k = k.replace(f"{prefix}model.layers.", f"{prefix}layers.", 1) + elif k == f"{prefix}model.norm.weight": + new_k = f"{prefix}norm.weight" + # Promote bare embed_tokens / lm_head saved by the training script to the + # native namespace expected by vLLM/SGLang. + elif k == "embed_tokens.weight": + new_k = f"{prefix}embed_tokens.weight" + elif k == "lm_head.weight": + new_k = f"{prefix}lm_head.weight" + else: + new_k = k + converted[new_k] = v + return _unshare_storage(converted) + + +def _unshare_storage(state: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """Clone tensors whose storage is aliased under another key. + + A tied target shares one Parameter between ``embed_tokens.weight`` and + ``mtp.lm_head.weight``; after promotion both keys keep aliasing the same + storage, which safetensors' ``save_file`` rejects ("tensors share + memory"). Cloning the later alias preserves the values while giving every + key its own storage. + """ + + seen: set[int] = set() + out: Dict[str, torch.Tensor] = {} + for key, value in state.items(): + ptr = value.untyped_storage().data_ptr() + if ptr in seen: + value = value.clone() + else: + seen.add(ptr) + out[key] = value + return out + + +def _find_base_key(state_dict: Dict[str, torch.Tensor], *candidates: str) -> str | None: + """Return the first candidate key that exists in ``state_dict``.""" + + for key in candidates: + if key in state_dict: + return key + return None + + +def _copy_shared_embeddings( + base_state: Dict[str, torch.Tensor], + mtp_state: Dict[str, torch.Tensor], + tie_word_embeddings: bool, + embed_key_candidates: List[str], + head_key_candidates: List[str], + prefix: str, +) -> Dict[str, torch.Tensor]: + """Copy base embed_tokens/lm_head into the MTP state if they are missing. + + During training the draft model typically shares ``embed_tokens`` and + ``lm_head`` with the target model, so the saved MTP checkpoint does not + contain those tensors. vLLM/SGLang, however, instantiate their own + ``mtp.embed_tokens`` (and a separate ``lm_head`` when weights are not tied), + and expect them in the checkpoint. Copying them from the base model keeps + the merged checkpoint self-contained and avoids random-initialization of the + MTP input/output embeddings at serving time. + """ + + embed_target = f"{prefix}embed_tokens.weight" + head_target = f"{prefix}lm_head.weight" + + if embed_target not in mtp_state: + embed_key = _find_base_key(base_state, *embed_key_candidates) + if embed_key: + mtp_state[embed_target] = base_state[embed_key] + print(f" copied {embed_key} -> {embed_target}") + else: + print( + " warning: base embed_tokens.weight not found; " + f"{embed_target} will be randomly initialized" + ) + + if not tie_word_embeddings and head_target not in mtp_state: + lm_head_key = _find_base_key(base_state, *head_key_candidates) + if lm_head_key: + mtp_state[head_target] = base_state[lm_head_key] + print(f" copied {lm_head_key} -> {head_target}") + else: + print( + " warning: base lm_head.weight not found; " + f"{head_target} will be randomly initialized" + ) + + return mtp_state + + +def _patch_text_config(base_config: dict, draft_config: dict) -> dict: + """Ensure base text_config contains MTP-critical dims from the draft config. + + Some Qwen3.5 base checkpoints omit ``head_dim`` in ``text_config``; vLLM's + ``Qwen3_5TextConfig`` then falls back to its default (``head_dim=256``), + which mismatches the trained MTP weights (e.g. q_norm/k_norm shape 128). + Only the structural dims that must agree between base and draft are synced. + """ + + keys_to_sync = [ + "head_dim", + "hidden_size", + "intermediate_size", + "num_attention_heads", + "num_key_value_heads", + ] + + target = base_config + if "text_config" in base_config: + target = base_config["text_config"] + + source = draft_config + if "text_config" in draft_config: + source = draft_config["text_config"] + + for key in keys_to_sync: + if key not in source: + continue + old = target.get(key) + new = source[key] + if old != new: + target[key] = new + print(f" overriding text_config.{key}: {old} -> {new}") + + return base_config + + +def _load_first_checkpoint(checkpoint_dir: str) -> Dict[str, torch.Tensor]: + """Load every tensor of a single-file checkpoint directory.""" + + safetensors = glob.glob(os.path.join(checkpoint_dir, "*.safetensors")) + bins = glob.glob(os.path.join(checkpoint_dir, "*.bin")) + if safetensors: + return load_selected_tensors(checkpoint_dir, lambda _key: True) + if bins: + return torch.load(bins[0], map_location="cpu", weights_only=True) + raise FileNotFoundError(f"No safetensors/bin weights found in {checkpoint_dir}") + + +def _has_model_weights(path: str) -> bool: + """Return whether ``path`` is already an exported model directory.""" + + if not os.path.isdir(path): + return False + patterns = ("model*.safetensors", "pytorch_model*.bin") + return any(glob.glob(os.path.join(path, pattern)) for pattern in patterns) + + +def _load_mtp_source( + checkpoint_path: str, + draft_config_path: Optional[str], +) -> Tuple[Dict[str, torch.Tensor], str, Optional[str]]: + """Load MTP weights from either runtime state or an exported draft. + + Returns ``(state_dict, config_source, model_source_dir)``. The last item is + set only for an exported model directory, where companion modeling files + may also need to be copied. + """ + + path = checkpoint_path + if path.startswith("file://"): + path = path[len("file://") :] + if _has_model_weights(path): + return _load_first_checkpoint(path), path, path + + from specforge.export.checkpoint_io import resolve_training_state + + state = resolve_training_state(checkpoint_path) + if state.get("strategy") != "mtp": + raise ValueError( + "MTP merge requires a training checkpoint written by strategy='mtp'; " + f"got strategy={state.get('strategy')!r}" + ) + draft_state = state.get("draft_state_dict") + if not isinstance(draft_state, dict): + raise ValueError("MTP training checkpoint has no draft_state_dict") + if not draft_config_path: + raise ValueError( + "draft_config_path is required when merging a runtime training " + "checkpoint" + ) + if not os.path.isfile(draft_config_path): + raise FileNotFoundError(f"draft config not found: {draft_config_path}") + return dict(draft_state), draft_config_path, None + + +def merge_mtp_into_base( + base_model_path: str, + mtp_checkpoint_path: str, + output_path: str, + key_format: str = "sglang", + *, + draft_config_path: Optional[str] = None, +) -> None: + """Merge trained MTP weights into a copy of the base checkpoint. + + The output directory is a self-contained HF checkpoint loadable directly by + SGLang's native MTP modules (no separate draft-model path). Runtime + checkpoints require ``draft_config_path``; an exported HF draft supplies its + own ``config.json``. + """ + + mtp_state, config_source, model_source_dir = _load_mtp_source( + mtp_checkpoint_path, draft_config_path + ) + embed_key_candidates, head_key_candidates, prefix = _resolve_key_candidates( + config_source + ) + os.makedirs(output_path, exist_ok=True) + + mtp_state = convert_mtp_keys(mtp_state, key_format, prefix) + + # Determine whether word embeddings are tied to decide whether a separate + # lm_head must be materialized for the MTP module. + tie_word_embeddings = True + base_config_path = os.path.join(base_model_path, "config.json") + if os.path.exists(base_config_path): + with open(base_config_path, "r") as f: + base_cfg = json.load(f) + # VLM checkpoints nest text config under "text_config". + text_cfg = base_cfg.get("text_config", base_cfg) + tie_word_embeddings = text_cfg.get("tie_word_embeddings", True) + + # If the trained checkpoint did not save shared embeddings, copy them from + # the base checkpoint so vLLM/SGLang can initialise the MTP embed_tokens/ + # lm_head from the merged checkpoint. + embed_target = f"{prefix}embed_tokens.weight" + head_target = f"{prefix}lm_head.weight" + if embed_target not in mtp_state or ( + not tie_word_embeddings and head_target not in mtp_state + ): + base_state = load_tensors_by_keys( + base_model_path, embed_key_candidates + head_key_candidates + ) + mtp_state = _copy_shared_embeddings( + base_state, + mtp_state, + tie_word_embeddings, + embed_key_candidates, + head_key_candidates, + prefix, + ) + + # The generic merge machinery (copy, prefix-key replacement, shard/index + # writing) lives in modeling/target/checkpoint.py. + merge_state_into_checkpoint( + base_model_path, + mtp_state, + output_path, + shard_name="mtp-merged.safetensors", + drop_prefixes=(prefix,), + ) + + # Ensure the merged config exposes the MTP structural dims. vLLM/SGLang + # use these values to build the MTP module; if the base config omits + # ``head_dim`` (common for some Qwen3.5 checkpoints), the loader will use + # its default and fail with a shape mismatch. + resolved_draft_config_path = ( + config_source + if os.path.isfile(config_source) + else os.path.join(config_source, "config.json") + ) + output_config_path = os.path.join(output_path, "config.json") + if os.path.exists(resolved_draft_config_path) and os.path.exists( + output_config_path + ): + with open(resolved_draft_config_path, "r") as f: + draft_config = json.load(f) + with open(output_config_path, "r") as f: + base_config = json.load(f) + patched_config = _patch_text_config(base_config, draft_config) + with open(output_config_path, "w") as f: + json.dump(patched_config, f, indent=2) + + # Copy over the MTP modeling file if present; some loaders need it for + # trust_remote_code / auto_map resolution. + if model_source_dir is not None: + mtp_py_src = os.path.join(model_source_dir, "mtp.py") + if os.path.exists(mtp_py_src): + shutil.copy2(mtp_py_src, os.path.join(output_path, "mtp.py")) + + print(f"Merged checkpoint saved to {output_path}") + print(f" key format: {key_format}") + print(f" MTP tensors merged: {len(mtp_state)}") + + +__all__ = ["convert_mtp_keys", "merge_mtp_into_base"] diff --git a/specforge/modeling/draft/__init__.py b/specforge/modeling/draft/__init__.py index 839869dd6..ce7c43f22 100644 --- a/specforge/modeling/draft/__init__.py +++ b/specforge/modeling/draft/__init__.py @@ -8,6 +8,7 @@ from .domino import DominoDraftModel from .dspark import DSparkDraftModel from .llama3_eagle import LlamaForCausalLMEagle3 +from .mtp import Qwen3_5MTPDraftModel from .peagle import PEagleDraftModel from .registry import DRAFT_REGISTRY, available_drafts, register_draft, resolve_draft @@ -18,6 +19,7 @@ "DSparkDraftModel", "LlamaForCausalLMEagle3", "PEagleDraftModel", + "Qwen3_5MTPDraftModel", "build_target_layer_ids", "extract_context_feature", "sample", diff --git a/specforge/modeling/draft/mtp/__init__.py b/specforge/modeling/draft/mtp/__init__.py new file mode 100644 index 000000000..77ee2881c --- /dev/null +++ b/specforge/modeling/draft/mtp/__init__.py @@ -0,0 +1,7 @@ +# coding=utf-8 +"""MTP draft architectures: one registered module per model family.""" + +from specforge.modeling.draft.mtp.base import MTPDraftModel +from specforge.modeling.draft.mtp.qwen3_5 import Qwen3_5MTPDraftModel + +__all__ = ["MTPDraftModel", "Qwen3_5MTPDraftModel"] diff --git a/specforge/modeling/draft/mtp/base.py b/specforge/modeling/draft/mtp/base.py new file mode 100644 index 000000000..e79b63606 --- /dev/null +++ b/specforge/modeling/draft/mtp/base.py @@ -0,0 +1,120 @@ +# coding=utf-8 +"""Architecture-independent MTP draft contract. + +One registered subclass per model family (e.g. ``qwen3_5.Qwen3_5MTPDraftModel``) +owns the actual trainable network and the native checkpoint key layout for +that family. This base owns the shared contract the MTP algorithm code relies +on: + +- ``embed_tokens`` plus a trainable ``mtp`` module with an ``lm_head`` +- ``forward(input_ids, hidden_states, ...)`` -> object exposing ``logits`` +- the native checkpoint key prefix (``mtp.*``) used for native-head init and + export round-trips +- sharing/freezing the target checkpoint's embedding (and optional lm_head) +""" + +from __future__ import annotations + +from typing import Optional + +import torch +from torch import nn + + +class MTPDraftModel(nn.Module): + """Contract shared by all MTP draft architectures.""" + + #: Native MTP weights live under this checkpoint key prefix. + NATIVE_KEY_PREFIX = "mtp." + + #: Target-checkpoint key candidates consulted when merging trained MTP + #: weights back into a base checkpoint (see ``specforge/export/mtp.py``). + #: Families override these when the target nests its text decoder + #: differently (VLM ``model.language_model.*`` vs plain ``model.*``). + TARGET_EMBED_KEY_CANDIDATES = [ + "model.language_model.embed_tokens.weight", + "model.embed_tokens.weight", + "embed_tokens.weight", + ] + TARGET_HEAD_KEY_CANDIDATES = [ + "model.language_model.lm_head.weight", + "model.lm_head.weight", + "lm_head.weight", + ] + + embed_tokens: nn.Embedding + mtp: nn.Module + + def forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + ): + """Run the draft on shifted tokens plus target last hidden states. + + Returns an object exposing ``logits`` of shape [batch, seq, vocab]. + """ + raise NotImplementedError + + def share_target_embeddings( + self, + embed_weight: torch.Tensor, + lm_head_weight: Optional[torch.Tensor] = None, + ) -> None: + """Share and freeze the target checkpoint's embedding and lm_head. + + lm_head sharing follows the family config's ``mtp_config.share_lm_head`` + flag (default True) and then requires ``lm_head_weight``. + """ + self.embed_tokens.weight = embed_weight + self.embed_tokens.requires_grad_(False) + mtp_config = getattr(self.config, "mtp_config", None) or {} + if mtp_config.get("share_lm_head", True): + if lm_head_weight is None: + raise ValueError( + "share_lm_head is enabled but no target lm_head weight given" + ) + self.mtp.lm_head.weight = lm_head_weight + self.mtp.lm_head.requires_grad_(False) + + def native_state_dict(self) -> dict[str, torch.Tensor]: + """Return the draft weights in the native ``mtp.*`` serving layout.""" + return { + key: value + for key, value in self.state_dict().items() + if key.startswith(self.NATIVE_KEY_PREFIX) + } + + def required_native_state_keys(self) -> set[str]: + """Return native keys that must exist for safe target initialization. + + A shared lm_head is deliberately reconstructed from the target model, + and native Qwen3.5 checkpoints do not need to duplicate it under the + MTP prefix. Every other native draft tensor must be present; accepting a + partial prefix match would leave part of the trainable head randomized. + """ + required = set(self.native_state_dict()) + mtp_config = getattr(self.config, "mtp_config", None) or {} + if mtp_config.get("share_lm_head", True): + required.discard(f"{self.NATIVE_KEY_PREFIX}lm_head.weight") + return required + + def allowed_extra_native_state_keys(self) -> set[str]: + """Native keys a merged checkpoint may carry that the draft never owns. + + ``export/mtp.merge_mtp_into_base`` backfills shared embeddings into the + native namespace so serving can instantiate ``mtp.embed_tokens`` (and a + separate ``mtp.lm_head`` for untied targets). Re-finetuning a merged + checkpoint must tolerate those keys instead of rejecting the + checkpoint as incompatible. + """ + extra = {f"{self.NATIVE_KEY_PREFIX}embed_tokens.weight"} + mtp_config = getattr(self.config, "mtp_config", None) or {} + if mtp_config.get("share_lm_head", True): + extra.add(f"{self.NATIVE_KEY_PREFIX}lm_head.weight") + return extra + + +__all__ = ["MTPDraftModel"] diff --git a/specforge/modeling/draft/mtp/qwen3_5.py b/specforge/modeling/draft/mtp/qwen3_5.py new file mode 100644 index 000000000..d1e023cfb --- /dev/null +++ b/specforge/modeling/draft/mtp/qwen3_5.py @@ -0,0 +1,543 @@ +# coding=utf-8 +"""Multi-Token Prediction (MTP) draft model for Qwen3.5. + +Architecture follows the Qwen3.5 MTP design: + 1. Normalize input embeddings and target last hidden states separately. + 2. Concatenate and project via fc( [norm(emb); norm(hidden)] ). + 3. Run a 1-layer Qwen3 transformer. + 4. Compute logits with a (shared) lm_head. + +Weight key layout matches SGLang's Qwen3_5ForCausalLMMTP: + mtp.pre_fc_norm_embedding.weight + mtp.pre_fc_norm_hidden.weight + mtp.fc.weight + mtp.layers.0.self_attn.q_proj.weight + mtp.layers.0.mlp.gate_proj.weight + mtp.lm_head.weight +""" + +import copy +from typing import Optional, Tuple + +import torch +from torch import nn +from transformers.cache_utils import Cache +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.models.qwen3.modeling_qwen3 import ( + ALL_ATTENTION_FUNCTIONS, + FlashAttentionKwargs, + GradientCheckpointingLayer, + Qwen3Config, + Qwen3MLP, + Qwen3PreTrainedModel, + eager_attention_forward, + rotate_half, +) +from typing_extensions import Unpack + +from specforge.modeling._mask_utils import _expand_mask, _make_causal_mask +from specforge.modeling.draft.mtp.base import MTPDraftModel +from specforge.modeling.draft.registry import register_draft + + +class Qwen3_5RMSNorm(nn.Module): + """Gemma-style RMSNorm used by Qwen3.5 in vLLM. + + The official Qwen3.5 checkpoints and vLLM use ``x * (1 + weight)`` instead of + the standard HuggingFace ``x * weight``. The parameter is initialized to + zeros so that ``1 + weight`` starts at one. This must match the inference + implementation, otherwise every RMSNorm weight trained in SpecForge would be + off by +1.0 when loaded into vLLM. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + # Gemma-style: multiply by (1 + weight). + hidden_states = hidden_states * (1.0 + self.weight.float()) + return hidden_states.to(input_dtype) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Apply rotary positional embeddings with partial rotary support. + + When partial_rotary_factor < 1.0, only the first ``rotary_dim`` dimensions + of q/k are rotated; the rest pass through unchanged. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) + k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) + q_embed = torch.cat([q_embed, q_pass], dim=-1) + k_embed = torch.cat([k_embed, k_pass], dim=-1) + return q_embed, k_embed + + +class PartialRotaryEmbedding(nn.Module): + """Rotary embedding that computes inv_freq for only a fraction of head_dim. + + Matches the official Qwen3.5 ``partial_rotary_factor`` behaviour where + ``dim = int(head_dim * partial_rotary_factor)``. + """ + + inv_freq: torch.Tensor + + def __init__(self, config, head_dim): + super().__init__() + rope_theta = getattr(config, "rope_theta", 10000.0) + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / ( + rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, x, position_ids): + inv_freq_expanded = ( + self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + ) + position_ids_expanded = position_ids[:, None, :].float() + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Qwen3MTPAttention(nn.Module): + """Causal self-attention with optional output gate (Qwen3.5 style). + + When ``attn_output_gate`` is True (the default for Qwen3.5), ``q_proj`` + outputs ``num_heads * head_dim * 2`` and the extra half is used as a + sigmoid gate applied to the attention output before ``o_proj``. + """ + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + self.num_key_value_groups = ( + config.num_attention_heads // config.num_key_value_heads + ) + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.attn_output_gate = getattr(config, "attn_output_gate", False) + + if self.attn_output_gate: + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim * 2, + bias=config.attention_bias, + ) + else: + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim, + bias=config.attention_bias, + ) + self.k_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + self.q_norm = Qwen3_5RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3_5RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + bsz, q_len, _ = hidden_states.size() + + if self.attn_output_gate: + query_states, gate = torch.chunk( + self.q_proj(hidden_states).view(bsz, q_len, -1, self.head_dim * 2), + 2, + dim=-1, + ) + gate = gate.reshape(bsz, q_len, -1) + else: + query_states = self.q_proj(hidden_states).view( + bsz, q_len, -1, self.head_dim + ) + + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view( + bsz, q_len, self.config.num_attention_heads, self.head_dim + ).transpose(1, 2) + key_states = key_states.view( + bsz, q_len, self.config.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = value_states.view( + bsz, q_len, self.config.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + attn_fn = eager_attention_forward + if self.config._attn_implementation != "eager": + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attn_fn( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + + if self.attn_output_gate: + attn_output = attn_output * torch.sigmoid(gate) + + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Qwen3MTPDecoderLayer(GradientCheckpointingLayer): + """A single Qwen3-style decoder layer for the MTP draft model.""" + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3MTPAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3MLP(config) + self.input_layernorm = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + past_key_value=past_key_value, + cache_position=cache_position, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + return outputs + + +class Qwen3_5MTPModel(nn.Module): + """The core MTP module wrapped under the ``mtp.`` prefix. + + SGLang's ``Qwen3_5ForCausalLMMTP`` wraps a *flat* ``Qwen3_5ForCausalLM`` + (``self.layers`` directly on the ForCausalLM, not nested under + ``self.model``), so after the ``mtp.`` -> ``model.`` remap the flat keys + ``mtp.layers.0.*`` / ``mtp.norm.weight`` become ``model.layers.0.*`` / + ``model.norm.weight``, matching ``self.model.layers`` / ``self.model.norm``. + """ + + def __init__(self, config: Qwen3Config): + super().__init__() + self.config = config + + # Fusion projection: fc( concat( norm(input_embeds), norm(target_hidden) ) ) + self.pre_fc_norm_embedding = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_fc_norm_hidden = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + + # Single-layer Qwen3 transformer, flat under `mtp.layers.*` / `mtp.norm` + # to match the native Qwen3.5 checkpoint layout consumed by SGLang. + mtp_config = copy.deepcopy(config) + mtp_config.num_hidden_layers = 1 + self.layers = nn.ModuleList([Qwen3MTPDecoderLayer(mtp_config, layer_idx=0)]) + self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = PartialRotaryEmbedding( + mtp_config, + getattr( + mtp_config, + "head_dim", + mtp_config.hidden_size // mtp_config.num_attention_heads, + ), + ) + + # LM head (shared with target model during training) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward( + self, + inputs_embeds: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + # Fusion + normed_emb = self.pre_fc_norm_embedding(inputs_embeds) + normed_hidden = self.pre_fc_norm_hidden(hidden_states) + hidden_states = self.fc(torch.cat([normed_emb, normed_hidden], dim=-1)) + + bsz, seq_len, _ = hidden_states.size() + if position_ids is None: + device = hidden_states.device + position_ids = ( + torch.arange(seq_len, dtype=torch.long, device=device) + .unsqueeze(0) + .expand(bsz, -1) + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # Causal mask + if attention_mask is not None and attention_mask.dim() == 2: + # [bsz, seq_len] -> [bsz, 1, seq_len, seq_len] + combined_mask = _make_causal_mask( + (bsz, seq_len), hidden_states.dtype, device=hidden_states.device + ) + expanded_mask = _expand_mask( + attention_mask, hidden_states.dtype, tgt_len=seq_len + ).to(hidden_states.device) + attention_mask = expanded_mask + combined_mask + elif attention_mask is None and seq_len > 1: + attention_mask = _make_causal_mask( + (bsz, seq_len), hidden_states.dtype, device=hidden_states.device + ) + + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + position_embeddings=position_embeddings, + )[0] + + hidden_states = self.norm(hidden_states) + return hidden_states + + +@register_draft +class Qwen3_5MTPDraftModel(MTPDraftModel, Qwen3PreTrainedModel): + """ + Qwen3.5 MTP draft model for SpecForge training. + + The embed_tokens table is loaded from the target model and frozen by default; + the lm_head is optionally shared with the target model. All trainable MTP + parameters live under the `mtp.*` prefix so that checkpoints can be loaded + directly by SGLang's Qwen3_5ForCausalLMMTP. + """ + + config_class = Qwen3Config + _no_split_modules = ["Qwen3MTPDecoderLayer"] + + def __init__(self, config: Qwen3Config) -> None: + super().__init__(config) + self.config = config + + # Shared embedding with the target model (loaded externally, frozen) + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id + ) + + # All trainable MTP weights under `mtp.*` + self.mtp = Qwen3_5MTPModel(config) + + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def get_output_embeddings(self): + return self.mtp.lm_head + + def set_output_embeddings(self, value): + self.mtp.lm_head = value + + def forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + **kwargs, + ) -> CausalLMOutputWithPast: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = self.mtp( + inputs_embeds=inputs_embeds, + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + ) + logits = self.mtp.lm_head(hidden_states) + return CausalLMOutputWithPast(logits=logits) + + @torch.inference_mode() + def spec_generate( + self, + target: nn.Module, + input_ids: torch.LongTensor, + max_new_tokens: int, + stop_token_ids: Optional[list[int]] = None, + temperature: float = 0.0, + ) -> torch.LongTensor: + """Sequential MTP speculative generation (single MTP layer).""" + self.eval() + device = input_ids.device + num_input_tokens = input_ids.shape[1] + max_length = num_input_tokens + max_new_tokens + output_ids = input_ids.clone() + + from transformers.cache_utils import DynamicCache + + past_key_values_target = DynamicCache() + + # Prefill target once to get initial last hidden state + target_out = target( + input_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + next_token_logits = target_out.logits[:, -1, :] + if temperature < 1e-5: + next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) + else: + next_token = torch.multinomial( + torch.softmax(next_token_logits / temperature, dim=-1), num_samples=1 + ) + output_ids = torch.cat([output_ids, next_token], dim=1) + target_hidden = target_out.hidden_states[-1][:, -1:, :] + + # Loop invariant: the target KV cache covers output_ids[:-1], and + # target_hidden is the post-norm hidden state at position + # len(output_ids)-2 (driving the draft's view of the newest token). + while output_ids.shape[1] < max_length: + committed = output_ids.shape[1] + draft_input_ids = output_ids[:, -1:] + draft_position_ids = torch.tensor( + [[committed - 1]], dtype=torch.long, device=device + ) + draft_embeds = self.embed_tokens(draft_input_ids) + draft_hidden = self.mtp( + inputs_embeds=draft_embeds, + hidden_states=target_hidden, + position_ids=draft_position_ids, + ) + draft_logits = self.mtp.lm_head(draft_hidden) + if temperature < 1e-5: + draft_token = torch.argmax(draft_logits[:, -1, :], dim=-1, keepdim=True) + else: + draft_token = torch.multinomial( + torch.softmax(draft_logits[:, -1, :] / temperature, dim=-1), + num_samples=1, + ) + + # Verify against the target. logits[:, 0] is the target's prediction + # for the position the draft token occupies; logits[:, -1] is already + # conditioned on the (possibly rejected) draft token. + verify_input_ids = torch.cat([draft_input_ids, draft_token], dim=1) + verify_position_ids = torch.arange( + committed - 1, + committed + 1, + dtype=torch.long, + device=device, + ).unsqueeze(0) + target_out = target( + verify_input_ids, + position_ids=verify_position_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + target_token = torch.argmax( + target_out.logits[:, 0:1, :], dim=-1, keepdim=True + ) + # h at the newest committed token drives the next draft round. + target_hidden = target_out.hidden_states[-1][:, 0:1, :] + # Trim the scored draft token's KV again: wrong on rejection, and + # re-appended by next round's verify on acceptance. + past_key_values_target.crop(committed) + + if torch.equal(draft_token, target_token): + output_ids = torch.cat([output_ids, draft_token], dim=1) + else: + output_ids = torch.cat([output_ids, target_token], dim=1) + + if ( + stop_token_ids is not None + and output_ids[0, -1].item() in stop_token_ids + ): + break + + return output_ids diff --git a/specforge/modeling/target/checkpoint.py b/specforge/modeling/target/checkpoint.py new file mode 100644 index 000000000..e27dfec6b --- /dev/null +++ b/specforge/modeling/target/checkpoint.py @@ -0,0 +1,209 @@ +# coding=utf-8 +"""Model-agnostic selective loading from local or Hugging Face checkpoints. + +These helpers know nothing about any model family or key naming convention; +callers provide the keys or the predicate. Both sharded checkpoints +(``*.safetensors.index.json``) and single-file checkpoints are supported. +""" + +from __future__ import annotations + +import glob +import json +import os +from typing import Callable, Dict, Iterable, List, Optional + +import torch +from safetensors import safe_open + + +def resolve_checkpoint_dir( + path_or_repo: str, + cache_dir: Optional[str] = None, + allow_patterns: Optional[List[str]] = None, +) -> str: + """Return a local checkpoint directory, downloading from the Hub if needed.""" + + if os.path.exists(path_or_repo): + return path_or_repo + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id=path_or_repo, + cache_dir=cache_dir, + allow_patterns=allow_patterns or ["*.json", "*.safetensors", "*.bin"], + ) + + +def read_weight_map(checkpoint_dir: str) -> Dict[str, str]: + """Return the ``weight_map`` of a sharded checkpoint, or {} if unsharded.""" + + index_files = glob.glob(os.path.join(checkpoint_dir, "*.index.json")) + if not index_files: + return {} + with open(index_files[0], "r") as f: + index = json.load(f) + return index.get("weight_map", {}) + + +def list_checkpoint_keys(checkpoint_dir: str) -> List[str]: + """List all tensor keys without loading tensor payloads.""" + + weight_map = read_weight_map(checkpoint_dir) + if weight_map: + return sorted(weight_map.keys()) + for pattern in ("*.safetensors", "*.bin"): + files = sorted(glob.glob(os.path.join(checkpoint_dir, pattern))) + if files: + target = files[0] + if target.endswith(".safetensors"): + with safe_open(target, framework="pt") as f: + return sorted(f.keys()) + state = torch.load(target, map_location="cpu", weights_only=True) + return sorted(state.keys()) + raise FileNotFoundError(f"No checkpoint found in {checkpoint_dir}") + + +def load_selected_tensors( + checkpoint_dir: str, + predicate: Callable[[str], bool], +) -> Dict[str, torch.Tensor]: + """Load only the tensors whose key matches ``predicate``. + + Sharded checkpoints open just the shards that hold selected keys. + """ + + weight_map = read_weight_map(checkpoint_dir) + selected: Dict[str, torch.Tensor] = {} + if weight_map: + shards = sorted({weight_map[k] for k in weight_map if predicate(k)}) + for shard in shards: + shard_path = os.path.join(checkpoint_dir, shard) + if not os.path.exists(shard_path): + continue + with safe_open(shard_path, framework="pt") as f: + for key in f.keys(): + if predicate(key): + selected[key] = f.get_tensor(key) + return selected + + for pattern in ("*.safetensors", "*.bin"): + files = sorted(glob.glob(os.path.join(checkpoint_dir, pattern))) + if files: + target = files[0] + if target.endswith(".safetensors"): + with safe_open(target, framework="pt") as f: + for key in f.keys(): + if predicate(key): + selected[key] = f.get_tensor(key) + else: + state = torch.load(target, map_location="cpu", weights_only=True) + for key, value in state.items(): + if predicate(key): + selected[key] = value + return selected + raise FileNotFoundError(f"No checkpoint found in {checkpoint_dir}") + + +def load_tensors_by_keys( + checkpoint_dir: str, keys: Iterable[str] +) -> Dict[str, torch.Tensor]: + """Load exactly ``keys`` (missing keys are simply absent from the result).""" + + wanted = set(keys) + return load_selected_tensors(checkpoint_dir, lambda key: key in wanted) + + +def merge_state_into_checkpoint( + base_checkpoint_dir: str, + state: Dict[str, torch.Tensor], + output_dir: str, + *, + shard_name: str, + drop_prefixes: Iterable[str] = (), +) -> None: + """Merge a state dict into a copy of a base checkpoint (model-agnostic). + + Copies non-weight files, drops base weight entries under ``drop_prefixes``, + and merges ``state``. Sharded bases get ``state`` written to a new + ``shard_name`` shard with the index ``weight_map`` updated in place (the + large base shards are never rewritten); single-file bases are rewritten + whole under their original file name. + """ + + import shutil + + from safetensors.torch import save_file + + os.makedirs(output_dir, exist_ok=True) + prefixes = tuple(drop_prefixes) + + # Copy non-weight files so the output directory is self-contained. + for fname in os.listdir(base_checkpoint_dir): + src = os.path.join(base_checkpoint_dir, fname) + if os.path.isfile(src): + shutil.copy2(src, os.path.join(output_dir, fname)) + + index_files = glob.glob(os.path.join(base_checkpoint_dir, "*.index.json")) + if index_files: + with open(index_files[0], "r") as f: + index = json.load(f) + weight_map = index.get("weight_map", {}) + + old_keys = [k for k in weight_map if k.startswith(prefixes)] + for key in old_keys: + del weight_map[key] + if old_keys: + print( + f"Replaced {len(old_keys)} weight entries under {prefixes} " + "from base model." + ) + + # Write the incoming tensors to a dedicated shard; base shards untouched. + save_file(state, os.path.join(output_dir, shard_name)) + for key in state.keys(): + weight_map[key] = shard_name + + index["weight_map"] = weight_map + with open(os.path.join(output_dir, os.path.basename(index_files[0])), "w") as f: + json.dump(index, f, indent=2) + return + + # Single-file base: load, drop, merge, rewrite under the original name. + base_safetensors = glob.glob(os.path.join(base_checkpoint_dir, "*.safetensors")) + base_bins = glob.glob(os.path.join(base_checkpoint_dir, "*.bin")) + if not base_safetensors and not base_bins: + raise FileNotFoundError(f"No checkpoint found in {base_checkpoint_dir}") + base_state = ( + load_selected_tensors(base_checkpoint_dir, lambda _key: True) + if base_safetensors + else torch.load(base_bins[0], map_location="cpu", weights_only=True) + ) + out_name = os.path.basename( + base_safetensors[0] if base_safetensors else base_bins[0] + ) + + old_keys = [k for k in base_state if k.startswith(prefixes)] + for key in old_keys: + del base_state[key] + if old_keys: + print( + f"Replaced {len(old_keys)} weight entries under {prefixes} " + "from base model." + ) + + merged = {**base_state, **state} + if out_name.endswith(".safetensors"): + save_file(merged, os.path.join(output_dir, out_name)) + else: + torch.save(merged, os.path.join(output_dir, out_name)) + + +__all__ = [ + "list_checkpoint_keys", + "load_selected_tensors", + "load_tensors_by_keys", + "merge_state_into_checkpoint", + "read_weight_map", + "resolve_checkpoint_dir", +] diff --git a/specforge/runtime/contracts.py b/specforge/runtime/contracts.py index 7b6958b06..6d30421eb 100644 --- a/specforge/runtime/contracts.py +++ b/specforge/runtime/contracts.py @@ -33,7 +33,7 @@ SCHEMA_VERSION = 1 RunMode = Literal["online", "offline"] -DraftStrategyName = Literal["eagle3", "dflash", "domino", "dspark", "peagle"] +DraftStrategyName = Literal["eagle3", "dflash", "domino", "dspark", "peagle", "mtp"] # Tagged union for the EAGLE3 target feature. The *strategy* owns the # projection so the trainer core stays branch-free: # - pruned_logits: rollout applied the t2d vocab map; stored (seq, draft_vocab) diff --git a/specforge/training/strategies/base.py b/specforge/training/strategies/base.py index 4d262cec9..69495a295 100644 --- a/specforge/training/strategies/base.py +++ b/specforge/training/strategies/base.py @@ -563,6 +563,63 @@ def checkpoint_state_filter(self, state_dict: Dict[str, Any]) -> Dict[str, Any]: } +class MTPTrainStrategy(DraftTrainStrategy): + """MTP strategy over ``OnlineMTPModel`` with final-hidden supervision.""" + + name = "mtp" + required_features = { + "input_ids", + "loss_mask", + "target_last_hidden_states", + } + + def __init__(self, mtp_model: nn.Module) -> None: + self.mtp_model = mtp_model + + def trainable_module(self) -> nn.Module: + return self.mtp_model + + def _device(self) -> torch.device: + return next(self.mtp_model.parameters()).device + + def forward_loss( + self, batch: TrainBatch, ctx: Optional[StepContext] = None + ) -> StepOutput: + self.validate_batch(batch) + t = batch.tensors + device = self._device() + # OnlineMTPModel performs the next-token shift internally and returns + # per-position correct/denominator tensors (single-layer: length-1 lists). + loss, corrects, denoms = self.mtp_model( + input_ids=t["input_ids"].to(device), + hidden_states=t["target_last_hidden_states"].to(device), + loss_mask=t["loss_mask"].to(device), + ) + correct_sum = corrects[0].sum() + denom_sum = denoms[0].sum() + metrics = { + "accuracy": (correct_sum / denom_sum.clamp_min(1)).detach(), + "accuracy_denom": denom_sum.detach(), + } + return StepOutput( + loss=loss, + metrics=metrics, + ratio_metrics={"accuracy": (correct_sum, denom_sum)}, + # TrainerCore backpropagates additive numerators and divides by the + # global token denominator across accumulation steps / DP ranks. + loss_terms=(loss * denom_sum, denom_sum), + ) + + def checkpoint_state_filter(self, state_dict: Dict[str, Any]) -> Dict[str, Any]: + # Everything trainable lives under draft_model.; persisting the stripped + # keys (embed_tokens.* + mtp.*) matches the native serving layout. + return { + k.replace("draft_model.", ""): v + for k, v in state_dict.items() + if "draft_model." in k + } + + class DominoTrainStrategy(DraftTrainStrategy): """Domino block-parallel strategy wrapping ``OnlineDominoModel``. diff --git a/tests/test_algorithms/test_builtin_providers.py b/tests/test_algorithms/test_builtin_providers.py index dc92b1d91..11a99b5fa 100644 --- a/tests/test_algorithms/test_builtin_providers.py +++ b/tests/test_algorithms/test_builtin_providers.py @@ -22,7 +22,7 @@ from specforge.algorithms.contracts import AlgorithmSpec, FeatureMode REPO_ROOT = Path(__file__).resolve().parents[2] -BUILTINS = ("dflash", "domino", "dspark", "eagle3", "peagle") +BUILTINS = ("dflash", "domino", "dspark", "eagle3", "mtp", "peagle") class BuiltinProviderContractTest(unittest.TestCase): @@ -149,7 +149,7 @@ def test_target_derived_defaults_and_overrides_are_provider_owned(self): self.assertEqual(vocab_size, defaults.draft_vocab_size) self.assertEqual(has_override, policy.apply_overrides is not None) - for name in ("domino", "dspark"): + for name in ("domino", "dspark", "mtp"): with self.subTest(algorithm=name): policy = self.registry.resolve(name).providers.model.draft_config self.assertIsNone(policy.target_defaults) @@ -209,6 +209,7 @@ def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): "dflash": dflash_family, "domino": dflash_family, "dspark": dflash_family, + "mtp": SimpleNamespace(), } expected_keys = { "eagle3": { @@ -244,6 +245,12 @@ def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): "dspark_l1_loss_alpha", "dspark_confidence_head_alpha", }, + "mtp": { + "mtp_draft_num_hidden_layers", + "mtp_draft_vocab_size", + "mtp_share_lm_head", + "mtp_attention_backend", + }, } for name in BUILTINS: @@ -430,7 +437,7 @@ def test_building_catalog_does_not_import_training_or_torch(self): code = ( "import sys; " "from specforge.algorithms.builtin import builtin_algorithm_registry; " - "r=builtin_algorithm_registry(); assert len(r)==5; " + "r=builtin_algorithm_registry(); assert len(r)==6; " "assert 'torch' not in sys.modules; " "assert 'specforge.training.strategies.registry' not in sys.modules" ) diff --git a/tests/test_config/test_launch_topology.py b/tests/test_config/test_launch_topology.py index e58cc9e03..603655efb 100644 --- a/tests/test_config/test_launch_topology.py +++ b/tests/test_config/test_launch_topology.py @@ -69,6 +69,7 @@ "qwen3.5-4b-dflash-online-npu.yaml": 8, "qwen3.5-4b-dflash-disaggregated-npu.yaml": 14, "qwen3.5-4b-domino-online-npu.yaml": 8, + "qwen3.5-4b-mtp-disaggregated-npu.yaml": 10, "qwen3.6-27b-dflash-disaggregated.yaml": 2, "qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml": 2, "qwen3.6-27b-dflash-multiserver-disaggregated.yaml": 2, @@ -262,6 +263,60 @@ ], }, }, + "qwen3.5-4b-mtp-disaggregated-npu.yaml": { + "control_dir": "outputs/qwen3.5-4b-mtp-npu-managed/control", + "consumer_state_dir": "outputs/qwen3.5-4b-mtp-npu-managed/consumer-state", + "backend": "mooncake", + "managed_local": { + "trainer_cuda_visible_devices": [ + "6", + "7", + "8", + "9", + "10", + "11", + "12", + "13", + "14", + "15", + ], + "mooncake": { + "global_segment_size_bytes": 17179869184, + }, + "capture_servers": [ + { + "port": 30000, + "cuda_visible_devices": ["0"], + "tp_size": 1, + }, + { + "port": 30001, + "cuda_visible_devices": ["1"], + "tp_size": 1, + }, + { + "port": 30002, + "cuda_visible_devices": ["2"], + "tp_size": 1, + }, + { + "port": 30003, + "cuda_visible_devices": ["3"], + "tp_size": 1, + }, + { + "port": 30004, + "cuda_visible_devices": ["4"], + "tp_size": 1, + }, + { + "port": 30005, + "cuda_visible_devices": ["5"], + "tp_size": 1, + }, + ], + }, + }, "qwen3.6-27b-dflash-disaggregated.yaml": { "control_dir": "outputs/qwen3.6-27b-dflash-disaggregated/control", "backend": "mooncake", @@ -356,7 +411,7 @@ def _recipes() -> dict[str, Path]: class ExampleLaunchTopologyTest(unittest.TestCase): def test_every_recipe_has_the_explicit_golden_topology(self): recipes = _recipes() - self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 66) + self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 67) self.assertEqual(set(recipes), set(EXPECTED_NPROC_PER_NODE)) for filename, nproc_per_node in EXPECTED_NPROC_PER_NODE.items(): diff --git a/tests/test_config/test_unified_feature_reachability.py b/tests/test_config/test_unified_feature_reachability.py index a9518bf12..8c40b986a 100644 --- a/tests/test_config/test_unified_feature_reachability.py +++ b/tests/test_config/test_unified_feature_reachability.py @@ -146,7 +146,7 @@ def test_all_example_configs_validate_through_the_typed_entry(self): for path in EXAMPLE_CONFIG_DIR.glob("*.yaml") if not path.name.startswith(".") ) - self.assertEqual(len(paths), 66) + self.assertEqual(len(paths), 67) resolved_runs = { path.name: resolve_run(Config.from_file(str(path))) for path in paths diff --git a/tests/test_utils/test_mtp.py b/tests/test_utils/test_mtp.py new file mode 100644 index 000000000..2fcc7cd33 --- /dev/null +++ b/tests/test_utils/test_mtp.py @@ -0,0 +1,619 @@ +# coding=utf-8 +"""CPU unit tests for the MTP (native-head fine-tune) algorithm. + +Covers the pieces that do not need a GPU or a real target checkpoint: + - OnlineMTPModel forward/shift/loss/accuracy plumbing + - MTPTrainStrategy batch -> StepOutput adaptation + - strict native ``mtp.*`` weight initialization from a target checkpoint + - selective checkpoint loading (modeling/target/checkpoint.py) + - merge-back round trip (export/mtp.py) + - draft-architecture and built-in algorithm registration +""" + +from __future__ import annotations + +import json +import os +import tempfile +import unittest +from types import SimpleNamespace + +import torch +from transformers.models.qwen3.modeling_qwen3 import Qwen3Config + +from specforge.algorithms.builtin import builtin_algorithm_registry +from specforge.algorithms.mtp.providers import _init_from_native_mtp +from specforge.core.mtp import OnlineMTPModel +from specforge.modeling.draft import available_drafts, resolve_draft +from specforge.modeling.draft.mtp import Qwen3_5MTPDraftModel +from specforge.training.strategies.base import MTPTrainStrategy + + +def _tiny_config(**overrides) -> Qwen3Config: + payload = dict( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=128, + max_position_embeddings=512, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + rms_norm_eps=1e-6, + attn_output_gate=True, + partial_rotary_factor=0.25, + mtp_config={"share_lm_head": True}, + tie_word_embeddings=False, + ) + payload.update(overrides) + config = Qwen3Config(**payload) + config._attn_implementation = "eager" + return config + + +def _tiny_batch(config: Qwen3Config, seq_len: int = 16, batch: int = 2): + input_ids = torch.randint(0, config.vocab_size, (batch, seq_len)) + hidden_states = torch.randn(batch, seq_len, config.hidden_size) + loss_mask = torch.ones(batch, seq_len) + return input_ids, hidden_states, loss_mask + + +def _cfg(target_model_path: str, draft_checkpoint_path: str = ""): + return SimpleNamespace( + model=SimpleNamespace( + target_model_path=target_model_path, + draft_checkpoint_path=draft_checkpoint_path, + cache_dir=None, + ) + ) + + +class OnlineMTPModelTest(unittest.TestCase): + def test_forward_returns_finite_loss_and_accuracy_lists(self): + config = _tiny_config() + model = OnlineMTPModel(Qwen3_5MTPDraftModel(config)) + input_ids, hidden_states, loss_mask = _tiny_batch(config) + + loss, corrects, denoms = model( + input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask + ) + + self.assertEqual(0, loss.dim()) + self.assertTrue(torch.isfinite(loss)) + self.assertEqual(1, len(corrects)) + self.assertEqual(1, len(denoms)) + # next-token shift drops one position + self.assertEqual((2, 15), corrects[0].shape) + self.assertEqual((2, 15), denoms[0].shape) + + def test_loss_backward_populates_mtp_grads(self): + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + model = OnlineMTPModel(draft) + input_ids, hidden_states, loss_mask = _tiny_batch(config) + + loss, _, _ = model( + input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask + ) + loss.backward() + + self.assertIsNotNone(draft.mtp.fc.weight.grad) + self.assertTrue(torch.isfinite(draft.mtp.fc.weight.grad).all()) + + def test_shift_for_next_token_matches_serving_alignment(self): + model = OnlineMTPModel(Qwen3_5MTPDraftModel(_tiny_config())) + logits = torch.zeros(1, 5, 7) + input_ids = torch.tensor([[10, 11, 12, 13, 14]]) + loss_mask = torch.ones(1, 5) + + shift_logits, shift_labels, shift_mask = model._shift_for_next_token( + logits, input_ids, loss_mask + ) + + self.assertEqual((1, 4, 7), shift_logits.shape) + # labels are x_2..x_T padded with one ignore index + self.assertEqual([12, 13, 14, -100], shift_labels[0].tolist()) + # the padded position is masked out + self.assertEqual([1, 1, 1, 0], shift_mask[0].tolist()) + + def test_forward_shifts_position_ids_with_draft_tokens(self): + class _RecordingDraft(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(pad_token_id=0) + self.position_ids = None + + def forward( + self, + input_ids, + hidden_states, + attention_mask=None, + position_ids=None, + ): + self.position_ids = position_ids.detach().clone() + logits = torch.zeros( + input_ids.shape[0], input_ids.shape[1], 32, requires_grad=True + ) + return SimpleNamespace(logits=logits) + + draft = _RecordingDraft() + model = OnlineMTPModel(draft) + model( + input_ids=torch.tensor([[10, 11, 12, 13]]), + hidden_states=torch.zeros(1, 4, 8), + loss_mask=torch.ones(1, 4), + position_ids=torch.tensor([[4, 5, 6, 7]]), + ) + + # x[t+1] is fused with h[t], but RoPE must use x[t+1]'s serving + # position. The synthetic final token is assigned the next position. + self.assertEqual([[5, 6, 7, 8]], draft.position_ids.tolist()) + + +class MTPTrainStrategyTest(unittest.TestCase): + def test_forward_loss_adapts_model_outputs(self): + loss = torch.tensor(1.5, requires_grad=True) + corrects = [torch.tensor([[1.0, 0.0, 1.0]])] + denoms = [torch.tensor([[1.0, 1.0, 1.0]])] + + class _Stub(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(1, 1) + + def forward(self, input_ids, hidden_states, loss_mask): + return loss, corrects, denoms + + strategy = MTPTrainStrategy(_Stub()) + batch = SimpleNamespace( + tensors={ + "input_ids": torch.zeros(1, 3, dtype=torch.long), + "loss_mask": torch.ones(1, 3), + "target_last_hidden_states": torch.zeros(1, 3, 4), + } + ) + + out = strategy.forward_loss(batch) + + self.assertIs(out.loss, loss) + self.assertAlmostEqual(2.0 / 3.0, out.metrics["accuracy"].item()) + self.assertEqual(3.0, out.metrics["accuracy_denom"].item()) + num, den = out.ratio_metrics["accuracy"] + self.assertEqual(2.0, num.item()) + self.assertEqual(3.0, den.item()) + loss_num, loss_den = out.loss_terms + self.assertIsNotNone(loss_num.grad_fn) + self.assertEqual(4.5, loss_num.item()) + self.assertEqual(3.0, loss_den.item()) + + def test_forward_loss_rejects_missing_features(self): + strategy = MTPTrainStrategy(torch.nn.Linear(1, 1)) + batch = SimpleNamespace(tensors={"input_ids": torch.zeros(1, 3)}) + with self.assertRaisesRegex(ValueError, "missing required features"): + strategy.forward_loss(batch) + + def test_checkpoint_state_filter_strips_draft_prefix(self): + strategy = MTPTrainStrategy(torch.nn.Linear(1, 1)) + state = { + "draft_model.mtp.fc.weight": torch.zeros(2, 2), + "draft_model.embed_tokens.weight": torch.zeros(4, 2), + "other.weight": torch.zeros(1), + } + filtered = strategy.checkpoint_state_filter(state) + self.assertEqual({"mtp.fc.weight", "embed_tokens.weight"}, set(filtered)) + + +class NativeMTPInitTest(unittest.TestCase): + def test_loads_native_mtp_weights_from_target_checkpoint(self): + from safetensors.torch import save_file + + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + replacement = torch.ones_like(draft.mtp.fc.weight) + native_state = { + key: torch.ones_like(value) + for key, value in draft.native_state_dict().items() + if key in draft.required_native_state_keys() + } + + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + save_file(native_state, f"{tmpdir}/model.safetensors") + _init_from_native_mtp(_cfg(tmpdir), draft) + + self.assertTrue(torch.equal(draft.mtp.fc.weight, replacement)) + + def test_partial_native_weights_raise(self): + from safetensors.torch import save_file + + draft = Qwen3_5MTPDraftModel(_tiny_config()) + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + save_file( + {"mtp.fc.weight": torch.ones_like(draft.mtp.fc.weight)}, + f"{tmpdir}/model.safetensors", + ) + with self.assertRaisesRegex(RuntimeError, "missing required native"): + _init_from_native_mtp(_cfg(tmpdir), draft) + + def test_missing_native_weights_raise_by_default(self): + draft = Qwen3_5MTPDraftModel(_tiny_config()) + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + with self.assertRaisesRegex(RuntimeError, "no native mtp"): + _init_from_native_mtp(_cfg(tmpdir), draft) + + def test_missing_native_weights_allowed_for_warm_start(self): + draft = Qwen3_5MTPDraftModel(_tiny_config()) + before = draft.mtp.fc.weight.detach().clone() + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + # warm start from a trained draft checkpoint skips the strict check + _init_from_native_mtp( + _cfg(tmpdir, draft_checkpoint_path="some/ckpt"), draft + ) + self.assertTrue(torch.equal(draft.mtp.fc.weight, before)) + + def test_native_init_covers_all_mtp_parameters(self): + """A native checkpoint shipping the full mtp.* key set must overwrite + every draft native parameter — none may keep its random init.""" + from safetensors.torch import save_file + + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + native_keys = draft.required_native_state_keys() + replacement = { + key: torch.ones_like(value) + for key, value in draft.native_state_dict().items() + if key in native_keys + } + + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + save_file(replacement, f"{tmpdir}/model.safetensors") + _init_from_native_mtp(_cfg(tmpdir), draft) + + after = draft.native_state_dict() + self.assertTrue(native_keys.issubset(after)) + for key in native_keys: + value = after[key] + self.assertTrue( + torch.all(value == 1), f"{key} was not loaded from native weights" + ) + + def test_native_init_tolerates_merged_checkpoint_keys(self): + """A previously merged checkpoint carries backfilled shared embeddings + (mtp.embed_tokens.weight / mtp.lm_head.weight); re-finetuning it must + not be rejected as an incompatible native state.""" + from safetensors.torch import save_file + + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + state = { + key: torch.ones_like(value) + for key, value in draft.native_state_dict().items() + if key in draft.required_native_state_keys() + } + state["mtp.embed_tokens.weight"] = torch.randn( + config.vocab_size, config.hidden_size + ) + state["mtp.lm_head.weight"] = torch.randn(config.vocab_size, config.hidden_size) + + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + save_file(state, f"{tmpdir}/model.safetensors") + _init_from_native_mtp(_cfg(tmpdir), draft) # must not raise + + self.assertTrue(torch.all(draft.mtp.fc.weight == 1)) + + +class DraftBaseContractTest(unittest.TestCase): + def test_share_target_embeddings_freezes_and_shares(self): + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + embed_w = torch.nn.Parameter(torch.randn(config.vocab_size, config.hidden_size)) + head_w = torch.nn.Parameter(torch.randn(config.vocab_size, config.hidden_size)) + + draft.share_target_embeddings(embed_w, lm_head_weight=head_w) + + self.assertIs(draft.embed_tokens.weight, embed_w) + self.assertIs(draft.mtp.lm_head.weight, head_w) + self.assertFalse(draft.embed_tokens.weight.requires_grad) + self.assertFalse(draft.mtp.lm_head.weight.requires_grad) + + def test_share_lm_head_disabled_keeps_own_head(self): + config = _tiny_config(mtp_config={"share_lm_head": False}) + draft = Qwen3_5MTPDraftModel(config) + own_head = draft.mtp.lm_head.weight + embed_w = torch.nn.Parameter(torch.randn(config.vocab_size, config.hidden_size)) + + draft.share_target_embeddings(embed_w) + + self.assertIs(draft.embed_tokens.weight, embed_w) + self.assertIs(draft.mtp.lm_head.weight, own_head) + + def test_native_state_dict_uses_native_prefix(self): + draft = Qwen3_5MTPDraftModel(_tiny_config()) + native = draft.native_state_dict() + self.assertTrue(native) + self.assertTrue(all(key.startswith(draft.NATIVE_KEY_PREFIX) for key in native)) + self.assertIn("mtp.fc.weight", native) + + def test_required_native_state_respects_shared_lm_head(self): + shared = Qwen3_5MTPDraftModel(_tiny_config()) + own = Qwen3_5MTPDraftModel(_tiny_config(mtp_config={"share_lm_head": False})) + + self.assertNotIn("mtp.lm_head.weight", shared.required_native_state_keys()) + self.assertIn("mtp.lm_head.weight", own.required_native_state_keys()) + + +class SelectiveCheckpointLoadingTest(unittest.TestCase): + def test_sharded_selective_loading(self): + from safetensors.torch import save_file + + from specforge.modeling.target.checkpoint import ( + list_checkpoint_keys, + load_selected_tensors, + read_weight_map, + ) + + with tempfile.TemporaryDirectory(prefix="mtp-ckpt-") as tmpdir: + save_file( + { + "mtp.fc.weight": torch.zeros(4, 4), + "model.embed_tokens.weight": torch.zeros(2, 2), + }, + os.path.join(tmpdir, "model-00001-of-00002.safetensors"), + ) + save_file( + { + "mtp.norm.weight": torch.ones(4), + "lm_head.weight": torch.zeros(2, 2), + }, + os.path.join(tmpdir, "model-00002-of-00002.safetensors"), + ) + weight_map = { + "mtp.fc.weight": "model-00001-of-00002.safetensors", + "model.embed_tokens.weight": "model-00001-of-00002.safetensors", + "mtp.norm.weight": "model-00002-of-00002.safetensors", + "lm_head.weight": "model-00002-of-00002.safetensors", + } + with open(os.path.join(tmpdir, "model.safetensors.index.json"), "w") as f: + json.dump({"weight_map": weight_map}, f) + + self.assertEqual(weight_map, read_weight_map(tmpdir)) + self.assertEqual(4, len(list_checkpoint_keys(tmpdir))) + selected = load_selected_tensors(tmpdir, lambda key: key.startswith("mtp.")) + + self.assertEqual({"mtp.fc.weight", "mtp.norm.weight"}, set(selected)) + self.assertTrue(torch.equal(selected["mtp.norm.weight"], torch.ones(4))) + + def test_single_file_selective_loading(self): + from safetensors.torch import save_file + + from specforge.modeling.target.checkpoint import load_selected_tensors + + with tempfile.TemporaryDirectory(prefix="mtp-ckpt-") as tmpdir: + save_file( + {"mtp.fc.weight": torch.zeros(4, 4), "other.weight": torch.zeros(1)}, + os.path.join(tmpdir, "model.safetensors"), + ) + selected = load_selected_tensors(tmpdir, lambda key: key.startswith("mtp.")) + self.assertEqual({"mtp.fc.weight"}, set(selected)) + + +class ExportRoundTripTest(unittest.TestCase): + """merge_mtp_into_base round trip on synthetic single-file checkpoints.""" + + def test_merge_replaces_native_and_copies_embeddings(self): + from safetensors.torch import save_file + + from specforge.export.mtp import merge_mtp_into_base + from specforge.modeling.target.checkpoint import load_selected_tensors + + with tempfile.TemporaryDirectory() as tmpdir: + base = os.path.join(tmpdir, "base") + draft = os.path.join(tmpdir, "draft") + out = os.path.join(tmpdir, "out") + os.makedirs(base) + os.makedirs(draft) + + base_embed = torch.randn(128, 64) + stale_native = torch.zeros(64, 128) + save_file( + { + "model.embed_tokens.weight": base_embed, + "mtp.fc.weight": stale_native, + }, + os.path.join(base, "model.safetensors"), + ) + with open(os.path.join(base, "config.json"), "w") as f: + json.dump({"hidden_size": 64, "tie_word_embeddings": True}, f) + + trained = torch.ones(64, 128) + save_file( + {"mtp.fc.weight": trained}, + os.path.join(draft, "model.safetensors"), + ) + with open(os.path.join(draft, "config.json"), "w") as f: + json.dump( + { + "architectures": ["Qwen3_5MTPDraftModel"], + "hidden_size": 64, + "head_dim": 16, + }, + f, + ) + + merge_mtp_into_base(base, draft, out) + + merged = load_selected_tensors(out, lambda _key: True) + # trained weights replace the stale native ones + self.assertTrue(torch.equal(merged["mtp.fc.weight"], trained)) + # shared embedding copied into the native namespace + self.assertTrue(torch.equal(merged["mtp.embed_tokens.weight"], base_embed)) + # base weights untouched + self.assertTrue( + torch.equal(merged["model.embed_tokens.weight"], base_embed) + ) + # config patched with the draft's structural dims + with open(os.path.join(out, "config.json")) as f: + merged_config = json.load(f) + self.assertEqual(16, merged_config["head_dim"]) + + def test_merge_accepts_runtime_training_checkpoint(self): + from safetensors.torch import save_file + + from specforge.export.mtp import merge_mtp_into_base + from specforge.modeling.target.checkpoint import load_selected_tensors + + with tempfile.TemporaryDirectory() as tmpdir: + base = os.path.join(tmpdir, "base") + runtime = os.path.join(tmpdir, "run-step1") + out = os.path.join(tmpdir, "out") + draft_config = os.path.join(tmpdir, "draft-config.json") + os.makedirs(base) + os.makedirs(runtime) + + base_embed = torch.randn(128, 64) + save_file( + { + "model.embed_tokens.weight": base_embed, + "mtp.fc.weight": torch.zeros(64, 128), + }, + os.path.join(base, "model.safetensors"), + ) + with open(os.path.join(base, "config.json"), "w") as f: + json.dump({"hidden_size": 64, "tie_word_embeddings": True}, f) + with open(draft_config, "w") as f: + json.dump( + { + "architectures": ["Qwen3_5MTPDraftModel"], + "hidden_size": 64, + "head_dim": 16, + }, + f, + ) + + trained = torch.ones(64, 128) + torch.save( + { + "strategy": "mtp", + "draft_state_dict": {"mtp.fc.weight": trained}, + }, + os.path.join(runtime, "training_state.pt"), + ) + + merge_mtp_into_base( + base, + runtime, + out, + draft_config_path=draft_config, + ) + + merged = load_selected_tensors(out, lambda _key: True) + self.assertTrue(torch.equal(merged["mtp.fc.weight"], trained)) + self.assertTrue(torch.equal(merged["mtp.embed_tokens.weight"], base_embed)) + with open(os.path.join(out, "config.json")) as f: + merged_config = json.load(f) + self.assertEqual(16, merged_config["head_dim"]) + + def test_runtime_checkpoint_requires_draft_config(self): + from specforge.export.mtp import merge_mtp_into_base + + with tempfile.TemporaryDirectory() as tmpdir: + runtime = os.path.join(tmpdir, "run-step1") + os.makedirs(runtime) + torch.save( + { + "strategy": "mtp", + "draft_state_dict": {"mtp.fc.weight": torch.ones(1)}, + }, + os.path.join(runtime, "training_state.pt"), + ) + + with self.assertRaisesRegex(ValueError, "draft_config_path is required"): + merge_mtp_into_base("unused", runtime, os.path.join(tmpdir, "out")) + + def test_merge_runtime_checkpoint_with_shared_tied_weights(self): + """Regression: a tied target shares one storage between the draft's + embed_tokens.weight and mtp.lm_head.weight; safetensors must not choke + on the aliased pair when writing the merged checkpoint.""" + from safetensors.torch import save_file + + from specforge.export.mtp import merge_mtp_into_base + from specforge.modeling.target.checkpoint import load_selected_tensors + + with tempfile.TemporaryDirectory() as tmpdir: + base = os.path.join(tmpdir, "base") + runtime = os.path.join(tmpdir, "run-step1") + out = os.path.join(tmpdir, "out") + draft_config = os.path.join(tmpdir, "draft-config.json") + os.makedirs(base) + os.makedirs(runtime) + + save_file( + {"model.embed_tokens.weight": torch.randn(128, 64)}, + os.path.join(base, "model.safetensors"), + ) + with open(os.path.join(base, "config.json"), "w") as f: + json.dump({"hidden_size": 64, "tie_word_embeddings": True}, f) + with open(draft_config, "w") as f: + json.dump( + { + "architectures": ["Qwen3_5MTPDraftModel"], + "hidden_size": 64, + "head_dim": 16, + }, + f, + ) + + shared = torch.randn(128, 64) + trained = torch.ones(64, 128) + torch.save( + { + "strategy": "mtp", + "draft_state_dict": { + "embed_tokens.weight": shared, + "mtp.lm_head.weight": shared, + "mtp.fc.weight": trained, + }, + }, + os.path.join(runtime, "training_state.pt"), + ) + + merge_mtp_into_base(base, runtime, out, draft_config_path=draft_config) + + merged = load_selected_tensors(out, lambda _key: True) + self.assertTrue(torch.equal(merged["mtp.embed_tokens.weight"], shared)) + self.assertTrue(torch.equal(merged["mtp.lm_head.weight"], shared)) + self.assertTrue(torch.equal(merged["mtp.fc.weight"], trained)) + + +class MTPRegistrationTest(unittest.TestCase): + def test_draft_architecture_is_registered(self): + self.assertIn("Qwen3_5MTPDraftModel", available_drafts()) + self.assertIs(resolve_draft("Qwen3_5MTPDraftModel"), Qwen3_5MTPDraftModel) + + def test_builtin_registry_resolves_mtp(self): + registration = builtin_algorithm_registry().resolve("mtp") + self.assertEqual("mtp", registration.spec.name) + self.assertEqual( + "Qwen3_5MTPDraftModel", + registration.providers.model.draft_config.architecture, + ) + + def test_offline_layout_persists_only_final_hidden(self): + providers = builtin_algorithm_registry().resolve("mtp").providers + layout = providers.offline_for("text").capture_layout + self.assertEqual( + ("input_ids", "loss_mask", "target_last_hidden_states"), + layout.output_names, + ) + self.assertIsNone(layout.aux_feature) + + def test_streaming_layout_exposes_final_hidden(self): + providers = builtin_algorithm_registry().resolve("mtp").providers + layout = providers.server_streaming_for("text").layout + self.assertEqual("target_last_hidden_states", layout.last_hidden_feature) + + +if __name__ == "__main__": + unittest.main()