diff --git a/tests/test_dit.py b/tests/test_dit.py index b01d724..8cebf98 100644 --- a/tests/test_dit.py +++ b/tests/test_dit.py @@ -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: diff --git a/trellis_sparse_mlx/__init__.py b/trellis_sparse_mlx/__init__.py index d9a4113..7f8bc5d 100644 --- a/trellis_sparse_mlx/__init__.py +++ b/trellis_sparse_mlx/__init__.py @@ -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" diff --git a/trellis_sparse_mlx/dit.py b/trellis_sparse_mlx/dit.py index 70b78e3..e4fc41f 100644 --- a/trellis_sparse_mlx/dit.py +++ b/trellis_sparse_mlx/dit.py @@ -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)