| |
|
|
| import argparse |
| import json |
| import os |
| import sys |
|
|
| from model_reader import read_model |
| from imatrix_reader import read_imatrix, detect_tied_groups, build_importance_table |
| from classifier import optimal_classify, compute_stats |
| from config_generator import generate_flags, format_flags |
| from quantizer import run_dry_run, run_quantization |
| from constants import CLASS_HARD_FLOORS |
|
|
|
|
| def _get_base_type(model: dict) -> str: |
| is_qat = model.get("features", {}).get("is_qat", False) |
| return "IQ4_XS" if is_qat else "Q5_K_M" |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser( |
| description="SHQ-program: imatrix-driven hybrid quantization" |
| ) |
| parser.add_argument("--model", help="BF16 GGUF model path") |
| parser.add_argument("--imatrix", action="append", default=[], |
| help="Imatrix GGUF path (can be specified multiple times)") |
| parser.add_argument("--imatrix-method", choices=["max", "mean"], default="max", |
| help="How to combine multiple imatrix: max (conservative) or mean (default: max)") |
| parser.add_argument("--size", type=float, default=6800, |
| help="Target file size in MiB (default: 6800 = ~6.6 GB)") |
| parser.add_argument("--output", default=None, help="Output GGUF path") |
| parser.add_argument("--run", action="store_true", help="Execute quantization") |
| parser.add_argument("--show-config", action="store_true", help="Print config and exit") |
| parser.add_argument("--verbose", action="store_true", help="Detailed output") |
| parser.add_argument("--allow-q3-or-lower", action="store_true", |
| help="Allow Q3_K for low-importance tensors (risk of quality loss)") |
| parser.add_argument("--aggro", type=float, default=None, |
| help="[deprecated] Use --size instead") |
| parser.add_argument("--show-floors", action="store_true", |
| help="Print class hard floors and exit") |
|
|
| args = parser.parse_args() |
|
|
| if args.show_floors: |
| _show_floors() |
| return |
|
|
| if not args.model or not args.imatrix: |
| parser.print_usage() |
| print("main.py: error: --model and --imatrix are required") |
| sys.exit(1) |
|
|
| target_mib = args.size |
|
|
| print("=== SHQ-program ===") |
| print(f"Model: {args.model}") |
| if len(args.imatrix) == 1: |
| print(f"Imatrix: {args.imatrix[0]}") |
| else: |
| print(f"Imatrix: {len(args.imatrix)} files ({args.imatrix_method})") |
| for p in args.imatrix: |
| print(f" - {p}") |
| print(f"Target: {target_mib:.0f} MiB ({target_mib / 1024:.2f} GB)") |
| if args.allow_q3_or_lower: |
| print(" --allow-q3-or-lower: low-importance tensors may go to Q3_K") |
| print() |
|
|
| print("[1/4] Reading model...") |
| model = read_model(args.model) |
| print(f" Architecture: {model['architecture']}") |
| print(f" Tensors: {model['n_tensors']}") |
| print(f" Features: {json.dumps(model['features'], indent=2)}") |
|
|
| print("\n[2/4] Reading imatrix...") |
| imatrix_list = [read_imatrix(p) for p in args.imatrix] |
| for im in imatrix_list: |
| print(f" {im['path']}: {im['n_tensors']} tensors, datasets={im['meta'].get('imatrix.datasets', '?')}") |
|
|
| from imatrix_reader import combine_imatrix |
| imatrix = combine_imatrix(imatrix_list, method=args.imatrix_method) |
| print(f" Combined: {imatrix['n_tensors']} tensors") |
|
|
| print("\n[3/4] Detecting tied groups...") |
| tied_groups = detect_tied_groups(imatrix) |
| print(f" Found {len(tied_groups)} tied groups:") |
| for g in tied_groups: |
| if len(g) > 1: |
| print(f" TIED ({len(g)}): {g[0].replace('.weight', '')} = " |
| f"{g[1].replace('.weight', '')}") |
|
|
| imp_table = build_importance_table(imatrix, model) |
|
|
| print("\n[4/4] Classifying tensors (greedy imatrix-driven)...") |
|
|
| |
| assignments, padded_ne_map = optimal_classify( |
| imp_table, tied_groups, model, |
| target_size_mib=target_mib, |
| allow_q3=args.allow_q3_or_lower, |
| ) |
| |
| ne_map = {k: v["n_elements"] for k, v in model.get("tensors", {}).items()} |
| for tname, info in imp_table.items(): |
| if tname not in ne_map: |
| ne_map[tname] = info["n_elements"] |
|
|
| |
| _show_tier_summary(assignments, imp_table, ne_map, padded_ne_map) |
|
|
| base_type = _get_base_type(model) |
| flags = generate_flags(assignments, model, base_type, target_mib) |
| flags["imatrix"] = args.imatrix |
|
|
| print(f"\nConfig (base={flags['base_type']}):") |
| print(format_flags(flags)) |
|
|
| if args.show_config: |
| return |
|
|
| print("\n--- Dry Run ---") |
| dry_size = run_dry_run(flags, args.model) |
| _show_size_result(dry_size, target_mib) |
|
|
| if not args.run: |
| print("\nDry run only. Use --run to execute quantization.") |
| return |
|
|
| if not args.output: |
| base = os.path.splitext(os.path.basename(args.model))[0] |
| args.output = base + "-SHQ.gguf" |
|
|
| print(f"\n--- Running quantization: {args.output} ---") |
| success = run_quantization(flags, args.model, args.output) |
| if success: |
| print("Done!") |
| else: |
| print("Failed!") |
| sys.exit(1) |
|
|
|
|
| def _show_tier_summary(assignments, imp_table, ne_map, padded_ne_map=None): |
| stats = compute_stats(assignments, ne_map, padded_ne_map) |
| |
| print("\n Tier distribution:") |
| for tier in sorted(stats["by_tier_count"].keys()): |
| count = stats["by_tier_count"][tier] |
| mib = stats["by_tier_mib"].get(tier, 0.0) |
| print(f" {tier}: {count} tensors ({mib:.1f} MiB)") |
| print(f" Total estimated size: {stats['total_mib']:.1f} MiB") |
|
|
| ranked = sorted( |
| [(n, v) for n, v in imp_table.items()], |
| key=lambda x: -x[1]["importance_mean"], |
| ) |
| print("\n Top 10 by importance:") |
| for n, v in ranked[:10]: |
| tier = assignments.get(n, "base") |
| display = n.replace(".weight", "").replace(".bias", "") |
| print(f" {display[:52]:52s} imp={v['importance_mean']:10.0f} tier={tier}") |
|
|
|
|
| def _show_size_result(dry_size, target_mib): |
| if dry_size: |
| print(f" Estimated size: {dry_size:.0f} MiB ({dry_size / 1024:.2f} GB)") |
| diff = dry_size - target_mib |
| if diff > 0: |
| print(f" ⚠ Over target by {diff:.0f} MiB") |
| else: |
| print(f" ✓ Under target by {-diff:.0f} MiB") |
| else: |
| print(" ⚠ Could not parse size from dry-run output") |
|
|
|
|
| def _show_floors(): |
| print(" Class hard floors (never below without --allow-q3-or-lower):\n") |
| max_n = max(len(c) for c in CLASS_HARD_FLOORS) |
| for cls, floor in sorted(CLASS_HARD_FLOORS.items()): |
| print(f" {cls:<{max_n}} → {floor}") |
| print(f"\n Default floor (unknown class): Q4_K") |
| print(f" --allow-q3-or-lower enables Q3_K for: ffn_down, attn_output, ssm_out") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
