Fix Ideogram4MRoPE collapsing under torch.autocast (compute rotary in float32)

[HaozheZhang6]

What does this PR do?

Fixes #13920

Ideogram4MRoPE produces collapsed rotary embeddings under torch.autocast, so denoising inside an autocast context (common in training, and when users wrap pipeline calls) renders a flat single-color image.

Root cause

Image-token positions are IMAGE_POSITION_OFFSET (65536) + (t, h, w). Ideogram4MRoPE.forward casts its operands to float32, but the frequency matmul is on autocast's downcast list, so under torch.autocast("cuda", torch.bfloat16) it executes in bfloat16 anyway. bfloat16's representable step at 65536 is 512, so every image position in a ≤512-wide grid rounds to the same value — all image tokens get identical rotary embeddings, spatial information is lost, and sampling degenerates to a flat field.

Reproduced with the weight-free snippet from the issue (max |cos_autocast − cos_fp32| ≈ 1.93, distinct positions become equal).

Fix

Wrap the frequency computation in torch.autocast(device_type=..., enabled=False) so the rotary embeddings are always computed in float32 regardless of an ambient autocast — the same guard transformers applies to its RoPE modules. After the fix the autocast and float32 paths are bit-identical (max |Δ| = 0.0).

Scope is Ideogram4MRoPE, the catastrophic case (others noted in the issue are far milder without the 65536 offset). Happy to extend the same guard to the sibling RoPE modules in a follow-up if you'd like.

Tests

Added test_ideogram4_mrope_is_autocast_invariant — it fails on main (collapsed positions) and passes with the fix. Full file green:

pytest tests/models/transformers/test_models_transformer_ideogram4.py
38 passed, 4 skipped

Before submitting

• Did you read the contributor guideline?
• Did you read our philosophy doc?
• Was this discussed/approved via a GitHub issue? Ideogram4: Ideogram4MRoPE breaks under torch.autocast: all image positions collapse, producing flat single-color images #13920
• Did you make sure to update the documentation with your changes? (n/a — bug fix)
• Did you write any new necessary tests?

Who can review?

@DN6 @sayakpaul

[dxqb]

before committing that (and thereby closing my report), please consider that other modules might be affected, just not as bad. bfloat16 becomes inaccurate for integers starting 257.0 (which is rounded to 256.0).

that's within the range of text token ids

[HaozheZhang6]

You're right — confirmed bf16 rounds 257→256, 259→260, so text positions past 256 lose precision in any RoPE that matmuls raw position ids under autocast. Ideogram4 is just the pathological case: the 65536 offset collapses a whole ≤512-wide grid onto a single value, where the others degrade gradually instead of all-at-once.

I'd checked the other diffusers transformers — Ideogram4 is the only RoPE with a large position offset, so the only catastrophic one — but the gradual loss you describe is real for the rest. I can extend the same autocast(enabled=False) guard to the other RoPE forwards in this PR, or keep this one targeted at the Ideogram4 regression and do a follow-up sweep, whichever you and the maintainers prefer. Either way it shouldn't close your report until the broader case is covered.

[sayakpaul]

We don't use autocast within our modeling implementation like this.

[HaozheZhang6]

Good catch — dropped the autocast guard and compute the freqs in float64 instead, which autocast doesn't downcast (matching the float64 rope path Flux uses). The autocast and float32 paths come out bit-identical (max|Δ| = 0), and the regression test still passes.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.