Upload folder using huggingface_hub

fdc6474 verified 3 months ago

11.5 kB

#!/usr/bin/env python3
"""AQLM quantization of GLM-5.2 routed experts from the NVFP4 checkpoint.
For each AQLM-assigned layer (see hybrid_plan.json):
  1. Load the layer's 256 experts' NVFP4 tensors from the downloaded repo,
     dequantize to bf16 on GPU (u8-packed fp4 * fp8 block scale * fp32
     global scale).
  2. Per projection (w13 = gate rows then up rows; w2 = down):
     a. per-output-channel scale s[e,o] (mean abs of the row);
     b. fit one 65536 x 8 codebook per book with GPU k-means on a row
        sample (codebook shared across all 256 experts);
     c. encode all groups-of-8 by chunked nearest-centroid search
        (books=2 encodes the residual against a second codebook);
     d. least-squares refit of the per-channel scales.
  3. Save codes/codebooks/scales + reconstruction stats per layer.
Everything is chunked: peak GPU memory is the bf16 expert tensor plus a
few GB of scratch, so 8 workers (one per GPU) run comfortably.
Usage: aqlm_quantize.py [--layers 11,12] [--gpus 0,1,...] [--out DIR]
"""

import argparse
import json
import os
import time
from concurrent.futures import ProcessPoolExecutor

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
NVFP4_DIR = os.environ.get("GLM52_NVFP4_DIR", "/tmp/glm52-dl/nvfp4-full")
OUT_DIR = "/data/glm52-aqlm-parts"
ENTRIES = 65536
GDIM = 8
SAMPLE_VECS = 4_000_000
KMEANS_ITERS = 12
CHUNK_VECS = 32_768  # scores are [CHUNK_VECS, 65536] fp16 (~4.3GB)

FP4_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
           -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]


def load_plan():
    plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))

    def books(tier):
        t = plan["aqlm"][tier]
        return {"w13": t["w13_books"], "w2": t["w2_books"]}

    layers = {}
    for li in plan["aqlm_mixed_layers"]:
        layers[li] = books("mixed")
    for li in plan["aqlm_cold_layers"]:
        layers[li] = books("cold")
    return layers


class ShardReader:
    """Random access to tensors across the downloaded safetensors shards."""

    def __init__(self, repo_dir: str):
        idx = json.load(
            open(os.path.join(repo_dir, "model.safetensors.index.json"))
        )
        self.weight_map = idx["weight_map"]
        self.repo_dir = repo_dir
        self._open = {}

    def get(self, name: str):
        from safetensors import safe_open

        shard = self.weight_map[name]
        if shard not in self._open:
            self._open[shard] = safe_open(
                os.path.join(self.repo_dir, shard), framework="pt"
            )
        return self._open[shard].get_tensor(name)


def dequant_nvfp4(reader, prefix: str, device, lut) -> "torch.Tensor":
    """Dequantize one NVFP4 linear weight to bf16 [out, in]."""
    import torch

    packed = reader.get(f"{prefix}.weight").to(device)          # u8 [out, in/2]
    wscale = reader.get(f"{prefix}.weight_scale").to(device)    # fp8 [out, in/16]
    wscale2 = reader.get(f"{prefix}.weight_scale_2").to(device)  # fp32 scalar

    lo = (packed & 0x0F).long()
    hi = (packed >> 4).long()
    vals = torch.empty(
        packed.shape[0], packed.shape[1] * 2, dtype=torch.float32, device=device
    )
    # NVFP4 packs the first element of each pair in the low nibble.
    vals[:, 0::2] = lut[lo]
    vals[:, 1::2] = lut[hi]
    del lo, hi
    scale = wscale.to(torch.float32).repeat_interleave(16, dim=1)
    return (vals * scale * wscale2.to(torch.float32)).to(torch.bfloat16)


def _argmin_codes(v16, cent16, cn16):
    """v16 [n,8] fp16, cent16 [ENTRIES,8] fp16 -> nearest ids int32.
    argmax of 2*v.c - |c|^2 (|v|^2 constant per vector)."""
    import torch

    d = v16 @ cent16.t()
    d.mul_(2).sub_(cn16)
    return d.argmax(-1).to(torch.int32)


