File size: 2,343 Bytes
c99f13f | 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 | import subprocess
import re
import os
from utils import parse_size_line, parse_quant_size
QUANTIZER_PATH = "/home/maxyag27/llm-tools/llama.cpp/build/bin/llama-quantize"
def _build_cmd(flags: dict, model_in: str, model_out: str, dry_run: bool = False) -> list:
"""Build llama-quantize command. Options before positional args."""
cmd = [QUANTIZER_PATH]
if dry_run:
cmd.append("--dry-run")
if flags.get("imatrix"):
imatrix = flags["imatrix"]
if isinstance(imatrix, list):
imatrix = imatrix[0] # llama-quantize accepts one imatrix; we combined in importance table
cmd.extend(["--imatrix", imatrix])
cmd.extend(["--output-tensor-type", flags["output_tensor_type"]])
cmd.extend(["--token-embedding-type", flags["token_embedding_type"]])
for rule in flags["tensor_type_rules"]:
parts = rule.split(" ", 1)
if len(parts) == 2:
cmd.extend(["--tensor-type", parts[1].strip('"')])
# Positional args: model_in model_out type
cmd.append(model_in)
cmd.append(model_out)
cmd.append(flags["base_type"])
return cmd
def run_dry_run(flags: dict, model_in: str) -> float | None:
"""Run llama-quantize --dry-run and return quant size in MiB."""
cmd = _build_cmd(flags, model_in, "/dev/null", dry_run=True)
result = subprocess.run(
cmd,
capture_output=True,
text=True,
timeout=300,
)
output = (result.stdout or "") + (result.stderr or "")
size = parse_quant_size(output)
if size is not None:
return size
print("STDERR:", result.stderr[:2000])
return None
def run_quantization(flags: dict, model_in: str, model_out: str, dry_run: bool = False) -> bool:
"""Run actual quantization or dry run."""
cmd = _build_cmd(flags, model_in, model_out, dry_run=dry_run)
print("Running:", " ".join(cmd[:6]) + " ...")
result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
if dry_run:
output = (result.stdout or "") + (result.stderr or "")
return parse_quant_size(output)
print(result.stderr[:500])
success = result.returncode == 0
if success:
import os
size_mib = os.path.getsize(model_out) / 1024 / 1024
print(f"Done: {model_out} ({size_mib:.0f} MiB)")
return success
|