Add vLLM NVFP4 lm_head patch (ParallelLMHead -> ReplicatedLinear)

vllm-patches/patch_nvfp4_lm_head.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+"""Patch vLLM's qwen3_5.py to support NVFP4-quantized lm_head.
+
+Problem: vLLM's ParallelLMHead inherits from VocabParallelEmbedding, which
+doesn't support quantized weight loading (NVFP4 packed format). The
+compressed-tensors quantization method can't match "Linear" target to
+"ParallelLMHead" class name, and even if it could, VocabParallelEmbedding's
+weight_loader doesn't handle packed weights.
+
+Fix: Replace ParallelLMHead with ReplicatedLinear for lm_head, which goes
+through the standard quantized linear weight loading path. Also hardcode
+hidden_size=3072 because the inner CausalLM class sees text_config defaults
+(2048) instead of the actual model hidden_size.
+
+Usage (inside container):
+    python3 /workspace/patch_nvfp4_lm_head.py
+
+Or from host:
+    docker cp patch_nvfp4_lm_head.py CONTAINER:/tmp/
+    docker exec CONTAINER python3 /tmp/patch_nvfp4_lm_head.py
+    docker restart CONTAINER
+"""
+
+import sys
+
+TARGET = "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/qwen3_5.py"
+
+with open(TARGET) as f:
+    content = f.read()
+
+changes = 0
+
+# 1. Add ReplicatedLinear import
+old_import = "from vllm.model_executor.layers.vocab_parallel_embedding import ("
+new_import = "from vllm.model_executor.layers.linear import ReplicatedLinear\nfrom vllm.model_executor.layers.vocab_parallel_embedding import ("
+if old_import in content and "ReplicatedLinear" not in content:
+    content = content.replace(old_import, new_import, 1)
+    changes += 1
+    print("[1/2] Added ReplicatedLinear import")
+
+# 2. Replace ParallelLMHead with ReplicatedLinear for lm_head
+old_lm = """                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )"""
+
+new_lm = """                self.lm_head = ReplicatedLinear(
+                    3072,
+                    config.vocab_size,
+                    bias=False,
+                    quant_config=self.quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )"""
+
+if old_lm in content:
+    content = content.replace(old_lm, new_lm)
+    changes += 1
+    print("[2/2] Replaced ParallelLMHead with ReplicatedLinear(3072, vocab_size)")
+
+if changes == 0:
+    print("ERROR: No patterns matched. File may already be patched or version mismatch.")
+    sys.exit(1)
+
+with open(TARGET, "w") as f:
+    f.write(content)
+
+print(f"\nPatched {TARGET} ({changes} changes)")
+print("Restart the container to apply.")