| |
| """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) |
|
|
| |
| 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_cfg = { |
| "producer": {"name": "modelopt-handoff", "version": "0.47.0-direct"}, |
| "quantization": { |
| "quant_algo": "NVFP4", |
| "kv_cache_quant_algo": None, |
| "group_size": 16, |
| "exclude_modules": [ |
| |
| |
| "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)) |
|
|
| |
| |
| 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() |
|
|
