scripts/quantize_veriloop.py

"""NVFP4 W4A4 quantization for VeriLoop-E2 (Qwen3.8-27B based).

Uses NVIDIA Model Optimizer's canonical recipe
    (NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG: local Hessian + fp8 scale sweep,
    static weight scales + dynamic input scales) with linear_attn blocks and
    self-attention projections kept in BF16, matching validated NVFP4
    releases for this architecture family.

Calibration: nvidia/Nemotron-Competitive-Programming-v1 (streaming), defaults
512 samples x 512 tokens (262144 tokens total).

Requirements:
  - 3+ CUDA GPUs holding ~55 GB total for a 27B BF16 model
    (tuned on a 17/34/17 GB layout; adjust --layers-split otherwise).
  - A C compiler for the Triton JIT (on Windows: run inside a VS
    Native Tools prompt, i.e. vcvars64, with CC pointing at cl.exe).
  - Enough RAM to hold the model for the CPU-side export (~80 GB for 27B).

Usage:
  python quantize_veriloop.py --model ./model-bf16 --output ./model-nvfp4
"""
import argparse
import os
import sys


def parse_args():
    p = argparse.ArgumentParser(description="VeriLoop-E2 -> NVFP4 quantization")
    p.add_argument("--model", default="./model-bf16",
                   help="Source BF16 HuggingFace model dir (default: ./model-bf16)")
    p.add_argument("--output", default="./model-nvfp4",
                   help="Output dir for the NVFP4 checkpoint (default: ./model-nvfp4)")
    p.add_argument("--calib-size", type=int, default=512,
                   help="Calibration samples (default: 512)")
    p.add_argument("--calib-seq-len", type=int, default=512,
                   help="Calibration sequence length (default: 512)")
    p.add_argument("--layers-split", default="16,32,16",
                   help="Comma-separated layer counts per visible GPU, must sum to 64 "
                        "(default tuned for a 17/34/17 GB VRAM layout: 16,32,16)")
    p.add_argument("--gpu-order", default=None,
                   help="Optional CUDA_VISIBLE_DEVICES value, e.g. '2,0,1' to make a "
                        "specific physical GPU cuda:0 (default: natural order)")
    return p.parse_args()


ARGS = parse_args()

if ARGS.gpu_order:
    os.environ["CUDA_VISIBLE_DEVICES"] = ARGS.gpu_order
if os.name == "nt":
    # Triton JIT needs a C compiler; run from a VS Native Tools prompt.
    os.environ.setdefault("CC", "cl")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import time
import copy
import torch
from pathlib import Path

from transformers import AutoModelForCausalLM, AutoTokenizer
import modelopt.torch.quantization as mtq

MODEL_PATH = ARGS.model
OUTPUT_PATH = ARGS.output
CALIB_SPLIT = "competitive_coding_python_part00"
CALIB_SIZE = ARGS.calib_size
CALIB_SEQ_LEN = ARGS.calib_seq_len

_split = [int(x) for x in ARGS.layers_split.split(",")]
assert sum(_split) == 64, "--layers-split must sum to 64"
assert len(_split) == torch.cuda.device_count(), \
    "--layers-split must have one entry per visible GPU"


def build_device_map():
    # Explicit map: all 64 transformer layers (~48.7 GB) must live on CUDA,
    # because the dynamic block quantizer requires CUDA amax tensors.
    # Modules with disabled quantizers (embed/head/norm) stay on CPU.
    bounds, acc = [], 0
    for count in _split:
        bounds.append((acc, acc + count))
        acc += count
    dm = {
        "model.embed_tokens": "cpu",
        "model.norm": "cpu",
        "model.rotary_emb": "cpu",
        "lm_head": "cpu",
    }
    for i in range(64):
        for dev, (lo, hi) in enumerate(bounds):
            if lo <= i < hi:
                dm[f"model.layers.{i}"] = dev
                break
    return dm


def build_quant_cfg():
    # NVIDIA canonical recipe + granularity adjustments: linear_attn (GDN)
    # fully BF16 plus BF16 self-attention, matching validated NVFP4 releases
    # for this architecture family (MLP-only NVFP4). NVFP4 attention/GDN
    # weights produce degenerate output on some stacks; MLP-only is the
    # widely-deployed pattern (conv1d/in_proj_a/in_proj_b already disabled
    # in the base recipe). Appended last: entries apply in list order,
    # later overrides earlier.
    cfg = copy.deepcopy(mtq.NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG)
    for name in ["*linear_attn.in_proj_qkv*", "*linear_attn.in_proj_z*",
                 "*linear_attn.out_proj*",
                 "*self_attn.q_proj*", "*self_attn.k_proj*",
                 "*self_attn.v_proj*", "*self_attn.o_proj*"]:
        cfg["quant_cfg"].append({"quantizer_name": name, "enable": False})
    return cfg


