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))