a44a0d7360
ss_flow ships rope_phases precomputed; the SLAT flows do not, because their positions are the input SparseTensor's own coordinates and vary per input. rope_phases_from_coords reproduces RotaryPositionEmbedder, including the pad detail: 3*21=63 frequencies fall one short of head_dim//2=64 and upstream right-pads with 1+0j, so dropping it would shift every later pair by a slot. Verified against ss_flow's own shipped tensor to 9.6e-7 - a real ground truth, since the same embedder produced it.
221 lines
8.7 KiB
Python
221 lines
8.7 KiB
Python
"""DiT block tests against torch.
|
|
|
|
These layers all have torch counterparts, so unlike the submanifold conv they are checked
|
|
against upstream's real semantics. The upstream forward bodies are reproduced verbatim
|
|
from pixal3d/modules/{attention/modules.py,transformer/modulated.py}.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from trellis_sparse_mlx.dit import ( # noqa: E402
|
|
DiTAttention,
|
|
ModulatedTransformerCrossBlock,
|
|
MultiHeadRMSNorm,
|
|
TimestepEmbedder,
|
|
apply_rope,
|
|
rope_phases_from_coords,
|
|
)
|
|
|
|
|
|
def test_rms_norm_matches_upstream(dim=32, heads=4, n=17):
|
|
"""Upstream writes it as F.normalize(x)*gamma*sqrt(dim); we implement RMS directly."""
|
|
rng = np.random.default_rng(0)
|
|
x = rng.standard_normal((n, heads, dim)).astype(np.float32)
|
|
g = rng.standard_normal((heads, dim)).astype(np.float32)
|
|
|
|
m = MultiHeadRMSNorm(dim, heads)
|
|
m.gamma = mx.array(g)
|
|
got = np.asarray(m(mx.array(x)))
|
|
|
|
tx = torch.tensor(x)
|
|
want = (F.normalize(tx.float(), dim=-1) * torch.tensor(g) * dim**0.5).numpy()
|
|
|
|
err = np.abs(got - want).max()
|
|
assert err < 2e-4, f"rms norm err {err:.3g}"
|
|
return err
|
|
|
|
|
|
def test_rms_norm_is_scale_invariant():
|
|
"""RMS norm must remove input magnitude — catches a plain scale-by-gamma stub."""
|
|
m = MultiHeadRMSNorm(8, 2)
|
|
x = np.random.default_rng(1).standard_normal((5, 2, 8)).astype(np.float32)
|
|
a = np.asarray(m(mx.array(x)))
|
|
b = np.asarray(m(mx.array(x * 37.0)))
|
|
err = np.abs(a - b).max()
|
|
assert err < 1e-3, f"not scale-invariant: {err:.3g}"
|
|
return err
|
|
|
|
|
|
def test_rope_is_a_rotation():
|
|
"""RoPE must preserve the norm of each rotary pair."""
|
|
rng = np.random.default_rng(2)
|
|
q = rng.standard_normal((2, 6, 4, 16)).astype(np.float32)
|
|
ph = rng.uniform(0, 2 * np.pi, (2, 6, 4, 8)).astype(np.float32)
|
|
rq, _ = apply_rope(mx.array(q), mx.array(q), mx.array(ph))
|
|
rq = np.asarray(rq)
|
|
n0 = np.sqrt(q[..., 0::2] ** 2 + q[..., 1::2] ** 2)
|
|
n1 = np.sqrt(rq[..., 0::2] ** 2 + rq[..., 1::2] ** 2)
|
|
err = np.abs(n0 - n1).max()
|
|
assert err < 1e-4, f"rope changed pair norms by {err:.3g}"
|
|
# and it must actually rotate
|
|
assert np.abs(rq - q).max() > 1e-3, "rope was a no-op"
|
|
return err
|
|
|
|
|
|
def test_attention_matches_torch(ch=64, heads=8, n=13, b=2):
|
|
rng = np.random.default_rng(3)
|
|
x = rng.standard_normal((b, n, ch)).astype(np.float32)
|
|
wq = rng.standard_normal((ch * 3, ch)).astype(np.float32) * 0.05
|
|
bq = rng.standard_normal((ch * 3,)).astype(np.float32) * 0.05
|
|
wo = rng.standard_normal((ch, ch)).astype(np.float32) * 0.05
|
|
bo = rng.standard_normal((ch,)).astype(np.float32) * 0.05
|
|
|
|
a = DiTAttention(ch, heads)
|
|
a.to_qkv.weight, a.to_qkv.bias = mx.array(wq), mx.array(bq)
|
|
a.to_out.weight, a.to_out.bias = mx.array(wo), mx.array(bo)
|
|
got = np.asarray(a(mx.array(x)))
|
|
|
|
d = ch // heads
|
|
tx = torch.tensor(x)
|
|
qkv = F.linear(tx, torch.tensor(wq), torch.tensor(bq)).reshape(b, n, 3, heads, d)
|
|
q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
|
|
o = F.scaled_dot_product_attention(
|
|
q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
|
|
)
|
|
o = o.permute(0, 2, 1, 3).reshape(b, n, ch)
|
|
want = F.linear(o, torch.tensor(wo), torch.tensor(bo)).detach().numpy()
|
|
|
|
err = np.abs(got - want).max()
|
|
assert err < 2e-4, f"attention err {err:.3g}"
|
|
return err
|
|
|
|
|
|
def test_modulation_gates_actually_gate(ch=32, heads=4, n=7, b=2):
|
|
"""gate=0 must make the block an identity on the modulated paths."""
|
|
blk = ModulatedTransformerCrossBlock(ch, ch, heads, share_mod=True)
|
|
x = mx.array(np.random.default_rng(4).standard_normal((b, n, ch)).astype(np.float32))
|
|
ctx = mx.array(np.zeros((b, 5, ch), dtype=np.float32))
|
|
# zero everything, then force the cross-attn output to zero via zero out-proj
|
|
blk.modulation = mx.zeros((6 * ch,))
|
|
blk.cross_attn.to_out.weight = mx.zeros_like(blk.cross_attn.to_out.weight)
|
|
blk.cross_attn.to_out.bias = mx.zeros_like(blk.cross_attn.to_out.bias)
|
|
mod = mx.zeros((b, 6 * ch))
|
|
out = np.asarray(blk(x, mod, ctx))
|
|
err = np.abs(out - np.asarray(x)).max()
|
|
assert err < 1e-4, f"gates did not zero the residual branches: {err:.3g}"
|
|
return err
|
|
|
|
|
|
def test_norm_affine_asymmetry():
|
|
"""norm1/norm3 non-affine, norm2 affine — silent if swapped."""
|
|
blk = ModulatedTransformerCrossBlock(16, 16, 2, share_mod=True)
|
|
assert not blk.norm1.affine, "norm1 should be non-affine"
|
|
assert blk.norm2.affine, "norm2 SHOULD be affine"
|
|
assert not blk.norm3.affine, "norm3 should be non-affine"
|
|
return 0.0
|
|
|
|
|
|
def test_timestep_embedder_shape():
|
|
te = TimestepEmbedder(64, 32)
|
|
out = te(mx.array(np.array([0.0, 0.5, 1.0], dtype=np.float32)))
|
|
assert out.shape == (3, 64), out.shape
|
|
assert np.isfinite(np.asarray(out)).all()
|
|
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
|
|
|
|
|
|
def test_derived_rope_matches_shipped_tensor():
|
|
"""The derivation must reproduce Pixal3D's own precomputed rope_phases.
|
|
|
|
ss_flow ships `rope_phases` for a 16^3 grid, so there is a ground-truth tensor to
|
|
check against — including the detail that 3*21=63 frequencies fall one short of
|
|
head_dim//2=64 and upstream right-pads with 1+0j.
|
|
"""
|
|
ck = Path(
|
|
"/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts/"
|
|
"ss_flow_img_dit_1_3B_64_bf16.safetensors"
|
|
)
|
|
if not ck.exists():
|
|
print(" (skipped: checkpoint not present)")
|
|
return 0.0
|
|
ref = np.asarray(mx.load(str(ck))["rope_phases"])
|
|
res = 16
|
|
zz, yy, xx = np.meshgrid(*[np.arange(res)] * 3, indexing="ij")
|
|
coords = np.stack([zz, yy, xx], -1).reshape(-1, 3).astype(np.int32)
|
|
got = np.asarray(rope_phases_from_coords(mx.array(coords), head_dim=128))
|
|
err = np.abs(got - ref).max()
|
|
assert err < 1e-5, f"derived phases differ from shipped by {err:.3g}"
|
|
return err
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tests = [
|
|
("rms norm vs upstream", test_rms_norm_matches_upstream),
|
|
("rms scale invariance", test_rms_norm_is_scale_invariant),
|
|
("rope is a rotation", test_rope_is_a_rotation),
|
|
("dit attention vs torch", test_attention_matches_torch),
|
|
("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),
|
|
("derived rope vs shipped", test_derived_rope_matches_shipped_tensor),
|
|
]
|
|
failed = 0
|
|
for name, fn in tests:
|
|
try:
|
|
err = fn()
|
|
print(f" PASS {name:26s} (max err {err:.2e})")
|
|
except AssertionError as e:
|
|
print(f" FAIL {name:26s} {e}")
|
|
failed += 1
|
|
except Exception as e: # noqa: BLE001
|
|
print(f" ERROR {name:26s} {type(e).__name__}: {e}")
|
|
failed += 1
|
|
print(f"\n{len(tests)-failed}/{len(tests)} passed")
|
|
sys.exit(1 if failed else 0)
|