def messages_to_text(messages):
    parts = []
    for msg in messages:
        role = msg.get("role", "")
        content = msg.get("content", "")
        if content:
            parts.append(f"{role}: {content}")
    return "\n".join(parts)


os.makedirs(OUTPUT_PATH, exist_ok=True)

print("=" * 60)
print("NVFP4 W4A4 Quantization")
print("Model: VeriLoop-E2 (Qwen3.8-27B)")
for i in range(torch.cuda.device_count()):
    p = torch.cuda.get_device_properties(i)
    free, _ = torch.cuda.mem_get_info(i)
    print(f"  cuda:{i}: {p.name}, total {p.total_memory/1e9:.1f} GB, free {free/1e9:.1f} GB")
print(f"Calibration: {CALIB_SIZE}x{CALIB_SEQ_LEN} from Nemotron-Competitive-Programming-v1")
print("=" * 60)

# Step 1: Load model across GPUs
print("\n[1/5] Loading model (multi-GPU)...")
t0 = time.time()
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    dtype=torch.bfloat16,
    device_map=build_device_map(),
    trust_remote_code=True,
)
print(f"  Model loaded in {time.time()-t0:.1f}s")
model.config.use_cache = False  # saves activation memory during calibration
n_cpu = sum(1 for v in getattr(model, "hf_device_map", {}).values() if v == "cpu")
print(f"  Modules on CPU: {n_cpu} (expected: lm_head/embed/norm with quantizer OFF)")

# Hook fix: "cpu" modules may arrive with execution_device=cuda:N, which makes
# modelopt's writeback (pre_forward) materialize multi-GB weights (embed/head)
# on the GPU. With exec=cpu the weight materializes on CPU (where it already
# lives in the weights map) and forward stays correct (the next layer's hook
# moves activations to CUDA).
from accelerate.hooks import AlignDevicesHook
_fixed = 0
for mod_name, dev in (getattr(model, "hf_device_map", {}) or {}).items():
    if dev != "cpu":
        continue
    m = model
    for p in mod_name.split("."):
        m = getattr(m, p)
    for sub in m.modules():
        hook = getattr(sub, "_hf_hook", None)
        if isinstance(hook, AlignDevicesHook) and hook.execution_device != torch.device("cpu"):
            hook.execution_device = torch.device("cpu")
            _fixed += 1
print(f"  Hooks redirected to CPU: {_fixed}")
try:
    print(f"  HF device map: {model.hf_device_map}")
except Exception:
    pass
for i in range(torch.cuda.device_count()):
    print(f"  cuda:{i} allocated={torch.cuda.memory_allocated(i)/1e9:.1f} GB")

# Step 2: Calibration data
print(f"\n[2/5] Loading calibration data ({CALIB_SIZE} samples)...")
t0 = time.time()
from datasets import load_dataset

calib_data = []
ds = load_dataset("nvidia/Nemotron-Competitive-Programming-v1", split=CALIB_SPLIT, streaming=True)
for item in ds:
    if len(calib_data) >= CALIB_SIZE:
        break
    text = messages_to_text(item.get("messages", []))
    if not text or len(text.strip()) < 100:
        continue
    encoded = tokenizer.encode(text, truncation=True, max_length=CALIB_SEQ_LEN)
    if len(encoded) < 32:
        continue
    if len(encoded) < CALIB_SEQ_LEN:
        encoded = encoded + [0] * (CALIB_SEQ_LEN - len(encoded))
    calib_data.append(torch.tensor(encoded[:CALIB_SEQ_LEN], dtype=torch.long))
    if len(calib_data) % 16 == 0:
        print(f"  Collected {len(calib_data)}/{CALIB_SIZE}...")
print(f"  Calibration data: {len(calib_data)} seqs in {time.time()-t0:.1f}s")

# Step 3: Forward loop
print("\n[3/5] Running calibration...")


@torch.no_grad()
def forward_loop(m):
    m.eval()
    # cpu/meta params first in line -> use the first CUDA param's device
    dev = next(p.device for p in m.parameters() if p.device.type == "cuda")
    torch.cuda.empty_cache()  # fights fragmentation (no expandable_segments on Windows)
    print(f"  forward_loop device: {dev}")
    for i, ids in enumerate(calib_data):
        try:
            m(input_ids=ids.unsqueeze(0).to(dev), labels=ids.unsqueeze(0).to(dev))
        except Exception as e:
            print(f"  Warning sample {i}: {type(e).__name__}: {e}")
            continue
        if (i + 1) % 16 == 0:
            print(f"    Calibrated {i+1}/{len(calib_data)}...")


t0 = time.time()
forward_loop(model)
print(f"  Calibration took {time.time()-t0:.1f}s")

# Step 4: Quantize
print("\n[4/5] Applying NVFP4 W4A4 quantization...")
t0 = time.time()
model = mtq.quantize(model, build_quant_cfg(), forward_loop)
print(f"  Quantization took {time.time()-t0:.1f}s")
mtq.print_quant_summary(model)

