scripts/convert_to_gguf.py

"""
convert_to_gguf.py
------------------
Converts the BF16 safetensors shards stored in a HF dataset repo into a
single Q8_0 GGUF file by:
  1. Downloading one shard at a time
  2. Writing its tensors into the GGUF file (appending)
  3. Deleting the shard immediately to free disk

Peak local disk = one shard (~8.6 GB) + growing GGUF output (~380 GB final).

Uses llama.cpp's gguf Python library for writing, and the same Q8_0
quantisation logic as llama-quantize.

Usage:
    python convert_to_gguf.py \
        --src-repo  memmywinks/DeepSeek_Math_V2 \
        --out-file  /workspace/output/deepseek-math-v2-q8_0.gguf \
        --work-dir  /workspace/scratch
"""

import argparse
import gc
import json
import os
import struct
import sys
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import hf_hub_download, HfApi
from safetensors import safe_open

# gguf library ships with llama.cpp clone
LLAMA_CPP_DIR = os.environ.get("LLAMA_CPP_DIR", "/opt/llama.cpp")
sys.path.insert(0, str(Path(LLAMA_CPP_DIR) / "gguf-py"))
import gguf


# ── Q8_0 quantisation ────────────────────────────────────────────────────────

def quantize_q8_0(tensor: torch.Tensor) -> tuple[np.ndarray, np.ndarray]:
    """Quantise a 2-D BF16/F32 tensor to Q8_0 blocks (32 values per block)."""
    BLOCK = 32
    t = tensor.to(torch.float32).reshape(-1)
    # Pad to multiple of BLOCK
    pad = (BLOCK - len(t) % BLOCK) % BLOCK
    if pad:
        t = torch.cat([t, torch.zeros(pad)])
    blocks = t.reshape(-1, BLOCK)
    amax   = blocks.abs().amax(dim=1, keepdim=True).clamp(min=1e-9)
    scale  = amax / 127.0
    quant  = (blocks / scale).round().clamp(-128, 127).to(torch.int8)
    return quant.numpy(), scale.squeeze(1).to(torch.float16).numpy()


def should_quantize(name: str, shape: torch.Size) -> bool:
    """Quantise weight matrices; keep everything else in F32."""
    if len(shape) < 2:
        return False
    # Skip small matrices and embeddings
    if shape[0] < 32 or shape[1] < 32:
        return False
    return True


# ── Tensor name mapping (HF → GGUF) ─────────────────────────────────────────

