[bugfix] Cosmos3 edge fix config use_k_norm_und_for_gen name (#5239)

Signed-off-by: Bartosz Stefaniak <bstefaniak@nvidia.com>

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)
 
 

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

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.",