AlexWortega

/

SIQ-1-35B

like 94

Recipe in https://huggingface.co/AlexWortega/SIQ-1-35B/discussions/2 was somewhat broken for me.
so, I use steps 1-2, then different script (it quantizes most of the sidecar weights to Q8_0):

#!/usr/bin/env python3
"""
merge_mtp_gguf.py — Merge an MTP (Multi-Token Prediction) sidecar GGUF into a base GGUF model.

This script was developed for Qwen3.6-MoE (e.g. SIQ-1-35B-APEX) to attach an MTP layer
(block 40) from a sidecar file to a quantized base model, preserving the base model's
quantization while re-quantizing MTP tensors to optimal types for inference speed.

Key features:
1. Preserves base model tensor quantization (raw bytes, no dequantization).
2. Splits fused ffn_gate_up_exps into separate ffn_gate_exps + ffn_up_exps.
3. Updates metadata: block_count, nextn_predict_layers, vocab_size.
4. Re-quantizes MTP tensors: heavy weights → Q8_0/BF16, norms → F32 (for SYCL compatibility).

Usage:
    python merge_mtp_gguf.py base.gguf mtp.gguf output.gguf

Requirements:
    pip install gguf numpy

Author: Kimi (Moonshot AI) — 2026-06-30
"""

import sys
import argparse
import numpy as np
from pathlib import Path

from gguf import GGUFWriter, GGMLQuantizationType, GGUFValueType
from gguf.gguf_reader import GGUFReader
from gguf.quants import quantize, dequantize, QuantError


def to_py_shape(shape):
    """Convert numpy shape to tuple of native Python ints.

    struct.pack chokes on np.int64, so we cast everything to int().
    """
    return tuple(int(x) for x in shape)


def read_gguf(path):
    """Read a GGUF file and return (reader, metadata_dict, tensors_dict).

    Uses field.contents() instead of field.parts[field.data[0]] because the
    latter only returns the first element for arrays (e.g. rope.dimension_sections
    would come back as length-1 instead of length-4).
    """
    reader = GGUFReader(path)

    metadata = {}
    for key, field in reader.fields.items():
        metadata[key] = field.contents()

    tensors = {}
    for tensor in reader.tensors:
        tensors[tensor.name] = {
            'data': tensor.data,
            'shape': to_py_shape(tensor.shape),        # physical shape from file header
            'data_shape': to_py_shape(tensor.data.shape),  # logical shape (reversed dims; byte-shape for quantized)
            'tensor_type': tensor.tensor_type,
            'n_bytes': tensor.n_bytes,
            'n_elements': tensor.n_elements,
        }

    return reader, metadata, tensors


def get_arch(meta):
    """Decode architecture string from metadata (handles bytes, ndarray, str)."""
    arch = meta.get('general.architecture', '')
    if isinstance(arch, bytes):
        return arch.decode('utf-8')
    elif isinstance(arch, np.ndarray):
        if arch.dtype == object:
            return str(arch[0])
        return bytes(arch).decode('utf-8')
    return str(arch)


def get_int(meta, key, default=0):
    """Extract an integer value from metadata."""
    val = meta.get(key, default)
    if isinstance(val, (list, tuple, np.ndarray)):
        return int(val[0]) if len(val) > 0 else default
    return int(val)


def get_str(meta, key, default=''):
    """Extract a string value from metadata."""
    val = meta.get(key, default)
    if isinstance(val, bytes):
        return val.decode('utf-8')
    return str(val)


def copy_metadata(writer, reader, skip_keys=None):
    """Copy all KV pairs from reader to writer with proper type conversion.

    Skips system keys (GGUF.*) and keys that the writer will set manually.
    """
    skip_keys = skip_keys or set()

    for key, field in reader.fields.items():
        if key in skip_keys:
            continue

        val = field.contents()
        vtype = field.types[0]

        if vtype == GGUFValueType.STRING:
            if val:
                writer.add_string(key, val)
        elif vtype == GGUFValueType.ARRAY:
            if len(val) > 0:
                writer.add_array(key, val)
        elif vtype == GGUFValueType.UINT8:
            writer.add_uint8(key, val)
        elif vtype == GGUFValueType.INT8:
            writer.add_int8(key, val)
        elif vtype == GGUFValueType.UINT16:
            writer.add_uint16(key, val)
        elif vtype == GGUFValueType.INT16:
            writer.add_int16(key, val)
        elif vtype == GGUFValueType.UINT32:
            writer.add_uint32(key, val)
        elif vtype == GGUFValueType.INT32:
            writer.add_int32(key, val)
        elif vtype == GGUFValueType.FLOAT32:
            writer.add_float32(key, val)
        elif vtype == GGUFValueType.BOOL:
            writer.add_bool(key, val)
        elif vtype == GGUFValueType.UINT64:
            writer.add_uint64(key, val)
        elif vtype == GGUFValueType.INT64:
            writer.add_int64(key, val)
        elif vtype == GGUFValueType.FLOAT64:
            writer.add_float64(key, val)


