source/src/speculators/cli/stitch.py

"""
Stitch finetuned MTP weights back into a verifier checkpoint.

Takes a finetuned MTP speculator checkpoint (speculators format) and merges
the trained weights back into the original verifier checkpoint, producing a
self-contained checkpoint directory deployable on vLLM.

Frozen weights (embed_tokens, lm_head) are skipped -- only the MTP layer
weights are replaced.

Usage::

    speculators stitch-mtp ./finetuned-mtp ./Qwen3-Next-80B-A3B-Instruct

    # verifier can be a HuggingFace model ID:
    speculators stitch-mtp ./finetuned-mtp Qwen/Qwen3-Next-80B-A3B-Instruct

    # custom output path (defaults to {verifier-name}-stitched):
    speculators stitch-mtp ./finetuned-mtp ./verifier --output-path ./out
"""

import json
import re
import shutil
from pathlib import Path
from typing import Annotated

import torch
import typer
from huggingface_hub import snapshot_download
from rich.console import Console
from rich.panel import Panel
from rich.progress import (
    BarColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
)
from safetensors import safe_open
from safetensors.torch import save_file

from speculators.convert.mtp import MTP_EXACT_REMAP, MTP_PREFIX_REMAP

console = Console()

INVERSE_MTP_EXACT_REMAP = {v: k for k, v in MTP_EXACT_REMAP.items()}
INVERSE_MTP_PREFIX_REMAP = [(dst, src) for src, dst in MTP_PREFIX_REMAP]

_FROZEN_KEYS = {"embed_tokens.weight", "lm_head.weight"}

_FUSED_GATE_UP_PATTERN = re.compile(r"^(.+\.experts)\.gate_up_proj$")
_FUSED_DOWN_PATTERN = re.compile(r"^(.+\.experts)\.down_proj$")


def _spinner() -> Progress:
    return Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        TimeElapsedColumn(),
        console=console,
    )


def _bar() -> Progress:
    return Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("{task.completed}/{task.total} shards"),
        TimeElapsedColumn(),
        console=console,
    )


def _remap_key(key: str) -> str:
    """Map a speculators-format weight key back to its native HF name."""
    if key in INVERSE_MTP_EXACT_REMAP:
        return INVERSE_MTP_EXACT_REMAP[key]
    for src, dst in INVERSE_MTP_PREFIX_REMAP:
        if key.startswith(src):
            return dst + key[len(src) :]
    return key


