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>