File size: 3,736 Bytes
491cce7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | #!/usr/bin/env python3
"""Force ffn_down to specific tier, test PPL."""
import argparse
import os
import subprocess
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from model_reader import read_model
from imatrix_reader import read_imatrix, detect_tied_groups, build_importance_table
from config_generator import generate_flags
from quantizer import run_dry_run, run_quantization
from constants import (
get_tensor_class, get_tensor_type, strip_weight
)
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LLAMA_PPL = os.environ.get(
"LLAMA_PPL", "/home/wepiqx/llama.cpp/build/bin/llama-perplexity"
)
WIKITEXT_DATA = os.environ.get(
"WIKITEXT_DATA", "/mnt/Vsio/wikitext-2-raw/wiki.test.raw"
)
def build_assignments(all_names, ffn_down_tier):
assignments = {}
for name in all_names:
cls = get_tensor_class(get_tensor_type(name))
if cls == "ffn_down":
assignments[name] = ffn_down_tier
elif cls in ("norms", "ssm_params"):
assignments[name] = "F16"
elif cls == "embd":
assignments[name] = "Q5_K"
elif cls == "mtp":
assignments[name] = "Q5_K"
else:
assignments[name] = "Q5_K"
return assignments
def test_tier(model, imatrix_path, model_path, all_names, out_dir, tier):
print(f"\n=== Testing ffn_down = {tier} ===")
assignments = build_assignments(all_names, tier)
flags = generate_flags(assignments, model, "Q5_K_M", 100000)
flags["imatrix"] = imatrix_path
dry_size = run_dry_run(flags, model_path)
print(f" Dry-run: {dry_size:.0f} MiB")
out_path = os.path.join(out_dir, f"ffn_down_{tier}.gguf")
print(f" Quantizing...")
success = run_quantization(flags, model_path, out_path)
if not success:
print(" FAILED")
return None
print(f" Running PPL (GPU)...")
cmd = [
LLAMA_PPL,
"-m", out_path,
"-f", WIKITEXT_DATA,
"-c", "1024", "-ngl", "99"
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=900)
ppl = None
for line in result.stdout.split("\n"):
if "Final estimate: PPL" in line:
ppl = float(line.split("PPL = ")[1].split(" ")[0])
break
print(f" PPL: {ppl}")
os.remove(out_path)
return ppl
def main():
parser = argparse.ArgumentParser(description="Force ffn_down to tier, test PPL")
parser.add_argument("--model", required=True, help="BF16 GGUF model path")
parser.add_argument("--imatrix", required=True, help="Imatrix GGUF path")
parser.add_argument("--out-dir", default=None,
help="Defaults to <repo>/output/ffn_down_test")
parser.add_argument("--tiers", default="Q3_K,IQ4_XS,Q4_K,Q5_K,Q6_K,Q8_0")
args = parser.parse_args()
if args.out_dir is None:
args.out_dir = os.path.join(REPO_ROOT, "output", "ffn_down_test")
os.makedirs(args.out_dir, exist_ok=True)
model = read_model(args.model)
imatrix = read_imatrix(args.imatrix)
tied_groups = detect_tied_groups(imatrix)
imp_table = build_importance_table(imatrix, model)
# Get all tensor names
all_names = set()
for tname in model.get("tensors", {}):
all_names.add(strip_weight(tname))
for tname in imatrix["tensors"]:
all_names.add(strip_weight(tname))
for tier in args.tiers.split(","):
ppl = test_tier(model, args.imatrix, args.model, all_names,
args.out_dir, tier.strip())
if ppl:
print(f" Result: {tier} -> PPL {ppl}")
else:
print(f" Result: {tier} -> FAILED")
if __name__ == "__main__":
main()
|