From 6008cf128d591068a5fe21438f581331e5f0c842 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 7 May 2026 15:15:32 +0000
Subject: [PATCH] Add model_opt_nvfp4_experts_only.py

Quantizes only MoE expert weights to NVFP4, leaving attention untouched.
Includes comments documenting all available NVFP4 strategies.
Copy to model_opt_nvfp4_<strategy>.py for each new strategy.
---
 scripts/model_opt_nvfp4_experts_only.py | 65 +++++++++++++++++++++++++
 1 file changed, 65 insertions(+)
 create mode 100644 scripts/model_opt_nvfp4_experts_only.py

diff --git a/scripts/model_opt_nvfp4_experts_only.py b/scripts/model_opt_nvfp4_experts_only.py
new file mode 100644
index 0000000..70d080a
--- /dev/null
+++ b/scripts/model_opt_nvfp4_experts_only.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+"""
+ModelOpt NVFP4 quantization — experts only.
+
+Quantizes only the MoE expert weights (gate_up_proj, down_proj) to NVFP4,
+leaving attention and shared MLP layers untouched. This avoids issues with
+FP8 attention kernels on Blackwell (DeepGEMM unsupported, Triton finegrained
+FP8 matmul shape mismatches).
+
+Available NVFP4 quantization strategies (from modelopt huggingface_example.sh):
+  - nvfp4               : Full model NVFP4 quantization
+  - nvfp4_experts_only  : Only MoE expert weights (this script)
+  - nvfp4_mlp_only      : Only MLP layers (experts + shared MLP)
+  - nvfp4_omlp_only     : Only output + MLP layers
+  - nvfp4_awq           : NVFP4 with AWQ calibration
+  - nvfp4_mse           : NVFP4 with MSE calibration
+  - w4a8_nvfp4_fp8      : W4A8 NVFP4 weights + FP8 activations
+  - w4a8_mxfp4_fp8      : W4A8 MXFP4 weights + FP8 activations
+  - nvfp4_svdquant      : NVFP4 with SVDQuant
+  - nvfp4_local_hessian : NVFP4 with local Hessian calibration
+
+Strategy: Copy this file to model_opt_nvfp4_<strategy>.py and tweak as needed.
+By the end, we'll have working quantized weights for each successful strategy.
+
+Output dir naming: DeepSeek-V4-Pro_NVFP4-<strategy>_kv_fp8_cast
+"""
+
+import subprocess
+import sys
+import os
+
+# ── Config ──────────────────────────────────────────────────────────────────
+MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro"  # Mixed-precision source
+QUANT = "nvfp4_experts_only"
+TP = 8
+CALIB = 256
+KV_CACHE_QUANT = "fp8_cast"
+EXTRA_FLAGS = "--trust_remote_code --use_seq_device_map"
+
+# Output dir follows modelopt convention: <model>_<quant>_kv_<kv_quant>
+# We override the model name to make the strategy clear
+OUTPUT_NAME = f"DeepSeek-V4-Pro_NVFP4-{QUANT}_kv_{KV_CACHE_QUANT}"
+
+SCRIPT_DIR = "/root/nvidia-meeting/modelopt-repo/examples/llm_ptq"
+LOG_FILE = f"/root/nvidia-meeting/modelopt_{QUANT}.log"
+
+# ── Run ─────────────────────────────────────────────────────────────────────
+cmd = f"""cd {SCRIPT_DIR} && \\
+source /root/nvidia-meeting/venv/bin/activate && \\
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \\
+bash scripts/huggingface_example.sh \\
+    --model {MODEL} \\
+    --quant {QUANT} \\
+    --tp {TP} \\
+    --calib {CALIB} \\
+    --kv_cache_quant {KV_CACHE_QUANT} \\
+    {EXTRA_FLAGS} 2>&1 | tee {LOG_FILE}"""
+
+print(f"Running: {QUANT} quantization on {MODEL}")
+print(f"Output: {OUTPUT_NAME}")
+print(f"Log: {LOG_FILE}")
+print(f"Command:\n{cmd}\n")
+
+ret = subprocess.call(cmd, shell=True)
+sys.exit(ret)