[Bugfix][Quantization] Remap ModelOpt NVFP4 scale tensors (#5087)

Signed-off-by: Alex Steiner <asteiner@nvidia.com>

tests/diffusion/model_loader/test_modelopt_fp8_adapter.py

@@ -9,6 +9,7 @@ import torch.nn as nn
 
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     ModelOptFp8CheckpointAdapter,
+    ModelOptNvFp4CheckpointAdapter,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -89,6 +90,23 @@ class _QuantizedRemappedModelOptModel(nn.Module):
         return name
 
 
+class _RemappedNvFp4Model(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Module()
+        self.runtime.proj.register_parameter(
+            "weight",
+            nn.Parameter(torch.empty(1, 1, dtype=torch.uint8), requires_grad=False),
+        )
+        for name in ("input_scale", "weight_scale", "weight_scale_2"):
+            self.runtime.proj.register_parameter(name, nn.Parameter(torch.empty(1), requires_grad=False))
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        return name.replace("transformer.orig.proj.", "runtime.proj.")
+
+
 def _make_source() -> SimpleNamespace:
     return SimpleNamespace(
         subfolder="transformer",
@@ -188,13 +206,36 @@ def test_modelopt_adapter_remaps_and_keeps_scales_for_quantized_target():
     emitted = dict(adapted)
     assert len(adapted) == len(emitted) == 3
 
-    # Both scales are retained (the dropped-scale regression). They keep their
-    # source-checkpoint names; the model's load_weights remaps them downstream.
-    assert torch.equal(emitted["transformer.orig.proj.weight_scale"], weight_scale)
-    assert torch.equal(emitted["transformer.orig.proj.input_scale"], input_scale)
+    # Both scales are retained and emitted under the remapped runtime names.
+    assert torch.equal(emitted["runtime.proj.weight_scale"], weight_scale)
+    assert torch.equal(emitted["runtime.proj.input_scale"], input_scale)
 
     # The FP8 weight is passed through unchanged under the remapped target name
     # (the quantized param is FP8, so it must NOT be dequantized). Compare via
     # float32 since torch.equal has no FP8 CPU kernel.
     assert emitted["runtime.proj.weight"].dtype == torch.float8_e4m3fn
     assert torch.equal(emitted["runtime.proj.weight"].to(torch.float32), fp8_weight.to(torch.float32))
+
+
+def test_modelopt_nvfp4_adapter_remaps_quantized_weights_and_scales():
+    adapter = ModelOptNvFp4CheckpointAdapter(_RemappedNvFp4Model(), _make_source())
+    prefix = "transformer.orig.proj"
+    checkpoint_tensors = [
+        (f"{prefix}.input_scale", torch.tensor([1.0])),
+        (f"{prefix}.weight_scale", torch.tensor([0.5])),
+        (f"{prefix}.weight_scale_2", torch.tensor([0.25])),
+        (f"{prefix}.weight", torch.tensor([[1]], dtype=torch.uint8)),
+    ]
+
+    adapted = list(adapter.adapt(iter(checkpoint_tensors)))
+
+    assert [name for name, _ in adapted] == [
+        name.replace(f"{prefix}.", "runtime.proj.") for name, _ in checkpoint_tensors
+    ]
+
+
+def test_modelopt_nvfp4_adapter_rejects_unconsumed_pre_quant_scale():
+    adapter = ModelOptNvFp4CheckpointAdapter(nn.Module(), _make_source())
+
+    with pytest.raises(ValueError, match="does not consume pre_quant_scale"):
+        list(adapter.adapt(iter([("transformer.proj.pre_quant_scale", torch.tensor([1.0]))])))

vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py

@@ -11,7 +11,12 @@ from vllm.model_executor.utils import get_packed_modules_mapping
 
 logger = init_logger(__name__)
 
-MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_2", ".weight_scale_inv")
+MODEL_OPT_SCALE_SUFFIXES = (
+    ".input_scale",
+    ".weight_scale",
+    ".weight_scale_2",
+    ".weight_scale_inv",
+)
 DEFAULT_PACKED_MODULES_MAPPING = {
     "to_qkv": ("to_q", "to_k", "to_v"),
     "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"),
@@ -122,6 +127,7 @@ class ModelOptFp8CheckpointAdapter:
 
         for candidate in self._weights_mapper.apply_list([name]):
             if candidate != name and candidate in self._loadable_tensors:
+                # Preserve shard names so downstream packed-weight loaders can route them.
                 return candidate, name
         return None, name
 
@@ -178,6 +184,7 @@ class ModelOptFp8CheckpointAdapter:
     def _handle_scale_tensor(
         self,
         name: str,
+        output_name: str,
         tensor: torch.Tensor,
         target_name: str | None,
         state: _AdaptState,
@@ -186,7 +193,7 @@ class ModelOptFp8CheckpointAdapter:
         if target_name is None:
             state.skipped_scales += 1
         else:
-            yield name, tensor
+            yield output_name, tensor
         yield from self._flush_pending_weights(name, state)
 
     def _target_dtype_for_dequantization(
@@ -249,7 +256,7 @@ class ModelOptFp8CheckpointAdapter:
         for name, tensor in weights:
             target_name, output_name = self._resolve_target_and_output_names(name)
             if self._is_scale(name):
-                yield from self._handle_scale_tensor(name, tensor, target_name, state)
+                yield from self._handle_scale_tensor(name, output_name, tensor, target_name, state)
                 continue
 
             target_dtype = self._target_dtype_for_dequantization(tensor, target_name)
@@ -264,6 +271,8 @@ class ModelOptFp8CheckpointAdapter:
 
 
 class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
+    _PRE_QUANT_SCALE_SUFFIX = ".pre_quant_scale"
+
     @staticmethod
     def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
         return (
@@ -273,6 +282,22 @@ class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
             and bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False))
         )
 
+    def adapt(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        def validated_weights() -> Generator[tuple[str, torch.Tensor], None, None]:
+            for name, tensor in weights:
+                if name.endswith(self._PRE_QUANT_SCALE_SUFFIX):
+                    raise ValueError(
+                        f"ModelOpt NVFP4 checkpoint tensor {name!r} is unsupported: "
+                        "vLLM 0.25.0 does not consume pre_quant_scale. Export the checkpoint "
+                        "with pre-quant scales folded into the weights."
+                    )
+                yield name, tensor
+
+        yield from super().adapt(validated_weights())
+
 
 class ModelOptMixedPrecisionCheckpointAdapter(ModelOptFp8CheckpointAdapter):
     @staticmethod