Commit Graph

1 Commits

Author SHA1 Message Date
John
384e87fca1 DiT blocks: RoPE, per-head RMS norm, AdaLN modulation (23 tests total)
Unlocks all four Pixal3D flow checkpoints (~20GB) at once - they are pure transformers
with no sparse conv. LATO.2 has flow models too (vertex_structured_flow, topo_flow), so
these belong in the shared core rather than either port.

Three details taken from upstream rather than assumed, each silent when wrong:
- norm1/norm3 are NON-affine but norm2 IS affine in the modulated cross block. There is
  an explicit test asserting that asymmetry.
- MultiHeadRMSNorm is written upstream as F.normalize(x)*gamma*sqrt(dim). F.normalize is
  L2, and the sqrt(d) turns it into RMS - implemented directly as RMS and verified equal
  to the upstream formulation to 9.5e-7.
- RoPE phases are NOT derived: Pixal3D ships rope_phases as a stored tensor, so they are
  passed in. Tested that the rotation preserves per-pair norms and is not a no-op.

Also tested: gates at zero make the block an identity on its residual branches.
2026-08-02 11:33:04 +10:00