scripts/patch_modelopt_guard.py

"""Replace debug instrumentation with the clean placeholder guard in vLLM modelopt.py.

Root cause (stage D, runs/stage_d logs): the pooling adapter replaces lm_head with
StageMissingLayer, but ModelOptNvFp4LinearMethod.process_weights_after_loading still
runs on it and calls torch.unique on never-materialized (meta) scale params. vLLM's
weight loader guards these placeholder classes; the ModelOpt quant path doesn't.
Upstream-worthy fix: early-return on placeholder layers.
"""
import sys
from pathlib import Path

if len(sys.argv) > 1:
    p = Path(sys.argv[1])
else:
    # Resolve against whatever vllm is importable here (venv or container).
    import vllm.model_executor.layers.quantization.modelopt as _m
    p = Path(_m.__file__)
src = p.read_text()

if "NQ_GUARD:" in src:
    print("already patched (NQ_GUARD present)")
    sys.exit(0)

# strip any previous NQ_DEBUG block (v2 or v3)
for tag in ("        # NQ_DEBUG2\n", "        # NQ_DEBUG3\n"):
    while tag in src:
        start = src.index(tag)
        end = src.index("        if (\n", start)
        src = src[:start] + src[end:]
        print(f"removed {tag.strip()} block")

anchor = """    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        if (
            torch.unique(layer.input_scale).numel() != 1"""

guard = """    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        # NQ_GUARD: pooling/PP adapters swap absent modules (e.g. lm_head in encoder
        # checkpoints) for placeholder layers whose quant params are never
        # materialized; vLLM's weight loader skips them but this path did not.
        if layer.__class__.__name__ in ("StageMissingLayer", "PPMissingLayer"):
            return
        if (
            torch.unique(layer.input_scale).numel() != 1"""

assert anchor in src, "anchor not found"
src = src.replace(anchor, guard)
p.write_text(src)
print("NQ_GUARD applied")