repack_experts_tiered.py

#!/usr/bin/env python3
"""Repack experts into tiered quantization: hot @ 4-bit, cold @ 2-bit.

Reads packed_experts/ (all 4-bit) and hot_experts.json,
produces packed_experts_tiered/ with mixed-quant layer files + manifest.

File format per layer:
  [expert_0_data][expert_1_data]...[expert_N_data]
  Each expert is either 4-bit or 2-bit format (variable size).
  tiered_manifest.json records per-expert: {offset, size, bits}

Usage:
    python repack_experts_tiered.py \
      --model ~/.cache/huggingface/hub/models--mlx-community--Qwen3.5-35B-A3B-4bit \
      --hot-experts hot_experts.json
"""
import argparse
import json
import numpy as np
import os
import sys
from pathlib import Path


def bf16_to_f32(bf16_u16):
    return (bf16_u16.astype(np.uint32) << 16).view(np.float32)


def f32_to_bf16(f32):
    return (f32.view(np.uint32) >> 16).astype(np.uint16)


def unpack_4bit(packed):
    """Extract 8 x 4-bit values from each uint32 (LSB first)."""
    flat = packed.ravel()
    out = np.empty(flat.size * 8, dtype=np.uint8)
    for i in range(8):
        out[i::8] = ((flat >> (i * 4)) & 0xF).astype(np.uint8)
    return out.reshape(packed.shape[:-1] + (packed.shape[-1] * 8,))


def pack_2bit(vals):
    """Pack 16 x 2-bit values into each uint32 (LSB first)."""
    shape = vals.shape
    n_packed = shape[-1] // 16
    flat = vals.reshape(-1, shape[-1])
    out = np.zeros((flat.shape[0], n_packed), dtype=np.uint32)
    for i in range(16):
        out |= flat[:, i::16].astype(np.uint32) << (i * 2)
    return out.reshape(shape[:-1] + (n_packed,))


def requantize_projection(packed_4bit, scales_bf16, biases_bf16, out_dim, in_dim, group_size=64):
    """4-bit -> 2-bit with optimal per-group quantization."""
    num_groups = in_dim // group_size
    vals_4bit = unpack_4bit(packed_4bit).reshape(out_dim, in_dim).astype(np.float32)
    scales_f32 = bf16_to_f32(scales_bf16)
    biases_f32 = bf16_to_f32(biases_bf16)

    vals_grouped = vals_4bit.reshape(out_dim, num_groups, group_size)
    s = scales_f32[:, :, np.newaxis]
    b = biases_f32[:, :, np.newaxis]
    dequant = vals_grouped * s + b

    f_min = dequant.min(axis=2, keepdims=True)
    f_max = dequant.max(axis=2, keepdims=True)
    s2 = (f_max - f_min) / 3.0
    b2 = f_min
    s2_safe = np.where(s2 == 0.0, 1.0, s2)
    vals_2bit = np.clip(np.round((dequant - b2) / s2_safe), 0, 3).astype(np.uint8)

    packed_2bit = pack_2bit(vals_2bit.reshape(out_dim, in_dim))
    new_scales = f32_to_bf16(s2.squeeze(axis=2))
    new_biases = f32_to_bf16(b2.squeeze(axis=2))

    return packed_2bit, new_scales, new_biases


