quantization/quantize_motif_fp8.py

#!/usr/bin/env python3
import argparse
import json
import math
import os
import shutil
import time

import torch
from safetensors import safe_open
from safetensors.torch import save_file

BLOCK = 128
FP8_MAX = 448.0
EXPERT_CHUNK = 96
SOURCE_COMMIT = "1695f5aa6d97cccc8623a03f93c9c7d5fbb14d45"  # Motif-Technologies/Motif-3 commit used in this environment
TARGET_SUFFIXES = (".moe.experts.gate_up_proj", ".moe.experts.down_proj")


def is_quant_target(name, tensor_shape, tensor_dtype):
    if not name.startswith("model.layers."):
        return False
    if not name.endswith(TARGET_SUFFIXES):
        return False
    if len(tensor_shape) != 3:
        return False
    if tensor_dtype not in (torch.bfloat16, torch.float32, torch.float16):
        return False
    assert tensor_shape[1] % BLOCK == 0 and tensor_shape[2] % BLOCK == 0, \
        f"{name}: {tensor_shape} not divisible by {BLOCK}"
    return True


@torch.no_grad()
def quantize_block_fp8(w, device):
    E, N, K = w.shape
    nb, kb = N // BLOCK, K // BLOCK
    q_out = torch.empty((E, N, K), dtype=torch.float8_e4m3fn)
    s_out = torch.empty((E, nb, kb), dtype=torch.float32)

    sq_err_sum = 0.0
    sq_ref_sum = 0.0
    dot_sum = 0.0
    qq_sum = 0.0
    max_abs_err = 0.0
    nan_inf = 0

    for i in range(0, E, EXPERT_CHUNK):
        chunk = w[i : i + EXPERT_CHUNK].to(device, non_blocking=True).float()
        e = chunk.shape[0]
        blocks = chunk.reshape(e, nb, BLOCK, kb, BLOCK)
        amax = blocks.abs().amax(dim=(2, 4))
        scale_inv = (amax / FP8_MAX).clamp(min=1e-12)
        q = (blocks / scale_inv[:, :, None, :, None]).clamp(-FP8_MAX, FP8_MAX)
        q_fp8 = q.to(torch.float8_e4m3fn)

        deq = q_fp8.float() * scale_inv[:, :, None, :, None]
        err = deq - blocks
        sq_err_sum += err.pow(2).sum().item()
        sq_ref_sum += blocks.pow(2).sum().item()
        dot_sum += (deq * blocks).sum().item()
        qq_sum += deq.pow(2).sum().item()
        max_abs_err = max(max_abs_err, err.abs().max().item())
        nan_inf += int(torch.isnan(deq).sum().item() + torch.isinf(deq).sum().item())
        nan_inf += int(torch.isnan(scale_inv).sum().item() + torch.isinf(scale_inv).sum().item())

        q_out[i : i + e] = q_fp8.reshape(e, N, K).cpu()
        s_out[i : i + e] = scale_inv.cpu()
        del chunk, blocks, amax, scale_inv, q, q_fp8, deq, err

    stats = {
        "rel_rmse": math.sqrt(sq_err_sum / max(sq_ref_sum, 1e-30)),
        "cosine": dot_sum / max(math.sqrt(sq_ref_sum) * math.sqrt(qq_sum), 1e-30),
        "max_abs_err": max_abs_err,
        "nan_inf": nan_inf,
    }
    return q_out, s_out, stats


def process_shard(src, dst, shard_name, device, manifest_dir):
    mpath = os.path.join(manifest_dir, shard_name + ".json")
    opath = os.path.join(dst, shard_name)
    if os.path.exists(mpath) and os.path.exists(opath):
        with open(mpath) as f:
            return json.load(f)

    t0 = time.time()
    out_tensors = {}
    entry = {"shard": shard_name, "tensors": {}, "quantized": [], "stats": {}}

    with safe_open(os.path.join(src, shard_name), framework="pt", device="cpu") as f:
        for name in f.keys():
            t = f.get_tensor(name)
            if is_quant_target(name, tuple(t.shape), t.dtype):
                q, s, stats = quantize_block_fp8(t, device)
                out_tensors[name] = q
                out_tensors[name + "_scale_inv"] = s
                entry["quantized"].append(name)
                entry["stats"][name] = stats
                entry["tensors"][name] = {"dtype": "F8_E4M3", "shape": list(q.shape)}
                entry["tensors"][name + "_scale_inv"] = {"dtype": "F32", "shape": list(s.shape)}
            else:
                out_tensors[name] = t
                entry["tensors"][name] = {"dtype": str(t.dtype).replace("torch.", ""), "shape": list(t.shape)}

    tmp = opath + ".tmp"
    save_file(out_tensors, tmp, metadata={"format": "pt"})
    os.replace(tmp, opath)
    entry["seconds"] = round(time.time() - t0, 2)
    entry["out_bytes"] = os.path.getsize(opath)

    os.makedirs(manifest_dir, exist_ok=True)
    with open(mpath, "w") as f:
        json.dump(entry, f)
    return entry


