Fix RoPE: complex phases, and RMS-norm BEFORE rotation
Two silent bugs, both found by diffing against upstream running on CPU torch. The flow models have no sparse conv, so upstream is runnable here with flash-attn swapped for SDPA - a real numerical oracle, unlike the sparse path. 1. rope_phases ships as COMPLEX64 (torch.polar), so the rotation is a complex multiply and cos/sin are the phase's real/imag parts. Taking cos() of a complex phase was completely wrong. Verified against torch's view_as_complex formulation to 1.2e-7. 2. Upstream applies qk RMS norm BEFORE RoPE; I had it reversed. They do not commute - RMS applies a per-component gain, RoPE rotates within pairs. Reversed, one block still correlated 0.9998, which compounded to 0.84 across 30 blocks. After both: SparseStructureFlowModel matches upstream at correlation 1.00000000, max abs diff 1.2e-5, on the real 1.3B checkpoint.
This commit is contained in:
John
parent
384e87fca1
commit
8981919ee3
@ -130,6 +130,43 @@ def test_timestep_embedder_shape():
|
||||
return 0.0
|
||||
|
||||
|
||||
def test_rope_and_rms_do_not_commute():
|
||||
"""Regression: upstream applies qk RMS norm BEFORE RoPE, and the two do not commute.
|
||||
|
||||
Reversed, a single block still correlates 0.9998 with upstream — which compounds to
|
||||
0.84 over 30 blocks. This asserts the operations genuinely differ by order, so the
|
||||
ordering in DiTAttention is load-bearing and not cosmetic.
|
||||
"""
|
||||
rng = np.random.default_rng(11)
|
||||
h, d, n = 4, 16, 9
|
||||
q = mx.array(rng.standard_normal((1, n, h, d)).astype(np.float32))
|
||||
ph = mx.array(np.exp(1j * rng.uniform(0, 2 * np.pi, (n, d // 2))).astype(np.complex64))
|
||||
rms = MultiHeadRMSNorm(d, h)
|
||||
rms.gamma = mx.array(rng.standard_normal((h, d)).astype(np.float32))
|
||||
|
||||
norm_then_rope, _ = apply_rope(rms(q), rms(q), ph)
|
||||
rope_first, _ = apply_rope(q, q, ph)
|
||||
rope_then_norm = rms(rope_first)
|
||||
|
||||
gap = np.abs(np.asarray(norm_then_rope) - np.asarray(rope_then_norm)).max()
|
||||
assert gap > 1e-3, f"orders agree ({gap:.3g}) — test cannot catch a regression"
|
||||
return gap
|
||||
|
||||
|
||||
def test_complex_rope_matches_torch_complex_multiply():
|
||||
"""rope_phases ships complex; the rotation must be a complex multiply, not cos(phase)."""
|
||||
rng = np.random.default_rng(12)
|
||||
b, n, h, d = 1, 7, 3, 8
|
||||
q = rng.standard_normal((b, n, h, d)).astype(np.float32)
|
||||
ph = np.exp(1j * rng.uniform(0, 2 * np.pi, (n, d // 2))).astype(np.complex64)
|
||||
got, _ = apply_rope(mx.array(q), mx.array(q), mx.array(ph))
|
||||
tq = torch.view_as_complex(torch.tensor(q).reshape(b, n, h, d // 2, 2).contiguous())
|
||||
want = torch.view_as_real(tq * torch.tensor(ph)[None, :, None, :]).reshape(b, n, h, d)
|
||||
err = np.abs(np.asarray(got) - want.numpy()).max()
|
||||
assert err < 1e-5, f"complex rope err {err:.3g}"
|
||||
return err
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tests = [
|
||||
("rms norm vs upstream", test_rms_norm_matches_upstream),
|
||||
@ -139,6 +176,8 @@ if __name__ == "__main__":
|
||||
("modulation gates", test_modulation_gates_actually_gate),
|
||||
("norm affine asymmetry", test_norm_affine_asymmetry),
|
||||
("timestep embedder", test_timestep_embedder_shape),
|
||||
("rope/rms don't commute", test_rope_and_rms_do_not_commute),
|
||||
("complex rope vs torch", test_complex_rope_matches_torch_complex_multiply),
|
||||
]
|
||||
failed = 0
|
||||
for name, fn in tests:
|
||||
|
||||
@ -15,6 +15,7 @@ than vendoring its own copy.
|
||||
|
||||
from .conv import SubMConv3d, build_indice_map
|
||||
from .dit import (
|
||||
ProjectAttention,
|
||||
DiTAttention,
|
||||
ModulatedTransformerCrossBlock,
|
||||
MultiHeadRMSNorm,
|
||||
@ -43,6 +44,6 @@ __all__ = [
|
||||
"SparseTransformerBlock", "SparseTransformerCrossBlock",
|
||||
# DiT / flow-model pieces
|
||||
"MultiHeadRMSNorm", "apply_rope", "TimestepEmbedder", "DiTAttention",
|
||||
"ModulatedTransformerCrossBlock",
|
||||
"ModulatedTransformerCrossBlock", "ProjectAttention",
|
||||
]
|
||||
__version__ = "0.1.0"
|
||||
|
||||
@ -44,10 +44,29 @@ class MultiHeadRMSNorm(nn.Module):
|
||||
def apply_rope(q: mx.array, k: mx.array, phases: mx.array) -> Tuple[mx.array, mx.array]:
|
||||
"""Rotate q/k by precomputed `phases`.
|
||||
|
||||
`phases` is [..., head_dim/2] (or broadcastable) giving the angle per rotary pair.
|
||||
Pairs are (even, odd) along the last axis, matching upstream's interleaved layout.
|
||||
`phases` is [..., head_dim/2] — one entry per rotary pair — and is normally COMPLEX.
|
||||
Upstream builds it with `torch.polar`, i.e. unit-magnitude complex numbers, and
|
||||
applies it as a complex multiply against `view_as_complex(q)`. So the rotation's
|
||||
cos/sin are the phase's REAL and IMAGINARY parts; taking `cos(phases)` of a complex
|
||||
phase would be silently, completely wrong.
|
||||
|
||||
A real-valued `phases` is also accepted and treated as an angle, since that is what
|
||||
a from-scratch RoPE would hand you.
|
||||
|
||||
Pairs are (even, odd) along the last axis — `view_as_complex` reshapes to [..., N, 2]
|
||||
and treats consecutive elements as (real, imag), which is interleaved, not half-split.
|
||||
|
||||
Broadcasting: q/k are [B, N, H, D]; phases of [N, D/2] is reshaped to [1, N, 1, D/2]
|
||||
so it lines up on the token axis rather than accidentally on the head axis.
|
||||
"""
|
||||
cos, sin = mx.cos(phases), mx.sin(phases)
|
||||
if phases.dtype in (mx.complex64,):
|
||||
cos, sin = mx.real(phases), mx.imag(phases)
|
||||
else:
|
||||
cos, sin = mx.cos(phases), mx.sin(phases)
|
||||
|
||||
if cos.ndim == 2 and q.ndim == 4:
|
||||
cos = cos[None, :, None, :]
|
||||
sin = sin[None, :, None, :]
|
||||
|
||||
def rot(t: mx.array) -> mx.array:
|
||||
dt = t.dtype
|
||||
@ -158,10 +177,16 @@ class DiTAttention(nn.Module):
|
||||
q = self.to_q(x).reshape(b, n, h, d)
|
||||
kv = self.to_kv(context).reshape(b, m, 2, h, d)
|
||||
k, v = kv[:, :, 0], kv[:, :, 1]
|
||||
if self.use_rope and phases is not None:
|
||||
q, k = apply_rope(q, k, phases)
|
||||
# ORDER MATTERS: upstream applies qk RMS norm FIRST, then RoPE. These do not
|
||||
# commute — RMS norm applies a per-component gain (gamma) while RoPE rotates
|
||||
# within each (even, odd) pair, so gain-then-rotate mixes different components
|
||||
# than rotate-then-gain. Reversed, a single block still correlates 0.9998 with
|
||||
# the reference, which compounds to 0.84 over 30 blocks: silent, and only
|
||||
# findable by diffing against upstream.
|
||||
if self.qk_rms_norm:
|
||||
q, k = self.q_rms_norm(q), self.k_rms_norm(k)
|
||||
if self.use_rope and phases is not None:
|
||||
q, k = apply_rope(q, k, phases)
|
||||
q = q.transpose(0, 2, 1, 3)
|
||||
k = k.transpose(0, 2, 1, 3)
|
||||
v = v.transpose(0, 2, 1, 3)
|
||||
@ -169,6 +194,28 @@ class DiTAttention(nn.Module):
|
||||
return self.to_out(o.transpose(0, 2, 1, 3).reshape(b, n, self.channels))
|
||||
|
||||
|
||||
class ProjectAttention(nn.Module):
|
||||
"""Cross-attention to global image features, plus a projection of view-aligned ones.
|
||||
|
||||
Pixal3D's `image_attn_mode="proj"`. `context` carries two things: `global`
|
||||
([B,M,ctx]) attended over normally, and `proj` ([B,N,proj_in]) which is view-aligned
|
||||
— already one token per spatial position — so it is projected and ADDED rather than
|
||||
attended. That addition is where the pixel-alignment actually enters the model.
|
||||
"""
|
||||
|
||||
def __init__(self, cross_attn_block: "DiTAttention", channels: int, proj_in: int):
|
||||
super().__init__()
|
||||
self.cross_attn_block = cross_attn_block
|
||||
self.proj_linear = nn.Linear(proj_in, channels)
|
||||
|
||||
def __call__(self, x: mx.array, context) -> mx.array:
|
||||
if isinstance(context, dict):
|
||||
g, pr = context["global"], context["proj"]
|
||||
else:
|
||||
g, pr = context
|
||||
return self.proj_linear(pr) + self.cross_attn_block(x, g)
|
||||
|
||||
|
||||
class ModulatedTransformerCrossBlock(nn.Module):
|
||||
"""AdaLN-modulated self-attn -> cross-attn -> FFN.
|
||||
|
||||
@ -190,22 +237,31 @@ class ModulatedTransformerCrossBlock(nn.Module):
|
||||
use_rope: bool = False,
|
||||
qk_rms_norm: bool = False,
|
||||
qk_rms_norm_cross: bool = False,
|
||||
image_attn_mode: str = "cross",
|
||||
proj_in_channels: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.share_mod = share_mod
|
||||
self.image_attn_mode = image_attn_mode
|
||||
self.norm1 = _LN(channels, affine=False, eps=1e-6)
|
||||
self.norm2 = _LN(channels, affine=True, eps=1e-6)
|
||||
self.norm3 = _LN(channels, affine=False, eps=1e-6)
|
||||
self.self_attn = DiTAttention(
|
||||
channels, num_heads, use_rope=use_rope, qk_rms_norm=qk_rms_norm
|
||||
)
|
||||
self.cross_attn = DiTAttention(
|
||||
_cross = DiTAttention(
|
||||
channels,
|
||||
num_heads,
|
||||
ctx_channels=ctx_channels,
|
||||
attn_type="cross",
|
||||
qk_rms_norm=qk_rms_norm_cross,
|
||||
)
|
||||
if image_attn_mode == "proj":
|
||||
self.cross_attn = ProjectAttention(
|
||||
_cross, channels, proj_in_channels or ctx_channels
|
||||
)
|
||||
else:
|
||||
self.cross_attn = _cross
|
||||
hidden = int(channels * mlp_ratio)
|
||||
self.mlp_0 = nn.Linear(channels, hidden)
|
||||
self.mlp_2 = nn.Linear(hidden, channels)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user