From ae949d129f4d6dbf58640397b926dda252f30122 Mon Sep 17 00:00:00 2001 From: Alex Steiner Date: Thu, 16 Jul 2026 12:33:57 -0700 Subject: [PATCH] [Bugfix][Quantization] Remap ModelOpt NVFP4 scale tensors (#5087) Signed-off-by: Alex Steiner --- .../model_loader/test_modelopt_fp8_adapter.py | 49 +++++++++++++++++-- .../checkpoint_adapters/modelopt.py | 31 ++++++++++-- 2 files changed, 73 insertions(+), 7 deletions(-) diff --git a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py b/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py index 66fa87a3c..8c429a00f 100644 --- a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py +++ b/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]))]))) diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py index d6f7b3e0c..68ec98d5a 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py +++ b/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