| """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]}") |
|
|
| |
| 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 |
|
|