def verify_compatibility(base_meta, mtp_meta, arch):
    """Check that base and MTP models have matching hyperparameters."""
    critical_checks = [
        ("embedding_length", f"{arch}.embedding_length"),
        ("expert_count", f"{arch}.expert_count"),
        ("expert_feed_forward_length", f"{arch}.expert_feed_forward_length"),
        ("attention.head_count", f"{arch}.attention.head_count"),
        ("attention.head_count_kv", f"{arch}.attention.head_count_kv"),
        ("attention.key_length", f"{arch}.attention.key_length"),
        ("attention.value_length", f"{arch}.attention.value_length"),
    ]

    all_ok = True
    for label, key in critical_checks:
        base_val = get_int(base_meta, key)
        mtp_val = get_int(mtp_meta, key)
        status = "OK" if base_val == mtp_val else "MISMATCH"
        if base_val != mtp_val:
            all_ok = False
        print(f"      {status} {label}: base={base_val}, mtp={mtp_val}")

    vocab_key = f"{arch}.vocab_size"
    base_vocab = get_int(base_meta, vocab_key, 0)
    mtp_vocab = get_int(mtp_meta, vocab_key, 0)
    if base_vocab == 0 and mtp_vocab > 0:
        print(f"      INFO vocab_size missing in base, using mtp: {mtp_vocab}")
    elif base_vocab != mtp_vocab and mtp_vocab > 0:
        print(f"      WARN vocab_size differs: base={base_vocab}, mtp={mtp_vocab}")
    else:
        print(f"      OK vocab_size: {base_vocab}")

    return all_ok


def split_gate_up(gate_up_data):
    """Split fused ffn_gate_up_exps into separate gate and up tensors.

    gate_up shape: [H, 2*I, E] = [2048, 1024, 256]
    gate: [H, I, E] = [2048, 512, 256] — first half along dim 1
    up:   [H, I, E] = [2048, 512, 256] — second half along dim 1
    """
    assert gate_up_data.shape[1] == 1024, f"Expected dim 1 = 1024, got {gate_up_data.shape[1]}"
    gate = gate_up_data[:, :512, :]
    up = gate_up_data[:, 512:, :]
    return gate, up


# ---------------------------------------------------------------------------
# Quantization mapping for MTP tensors
# Heavy weights → Q8_0 or BF16 for speed
# Norms (1D/2D) → F32 for SYCL backend compatibility
#   (SYCL ggml_sycl_op_bin_bcast does not support f16 + f32 broadcast)
# ---------------------------------------------------------------------------
MTP_QUANT_MAP = {
    'attn_k.weight': GGMLQuantizationType.Q8_0,
    'attn_q.weight': GGMLQuantizationType.Q8_0,
    'attn_v.weight': GGMLQuantizationType.Q8_0,
    'attn_output.weight': GGMLQuantizationType.Q8_0,
    'ffn_down_exps.weight': GGMLQuantizationType.Q8_0,
    'ffn_gate_exps.weight': GGMLQuantizationType.Q8_0,
    'ffn_up_exps.weight': GGMLQuantizationType.Q8_0,
    'ffn_gate_inp.weight': GGMLQuantizationType.BF16,
    'ffn_gate_inp_shexp.weight': GGMLQuantizationType.BF16,
    'ffn_down_shexp.weight': GGMLQuantizationType.Q8_0,
    'ffn_gate_shexp.weight': GGMLQuantizationType.Q8_0,
    'ffn_up_shexp.weight': GGMLQuantizationType.Q8_0,
    'nextn.eh_proj.weight': GGMLQuantizationType.Q8_0,
    # Norms → F32 (SYCL broadcast compatibility)
    'nextn.enorm.weight': GGMLQuantizationType.F32,
    'nextn.hnorm.weight': GGMLQuantizationType.F32,
    'nextn.shared_head_norm.weight': GGMLQuantizationType.F32,
    'attn_norm.weight': GGMLQuantizationType.F32,
    'attn_q_norm.weight': GGMLQuantizationType.F32,
    'attn_k_norm.weight': GGMLQuantizationType.F32,
    'post_attention_norm.weight': GGMLQuantizationType.F32,
}


