[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()