scripts/patch_modelopt_fp8_guard.py

"""Extend the NQ_GUARD placeholder guard to ModelOptFp8LinearMethod.

Same root cause as the NVFP4 case (see patch_modelopt_guard.py): vLLM's pooling
adapter swaps absent modules (lm_head in encoder checkpoints) for StageMissingLayer
placeholders whose quant params are meta tensors; ModelOptFp8LinearMethod.
process_weights_after_loading then calls .all() on a meta tensor and crashes.
Idempotent: skips if the FP8 guard is already present.
"""

import re
import sys
from pathlib import Path

if len(sys.argv) > 1:
    path = Path(sys.argv[1])
else:
    # Resolve against whatever vllm is importable here — works in any venv or
    # container (the old hardcoded arcade path broke the first AWS run).
    import vllm.model_executor.layers.quantization.modelopt as _m

    path = Path(_m.__file__)
src = path.read_text()

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

# Target: the process_weights_after_loading whose first body line is
# "weight = layer.weight" followed by "max_w_scale = layer.weight_scale.max()"
anchor = (
    "    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:\n"
    "        weight = layer.weight\n"
    "        max_w_scale = layer.weight_scale.max()\n"
)
guard = (
    "    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:\n"
    "        # NQ_GUARD_FP8: same placeholder guard as the NVFP4 method above --\n"
    "        # StageMissingLayer/PPMissingLayer carry meta-tensor quant params.\n"
    '        if layer.__class__.__name__ in ("StageMissingLayer", "PPMissingLayer"):\n'
    "            return\n"
    "        weight = layer.weight\n"
    "        max_w_scale = layer.weight_scale.max()\n"
)

count = src.count(anchor)
if count != 1:
    print(f"FATAL: anchor matched {count} times (expected 1); aborting")
    sys.exit(1)

path.write_text(src.replace(anchor, guard, 1))
print("patched: NQ_GUARD_FP8 inserted into ModelOptFp8LinearMethod")