From 2f33df4678b2b13c70abd9ddffffc314908fb23a Mon Sep 17 00:00:00 2001 From: bastefaniak Date: Tue, 21 Jul 2026 09:56:31 +0200 Subject: [PATCH] [bugfix] Cosmos3 edge fix config use_k_norm_und_for_gen name (#5239) Signed-off-by: Bartosz Stefaniak --- .../models/cosmos3/test_cosmos3_transformer.py | 12 ++++++------ .../models/cosmos3/transformer_cosmos3.py | 2 +- .../models/cosmos3/transformer_cosmos3_edge.py | 16 ++++++++-------- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index cb59dffe5..5dd8a875a 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -193,7 +193,7 @@ def test_edge_und_and_gen_use_relu2_and_no_und_qk_norm() -> None: SimpleNamespace( tf_model_config=_tiny_cosmos3_edge_config( num_hidden_layers=1, - use_k_norm_und_for_gen=False, + use_und_k_norm_for_gen=False, ), dtype=torch.float32, ) @@ -221,14 +221,14 @@ def test_edge_creates_gen_k_norm_when_enabled() -> None: tf_model_config=_tiny_cosmos3_edge_config( num_hidden_layers=1, qk_norm_for_diffusion=True, - use_k_norm_und_for_gen=True, + use_und_k_norm_for_gen=True, ), dtype=torch.float32, ) ) k_norm = model.language_model.layers[0].self_attn.k_norm_und_for_gen - assert model.use_k_norm_und_for_gen is True + assert model.use_und_k_norm_for_gen is True assert isinstance(k_norm, RMSNorm) assert isinstance(model.gen_layers[0].cross_attention.norm_q, RMSNorm) assert isinstance(model.gen_layers[0].cross_attention.norm_k, RMSNorm) @@ -244,7 +244,7 @@ def test_edge_supports_no_gen_qk_norm_variant() -> None: tf_model_config=_tiny_cosmos3_edge_config( num_hidden_layers=1, qk_norm_for_diffusion=False, - use_k_norm_und_for_gen=False, + use_und_k_norm_for_gen=False, ), dtype=torch.float32, ) @@ -266,7 +266,7 @@ def test_edge_validates_required_relu2_weights() -> None: nn.Module.__init__(model) model.num_hidden_layers = 1 model.qk_norm_for_diffusion = True - model.use_k_norm_und_for_gen = True + model.use_und_k_norm_for_gen = True complete = { "transformer.language_model.layers.0.mlp.up_proj.weight", @@ -288,7 +288,7 @@ def test_edge_validates_required_relu2_weights() -> None: with pytest.raises(ValueError, match=r"self_attn\.k_norm_und_for_gen"): model.validate_loaded_weights(missing_k_norm) - model.use_k_norm_und_for_gen = False + model.use_und_k_norm_for_gen = False model.validate_loaded_weights(missing_k_norm) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index de9767d6f..88d280bd0 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -1123,7 +1123,7 @@ class Cosmos3VFMTransformer(nn.Module): ) self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size - self.use_k_norm_und_for_gen = _tf_config_get(model_config, "use_k_norm_und_for_gen", None) + self.use_und_k_norm_for_gen = _tf_config_get(model_config, "use_und_k_norm_for_gen", None) dtype = od_config.dtype quant_config = getattr(od_config, "quantization_config", None) if od_config else None diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py index 6a1b4313b..c9a7d856e 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py @@ -93,7 +93,7 @@ class Cosmos3EdgeCausalAttention(nn.Module): num_key_value_heads: int, head_dim: int, rms_norm_eps: float, - use_k_norm_und_for_gen: bool, + use_und_k_norm_for_gen: bool, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: @@ -144,7 +144,7 @@ class Cosmos3EdgeCausalAttention(nn.Module): prefix=f"{prefix}.to_out", ) - self.k_norm_und_for_gen = RMSNorm(self.head_dim, eps=rms_norm_eps) if use_k_norm_und_for_gen else None + self.k_norm_und_for_gen = RMSNorm(self.head_dim, eps=rms_norm_eps) if use_und_k_norm_for_gen else None self.attn = FrameworkAttention( num_heads=self.num_heads, head_size=self.head_dim, @@ -195,7 +195,7 @@ class Cosmos3EdgeUndDecoderLayer(nn.Module): num_key_value_heads: int, head_dim: int, rms_norm_eps: float, - use_k_norm_und_for_gen: bool, + use_und_k_norm_for_gen: bool, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: @@ -206,7 +206,7 @@ class Cosmos3EdgeUndDecoderLayer(nn.Module): num_key_value_heads=num_key_value_heads, head_dim=head_dim, rms_norm_eps=rms_norm_eps, - use_k_norm_und_for_gen=use_k_norm_und_for_gen, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, quant_config=quant_config, prefix=f"{prefix}.self_attn", ) @@ -256,7 +256,7 @@ class Cosmos3EdgeLanguageModel(nn.Module): rms_norm_eps: float, rope_theta: float, mrope_section: list[int], - use_k_norm_und_for_gen: bool, + use_und_k_norm_for_gen: bool, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: @@ -276,7 +276,7 @@ class Cosmos3EdgeLanguageModel(nn.Module): num_key_value_heads=num_key_value_heads, head_dim=head_dim, rms_norm_eps=rms_norm_eps, - use_k_norm_und_for_gen=use_k_norm_und_for_gen, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, quant_config=quant_config, prefix=f"{prefix}.layers.{i}", ) @@ -368,7 +368,7 @@ class Cosmos3EdgeVFMTransformer(Cosmos3VFMTransformer): return list(mrope_section) def _language_model_kwargs(self) -> dict[str, Any]: - return {"use_k_norm_und_for_gen": self.use_k_norm_und_for_gen} + return {"use_und_k_norm_for_gen": self.use_und_k_norm_for_gen} def validate_loaded_weights(self, loaded: set[str]) -> None: missing: list[str] = [] @@ -379,7 +379,7 @@ class Cosmos3EdgeVFMTransformer(Cosmos3VFMTransformer): f"gen_layers.{layer_idx}.mlp.up_proj.", f"gen_layers.{layer_idx}.mlp.down_proj.", ) - if self.use_k_norm_und_for_gen: + if self.use_und_k_norm_for_gen: required_markers = ( *required_markers, f"language_model.layers.{layer_idx}.self_attn.k_norm_und_for_gen.",