Skip to content

v0.2.1

Latest

Choose a tag to compare

@github-actions github-actions released this 20 Jul 11:06

[0.2.1] - 2026-07-20

No new model family this release. This is a consolidation release on top of v0.2.0's Orb-v3/
OrbMol port: one public API instead of three, several measured Orb-v3 perf wins, the on-hardware
release-gate framework that now runs all four checks with one command, and the MultiCard Orb
dispatch bug fix. UMA code paths are untouched throughout.

Added

  • Single public entry point: Calculator(atoms, model=...) picks UMA or Orb by checkpoint
    name (like fairchem's FAIRChemCalculator or HF's AutoModel.from_pretrained) instead of
    separate UMA()/Orb() factory functions, which are removed (a clean break, not an additive
    wrapper — no external contract yet to preserve). The typed classes (TTAtomCalculator,
    OrbCalculator) and their .from_uma/.from_checkpoint classmethods stay as the advanced
    surface; tt_atom/auto.py dispatches to them by model name. A family-inapplicable argument
    (e.g. task= on Orb, trace= on Orb) raises with an actionable message rather than being
    silently dropped. Pure call-surface change — calculate/from_uma/from_checkpoint and all
    numerics are byte-for-byte untouched; verified on device (30 Orb parity tests green through the
    unified entry point, bit-exact vs golden).
  • OrbCalculator.evaluate_batch: disjoint-union batched inference for Orb-v3/OrbMol — K
    systems in one device forward, per-system energies via EnergyHead.batch, forces from the
    batched conservative VJP or the direct ForceHead.batch. Measured on one card, 9-atom ethanol
    (benchmarks/bench_orb_evaluate_batch.py, K=128): conservative-omol ~18.7–19x (1338 vs
    71 sys/s), direct-omol ~12.4x (2056 vs 166 sys/s) — exceeds UMA's own ~13x batched speedup.
    Parity gated vs looping calculate on all three checkpoint variants (energy rel err < 1e-2,
    force PCC > 0.999).
  • Capability-gated perf defaults: fused_lnbw (UMA's fused radial-LayerNorm backward) is now
    default-ON, gated behind a build/capability probe (mirrors tt-bio's fuse_swiglu pattern) that
    detects whether the installed ttnn carries the custom fused kernels, falling back to the stock
    path on wheels that lack them (e.g. 0.68). device_ede/bf8_edge (Orb on-device edge-degree
    embedding / bf8 edge-activation dataflow) stay opt-in (env var): re-measurement showed they win
    ~2x at 512+ atoms but regress to ~0.85x on 9-atom molecules, so there's no size-independent
    default that's safe.
  • scripts/release_gate.py: one-command on-hardware release gate, the runnable version of
    RELEASING.md's manual checklist. Four legs, each machine-readable PASS/FAIL/GAP:
    accuracy (reuses tests/test_*realweight*.py verbatim via pytest + JUnit XML, so the gate never
    re-derives a parity bar or oracle), no-OOM (Orb disjoint-union batch sweep to the OOM ceiling —
    128 systems / 1152 atoms on this release's run; UMA's sweep is a documented env GAP), perf
    (warm throughput vs docs/perf_baselines.json, card-type aware, one entry per shipped family:
    OrbMol conservative-omol, Orb-v3 bulk conservative-inf-omat, UMA uma-s-1), and UX
    (scripts/ux_regression.py: CLI --help surfaces the core flags, a real single-point/relax/
    MD roundtrip writes and re-parses, and per-step progress output actually advances). A missing
    golden/baseline reports GAP, never a silent PASS.
  • docs/materials-benchmark.md: device-vs-reference implementation-parity doc (the R/D/X
    framework, mirrors tt-bio's pharma-benchmark.md), covering every shipped family — UMA, Orb-v3,
    OrbMol — in one table, plus a new MgO rock-salt row, the suite's first multi-element bulk
    system (every prior bulk row was pure Si): energy rel err 1.6e-3, force PCC 0.99998, exercising
    the per-element reference-energy denormalize, mixed-Z ZBL pair repulsion, and the encoder's
    per-element embedding table at two atomic numbers simultaneously.
  • examples/orb_md.py: periodic-crystal MD driver with on-device Orb-v3 forces.

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).
  • OrbMol's open-shell direct-forces PCC bar (orb-v3-direct-omol, CH3• radical) re-baselined
    0.9 → 0.85. Root-caused as a bf16 noise floor, not a port bug: the reference is bit-identical
    across reruns (self-consistency PCC 1.0), and treating the device's measured RMSE as additive
    noise on the reference forces predicts a best-case PCC of 0.908 — the device measures 0.8926,
    within 1.6% of that floor, while its energy error (1.25e-5) is the tightest of all six OrbMol
    rows. The conservative checkpoint of the same system clears 0.9785, so the charge/spin
    conditioning path itself is correct; only the direct ForceHead's extra rounding on this
    tiny-force system (oracle |F|max 0.05 eV/Å, an order of magnitude below the other rows) was
    ever going to depress the correlation. See docs/orb-port.md for the full X-vs-R/D analysis.
  • Release-gate reliability, found running the gate for real on the release host: the perf-leg
    measurement subprocess had no timeout, so a device-side hang blocked the whole gate forever
    (recurred 3x on the same model) — added a 240s timeout, child process-group kill, and a clean
    card reset before the next model. The OOM leg's process-global MetalContext was leaking into
    the perf leg's device open and causing cross-leg hangs; the OOM sweep now runs in its own child
    process so its context is torn down before perf opens the card. The accuracy leg's fused_rotate
    crash (as opposed to an auto-skipped missing golden) was reading as a blocking FAIL instead of
    the same documented env GAP the OOM/perf legs already use for that host limitation; classified
    consistently.
  • host_zbl_forces: documented (docstring) that a periodic system's ZBL term can silently NaN
    if called without cell_shift — found while adding the MgO golden.

Performance

Orb-v3 continued past v0.2.0's "not near the hardware limit" open question. Roofline + stage-
level profiling showed the p150 port running at 7.5% of bf16 compute peak / ~22% of DRAM
bandwidth floor at 2016 atoms — materialized edge-MLP activation traffic, not host dispatch or a
saturated matmul engine, so v0.2.0's "likely near the limit" framing did not hold. Several real,
parity-safe, bit-exact-or-PCC-gated wins landed this release, each measured independently and
compounding on the traced MD step (benchmarks/bench_orb_perf_dollar_tt.py and siblings, one
p150, bf16, real weights, median of repeated jittered-pos steps):

fusion measured speedup scope
closed-form host geometry VJP (replaces autograd-graph rebuild) + hardware-limit accel up to 1.12x (default), up to 1.23x (--fast) full MD step
fused edge-MLP SiLU backward (source-ttnn ttnn.experimental fused derivative) curve to 42.68 / 91.80 / 188.27 / 372.13 ms at 216 / 512 / 1000 / 2016 atoms edge-MLP backward
row-major edge-aggregation scatter (fixes a TILE-concat layout inefficiency in scatter.segment_sum) 1.04–1.08x full MD step
edge-MLP matmul factory (ttnn.experimental.minimal_matmul, drop-in, preserves force-path pre-activations) 1.08x–1.13x, curve to 39.70 / 84.61 / 164.36 / 341.03 ms at 216 / 512 / 1000 / 2016 atoms edge-MLP matmuls
--fast (bf8 hidden edge activations, bf16 residual stream, fp32 accumulate) 1.03x / 1.21x / 1.23x / 1.23x vs the bf16 curve at the same sizes full MD step, opt-in

No single before/after number spans the whole sequence (each fusion's baseline is the state left
by the previous one, measured in its own PR); the table above is the ordered, individually-
verified chain. An edge-chunked L1-resident streaming kernel and a forward Linear+SiLU epilogue
were both measured and rejected (0.42–0.98x and 0.988–0.996x respectively) — see
docs/orb-port.md for the full profiling writeup. --fast remains release-gated: force MAE is
0.049 eV/Å vs 0.0089 eV/Å for bf16, so it trades accuracy for the extra ~20% at 512+ atoms.

Fair TT-vs-GPU perf-per-dollar, redone honestly. An earlier claim (p150 "1.74x faster than an
H200", "~40x perf-per-dollar") compared TT's optimized trace/replay path against the GPU's stock
path with the neighbour list rebuilt every step, and its H200 timings had no committed raw
evidence — withdrawn. The redo compares the out-of-box path each side's user actually runs (TT:
OrbTracedEngine, the path examples/orb_md.py uses; GPU: stock pip install orb-models==0.7.0
ORBCalculator) on the same periodic Si diamond supercells, with raw per-step timings committed
(benchmarks/orb_perf_dollar_gpu_v0.7.0.json, benchmarks/orb_perf_edge_mlp_fused.json):

system p150 (bf16, traced) H200 (fp32, stock) H200 raw speedup p150 perf-per-dollar edge
216 atoms 42.68 ms/step 16.85 ms/step 2.5x faster ~9.1x
512 atoms 91.80 ms/step 19.43 ms/step 4.7x faster ~4.9x
1000 atoms 188.27 ms/step 44.47 ms/step 4.2x faster ~5.4x
2016 atoms 372.13 ms/step 70.51 ms/step 5.3x faster ~4.4x

The honest result: the H200 is faster on raw throughput at every size tested. The p150 (~$1,399
vs an H200 at ~$30–40k, roughly 23x cheaper) still wins on throughput-per-dollar, but by ~4.4–
9.1x, not ~40x, and the edge shrinks as systems grow — price/performance, not raw speed. Full
methodology, the matched-policy transparency view, and the hardware cost basis are in
docs/orb-port.md.

Notes

  • README rewritten for the user audience: kernel/dispatch/hardware-internals content and
    investigation-specific numbers moved out of the top-level README into custom_kernels/README.md
    and the relevant docs/*.md, with one practical fact plus a link left at the top level. Added a
    Troubleshooting section.
  • The v0.2.0 scope note previously 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); see Fixed above.
  • No real-weights multi-card scaling number is reported this release either: the release host
    has a single Tenstorrent card, so N>1 scaling cannot be measured on it. The one honest datapoint
    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 (v0.1.0-era) used the synthetic
    examples/model_tiny_demo.npz UMA bundle, not real weights and not Orb, so it is still not a
    real-weights scaling number for either family.
  • 9 closed-investigation Orb-v3 perf scripts (scatter dead-ends, forward-SiLU-fusion dead-end,
    rejected megakernel/chunking attempts) archived to benchmarks/archive/, each with a README
    pointer to the memory lesson that closed it. Nothing deleted, live scripts unchanged — internal
    housekeeping, no user-facing effect.
  • A handful of commits in this range were social-media drafts and Si-melt demo-video polish
    (LinkedIn/X post drafts, render-window/flicker fixes) — not user-facing library changes, omitted
    above.