def compute_offsets(moe_intermediate, hidden_dim, group_size, bits):
    """Compute expert byte offsets for given bit width. Mirrors compute_expert_offsets() in infer.m."""
    vals_per_u32 = 32 // bits
    mid, hid, gs = moe_intermediate, hidden_dim, group_size

    def proj_sizes(out_d, in_d):
        w = out_d * ((in_d + vals_per_u32 - 1) // vals_per_u32) * 4
        s = out_d * ((in_d + gs - 1) // gs) * 2
        return w, s, s  # weight, scales, biases

    gw, gs_sz, gb = proj_sizes(mid, hid)
    uw, us, ub = proj_sizes(mid, hid)
    dw, ds, db = proj_sizes(hid, mid)

    offsets = {}
    off = 0
    for name, w, s, b in [("gate", gw, gs_sz, gb), ("up", uw, us, ub), ("down", dw, ds, db)]:
        offsets[f"{name}_w"] = off; off += w
        offsets[f"{name}_s"] = off; off += s
        offsets[f"{name}_b"] = off; off += b

    return offsets, off  # offsets dict + total expert size


def requantize_expert(expert_4bit_blob, off4, size4, off2, size2,
                       moe_intermediate, hidden_dim, group_size=64):
    """Requantize a single expert from 4-bit to 2-bit."""
    output = bytearray(size2)

    projs = [
        ("gate", moe_intermediate, hidden_dim),
        ("up", moe_intermediate, hidden_dim),
        ("down", hidden_dim, moe_intermediate),
    ]

    for name, out_dim, in_dim in projs:
        # Read 4-bit components
        w_start = off4[f"{name}_w"]
        s_start = off4[f"{name}_s"]
        b_start = off4[f"{name}_b"]

        packed_cols_4 = (in_dim + 7) // 8
        w_bytes = out_dim * packed_cols_4 * 4
        num_groups = (in_dim + group_size - 1) // group_size
        s_bytes = out_dim * num_groups * 2

        packed_4bit = np.frombuffer(
            expert_4bit_blob[w_start:w_start + w_bytes], dtype=np.uint32
        ).reshape(out_dim, packed_cols_4)
        scales = np.frombuffer(
            expert_4bit_blob[s_start:s_start + s_bytes], dtype=np.uint16
        ).reshape(out_dim, num_groups)
        biases = np.frombuffer(
            expert_4bit_blob[b_start:b_start + s_bytes], dtype=np.uint16
        ).reshape(out_dim, num_groups)

        packed_2bit, new_scales, new_biases = requantize_projection(
            packed_4bit, scales, biases, out_dim, in_dim, group_size
        )

        # Write 2-bit components
        w2_start = off2[f"{name}_w"]
        s2_start = off2[f"{name}_s"]
        b2_start = off2[f"{name}_b"]
        output[w2_start:w2_start + packed_2bit.nbytes] = packed_2bit.tobytes()
        output[s2_start:s2_start + new_scales.nbytes] = new_scales.tobytes()
        output[b2_start:b2_start + new_biases.nbytes] = new_biases.tobytes()

    return bytes(output)


def main():
    parser = argparse.ArgumentParser(description="Repack experts with tiered quantization")
    parser.add_argument("--model", required=True, help="Path to model directory with packed_experts/")
    parser.add_argument("--hot-experts", required=True, help="Path to hot_experts.json from profile_experts.py")
    parser.add_argument("--group-size", type=int, default=64)
    parser.add_argument("--dry-run", action="store_true", help="Compute sizes without writing files")
    args = parser.parse_args()

    model_path = Path(args.model)
    hot_data = json.loads(Path(args.hot_experts).read_text())

    # Read model config
    config_path = model_path / "config.json"
    if not config_path.exists():
        # Try snapshots subdir (HF cache layout)
        candidates = list(model_path.glob("snapshots/*/config.json"))
        config_path = candidates[0] if candidates else None
    if config_path is None or not config_path.exists():
        print("ERROR: config.json not found"); sys.exit(1)
    config = json.loads(config_path.read_text())
    tc = config.get("text_config", config)

    num_experts = tc["num_experts"]
    num_layers = tc["num_hidden_layers"]
    moe_intermediate = tc.get("moe_intermediate_size", tc.get("moe_intermediate"))
    hidden_dim = tc["hidden_size"]
    group_size = args.group_size

    # Compute offsets for both formats
    off4, size4 = compute_offsets(moe_intermediate, hidden_dim, group_size, 4)
    off2, size2 = compute_offsets(moe_intermediate, hidden_dim, group_size, 2)

    print(f"Model: {model_path}")
    print(f"  {num_layers} layers, {num_experts} experts/layer")
    print(f"  4-bit expert: {size4:,} bytes ({size4/1024/1024:.2f} MB)")
    print(f"  2-bit expert: {size2:,} bytes ({size2/1024/1024:.2f} MB)")
    print(f"  Reduction per cold expert: {(size4-size2)/1024/1024:.2f} MB ({100*(size4-size2)/size4:.1f}%)")

    # Compute total sizes
    total_4bit = num_layers * num_experts * size4
    total_hot = 0
    total_cold = 0
    for l in range(num_layers):
        hot_set = set(hot_data["hot_experts"].get(str(l), []))
        n_hot = len(hot_set)
        n_cold = num_experts - n_hot
        total_hot += n_hot * size4
        total_cold += n_cold * size2

    total_tiered = total_hot + total_cold
    print(f"\n  All 4-bit:    {total_4bit/1024/1024/1024:.2f} GB")
    print(f"  Tiered total: {total_tiered/1024/1024/1024:.2f} GB ({100*total_tiered/total_4bit:.1f}%)")
    print(f"  Savings:      {(total_4bit-total_tiered)/1024/1024/1024:.2f} GB")

    if args.dry_run:
        print("\n[dry-run] No files written.")
        return

    # Create output directory
    experts_dir = model_path / "packed_experts"
    tiered_dir = model_path / "packed_experts_tiered"
    tiered_dir.mkdir(exist_ok=True)

    manifest = {
        "expert_size_4bit": size4,
        "expert_size_2bit": size2,
        "num_layers": num_layers,
        "num_experts": num_experts,
        "threshold": hot_data["threshold"],
        "layers": {},
    }

    for l in range(num_layers):
        src_path = experts_dir / f"layer_{l:02d}.bin"
        dst_path = tiered_dir / f"layer_{l:02d}.bin"

        if not src_path.exists():
            print(f"  WARNING: {src_path} not found, skipping layer {l}")
            continue

        hot_set = set(hot_data["hot_experts"].get(str(l), []))
        n_hot = len(hot_set)
        n_cold = num_experts - n_hot
        print(f"  Layer {l:2d}: {n_hot} hot (4-bit), {n_cold} cold (2-bit)...",
              end="", flush=True)

        # Read entire source layer file
        src_data = src_path.read_bytes()

        # Build tiered layer: sequential experts with variable sizes
        layer_manifest = {"experts": []}
        output_offset = 0
        layer_chunks = []

        for e in range(num_experts):
            expert_start = e * size4
            expert_blob = src_data[expert_start:expert_start + size4]

            if e in hot_set:
                # Keep 4-bit
                layer_chunks.append(expert_blob)
                layer_manifest["experts"].append({
                    "offset": output_offset,
                    "size": size4,
                    "bits": 4,
                })
                output_offset += size4
            else:
                # Requantize to 2-bit
                blob_2bit = requantize_expert(
                    expert_blob, off4, size4, off2, size2,
                    moe_intermediate, hidden_dim, group_size
                )
                layer_chunks.append(blob_2bit)
                layer_manifest["experts"].append({
                    "offset": output_offset,
                    "size": size2,
                    "bits": 2,
                })
                output_offset += size2

        # Write tiered layer file
        with open(dst_path, "wb") as f:
            for chunk in layer_chunks:
                f.write(chunk)

        layer_manifest["file_size"] = output_offset
        manifest["layers"][str(l)] = layer_manifest

        print(f" {output_offset/1024/1024:.1f} MB (was {num_experts * size4 / 1024/1024:.1f} MB)")

    # Write manifest
    manifest_path = tiered_dir / "tiered_manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2))
    print(f"\nWrote manifest: {manifest_path}")
    print(f"Total tiered: {total_tiered/1024/1024/1024:.2f} GB (was {total_4bit/1024/1024/1024:.2f} GB)")


if __name__ == "__main__":
    main()