From 23bcf880266eadc2efd7226509542d16ee7909a7 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 5 Jul 2026 13:11:42 -0700 Subject: [PATCH] [bugfix]: keep _rms_norm_dispatch output in the input dtype under autocast --- fastvideo/models/dits/ltx2.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index bdfeef0c7..c3b31674e 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -72,8 +72,12 @@ 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) + # torch 2.12 added rms_norm to autocast's float32 cast policy, so under + # autocast(bfloat16) this norm now returns float32 where torch <= 2.11 + # returned bfloat16. Restore the input dtype so autocast-blind consumers + # downstream don't receive upcast activations. 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):