[Bugfix] Register LTX RMSNorm identity weight as buffer (#4278)

Signed-off-by: mglyn <1203789601@qq.com>

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from torch import nn
+
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm
+
+
+def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
+    norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
+
+    assert "weight" not in dict(norm.named_parameters())
+    assert "weight" in dict(norm.named_buffers())
+    assert "weight" not in norm.state_dict()
+
+
+def test_ltx_rms_norm_affine_weight_remains_parameter():
+    norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=True)
+
+    assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter)
+    assert "weight" not in dict(norm.named_buffers())
+    assert "weight" in norm.state_dict()

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -62,7 +62,17 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) ->
         kwargs["has_weight"] = elementwise_affine
     elif not elementwise_affine:
         raise TypeError("RMSNorm backend does not support disabling affine weights.")
-    return RMSNorm(hidden_size, **kwargs)
+    norm = RMSNorm(hidden_size, **kwargs)
+    weight = getattr(norm, "weight", None)
+    if (
+        not elementwise_affine
+        and isinstance(weight, torch.Tensor)
+        and not isinstance(weight, nn.Parameter)
+        and "weight" not in dict(norm.named_buffers())
+    ):
+        delattr(norm, "weight")
+        norm.register_buffer("weight", weight, persistent=False)
+    return norm
 
 
 def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: