tt-atom
affiliated
Meta's UMA interatomic potential running on Tenstorrent Blackhole — energy, forces, and stress for molecules and periodic materials behind an ASE calculator. Its per-edge Wigner rotation runs as a custom tt-metal kernel for a highest-performance uma-s build.
Links
📦
Repo
📋 Changelog
# Changelog All notable changes to TT-Atom are recorded here. Versioning is [SemVer](https://semver.org); releases are cut only from a commit that has passed the on-hardware release gate — accuracy parity, no OOM across the supported size range, and no perf regression (see `RELEASING.md`). ## [Unreleased] ### Fixed - `tt_atom.batch.MultiCard` now builds the Orb-v3/OrbMol backbone when given Orb weights. The worker previously hardcoded the UMA path (`WeightBundle` + eSCN-MD `Backbone`), so pointing it at an Orb weights file built the wrong model silently. It now dispatches on the loaded bundle's `config` (the same UMA/Orb family split `tt_atom.auto` exposes by name) and runs the `Encoder`/`AttentionInteractionLayer`/`EnergyHead` forward for Orb. Verified bit-exact vs the single-card `OrbCalculator` on `orb-v3-conservative-inf-omat` (energy diff 0 eV on H2O / ethanol / benzene). `tests/test_multicard_orb.py` mirrors the UMA `test_multicard.py` sharded-vs- sequential parity shape (auto-skips below 2 cards). ### Notes - The v0.2.0 scope note below flagged Orb multi-card as "not independently re-run — same scheduler as UMA". That understated the gap: at v0.2.0 the worker was UMA-only, so Orb multi-card did not work at all (not merely unmeasured). Fixed here. - No real-weights multi-card *scaling* number is re-reported this pass: pc has a single Tenstorrent card, so N>1 scaling cannot be measured on it. The one honest datapoint measured here is the per-card baseline: Orb (`conservative-inf-omat`, real weights) on one card at ~128-atom Si supercells — 0.37 Medges/s. The earlier 2.95x@4cards figure (commit 43e981b) used the synthetic `examples/model_tiny_demo.npz` UMA bundle, not real weights and not Orb, so it is not a real-weights scaling number for either family. ## [0.2.0] - 2026-07-11 A second model family, additive to v0.1.0: **Orb-v3** (Orbital Materials) and **OrbMol**, its charge/spin-conditioned molecular variant. UMA/eSEN code paths are untouched (byte-identical to v0.1.0) — see full history and numbers in `docs/orb-port.md`. ### Added - **Orb-v3** (`orb-v3-conservative-inf-omat`, `orb-v3-direct-20-omat`): a non-equivariant, attention-MPNN backbone, ported bottom-up (encoder, 5-layer backbone, energy/force/stress heads, ZBL pair repulsion, periodic images, disjoint-union batching). None of UMA's four custom kernels transfer (Orb has no equivariant hidden representation) — this path runs on stock `ttnn` ops only, no source tt-metal build required for Orb-only use. - **OrbMol** (`orb-v3-conservative-omol`, `orb-v3-direct-omol`): the OMol25-trained, charge/spin- conditioned checkpoints. Reuses the Orb-v3 backbone unmodified plus a closed-form, node-only charge/spin embedding (zero learned matmuls) — no new forward/backward machinery. - `OrbTracedEngine` (`tt_atom/orb_trace.py`): trace-capture for the Orb-v3 forward(+analytic-VJP backward), refreshing only the two pos-dependent device inputs per MD/
Works on
blackhole