Add UMA-S-1.2 support (charge-balanced channels) + edge-frame & energy-head fixes

[janweinreich]

Summary

Adds uma-s-1.2 support to TT-Atom (run via UMA(atoms, checkpoint="uma-s-1p2.pt")), plus fixes uncovered while validating it — rebased onto current master (after its recent perf/kernel overhaul). Full 757-system A/B vs fairchem: no genuine force errors, and 691/757 pass the strict bar (above the 672 the pre-perf-pass backbone reached). Draft, for review.

Six commits on current master:

Feature

1. uma-s-1.2 support: charge_balanced_channels + fp32 energy head — uma-s-1.2 needs fairchem's charge_balanced_channels (missing → force PCC ~0.83): a self-adjoint per-system l=0 mean-subtraction of channels [cs:ce], re-balanced to the system charge after each block. Threaded through the forward (Backbone.node_embedding/GraphContext.node_meanM), the reverse VJP (backbone_bw), the disjoint batch, and the traced loop (single + batched); export_weights writes charge_channel_start/end. calculate/from_uma already source charge/spin, so this threads the numeric charge into the force calls. A self-review also hardened the batched/multi-card paths: MultiCard builds the mean operator (no hang on uma-s-1.2), the traced batch uses the true per-system natoms, charged batches of unequal atom counts are rejected with a clear error, and the vendored Wigner-D asset now ships in the wheel (pyproject package-data). Also emits the energy head + node-sum in fp32 (the bf16-rounded reference-subtracted energy was the dominant device energy error on high-|raw| systems).

Correctness fixes

2. smooth quaternion edge frame + generic eval rotation — replaces the singular ZYZ-Euler edge→+Y frame (±Y-axis singularity → ~300 meV/Å forces at exactly-symmetric geometries) with fairchem's default two-chart quaternion frame (vendored, MIT). Symmetric-molecule force MAE ~300 → <8 meV/Å. (The trace-width crash the Euler zeros caused is independently handled on master now via rotation.gather_coef at the cached sparsity pattern, so only the frame is ported.)
3. fp32 radial-MLP backward — fixes a force regression in master's on-device-backward work. The radial-MLP LayerNorm VJP (RadialMLP._ln_bw) runs in bf16 by default, unlike the SH-norm/gate backwards which reuse fp32-accurate stats. On an OOD geometry the bf16 radial VJP mis-directs the force: el_Sn_cmp (2-atom compressed β-Sn) dropped a whole component → traj PCC 0.73 / 230 meV/Å, energy still correct. Diagnosed vs an fp64 mirror.radial_mlp oracle (device rad.bw PCC 0.35; same math in fp64 from the same bf16 cache → PCC 1.0, so it's the bf16 arithmetic). Fix: run the small radial backward in fp32 (fp32 weight copies incl. the transformed w6; forward stays bf16; the opt-in bf16-only fused_ln_bw kernel path is unchanged). Result: el_Sn_cmp → 1.0000, ~+4% eager. Adds the missing test_forces.py assertion that the device radial finish matches the oracle.

Docs / tests

4. chore — track the vendored tt_atom/assets/wigner_d_coefficients.pt.
5. test — gated uma-s-1.2 real-weight parity (auto-skips without the checkpoint) + gen_golden_real.py --ckpt-path.
6. docs — README note + docs/uma-s-1p2-validation.md.

Validation — full 757-system A/B vs fairchem on CPU (uma-s-1p2.pt), on current master

757 systems, 2–99 atoms, ~63 elements, molecular + 491 periodic (omol/omat/oc20/odac/omc; 175 oc20 surfaces, 40 nanoparticles, S22, radicals, charged TM complexes, EOS curves).

• No genuine force errors. 741/757 keep pooled force PCC > 0.999 (754/757 > 0.99); the 3 below 0.99 are the zero-force perfect-lattice artifact (|F|≈0 by symmetry, traj Fmae < 1.6 meV/Å). Showcase systems — Pt(111) CO-oxidation, Cu₅₅+CO, [Cu(EDTA)]²⁻ bare + solvated — are all PCC 1.0000. The one regression master's bf16 radial backward introduced (el_Sn_cmp) is fixed by commit 3.
• 691 / 757 pass the strict bar (force PCC > 0.99, energy < 5 meV/atom, stress PCC > 0.99), above the 672 the earlier backbone reached — master's numerics plus the fp32 head net lower the bf16 energy floor. The misses are that floor on the highest-|raw| systems (heavy-metal diatomics/crystals, up to ~11 meV/atom), forces PCC ~1.0, all under chemical accuracy (~43).
• Stress (periodic): median PCC 0.99996, 488/491 ≥ 0.99.

Performance — MgO supercells, ms/call, TT (Wormhole, ttnn 0.68.0) vs 16-thread CPU

atoms	TT eager	TT trace	CPU (16t)	CPU / TT-best
8	103	75	138	1.8×
64	317	319	860	2.7×
216	1059	973	5326	5.5×
512	2794	OOM	—	—

Up to 5.5× at 216 atoms (trace); ~2.2× faster than the pre-perf-pass backbone. master's kernel work restores the trace path at scale (captures through 216-atom periodic / 192-atom molecular; molecular ethanol 90→47 ms trace, ~1.9×) and reaches the 512-atom cell in eager — only the 512 trace still exceeds single-card DRAM. The fp32 radial-backward fix adds ~4%. Full writeup in docs/uma-s-1p2-validation.md.

Testing

• Host tests/test_symmetry.py (no card). Card random-weight suite green; test_forces.py now also asserts the device radial finish vs the fp64 oracle. Gated tests/test_realweight_uma_s_1p2.py: 3 passed with a generated golden; auto-skips otherwise.

Notes

• ttnn 0.68.0. uma-s-1.2 via --checkpoint / UMA(checkpoint=...); bundle_cache auto-build stays uma-s-1. tt_atom/quaternion.py vendored verbatim from fairchem (MIT). UMA weights gated, not included.

🤖 Generated with Claude Code

[janweinreich]

Rebased onto current master (conflict-free — it merges cleanly; PR now shows mergeable) and re-validated end to end: full 757-system A/B vs fairchem gives no genuine force errors (741/757 force PCC > 0.999) and 691/757 pass the strict bar, above the 672 the pre-perf-pass backbone reached; up to 5.5× a 16-thread CPU at 216 atoms (master's kernels also restore the trace path at scale). Details in docs/uma-s-1p2-validation.md.

The one real force regression master's bf16 radial-MLP backward introduced (el_Sn_cmp, 230 meV/Å) is fixed here in fp32 (verified vs an fp64 mirror.radial_mlp oracle → PCC 1.0000).

I also ran a max-effort self-review and hardened the new charge-balancing across the batched paths it hadn't reached: MultiCard now builds the per-system mean operator (was a hang on uma-s-1.2), the traced batch uses the true per-system natoms, charged batches of unequal atom counts are rejected with a clear error rather than silently mis-targeting, and the vendored Wigner-D asset ships in the wheel (pyproject package-data, so non-editable installs don't FileNotFoundError).

Still a draft — happy to split the radial-backward fix into its own PR if you'd prefer it separate from the uma-s-1.2 feature.