Qwen2.5-Omni Token2Wav DiT uses mismatched RoPE layout (half-split cos/sin with interleaved rotate) — regression from #39847
Author: HenryVarro666
Pull request: no

### System Info
- `transformers`: main (reproduced at `ab1771c`); bug present since the v5 RoPE refactor (#39847, merged 2025-10-17)
- Code-level bug, weight-independent — reproducible on CPU, no GPU or pretrained weights needed

### Who can help?
@eustlb @ebezzam

### Description
`Qwen2_5OmniDiTRotaryEmbedding.forward` builds the rotary `cos`/`sin` with a **half-split** layout:
```python
# modeling_qwen2_5_omni.py:2503
emb = torch.cat((freqs, freqs), dim=-1)   # [f0, f1, …, f0, f1, …]
```
but the Token2Wav DiT applies an **interleaved** rotate function, `rotate_half_codec`, which pairs adjacent channels `(2i, 2i+1)`:
```python
# modeling_qwen2_5_omni.py:2940
def rotate_half_codec(x):
    x = x.reshape(*x.shape[:-1], -1, 2)
    x1, x2 = x.unbind(dim=-1)
    x = torch.stack((-x2, x1), dim=-1)
    return x.reshape(*x.shape[:-2], -1)
```
A half-split cos/sin combined with an interleaved rotate is **not a valid rotation**: RoPE's defining property `⟨R_m·q, R_n·k⟩` depending only on the relative offset `m − n` is destroyed. This is applied to head 0 of every DiT attention layer (`query[:, :1], key[:, :1] = apply_rotary_pos_emb(...)`), silently degrading the predicted mel-spectrogram and thus the synthesized audio for `return_audio=True`.

This is a **regression**. The DiT rotary originally used an interleaved layout (`torch.stack((freqs, freqs), dim=-1).reshape(...)`, i.e. `repeat_interleave`), which matches `rotate_half_codec`. PR #39847 ("🚨 [v5] Refactor RoPE for layer types") made the class inherit `LlamaRotaryEmbedding`, replacing `stack` with `cat` while leaving `rotate_half_codec` interleaved. The original add PR #36752 (integration-tested against `Qwen/Qwen2.5-Omni-7B`) used the interleaved layout, so the released weights expect it. (The text model's mrope at line ~1380 correctly uses `cat`, because it uses the half-split `rotate_half` — that one is fine.)

### Reproduction (CPU, no weights)
```python
import torch
dim, base = 16, 10000.0
positions = torch.arange(0, 64, dtype=torch.float)
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
freqs = positions[:, None] * inv_freq[None, :]

def rotate_half_codec(x):  # verbatim from modeling_qwen2_5_omni.py
    x = x.reshape(*x.shape[:-1], -1, 2); x1, x2 = x.unbind(dim=-1)
    x = torch.stack((-x2, x1), dim=-1); return x.reshape(*x.shape[:-2], -1)
def apply(q, cos, sin): return q * cos + rotate_half_codec(q) * sin

def spread(emb, delta=5):  # max drift of <R_m q, R_{m+delta} k> over absolute m; ~0 == valid RoPE
    cos, sin = emb.cos(), emb.sin(); q = torch.randn(dim); k = torch.randn(dim); vals = []
    for m in range(40):
        vals.append((apply(q, cos[m], sin[m]) * apply(k, cos[m+delta], sin[m+delta])).sum().item())
    v = torch.tensor(vals); return (v.max() - v.min()).item()

print("current (cat):            ", spread(torch.cat((freqs, freqs), dim=-1)))            # ~5.1  (broken)
print("fixed   (repeat_interleave):", spread(torch.repeat_interleave(freqs, 2, dim=-1)))  # ~1e-6 (valid)
```

### Expected behavior
The DiT rotary cos/sin layout should match its interleaved rotate function, so RoPE is translation-invariant and consistent with the released checkpoint.

### Suggested fix (one line, checkpoint-faithful)
```diff
-            emb = torch.cat((freqs, freqs), dim=-1)
+            emb = torch.repeat_interleave(freqs, 2, dim=-1)
```
(in the modular source, override `forward` on `Qwen2_5OmniDiTRotaryEmbedding` instead of inheriting Llama's, so the generated forward emits the interleaved layout). I have this fix ready and am happy to open a PR if you agree with the direction.

> Disclosure: the root-cause analysis, regression bisect, and reproduction in this report were prepared with AI assistance and independently verified by me.


Comments:
Excellent analysis @HenryVarro666 — the half-split/interleaved mismatch, the bisect to #39847, and the translation-invariance repro all check out on current main (I measured the same ~5.1 → ~1e-6 drift with your script, and 9.9 → 6.7e-6 when driving the actual `Qwen2_5OmniDiTRotaryEmbedding` + `apply_rotary_pos_emb`). Also confirmed `qwen3_omni_moe` is consistent and unaffected.

Since you mentioned having the fix ready: happy to defer to your PR — the diagnosis and direction are entirely yours. In case it helps move things faster, I have a branch (prepared with AI assistance, per the agentic-contribution guidelines) with the `repeat_interleave` override in the modular file plus CPU regression tests (interleaved-layout check + translation-invariance drift < 1e-4, both failing on main and passing with the fix) ready to go. @eustlb @ebezzam — does the checkpoint-faithful direction (fix the cos/sin layout, keep `rotate_half_codec`) sound right to you?


Thanks so much, @pepepeboom — I really appreciate you taking the time to reproduce this end-to-end. It's genuinely reassuring to have an independent confirmation: your ~9.9 → 6.7e-6 drift driving the actual `Qwen2_5OmniDiTRotaryEmbedding` + `apply_rotary_pos_emb` lines up nicely with the ~5.1 → 1e-6 I saw on the self-contained repro, and you're right that `qwen3_omni_moe` is consistent and out of scope here.

That's also very kind of you to offer your branch. I happen to have a fix ready on my end as well, so I'll go ahead and open the PR to keep things moving — but I'd genuinely value your review on it, and I'll be sure to credit your independent verification. It overrides `forward` on `Qwen2_5OmniDiTRotaryEmbedding` in the modular file to emit the interleaved layout (`repeat_interleave(freqs, 2)`, restoring the pre-#39847 `stack` behavior that matches the #36752 integration-tested weights), with CPU regression tests (interleaved-layout assertion + translation-invariance drift `< 1e-4`) that fail on `main` and pass with the fix.

@eustlb @ebezzam — on direction: I've gone with the checkpoint-faithful fix (correct the `cos`/`sin` layout to interleaved, matching the released weights, and leave `rotate_half_codec` unchanged) rather than rewriting the rotate to half-split. More than happy to adjust if you'd prefer a different approach.

> Disclosure: the root-cause analysis, regression bisect, and reproduction were prepared with AI assistance and independently verified by me.