runtime/experimental/nvfp4/checkpoint.py

"""Standalone loader for project-local mixed NVFP4/MXFP8 S2-Pro checkpoints."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any

import torch
from safetensors.torch import load_file

from fish_speech.models.text2semantic.llama import (
    BaseModelArgs,
    DualARTransformer,
    precompute_freqs_cis,
)
from fish_speech.tokenizer import FishTokenizer

from experimental.fp8 import MXFP8Linear
from .modules import NVFP4Linear


CHECKPOINT_FORMAT = "fish-s2-pro-project-local-nvfp4-mixed"


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(8 * 1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def _install_empty_projection(
    model: DualARTransformer,
    record: dict[str, Any],
    *,
    w4a16_max_m: int,
) -> None:
    name = str(record["name"])
    in_features = int(record["in_features"])
    out_features = int(record["out_features"])
    precision = str(record["precision"])
    original = model.get_submodule(name)
    if not isinstance(original, torch.nn.Linear):
        raise TypeError(f"Expected an unmodified Linear at {name}, got {type(original)}")
    if (original.in_features, original.out_features) != (in_features, out_features):
        raise ValueError(f"Checkpoint shape metadata does not match config at {name}")

    if precision.startswith("nvfp4_w4a16_through_m"):
        unsupported = (
            "low_rank_corrected",
            "sparse_channel_corrected",
            "hadamard",
        )
        if any(marker in precision for marker in unsupported):
            raise ValueError(
                f"This standalone loader does not support corrected/rotated NVFP4: {name}"
            )
        replacement: torch.nn.Module = NVFP4Linear(
            torch.empty(
                out_features,
                in_features // 2,
                dtype=torch.uint8,
                device="meta",
            ),
            torch.empty(
                out_features,
                in_features // 16,
                dtype=torch.float8_e4m3fn,
                device="meta",
            ),
            torch.empty((), dtype=torch.float32, device="meta"),
            in_features=in_features,
            out_features=out_features,
            w4a16_max_m=w4a16_max_m,
        )
    elif precision == "mxfp8_w8a8":
        replacement = MXFP8Linear(
            torch.empty(
                out_features,
                in_features,
                dtype=torch.float8_e4m3fn,
                device="meta",
            ),
            torch.empty(
                in_features // 128,
                out_features,
                dtype=torch.int32,
                device="meta",
            ),
            in_features=in_features,
            out_features=out_features,
        )
    else:
        raise ValueError(f"Unsupported precision record for {name}: {precision}")

    parent_name, attribute = name.rsplit(".", 1)
    setattr(model.get_submodule(parent_name), attribute, replacement)


@torch.inference_mode()
def load_mixed_nvfp4_checkpoint(
    path: str | Path,
    *,
    device: str | torch.device = "cuda:0",
    max_length: int = 3072,
    verify_checksums: bool = False,
) -> DualARTransformer:
    """Load a mixed checkpoint without materializing its BF16 source projections."""
    path = Path(path)
    metadata = json.loads((path / "quantization.json").read_text())
    if metadata.get("format") != CHECKPOINT_FORMAT:
        raise ValueError(f"Unsupported checkpoint format: {metadata.get('format')}")
    device = torch.device(device)
    if device.type != "cuda" or torch.cuda.get_device_capability(device)[0] != 12:
        raise RuntimeError("This mixed NVFP4/MXFP8 artifact currently requires sm_120")
    if verify_checksums:
        for filename, record in metadata["checksums"].items():
            file_path = path / filename
            if file_path.stat().st_size != int(record["bytes"]):
                raise RuntimeError(f"Size mismatch for {filename}")
            if _sha256(file_path) != record["sha256"]:
                raise RuntimeError(f"SHA256 mismatch for {filename}")

    conversion = metadata["conversion"]
    records = conversion["records"]
    if len(records) != 180:
        raise ValueError(f"Expected 180 mixed projection records, found {len(records)}")
    if int(conversion["correction_parameters"]) != 0:
        raise ValueError("This loader intentionally rejects correction-bearing artifacts")

    config = BaseModelArgs.from_pretrained(str(path))
    config.max_seq_len = max_length
    with torch.device("meta"):
        model = DualARTransformer(config)
    model.tokenizer = FishTokenizer.from_pretrained(path)
    for record in records:
        _install_empty_projection(
            model,
            record,
            w4a16_max_m=int(conversion["w4a16_max_m"]),
        )

    index_path = path / "model.safetensors.index.json"
    if index_path.is_file():
        index = json.loads(index_path.read_text())
        shard_names = sorted(set(index["weight_map"].values()))
    else:
        shard_names = ["model.safetensors"]
    expected_keys = set(model.state_dict())
    loaded_keys: set[str] = set()
    for shard_name in shard_names:
        shard = load_file(path / shard_name, device="cpu")
        unexpected = set(shard) - expected_keys
        if unexpected:
            raise RuntimeError(
                f"Unexpected checkpoint tensors in {shard_name}: {sorted(unexpected)[:5]}"
            )
        model.load_state_dict(shard, strict=False, assign=True)
        loaded_keys.update(shard)
    missing = expected_keys - loaded_keys
    if missing:
        raise RuntimeError(f"Missing checkpoint tensors: {sorted(missing)[:5]}")

    # These buffers are non-persistent, so reconstruct them after the meta load.
    model.freqs_cis = precompute_freqs_cis(
        config.max_seq_len,
        config.head_dim,
        config.rope_base,
    )
    model.causal_mask = torch.tril(
        torch.ones(config.max_seq_len, config.max_seq_len, dtype=torch.bool)
    )
    model.fast_freqs_cis = precompute_freqs_cis(
        config.num_codebooks,
        config.fast_head_dim,
        config.rope_base,
    )
    model = model.to(device=device).eval()
    sampling = metadata.get("qualified_sampling", {})
    model.fixed_temperature = torch.tensor(
        sampling.get("temperature", 1.0), device=device, dtype=torch.float
    )
    model.fixed_top_p = torch.tensor(
        sampling.get("top_p", 0.85), device=device, dtype=torch.float
    )
    model.fixed_repetition_penalty = torch.tensor(1.5, device=device, dtype=torch.float)
    model._cache_setup_done = False

    nvfp4_count = sum(isinstance(module, NVFP4Linear) for module in model.modules())
    mxfp8_count = sum(isinstance(module, MXFP8Linear) for module in model.modules())
    if (nvfp4_count, mxfp8_count) != (60, 120):
        raise RuntimeError(
            f"Expected 60 NVFP4 and 120 MXFP8 modules, got {nvfp4_count}/{mxfp8_count}"
        )
    meta_tensors = [
        name for name, tensor in model.state_dict().items() if tensor.device.type == "meta"
    ]
    if meta_tensors:
        raise RuntimeError(f"Checkpoint left meta tensors: {meta_tensors[:5]}")
    return model