Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
175 lines
9.6 KiB
Python
175 lines
9.6 KiB
Python
# /// script
|
|
# requires-python = ">=3.10"
|
|
# dependencies = ["gguf>=0.19", "numpy>=1.26", "safetensors>=0.5", "torch>=2.6"]
|
|
# [tool.uv.sources]
|
|
# torch = { index = "pytorch-cpu" }
|
|
# [[tool.uv.index]]
|
|
# name = "pytorch-cpu"
|
|
# url = "https://download.pytorch.org/whl/cpu"
|
|
# explicit = true
|
|
# ///
|
|
"""Convert Drama 3's scaled FP8 experts into a native audio.cpp GGUF package.
|
|
|
|
Uses CPU memory one expert bank at a time. The optional native converter stage
|
|
folds codec weight norms, embeds tokenizer/config files, and applies mixed
|
|
quantization. It never loads weights onto a GPU.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import struct
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
import gguf
|
|
import numpy as np
|
|
import torch
|
|
from safetensors import safe_open
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--checkpoint", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, help="Intermediate FP16 GGUF (about 68 GiB)")
|
|
parser.add_argument("--input-gguf", type=Path, help="Reuse a prepared FP16 GGUF instead of exporting again")
|
|
parser.add_argument("--quantized-output", type=Path, help="Standalone mixed-quantization GGUF")
|
|
parser.add_argument("--gguf-tool", type=Path, help="Path to the native audiocpp_gguf executable")
|
|
parser.add_argument("--expert-type", choices=("q2_k", "q3_k", "q4_k", "q5_k", "q6_k", "q8_0"), default="q4_k")
|
|
parser.add_argument("--backbone-type", choices=("q2_k", "q3_k", "q4_k", "q5_k", "q6_k", "q8_0"), default="q8_0",
|
|
help="Dense backbone matrix quantization; audio tower and codec retain their preparation types")
|
|
parser.add_argument("--embedding-type", choices=("f16", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0"), default="q8_0",
|
|
help="Main text embedding storage; defaults to Q8_0 and exposes CUDA-supported row lookup types")
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
arguments = parser.parse_args()
|
|
if arguments.threads <= 0:
|
|
parser.error("--threads must be positive")
|
|
if arguments.input_gguf and arguments.output:
|
|
parser.error("choose --output or --input-gguf")
|
|
if not arguments.input_gguf and not arguments.output:
|
|
parser.error("provide --output or --input-gguf")
|
|
if arguments.input_gguf and not arguments.quantized_output:
|
|
parser.error("--input-gguf requires --quantized-output")
|
|
if arguments.quantized_output and (not arguments.gguf_tool or not arguments.gguf_tool.is_file()):
|
|
parser.error("--quantized-output requires an existing --gguf-tool executable")
|
|
if arguments.output and arguments.quantized_output and arguments.output.resolve() == arguments.quantized_output.resolve():
|
|
parser.error("intermediate and quantized outputs must be different files")
|
|
for destination in (arguments.output, arguments.quantized_output):
|
|
if destination and destination.exists():
|
|
parser.error(f"output already exists: {destination}")
|
|
torch.set_num_threads(arguments.threads)
|
|
configuration = json.loads((arguments.checkpoint / "config.json").read_text())
|
|
if configuration.get("model_type") != "fish_omni":
|
|
parser.error("checkpoint config must describe Fish Drama 3 (fish_omni)")
|
|
|
|
|
|
def quantize_archive(archive):
|
|
if not arguments.quantized_output:
|
|
return
|
|
repository = Path(__file__).resolve().parents[2]
|
|
command = [str(arguments.gguf_tool.resolve()), "--input", str(archive.resolve()),
|
|
"--output", str(arguments.quantized_output.resolve()), "--type", arguments.backbone_type,
|
|
"--family", "fish_drama3", "--model-spec", str(repository / "model_specs/fish_drama3.json"),
|
|
"--root", str(arguments.checkpoint.resolve()),
|
|
"--fold-weight-norm", "codec_weights*", "--keep-type", "codec_weights*=f16",
|
|
"--keep-type", "model_weights/audio_decoder*=orig",
|
|
"--keep-type", "model_weights/model.language_model.embed_tokens.weight=" + arguments.embedding_type]
|
|
layers = configuration["text_config"]["num_hidden_layers"]
|
|
for layer in range(layers):
|
|
command += ["--keep-type", f"model_weights/model.language_model.layers.{layer}.mlp.experts.*={arguments.expert_type}"]
|
|
for name in ("config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "LICENSE", "chat_template.jinja"):
|
|
path = arguments.checkpoint / name
|
|
if path.is_file():
|
|
command += ["--sidecar", f"{path.resolve()}={name}"]
|
|
subprocess.run(command, check=True)
|
|
|
|
|
|
if arguments.input_gguf:
|
|
if not arguments.input_gguf.is_file():
|
|
parser.error(f"input GGUF does not exist: {arguments.input_gguf}")
|
|
quantize_archive(arguments.input_gguf)
|
|
raise SystemExit(0)
|
|
|
|
weight_map = json.loads((arguments.checkpoint / "model.safetensors.index.json").read_text())["weight_map"]
|
|
layer_count = configuration["text_config"]["num_hidden_layers"]
|
|
expert_count = configuration["text_config"]["num_experts"]
|
|
expert_width = configuration["text_config"]["moe_intermediate_size"]
|
|
hidden_width = configuration["text_config"]["hidden_size"]
|
|
shards = {}
|
|
headers = {}
|
|
for shard_name in dict.fromkeys(weight_map.values()):
|
|
shards[shard_name] = safe_open(arguments.checkpoint / shard_name, framework="pt", device="cpu")
|
|
with (arguments.checkpoint / shard_name).open("rb") as shard_file:
|
|
header_length = struct.unpack("<Q", shard_file.read(8))[0]
|
|
headers.update(json.loads(shard_file.read(header_length)))
|
|
|
|
|
|
def load_tensor(tensor_name):
|
|
tensor = shards[weight_map[tensor_name]].get_tensor(tensor_name)
|
|
if tensor.dtype == torch.float8_e4m3fn:
|
|
scale = shards[weight_map[tensor_name + "_scale_inv"]].get_tensor(tensor_name + "_scale_inv")
|
|
tensor = tensor.float() * scale.repeat_interleave(128, 0).repeat_interleave(128, 1)[:tensor.shape[0], :tensor.shape[1]]
|
|
return tensor.to(torch.float32 if tensor.ndim <= 1 or tensor.dtype == torch.float32 else torch.float16).numpy()
|
|
|
|
|
|
tensor_plan = []
|
|
for tensor_name in weight_map:
|
|
if ".experts." in tensor_name or tensor_name.endswith("_scale_inv"):
|
|
continue
|
|
tensor_header = headers[tensor_name]
|
|
tensor_shape = tensor_header["shape"]
|
|
tensor_dtype = np.float32 if len(tensor_shape) <= 1 or tensor_header["dtype"] == "F32" else np.float16
|
|
tensor_plan.append(("model_weights/" + tensor_name, tensor_shape, tensor_dtype, tensor_name))
|
|
for layer_index in range(layer_count):
|
|
for projection_name in ("gate_proj", "up_proj", "down_proj"):
|
|
matrix_shape = (expert_width, hidden_width) if projection_name != "down_proj" else (hidden_width, expert_width)
|
|
tensor_name = f"model.language_model.layers.{layer_index}.mlp.experts.{projection_name}.weight"
|
|
tensor_plan.append(("model_weights/" + tensor_name, (expert_count, *matrix_shape), np.float16, (layer_index, projection_name)))
|
|
|
|
codec = torch.load(arguments.checkpoint / "codec.pth", map_location="cpu", weights_only=True)
|
|
codec = codec.get("state_dict", codec)
|
|
codec = {name.removeprefix("generator."): value for name, value in codec.items()}
|
|
folded_codec = {}
|
|
for tensor_name, tensor in codec.items():
|
|
if tensor_name.endswith("weight_g"):
|
|
continue
|
|
if tensor_name.endswith("weight_v"):
|
|
scale = codec[tensor_name.removesuffix("weight_v") + "weight_g"]
|
|
tensor = tensor * (scale / torch.linalg.vector_norm(tensor, dim=tuple(range(1, tensor.ndim)), keepdim=True))
|
|
tensor_name = tensor_name.removesuffix("weight_v") + "weight"
|
|
if not isinstance(tensor, torch.Tensor):
|
|
continue
|
|
tensor = tensor.to(torch.float32 if tensor.ndim <= 1 else torch.float16).contiguous()
|
|
folded_codec[tensor_name] = tensor.numpy()
|
|
tensor_plan.append(("codec_weights/" + tensor_name, tuple(tensor.shape), folded_codec[tensor_name].dtype, None))
|
|
|
|
arguments.output.parent.mkdir(parents=True, exist_ok=True)
|
|
writer = gguf.GGUFWriter(arguments.output, "fish_drama3")
|
|
writer.add_name("Fish Audio Drama 3")
|
|
writer.add_string("audiocpp.family", "fish_drama3")
|
|
writer.add_string("audiocpp.source_model", "fishaudio/drama-3")
|
|
writer.add_array("audiocpp.tensor_names", [tensor_name for tensor_name, _, _, _ in tensor_plan])
|
|
writer.add_key_value("audiocpp.tensor_ranks", [len(shape) for _, shape, _, _ in tensor_plan], gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.INT32)
|
|
writer.add_key_value("audiocpp.tensor_shapes", [dimension for _, shape, _, _ in tensor_plan for dimension in shape], gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.INT64)
|
|
for tensor_index, (tensor_name, tensor_shape, tensor_dtype, _) in enumerate(tensor_plan):
|
|
writer.add_tensor_info(f"_audiocpp.{tensor_index}", tensor_shape, np.dtype(tensor_dtype), math.prod(tensor_shape) * np.dtype(tensor_dtype).itemsize)
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_ti_data_to_file()
|
|
for tensor_index, (tensor_name, tensor_shape, tensor_dtype, source_name) in enumerate(tensor_plan):
|
|
if tensor_name.startswith("codec_weights/"):
|
|
tensor_values = folded_codec[tensor_name.removeprefix("codec_weights/")]
|
|
elif isinstance(source_name, tuple):
|
|
layer_index, projection_name = source_name
|
|
tensor_values = np.empty(tensor_shape, dtype=tensor_dtype)
|
|
for expert_index in range(expert_count):
|
|
source_tensor_name = f"model.language_model.layers.{layer_index}.mlp.experts.{expert_index}.{projection_name}.weight"
|
|
tensor_values[expert_index] = load_tensor(source_tensor_name)
|
|
else:
|
|
tensor_values = load_tensor(source_name)
|
|
if not np.isfinite(tensor_values).all():
|
|
raise RuntimeError(f"Non-finite values in {tensor_name}")
|
|
writer.write_tensor_data(tensor_values)
|
|
if tensor_index % 20 == 0 or isinstance(source_name, tuple):
|
|
print(f"{tensor_index + 1}/{len(tensor_plan)} {tensor_name}", flush=True)
|
|
writer.close()
|
|
print(f"Wrote {arguments.output} with {len(tensor_plan)} tensors", flush=True)
|
|
quantize_archive(arguments.output)
|