Fix loading of ModelOpt FP8 checkpoints for Cosmos3 (#5076)

Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>

tests/diffusion/model_loader/test_modelopt_fp8_adapter.py

@@ -42,6 +42,53 @@ class _QuantizedPackedModelOptModel(nn.Module):
         )
 
 
+class _RemappedModelOptModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Linear(2, 2, bias=False)
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        return {
+            "transformer.orig.proj.weight": "runtime.proj.weight",
+        }.get(name, name)
+
+
+class _QuantizedRemappedModelOptModel(nn.Module):
+    """FP8 target whose params live under a namespace the raw checkpoint keys
+    never use, reachable only via ``remap_checkpoint_key``.
+
+    Mirrors the Cosmos3 case that regressed: the generic weights-mapper cannot
+    resolve ``transformer.orig.*`` to ``runtime.*``, so without the remap hook
+    the adapter treats every tensor as unresolved and silently drops the scales.
+    """
+
+    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(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "weight_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "input_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        prefix = "transformer.orig.proj."
+        if name.startswith(prefix):
+            return "runtime.proj." + name[len(prefix) :]
+        return name
+
+
 def _make_source() -> SimpleNamespace:
     return SimpleNamespace(
         subfolder="transformer",
@@ -92,3 +139,62 @@ def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target():
         "transformer.block.to_q.weight_scale",
         "transformer.block.to_q.input_scale",
     ]
+
+
+def test_modelopt_adapter_uses_checkpoint_key_remap_for_target_dtype():
+    model = _RemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    scale = torch.tensor([0.5], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight", fp8_weight),
+                    ("transformer.orig.proj.weight_scale", scale),
+                ]
+            )
+        )
+    )
+
+    assert [name for name, _ in adapted] == ["runtime.proj.weight"]
+    assert adapted[0][1].dtype == model.runtime.proj.weight.dtype
+    assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale)
+
+
+def test_modelopt_adapter_remaps_and_keeps_scales_for_quantized_target():
+    """Regression: scales were dropped for a quantized (FP8) target that only the
+    model's ``remap_checkpoint_key`` hook can resolve. Assert both scales are
+    emitted and the FP8 weight passes through unchanged (no dequantization)."""
+    model = _QuantizedRemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    weight_scale = torch.tensor([0.5], dtype=torch.float32)
+    input_scale = torch.tensor([1.0], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight_scale", weight_scale),
+                    ("transformer.orig.proj.input_scale", input_scale),
+                    ("transformer.orig.proj.weight", fp8_weight),
+                ]
+            )
+        )
+    )
+
+    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)
+
+    # 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))

vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py

@@ -32,7 +32,7 @@ FP8_DTYPES = tuple(
 @dataclass
 class _AdaptState:
     scale_tensors: dict[str, torch.Tensor] = field(default_factory=dict)
-    pending_weights: dict[str, list[tuple[str, torch.Tensor, torch.dtype]]] = field(default_factory=dict)
+    pending_weights: dict[str, list[tuple[str, str, torch.Tensor, torch.dtype]]] = field(default_factory=dict)
     skipped_scales: int = 0
     dequantized_weights: int = 0
 
@@ -41,6 +41,7 @@ class ModelOptFp8CheckpointAdapter:
     def __init__(self, model: nn.Module, source: object):
         self._loadable_tensors = self._get_model_loadable_tensors(model)
         self._weights_mapper = self._get_weights_mapper(model)
+        self._checkpoint_key_mapper = getattr(model, "remap_checkpoint_key", None)
         self._source_label = getattr(source, "prefix", "") or getattr(source, "subfolder", None) or "model"
 
     @classmethod
@@ -110,14 +111,19 @@ class ModelOptFp8CheckpointAdapter:
             orig_to_new_prefix=orig_to_new_prefix,
         )
 
-    def _resolve_target_name(self, name: str) -> str | None:
+    def _resolve_target_and_output_names(self, name: str) -> tuple[str | None, str]:
         if name in self._loadable_tensors:
-            return name
+            return name, name
+
+        if callable(self._checkpoint_key_mapper):
+            candidate = self._checkpoint_key_mapper(name)
+            if candidate in self._loadable_tensors:
+                return candidate, candidate
 
         for candidate in self._weights_mapper.apply_list([name]):
             if candidate != name and candidate in self._loadable_tensors:
-                return candidate
-        return None
+                return candidate, name
+        return None, name
 
     @staticmethod
     def _reshape_weight_scale(scale: torch.Tensor, weight_shape: torch.Size) -> torch.Tensor:
@@ -160,8 +166,13 @@ class ModelOptFp8CheckpointAdapter:
         scale_name: str,
         state: _AdaptState,
     ) -> Generator[tuple[str, torch.Tensor], None, None]:
-        for weight_name, weight_tensor, target_dtype in state.pending_weights.pop(scale_name, []):
-            yield weight_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype)
+        for (
+            weight_name,
+            output_name,
+            weight_tensor,
+            target_dtype,
+        ) in state.pending_weights.pop(scale_name, []):
+            yield output_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype)
             state.dequantized_weights += 1
 
     def _handle_scale_tensor(
@@ -194,6 +205,7 @@ class ModelOptFp8CheckpointAdapter:
     def _maybe_dequantize_or_defer_weight(
         self,
         name: str,
+        output_name: str,
         tensor: torch.Tensor,
         target_dtype: torch.dtype,
         state: _AdaptState,
@@ -203,7 +215,7 @@ class ModelOptFp8CheckpointAdapter:
             raise ValueError(f"Missing ModelOpt FP8 weight_scale name for weight {name!r}")
 
         if scale_name not in state.scale_tensors:
-            state.pending_weights.setdefault(scale_name, []).append((name, tensor, target_dtype))
+            state.pending_weights.setdefault(scale_name, []).append((name, output_name, tensor, target_dtype))
             return None
 
         state.dequantized_weights += 1
@@ -235,17 +247,17 @@ class ModelOptFp8CheckpointAdapter:
         state = _AdaptState()
 
         for name, tensor in weights:
-            target_name = self._resolve_target_name(name)
+            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)
                 continue
 
             target_dtype = self._target_dtype_for_dequantization(tensor, target_name)
             if target_dtype is not None:
-                tensor = self._maybe_dequantize_or_defer_weight(name, tensor, target_dtype, state)
+                tensor = self._maybe_dequantize_or_defer_weight(name, output_name, tensor, target_dtype, state)
                 if tensor is None:
                     continue
-            yield name, tensor
+            yield output_name, tensor
 
         self._check_pending_weights(state)
         self._log_adaptation_summary(state)

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -1018,6 +1018,9 @@ class Cosmos3OmniDiffusersPipeline(
 
         return None
 
+    # Checkpoint adapters use this hook before model-specific weight loading.
+    remap_checkpoint_key = _remap_ckpt_key
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Stream-remap checkpoint weights and load via AutoWeightsLoader.
 
@@ -1032,6 +1035,10 @@ class Cosmos3OmniDiffusersPipeline(
             total = kept = 0
             for name, tensor in weights:
                 total += 1
+                if name in allowed or name in tp_aware:
+                    kept += 1
+                    yield name, tensor
+                    continue
                 remapped = self._remap_ckpt_key(name)
                 if remapped is not None and (remapped in allowed or remapped in tp_aware):
                     kept += 1