def _filter_frozen_keys(
    weights: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    return {k: v for k, v in weights.items() if k not in _FROZEN_KEYS}


def _unfuse_moe_experts(  # noqa: C901
    weights: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Unfuse packed 3D expert tensors back to per-expert format.

    Inverse of ``MTPConverter._fuse_moe_experts``.
    """
    result: dict[str, torch.Tensor] = {}
    gate_up_keys: dict[str, torch.Tensor] = {}
    down_keys: dict[str, torch.Tensor] = {}

    for key, tensor in weights.items():
        m_gu = _FUSED_GATE_UP_PATTERN.match(key)
        m_d = _FUSED_DOWN_PATTERN.match(key)
        if m_gu is not None:
            gate_up_keys[m_gu.group(1)] = tensor
        elif m_d is not None:
            down_keys[m_d.group(1)] = tensor
        else:
            result[key] = tensor

    if not gate_up_keys:
        if down_keys:
            raise ValueError(
                f"Found down_proj without matching gate_up_proj at: {sorted(down_keys)}"
            )
        return weights

    for prefix, gate_up in gate_up_keys.items():
        if prefix not in down_keys:
            raise ValueError(f"Found gate_up_proj at '{prefix}' but missing down_proj")
        down = down_keys[prefix]

        num_experts = gate_up.shape[0]
        if down.shape[0] != num_experts:
            raise ValueError(
                f"Expert count mismatch at '{prefix}': "
                f"gate_up_proj has {num_experts} but down_proj has {down.shape[0]}"
            )
        half = gate_up.shape[1] // 2

        for i in range(num_experts):
            result[f"{prefix}.{i}.gate_proj.weight"] = gate_up[i, :half].contiguous()
            result[f"{prefix}.{i}.up_proj.weight"] = gate_up[i, half:].contiguous()
            result[f"{prefix}.{i}.down_proj.weight"] = down[i].contiguous()

        console.print(f"  Unfused [cyan]{num_experts}[/] experts at [dim]{prefix}[/]")

    orphaned = set(down_keys) - set(gate_up_keys)
    if orphaned:
        raise ValueError(
            f"Found down_proj without matching gate_up_proj at: {sorted(orphaned)}"
        )

    return result


def _resolve_verifier_path(verifier_path: Path) -> Path:
    """Return a local directory, downloading from HF if needed."""
    if verifier_path.exists():
        return verifier_path

    model_id = str(verifier_path)
    console.print(
        f"Verifier path [cyan]{model_id}[/] not found locally, "
        "downloading from HuggingFace..."
    )
    with _spinner() as progress:
        progress.add_task(f"Downloading {model_id}", total=None)
        local_path = snapshot_download(
            repo_id=model_id,
            allow_patterns=[
                "*.json",
                "*.safetensors",
                "*.bin",
                "*.index.json",
            ],
        )
    console.print(f"  Downloaded to [dim]{local_path}[/]")
    return Path(local_path)


def _load_finetuned_weights(
    checkpoint_dir: Path,
) -> dict[str, torch.Tensor]:
    weights: dict[str, torch.Tensor] = {}

    index_path = checkpoint_dir / "model.safetensors.index.json"
    if index_path.exists():
        with index_path.open() as f:
            weight_map = json.load(f)["weight_map"]
        shards = set(weight_map.values())
        with _bar() as progress:
            task = progress.add_task("Loading finetuned weights", total=len(shards))
            for shard in shards:
                with safe_open(str(checkpoint_dir / shard), framework="pt") as f:
                    for key in f.keys():  # noqa: SIM118
                        weights[key] = f.get_tensor(key)
                progress.advance(task)
        return weights

    single = checkpoint_dir / "model.safetensors"
    if single.exists():
        with safe_open(str(single), framework="pt") as f:
            for key in f.keys():  # noqa: SIM118
                weights[key] = f.get_tensor(key)
        return weights

    raise FileNotFoundError(f"No safetensors found at {checkpoint_dir}")


def _stitch_sharded(
    output_dir: Path,
    native_weights: dict[str, torch.Tensor],
) -> None:
    index_path = output_dir / "model.safetensors.index.json"
    with index_path.open() as f:
        index_data = json.load(f)
    weight_map: dict[str, str] = index_data["weight_map"]

    shard_to_new: dict[str, dict[str, torch.Tensor]] = {}
    for key, tensor in native_weights.items():
        shard = weight_map.get(key)
        if shard is None:
            raise ValueError(
                f"Finetuned key '{key}' not found in verifier weight "
                "map. The finetuned checkpoint may not match the "
                "verifier."
            )
        shard_to_new.setdefault(shard, {})[key] = tensor

    with _bar() as progress:
        task = progress.add_task("Stitching shards", total=len(shard_to_new))
        for shard_filename, new_weights in shard_to_new.items():
            shard_path = output_dir / shard_filename
            existing: dict[str, torch.Tensor] = {}
            metadata = None
            with safe_open(str(shard_path), framework="pt") as f:
                metadata = f.metadata()
                for k in f.keys():  # noqa: SIM118
                    existing[k] = f.get_tensor(k)
            for key, tensor in new_weights.items():
                if key in existing and existing[key].shape != tensor.shape:
                    raise ValueError(
                        f"Shape mismatch for '{key}': verifier has "
                        f"{existing[key].shape} but finetuned has {tensor.shape}"
                    )
            existing.update(new_weights)
            save_file(existing, str(shard_path), metadata=metadata)
            progress.advance(task)


def _stitch_single(
    safetensors_path: Path,
    native_weights: dict[str, torch.Tensor],
) -> None:
    existing: dict[str, torch.Tensor] = {}
    metadata = None
    with safe_open(str(safetensors_path), framework="pt") as f:
        metadata = f.metadata()
        for k in f.keys():  # noqa: SIM118
            existing[k] = f.get_tensor(k)
    for key, tensor in native_weights.items():
        if key in existing and existing[key].shape != tensor.shape:
            raise ValueError(
                f"Shape mismatch for '{key}': verifier has "
                f"{existing[key].shape} but finetuned has {tensor.shape}"
            )
    existing.update(native_weights)
    save_file(existing, str(safetensors_path), metadata=metadata)


def stitch(
    finetuned_checkpoint: Path,
    verifier_path: Path,
    output_path: Path,
) -> Path:
    """Stitch finetuned MTP weights back into a verifier checkpoint."""
    console.print(
        Panel(
            f"[bold]Finetuned:[/] {finetuned_checkpoint}\n"
            f"[bold]Verifier:[/]  {verifier_path}\n"
            f"[bold]Output:[/]    {output_path}",
            title="[bold green]MTP Stitch[/]",
            border_style="green",
        )
    )

    verifier_path = _resolve_verifier_path(verifier_path)

    weights = _load_finetuned_weights(finetuned_checkpoint)
    console.print(f"  Loaded [cyan]{len(weights)}[/] finetuned weight tensors")

    frozen_count = sum(1 for k in weights if k in _FROZEN_KEYS)
    weights = _filter_frozen_keys(weights)
    if frozen_count:
        console.print(
            f"  Skipped [yellow]{frozen_count}[/] frozen key(s) (embed_tokens, lm_head)"
        )

    weights = _unfuse_moe_experts(weights)
    native_weights = {_remap_key(k): v for k, v in weights.items()}
    console.print(f"  Remapped [cyan]{len(native_weights)}[/] keys to native format")

    if output_path.exists():
        console.print(
            f"[yellow]Warning:[/] Output directory [dim]{output_path}[/] already "
            "exists; stale files from a previous run may remain."
        )
    with _spinner() as progress:
        progress.add_task("Copying verifier checkpoint", total=None)
        shutil.copytree(verifier_path, output_path, dirs_exist_ok=True)

    index_path = output_path / "model.safetensors.index.json"
    if index_path.exists():
        _stitch_sharded(output_path, native_weights)
    else:
        single = output_path / "model.safetensors"
        if single.exists():
            _stitch_single(single, native_weights)
        else:
            raise FileNotFoundError(f"No safetensors checkpoint found at {output_path}")

    console.print(
        Panel(
            f"[bold]{output_path}[/]",
            title="[bold green]Stitched checkpoint saved[/]",
            border_style="green",
        )
    )
    return output_path


def stitch_command(
    finetuned_checkpoint: Annotated[
        Path,
        typer.Argument(
            help="Path to the finetuned MTP speculator checkpoint.",
        ),
    ],
    verifier_path: Annotated[
        Path,
        typer.Argument(
            help=(
                "Path to the verifier checkpoint, or a HuggingFace "
                "model ID (e.g. Qwen/Qwen3-Next-80B-A3B-Instruct)."
            ),
        ),
    ],
    output_path: Annotated[
        Path | None,
        typer.Option(
            help=(
                "Output directory for the stitched checkpoint. "
                "Defaults to {verifier-name}-stitched."
            ),
        ),
    ] = None,
) -> None:
    """Stitch finetuned MTP weights back into a verifier checkpoint."""
    if output_path is None:
        output_path = Path.cwd() / f"{verifier_path.name}-stitched"

    stitch(
        finetuned_checkpoint=finetuned_checkpoint,
        verifier_path=verifier_path,
        output_path=output_path,
    )