File size: 6,018 Bytes
8045270 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | #!/usr/bin/env python3
"""Export an NVFP4 draft (W4A16: quantized routed experts, BF16 the rest)
from the BF16 GLM-5.3-Flash mtp/ folder.
Uses ModelOpt's own NVFP4QTensor.quantize (same math as the target
checkpoint; packing: byte k = (even nibble | odd nibble<<4), one-level
e4m3 block scales, scale_2 = 1.0). Declares input_activations: null so
vLLM routes the experts through its W4A16 NVFP4 path (no input_scale
needed, mirroring the target's expert behavior).
Output: ~/github/GLM5.3-Flash-E224-DGX-Spark/mtp-nvfp4/
"""
import json
import shutil
import struct
import sys
import time
from pathlib import Path
import torch
from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor
from safetensors.torch import save_file
torch.set_grad_enabled(False)
# Usage: python3 export_mtp_nvfp4.py /path/to/GLM5.3-Flash-E224-DGX-Spark
ROOT = (Path(sys.argv[1]).expanduser() if len(sys.argv) > 1
else Path.home() / "github/GLM5.3-Flash-E224-DGX-Spark")
SRC = ROOT / "mtp"
DST = ROOT / "mtp-nvfp4"
def read_shard(p: Path):
with open(p, "rb") as f:
n = struct.unpack("<Q", f.read(8))[0]
hdr = json.loads(f.read(n))
base = f.tell()
blobs = {}
for k, v in hdr.items():
if k == "__metadata__":
continue
s, e = v["data_offsets"]
f.seek(base + s)
blobs[k] = f.read(e - s)
return hdr, blobs
def main():
DST.mkdir(exist_ok=True)
for f in SRC.iterdir():
if f.suffix in (".json", ".jinja") and f.name not in (
"hf_quant_config.json",
"model.safetensors.index.json",
):
shutil.copy(f, DST / f.name)
out_tensors = {}
n_quant = 0
n_bf16 = 0
t0 = time.time()
for si in range(1, 5):
p = SRC / f"model-{si:05d}-of-00004.safetensors"
hdr, blobs = read_shard(p)
for name, blob in blobs.items():
short = name.replace("model.language_model.layers.45.", "")
quantize_this = (
short.startswith("mlp.experts.")
or short.startswith("mlp.shared_experts.")
) and short.endswith(("gate_proj.weight", "up_proj.weight", "down_proj.weight"))
if quantize_this:
shape = hdr[name]["shape"]
w = torch.frombuffer(bytearray(blob), dtype=torch.bfloat16).clone().view(shape)
qt, block_scale, scale_2 = NVFP4QTensor.quantize(
w, 16, weights_scaling_factor_2=torch.tensor(1.0), try_tensorrt=False
)
packed = qt.quantized_data if hasattr(qt, "quantized_data") else qt._quantized_data
base = name[: -len(".weight")]
out_tensors[base + ".weight"] = packed.contiguous()
out_tensors[base + ".weight_scale"] = block_scale.contiguous()
out_tensors[base + ".input_scale"] = torch.ones((), dtype=torch.float32)
out_tensors[base + ".weight_scale_2"] = torch.ones((), dtype=torch.float32)
n_quant += 1
if short == "eh_proj.weight":
out_tensors["model.layers.45.eh_proj.input_scale"] = torch.ones((), dtype=torch.float32)
out_tensors["model.layers.45.eh_proj.weight_scale_2"] = torch.ones((), dtype=torch.float32)
if n_quant % 100 == 0:
print(f" {n_quant} expert linears quantized ({time.time()-t0:.0f}s)")
else:
shape = hdr[name]["shape"]
dt = {"BF16": torch.bfloat16, "F32": torch.float32,
"F8_E4M3": torch.float8_e4m3fn, "U8": torch.uint8}[hdr[name]["dtype"]]
w = torch.frombuffer(bytearray(blob), dtype=dt).clone()
out_tensors[name] = w.view(shape) if w.numel() == torch.Size(shape).numel() else w[: torch.Size(shape).numel()].view(shape)
n_bf16 += 1
print(f"quantized: {n_quant} expert linears; kept BF16: {n_bf16} tensors")
save_file(out_tensors, str(DST / "model.safetensors"))
# quant config: W4A16 NVFP4 (input_activations: null)
quant_cfg = {
"producer": {"name": "modelopt-handoff", "version": "0.47.0-direct"},
"quantization": {
"quant_algo": "NVFP4",
"kv_cache_quant_algo": None,
"group_size": 16,
"exclude_modules": [
# keep attention + indexer + norms + embed + head + eh_proj
# in BF16; experts and shared_experts ARE quantized (NVFP4)
"model.layers.45.mtp_block.mlp.gate",
"model.layers.45.mtp_block.self_attn*",
"model.layers.45.mtp_block.input_layernorm*",
"model.layers.45.mtp_block.post_attention_layernorm*",
"model.layers.45.eh_proj",
"model.layers.45.enorm",
"model.layers.45.hnorm",
"model.layers.45.shared_head*",
"model.layers.45.embed_tokens*",
],
"config_groups": {
"group_0": {
"weights": {"num_bits": 4, "type": "float", "dynamic": False, "group_size": 16},
"input_activations": {"dynamic": False, "num_bits": 4, "type": "float", "group_size": 16},
"targets": ["Linear"],
}
},
"quant_method": "modelopt",
},
}
(DST / "hf_quant_config.json").write_text(json.dumps(quant_cfg, indent=2))
# get_quant_config() reads hf_config.quantization_config (config.json),
# not hf_quant_config.json — keep both consistent.
cfg = json.loads((DST / "config.json").read_text())
cfg["quantization_config"] = quant_cfg["quantization"]
(DST / "config.json").write_text(json.dumps(cfg, indent=2, ensure_ascii=False))
print("wrote", DST / "model.safetensors",
f"({sum(t.numel() * t.element_size() for t in out_tensors.values()) / 1e9:.2f} GB)")
if __name__ == "__main__":
main()
|