"""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 # Note: the sglang loader smoke check in selftest is skipped when sglang is not installed. import torch from safetensors import safe_open # safetensors writers (save_file / serialize) proved version-fragile and mmap-EFAULT-prone on the # job runners, so shards are written with the manual writer below; safe_open is only used for reading. _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 bf16 linears to quantize (full-attention layers only for self_attn). 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) # int32 [K/8, N] scl = sc.t().contiguous().half() # fp16 [K/G, N] qzeros = torch.full((K // G, math.ceil(N / 8)), 0x88888888, dtype=torch.int64, device=q.device).to(torch.int32) # wraps to -2004318072 g_idx = torch.arange(K, dtype=torch.int32, device=q.device) // G # ascending groups (desc_act False) 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) # --- manual safetensors writer round-trip (all dtypes the shards contain) ------------------- 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) # --- expert path (moe_wna16 packing round-trip) ------------------------------------------- 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) # --- dense path: PRE-REPACK verification of the on-disk GPTQ v1 fields -------------------- # The Marlin repack (transform_w_q / gptq_marlin_repack) is GPU-only in sglang's kernel # path, so here we verify exactly what lands in the checkpoint: field shapes/dtypes match # what GPTQMarlinLinearMethod.create_weights registers, and a reference unpack (same # nibble order as the expert loader-view check) reconstructs W. shapes = [ # (N, K) from the fishaudio/drama-3 checkpoint (8192, 2048), # self_attn.q_proj, linear_attn.in_proj_qkv (512, 2048), # self_attn.k_proj / v_proj, mlp.shared_expert.gate_proj / up_proj (2048, 4096), # self_attn.o_proj, linear_attn.out_proj (4096, 2048), # linear_attn.in_proj_z (2048, 512), # mlp.shared_expert.down_proj ] 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) # exact shapes/dtypes registered by the GPTQ-Marlin scheme's create_weights 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") # sym format invariants: qzeros = packed 8s (zp=8), g_idx = arange//G (sorted => desc_act False) 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") # reference unpack: same nibble order the expert loader-view check uses 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) # smoke check only: the real sglang loader must accept our layer (create_weights + copy # into the registered params). The Marlin repack inside process_weights_after_loading # runs on GPU in production; on CPU it may choke, so any failure here is logged, not # fatal. Repack numerics + apply are verified on GPU by the L40S serve test. 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 # smallest real dense shape 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) # mmap reads over the fuse volume SIGBUS; stage the checkpoint to local disk once and read locally 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 # consumed with its weight 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: # scale in another shard 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() # rename shards to the HF convention 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()