def finalize(src, dst, entries):
    with open(os.path.join(src, "model.safetensors.index.json")) as f:
        old_index = json.load(f)
    weight_map = dict(old_index["weight_map"])
    for e in entries:
        for qname in e["quantized"]:
            weight_map[qname + "_scale_inv"] = e["shard"]
    total = sum(e["out_bytes"] for e in entries)
    with open(os.path.join(dst, "model.safetensors.index.json"), "w") as f:
        json.dump({"metadata": {"total_size": total}, "weight_map": weight_map}, f, indent=2)

    with open(os.path.join(src, "config.json")) as f:
        config = json.load(f)
    config["quantization_config"] = {
        "quant_method": "fp8",
        "fmt": "e4m3",
        "activation_scheme": "dynamic",
        "weight_block_size": [BLOCK, BLOCK],
        "modules_to_not_convert": [
            "model.embed_tokens", "lm_head",
            "model.layers.*.self_attn", "model.layers.*.mlp",
            "model.layers.*.input_layernorm", "model.layers.*.post_attention_layernorm",
            "model.layers.*.mhc_attn", "model.layers.*.mhc_ffn",
            "model.layers.*.moe.router", "model.layers.*.moe.shared_experts",
            "model.layers.*.moe.experts.act_fn",
            "model.mtp_layers.*",
        ],
    }
    with open(os.path.join(dst, "config.json"), "w") as f:
        json.dump(config, f, indent=2)

    for fn in os.listdir(src):
        if fn.endswith((".py", ".jinja", ".txt", ".md")) or fn in (
            "tokenizer.json", "tokenizer_config.json", "generation_config.json",
            "special_tokens_map.json", "chat_template.json",
        ):
            src_f = os.path.join(src, fn)
            if os.path.isfile(src_f):
                shutil.copy2(src_f, os.path.join(dst, fn))

    n_quant_params = sum(
        int(torch.tensor(e["tensors"][q]["shape"]).prod())
        for e in entries for q in e["quantized"]
    )
    all_stats = [s for e in entries for s in e["stats"].values()]
    summary = {
        "source_model": "Motif-Technologies/Motif-3",
        "source_commit": SOURCE_COMMIT,
        "recipe": {
            "method": "expert-only FP8 E4M3, 128x128 block, W8A16",
            "targets": list(TARGET_SUFFIXES),
            "scale_naming": "<param>_scale_inv (dequant = fp8_block * scale_inv)",
        },
        "quantized_tensor_count": sum(len(e["quantized"]) for e in entries),
        "quantized_params": n_quant_params,
        "output_size_bytes": total,
        "nan_inf_total": sum(s["nan_inf"] for s in all_stats),
        "rel_rmse": {
            "mean": sum(s["rel_rmse"] for s in all_stats) / max(len(all_stats), 1),
            "max": max((s["rel_rmse"] for s in all_stats), default=0),
        },
        "cosine": {
            "mean": sum(s["cosine"] for s in all_stats) / max(len(all_stats), 1),
            "min": min((s["cosine"] for s in all_stats), default=0),
        },
        "max_abs_err": max((s["max_abs_err"] for s in all_stats), default=0),
    }
    with open(os.path.join(dst, "quantization_recipe.json"), "w") as f:
        json.dump(summary, f, indent=2)
    return summary


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default="./Motif-3-Beta")
    ap.add_argument("--dst", default="./Motif-3-Beta-FP8")
    ap.add_argument("--only")
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--skip-finalize", action="store_true")
    args = ap.parse_args()

    os.makedirs(args.dst, exist_ok=True)
    manifest_dir = os.path.join(args.dst, ".manifest")
    os.makedirs(manifest_dir, exist_ok=True)

    with open(os.path.join(args.src, "model.safetensors.index.json")) as f:
        shards = sorted(set(json.load(f)["weight_map"].values()))
    if args.only:
        shards = [args.only]

    print(f"[start] {len(shards)} shards, device={args.device}", flush=True)
    entries = []
    t0 = time.time()
    for i, shard in enumerate(shards):
        e = process_shard(args.src, args.dst, shard, args.device, manifest_dir)
        entries.append(e)
        done_q = sum(len(x["quantized"]) for x in entries)
        print(f"[{i+1}/{len(shards)}] {shard} quantized={len(e['quantized'])} "
              f"({e.get('seconds', 0)}s) total_q={done_q}", flush=True)

    if not args.skip_finalize and not args.only:
        summary = finalize(args.src, args.dst, entries)
        print(json.dumps(summary, indent=2), flush=True)
    print(f"[done] {(time.time()-t0)/60:.1f} min", flush=True)


if __name__ == "__main__":
    main()