From ca8a9471a6aa75c40bbc8542365cb3601eb4af13 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Thu, 22 Oct 2026 04:16:47 -0700 Subject: [PATCH] Fix Cosmos3 warnings for auxiliary FP8 export scales --- .../transformers/transformer_cosmos3.py | 3 ++ .../test_models_transformer_cosmos3.py | 32 ++++++++++++++++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py index 0295e1c0567b..57dcb8b68bfd 100644 --- a/src/diffusers/models/transformers/transformer_cosmos3.py +++ b/src/diffusers/models/transformers/transformer_cosmos3.py @@ -377,6 +377,9 @@ class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, Attentio _repeated_blocks = ["Cosmos3VLTextMoTDecoderLayer"] _skip_layerwise_casting_patterns = ["embed_tokens", "time_embedder", "norm"] _keep_in_fp32_modules = ["time_embedder"] + # These scale tensors are exported for vLLM-Omni/SGLang and are unused + # by Diffusers. ModelOpt restores its quantizer state separately. + _keys_to_ignore_on_load_unexpected = [r"\.(?:input_scale|weight_scale)$"] # Optional context-parallelism seams. They default to ``None`` (no-op) so the # model itself carries no CP logic. `forward` applies `_cp_shard_fn` to the # per-pathway hidden states + rotary embeddings before the decoder layers, and diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index c897d5bfdc69..3d42aeecfd8d 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -15,6 +15,7 @@ import pytest import torch +from safetensors.torch import save_file from diffusers import Cosmos3OmniTransformer, SeaCacheConfig from diffusers.hooks import sea_cache as sea_cache_module @@ -27,9 +28,10 @@ Cosmos3PackedMoTAttention, Cosmos3VLTextMoTDecoderLayer, ) +from diffusers.utils import logging from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import CaptureLogger, enable_full_determinism, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -524,6 +526,34 @@ def test_cosmos3_sea_cache_regional_compile_fullgraph_without_recompile(self): class TestCosmos3OmniTransformerModel(Cosmos3OmniTransformerTesterConfig, ModelTesterMixin): + @pytest.mark.parametrize("include_unexpected_weight", [False, True]) + def test_load_ignores_export_scales(self, tmp_path, include_unexpected_weight): + model = self.model_class(**self.get_init_dict()) + model.save_pretrained(tmp_path) + state_dict = model.state_dict() + scale_keys = ["layers.0.self_attn.to_q.input_scale", "layers.0.self_attn.to_q.weight_scale"] + for key in scale_keys: + state_dict[key] = torch.ones(1) + unexpected_key = "layers.0.self_attn.to_q.unexpected_weight" + if include_unexpected_weight: + state_dict[unexpected_key] = torch.ones(1) + save_file(state_dict, tmp_path / "diffusion_pytorch_model.safetensors", metadata={"format": "pt"}) + + logger = logging.get_logger("diffusers.models.modeling_utils") + with CaptureLogger(logger) as captured: + _, loading_info = self.model_class.from_pretrained(tmp_path, output_loading_info=True) + + assert loading_info["unexpected_keys"] == ([unexpected_key] if include_unexpected_weight else []) + assert loading_info["missing_keys"] == [] + assert loading_info["mismatched_keys"] == [] + for key in scale_keys: + assert key not in captured.out + if include_unexpected_weight: + assert unexpected_key in captured.out + assert "were not used when initializing" in captured.out + else: + assert "were not used when initializing" not in captured.out + def test_cosmos3_decoder_layer_cache_metadata_tracks_generation_stream(self): metadata = TransformerBlockRegistry.get(Cosmos3VLTextMoTDecoderLayer)