def quantize_mtp_tensor(data, src_type, target_type):
    """Re-quantize an MTP tensor from src_type to target_type.

    Steps:
      1. Dequantize to float32 (if needed)
      2. Quantize to target_type via gguf.quants.quantize()

    Falls back to src_type if quantization fails.
    """
    if src_type == target_type:
        return data, target_type

    # Step 1: dequantize to float32
    if src_type == GGMLQuantizationType.F16:
        f32_data = data.astype(np.float32)
    elif src_type == GGMLQuantizationType.F32:
        f32_data = data
    elif src_type == GGMLQuantizationType.BF16:
        f32_data = dequantize(data, src_type)
    else:
        f32_data = dequantize(data, src_type)

    # Step 2: quantize to target
    try:
        qdata = quantize(f32_data, target_type)
        return qdata, target_type
    except (QuantError, NotImplementedError) as e:
        print(f"      WARN quantize to {target_type.name} failed: {e}, keeping {src_type.name}")
        return data, src_type


def get_target_quant(name):
    """Look up the target quantization type for an MTP tensor by its short name."""
    short_name = name.split('.', 2)[-1] if name.count('.') >= 2 else name
    return MTP_QUANT_MAP.get(short_name)


def merge_mtp(base_path, mtp_path, out_path):
    print(f"[1/8] Reading base model: {base_path}")
    base_reader, base_meta, base_tensors = read_gguf(base_path)
    print(f"      Tensors: {len(base_tensors)}")

    print(f"[2/8] Reading MTP sidecar: {mtp_path}")
    mtp_reader, mtp_meta, mtp_tensors = read_gguf(mtp_path)
    print(f"      Tensors: {len(mtp_tensors)}")

    base_arch = get_arch(base_meta)
    mtp_arch = get_arch(mtp_meta)
    print(f"[3/8] Architecture: base={base_arch}, mtp={mtp_arch}")

    if base_arch != mtp_arch:
        print(f"      FAIL architectures do not match!")
        sys.exit(1)

    print(f"[4/8] Checking compatibility...")
    if not verify_compatibility(base_meta, mtp_meta, base_arch):
        print(f"      FAIL critical parameters mismatch!")
        sys.exit(1)
    print(f"      OK models are compatible")

    block_key = f"{base_arch}.block_count"
    base_blocks = get_int(base_meta, block_key)
    mtp_blocks = get_int(mtp_meta, block_key)
    print(f"[5/8] block_count: base={base_blocks}, mtp={mtp_blocks}")

    if mtp_blocks != base_blocks + 1:
        print(f"      WARN expected block_count={base_blocks + 1}, got {mtp_blocks}")

    writer = GGUFWriter(out_path, base_arch)

    print(f"[6/8] Copying metadata...")
    skip = {
        "GGUF.version",
        "GGUF.tensor_count",
        "GGUF.kv_count",
        "general.architecture",  # GGUFWriter sets this in __init__
        block_key,
        f"{base_arch}.nextn_predict_layers",
        f"{base_arch}.vocab_size",
        "general.file_type",
        "general.name",
    }
    copy_metadata(writer, base_reader, skip)

    new_blocks = base_blocks + 1
    writer.add_uint32(block_key, new_blocks)
    print(f"      {block_key}: {base_blocks} → {new_blocks}")

    nextn_key = f"{base_arch}.nextn_predict_layers"
    nextn_val = get_int(mtp_meta, nextn_key)
    if nextn_val == 0:
        nextn_val = 1
    writer.add_uint32(nextn_key, nextn_val)
    print(f"      {nextn_key}: 0 → {nextn_val}")

    vocab_key = f"{base_arch}.vocab_size"
    base_vocab = get_int(base_meta, vocab_key, 0)
    mtp_vocab = get_int(mtp_meta, vocab_key, 0)
    if base_vocab == 0 and mtp_vocab > 0:
        writer.add_uint32(vocab_key, mtp_vocab)
        print(f"      {vocab_key}: added from MTP: {mtp_vocab}")
    elif base_vocab > 0:
        writer.add_uint32(vocab_key, base_vocab)
        print(f"      {vocab_key}: {base_vocab} (from base)")

    base_ftype = get_int(base_meta, "general.file_type")
    writer.add_uint32("general.file_type", base_ftype)
    print(f"      general.file_type: {base_ftype}")

    base_name = get_str(base_meta, "general.name", "Merged Model")
    writer.add_string("general.name", f"{base_name} + MTP")
    print(f"      general.name: {base_name} + MTP")

    print(f"[7/8] Writing tensors...")

    # 7a. Copy all base tensors (preserve their quantization)
    base_count = 0
    for name, info in base_tensors.items():
        writer.add_tensor(name, info['data'], raw_shape=info['data_shape'],
                          raw_dtype=info['tensor_type'])
        base_count += 1
    print(f"      From base model: {base_count} (quantization preserved)")

    # 7b. Add MTP tensors for block 40
    mtp_count = 0
    skipped = 0
    gate_up_split = 0
    quant_count = 0

    for name, info in mtp_tensors.items():
        if name in base_tensors:
            print(f"      SKIP duplicate: {name}")
            skipped += 1
            continue

        # Split fused gate_up into separate gate + up
        if 'ffn_gate_up_exps' in name:
            gate_data, up_data = split_gate_up(info['data'])

            gate_name = name.replace('ffn_gate_up_exps', 'ffn_gate_exps')
            up_name = name.replace('ffn_gate_up_exps', 'ffn_up_exps')

            gate_target = MTP_QUANT_MAP.get('ffn_gate_exps.weight')
            up_target = MTP_QUANT_MAP.get('ffn_up_exps.weight')

            gate_data, gate_type = quantize_mtp_tensor(gate_data, info['tensor_type'], gate_target)
            up_data, up_type = quantize_mtp_tensor(up_data, info['tensor_type'], up_target)

            writer.add_tensor(gate_name, gate_data, raw_shape=to_py_shape(gate_data.shape),
                              raw_dtype=gate_type)
            writer.add_tensor(up_name, up_data, raw_shape=to_py_shape(up_data.shape),
                              raw_dtype=up_type)

            if gate_type != info['tensor_type']:
                quant_count += 1
            print(f"      OK {name} → split into {gate_name} ({gate_type.name}) + {up_name} ({up_type.name})")
            gate_up_split += 1
            mtp_count += 2
            continue

        # Other MTP tensors
        if name.startswith(f'blk.{base_blocks}.'):
            target_type = get_target_quant(name)
            if target_type is not None:
                data, tensor_type = quantize_mtp_tensor(info['data'], info['tensor_type'], target_type)
                if tensor_type != info['tensor_type']:
                    quant_count += 1
            else:
                data, tensor_type = info['data'], info['tensor_type']

            writer.add_tensor(name, data, raw_shape=to_py_shape(data.shape),
                              raw_dtype=tensor_type)
            mtp_count += 1

    print(f"      From MTP model: {mtp_count} (including {gate_up_split} split gate/up)")
    print(f"      Re-quantized: {quant_count}")
    print(f"      Duplicates skipped: {skipped}")
    print(f"      Total tensors: {base_count + mtp_count}")

    nextn_list = [n for n in mtp_tensors if '.nextn.' in n]
    if nextn_list:
        print(f"      NextN tensors added: {len(nextn_list)}")

    print(f"[8/8] Writing file...")
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file()
    writer.close()

    print(f"\n[DONE] Output: {out_path}")
    print(f"       Size: {Path(out_path).stat().st_size / 1024**3:.2f} GB")


def main():
    parser = argparse.ArgumentParser(
        description='Merge an MTP GGUF sidecar into a base GGUF model (Qwen3.6-MoE)'
    )
    parser.add_argument('base', help='Base GGUF model (e.g. SIQ-1-35B)')
    parser.add_argument('mtp', help='MTP GGUF sidecar (19 tensors, block 40)')
    parser.add_argument('output', help='Output merged file')
    args = parser.parse_args()
    merge_mtp(args.base, args.mtp, args.output)


if __name__ == '__main__':
    main()

Hope it helps someone as it helped me