convert_qwen3_v2.py

#!/usr/bin/env python3
"""
Memory-efficient unary converter for Qwen3.
Processes one safetensors shard at a time to avoid OOM.
(c) 2026 OpenTransformers Ltd / Scott Bisset
"""
import numpy as np
import os, sys, json, time, gc

def quantize_unary(w_fp32, n_planes):
    """Quantize weight matrix to unary bitplane format"""
    out_dim, in_dim = w_fp32.shape
    max_val = n_planes

    abs_max = np.abs(w_fp32).max(axis=1, keepdims=True)
    abs_max = np.where(abs_max == 0, 1.0, abs_max)
    scaled = w_fp32 / abs_max * max_val
    rounded = np.clip(np.round(scaled), -max_val, max_val).astype(np.int32)

    scales = (abs_max.flatten() / max_val).astype(np.float32)
    signs = (rounded < 0)
    magnitudes = np.abs(rounded)

    chunks = (in_dim + 63) // 64
    padded = chunks * 64

    if padded > in_dim:
        signs = np.pad(signs, ((0,0),(0,padded-in_dim)), constant_values=False)
        magnitudes = np.pad(magnitudes, ((0,0),(0,padded-in_dim)), constant_values=0)

    sign_bits = np.packbits(signs.astype(np.uint8), axis=1, bitorder='little')
    sign_u64 = sign_bits.view(np.uint64)[:, :chunks]

    plane_bits = np.zeros((n_planes, out_dim, chunks), dtype=np.uint64)
    for p in range(n_planes):
        mask = (magnitudes > p)
        packed = np.packbits(mask.astype(np.uint8), axis=1, bitorder='little')
        plane_bits[p] = packed.view(np.uint64)[:, :chunks]

    return sign_u64, plane_bits, scales

def convert_model(model_dir, output_dir, n_planes=7):
    os.makedirs(output_dir, exist_ok=True)

    config = json.load(open(os.path.join(model_dir, "config.json")))
    n_layers = config["num_hidden_layers"]
    hidden = config["hidden_size"]

    # Load index to know which keys are in which shard
    index_file = os.path.join(model_dir, "model.safetensors.index.json")
    if os.path.exists(index_file):
        index = json.load(open(index_file))
        weight_map = index["weight_map"]
        shards = sorted(set(weight_map.values()))
    else:
        # Single shard
        shards = [f for f in os.listdir(model_dir) if f.endswith('.safetensors')]
        weight_map = None

    print(f"Model: {n_layers} layers, hidden={hidden}, n_planes={n_planes}")
    print(f"Shards: {len(shards)}")

    manifest = {"unary": {}, "fp16": {}, "n_planes": n_planes, "n_layers": n_layers, "config": config}
    total_converted = 0
    total_linear = 0

    # Count total linear layers
    if weight_map:
        total_linear = sum(1 for k in weight_map if k.endswith(".weight") and "proj" in k)
    print(f"Total linear layers to convert: {total_linear}")

    import torch
    from safetensors import safe_open

    for shard_idx, shard in enumerate(shards):
        shard_path = os.path.join(model_dir, shard)
        print(f"\n=== Shard {shard_idx+1}/{len(shards)}: {shard} ===")

        with safe_open(shard_path, framework="pt") as f:
            keys = list(f.keys())
            print(f"  {len(keys)} tensors in shard")

            for key in sorted(keys):
                fname = key.replace(".", "_")
                is_linear = key.endswith(".weight") and "proj" in key and f.get_tensor(key).dim() == 2

                if is_linear:
                    # Check if already converted
                    sign_path = os.path.join(output_dir, f"{fname}.sign")
                    if os.path.exists(sign_path):
                        w = f.get_tensor(key)
                        manifest["unary"][key] = list(w.shape)
                        total_converted += 1
                        print(f"  [SKIP] {key} already converted")
                        continue

                    w = f.get_tensor(key).float().numpy()
                    t0 = time.time()
                    sign, planes, scales = quantize_unary(w, n_planes)
                    dt = time.time() - t0

                    np.array(sign).tofile(os.path.join(output_dir, f"{fname}.sign"))
                    np.array(planes).tofile(os.path.join(output_dir, f"{fname}.planes"))
                    np.array(scales).tofile(os.path.join(output_dir, f"{fname}.scales"))

                    orig_mb = w.nbytes / 1e6
                    comp_mb = (sign.nbytes + planes.nbytes + scales.nbytes) / 1e6
                    total_converted += 1
                    manifest["unary"][key] = list(w.shape)
                    print(f"  [{total_converted}/{total_linear}] {key}: {list(w.shape)} -> {comp_mb:.1f}MB ({orig_mb/comp_mb:.1f}x) [{dt:.1f}s]")

                    del w, sign, planes, scales
                else:
                    # FP16 weight (norms, embeddings, etc)
                    fp16_path = os.path.join(output_dir, f"{fname}.fp16")
                    if os.path.exists(fp16_path):
                        w = f.get_tensor(key)
                        manifest["fp16"][key] = list(w.shape)
                        print(f"  [SKIP] {key} already saved")
                        continue

                    w = f.get_tensor(key).float().numpy()
                    w_fp16 = w.astype(np.float16)
                    w_fp16.view(np.uint16).tofile(fp16_path)
                    manifest["fp16"][key] = list(w.shape)
                    print(f"  [FP16] {key}: {list(w.shape)} ({w_fp16.nbytes/1e6:.1f}MB)")
                    del w, w_fp16

        # Force GC between shards
        gc.collect()
        print(f"  Shard done, memory freed")

    # Save manifest
    with open(os.path.join(output_dir, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)

    # Copy config
    import shutil
    for cf in ["config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]:
        src = os.path.join(model_dir, cf)
        if os.path.exists(src):
            shutil.copy(src, os.path.join(output_dir, cf))

    # Summary
    total_unary = sum(os.path.getsize(os.path.join(output_dir, f))
                      for f in os.listdir(output_dir)
                      if f.endswith((".sign", ".planes", ".scales")))
    total_fp16 = sum(os.path.getsize(os.path.join(output_dir, f))
                     for f in os.listdir(output_dir)
                     if f.endswith(".fp16"))

    print(f"\n=== CONVERSION COMPLETE ===")
    print(f"Unary linear:  {total_unary/1e9:.2f} GB")
    print(f"FP16 other:    {total_fp16/1e9:.2f} GB")
    print(f"Total:         {(total_unary+total_fp16)/1e9:.2f} GB")

if __name__ == "__main__":
    model_dir = sys.argv[1] if len(sys.argv) > 1 else "qwen3-4b-thinking-hf"
    output_dir = sys.argv[2] if len(sys.argv) > 2 else "qwen3-4b-thinking-unary"
    n_planes = int(sys.argv[3]) if len(sys.argv) > 3 else 7
    convert_model(model_dir, output_dir, n_planes)