| |
| """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 |
|
|
| |
| |
| |
|
|
| 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()} |
|
|
| |
| 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", |
| ] |
|
|
| |
| 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)} |
|
|
| |
| CLASS_HARD_FLOORS={ |
| "gate":"Q5_K", |
| "attn_proj":"IQ3_S", |
| "ffn_gate_up":"IQ2_XXS", |
| "ffn_down":"IQ3_S", |
| "norms":"F16", |
| "ssm_params":"F16", |
| |
| "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", |
| } |
| } |
|
|
| |
| PROFILE_MTP_TIER={"nano":"Q5_K","mini":"Q6_K","compact":"Q8_0"} |
| MTP_DEPLOY_TIER="Q8_0" |
|
|
| |
| 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"} |
| |
| 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 |
|
|
| |
| |
| |
|
|
| 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" |
|
|
| |
| |
| |
|
|
| 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, |
| } |
|
|
| |
| |
| |
|
|
| 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 |
|
|
| |
| |
| |
|
|
| 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<=idx<len(TIER_ORDER)):raise IndexError(f"Tier index {idx} out of range") |
| return TIER_ORDER[idx] |
|
|
| def _size_mib(tier:str,n_elements:int)->float: |
| 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 idx<len(TIER_ORDER): |
| if tier_fits(shape,TIER_ORDER[idx]):return TIER_ORDER[idx] |
| idx+=1 |
| return "F16" |
|
|
| def _free_embd_assignments(model:dict,all_names:set,enabled:bool,embd_cap:str | None=None)->dict: |
| |
| 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: |
| |
| 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 |
|
|
| |
| 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<best_diff: |
| best_diff=diff_pct |
| best=(dict(assignments),dict(padded_ne_map),current_size) |
| if diff_pct<=tolerance_pct: |
| print(f" [converge] iter {iteration+1}: {current_size:.0f} MiB (diff {diff_pct:.2f}%) — within tolerance {tolerance_pct}%") |
| return assignments,padded_ne_map |
| print(f" [converge] iter {iteration+1}: {current_size:.0f} MiB (diff {diff_pct:.2f}%) — adjusting effective budget…") |
| if current_size<=0:break |
| correction=target_size_mib/current_size |
| effective_target=min(max(effective_target*correction,target_size_mib*0.95),target_size_mib*1.05) |
| key=round(effective_target,1) |
| if key in seen_targets: |
| print(" [converge] Effective budget stationary — keeping best result") |
| break |
| seen_targets.add(key) |
| print(f" [converge] Selected best result: {best[2]:.0f} MiB (diff {best_diff:.2f}%)") |
| return best[0],best[1] |
|
|
| def compute_mtp_assignments_check(non_mtp_names,mtp_names,importance_table,allow_q3,is_qat,profile,int4_lineage): |
| has_im=any("importance_mean" in importance_table.get(n,{}) for n in mtp_names) |
| mtp_tier=PROFILE_MTP_TIER.get(profile,MTP_DEPLOY_TIER) if has_im else CLASS_HARD_FLOORS["mtp"] |
| if int4_lineage and _tier_index(mtp_tier)>_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: |
| |
| 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 |
|
|
| |
| |
| |
|
|
| 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) |
|
|
| |
| |
| |
|
|
| 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): |
| |
| 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 |
|
|
| |
| |
| |
|
|
| 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}") |
|
|
| |
| 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))}") |
|
|
| |
| |
| |
|
|
| 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): |
| |
| 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: |
| |
| 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) |
|
|