def hf_to_gguf_name(name: str) -> str:
    """Best-effort HF tensor name → llama.cpp GGUF name mapping for DeepSeek."""
    n = name
    # token embeddings
    n = n.replace("model.embed_tokens.weight", "token_embd.weight")
    n = n.replace("model.norm.weight", "output_norm.weight")
    n = n.replace("lm_head.weight", "output.weight")
    # layers
    import re
    n = re.sub(r"model\.layers\.(\d+)\.", lambda m: f"blk.{m.group(1)}.", n)
    n = n.replace(".self_attn.q_proj.weight",   ".attn_q.weight")
    n = n.replace(".self_attn.k_proj.weight",   ".attn_k.weight")
    n = n.replace(".self_attn.v_proj.weight",   ".attn_v.weight")
    n = n.replace(".self_attn.o_proj.weight",   ".attn_output.weight")
    n = n.replace(".self_attn.q_a_proj.weight", ".attn_q_a.weight")
    n = n.replace(".self_attn.q_b_proj.weight", ".attn_q_b.weight")
    n = n.replace(".self_attn.kv_a_proj_with_mqa.weight", ".attn_kv_a_mqa.weight")
    n = n.replace(".self_attn.kv_b_proj.weight", ".attn_kv_b.weight")
    n = n.replace(".self_attn.q_a_layernorm.weight", ".attn_q_a_norm.weight")
    n = n.replace(".self_attn.kv_a_layernorm.weight", ".attn_kv_a_norm.weight")
    n = n.replace(".mlp.gate_proj.weight",  ".ffn_gate.weight")
    n = n.replace(".mlp.up_proj.weight",    ".ffn_up.weight")
    n = n.replace(".mlp.down_proj.weight",  ".ffn_down.weight")
    n = n.replace(".mlp.gate.weight",       ".ffn_gate_inp.weight")
    n = n.replace(".input_layernorm.weight",       ".attn_norm.weight")
    n = n.replace(".post_attention_layernorm.weight", ".ffn_norm.weight")
    # MoE experts
    n = re.sub(r"\.mlp\.experts\.(\d+)\.gate_proj\.weight",
               lambda m: f".ffn_gate_exps.weight", n)
    n = re.sub(r"\.mlp\.experts\.(\d+)\.up_proj\.weight",
               lambda m: f".ffn_up_exps.weight", n)
    n = re.sub(r"\.mlp\.experts\.(\d+)\.down_proj\.weight",
               lambda m: f".ffn_down_exps.weight", n)
    return n


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--src-repo",  required=True)
    parser.add_argument("--out-file",  required=True)
    parser.add_argument("--work-dir",  default="/workspace/scratch")
    args = parser.parse_args()

    token    = os.environ.get("HF_TOKEN")
    work_dir = Path(args.work_dir)
    work_dir.mkdir(parents=True, exist_ok=True)
    out_file = Path(args.out_file)
    out_file.parent.mkdir(parents=True, exist_ok=True)

    # ── Fetch index ───────────────────────────────────────────────────────────
    print(f"Fetching index from {args.src_repo}...")
    index_path = hf_hub_download(
        repo_id=args.src_repo,
        filename="model.safetensors.index.json",
        repo_type="dataset",
        token=token,
    )
    with open(index_path) as f:
        index = json.load(f)

    config_path = hf_hub_download(
        repo_id=args.src_repo,
        filename="config.json",
        repo_type="dataset",
        token=token,
    )
    with open(config_path) as f:
        config = json.load(f)

    shard_names = sorted(set(index["weight_map"].values()))
    print(f"Shards: {len(shard_names)}")

    # ── Build ordered tensor list from index ──────────────────────────────────
    # Group tensors by shard so we only open each shard once
    shard_to_tensors: dict[str, list[str]] = {}
    for tensor_name, shard in index["weight_map"].items():
        shard_to_tensors.setdefault(shard, []).append(tensor_name)

    # ── Init GGUF writer ──────────────────────────────────────────────────────
    print(f"Writing GGUF to {out_file} ...")
    writer = gguf.GGUFWriter(str(out_file), arch="deepseek2")

    # Write metadata from config
    writer.add_name("DeepSeek-Math-V2")
    writer.add_context_length(config.get("max_position_embeddings", 163840))
    writer.add_embedding_length(config.get("hidden_size", 7168))
    writer.add_block_count(config.get("num_hidden_layers", 61))
    writer.add_feed_forward_length(config.get("intermediate_size", 18432))
    writer.add_head_count(config.get("num_attention_heads", 128))
    writer.add_head_count_kv(config.get("num_key_value_heads", 128))
    writer.add_rope_freq_base(config.get("rope_theta", 10000.0))

    writer.write_header_to_file()
    writer.write_kv_data_to_file()

    # ── Stream shards and write tensors ──────────────────────────────────────
    for i, shard_name in enumerate(shard_names, 1):
        shard_path = work_dir / shard_name
        print(f"\n[{i}/{len(shard_names)}] {shard_name}")

        # Download
        print("  Downloading...")
        hf_hub_download(
            repo_id=args.src_repo,
            filename=shard_name,
            repo_type="dataset",
            local_dir=str(work_dir),
            token=token,
        )

        # Process tensors in this shard
        tensor_names = shard_to_tensors.get(shard_name, [])
        print(f"  Writing {len(tensor_names)} tensors...")
        with safe_open(str(shard_path), framework="pt", device="cpu") as f:
            for hf_name in tensor_names:
                if hf_name not in f.keys():
                    continue
                tensor = f.get_tensor(hf_name)
                gguf_name = hf_to_gguf_name(hf_name)

                if should_quantize(hf_name, tensor.shape):
                    q, scales = quantize_q8_0(tensor)
                    writer.add_tensor_info(
                        gguf_name, tensor.shape,
                        np.int8, gguf.GGMLQuantizationType.Q8_0
                    )
                    # Write raw Q8_0 block data
                    # Each block: 2 bytes scale (f16) + 32 bytes quants
                    n_blocks = q.reshape(-1, 32).shape[0]
                    block_data = bytearray()
                    q_flat = q.reshape(-1, 32)
                    for b in range(n_blocks):
                        block_data += struct.pack('e', float(scales[b]))  # f16
                        block_data += q_flat[b].tobytes()
                    writer.write_tensor_data(bytearray(block_data))
                else:
                    np_tensor = tensor.to(torch.float32).numpy()
                    writer.add_tensor(gguf_name, np_tensor)

                del tensor
            gc.collect()

        # Delete shard to free disk
        shard_path.unlink(missing_ok=True)
        print(f"  Done. GGUF size so far: {out_file.stat().st_size/1e9:.1f} GB")

    writer.write_tensors_to_file()
    writer.close()
    print(f"\n✔ GGUF written: {out_file}  ({out_file.stat().st_size/1e9:.1f} GB)")


if __name__ == "__main__":
    main()