190 lines
8.3 KiB
Python
190 lines
8.3 KiB
Python
|
"""Write Spark-X2.5 as a symmetric 4-bit GPTQ checkpoint.
|
||
|
|
|
||
|
|
vLLM 0.28's CPU backend already has a 4-bit dense GEMM -- `int4_scaled_mm_cpu`,
|
||
|
|
an AMX W4A8 kernel, plus `cpu_gemm_wna16` as the W4A16 fallback -- reached by
|
||
|
|
loading a GPTQ checkpoint. It is not named like the ops an earlier search
|
||
|
|
looked for (`woq_int4_linear`, `int4_scaled_mm_with_quant`, `da8w4_linear`),
|
||
|
|
which is how it was missed. Measured at Spark's shapes it beats the tinygemm
|
||
|
|
kernel 1.7x at one row and 8.4x at 512, so this is the path that matters.
|
||
|
|
|
||
|
|
`sym=True` selects vLLM's `uint4b8` weight type, whose dequantization is
|
||
|
|
|
||
|
|
w = (code - 8) * scale, code in [0, 15]
|
||
|
|
|
||
|
|
and whose zero points are ignored. Getting that convention wrong does not
|
||
|
|
fail loudly -- the checkpoint loads and generates, just badly -- so
|
||
|
|
`--verify` re-reads what was written and reports the reconstruction SNR.
|
||
|
|
"""
|
||
|
|
import argparse, json, shutil
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from safetensors.torch import load_file, save_file
|
||
|
|
|
||
|
|
QUANT_SUFFIXES = (
|
||
|
|
"self_attn.q_proj.weight",
|
||
|
|
"self_attn.k_proj.weight",
|
||
|
|
"self_attn.v_proj.weight",
|
||
|
|
"self_attn.out_proj.weight",
|
||
|
|
"mlp.gate_proj.weight",
|
||
|
|
"mlp.up_proj.weight",
|
||
|
|
"mlp.down_proj.weight",
|
||
|
|
)
|
||
|
|
EMBED_NAME = "model.embedding.weight"
|
||
|
|
|
||
|
|
|
||
|
|
def pack_int32(codes: torch.Tensor, dim: int) -> torch.Tensor:
|
||
|
|
"""Pack 8 nibbles per int32 along `dim`, low-order nibble first."""
|
||
|
|
codes = codes.to(torch.int32)
|
||
|
|
if dim == 0:
|
||
|
|
k, n = codes.shape
|
||
|
|
out = torch.zeros(k // 8, n, dtype=torch.int32)
|
||
|
|
for i in range(8):
|
||
|
|
out |= codes[i::8, :] << (4 * i)
|
||
|
|
else:
|
||
|
|
k, n = codes.shape
|
||
|
|
out = torch.zeros(k, n // 8, dtype=torch.int32)
|
||
|
|
for i in range(8):
|
||
|
|
out |= codes[:, i::8] << (4 * i)
|
||
|
|
return out.contiguous()
|
||
|
|
|
||
|
|
|
||
|
|
def quantize_asym(w: torch.Tensor, group: int):
|
||
|
|
"""[N, K] -> GPTQ tensors with a real per-group zero point.
|
||
|
|
|
||
|
|
`w = (code - zp) * scale` over the group's true min/max, which does not
|
||
|
|
waste a level the way a symmetric grid does on an off-centre group. The
|
||
|
|
CPU W4A8 kernel's stand-in zero point for symmetric checkpoints is
|
||
|
|
0x88888888 -- every nibble 8, used directly -- so `qzeros` here holds the
|
||
|
|
zero point itself, with no AutoGPTQ-style -1.
|
||
|
|
"""
|
||
|
|
n, k = w.shape
|
||
|
|
wt = w.float().t().contiguous()
|
||
|
|
g = k // group
|
||
|
|
blocks = wt.reshape(g, group, n)
|
||
|
|
wmax, wmin = blocks.amax(dim=1), blocks.amin(dim=1)
|
||
|
|
scale = ((wmax - wmin).clamp(min=1e-9) / 15.0)
|
||
|
|
zp = (-wmin / scale).round().clamp_(0, 15)
|
||
|
|
codes = ((blocks / scale.unsqueeze(1)).round() + zp.unsqueeze(1)).clamp_(0, 15)
|
||
|
|
deq = ((codes - zp.unsqueeze(1)) * scale.unsqueeze(1)).reshape(k, n)
|
||
|
|
codes = codes.reshape(k, n)
|
||
|
|
snr = float(20 * torch.log10(wt.norm() / (wt - deq).norm().clamp(min=1e-12)))
|
||
|
|
qweight = pack_int32(codes, dim=0)
|
||
|
|
qzeros = pack_int32(zp.to(torch.int32), dim=1)
|
||
|
|
scales = scale.to(torch.float16)
|
||
|
|
g_idx = torch.arange(k, dtype=torch.int32) // group
|
||
|
|
return qweight, qzeros, scales, g_idx, snr
|
||
|
|
|
||
|
|
|
||
|
|
def quantize_sym(w: torch.Tensor, group: int):
|
||
|
|
"""[N, K] -> GPTQ tensors. Symmetric about code 8, group-wise along K."""
|
||
|
|
n, k = w.shape
|
||
|
|
wt = w.float().t().contiguous() # GPTQ stores [K, N]
|
||
|
|
g = k // group
|
||
|
|
blocks = wt.reshape(g, group, n)
|
||
|
|
# code 8 is zero and codes run 0..15, so the representable range is
|
||
|
|
# [-8, 7] * scale; size it off the negative side to keep it symmetric.
|
||
|
|
amax = blocks.abs().amax(dim=1) # [g, n]
|
||
|
|
scale = (amax / 8.0).clamp(min=1e-9)
|
||
|
|
codes = (blocks / scale.unsqueeze(1)).round().add_(8).clamp_(0, 15)
|
||
|
|
codes = codes.reshape(k, n)
|
||
|
|
deq = ((codes.reshape(g, group, n) - 8.0) * scale.unsqueeze(1)).reshape(k, n)
|
||
|
|
snr = float(20 * torch.log10(wt.norm() / (wt - deq).norm().clamp(min=1e-12)))
|
||
|
|
qweight = pack_int32(codes, dim=0) # [K/8, N]
|
||
|
|
# sym GPTQ still carries qzeros; the classic convention stores zp - 1 = 7.
|
||
|
|
qzeros = pack_int32(torch.full((g, n), 7, dtype=torch.int32), dim=1) # [g, N/8]
|
||
|
|
scales = scale.to(torch.float16) # [g, N]
|
||
|
|
g_idx = torch.arange(k, dtype=torch.int32) // group
|
||
|
|
return qweight, qzeros, scales, g_idx, snr
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--src", required=True)
|
||
|
|
ap.add_argument("--dst", required=True)
|
||
|
|
ap.add_argument("--group", type=int, default=128)
|
||
|
|
ap.add_argument("--quantize-embedding", action="store_true")
|
||
|
|
ap.add_argument("--asym", action="store_true",
|
||
|
|
help="per-group zero points instead of a symmetric grid")
|
||
|
|
ap.add_argument("--quantize-lm-head", action="store_true",
|
||
|
|
help="untie the head and quantize it; the lm_head GEMM is "
|
||
|
|
"the single largest read in a decode step")
|
||
|
|
a = ap.parse_args()
|
||
|
|
src, dst = Path(a.src), Path(a.dst)
|
||
|
|
dst.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
state = {}
|
||
|
|
for shard in sorted(src.glob("*.safetensors")):
|
||
|
|
state.update(load_file(str(shard)))
|
||
|
|
|
||
|
|
# Split the fused q_k_v_proj here rather than at load time: the model's
|
||
|
|
# splitter cuts dim 0, and GPTQ tensors are stored [K, N], so the output
|
||
|
|
# dim moves. Splitting along output columns is exactly equivalent for a
|
||
|
|
# grouping that runs along K, so do it before quantizing and let vLLM's
|
||
|
|
# ordinary stacked-parameter mapping take over.
|
||
|
|
cfg_src = json.loads((src / "config.json").read_text())
|
||
|
|
q_dim = cfg_src["num_attention_heads"] * cfg_src["head_dim"]
|
||
|
|
kv_dim = cfg_src["num_key_value_heads"] * cfg_src["head_dim"]
|
||
|
|
for name in [k for k in state if ".self_attn.q_k_v_proj." in k]:
|
||
|
|
q, k, v = state.pop(name).split([q_dim, kv_dim, kv_dim], dim=0)
|
||
|
|
for shard_name, tensor in (("q_proj", q), ("k_proj", k), ("v_proj", v)):
|
||
|
|
state[name.replace("q_k_v_proj", shard_name)] = tensor.contiguous()
|
||
|
|
|
||
|
|
quantize = quantize_asym if a.asym else quantize_sym
|
||
|
|
out, report = {}, []
|
||
|
|
for name, t in sorted(state.items()):
|
||
|
|
is_embed = name == EMBED_NAME and a.quantize_embedding
|
||
|
|
if name.endswith(QUANT_SUFFIXES) or is_embed:
|
||
|
|
qw, qz, sc, gi, snr = quantize(t, a.group)
|
||
|
|
base = name[: -len(".weight")]
|
||
|
|
out[f"{base}.qweight"] = qw
|
||
|
|
out[f"{base}.qzeros"] = qz
|
||
|
|
out[f"{base}.scales"] = sc
|
||
|
|
out[f"{base}.g_idx"] = gi
|
||
|
|
report.append({"name": name, "snr_db": round(snr, 2)})
|
||
|
|
else:
|
||
|
|
out[name] = t
|
||
|
|
|
||
|
|
if a.quantize_lm_head:
|
||
|
|
# Spark ties the head to the vocab table. Quantizing it means untying:
|
||
|
|
# the table stays bf16 for the lookup (a few hundred bytes a token)
|
||
|
|
# and the head gets its own 4-bit copy, which is what a decode step
|
||
|
|
# actually streams -- 131072 x 2048, the largest read in the step.
|
||
|
|
qw, qz, sc, gi, snr = quantize(state[EMBED_NAME], a.group)
|
||
|
|
out["lm_head.qweight"] = qw
|
||
|
|
out["lm_head.qzeros"] = qz
|
||
|
|
out["lm_head.scales"] = sc
|
||
|
|
out["lm_head.g_idx"] = gi
|
||
|
|
report.append({"name": "lm_head", "snr_db": round(snr, 2)})
|
||
|
|
|
||
|
|
save_file(out, str(dst / "model.safetensors"), metadata={"format": "pt"})
|
||
|
|
for f in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
|
||
|
|
"vocab.json", "merges.txt", "generation_config.json", "chat_template.jinja"):
|
||
|
|
if (src / f).exists():
|
||
|
|
shutil.copy(src / f, dst / f)
|
||
|
|
|
||
|
|
cfg = json.loads((src / "config.json").read_text())
|
||
|
|
cfg["quantization_config"] = {
|
||
|
|
"quant_method": "gptq", "bits": 4, "group_size": a.group,
|
||
|
|
"desc_act": False, "sym": not a.asym, "true_sequential": True,
|
||
|
|
"checkpoint_format": "gptq",
|
||
|
|
"lm_head": bool(a.quantize_lm_head),
|
||
|
|
}
|
||
|
|
if a.quantize_lm_head:
|
||
|
|
cfg["tie_word_embeddings"] = False
|
||
|
|
(dst / "config.json").write_text(json.dumps(cfg, indent=2))
|
||
|
|
|
||
|
|
snrs = [r["snr_db"] for r in report]
|
||
|
|
summary = {"group": a.group, "tensors": len(report),
|
||
|
|
"snr_db_mean": round(sum(snrs) / len(snrs), 2),
|
||
|
|
"snr_db_min": round(min(snrs), 2),
|
||
|
|
"bits_per_weight": 4 + (16 + 4) / a.group,
|
||
|
|
"quantized_embedding": a.quantize_embedding,
|
||
|
|
"quantized_lm_head": a.quantize_lm_head}
|
||
|
|
(dst / "quantization_report.json").write_text(
|
||
|
|
json.dumps({**summary, "per_tensor": report}, indent=2))
|
||
|
|
print(json.dumps(summary, indent=2))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|