def encode_chunked(vecs16, cent16):
    import torch

    cn16 = (cent16.float() * cent16.float()).sum(-1).half()
    out = torch.empty(
        vecs16.shape[0], dtype=torch.int32, device=vecs16.device
    )
    for s in range(0, vecs16.shape[0], CHUNK_VECS):
        out[s : s + CHUNK_VECS] = _argmin_codes(
            vecs16[s : s + CHUNK_VECS], cent16, cn16
        )
    return out


def kmeans_fit(vecs16, iters: int, generator):
    """Lloyd's k-means on [N,8] fp16; returns [ENTRIES,8] fp16 centroids."""
    import torch

    n = vecs16.shape[0]
    device = vecs16.device
    perm = torch.randperm(n, generator=generator, device=device)[:ENTRIES]
    cent = vecs16[perm].float()
    for _ in range(iters):
        cent16 = cent.half()
        cn16 = (cent * cent).sum(-1).half()
        sums = torch.zeros(ENTRIES, GDIM, device=device)
        cnts = torch.zeros(ENTRIES, device=device)
        ones = None
        for s in range(0, n, CHUNK_VECS):
            v = vecs16[s : s + CHUNK_VECS]
            a = _argmin_codes(v, cent16, cn16).long()
            sums.index_add_(0, a, v.float())
            if ones is None or ones.shape[0] != a.shape[0]:
                ones = torch.ones(a.shape[0], device=device)
            cnts.index_add_(0, a, ones[: a.shape[0]])
        mask = cnts > 0
        cent[mask] = sums[mask] / cnts[mask].unsqueeze(-1)
        dead = int((~mask).sum())
        if dead:
            ridx = torch.randperm(n, generator=generator, device=device)[:dead]
            cent[~mask] = vecs16[ridx].float()
    return cent.half()


