| |
| """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) |
|
|
| |
| 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() |
|
|
