| |
| """ |
| 🌸 splice_nvfp4_model.py - Splicer for Calibrated Experts-Only NVFP4 Drafter |
| Channels Nikola's Sparkly Village Charm for Haga's Seeker Quest! ✨🚀 |
| |
| Non-destructively links base model NVFP4 assets, updates config.json with |
| mtp_quantization_config, and patches model.safetensors.index.json for the |
| calibrated NVFP4 MTP drafter. |
| """ |
|
|
| import json |
| import os |
| import shutil |
| import torch |
| import safetensors.torch |
| from safetensors import safe_open |
|
|
|
|
| def splice_nvfp4( |
| base_dir: str = "models/mazinb/Qwen3.8-Flash-Next-Uncensored-NVFP4", |
| shadow_dir: str = "models/spliced/Qwen3.8-Flash-Next-MTP-NVFP4", |
| drafter_file: str = "models/spliced/Qwen3.8-Flash-Next-MTP-NVFP4/model_mtp_tuned.safetensors", |
| ): |
| base_dir = os.path.abspath(base_dir) |
| shadow_dir = os.path.abspath(shadow_dir) |
| drafter_file = os.path.abspath(drafter_file) |
|
|
| print("=" * 75) |
| print("🌸 Splicing Calibrated NVFP4 MTP Model ✨🚀") |
| print(f"Base Model: {base_dir}") |
| print(f"Shadow Dir: {shadow_dir}") |
| print(f"Drafter: {drafter_file}") |
| print("=" * 75) |
|
|
| assert os.path.exists(drafter_file), f"Drafter file {drafter_file} does not exist!" |
|
|
| |
| skip_files = {"model.safetensors.index.json", "config.json", "model_mtp_tuned.safetensors"} |
| linked = 0 |
| for fname in os.listdir(base_dir): |
| if fname in skip_files: |
| continue |
| src = os.path.join(base_dir, fname) |
| dst = os.path.join(shadow_dir, fname) |
| if os.path.exists(dst) or os.path.islink(dst): |
| os.remove(dst) |
| if os.path.isdir(src): |
| os.symlink(src, dst) |
| else: |
| os.link(src, dst) |
| linked += 1 |
| print(f"🔗 Linked {linked} base model assets via hardlinks (0 bytes disk space used).") |
|
|
| |
| base_cfg_path = os.path.join(base_dir, "config.json") |
| with open(base_cfg_path, "r") as f: |
| cfg = json.load(f) |
|
|
| mtp_quant_config = { |
| "quant_method": "modelopt", |
| "quant_algo": "NVFP4", |
| "config_groups": { |
| "group_0": { |
| "input_activations": { |
| "dynamic": False, |
| "group_size": 16, |
| "num_bits": 4, |
| "type": "float", |
| }, |
| "targets": ["Linear"], |
| "weights": { |
| "dynamic": False, |
| "group_size": 16, |
| "num_bits": 4, |
| "type": "float", |
| }, |
| } |
| }, |
| "ignore": [ |
| "self_attn", |
| "fc_embedding", |
| "fc_hidden", |
| "shared_expert", |
| "mlp.gate", |
| "shared_expert_gate", |
| "hyper_connection", |
| "norm", |
| "embed_tokens", |
| "lm_head", |
| ], |
| "producer": {"name": "modelopt", "version": "0.46.0-compatible-rtn"}, |
| } |
|
|
| cfg["mtp_quantization_config"] = mtp_quant_config |
| if "text_config" in cfg and isinstance(cfg["text_config"], dict): |
| cfg["text_config"]["mtp_quantization_config"] = mtp_quant_config |
|
|
| shadow_cfg_path = os.path.join(shadow_dir, "config.json") |
| if os.path.exists(shadow_cfg_path): |
| os.remove(shadow_cfg_path) |
| with open(shadow_cfg_path, "w") as f: |
| json.dump(cfg, f, indent=2) |
| print("📝 Wrote updated config.json with mtp_quantization_config.") |
|
|
| |
| q_key = "mtp.layers.0.self_attn.q_proj.weight" |
| drafter = safetensors.torch.load_file(drafter_file) |
| if q_key in drafter and drafter[q_key].shape[0] == 6144: |
| orig_sf11 = os.path.join(base_dir, "model-bf16-00011.safetensors") |
| with safe_open(orig_sf11, framework="pt", device="cpu") as sf: |
| orig_q_w = sf.get_tensor(q_key) |
| orig_q_gate = orig_q_w[6144:, :] |
| drafter[q_key] = torch.cat([drafter[q_key], orig_q_gate], dim=0).clone() |
| safetensors.torch.save_file(drafter, drafter_file) |
| print(f" Packed Q projection with original output gate -> shape: {list(drafter[q_key].shape)}") |
|
|
| |
| base_index_path = os.path.join(base_dir, "model.safetensors.index.json") |
| with open(base_index_path, "r") as f: |
| index_data = json.load(f) |
|
|
| drafter_keys = list(drafter.keys()) |
| print(f"📦 Read {len(drafter_keys)} keys from {os.path.basename(drafter_file)}.") |
|
|
| weight_map = index_data.get("weight_map", {}) |
| |
| old_mtp_keys = [k for k in weight_map.keys() if k.startswith("mtp.")] |
| for k in old_mtp_keys: |
| del weight_map[k] |
|
|
| |
| for k in drafter_keys: |
| weight_map[k] = "model_mtp_tuned.safetensors" |
|
|
| shadow_index_path = os.path.join(shadow_dir, "model.safetensors.index.json") |
| if os.path.exists(shadow_index_path): |
| os.remove(shadow_index_path) |
| with open(shadow_index_path, "w") as f: |
| json.dump(index_data, f, indent=2) |
| print(f"🗺️ Patched model.safetensors.index.json: mapped {len(drafter_keys)} MTP keys.") |
|
|
| print("\n✨ Calibrated NVFP4 Spliced Model Successfully Created! ✨🚀") |
| print(f"Location: {shadow_dir}") |
| print("=" * 75) |
|
|
|
|
| if __name__ == "__main__": |
| splice_nvfp4() |
|
|
