From 02eb7a4ecb5cd602d1e3dc296b0ee03928e9218b Mon Sep 17 00:00:00 2001 From: Mu GuanLin <39453182+mglyn@users.noreply.github.com> Date: Tue, 9 Jun 2026 17:09:39 +0800 Subject: [PATCH] [Bugfix] Register LTX RMSNorm identity weight as buffer (#4278) Signed-off-by: mglyn <1203789601@qq.com> --- .../models/ltx2/test_ltx2_transformer.py | 22 +++++++++++++++++++ .../diffusion/models/ltx2/ltx2_transformer.py | 12 +++++++++- 2 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 tests/diffusion/models/ltx2/test_ltx2_transformer.py diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py new file mode 100644 index 000000000..e71253297 --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index a1d72cf66..4e5659c54 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/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: