diffusers
6e8c8c4c - Fix `Ideogram4MRoPE` collapsing under `torch.autocast` (compute rotary in float32) (#13922)

Commit
28 days ago
Fix `Ideogram4MRoPE` collapsing under `torch.autocast` (compute rotary in float32) (#13922) * Fix `Ideogram4MRoPE` collapsing under `torch.autocast` (compute rotary in float32) Ideogram4 builds image-token positions as IMAGE_POSITION_OFFSET (65536) + (t, h, w). `Ideogram4MRoPE.forward` casts its operands to float32, but the rotary matmul (and cos/sin) is on autocast's downcast list, so under torch.autocast("cuda", bfloat16) — common in training and pipeline code — it runs in bfloat16 anyway. bfloat16's step at 65536 is 512, so every image position in a <=512 grid rounds to the same value: all image tokens get identical rotary embeddings, spatial information is lost, and the decoded image degenerates to a flat color. Wrap the frequency computation in torch.autocast(enabled=False) so the rotary embeddings are always computed in float32, matching how transformers guards its RoPE modules. Added a regression test that fails on main and passes with the fix. Fixes #13920 * Compute the rotary frequencies in float64 instead of disabling autocast Per review: replace the torch.autocast(enabled=False) guard with a float64 computation, which autocast does not downcast — matching the float64 rope path used elsewhere (Flux). The autocast and float32 paths stay bit-identical (max|delta|=0). * Disable autocast for Ideogram4 rope matmul instead of using float64 Per review, use torch.autocast(enabled=False) around the rotary matmul (as the original implementation did) rather than computing in float64, and adopt the clearer comment describing the bfloat16 collapse at the 65536 offset. * Disable autocast for ernie_image and helios rope einsum Extend the Ideogram4 fix: ernie_image's `rope` and helios's `get_frequency_batched` build rotary freqs with an unguarded float32 einsum over raw position ids. Under an ambient autocast the einsum runs in bfloat16 on CUDA, which cannot represent consecutive integers past 256, so positions degrade — the same bug, matching the guards mochi/omnigen already have. Wrap each in torch.autocast(enabled=False). * Disable autocast for Cosmos3 VL-text rope matmul Cosmos3VLTextRotaryEmbedding builds its interleaved-mrope freqs with an unguarded position-id matmul (same shape as Ideogram4), so an ambient autocast downcasts it to bfloat16 and collapses positions past 256. Wrap in torch.autocast(enabled=False). * Tighten autocast(enabled=False) to just the rope matmul Per review, scope the guard to the precision-sensitive position-id matmul in Ideogram4 and Cosmos3 rather than the whole freqs block (ernie_image and helios already wrap only the einsum). --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Author
Parents
Loading