# Materialize meta weights of CPU modules (embed/head/norm): the export's dummy
# forward builds its fake input from next(model.parameters()).device, and meta
# tensors break everything ("Cannot copy out of meta tensor"). pre_forward with
# exec=cpu (fix above) brings the weight to CPU without touching VRAM;
# offload=False pins it there.
print("\n  Materializing meta weights on CPU...")
_n_meta = 0
for mod_name, dev in (getattr(model, "hf_device_map", {}) or {}).items():
    if dev != "cpu":
        continue
    m = model
    for p in mod_name.split("."):
        m = getattr(m, p)
    hook = getattr(m, "_hf_hook", None)
    if hook is None:
        continue
    if any(p.device.type == "meta" for p in m.parameters()):
        hook.pre_forward(m)
        _n_meta += 1
    hook.offload = False
print(f"  Materialized modules: {_n_meta}")
_n_meta_left = sum(1 for p in model.parameters() if p.device.type == "meta")
print(f"  Remaining meta params: {_n_meta_left}")

# Export on CPU: move the 64 layers (bf16 weights + NVFP4 scales) to RAM.
# Export quantizes one linear at a time and needs transient workspace on top of
# the resident base, which overflows smaller GPUs. RAM needs roughly:
# ~48.7 GB weights + ~9 GB scales + ~15 GB quantized output + temps.
print("\n  Moving layers to CPU for export...")
from accelerate.hooks import AlignDevicesHook as _ADH
for i in range(64):
    layer = model.model.layers[i]
    layer.to("cpu")
    for sub in layer.modules():
        hook = getattr(sub, "_hf_hook", None)
        if isinstance(hook, _ADH):
            hook.io_device = torch.device("cpu")
            hook.execution_device = torch.device("cpu")
    model.hf_device_map[f"model.layers.{i}"] = "cpu"
import gc as _gc
_gc.collect()
for _d in range(torch.cuda.device_count()):
    with torch.cuda.device(_d):
        torch.cuda.empty_cache()
print("  Layers on CPU. VRAM now:",
      " / ".join(f"cuda:{d}={torch.cuda.memory_allocated(d)/1e9:.1f}GB"
                 for d in range(torch.cuda.device_count())))

# Step 5: Export (workaround for a modelopt multimodal export bug where
# config.architectures ends up None and is_multimodal_model crashes on it)
print("\n[5/5] Exporting...")
try:
    import modelopt.torch.export.model_utils as mu
    _orig = mu.is_multimodal_model

    def _safe_is_mm(m):
        try:
            return _orig(m)
        except TypeError:
            archs = getattr(getattr(m, "config", None), "architectures", None)
            print(f"  is_multimodal_model fallback, architectures={archs} -> False")
            return False

    mu.is_multimodal_model = _safe_is_mm
    print("  Patched is_multimodal_model (None-safe)")
except Exception as e:
    print(f"  Patch skipped: {e}")

if getattr(model.config, "architectures", None) is None:
    model.config.architectures = ["Qwen3_5ForConditionalGeneration"]
    print(f"  Fixed config.architectures={model.config.architectures}")

from modelopt.torch.export import export_hf_checkpoint

# Free VRAM before export; calib buffers are no longer needed.
import gc
del calib_data
gc.collect()
for _d in range(torch.cuda.device_count()):
    with torch.cuda.device(_d):
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
print("  CUDA cache freed:",
      " / ".join(f"cuda:{d}={torch.cuda.memory_allocated(d)/1e9:.1f}GB"
                 for d in range(torch.cuda.device_count())))

t0 = time.time()
with torch.inference_mode():
    export_hf_checkpoint(model, export_dir=OUTPUT_PATH, max_shard_size="4GB")
print(f"  Export took {time.time()-t0:.1f}s")

print("\n  Copying tokenizer/config files...")
import shutil
for fname in [
    "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
    "config.json", "generation_config.json", "model.safetensors.index.json",
    "chat_template.jinja", "merges.txt", "vocab.json",
    "preprocessor_config.json", "configuration.json",
]:
    src = os.path.join(MODEL_PATH, fname)
    dst = os.path.join(OUTPUT_PATH, fname)
    if os.path.exists(src) and not os.path.exists(dst):
        shutil.copy2(src, dst)

files = list(Path(OUTPUT_PATH).rglob("*.safetensors"))
total_w = sum(f.stat().st_size for f in files)
total_all = sum(f.stat().st_size for f in Path(OUTPUT_PATH).rglob("*") if f.is_file())
print(f"\n{'=' * 60}")
print("Quantization complete!")
print(f"Output: {OUTPUT_PATH}")
print(f"Shards: {len(files)}, weights {total_w/1e9:.2f} GB, total {total_all/1e9:.2f} GB")
print(f"{'=' * 60}")