Conversation
bf16 and fp16 loads cast the persistent inv_freq buffer and built positions in that dtype, so neighbouring latents shared a rotary angle. Compute the outer product in fp32 and keep rotary_pos_emb in fp32 on from_pretrained. Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com>
github-actions
Bot
added
fixes-issue
size/S
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #14934
Loading
StableAudio3DiTModelwithtorch_dtype=bf16orfp16casts the persistentrotary_pos_emb.inv_freqbuffer, andforwardwas building positions in that dtype. bf16 only keeps integers exact up to 256 (fp16 up to 2048), so on longer clips neighbouring audio latents shared a rotary angle.The outer product is now computed in fp32, and
_keep_in_fp32_moduleskeepsrotary_pos_embin fp32 acrossfrom_pretrained(torch_dtype=...). A later.to(bf16)can still round the buffer itself; positions no longer collapse.test_low_precision_rope_keeps_distinct_positionscovers a ~2 min sequence in fp16 and bf16.