From ef15ea2391e9cbcc63dbe0b094b710d1e584fa20 Mon Sep 17 00:00:00 2001 From: William Lin Date: Mon, 13 Jul 2026 16:04:21 -0700 Subject: [PATCH] [bugfix]: keep LTX2 rms_norm outputs bf16 under torch 2.12 autocast (#1587) --- fastvideo/models/dits/ltx2.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index bdfeef0c7..4e5227e49 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -72,8 +72,13 @@ def _rms_norm_dispatch( """Use QuACK RMSNorm only for LTX-2 refine stage, else Torch RMSNorm.""" if _is_ltx2_refine_stage(): return _quack_rmsnorm(x, weight=weight, eps=eps) + # Official LTX-2 runs the DiT natively in bf16 with no autocast, so its + # rms_norm outputs stay bf16 on every torch version. torch 2.12 added + # rms_norm to autocast's float32 cast policy, which under our + # autocast(bfloat16) forward returns float32 instead. Restore the input + # dtype to match the official numerics (and torch <= 2.11 behavior). return torch.nn.functional.rms_norm( - x, (x.shape[-1], ), weight=weight, eps=eps) + x, (x.shape[-1], ), weight=weight, eps=eps).to(x.dtype) class StageAwareRMSNorm(nn.RMSNorm):