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: HaozheZhang6

src/diffusers/models/transformers/transformer_cosmos3.py

@@ -128,7 +128,12 @@ def forward(self, position_ids, device, dtype):
             self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1).to(device)
         )  # [3,B,head_dim//2,1]
         position_ids_expanded = position_ids[:, :, None, :].float()  # [3,B,1,N]
-        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3)  # [3,B,N,head_dim//2]
+        # Disable autocast so the position-id matmul runs in float32: under an ambient autocast it would run in
+        # bfloat16, which cannot represent consecutive integers past 256, collapsing positions onto the same
+        # frequency and degrading the rotary embedding.
+        with torch.autocast(device_type=position_ids.device.type, enabled=False):
+            freqs = inv_freq_expanded @ position_ids_expanded
+        freqs = freqs.transpose(2, 3)  # [3,B,N,head_dim//2]
         freqs = self.apply_interleaved_mrope(freqs, self.rope_axes_dim)  # [B,N,head_dim//2]
         emb = torch.cat((freqs, freqs), dim=-1)  # [B,N,head_dim]
         return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype)  # each: [B,N,head_dim]

src/diffusers/models/transformers/transformer_ernie_image.py

@@ -47,7 +47,11 @@ def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
     assert dim % 2 == 0
     scale = torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device) / dim
     omega = 1.0 / (theta**scale)
-    out = torch.einsum("...n,d->...nd", pos, omega)
+    # Disable autocast so the position-id einsum runs in float32: under an ambient autocast it would run in
+    # bfloat16, which cannot represent consecutive integers past 256, so position ids beyond that point would
+    # collapse onto the same frequency and degrade the rotary embedding.
+    with torch.autocast(device_type=pos.device.type, enabled=False):
+        out = torch.einsum("...n,d->...nd", pos, omega)
     return out.float()
 
 

src/diffusers/models/transformers/transformer_helios.py

@@ -337,7 +337,11 @@ def _get_freqs_base(self, dim):
 
     @torch.no_grad()
     def get_frequency_batched(self, freqs_base, pos):
-        freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos)
+        # Disable autocast so the position-grid einsum runs in float32: under an ambient autocast it would run
+        # in bfloat16, which cannot represent consecutive integers past 256, so positions beyond that point
+        # would collapse onto the same frequency and degrade the rotary embedding.
+        with torch.autocast(device_type=pos.device.type, enabled=False):
+            freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos)
         freqs = freqs.repeat_interleave(2, dim=0)
         return freqs.cos(), freqs.sin()
 

src/diffusers/models/transformers/transformer_ideogram4.py

@@ -70,9 +70,14 @@ def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso
             raise ValueError(f"`position_ids` must have shape (B, L, 3), got {tuple(position_ids.shape)}.")
         batch_size, seq_len, _ = position_ids.shape
 
+        # Ideogram4's image position ids start at IMAGE_POSITION_OFFSET (65536). If an ambient autocast downcasts the
+        # matmul to bfloat16, the image positions will collapse to only a few distinct values because bfloat16 cannot
+        # represent consecutive integers at this value (after pos 65536 each 512-integer block will collapse to the
+        # same value), which causes the image to become essentially flat. Therefore, we need to disable autocast here.
         pos = position_ids.permute(2, 0, 1).to(dtype=torch.float32)
         inv_freq = self.inv_freq.to(dtype=torch.float32)[None, None, :, None].expand(3, batch_size, -1, 1)
-        freqs = inv_freq @ pos.unsqueeze(2)
+        with torch.autocast(device_type=position_ids.device.type, enabled=False):
+            freqs = inv_freq @ pos.unsqueeze(2)
         freqs = freqs.transpose(2, 3)  # (3, B, L, inv_freq_size)
 
         # Interleaved mrope: pull H freqs into idx 1 mod 3, W freqs into idx 2 mod 3.
@@ -83,7 +88,7 @@ def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso
             freqs_t[..., idx] = freqs[axis][..., idx]
 
         emb = torch.cat((freqs_t, freqs_t), dim=-1)
-        return emb.cos(), emb.sin()
+        return emb.cos().float(), emb.sin().float()
 
 
 class Ideogram4AttnProcessor:

tests/models/transformers/test_models_transformer_ideogram4.py

@@ -21,6 +21,7 @@
     IMAGE_POSITION_OFFSET,
     LLM_TOKEN_INDICATOR,
     OUTPUT_IMAGE_INDICATOR,
+    Ideogram4MRoPE,
 )
 from diffusers.utils.torch_utils import randn_tensor
 
@@ -164,3 +165,20 @@ def test_gradient_checkpointing_is_applied(self):
 
 class TestIdeogram4TransformerAttention(Ideogram4TransformerTesterConfig, AttentionTesterMixin):
     """Attention processor tests for Ideogram 4 Transformer."""
+
+
+def test_ideogram4_mrope_is_autocast_invariant():
+    # Ideogram4's image positions start at IMAGE_POSITION_OFFSET (65536), so the rotary matmul must
+    # run in float32: under an ambient autocast it would otherwise execute in bfloat16 and round every
+    # image position to the same value, collapsing all spatial information (the decoded image goes flat).
+    rope = Ideogram4MRoPE(head_dim=256, base=5_000_000, mrope_section=(24, 20, 20)).to(torch_device)
+    position_ids = torch.tensor([[[0, 0, 0], [0, 0, 1], [0, 63, 63]]], device=torch_device) + IMAGE_POSITION_OFFSET
+
+    cos_ref, sin_ref = rope(position_ids)
+    with torch.autocast(device_type=torch.device(torch_device).type, dtype=torch.bfloat16):
+        cos_ac, sin_ac = rope(position_ids)
+
+    # Distinct image positions must keep distinct embeddings, identical to the float32 computation.
+    assert not torch.equal(cos_ac[0, 0], cos_ac[0, 1])
+    assert torch.equal(cos_ac, cos_ref)
+    assert torch.equal(sin_ac, sin_ref)