def quantize_matrix(w, books: int, generator, log):
    """w: [E, out, in] bf16 (GPU) -> dict with codes int16 [E,books,out,in/8],
    codebooks fp16 [books,ENTRIES,8], scales fp16 [E,out], rel mse."""
    import torch

    e, out, k = w.shape
    device = w.device
    rows = e * out
    k8 = k // GDIM
    wv = w.reshape(rows, k)

    # per-row scales, chunked
    scales = torch.empty(rows, dtype=torch.float32, device=device)
    ROWCH = max(1, (CHUNK_VECS * 64) // k)
    for s in range(0, rows, ROWCH):
        scales[s : s + ROWCH] = (
            wv[s : s + ROWCH].float().abs().mean(-1).clamp_min(1e-8)
        )

    # --- fit codebooks on a row sample ---
    n_sample_rows = min(rows, max(1, SAMPLE_VECS // k8))
    ridx = torch.randperm(rows, generator=generator, device=device)[
        :n_sample_rows
    ]
    samp = (
        (wv[ridx].float() / scales[ridx].unsqueeze(-1))
        .reshape(-1, GDIM)
        .half()
    )
    log(f"    kmeans book0 on {samp.shape[0]} vecs")
    cbs = [kmeans_fit(samp, KMEANS_ITERS, generator)]
    if books == 2:
        c0 = encode_chunked(samp, cbs[0])
        resid = (samp.float() - cbs[0].float()[c0.long()]).half()
        log(f"    kmeans book1 on residuals")
        cbs.append(kmeans_fit(resid, KMEANS_ITERS, generator))
        del c0, resid
    del samp
    torch.cuda.empty_cache()

    cb_f = [c.float() for c in cbs]
    cn16 = [(cf * cf).sum(-1).half() for cf in cb_f]

    # --- full encode, chunked by rows; accumulate scale-refit stats ---
    codes = torch.empty(
        (rows, books, k8), dtype=torch.int16, device=device
    )
    num = torch.zeros(rows, dtype=torch.float32, device=device)
    den = torch.zeros(rows, dtype=torch.float32, device=device)
    sse = 0.0
    stot = 0.0
    ROWCH = max(1, CHUNK_VECS // k8)
    for s in range(0, rows, ROWCH):
        wc = wv[s : s + ROWCH].float() / scales[s : s + ROWCH].unsqueeze(-1)
        v = wc.reshape(-1, GDIM).half()
        approx = torch.zeros(v.shape[0], GDIM, device=device)
        r = v.float()
        for b in range(books):
            cb = _argmin_codes(r.half(), cbs[b], cn16[b]).long()
            dec = cb_f[b][cb]
            approx += dec
            r -= dec
            codes[s : s + ROWCH, b] = (
                torch.where(cb >= 32768, cb - 65536, cb)
                .to(torch.int16)
                .reshape(-1, k8)
            )
        nrow = wc.shape[0]
        a2 = approx.reshape(nrow, k)
        w2 = wc
        num[s : s + ROWCH] = (w2 * a2).sum(-1)
        den[s : s + ROWCH] = (a2 * a2).sum(-1).clamp_min(1e-8)
        sse += float((r * r).sum())
        stot += float((w2 * w2).sum())
        del wc, v, approx, r, a2

    rel = sse / max(stot, 1e-9)
    scales = scales * (num / den).clamp(0.5, 2.0)
    log(f"    rel_mse={rel:.4f}")

    return {
        "codes": codes.reshape(e, out, books, k8)
        .permute(0, 2, 1, 3)
        .contiguous()
        .cpu(),
        "codebooks": torch.stack(cbs).cpu(),
        "scales": scales.reshape(e, out).half().cpu(),
        "rel_mse": rel,
    }


def process_layer(layer_idx: int, books: dict, gpu: int):
    import torch

    t0 = time.time()
    device = f"cuda:{gpu}"
    torch.cuda.set_device(gpu)
    gen = torch.Generator(device=device)
    gen.manual_seed(1234 + layer_idx)
    lut = torch.tensor(FP4_LUT, dtype=torch.float32, device=device)

    def log(msg):
        print(f"[L{layer_idx} gpu{gpu}] {msg}", flush=True)

    out_path = os.path.join(OUT_DIR, f"layer_{layer_idx}.pt")
    if os.path.exists(out_path):
        log("already done, skip")
        return layer_idx, None

    reader = ShardReader(NVFP4_DIR)
    n_exp, inter, hidden = 256, 2048, 6144

    log("dequantizing experts")
    w13 = torch.empty(
        n_exp, 2 * inter, hidden, dtype=torch.bfloat16, device=device
    )
    w2 = torch.empty(n_exp, hidden, inter, dtype=torch.bfloat16, device=device)
    for ei in range(n_exp):
        p = f"model.layers.{layer_idx}.mlp.experts.{ei}"
        w13[ei, :inter] = dequant_nvfp4(reader, f"{p}.gate_proj", device, lut)
        w13[ei, inter:] = dequant_nvfp4(reader, f"{p}.up_proj", device, lut)
        w2[ei] = dequant_nvfp4(reader, f"{p}.down_proj", device, lut)

    log(f"w13 {tuple(w13.shape)} books={books['w13']}")
    w13_out = quantize_matrix(w13, books["w13"], gen, log)
    del w13
    torch.cuda.empty_cache()
    log(f"w2 {tuple(w2.shape)} books={books['w2']}")
    w2_out = quantize_matrix(w2, books["w2"], gen, log)
    del w2
    torch.cuda.empty_cache()

    torch.save(
        {
            "layer": layer_idx,
            "books": books,
            "w13_codes": w13_out["codes"],
            "w13_codebooks": w13_out["codebooks"],
            "w13_scales": w13_out["scales"],
            "w13_rel_mse": w13_out["rel_mse"],
            "w2_codes": w2_out["codes"],
            "w2_codebooks": w2_out["codebooks"],
            "w2_scales": w2_out["scales"],
            "w2_rel_mse": w2_out["rel_mse"],
        },
        out_path,
    )
    log(f"saved ({time.time()-t0:.0f}s)")
    return layer_idx, (w13_out["rel_mse"], w2_out["rel_mse"])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--layers", default=None)
    ap.add_argument("--gpus", default="0,1,2,3,4,5,6,7")
    ap.add_argument("--out", default=OUT_DIR)
    args = ap.parse_args()

    globals()["OUT_DIR"] = args.out
    os.makedirs(args.out, exist_ok=True)
    plan = load_plan()
    layers = sorted(plan)
    if args.layers:
        want = {int(x) for x in args.layers.split(",")}
        layers = [li for li in layers if li in want]
    gpus = [int(g) for g in args.gpus.split(",")]

    print(f"{len(layers)} layers over {len(gpus)} GPUs", flush=True)
    jobs = [(li, plan[li], gpus[i % len(gpus)]) for i, li in enumerate(layers)]
    with ProcessPoolExecutor(max_workers=len(gpus)) as ex:
        futs = [ex.submit(process_layer, *j) for j in jobs]
        for f in futs:
            li, rel = f.result()
            if rel:
                print(
                    f"== layer {li}: rel_mse w13={rel[0]:.4f} w2={rel[1]:.4f}",
                    flush=True,
                )


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""AQLM quantization of GLM-5.2 routed experts from the NVFP4 checkpoint.

	For each AQLM-assigned layer (see hybrid_plan.json):
	1. Load the layer's 256 experts' NVFP4 tensors from the downloaded repo,
	dequantize to bf16 on GPU (u8-packed fp4 * fp8 block scale * fp32
	global scale).
	2. Per projection (w13 = gate rows then up rows; w2 = down):
	a. per-output-channel scale s[e,o] (mean abs of the row);
	b. fit one 65536 x 8 codebook per book with GPU k-means on a row
	sample (codebook shared across all 256 experts);
	c. encode all groups-of-8 by chunked nearest-centroid search
	(books=2 encodes the residual against a second codebook);
	d. least-squares refit of the per-channel scales.
	3. Save codes/codebooks/scales + reconstruction stats per layer.

	Everything is chunked: peak GPU memory is the bf16 expert tensor plus a
	few GB of scratch, so 8 workers (one per GPU) run comfortably.

	Usage: aqlm_quantize.py [--layers 11,12] [--gpus 0,1,...] [--out DIR]
	"""

	import argparse
	import json
	import os
	import time
	from concurrent.futures import ProcessPoolExecutor

	ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
	NVFP4_DIR = os.environ.get("GLM52_NVFP4_DIR", "/tmp/glm52-dl/nvfp4-full")
	OUT_DIR = "/data/glm52-aqlm-parts"
	ENTRIES = 65536
	GDIM = 8
	SAMPLE_VECS = 4_000_000
	KMEANS_ITERS = 12
	CHUNK_VECS = 32_768 # scores are [CHUNK_VECS, 65536] fp16 (~4.3GB)

	FP4_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
	-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]


	def load_plan():
	plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))

	def books(tier):
	t = plan["aqlm"][tier]
	return {"w13": t["w13_books"], "w2": t["w2_books"]}

	layers = {}
	for li in plan["aqlm_mixed_layers"]:
	layers[li] = books("mixed")
	for li in plan["aqlm_cold_layers"]:
	layers[li] = books("cold")
	return layers


	class ShardReader:
	"""Random access to tensors across the downloaded safetensors shards."""

	def __init__(self, repo_dir: str):
	idx = json.load(
	open(os.path.join(repo_dir, "model.safetensors.index.json"))
	)
	self.weight_map = idx["weight_map"]
	self.repo_dir = repo_dir
	self._open = {}

	def get(self, name: str):
	from safetensors import safe_open

	shard = self.weight_map[name]
	if shard not in self._open:
	self._open[shard] = safe_open(
	os.path.join(self.repo_dir, shard), framework="pt"
	)
	return self._open[shard].get_tensor(name)


	def dequant_nvfp4(reader, prefix: str, device, lut) -> "torch.Tensor":
	"""Dequantize one NVFP4 linear weight to bf16 [out, in]."""
	import torch

	packed = reader.get(f"{prefix}.weight").to(device) # u8 [out, in/2]
	wscale = reader.get(f"{prefix}.weight_scale").to(device) # fp8 [out, in/16]
	wscale2 = reader.get(f"{prefix}.weight_scale_2").to(device) # fp32 scalar

	lo = (packed & 0x0F).long()
	hi = (packed >> 4).long()
	vals = torch.empty(
	packed.shape[0], packed.shape[1] * 2, dtype=torch.float32, device=device
	)
	# NVFP4 packs the first element of each pair in the low nibble.
	vals[:, 0::2] = lut[lo]
	vals[:, 1::2] = lut[hi]
	del lo, hi
	scale = wscale.to(torch.float32).repeat_interleave(16, dim=1)
	return (vals * scale * wscale2.to(torch.float32)).to(torch.bfloat16)


	def _argmin_codes(v16, cent16, cn16):
	"""v16 [n,8] fp16, cent16 [ENTRIES,8] fp16 -> nearest ids int32.
	argmax of 2*v.c - \|c\|^2 (\|v\|^2 constant per vector)."""
	import torch

	d = v16 @ cent16.t()
	d.mul_(2).sub_(cn16)
	return d.argmax(-1).to(torch.int32)


	def encode_chunked(vecs16, cent16):
	import torch

	cn16 = (cent16.float() * cent16.float()).sum(-1).half()
	out = torch.empty(
	vecs16.shape[0], dtype=torch.int32, device=vecs16.device
	)
	for s in range(0, vecs16.shape[0], CHUNK_VECS):
	out[s : s + CHUNK_VECS] = _argmin_codes(
	vecs16[s : s + CHUNK_VECS], cent16, cn16
	)
	return out


	def kmeans_fit(vecs16, iters: int, generator):
	"""Lloyd's k-means on [N,8] fp16; returns [ENTRIES,8] fp16 centroids."""
	import torch

	n = vecs16.shape[0]
	device = vecs16.device
	perm = torch.randperm(n, generator=generator, device=device)[:ENTRIES]
	cent = vecs16[perm].float()
	for _ in range(iters):
	cent16 = cent.half()
	cn16 = (cent * cent).sum(-1).half()
	sums = torch.zeros(ENTRIES, GDIM, device=device)
	cnts = torch.zeros(ENTRIES, device=device)
	ones = None
	for s in range(0, n, CHUNK_VECS):
	v = vecs16[s : s + CHUNK_VECS]
	a = _argmin_codes(v, cent16, cn16).long()
	sums.index_add_(0, a, v.float())
	if ones is None or ones.shape[0] != a.shape[0]:
	ones = torch.ones(a.shape[0], device=device)
	cnts.index_add_(0, a, ones[: a.shape[0]])
	mask = cnts > 0
	cent[mask] = sums[mask] / cnts[mask].unsqueeze(-1)
	dead = int((~mask).sum())
	if dead:
	ridx = torch.randperm(n, generator=generator, device=device)[:dead]
	cent[~mask] = vecs16[ridx].float()
	return cent.half()


	def quantize_matrix(w, books: int, generator, log):
	"""w: [E, out, in] bf16 (GPU) -> dict with codes int16 [E,books,out,in/8],
	codebooks fp16 [books,ENTRIES,8], scales fp16 [E,out], rel mse."""
	import torch

	e, out, k = w.shape
	device = w.device
	rows = e * out
	k8 = k // GDIM
	wv = w.reshape(rows, k)

	# per-row scales, chunked
	scales = torch.empty(rows, dtype=torch.float32, device=device)
	ROWCH = max(1, (CHUNK_VECS * 64) // k)
	for s in range(0, rows, ROWCH):
	scales[s : s + ROWCH] = (
	wv[s : s + ROWCH].float().abs().mean(-1).clamp_min(1e-8)
	)

	# --- fit codebooks on a row sample ---
	n_sample_rows = min(rows, max(1, SAMPLE_VECS // k8))
	ridx = torch.randperm(rows, generator=generator, device=device)[
	:n_sample_rows
	]
	samp = (
	(wv[ridx].float() / scales[ridx].unsqueeze(-1))
	.reshape(-1, GDIM)
	.half()
	)
	log(f" kmeans book0 on {samp.shape[0]} vecs")
	cbs = [kmeans_fit(samp, KMEANS_ITERS, generator)]
	if books == 2:
	c0 = encode_chunked(samp, cbs[0])
	resid = (samp.float() - cbs[0].float()[c0.long()]).half()
	log(f" kmeans book1 on residuals")
	cbs.append(kmeans_fit(resid, KMEANS_ITERS, generator))
	del c0, resid
	del samp
	torch.cuda.empty_cache()

	cb_f = [c.float() for c in cbs]
	cn16 = [(cf * cf).sum(-1).half() for cf in cb_f]

	# --- full encode, chunked by rows; accumulate scale-refit stats ---
	codes = torch.empty(
	(rows, books, k8), dtype=torch.int16, device=device
	)
	num = torch.zeros(rows, dtype=torch.float32, device=device)
	den = torch.zeros(rows, dtype=torch.float32, device=device)
	sse = 0.0
	stot = 0.0
	ROWCH = max(1, CHUNK_VECS // k8)
	for s in range(0, rows, ROWCH):
	wc = wv[s : s + ROWCH].float() / scales[s : s + ROWCH].unsqueeze(-1)
	v = wc.reshape(-1, GDIM).half()
	approx = torch.zeros(v.shape[0], GDIM, device=device)
	r = v.float()
	for b in range(books):
	cb = _argmin_codes(r.half(), cbs[b], cn16[b]).long()
	dec = cb_f[b][cb]
	approx += dec
	r -= dec
	codes[s : s + ROWCH, b] = (
	torch.where(cb >= 32768, cb - 65536, cb)
	.to(torch.int16)
	.reshape(-1, k8)
	)
	nrow = wc.shape[0]
	a2 = approx.reshape(nrow, k)
	w2 = wc
	num[s : s + ROWCH] = (w2 * a2).sum(-1)
	den[s : s + ROWCH] = (a2 * a2).sum(-1).clamp_min(1e-8)
	sse += float((r * r).sum())
	stot += float((w2 * w2).sum())
	del wc, v, approx, r, a2

	rel = sse / max(stot, 1e-9)
	scales = scales * (num / den).clamp(0.5, 2.0)
	log(f" rel_mse={rel:.4f}")

	return {
	"codes": codes.reshape(e, out, books, k8)
	.permute(0, 2, 1, 3)
	.contiguous()
	.cpu(),
	"codebooks": torch.stack(cbs).cpu(),
	"scales": scales.reshape(e, out).half().cpu(),
	"rel_mse": rel,
	}


	def process_layer(layer_idx: int, books: dict, gpu: int):
	import torch

	t0 = time.time()
	device = f"cuda:{gpu}"
	torch.cuda.set_device(gpu)
	gen = torch.Generator(device=device)
	gen.manual_seed(1234 + layer_idx)
	lut = torch.tensor(FP4_LUT, dtype=torch.float32, device=device)

	def log(msg):
	print(f"[L{layer_idx} gpu{gpu}] {msg}", flush=True)

	out_path = os.path.join(OUT_DIR, f"layer_{layer_idx}.pt")
	if os.path.exists(out_path):
	log("already done, skip")
	return layer_idx, None

	reader = ShardReader(NVFP4_DIR)
	n_exp, inter, hidden = 256, 2048, 6144

	log("dequantizing experts")
	w13 = torch.empty(
	n_exp, 2 * inter, hidden, dtype=torch.bfloat16, device=device
	)
	w2 = torch.empty(n_exp, hidden, inter, dtype=torch.bfloat16, device=device)
	for ei in range(n_exp):
	p = f"model.layers.{layer_idx}.mlp.experts.{ei}"
	w13[ei, :inter] = dequant_nvfp4(reader, f"{p}.gate_proj", device, lut)
	w13[ei, inter:] = dequant_nvfp4(reader, f"{p}.up_proj", device, lut)
	w2[ei] = dequant_nvfp4(reader, f"{p}.down_proj", device, lut)

	log(f"w13 {tuple(w13.shape)} books={books['w13']}")
	w13_out = quantize_matrix(w13, books["w13"], gen, log)
	del w13
	torch.cuda.empty_cache()
	log(f"w2 {tuple(w2.shape)} books={books['w2']}")
	w2_out = quantize_matrix(w2, books["w2"], gen, log)
	del w2
	torch.cuda.empty_cache()

	torch.save(
	{
	"layer": layer_idx,
	"books": books,
	"w13_codes": w13_out["codes"],
	"w13_codebooks": w13_out["codebooks"],
	"w13_scales": w13_out["scales"],
	"w13_rel_mse": w13_out["rel_mse"],
	"w2_codes": w2_out["codes"],
	"w2_codebooks": w2_out["codebooks"],
	"w2_scales": w2_out["scales"],
	"w2_rel_mse": w2_out["rel_mse"],
	},
	out_path,
	)
	log(f"saved ({time.time()-t0:.0f}s)")
	return layer_idx, (w13_out["rel_mse"], w2_out["rel_mse"])


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--layers", default=None)
	ap.add_argument("--gpus", default="0,1,2,3,4,5,6,7")
	ap.add_argument("--out", default=OUT_DIR)
	args = ap.parse_args()

	globals()["OUT_DIR"] = args.out
	os.makedirs(args.out, exist_ok=True)
	plan = load_plan()
	layers = sorted(plan)
	if args.layers:
	want = {int(x) for x in args.layers.split(",")}
	layers = [li for li in layers if li in want]
	gpus = [int(g) for g in args.gpus.split(",")]

	print(f"{len(layers)} layers over {len(gpus)} GPUs", flush=True)
	jobs = [(li, plan[li], gpus[i % len(gpus)]) for i, li in enumerate(layers)]
	with ProcessPoolExecutor(max_workers=len(gpus)) as ex:
	futs = [ex.submit(process_layer, *j) for j in jobs]
	for f in futs:
	li, rel = f.result()
	if rel:
	print(
	f"== layer {li}: rel_mse w13={rel[0]:.4f} w2={rel[1]:.4f}",
	flush=True,
	)


	if __name__ == "__main__":
	main()