#!/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": "_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()