| """Quantize Fish Audio Drama 3 to int4 (W4A16): routed MoE experts from block-FP8 AND dense bf16 |
| linears (full-attention q/k/v/o, gated-deltanet in_proj_qkv/in_proj_z/out_proj, shared-expert MLP). |
| |
| Experts are packed exactly like the experts-only reference script (qweight int32 [K/8, N], |
| scales fp16 [K/128, N]) for SGLang's moe_wna16 kernel. Dense linears are written in the standard |
| GPTQ symmetric format (qweight int32 [K/8, N], qzeros int32 [K/128, N/8] filled with 8, |
| scales fp16 [K/128, N], g_idx int32 [K] sorted), which sglang's GPTQMarlinLinearMethod reads and |
| repacks in process_weights_after_loading. Everything else is copied byte-for-byte. |
| |
| python quantize.py SRC_DIR DST_DIR [--device cuda] [--selftest-only] |
| """ |
| import argparse, json, math, os, re, shutil, sys, time |
|
|
| |
|
|
| import torch |
| from safetensors import safe_open |
|
|
| |
| |
| _ST_DTYPE = { |
| torch.float16: "F16", torch.bfloat16: "BF16", torch.float32: "F32", |
| torch.int32: "I32", torch.int64: "I64", torch.uint8: "U8", |
| torch.float8_e4m3fn: "F8_E4M3", |
| } |
|
|
|
|
| def write_safetensors(tensors, path, metadata=None): |
| """Minimal safetensors writer: 8-byte LE header len + JSON header + raw bytes.""" |
| header, blobs, offset = {}, [], 0 |
| for k, t in tensors.items(): |
| t = t.contiguous() |
| dt = t.dtype |
| if dt == torch.bfloat16: |
| raw = t.view(torch.uint16).numpy().tobytes() |
| elif dt == torch.float8_e4m3fn: |
| raw = t.view(torch.uint8).numpy().tobytes() |
| else: |
| raw = t.numpy().tobytes() |
| header[k] = {"dtype": _ST_DTYPE[dt], "shape": list(t.shape), |
| "data_offsets": [offset, offset + len(raw)]} |
| blobs.append(raw) |
| offset += len(raw) |
| if metadata: |
| header["__metadata__"] = metadata |
| hj = json.dumps(header).encode() |
| pad = (8 - (len(hj) % 8)) % 8 |
| hj += b" " * pad |
| with open(path, "wb") as fh: |
| fh.write(len(hj).to_bytes(8, "little")) |
| fh.write(hj) |
| for b in blobs: |
| fh.write(b) |
|
|
| G = 128 |
| CLIPS = torch.linspace(1.0, 0.81, 20) |
| QUANT_METHOD = "drama3-int4" |
|
|
| EXPERT = re.compile(r"\.mlp\.experts\.\d+\.(gate|up|down)_proj\.weight$") |
| |
| DENSE_SELF_ATTN = re.compile( |
| r"^model\.language_model\.layers\.(3|7|11|15|19|23|27|31|35|39)\.self_attn\.(q|k|v|o)_proj\.weight$" |
| ) |
| DENSE_LINEAR_ATTN = re.compile( |
| r"^model\.language_model\.layers\.\d+\.linear_attn\.(in_proj_qkv|in_proj_z|out_proj)\.weight$" |
| ) |
| DENSE_SHARED_EXPERT = re.compile( |
| r"^model\.language_model\.layers\.\d+\.mlp\.shared_expert\.(gate|up|down)_proj\.weight$" |
| ) |
| DENSE_GROUPS = { |
| "attn": ((DENSE_SELF_ATTN,), |
| ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj"]), |
| "gdn": ((DENSE_LINEAR_ATTN,), |
| ["linear_attn.in_proj_qkv", "linear_attn.in_proj_z", "linear_attn.out_proj"]), |
| "shared": ((DENSE_SHARED_EXPERT,), |
| ["mlp.shared_expert.gate_proj", "mlp.shared_expert.up_proj", "mlp.shared_expert.down_proj"]), |
| } |
|
|
|
|
| def dense_class(k, patterns): |
| return any(pat.match(k) for pat in patterns) |
|
|
|
|
| def dequant_fp8(w, s): |
| N, K = w.shape |
| s = s.repeat_interleave(128, 0)[:N].repeat_interleave(128, 1)[:, :K] |
| return w.float() * s.float() |
|
|
|
|
| def quant_rtn(W): |
| """W [N, K] fp32 -> q uint8 [N, K] in 0..15 (offset 8), scales [N, K/G] fp32. Clip search per group.""" |
| N, K = W.shape |
| Wg = W.view(N, K // G, G) |
| amax = Wg.abs().amax(-1, keepdim=True).clamp_min(1e-12) |
| best_err = torch.full(amax.shape, float("inf"), device=W.device) |
| best_s = amax / 7.0 |
| for c in CLIPS.tolist(): |
| s = amax * c / 7.0 |
| q = (Wg / s).round().clamp(-8, 7) |
| err = ((q * s - Wg) ** 2).sum(-1, keepdim=True) |
| better = err < best_err |
| best_err = torch.where(better, err, best_err) |
| best_s = torch.where(better, s, best_s) |
| q = (Wg / best_s).round().clamp(-8, 7) + 8 |
| return q.view(N, K).to(torch.uint8), best_s.squeeze(-1) |
|
|
|
|
| def pack_gptq(q): |
| """q uint8 [N, K] (0..15) -> qweight int32 [K/8, N]: value k sits at bits 4*(k%8) of word k//8.""" |
| N, K = q.shape |
| qt = q.t().contiguous().view(K // 8, 8, N).to(torch.int64) |
| word = torch.zeros(K // 8, N, dtype=torch.int64, device=q.device) |
| for j in range(8): |
| word |= qt[:, j, :] << (4 * j) |
| return (word - (word >= 2**31).to(torch.int64) * 2**32).to(torch.int32) |
|
|
|
|
| def gptq_sym_fields(q, sc): |
| """Standard GPTQ sym checkpoint fields from q uint8 [N, K] (0..15, offset 8), sc [N, K/G] fp32.""" |
| N, K = q.shape |
| qw = pack_gptq(q) |
| scl = sc.t().contiguous().half() |
| qzeros = torch.full((K // G, math.ceil(N / 8)), 0x88888888, |
| dtype=torch.int64, device=q.device).to(torch.int32) |
| g_idx = torch.arange(K, dtype=torch.int32, device=q.device) // G |
| return qw, qzeros, scl, g_idx |
|
|
|
|
| def unpack_gptq(qweight): |
| """qweight int32 [K/8, N] -> q uint8 [N, K] (0..15).""" |
| K8, N = qweight.shape |
| w = qweight.t().contiguous().to(torch.int64) & 0xFFFFFFFF |
| q = torch.zeros(N, K8 * 8, dtype=torch.int64, device=qweight.device) |
| for j in range(8): |
| q[:, j::8] = (w >> (4 * j)) & 0xF |
| return q.to(torch.uint8) |
|
|
|
|
| def loader_view_dequant(qweight, scales): |
| """Mirror of SGLang moe_wna16 (gptq, sym): qweight.T.view(uint8) [N, K/2], low nibble = even k, |
| value - 8, times the group scale. Used to verify our expert packing.""" |
| b = qweight.t().contiguous().view(torch.uint8) |
| lo, hi = (b & 0xF).to(torch.int16), (b >> 4).to(torch.int16) |
| q = torch.stack([lo, hi], -1).view(b.size(0), -1) - 8 |
| s = scales.t().float().repeat_interleave(G, 1) |
| return q.float() * s |
|
|
|
|
| def selftest(device): |
| torch.manual_seed(0) |
| |
| import tempfile |
| sample = { |
| "f16": torch.randn(4, 8).half(), "bf16": torch.randn(4, 8).bfloat16(), |
| "f32": torch.randn(4, 8), "i32": torch.randint(-100, 100, (4, 8), dtype=torch.int32), |
| "fp8": torch.randn(4, 8).to(torch.float8_e4m3fn), |
| } |
| with tempfile.TemporaryDirectory() as td: |
| p = f"{td}/t.safetensors" |
| write_safetensors(sample, p, metadata={"format": "pt"}) |
| with safe_open(p, framework="pt", device="cpu") as f: |
| for k in sample: |
| a, b = sample[k], f.get_tensor(k) |
| av = a.view(torch.uint8) if a.dtype in (torch.bfloat16, torch.float8_e4m3fn) else a |
| bv = b.view(torch.uint8) if b.dtype in (torch.bfloat16, torch.float8_e4m3fn) else b |
| assert torch.equal(av, bv), f"writer round-trip mismatch on {k} ({a.dtype})" |
| print("selftest ok: manual safetensors writer round-trips f16/bf16/f32/i32/fp8", flush=True) |
| |
| W = torch.randn(512, 2048, device=device) * 0.02 |
| q, s = quant_rtn(W) |
| qw = pack_gptq(q) |
| sc = s.t().contiguous().half() |
| rec = loader_view_dequant(qw, sc) |
| ref = (q.float() - 8) * s.repeat_interleave(G, 1) |
| assert torch.allclose(rec, ref.half().float(), atol=1e-3, rtol=1e-2), (rec - ref).abs().max() |
| rel = ((rec - W).norm() / W.norm()).item() |
| assert rel < 0.12, rel |
| print(f"selftest ok: expert packing matches the moe_wna16 loader view; gaussian rel err {rel:.4f}", flush=True) |
|
|
| |
| |
| |
| |
| |
| shapes = [ |
| (8192, 2048), |
| (512, 2048), |
| (2048, 4096), |
| (4096, 2048), |
| (2048, 512), |
| ] |
| for N, K in shapes: |
| W = torch.randn(N, K, device=device) * 0.02 |
| q, s = quant_rtn(W) |
| qw, qz, scl, gi = gptq_sym_fields(q, s) |
| |
| assert qw.shape == (K // 8, N) and qw.dtype == torch.int32, (N, K, "qweight shape/dtype") |
| assert qz.shape == (K // G, (N + 7) // 8) and qz.dtype == torch.int32, (N, K, "qzeros shape/dtype") |
| assert scl.shape == (K // G, N) and scl.dtype == torch.float16, (N, K, "scales shape/dtype") |
| assert gi.shape == (K,) and gi.dtype == torch.int32, (N, K, "g_idx shape/dtype") |
| |
| assert (qz == -2004318072).all(), (N, K, "qzeros must be packed 8s (sym, zp=8)") |
| assert torch.equal(gi, torch.arange(K, dtype=torch.int32, device=gi.device) // G), (N, K, "g_idx sorted") |
| |
| uq = unpack_gptq(qw) |
| assert torch.equal(uq, q), (N, K, "qweight round-trip") |
| rec = (uq.float() - 8) * scl.float().t().repeat_interleave(G, 1) |
| relw = ((rec - W).norm() / W.norm()).item() |
| assert relw < 0.12, (N, K, relw) |
| print(f"selftest ok: dense GPTQ fields {N}x{K}: pre-repack weight rel err {relw:.4f}", flush=True) |
|
|
| |
| |
| |
| |
| try: |
| from sglang.srt.layers.linear import LinearBase |
| from sglang.srt.layers.quantization.gptq.gptq import ( |
| GPTQMarlinConfig, |
| GPTQMarlinLinearMethod, |
| ) |
| if device == "cpu" or not torch.cuda.is_available(): |
| _monkeypatch_marlin_for_cpu() |
| full = {"quant_method": "gptq_marlin", "bits": 4, "group_size": G, "sym": True, "desc_act": False} |
| cfg = GPTQMarlinConfig(4, G, False, True, False, {}, full) |
| method = GPTQMarlinLinearMethod(cfg) |
| N, K = 512, 2048 |
| W = torch.randn(N, K, device=device) * 0.02 |
| q, s = quant_rtn(W) |
| qw, qz, scl, gi = gptq_sym_fields(q, s) |
| layer = LinearBase(K, N, skip_bias_add=True, params_dtype=torch.bfloat16).to(device) |
| method.create_weights(layer, K, [N], K, N, torch.bfloat16, weight_loader=None) |
| layer.qweight.data.copy_(qw) |
| layer.qzeros.data.copy_(qz) |
| layer.scales.data.copy_(scl.to(layer.scales.dtype)) |
| layer.g_idx.data.copy_(gi) |
| with torch.no_grad(): |
| method.process_weights_after_loading(layer) |
| print("selftest ok: dense GPTQ-Marlin loader accepted layer and repacked weights", flush=True) |
| except Exception as e: |
| print(f"selftest: dense loader smoke check failed on CPU (repack numerics are GPU-only; " |
| f"logged, not fatal): {type(e).__name__}: {e}", flush=True) |
|
|
|
|
| def _monkeypatch_marlin_for_cpu(): |
| """Let the GPTQ-Marlin code path run on a CUDA-less box. |
| |
| The marlin repack/GEMM kernels themselves are CUDA-only, so on CPU we only validate |
| checkpoint-field shapes and the create_weights / process_weights_after_loading control |
| flow. We stub: the device-capability check in the scheme constructor, the shape-check in |
| the kernel, the workspace allocation, and the repack custom op (identity that keeps the |
| qweight tensor intact). |
| """ |
| from sglang.srt.layers.quantization import marlin_utils as _mu |
| from sglang.srt.layers.quantization.gptq.schemes import gptq_marlin as _sch |
| from sglang.srt.hardware_backend.gpu.quantization import gptq_kernels as _ker |
|
|
| _sch.verify_marlin_supported = lambda *a, **kw: None |
| _ker.check_marlin_supports_shape = lambda *a, **kw: (True, None) |
| _ker.marlin_make_workspace = lambda device, max_blocks_per_sm=1: torch.zeros( |
| 1, dtype=torch.int, device="cpu", requires_grad=False |
| ) |
| _ker.gptq_marlin_repack = lambda q, perm, size_k, size_n, num_bits: q |
| _ker.marlin_permute_scales = lambda s, size_k, size_n, group_size: s |
|
|
|
|
| def _acc(stats, cls): |
| return stats.setdefault(cls, {"tensors": 0, "rel_err_sum": 0.0, "rel_err_max": 0.0, |
| "src_bytes": 0, "int4_bytes": 0}) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("src", nargs="?", default=os.environ.get("QUANT_SRC", "/data/model")) |
| ap.add_argument("dst", nargs="?", default=os.environ.get("QUANT_DST", "/root/drama-3-int4")) |
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| ap.add_argument("--shard_gb", type=float, default=2.0) |
| ap.add_argument("--selftest-only", action="store_true") |
| ap.add_argument("--dense", default="attn,shared,gdn", |
| help="comma list of dense groups to quantize: attn,gdn,shared (or 'none')") |
| a = ap.parse_args() |
| groups = [g.strip() for g in a.dense.split(",") if g.strip() and g.strip() != "none"] |
| dense_patterns = [p for g in groups for p in DENSE_GROUPS[g][0]] |
| dense_targets = [t for g in groups for t in DENSE_GROUPS[g][1]] |
| print(f"dense groups: {groups} ({len(dense_targets)} targets)", flush=True) |
| dev = a.device |
| selftest(dev) |
| if a.selftest_only: |
| return |
| os.makedirs(a.dst, exist_ok=True) |
| |
| staging = os.environ.get("QUANT_STAGING", "/root/model-src") |
| if not os.path.exists(f"{staging}/model.safetensors.index.json"): |
| print(f"staging {a.src} -> {staging} (local disk) ...", flush=True) |
| t_st = time.time() |
| shutil.copytree(a.src, staging, dirs_exist_ok=True, copy_function=shutil.copy) |
| print(f"staged in {time.time()-t_st:.0f}s", flush=True) |
| a.src = staging |
| idx = json.load(open(f"{a.src}/model.safetensors.index.json")) |
| files = sorted(set(idx["weight_map"].values())) |
| out, out_bytes, shard_id, new_map = {}, 0, 0, {} |
| stats = {} |
| t0 = time.time() |
|
|
| def flush(): |
| nonlocal out, out_bytes, shard_id |
| if not out: |
| return |
| shard_id += 1 |
| name = f"model-{shard_id:05d}.safetensors" |
| write_safetensors(out, f"{a.dst}/{name}", metadata={"format": "pt"}) |
| for k in out: |
| new_map[k] = name |
| out, out_bytes = {}, 0 |
|
|
| def put(k, t): |
| nonlocal out_bytes |
| out[k] = t.contiguous().cpu() |
| out_bytes += t.numel() * t.element_size() |
| if out_bytes > a.shard_gb * 1e9: |
| flush() |
|
|
| for fi, fn in enumerate(files): |
| with safe_open(f"{a.src}/{fn}", framework="pt", device="cpu") as f: |
| keys = list(f.keys()) |
| for k in keys: |
| if k.endswith(".weight_scale_inv") and EXPERT.search(k[: -len("_scale_inv")]): |
| continue |
| t = f.get_tensor(k) |
| if EXPERT.search(k): |
| st = _acc(stats, "experts") |
| sk = k + "_scale_inv" |
| s = f.get_tensor(sk) if sk in keys else None |
| if s is None: |
| with safe_open(f"{a.src}/{idx['weight_map'][sk]}", framework="pt", device="cpu") as g: |
| s = g.get_tensor(sk) |
| W = dequant_fp8(t.to(dev), s.to(dev)) |
| q, sc = quant_rtn(W) |
| rec = (q.float() - 8) * sc.repeat_interleave(G, 1) |
| rel = ((rec - W).norm() / W.norm().clamp_min(1e-12)).item() |
| st["tensors"] += 1 |
| st["rel_err_sum"] += rel |
| st["rel_err_max"] = max(st["rel_err_max"], rel) |
| st["src_bytes"] += t.numel() + s.numel() * 4 |
| base = k[: -len(".weight")] |
| qw, scl = pack_gptq(q), sc.t().contiguous().half() |
| st["int4_bytes"] += qw.numel() * 4 + scl.numel() * 2 |
| put(base + ".qweight", qw) |
| put(base + ".scales", scl) |
| elif dense_patterns and dense_class(k, dense_patterns): |
| st = _acc(stats, "dense") |
| if t.shape[1] % G: |
| raise ValueError(f"dense target {k} has K={t.shape[1]} not divisible by {G}") |
| W = t.to(dev).float() |
| q, sc = quant_rtn(W) |
| rec = (q.float() - 8) * sc.repeat_interleave(G, 1) |
| rel = ((rec - W).norm() / W.norm().clamp_min(1e-12)).item() |
| st["tensors"] += 1 |
| st["rel_err_sum"] += rel |
| st["rel_err_max"] = max(st["rel_err_max"], rel) |
| st["src_bytes"] += t.numel() * t.element_size() |
| base = k[: -len(".weight")] |
| qw, qz, scl, gi = gptq_sym_fields(q, sc) |
| st["int4_bytes"] += qw.numel() * 4 + qz.numel() * 4 + scl.numel() * 2 + gi.numel() * 4 |
| put(base + ".qweight", qw) |
| put(base + ".qzeros", qz) |
| put(base + ".scales", scl) |
| put(base + ".g_idx", gi) |
| else: |
| put(k, t) |
| nt = sum(s["tensors"] for s in stats.values()) |
| if nt and nt % 2000 == 0: |
| print(f" [{time.time()-t0:6.0f}s] ...{nt} quantized tensors so far in shard {fi+1}", flush=True) |
| ex, de = stats.get("experts", {}), stats.get("dense", {}) |
| print(f"[{time.time()-t0:6.0f}s] {fn} done ({fi+1}/{len(files)}); " |
| f"experts {ex.get('tensors',0)} mean err {ex.get('rel_err_sum',0)/max(1,ex.get('tensors',1)):.4f}, " |
| f"dense {de.get('tensors',0)} mean err {de.get('rel_err_sum',0)/max(1,de.get('tensors',1)):.4f}", flush=True) |
| flush() |
| |
| n = shard_id |
| final_map = {} |
| for i in range(1, n + 1): |
| src_name, dst_name = f"model-{i:05d}.safetensors", f"model-{i:05d}-of-{n:05d}.safetensors" |
| os.replace(f"{a.dst}/{src_name}", f"{a.dst}/{dst_name}") |
| for k, v in new_map.items(): |
| final_map[k] = v.replace(".safetensors", f"-of-{n:05d}.safetensors") |
| total = sum(os.path.getsize(f"{a.dst}/{v}") for v in set(final_map.values())) |
| json.dump({"metadata": {"total_size": total}, "weight_map": final_map}, open(f"{a.dst}/model.safetensors.index.json", "w"), indent=2) |
| cfg = json.load(open(f"{a.src}/config.json")) |
| cfg["quantization_config"] = { |
| "quant_method": QUANT_METHOD, |
| "bits": 4, |
| "group_size": G, |
| "sym": True, |
| "desc_act": False, |
| "experts": {"quant_method": "gptq", "bits": 4, "group_size": G, "sym": True}, |
| "dense": {"quant_method": "gptq_marlin", "bits": 4, "group_size": G, "sym": True, |
| "desc_act": False, |
| "targets": dense_targets}, |
| } |
| json.dump(cfg, open(f"{a.dst}/config.json", "w"), indent=2) |
| for fn in os.listdir(a.src): |
| if not (fn.endswith(".safetensors") or fn in ("model.safetensors.index.json", "config.json")) and os.path.isfile(f"{a.src}/{fn}"): |
| shutil.copy2(f"{a.src}/{fn}", f"{a.dst}/{fn}") |
| for cls, st in stats.items(): |
| st["mean_rel_err"] = st["rel_err_sum"] / max(1, st["tensors"]) |
| json.dump(stats, open(f"{a.dst}/quant_stats.json", "w"), indent=2) |
| print("done", json.dumps(stats), f"total {total/1e9:.2f} GB", flush=True) |
| repo = os.environ.get("QUANT_UPLOAD_REPO", "akhaliq/drama-3-int4") |
| if repo: |
| from huggingface_hub import HfApi |
| HfApi(token=os.environ["HF_TOKEN"]).upload_folder( |
| repo_id=repo, folder_path=a.dst, |
| commit_message="int4 quant (experts + dense linears) of fishaudio/drama-3") |
| print(f"uploaded {a.dst} -> {repo}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
