apple / apple/coreai-models

Composite RoPE partial-rotary mode pairs contiguously — incompatible with transformers partial/proportional rotary (partial_rotary_factor models)

Open
#66 2 comments 0 reactions 1 assignee Claimed by @stikves View on GitHub
Dominant language
Swift
Stars
2.1k
Forks
202
Avg merge
2d 8h
Merged PRs (30d)
55

Description

## Summary

The composite `RoPE` partial-rotary mode (`coreai_torch.composite_ops.RoPE` with `dims < head_dim`, as wrapped by `coreai_models/primitives/macos/rope.py`) pairs dimensions in a **contiguous block** (dim `i` ↔ `i + dims/2`, *inside* the first `dims` dims, passing the rest through). HuggingFace `transformers`' **partial / "proportional" rotary** (used by any model with `partial_rotary_factor < 1`) instead pairs across the **full head_dim half-split** (dim `i` ↔ `i + head_dim/2`), with only the first `rope_angles` frequencies non-zero (`inv_freq` zero-padded). The **frequencies are identical; only the dim pairing differs**, so the result is silently wrong for partial-rotary models.

Surfaced while porting **Gemma-4** (`partial_rotary_factor = 0.25`, `global_head_dim = 512`) — its global/full-attention layers use partial rotary. (Gemma-4 isn't in this repo yet; filing because the issue is in the shared primitive and affects *any* partial-rotary model.)

## Where

A recipe that folds the proportional factor into the composite as `RoPE(base**prf, dims=2*rope_angles)` gets the contiguous-block partial mode (`_rope_with_cos_and_sin_impl`), not the transformers layout.

## Repro (bit-level)

Apply the composite partial RoPE vs `transformers` `Gemma4TextRotaryEmbedding(..., "full_attention")` + `apply_rotary_pos_emb` to the same tensor at any position > 0:

- Sliding (full-rotary) layers match bit-exact (PSNR ∞).
- Global (partial-rotary) layer **diverges: PSNR ≈ 21.6 dB, max-abs ≈ 8.2** (gemma-4-26B-A4B, `head_dim=512`, `partial_rotary_factor=0.25` → it rotates dims {0..63}∪{256..319}, not {0..127}).

## Impact

Silently degrades every full-attention layer of the exported model. Generation stays *coherent* (global layers are ~1/6 of the stack), so it passes a smoke test — but it isn't faithful to the reference, and it **breaks EAGLE/MTP speculative-draft acceptance** (a draft trained against the reference layout stops matching the mis-rotated target).

## Reference fix (verified bit-exact)

Replace the partial branch with a small module that reproduces transformers exactly:

```python
inv_freq = cat([1/base**(arange(0, 2*rope_angles, 2)/head_dim),
zeros(head_dim//2 - rope_angles)]) # zero-padded, full-head
ang = positions[..., None] * inv_freq
emb = cat([ang, ang], dim=-1) # full-head half-split
x_rot = x * emb.cos() + rotate_half(x) * emb.sin()
```

After this, both sliding and global RoPE match the reference bit-exact (PSNR ∞).

## Secondary (related precision footgun)

If `inv_freq` is stored as a **registered buffer**, `model.to(bfloat16)` downcasts it and bf16's ~3-digit mantissa corrupts the frequencies (cos error ≈ 0.35 at position 200). Recomputing `inv_freq` in fp32 inside `forward` avoids it. Worth a guard/note for any rotary buffer in a bf16-compute path.

---
*Found while building [caix](https://github.com/RedHillsMediaFL/caix) (native Core AI serving + speculative decoding). Happy to provide the full repro script or open a PR if Gemma-4 / a partial-rotary model lands upstream.*

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.