From a7b6a174bc215fd7b79121740f82c59482e5b0ac Mon Sep 17 00:00:00 2001 From: wkutak Date: Wed, 15 Jul 2026 02:33:27 +0200 Subject: [PATCH] Fix loading of ModelOpt FP8 checkpoints for Cosmos3 (#5076) Signed-off-by: Wojciech Kutak --- .../model_loader/test_modelopt_fp8_adapter.py | 106 ++++++++++++++++++ .../checkpoint_adapters/modelopt.py | 34 ++++-- .../models/cosmos3/pipeline_cosmos3.py | 7 ++ 3 files changed, 136 insertions(+), 11 deletions(-) diff --git a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py b/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py index 513e80d28..66fa87a3c 100644 --- a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py +++ b/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)) diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py index 04fc8ea32..d6f7b3e0c 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py +++ b/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) diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index cd8c12777..5d0e24789 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/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