#!/usr/bin/env python3 """ASHQ1 — All-in-One Imatrix-Driven Hybrid Quantization Engine.""" import argparse,heapq,json,math,os,re,subprocess,sys,warnings from bisect import bisect_left from collections import defaultdict from dataclasses import dataclass,field from functools import lru_cache from typing import Any,Dict,List,Set,Tuple import gguf import numpy as np # ═══════════════════════════════════════════════════════════════════════════ # 1. CONSTANTS & TENSOR CLASSIFICATION HELPERS # ═══════════════════════════════════════════════════════════════════════════ GGUF_TYPE_NAMES={ 0:"F32",1:"F16",2:"Q4_0",3:"Q4_1",6:"Q5_0",7:"Q5_1",8:"Q8_0",9:"Q8_1", 10:"Q2_K",11:"Q3_K",12:"Q4_K",13:"Q5_K",14:"Q6_K",15:"Q8_K", 16:"IQ2_XXS",17:"IQ2_XS",18:"IQ3_XXS",19:"IQ1_S",20:"IQ4_NL", 21:"IQ3_S",22:"IQ2_S",23:"IQ4_XS",29:"IQ1_M",30:"BF16", } GGUF_TYPE_NAMES_INV={v:k for k,v in GGUF_TYPE_NAMES.items()} # Ordered from lowest to highest precision TIER_ORDER=[ "IQ1_S","IQ2_XXS","IQ2_XS","IQ2_S", "IQ3_XXS","Q3_K","IQ3_S", "IQ4_XS","IQ4_NL","Q4_K","Q5_K","Q6_K","Q8_0","F16", ] # Exact bits per weight from ggml block structs (ggml_type_sizef * 8) TIER_BPW={ "IQ1_S":1.5625, "IQ2_XXS":2.0625, "IQ2_XS":2.3125, "IQ2_S":2.5, "IQ3_XXS":3.0625, "Q3_K":3.4375, "IQ3_S":3.44, "IQ4_XS":4.25, "IQ4_NL":4.5, "Q4_K":4.5, "Q5_K":5.5, "Q6_K":6.5625, "Q8_0":8.5, "F16":16.0, } OVERHEAD_CACHE="ashq1-overhead.json" def _overhead_cache_path(model_path:str)->str: return os.path.join(os.path.dirname(os.path.abspath(model_path)),OVERHEAD_CACHE) def _overhead_key(model_path:str,profile:str)->str: return f"{os.path.basename(model_path)}|{profile}" def load_overhead_factor(model_path:str,profile:str="quality")->float: p=_overhead_cache_path(model_path) if os.path.isfile(p): try: with open(p,"r",encoding="utf-8") as f: v=float(json.load(f).get(_overhead_key(model_path,profile),1.0)) if 0.9<=v<=1.1:return v except Exception:pass return 1.0 def save_overhead_factor(model_path:str,factor:float,profile:str="quality"): if not(0.9<=factor<=1.1):return p=_overhead_cache_path(model_path) data={} if os.path.isfile(p): try: with open(p,"r",encoding="utf-8") as f:data=json.load(f) except Exception:data={} data[_overhead_key(model_path,profile)]=round(factor,4) try: with open(p,"w",encoding="utf-8") as f:json.dump(data,f,indent=2) except Exception as e: print(f" [WARN] Overhead cache write failed: {e}") GGUF_OVERHEAD_FACTOR=1.000 QUANT_RANK={tier:i for i,tier in enumerate(TIER_ORDER)} # Hard floors per tensor class to guarantee stability across long context CLASS_HARD_FLOORS={ "gate":"Q5_K", "attn_proj":"IQ3_S", "ffn_gate_up":"IQ2_XXS", "ffn_down":"IQ3_S", "norms":"F16", "ssm_params":"F16", # Recurrent memory state requires Q8_0 to prevent state collapse in long contexts "gdn_state":"Q8_0", "mtp":"Q5_K", "embd":"IQ4_XS", "shexp":"Q4_K", } CLASS_MAX_TIER={ "gate":"F16", "attn_proj":"Q8_0", "ffn_gate_up":"F16", "ffn_down":"F16", "norms":"F16", "ssm_params":"F16", "gdn_state":"F16", "mtp":"F16", "embd":"F16", "shexp":"F16", } CAN_Q3={"ffn_gate","ffn_up","ffn_down","attn_output","ssm_out"} ALLOW_LOWER_FLOOR="IQ2_XXS" DEFAULT_FLOOR="Q4_K" TIER_FLOORS={ "nano":{ "gate":"Q6_K", "attn_proj":"IQ3_S", "ffn_gate_up":"IQ2_XXS", "ffn_down":"IQ2_S", "norms":"F16", "ssm_params":"F16", "gdn_state":"Q8_0", "mtp":"IQ4_XS", "embd":"IQ4_XS", "shexp":"Q4_K", "unknown":"IQ3_XXS", }, "mini":{ "gate":"Q6_K", "attn_proj":"IQ4_XS", "ffn_gate_up":"IQ3_S", "ffn_down":"IQ4_XS", "norms":"F16", "ssm_params":"F16", "gdn_state":"Q8_0", "mtp":"Q5_K", "embd":"IQ4_XS", "shexp":"Q5_K", "unknown":"IQ4_XS", }, "compact":{ "gate":"Q6_K", "attn_proj":"IQ4_XS", "ffn_gate_up":"IQ3_S", "ffn_down":"IQ4_XS", "norms":"F16", "ssm_params":"F16", "gdn_state":"Q8_0", "mtp":"Q8_0", "embd":"IQ4_XS", "shexp":"Q5_K", "unknown":"Q4_K", }, "quality":{ "gate":"Q6_K", "attn_proj":"Q4_K", "ffn_gate_up":"Q4_K", "ffn_down":"Q4_K", "norms":"F16", "ssm_params":"F16", "gdn_state":"Q8_0", "mtp":"Q8_0", "embd":"Q5_K", "shexp":"Q5_K", "unknown":"Q4_K", }, "fidelity":{ "gate":"Q8_0", "attn_proj":"Q5_K", "ffn_gate_up":"Q5_K", "ffn_down":"Q5_K", "norms":"F16", "ssm_params":"F16", "gdn_state":"Q8_0", "mtp":"Q8_0", "embd":"Q5_K", "shexp":"Q6_K", "unknown":"Q5_K", } } # Monotonic tier progression for the speculative draft head PROFILE_MTP_TIER={"nano":"Q5_K","mini":"Q6_K","compact":"Q8_0"} MTP_DEPLOY_TIER="Q8_0" # Ceilings applied when weights originate from an AutoRound int4 optimization INT4_LINEAGE_CAP="Q5_K" INT4_MTP_CAP="Q6_K" INT4_GDN_CAP="Q8_0" INT4_GATE_CAP="Q6_K" INT4_CAP_CLASSES={"attn_proj","ffn_gate_up","ffn_down","shexp"} # llama-quantize rejects sub-4-bit tiers on tensors absent from the imatrix NO_IMATRIX_MIN_TIER="IQ4_XS" TENSOR_CLASS={ "attn_gate":"gate", "ssm_alpha":"gdn_state", "ssm_beta":"gdn_state", "ssm_beta_alpha":"gdn_state", "ssm_ba":"gdn_state", "attn_q":"attn_proj", "attn_k":"attn_proj", "attn_v":"attn_proj", "attn_qkv":"attn_proj", "attn_output":"attn_proj", "attn_out":"attn_proj", "ffn_gate":"ffn_gate_up", "ffn_up":"ffn_gate_up", "ffn_down":"ffn_down", "ssm_in":"attn_proj", "ssm_out":"attn_proj", "ssm_d":"ssm_params", "ssm_norm":"norms", "ssm_conv1d":"norms", "router":"norms", "ssm_dt":"ssm_params", "ssm_a":"ssm_params", "nextn":"mtp", "ffn_gate_exps":"ffn_gate_up", "ffn_up_exps":"ffn_gate_up", "ffn_down_exps":"ffn_down", "ffn_gate_inp":"norms", "vocoder":"norms", "decoder_wave":"norms", "codec_decoder":"norms", "patch_embd":"norms", "patch_embedding":"norms", "position_embd":"norms", "post_ln":"norms", "pre_ln":"norms", "ln1":"norms", "ln2":"norms", "ln_q":"norms", "mm_input_norm":"norms", "mm_soft_emb_norm":"norms", "mm_0":"embd", "mm_1":"embd", "mm_2":"embd", "0":"embd", "1":"embd", "2":"embd", "mm_proj":"embd", "v_enc_embd":"norms", "q_proj":"attn_proj", "k_proj":"attn_proj", "v_proj":"attn_proj", "o_proj":"attn_proj", "in_proj_a":"gdn_state", "in_proj_b":"gdn_state", "in_proj_qkv":"attn_proj", "in_proj_z":"attn_proj", "out_proj":"attn_proj", "gate_proj":"ffn_gate_up", "up_proj":"ffn_gate_up", "down_proj":"ffn_down", "attn_sinks":"norms", "ffn_exp_probs_b":"norms", "exp_probs_b":"norms", "per_layer_token_embd":"embd", "ffn_gate_shexp":"shexp", "ffn_up_shexp":"shexp", "ffn_down_shexp":"shexp", "ffn_gate_inp_shexp":"norms", } ARCH_FEATURES={ "qwen35":{ "has_qkv":True,"has_ssm":True,"has_mtp":True,"has_moe":False, "is_qat":False,"prefix":"blk","n_layers":32, }, "mellum2":{ "has_qkv":False,"has_ssm":False,"has_mtp":False,"has_moe":True, "is_qat":False,"prefix":"blk","n_layers":28, }, "gemma4":{ "has_qkv":False,"has_ssm":False,"has_mtp":False,"has_moe":False, "is_qat":True,"prefix":"blk","n_layers":48, }, "dense":{ "has_qkv":False,"has_ssm":False,"has_mtp":False,"has_moe":False, "is_qat":False,"prefix":"blk","n_layers":32, }, } def strip_weight(name:str)->str: return name.lstrip(".").removesuffix(".weight").removesuffix(".bias") def get_tensor_type(name:str)->str: parts=strip_weight(name).split(".") if len(parts)>=3 and parts[0] in("v","a") and parts[1] in("blk","BLK"): return parts[3] if len(parts)>=4 else "unknown" if len(parts)>=2 and parts[0] in("blk","BLK"): return parts[2] if len(parts)>=3 else "unknown" if "token_embd" in name: return "token_embd" if name.startswith("output") and "norm" not in name: return "output" return name def get_tensor_class(ttype:str)->str: if ttype in TENSOR_CLASS: return TENSOR_CLASS[ttype] if "." in ttype: tail=ttype.split(".")[-1] if tail in TENSOR_CLASS: return TENSOR_CLASS[tail] for part in reversed(ttype.split(".")): if part in TENSOR_CLASS: return TENSOR_CLASS[part] if "norm" in ttype or "scale" in ttype or ttype.startswith(("ln","pre_ln","post_ln")): return "norms" if ttype.startswith("ssm_"): return "ssm_params" if ttype in("token_embd","output","embed_tokens","lm_head","vision_embedder","audio_embedder"): return "embd" return "unknown" def is_mtp_tensor(name:str,n_layers:int=32,mtp_layers:set | None=None)->bool: if "nextn" in name or name.startswith("mtp."): return True if mtp_layers is not None: return get_layer_number(name) in mtp_layers layer=get_layer_number(name) return layer is not None and layer>=n_layers def get_layer_number(name:str)->int | None: parts=strip_weight(name).split(".") if len(parts)>=3 and parts[0] in("v","a") and parts[1] in("blk","BLK"): try:return int(parts[2]) except ValueError:return None if len(parts)>=2 and parts[0] in("blk","BLK"): try:return int(parts[1]) except ValueError:return None return None # ═══════════════════════════════════════════════════════════════════════════ # 2. UTILITY FUNCTIONS # ═══════════════════════════════════════════════════════════════════════════ def type_name(type_id:int)->str: return GGUF_TYPE_NAMES.get(type_id,f"UNKNOWN({type_id})") def type_id(name:str)->int: return GGUF_TYPE_NAMES_INV.get(name,-1) def parse_size_line(line:str)->float | None: m=re.search(r"quant size\s*=\s*([0-9.]+)\s*MiB",line) if m:return float(m.group(1)) m=re.search(r"model size\s*=\s*([0-9.]+)\s*MiB",line) if m:return float(m.group(1)) return None def parse_quant_size(output:str)->float | None: m=re.search(r"quant size\s*=\s*([0-9.]+)\s*MiB",output) if m:return float(m.group(1)) m=re.search(r"model size\s*=\s*([0-9.]+)\s*MiB",output) if m:return float(m.group(1)) return None def parse_fallback_warnings(output:str)->int: return len(re.findall(r"converting to\s+(q[0-9]_[0-9KMS]|iq[0-9])",output)) def format_size(mib:float)->str: if mib>=1024:return f"{mib/1024:.2f} GB" return f"{mib:.0f} MiB" # ═══════════════════════════════════════════════════════════════════════════ # 3. MODEL READER # ═══════════════════════════════════════════════════════════════════════════ def detect_architecture(tensors:dict)->str: names=list(tensors.keys()) has_ssm=any("ssm_" in n for n in names) has_qkv=any("attn_qkv" in n for n in names) has_moe=any("exps" in n for n in names) has_gemma_specific=any(t in n for t in("layer_output_scale","post_attention_norm","post_ffw_norm") for n in names) if has_ssm and has_qkv:return "qwen35" if has_moe:return "mellum2" if has_gemma_specific:return "gemma4" return "dense" def _detect_prefix(tensors:dict)->str: for name in tensors: if name.startswith("BLK."):return "BLK" return "blk" def _estimate_layers(tensors:dict)->int: max_layer=-1 for name in tensors: parts=name.split(".") if len(parts)>=3 and parts[0] in("v","a") and parts[1] in("blk","BLK"): parts=parts[1:] if len(parts)>=2 and parts[0] in("blk","BLK"): try: layer=int(parts[1]) if layer>max_layer:max_layer=layer except ValueError:pass return max_layer+1 if max_layer>=0 else 0 def read_model(path:str)->dict: r=gguf.GGUFReader(path) tensors={} meta={} for k,v in r.fields.items(): try: data=v.data if isinstance(data,np.ndarray):data=data.tolist() elif isinstance(data,(np.generic,)):data=data.item() meta[k]=data except Exception:meta[k]=str(v) for t in r.tensors: shape=list(t.shape) name=t.name n_elements=int(np.prod(shape)) tensors[name]={"shape":shape,"n_elements":n_elements,"size_mib":n_elements*2/1024/1024} arch=detect_architecture(tensors) arch_features=ARCH_FEATURES.get(arch,{}).copy() prefix=_detect_prefix(tensors) n_layers=_estimate_layers(tensors) NEXTN_HINT=("nextn","eh_proj","mtp.fc") has_nextn=any(any(h in n for h in NEXTN_HINT) for n in tensors) BODY=("attn_qkv","attn_q","attn_k","attn_v","attn_output","ssm_","ffn_","eh_proj","exps") last=n_layers-1 last_names=[n for n in tensors if n.startswith(f"blk.{last}.") or n.startswith(f"BLK.{last}.")] last_has_body=any(any(h in n for h in BODY) for n in last_names) last_has_nextn=any(any(h in n for h in NEXTN_HINT) for n in last_names) has_blk_mtp=bool(last_names) and last_has_body and last_has_nextn mtp_layers=[last] if has_blk_mtp else [] has_mtp=has_nextn and(has_blk_mtp or any("nextn" in n for n in tensors)) has_moe=any("exps" in n for n in tensors) if last_names and not last_has_body: print(f" [WARN] blk.{last} contains only {len(last_names)} norm tensor(s) — truncated MTP head ignored") n_layers-=1 elif has_blk_mtp: n_layers-=1 if arch=="mellum2" and arch_features.get("moe_intermediate_size",0)==0: arch_features["moe_intermediate_size"]=896 arch_features["prefix"]=prefix if n_layers>0:arch_features["n_layers"]=n_layers if has_moe:arch_features["has_moe"]=True arch_features["has_mtp"]=has_mtp arch_features["mtp_layers"]=mtp_layers return{ "path":path, "architecture":arch, "features":arch_features, "tensors":tensors, "n_tensors":len(tensors), "meta":meta, } # ═══════════════════════════════════════════════════════════════════════════ # 4. IMATRIX READER # ═══════════════════════════════════════════════════════════════════════════ def read_imatrix(path:str)->dict: r=gguf.GGUFReader(path) raw={} meta={} for k,v in r.fields.items(): try:meta[k]=v.data except Exception:meta[k]=str(v) by_name={t.name:t for t in r.tensors} for t in r.tensors: name=t.name if name.endswith(".in_sum2"): base=name[:-8] raw.setdefault(base,{})["in_sum2"]=np.asarray(t.data,dtype=np.float64) src=by_name.get(base+".in_sum") if src is not None:raw[base]["in_sum"]=np.asarray(src.data,dtype=np.float64) elif name.endswith(".counts"): base=name[:-7] raw.setdefault(base,{})["counts"]=float(np.mean(np.asarray(t.data,dtype=np.float64))) result={} for base,data in raw.items(): if "in_sum2" not in data:continue arr=data["in_sum2"] n=arr.size imp_mean=float(np.mean(arr)) outlier_ratio=None if "in_sum" in data and imp_mean>0: in_sum_arr=data["in_sum"] total_sum=float(np.sum(in_sum_arr)) if total_sum>0 and n>0: total_sum2=float(np.sum(arr**2)) mean_sq=total_sum2/n mean_x=total_sum/n if mean_x>0: outlier_ratio=math.sqrt(mean_sq /(mean_x**2)-1) result[base]={ "importance_mean":imp_mean, "importance_sum":float(np.sum(arr)), "importance_max":float(np.max(arr)), "importance_min":float(np.min(arr)), "n_elements":n, "in_sum2_raw":arr, "outlier_ratio":outlier_ratio, } return{ "path":path, "tensors":result, "n_tensors":len(result), "meta":meta, } def combine_imatrix(imatrix_list:List[dict],method:str="max")->dict: if not imatrix_list:return{"path":"","tensors":{},"n_tensors":0,"meta":{}} if len(imatrix_list)==1:return imatrix_list[0] all_names=set() for im in imatrix_list:all_names.update(im["tensors"].keys()) combined_tensors={} for name in all_names: vals=[] in_sum2_raw=None n_elements=0 for im in imatrix_list: if name in im["tensors"]: t=im["tensors"][name] vals.append(t["importance_mean"]) if in_sum2_raw is None and "in_sum2_raw" in t:in_sum2_raw=t["in_sum2_raw"] n_elements=max(n_elements,t["n_elements"]) if not vals:continue if method=="max":imp_mean=max(vals) elif method in("mean","weighted_mean"):imp_mean=sum(vals)/len(vals) else:imp_mean=max(vals) combined_tensors[name]={ "importance_mean":imp_mean, "importance_sum":imp_mean*n_elements, "importance_max":max(v.get("importance_max",0) for im in imatrix_list if name in im["tensors"] for v in[im["tensors"][name]]), "importance_min":min(v.get("importance_min",float('inf')) for im in imatrix_list if name in im["tensors"] for v in[im["tensors"][name]]), "n_elements":n_elements, "in_sum2_raw":in_sum2_raw, "outlier_ratio":max((t.get("outlier_ratio") for im in imatrix_list if name in im["tensors"] for t in[im["tensors"][name]] if t.get("outlier_ratio") is not None),default=None), } combined_meta={} for im in imatrix_list: for k,v in im["meta"].items(): if k not in combined_meta:combined_meta[k]=v return{ "path":"+".join(im["path"] for im in imatrix_list), "tensors":combined_tensors, "n_tensors":len(combined_tensors), "meta":combined_meta, } def detect_tied_groups(imatrix:dict,atol:float=1e-5)->list: coarse={} no_raw=[] for name in sorted(imatrix["tensors"].keys()): arr=imatrix["tensors"][name].get("in_sum2_raw") if arr is None: no_raw.append(name) continue key=(arr.shape,hash(np.round(arr/atol).astype(np.int64).tobytes())) coarse.setdefault(key,[]).append(name) buckets={} for key,names in coarse.items(): if len(names)==1: buckets[key]=names continue n=len(names) groups=list(range(n)) def find(x): while groups[x]!=x: groups[x]=groups[groups[x]] x=groups[x] return x def union(a,b): ra,rb=find(a),find(b) if ra!=rb:groups[ra]=rb arrs=[imatrix["tensors"][n]["in_sum2_raw"] for n in names] for i in range(n): for j in range(i+1,n): a,b=arrs[i],arrs[j] if a.shape!=b.shape:continue if np.allclose(a,b,rtol=1e-5,atol=1e-6):union(i,j) comps=defaultdict(list) for i,name in enumerate(names):comps[find(i)].append(name) for comp in comps.values():buckets.setdefault(key +(id(comp),),[]).extend(comp) tied_groups=list(buckets.values()) tied_groups.extend([[n] for n in no_raw]) return tied_groups def _imatrix_type(name:str)->str: parts=name.split(".") if len(parts)>=3 and parts[0]=="blk":return parts[2] return name def build_importance_table(imatrix:dict,model:dict)->dict: table={} for tname,info in imatrix["tensors"].items(): ttype=_imatrix_type(tname) table[tname]={ "importance_mean":info["importance_mean"], "importance_sum":info["importance_sum"], "importance_max":info["importance_max"], "importance_min":info["importance_min"], "n_elements":info["n_elements"], "type":ttype, } for tname,info in list(table.items()): if tname.endswith("."): alt=tname.rstrip(".") table[alt]=info return table # ═══════════════════════════════════════════════════════════════════════════ # 5. CLASSIFIER & OPTIMIZER # ═══════════════════════════════════════════════════════════════════════════ BITS_IN_MIB=8*1024*1024.0 TIER_SIZE_MULTIPLIER={ tier:(bpw/BITS_IN_MIB)*GGUF_OVERHEAD_FACTOR for tier,bpw in TIER_BPW.items() } def apply_overhead_factor(factor:float): TIER_SIZE_MULTIPLIER.update({t:(b/BITS_IN_MIB)*factor for t,b in TIER_BPW.items()}) K_QUANTS={"Q3_K","Q4_K","Q5_K","Q6_K"} MOE_PAD_TYPES={"ffn_gate_exps","ffn_up_exps","ffn_down_exps","ffn_down"} MSE_BPW=dict(TIER_BPW) def _percentile_rank(values:Dict[str,float])->Dict[str,float]: if not values:return{} sorted_vals=sorted(values.values()) n=len(sorted_vals) return{name:bisect_left(sorted_vals,raw)/n for name,raw in values.items()} def _normalize_by_class( tensor_importance:Dict[str,float], importance_table:Dict[str,Any], )->Dict[str,float]: class_groups:Dict[str,Dict[str,float]]={} for name,raw_imp in tensor_importance.items(): rep_info=importance_table.get(name,{}) ttype=rep_info.get("type",get_tensor_type(name)) cls=get_tensor_class(ttype) class_groups.setdefault(cls,{})[name]=raw_imp normalized={} for cls,group_vals in class_groups.items(): pct=_percentile_rank(group_vals) normalized.update(pct) for name in tensor_importance: if name not in normalized:normalized[name]=0.5 return normalized @dataclass(order=True,slots=True) class TierMove: neg_utility:float group_id:int=field(compare=False) target_tier:str=field(compare=False) size_delta:float=field(compare=False) is_downgrade:bool=field(compare=False,default=False) from_tier:str=field(compare=False,default="") def _tier_index(tier:str)->int: if tier not in TIER_ORDER:raise ValueError(f"Unknown tier: {tier}") return TIER_ORDER.index(tier) def _tier_at(idx:int)->str: if not(0<=idxfloat: if n_elements<=0:return 0.0 return n_elements*TIER_SIZE_MULTIPLIER.get(tier,0.0) FREE_EMBD_TIER="Q8_0" TIER_BLOCK_ALIGN={ "IQ1_S":256,"IQ2_XXS":256,"IQ2_XS":256,"IQ2_S":256, "IQ3_XXS":256,"Q3_K":256,"IQ3_S":256, "IQ4_XS":256,"Q4_K":256,"Q5_K":256,"Q6_K":256, "IQ4_NL":32,"Q8_0":32, "F16":1, } def _row_len(shape)->int: return int(shape[0]) if shape else 1 def tier_fits(shape,tier:str)->bool: return _row_len(shape) % TIER_BLOCK_ALIGN.get(tier,1)==0 def highest_fitting_tier(shape,tier:str)->str: idx=_tier_index(tier) while idxdict: # Host CPU token embeddings live outside VRAM allocations unless tied to lm_head if not enabled:return {} types={get_tensor_type(n):n for n in all_names} tok,out=types.get("token_embd"),types.get("output") if not(tok and out):return {} tier=FREE_EMBD_TIER if embd_cap and _tier_index(embd_cap)<_tier_index(tier):tier=embd_cap return {tok:tier} def embd_int4_cap(model_path:str)->str | None: globals_=provenance(model_path).get("int4_globals") or [] return INT4_LINEAGE_CAP if any("embed" in g or "token_embd" in g for g in globals_) else None def _pinned_f32_names(model:dict)->set: # Keep 1-D tensors and rows unaligned with ggml block boundaries in F32 out=set() for n,info in model.get("tensors",{}).items(): sh=[int(d) for d in info.get("shape",[])] if len([d for d in sh if d>1])<2 or _row_len(sh)%32!=0: out.add(n) return out def _shape_map(model:dict)->dict: return {n:[int(d) for d in i.get("shape",[])] for n,i in model.get("tensors",{}).items()} def _clamp_to_alignment(assignments:dict,shape_map:dict)->int: fixed=0 for name,tier in list(assignments.items()): sh=shape_map.get(name) if not sh or tier_fits(sh,tier):continue assignments[name]=highest_fitting_tier(sh,tier) fixed+=1 return fixed # Relative efficiency bonus for learned codebook representations FAMILY_EFF_BITS={ "IQ4_XS":0.45,"IQ4_NL":0.30, "IQ3_S":0.40,"IQ3_XXS":0.35, "IQ2_S":0.30,"IQ2_XS":0.25,"IQ2_XXS":0.20, "IQ1_S":0.15, } def _eff_bpw(tier:str)->float: return MSE_BPW[tier]+FAMILY_EFF_BITS.get(tier,0.0) @lru_cache(maxsize=256) def _mse_delta(cur_tier:str,next_tier:str)->float: return(2**(-2*_eff_bpw(cur_tier))) -(2**(-2*_eff_bpw(next_tier))) def _group_size(group_registry:Dict[int,Tuple[List[str],int,int]],group_id:int,tier:str)->float: g_names,g_elements,g_elements_padded=group_registry[group_id] elements=g_elements_padded if tier in K_QUANTS else g_elements return _size_mib(tier,elements) def _push_upgrade(group_id:int,group_registry:Dict[int,Tuple[List[str],int,int]],assignments:Dict[str,str],tensor_importance:Dict[str,float],upgrade_queue:List[TierMove],importance_table:Dict[str,Any],cap_table:dict | None=None): g_names,_,_=group_registry[group_id] rep_name=g_names[0] cur_tier=assignments[rep_name] cur_idx=_tier_index(cur_tier) rep_info=importance_table.get(rep_name,{}) ttype=rep_info["type"] if "type" in rep_info else get_tensor_type(rep_name) cls=get_tensor_class(ttype) max_tier=CLASS_MAX_TIER.get(cls,"Q8_0") if cap_table and cls in cap_table and _tier_index(cap_table[cls])<_tier_index(max_tier): max_tier=cap_table[cls] if cur_idx>=_tier_index(max_tier) or cur_idx>=len(TIER_ORDER)-1: return next_idx=cur_idx+1 max_idx=min(_tier_index(max_tier),len(TIER_ORDER)-1) next_tier,cost_delta,quality_delta=None,0.0,0.0 while next_idx<=max_idx: next_tier=_tier_at(next_idx) cost_delta=_group_size(group_registry,group_id,next_tier)-_group_size(group_registry,group_id,cur_tier) quality_delta=_mse_delta(cur_tier,next_tier) if quality_delta>0:break next_idx+=1 if next_tier is None or next_idx>max_idx or quality_delta<=0 or cost_delta<0: return if cost_delta==0: utility_per_mb=float('inf') else: total_g_imp=sum(tensor_importance.get(n,0) for n in g_names) utility_per_mb=(total_g_imp*quality_delta)/cost_delta heapq.heappush(upgrade_queue,TierMove(-utility_per_mb,group_id,next_tier,cost_delta,from_tier=cur_tier)) def _push_downgrade(group_id:int,group_registry:Dict[int,Tuple[List[str],int,int]],assignments:Dict[str,str],tensor_importance:Dict[str,float],downgrade_queue:List[TierMove],importance_table:Dict[str,Any],floor_table:Dict[str,str]): g_names,_,_=group_registry[group_id] rep_name=g_names[0] cur_tier=assignments[rep_name] cur_idx=_tier_index(cur_tier) rep_info=importance_table.get(rep_name,{}) ttype=rep_info["type"] if "type" in rep_info else get_tensor_type(rep_name) cls=get_tensor_class(ttype) floor=floor_table.get(cls,"Q4_K") if "importance_mean" not in rep_info and _tier_index(floor)<_tier_index(NO_IMATRIX_MIN_TIER): floor=NO_IMATRIX_MIN_TIER if cur_idx<=_tier_index(floor) or cur_idx<=0: return next_tier=_tier_at(cur_idx-1) saved=_group_size(group_registry,group_id,cur_tier)-_group_size(group_registry,group_id,next_tier) quality_loss=_mse_delta(next_tier,cur_tier) if saved<=0 or quality_loss<=0: return total_g_imp=sum(tensor_importance.get(n,0) for n in g_names) loss_per_mb=(total_g_imp*quality_loss)/saved heapq.heappush(downgrade_queue,TierMove(loss_per_mb,group_id,next_tier,-saved,True,from_tier=cur_tier)) def converge_to_target(importance_table:dict,tied_groups:list,model:dict,target_size_mib:float,allow_q3:bool=False,tolerance_pct:float=2.0,max_iterations:int=5,profile:str="quality",int4_lineage:bool=False,free_embd:bool=True,embd_cap:str | None=None)->Tuple[dict,dict]: effective_target=target_size_mib ne_map={k:v["n_elements"] for k,v in model.get("tensors",{}).items()} if not ne_map: ne_map={k:v["n_elements"] for k,v in importance_table.items()} f32_names=_pinned_f32_names(model) best=None best_diff=float("inf") seen_targets=set() for iteration in range(max(max_iterations,1)): assignments,padded_ne_map=optimal_classify(importance_table,tied_groups,model,target_size_mib=effective_target,allow_q3=allow_q3,profile=profile,int4_lineage=int4_lineage,free_embd=free_embd,embd_cap=embd_cap) current_size=compute_stats(assignments,ne_map,padded_ne_map,f32_names)["total_mib"] diff_pct=abs(current_size-target_size_mib)/target_size_mib*100 if diff_pct_tier_index(INT4_MTP_CAP): mtp_tier=INT4_MTP_CAP assignments={n:mtp_tier for n in mtp_names} floors=TIER_FLOORS.get(profile,TIER_FLOORS["quality"]) for name in non_mtp_names: rep_info=importance_table.get(name,{}) has_im="importance_mean" in rep_info ttype=rep_info.get("type",get_tensor_type(name)) cls=get_tensor_class(ttype) tier=floors.get(cls,floors.get("unknown",DEFAULT_FLOOR)) hard_floor=CLASS_HARD_FLOORS.get(cls) if hard_floor and _tier_index(tier)<_tier_index(hard_floor) and not(allow_q3 and cls in CAN_Q3): tier=hard_floor if is_qat and cls not in("norms","ssm_params","gdn_state"): cap="Q4_K" if cls=="attn_proj" else "IQ4_XS" if _tier_index(cap)<_tier_index(tier):tier=cap if int4_lineage and cls in INT4_CAP_CLASSES and _tier_index(tier)>_tier_index(INT4_LINEAGE_CAP): tier=INT4_LINEAGE_CAP if allow_q3 and cls in CAN_Q3 and has_im: tier=ALLOW_LOWER_FLOOR if not has_im and _tier_index(tier)<_tier_index(NO_IMATRIX_MIN_TIER): tier=NO_IMATRIX_MIN_TIER assignments[name]=tier return assignments def compute_initial_assignments(non_mtp_names:Set[str],mtp_names:Set[str],importance_table:Dict,allow_q3:bool,is_qat:bool=False,profile:str="quality",int4_lineage:bool=False)->Dict[str,str]: return compute_mtp_assignments_check(non_mtp_names,mtp_names,importance_table,allow_q3,is_qat,profile,int4_lineage) def _tie_family(name:str,importance_table:Dict)->str: ttype=importance_table.get(name,{}).get("type",get_tensor_type(name)) if ttype in("ffn_gate","ffn_up"):return "ffn_gate_up" return ttype def build_groups(tied_groups:List[List[str]],non_mtp_names:Set[str],ne_map:Dict[str,int],padded_ne_map:Dict[str,int],importance_table:Dict)->Dict[int,Tuple[List[str],int,int]]: group_registry={} assigned_tensors=set() next_group_id=0 for tied_group in tied_groups: families={} for name in tied_group: if name in non_mtp_names: families.setdefault(_tie_family(name,importance_table),[]).append(name) for names in families.values(): group_registry[next_group_id]=(names,sum(ne_map.get(name,0) for name in names),sum(padded_ne_map.get(name,0) for name in names)) assigned_tensors.update(names) next_group_id+=1 for name in sorted(non_mtp_names-assigned_tensors): group_registry[next_group_id]=([name],ne_map.get(name,0),padded_ne_map.get(name,0)) next_group_id+=1 return group_registry def optimal_classify(importance_table:dict,tied_groups:list,model:dict,target_size_mib:float,allow_q3:bool=False,profile:str="quality",int4_lineage:bool=False,free_embd:bool=True,embd_cap:str | None=None)->Tuple[dict,dict]: if target_size_mib<=0:raise ValueError("target_size_mib must be positive") cap_table=None if int4_lineage: cap_table={cls:INT4_LINEAGE_CAP for cls in INT4_CAP_CLASSES} cap_table["gdn_state"]=INT4_GDN_CAP cap_table["gate"]=INT4_GATE_CAP features=model.get("features",{}) has_mtp=features.get("has_mtp",False) n_layers=features.get("n_layers",31) is_qat=features.get("is_qat",False) model_tensors=model.get("tensors",{}) ne_map={k:v["n_elements"] for k,v in model_tensors.items()} if not model_tensors: for tname,info in importance_table.items(): ne_map[tname]=info["n_elements"] moe_d_ff=features.get("moe_intermediate_size",0) padded_ne_map=dict(ne_map) if moe_d_ff>0 and moe_d_ff % 256!=0: aligned_d_ff=((moe_d_ff+255) // 256)*256 for name,n_el in ne_map.items(): ttype=importance_table.get(name,{}).get("type",get_tensor_type(name)) if ttype in MOE_PAD_TYPES: padded_ne_map[name]=(n_el // moe_d_ff)*aligned_d_ff all_names=set(ne_map.keys()) mtp_layer_set=set(features.get("mtp_layers",[])) f32_names=_pinned_f32_names(model) mtp_names=({n for n in all_names if is_mtp_tensor(n,n_layers,mtp_layer_set)}-f32_names) if has_mtp else set() free_map=_free_embd_assignments(model,all_names,free_embd,embd_cap) non_mtp_names=all_names-mtp_names-f32_names-set(free_map) tensor_importance={} for name in non_mtp_names: info=importance_table.get(name,{}) tensor_importance[name]=info.get("importance_mean",0.0) tensor_importance=_normalize_by_class(tensor_importance,importance_table) for name in tensor_importance: tensor_importance[name]*=max(ne_map.get(name,1),1) assignments=compute_mtp_assignments_check(non_mtp_names,mtp_names,importance_table,allow_q3,is_qat,profile,int4_lineage) assignments.update(free_map) group_registry=build_groups(tied_groups,non_mtp_names,ne_map,padded_ne_map,importance_table) profile_floors=TIER_FLOORS.get(profile,{}) floor_table={} for cls in set(list(CLASS_HARD_FLOORS)+list(profile_floors)): soft=profile_floors.get(cls,CLASS_HARD_FLOORS.get(cls,"Q4_K")) hard=CLASS_HARD_FLOORS.get(cls,"IQ1_S") floor_table[cls]=soft if _tier_index(soft)>=_tier_index(hard) else hard if profile in("mini","compact","quality"): if _tier_index(floor_table.get("embd","Q4_K"))>_tier_index("IQ4_XS"): floor_table["embd"]="IQ4_XS" mtp_cost=sum(_size_mib(assignments[n],ne_map.get(n,0)) for n in mtp_names) pinned_cost=sum(ne_map.get(n,0)*32.0/BITS_IN_MIB for n in f32_names) effective_target=target_size_mib-mtp_cost-pinned_cost current_size=sum( _size_mib(assignments[n],padded_ne_map.get(n,ne_map.get(n,0)) if assignments[n] in K_QUANTS else ne_map.get(n,0)) for n in non_mtp_names ) if current_size>effective_target: downgrade_queue=[] for g_id in group_registry: _push_downgrade(g_id,group_registry,assignments,tensor_importance,downgrade_queue,importance_table,floor_table) while downgrade_queue and current_size>effective_target: item=heapq.heappop(downgrade_queue) if assignments[group_registry[item.group_id][0][0]]!=item.from_tier: continue for n in group_registry[item.group_id][0]: assignments[n]=item.target_tier current_size+=item.size_delta _push_downgrade(item.group_id,group_registry,assignments,tensor_importance,downgrade_queue,importance_table,floor_table) if current_size>effective_target: warnings.warn( f"Cannot reach target: floors hold size at {current_size:.1f} MiB vs target {effective_target:.1f} MiB. " "Use --allow-q3-or-lower or increase --size.",RuntimeWarning) has_signal=any("importance_mean" in importance_table.get(n,{}) for n in non_mtp_names) upgrade_queue=[] if has_signal: for g_id in group_registry: _push_upgrade(g_id,group_registry,assignments,tensor_importance,upgrade_queue,importance_table,cap_table) while upgrade_queue: item=heapq.heappop(upgrade_queue) if assignments[group_registry[item.group_id][0][0]]!=item.from_tier: continue if item.size_delta>0 and current_size+item.size_delta>effective_target: # Skip the oversized move and keep draining: current_size only grows, so this group never fits later continue for n in group_registry[item.group_id][0]: assignments[n]=item.target_tier current_size+=item.size_delta _push_upgrade(item.group_id,group_registry,assignments,tensor_importance,upgrade_queue,importance_table,cap_table) fixed=_clamp_to_alignment(assignments,_shape_map(model)) if fixed:print(f" ℹ {fixed} tensor(s) promoted to closest aligned tier (ggml block)") return assignments,padded_ne_map def compute_stats(assignments:dict,ne_map:dict=None,padded_ne_map:dict=None,f32_names:set=None)->dict: stats={"by_tier_count":{},"by_tier_mib":{},"total_mib":0.0,"tensor_count":0} for name,tier in assignments.items(): if not isinstance(tier,str):continue stats["tensor_count"]+=1 stats["by_tier_count"][tier]=stats["by_tier_count"].get(tier,0)+1 if ne_map: if padded_ne_map and tier in K_QUANTS: elements=padded_ne_map.get(name,ne_map.get(name,0)) else: elements=ne_map.get(name,0) size=_size_mib(tier,elements) stats["by_tier_mib"][tier]=stats["by_tier_mib"].get(tier,0.0)+size stats["total_mib"]+=size if ne_map and f32_names: for name in f32_names: size=ne_map.get(name,0)*32.0/BITS_IN_MIB stats["tensor_count"]+=1 stats["by_tier_count"]["F32"]=stats["by_tier_count"].get("F32",0)+1 stats["by_tier_mib"]["F32"]=stats["by_tier_mib"].get("F32",0.0)+size stats["total_mib"]+=size return stats # ═══════════════════════════════════════════════════════════════════════════ # 6. CONFIG GENERATOR # ═══════════════════════════════════════════════════════════════════════════ def get_regex_priority(regex:str)->int: score=0 if "nextn" in regex:score+=200 if re.search(r"(blk|BLK)\\.3[0-2]\\.",regex):score+=100 if re.search(r"(blk|BLK)\\.0\\.",regex):score+=90 if re.search(r"(blk|BLK)\\.31\\.",regex):score+=80 if r"(blk|BLK)\.(" in regex:score+=50 elif r"(blk|BLK)\.\d" in regex:score+=30 if regex.startswith(".*"):score-=50 if regex.endswith(r"\.weight"):score+=10 return score def _is_contiguous(lst,low,high): if not lst:return False return len(lst)==(high-low+1) def _group_ranges(lst): if not lst:return start=lst[0] end=lst[0] for i in range(1,len(lst)): if lst[i]==end+1: end=lst[i] else: yield(start,end) start=end=lst[i] yield(start,end) def _range_to_regex(start:int,end:int)->str: if start==end:return str(start) if end<=9:return f"[{start}-{end}]" alt="|".join(str(i) for i in range(start,end+1)) return f"(?:{alt})" def generate_flags(assignments:dict,model:dict,base_type:str,target_size_mib:float=None)->dict: is_qat=model.get("features",{}).get("is_qat",False) output_type="Q5_K" token_embd_type="Q4_K" if is_qat else "Q5_K" for tname,tier in assignments.items(): ttype=get_tensor_type(tname) if ttype=="output":output_type=tier elif ttype=="token_embd":token_embd_type=tier max_layer=model.get("features",{}).get("n_layers",31) rules=[] type_tier_layers={} max_layer_seen=-1 for tname,tier in assignments.items(): parts=tname.split(".") tower="" if len(parts)>=4 and parts[0] in("v","a") and parts[1] in("blk","BLK"): tower=parts[0]+"." parts=parts[1:] if len(parts)>=3 and parts[0] in("blk","BLK"): try:layer=int(parts[1]) except ValueError:continue ttype=".".join(parts[2:]) key=(tower,ttype,tier) if key not in type_tier_layers:type_tier_layers[key]=[] type_tier_layers[key].append(layer) max_layer_seen=max(max_layer_seen,layer) max_layer=max_layer_seen if max_layer_seen>=0 else max_layer for(tower,ttype,tier),layers in sorted(type_tier_layers.items(),key=lambda x:-QUANT_RANK.get(x[0][2],0)): layers=sorted(set(layers)) if len(layers)>=8 and _is_contiguous(layers,0,max_layer): pattern=f"{tower}(blk|BLK)\\.\\d+\\.{re.escape(ttype)}={tier}" else: parts=[] for start,end in _group_ranges(layers): if start==end:parts.append(str(start)) else:parts.append(_range_to_regex(start,end)) desc="|".join(parts) pattern=f"{tower}(blk|BLK)\\.({desc})\\.{re.escape(ttype)}={tier}" prio=get_regex_priority(pattern) +(10 if tier=="Q8_0" else 5 if tier=="Q6_K" else 0) +(5 if len(layers)==1 else 0) +(3 if "ffn_down" in ttype else 0) rules.append((pattern,prio)) prefix=model.get("features",{}).get("prefix","blk") for tname,tier in sorted(assignments.items(),key=lambda kv:-QUANT_RANK.get(kv[1],0)): parts=tname.split(".") if len(parts)>=2 and parts[0].lower()==prefix.lower():continue if len(parts)>=4 and parts[0] in("v","a") and parts[1] in("blk","BLK"):continue ttype=get_tensor_type(tname) if ttype==tname:ttype=tname.removesuffix(".weight").removesuffix(".bias") if ttype in("token_embd","output"):continue ttype=re.sub(r"\.(?:weight|bias)\.\d+$","",ttype) body=f".*{re.escape(ttype)}\\.weight" pattern=f"{body}={tier}" prio=get_regex_priority(pattern) +(5 if tier=="Q8_0" else 0) if not any(p.rsplit("=",1)[0]==body for p,_ in rules):rules.append((pattern,prio)) rules.sort(key=lambda x:-x[1]) has_output=any(get_tensor_type(t)=="output" for t in assignments) has_embd=any(get_tensor_type(t)=="token_embd" for t in assignments) flags={ "imatrix":None, "output_tensor_type":output_type if has_output else None, "token_embedding_type":token_embd_type if has_embd else None, "tensor_type_rules":[r[0] for r in rules], "base_type":base_type, "target_size_mib":target_size_mib, } return flags def format_flags(flags:dict)->str: lines=[] if flags.get("output_tensor_type"):lines.append(" --output-tensor-type "+flags["output_tensor_type"]) if flags.get("token_embedding_type"):lines.append(" --token-embedding-type "+flags["token_embedding_type"]) for pattern in flags["tensor_type_rules"]:lines.append(f' --tensor-type "{pattern}"') return "\n".join(lines) # ═══════════════════════════════════════════════════════════════════════════ # 7. QUANTIZER EXECUTOR # ═══════════════════════════════════════════════════════════════════════════ def _find_binary(name:str)->str: exe=name +(".exe" if sys.platform=="win32" else "") env_var=f"LLAMA_{name.upper().replace('-','_')}_PATH" env_path=os.environ.get(env_var) if env_path: if os.path.isfile(env_path):return env_path cand=os.path.join(env_path,exe) if os.path.isfile(cand):return cand script_dir=os.path.dirname(os.path.abspath(__file__)) roots=[] env_root=os.environ.get("LLAMA_CPP_DIR") if env_root:roots.append(env_root) for base in(script_dir,os.getcwd()): roots.append(base) roots.append(os.path.normpath(os.path.join(base,"llama-cpp"))) roots.append(os.path.normpath(os.path.join(base,"..","llama-cpp"))) roots.append(os.path.normpath(os.path.join(base,"..","..","llama-cpp"))) subdirs=("","bin",os.path.join("build","bin"),os.path.join("build","bin","Release"),os.path.join("build","Release"),"build") seen=set() for root in roots: for sub in subdirs: p=os.path.normpath(os.path.join(root,sub,exe)) if p in seen:continue seen.add(p) if os.path.isfile(p):return p from shutil import which found=which(name) or which(exe) if found:return found return "" def _build_cmd(flags:dict,model_in:str,model_out:str,dry_run:bool=False)->list: cmd=[_find_binary("llama-quantize")] if dry_run:cmd.append("--dry-run") if flags.get("imatrix"): imatrix=flags["imatrix"] if isinstance(imatrix,list):imatrix=imatrix[0] cmd.extend(["--imatrix",imatrix]) if flags.get("output_tensor_type"):cmd.extend(["--output-tensor-type",flags["output_tensor_type"]]) if flags.get("token_embedding_type"):cmd.extend(["--token-embedding-type",flags["token_embedding_type"]]) for pattern in flags["tensor_type_rules"]:cmd.extend(["--tensor-type",pattern]) cmd.append(model_in) cmd.append(model_out) cmd.append(flags["base_type"]) return cmd def _missing_binary_msg(name:str): exe=name +(".exe" if sys.platform=="win32" else "") print(f" ⚠ Binary '{exe}' not found.") print(f" Set LLAMA_CPP_DIR to your llama.cpp folder,") print(f" or set LLAMA_{name.upper().replace('-','_')}_PATH to the exact path,") print(f" or place '{exe}' next to this script.") def run_dry_run(flags:dict,model_in:str)->float | None: binary=_find_binary("llama-quantize") if not binary: _missing_binary_msg("llama-quantize") return None cmd=_build_cmd(flags,model_in,os.devnull,dry_run=True) try:result=subprocess.run(cmd,capture_output=True,text=True,timeout=600) except FileNotFoundError: _missing_binary_msg("llama-quantize") return None except subprocess.TimeoutExpired: print(" ⚠ Dry run timed out after 600 s") return None output=(result.stdout or "") +(result.stderr or "") size=parse_quant_size(output) if size is not None:return size m=re.search(r"unsupported model architecture:'([^']+)'",output) if m: print(f" ⚠ llama-quantize rejected architecture '{m.group(1)}' — quantization aborted.") return None print("STDERR:",(result.stderr or "")[:2000]) return None def run_quantization(flags:dict,model_in:str,model_out:str)->bool: binary=_find_binary("llama-quantize") if not binary: _missing_binary_msg("llama-quantize") return False cmd=_build_cmd(flags,model_in,model_out) print("Running:"," ".join(cmd[:6])+" ...") try:result=subprocess.run(cmd) except FileNotFoundError: _missing_binary_msg("llama-quantize") return False success=result.returncode==0 if success and os.path.isfile(model_out): size_mib=os.path.getsize(model_out)/1024/1024 print(f"Done: {model_out} ({size_mib:.0f} MiB)") elif os.path.isfile(model_out): # llama-quantize leaves a truncated file behind: remove it so callers see a clean failure try: os.remove(model_out) print(f"Removed partial output: {model_out}") except OSError as e:print(f" [WARN] Partial output kept ({e})") return success # ═══════════════════════════════════════════════════════════════════════════ # 8. MAIN SHQ ENGINE (SINGLE TARGET RUNNER) # ═══════════════════════════════════════════════════════════════════════════ def _get_base_type(model:dict,profile:str="quality")->str: if model.get("features",{}).get("is_qat",False):return "IQ4_XS" return{ "nano":"IQ3_XXS","mini":"IQ4_XS","compact":"IQ4_XS","quality":"Q5_K_M", "fidelity":"Q6_K", }.get(profile,"Q5_K_M") def run_main_cli(args_list=None)->int: parser=argparse.ArgumentParser(description="ASHQ1: AutoRound-infused imatrix 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 imatrices: max or mean") parser.add_argument("--size",type=float,default=6800,help="Target file size in MiB") parser.add_argument("--profile",choices=["nano","mini","compact","quality","fidelity"],default="quality",help="Quantization profile") parser.add_argument("--lineage",choices=["auto","autoround","plain"],default="auto",help="Force output name lineage tag") 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 starting tiers down to IQ2_XXS") parser.add_argument("--no-free-embd",action="store_true",help="Keep token_embd inside the size budget (tied lm_head or strict file-size target)") parser.add_argument("--aggro",type=float,default=None,help="[deprecated] Use --size instead") parser.add_argument("--target-size",action="store_true",help="Iteratively converge to --size") parser.add_argument("--show-floors",action="store_true",help="Print class hard floors and exit") try:args=parser.parse_args(args_list) except SystemExit as e:return e.code if isinstance(e.code,int) else 1 if args.show_floors: _show_floors() return 0 if not args.model: parser.print_usage() print("main: error: --model is required") return 1 if not args.imatrix: print("Note: no --imatrix provided, using profile defaults only.") target_mib=args.size print("=== ASHQ1 (AutoRound-infused) ===") print(f"Model: {args.model}") if not args.imatrix: print("Imatrix: none (floor-driven profile)") elif 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)") int4_lineage=int4_lineage_of(args.model) if int4_lineage: print(f" Lineage: int4 AutoRound source — weight upgrades capped at {INT4_LINEAGE_CAP}, MTP at {INT4_MTP_CAP}") _ov=load_overhead_factor(args.model,args.profile) if abs(_ov-1.0)>1e-4: apply_overhead_factor(_ov) print(f" Calibrated overhead (previous run): ×{_ov:.4f}") 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) if not args.no_free_embd: _free=_free_embd_assignments(model,set(model.get("tensors",{})),True) _mib=sum(_size_mib(t,model["tensors"].get(n,{}).get("n_elements",0)) for n,t in _free.items()) if _mib>max(target_mib*0.15,128): print(f" ⚠ Embedding {_mib:.0f} MiB > 15% of target — disabling free-embd to preserve accurate tier ratios") args.no_free_embd=True no_free=args.no_free_embd 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','?')}") imatrix=combine_imatrix(imatrix_list,method=args.imatrix_method) if imatrix_list:print(f" Combined: {imatrix['n_tensors']} tensors") else:print(" Skipped: profile floors drive every assignment") print("\n[3/4] Detecting tied groups...") imp_table=build_importance_table(imatrix,model) tied_groups=detect_tied_groups(imatrix) print(f" Found {len(tied_groups)} tied groups:") for g in tied_groups: if len(g)<2:continue fam={} for n in g:fam.setdefault(_tie_family(n,imp_table),[]).append(n) for names in fam.values(): if len(names)<2:continue print(f" TIED({len(names)}): {names[0].replace('.weight','')} = {names[1].replace('.weight','')}") print("\n[4/4] Classifying tensors (greedy imatrix-driven)...") embd_cap=embd_int4_cap(args.model) if int4_lineage else None if embd_cap: print(f" Embeddings covered by int4 grid → host bump capped at {embd_cap}") if args.target_size: assignments,padded_ne_map=converge_to_target( imp_table,tied_groups,model, target_size_mib=target_mib, allow_q3=args.allow_q3_or_lower, profile=args.profile, int4_lineage=int4_lineage, free_embd=not no_free, embd_cap=embd_cap, ) else: assignments,padded_ne_map=optimal_classify( imp_table,tied_groups,model, target_size_mib=target_mib, allow_q3=args.allow_q3_or_lower, profile=args.profile, int4_lineage=int4_lineage, free_embd=not no_free, embd_cap=embd_cap, ) ne_map={k:v["n_elements"] for k,v in model.get("tensors",{}).items()} if not ne_map:ne_map={k:v["n_elements"] for k,v in imp_table.items()} free_mib=sum(_size_mib(t,ne_map.get(n,0)) for n,t in _free_embd_assignments(model,set(ne_map),not no_free,embd_cap).items()) if free_mib:print(f" Host-side embedding bump: {free_mib:.0f} MiB outside VRAM budget") _show_tier_summary(assignments,imp_table,ne_map,padded_ne_map,_pinned_f32_names(model)) base_type=_get_base_type(model,args.profile) 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 0 print("\n--- Dry Run ---") estimated=compute_stats(assignments,ne_map,padded_ne_map,_pinned_f32_names(model))["total_mib"] dry_size=run_dry_run(flags,args.model) if dry_size and estimated>0: ratio=dry_size/estimated save_overhead_factor(args.model,_ov*ratio,args.profile) if abs(ratio-1.0)>0.02: print(f" Calibration: actual/estimated overhead ×{ratio:.4f} — saved for subsequent runs") _show_size_result(dry_size,target_mib+free_mib) if not args.run: print("\nDry run only. Use --run to execute quantization.") return 0 if not args.output: base=clean_name(args.model) pct=resolve_pct(args.profile,args.model) args.output=f"{base}-{lineage_tag(args.model,args.lineage)}-{tier_label(args.profile)}-{pct}pc.gguf" print(f"\n--- Running quantization: {args.output} ---") success=run_quantization(flags,args.model,args.output) if success: print("Done!") return 0 else: print("Failed!") return 1 def _show_tier_summary(assignments,imp_table,ne_map,padded_ne_map=None,f32_names=None): stats=compute_stats(assignments,ne_map,padded_ne_map,f32_names) 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") unknown=[n for n in assignments if get_tensor_class(get_tensor_type(n))=="unknown"] if unknown: print(f"\n Unclassified tensors (floor-driven, {len(unknown)}) — extend TENSOR_CLASS if sensitive:") for n in unknown[:8]:print(f" {n}") if len(unknown)>8:print(f" … +{len(unknown)-8} more") 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}") # Metadata overhead (header + vocab) not tracked by tensor arrays META_OVERHEAD_MIB=12.0 def _show_size_result(dry_size,target_mib): if dry_size: on_disk=dry_size+META_OVERHEAD_MIB print(f" Estimated size: {dry_size:.0f} MiB ({dry_size/1024:.2f} GB)") print(f" Estimated on disk: {on_disk:.0f} MiB (+{META_OVERHEAD_MIB:.0f} MiB metadata)") diff=on_disk-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 IQ2_XXS start for: {','.join(sorted(CAN_Q3))}") # ═══════════════════════════════════════════════════════════════════════════ # 9. TIERS RUNNER & CLI ENTRYPOINT # ═══════════════════════════════════════════════════════════════════════════ TIER_RATIOS={ "nano":0.24, "mini":0.27, "compact":0.33, "quality":0.39, "fidelity":0.48, } INT4_TIER_RATIOS={"quality":0.36} def resolve_tier_ratio(tier,lineage): # int4 upgrades stall against INT4_LINEAGE_CAP; high tiers get a truthful target if lineage=="int4" and tier in INT4_TIER_RATIOS: return INT4_TIER_RATIOS[tier] return TIER_RATIOS[tier] ALL_RATIOS=dict(TIER_RATIOS) TIER_DISPLAY={"nano":"Nano"} def tier_label(tier:str)->str: return TIER_DISPLAY.get(tier.lower(),tier.capitalize()) def resolve_ratio(tier:str,int4:bool)->float: # int4 upgrades stall against INT4_LINEAGE_CAP: high tiers get a truthful target t=tier.lower() if int4 and t in INT4_TIER_RATIOS:return INT4_TIER_RATIOS[t] return ALL_RATIOS.get(t,0.0) def resolve_pct(tier:str,model_path:str)->int: return round(resolve_ratio(tier,int4_lineage_of(model_path))*100) def file_size_mib(path:str)->float: return os.path.getsize(path)/1024/1024 def clean_name(path:str)->str: base=os.path.splitext(os.path.basename(path))[0] for suffix in("-no-mtp-BF16","-no-mtp-F16","-BF16-no-mtp","-F16-no-mtp","-BF16","-F16","-bf16","-f16"): if base.endswith(suffix): base=base[:-len(suffix)] break return base PROVENANCE_MAP="quant-provenance.json" def provenance(model_path:str)->dict: sidecar=os.path.splitext(model_path)[0]+".provenance.json" if os.path.isfile(sidecar): try: with open(sidecar,"r",encoding="utf-8") as f:return json.load(f) except Exception:pass fmap=os.path.join(os.path.dirname(os.path.abspath(model_path)),PROVENANCE_MAP) if os.path.isfile(fmap): try: with open(fmap,"r",encoding="utf-8") as f:return json.load(f).get(os.path.basename(model_path),{}) except Exception:pass return {} def lineage_tag(model_path:str,override:str="auto")->str: if override=="autoround":return "AutoRound-ASHQ1" if override=="plain":return "ASHQ1" return "AutoRound-ASHQ1" if provenance(model_path).get("autoround") else "ASHQ1" def int4_lineage_of(model_path:str)->bool: prov=provenance(model_path) try:bits=int(prov.get("bits",16) or 16) except (TypeError,ValueError):bits=16 return bool(prov.get("autoround")) and bits<=4 def gen_imatrix(model:str,data:str,output:str,chunks:int)->bool: bin_path=_find_binary("llama-imatrix") if not bin_path: _missing_binary_msg("llama-imatrix") return False cmd=[bin_path,"-m",model,"-f",data,"-o",output,"--chunks",str(chunks)] print("Running:"," ".join(cmd[:4])+" ...") try:result=subprocess.run(cmd) except FileNotFoundError: _missing_binary_msg("llama-imatrix") return False return result.returncode==0 def print_tier_table(bf16_mib:float,model_name:str,int4:bool=False): print(f"\n Model: {model_name}") print(f" BF16 source: {bf16_mib:.0f} MiB ({bf16_mib/1024:.2f} GB)\n") print(f" {'Tier':<12}{'Ratio':>6}{'Target(MiB)':>14}{'Target(GB)':>12}") print(f" {'-'*12}{'-'*6}{'-'*14}{'-'*12}") for tier in ALL_RATIOS: ratio=resolve_ratio(tier,int4) target=bf16_mib*ratio print(f" {tier_label(tier):<12}{ratio*100:>5.0f}%{target:>13.0f}{target/1024:>11.2f}") print() def run_tier(model:str,imatrix_paths:list,tier:str,target_mib:float,output_dir:str,run:bool,extra_args:list,lineage:str="auto")->bool: name=clean_name(model) pct=resolve_pct(tier,model) output=os.path.join(output_dir,f"{name}-{lineage_tag(model,lineage)}-{tier_label(tier)}-{pct}pc.gguf") cmd_args=[ "--model",model, "--size",f"{target_mib:.0f}", "--output",output, "--profile",tier, ] for im in imatrix_paths:cmd_args.extend(["--imatrix",im]) cmd_args.extend(extra_args) if run:cmd_args.append("--run") print(f"\n{'='*60}") print(f" [{tier.upper()}] Target: {target_mib:.0f} MiB ({target_mib/1024:.2f} GB)") print(f" Output: {output}") print(f"{'='*60}\n") return run_main_cli(cmd_args)==0 def tiers_main(args_list=None): parser=argparse.ArgumentParser( description="ASHQ1 Tier Runner — standardized quantization tiers", formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument("--model",help="BF16/F16 GGUF model path") parser.add_argument("--imatrix",action="append",default=[],help="Imatrix file path") parser.add_argument("--tier",choices=list(TIER_RATIOS.keys())+["all"],default="all",help="Which tier to run") parser.add_argument("--lineage",choices=["auto","autoround","plain"],default="auto",help="Force output name lineage tag") parser.add_argument("--output-dir",default=None,help="Output directory") parser.add_argument("--run",action="store_true",help="Execute quantization") parser.add_argument("--show-sizes",action="store_true",help="Only print size table and exit") parser.add_argument("--allow-q3-or-lower",action="store_true",help="Allow Q3 or lower") parser.add_argument("--imatrix-method",choices=["max","mean"],default="max",help="Imatrix combination method") parser.add_argument("--verbose",action="store_true",help="Detailed output") parser.add_argument("--gen-imatrix",action="store_true",help="Generate imatrix before quantization") parser.add_argument("--data",help="Calibration data file for imatrix") parser.add_argument("--chunks",type=int,default=100,help="Number of chunks for imatrix") args,remaining=parser.parse_known_args(args_list) if not args.model: parser.print_usage() print("ashq1: error: --model is required") sys.exit(1) if not os.path.isfile(args.model): print(f"ERROR: Model not found: {args.model}") sys.exit(1) bf16_mib=file_size_mib(args.model) model_name=clean_name(args.model) output_dir=args.output_dir or os.path.dirname(os.path.abspath(args.model)) os.makedirs(output_dir,exist_ok=True) print("=== ASHQ1 Tier Runner ===") print_tier_table(bf16_mib,model_name,int4_lineage_of(args.model)) if args.show_sizes:return imatrix_paths=list(args.imatrix) if args.gen_imatrix: if not args.data: print("ERROR: --data required for --gen-imatrix") sys.exit(1) imatrix_path=os.path.join(output_dir,f"{model_name}.imatrix.dat") print(f"\n--- Generating imatrix: {imatrix_path} ---") if not gen_imatrix(args.model,args.data,imatrix_path,args.chunks): print("ERROR: Imatrix generation failed!") sys.exit(1) imatrix_paths.insert(0,imatrix_path) if not imatrix_paths: print("ERROR: --imatrix required (or use --gen-imatrix with --data)") sys.exit(1) for im in imatrix_paths: if not os.path.isfile(im): print(f"ERROR: Imatrix not found: {im}") sys.exit(1) tiers=["mini","compact","quality","fidelity","nano"] if args.tier=="all" else[args.tier] if args.tier=="all" and int4_lineage_of(args.model): if os.environ.get("ASHQ1_INCLUDE_FIDELITY")!="1" and "fidelity" in tiers: tiers.remove("fidelity") if os.environ.get("ASHQ1_INCLUDE_QUALITY")!="1" and "quality" in tiers: tiers.remove("quality") print(f" ℹ AutoRound int4 lineage detected — Fidelity & Quality skipped (information ceiling reached at Compact). Overrides: ASHQ1_INCLUDE_QUALITY=1, ASHQ1_INCLUDE_FIDELITY=1.") extra_args=[] if args.allow_q3_or_lower:extra_args.append("--allow-q3-or-lower") extra_args.extend(["--imatrix-method",args.imatrix_method,"--lineage",args.lineage]) if args.verbose:extra_args.append("--verbose") int4=int4_lineage_of(args.model) results={} for tier in tiers: target=bf16_mib*resolve_ratio(tier,int4) success=run_tier(args.model,imatrix_paths,tier,target,output_dir,args.run,extra_args,args.lineage) results[tier]=success print(f"\n{'='*60}") print(" SUMMARY") print(f"{'='*60}") mode="quantized" if args.run else "dry-run" for tier,success in results.items(): target=bf16_mib*resolve_ratio(tier,int4) status="✓ Done" if success else "✗ Failed" print(f" {tier_label(tier):<12}{target:>7.0f} MiB {status} ({mode})") print() if results and not all(results.values()):sys.exit(2) if __name__=="__main__": raw_args=sys.argv[1:] if any(arg in raw_args for arg in("--profile","--show-config","--show-floors","--size")) and "--tier" not in raw_args: sys.exit(run_main_cli(raw_args)) else: tiers_main(raw_args)