File size: 3,180 Bytes
c58f252 92af670 c58f252 92af670 c58f252 92af670 c58f252 | 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 | import subprocess
import re
import os
import warnings
from utils import parse_quant_size
QUANTIZER_PATH = os.environ.get(
"LLAMA_QUANTIZE", "/home/wepiqx/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):
if len(imatrix) > 1:
warnings.warn(
"llama-quantize accepts a single imatrix: passing only the "
f"first of {len(imatrix)} files ({imatrix[0]}). The importance "
"table used files combined, so sizes may diverge slightly.",
RuntimeWarning,
)
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_ex(flags: dict, model_in: str):
"""Run llama-quantize --dry-run.
Returns (size_mib | None, raw_output). The raw text is NOT noise: it
carries the per-tensor type the binary will actually write plus the
override log — the only ground truth about what lands in the file.
preflight.audit() consumes it; nothing else should discard it.
"""
cmd = _build_cmd(flags, model_in, "/dev/null", dry_run=True)
result = subprocess.run(
cmd,
capture_output=True,
text=True,
errors="replace", # binary may print non-UTF8 bytes on unknown-arch errors
timeout=300,
)
output = (result.stdout or "") + (result.stderr or "")
return parse_quant_size(output), output
def run_dry_run(flags: dict, model_in: str) -> float | None:
size, output = run_dry_run_ex(flags, model_in)
if size is None:
print("STDERR:", output[:2000])
return size
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, errors="replace", 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
|