port-bio-model-to-tenstorrent

name

port-bio-model-to-tenstorrent

description

Playbook for implementing/porting a new biomolecular model (protein folding, structure prediction, protein language model, diffusion structure head, MSA encoder, etc.) onto Tenstorrent hardware inside the tt-bio codebase. Use when the task is "port model X to ttnn/Tenstorrent", "add a new model to tt-bio", "implement <bio model> on TT", or extending an existing port. Encodes the non-negotiable requirements (inference-only, accuracy, performance, --fast, unification, no redundancy, vendored pip deps) and the hard-won methodology (component-by-component parity porting, on-device residency, warm profiling, bit-exact bucketing, sample-dim batching with OOM chunking, bf16 load). Distilled from the ESMC-300M/6B + ESMFold2 and Boltz-2 ports (the 58-commit ESMFold2 history).

Porting a bio model to Tenstorrent (tt-bio)

Canonical playbook for adding a model to tt-bio (package tt_bio/, CLI tt-bio predict ... --model <name>). Worked examples in-tree: Boltz-2 (tt_bio/boltz2.py) and ESMFold2 (tt_bio/esmfold2.py + esmfold2_runtime.py, vendored reference under tt_bio/_vendor/). Make a new model look exactly like these — symmetric, unified, no parallel "model-X folder with its own everything".

Non-negotiable requirements (the bar every model must clear)

1. Inference-only. Delete all training code, losses, optimizers, schedulers. No PyTorch Lightning. Filter checkpoint hparams to the inference signature.
2. Pure pip deps — no git clone at runtime. Vendor the small host-side reference into tt_bio/_vendor/<name>/ and depend on stock PyPI wheels (see Vendoring below). No sibling clones, no sys.path shims, no ESM_PATH.
3. Accuracy must not regress. Per-module PCC > 0.98 vs a random-weight reference; end-to-end Cα-RMSD vs ground truth (sub-Å is the target — ESMFold2 hit 0.40–0.43 Å on Trp-cage). ttnn is not bit-deterministic (~0.05 Å run-to-run) — compare distributions, and see the RNG-confound note below.
4. Performance must not regress — both load time and predict time.
5. Simplicity & elegance. Reuse shared Module/TorchWrapper primitives; never duplicate attention/triangle/linear code.
6. Unification, zero redundancy. Same I/O, scheduler/worker, and CLI as the existing models. After it works, delete the parallel/unused paths.
7. --fast mode (block-fp8). Route through set_fast_mode/_dtype() so --fast runs bfloat8_b where safe; keep fp32/bf16 where it hurts accuracy.
8. Complete functionality. Every input size (to the ceiling, e.g. L=1024, no OOM), every argument, every checkpoint variant (e.g. a "fast" checkpoint), MSA path if the model has one.
9. Consistent terminal output in normal (Rich) and --debug --log modes — identical look to the other models.
10. Unified, non-redundant README — fold into the existing --model table / sections; don't append a parallel prose section.
11. Step by step, simplest case first, testing in between every step.

Process (the order that worked)

Phase 0 — Understand three codebases first: the target reference (forward, module tree, checkpoint layout, the real inference defaults — not a demo wrapper's), the tt-bio ttnn patterns (tenstorrent.py, boltz2.py, esmfold2*.py), and tt-metal examples/.

Phase 1 — Start with the SMALLEST sub-model/variant. ESMFold2 began with a from-scratch ttnn port of ESMC-300M (sequence-only LM) validated to emb/logits PCC ~0.999 before touching the 6B model or the folding head. Get a small thing bit-faithful, then scale.

Build the reference harness as "step 0" (before any ttnn code). Load only the needed reference module file(s) under stub/fake namespace packages (stub unused deps) to avoid importing the whole framework or mutating the shared env; pin the arch config from the actual checkpoint keys. Add any missing shared primitive to tenstorrent.py (RoPE was the one tt-bio lacked).

Phase 2 — Port component-by-component, each "with parity" (PCC > 0.98 vs the random-weight reference via load_state_dict(strict=False); idiom in tests/test_esmfold2.py + tests/esmfold2_reference.py). The order that worked for a folding model: reference harness → folding trunk (triangle-mult) → token diffusion transformer (DiT) → diffusion conditioning → atom transformer (SWA + 3D RoPE) → atom encoder/decoder + diffusion module → distogram + confidence heads → inputs embedder + relpos → MSA encoder (optional) → LM shim + diffusion sampler → end-to-end orchestration. Don't advance until the current module passes. Parallelize this across the machine's TT cards — one experiment per free device (see Hardware discipline) — so parity tests, seeds, and variants run at once.

Phase 3 — End-to-end on device. Real weights, entire model on device (host = glue), validate Cα-RMSD vs ground truth. Only "working" when accuracy is good on device.

Phase 4 — Robustness. All sizes (bucketed), all args/variants, no OOM at the ceiling (size-adaptive release of resident tensors).

Phase 5 — Optimize (load first, then predict — see Performance). Accuracy fixed throughout.

Phase 6 — Unify & document. Route through the shared CLI/worker, delete dead code, fold into the README.

Architecture / wiring (grounded in tt_bio)

• ttnn modules subclass Module or TorchWrapper in tenstorrent.py. Use the shared helpers — never reimplement: torch_to_tt, _lin, _split_heads, _merge_heads, WeightScope, CORE_GRID_MAIN, and the shared TriangleMultiplication / TriangleAttention / AttentionPairBias blocks.
• Swap only the computationally heavy ops to ttnn; keep cheap host code as-is. Production runs the reference confidence head with only its O(L³) folding trunk swapped to ttnn — the full ttnn head was built, never paid off, and was deleted. Don't port what's cheap; it's simpler and the accuracy is free.
• Two integration styles: prefer vendored-model + built-in flag (like boltz2.py, builds ttnn at init). The external-model + post-hoc patch style (ESMFold2: esmfold2_runtime.py reassigns submodules to ttnn wrappers via a single declarative _SPEC table + one generic _Adapter) suits a vendored reference whose forward you reuse.
• CLI: add to the --model click.Choice in main.py (predict), reusing --fast/--debug/--log/--device_ids. Surface unsupported-input notes in yellow (warning, not error). Print the same "Done" summary line.
• Scheduler/worker: dispatch on cfg["model"] in _WorkerState.load/ predict_one (worker.py) — same scheduler, same multi-device fan-out, weights resident per device, MSA unified worker-side with a shared cache. No separate run loop.
• I/O unified: same input reading and structure/metrics writing; output contract (coords, plddt, ptm) matches so selection/writing is shared.

Vendoring the reference (how "no clones" is actually done)

• Copy the small host-side reference (featurization, MSA, mmCIF assembly, the model files) into tt_bio/_vendor/<name>/; rewrite absolute imports to the vendored namespace (esm. → tt_bio._vendor.esm.).
• Depend on the exact stock wheel the model targets (e.g. transformers==4.57.6), declared in pyproject.toml. Declare every import as a real dep (huggingface_hub, safetensors, zstd, …) — missing ones fail only at install/runtime on a clean machine.
• Ship licenses: per-file provenance headers + the upstream LICENSE texts in the vendored dirs + a top-level NOTICE; via package-data/license-files.
• Packaging gotcha: make sure .gitignore doesn't swallow vendored dirs (a generic msa/ output-cache rule once hid a vendored msa/ package → "No module named …" at runtime) and that setuptools packages.find covers them.

Performance methodology (hard-won — measured numbers in parens)

• Profile WARM, never cold. ttnn compiles each program on first use (seconds, one-time, cached per process); a cold single-shot per-stage profile misattributes it. Warm up, then time steady state.
• Profile to find the REAL hotspot — it's counterintuitive. In ESMFold2 the O(L³) folding trunk is ~67% of fold time while the 6B-param LM is only ~2%. Optimize the dominant stage, not the biggest-looking module.
• Transfers dominate loops, not compute. The reference re-transfers the big pair tensor (L²·256) host↔device every iteration (host-side tile-layout conversion in from_torch). Two residency patterns fix this:
  • Resident recurrence loop: keep the pair state on device across iterations; make deterministic-inference-invariant sub-modules compute-once (~32% faster trunk). The reference loop won't speed up warm — host tilization isn't a cacheable kernel — which proves the transfer cost is real.
  • Resident sampler context: split conditioning into step-invariant (f(z_trunk, relpos), compute once + keep resident) vs per-step (cheap, t-dependent); per step transfer only the tiny noisy coords. Pattern: prepare()/step()/release_cache(). (Diffusion 30.6 s → 2.5 s, ~12×.)
• Load straight to bf16 (not fp32-then-convert): ~2.6× faster load, bit-identical at default. Gate fp32 only for quantization-sensitive paths.
• Keep weights resident for batch/folder prediction (load once, fold many).
• --fast = block-fp8 on the heavy accuracy-safe op. Default the trunk's triangle-mul to bfloat8_b (trunk −40%, total fold −31%, accuracy unchanged). But keep high matmul accumulation fidelity (fp32 accum / HiFi4) — LoFi accumulation measurably degraded RMSD (0.61 → 0.73 Å) for ~0.5 s; the win is fp8 weights, not lossy accumulation.
• Bucketing (cap kernel-compile count for variable lengths) — bit-exact ONLY where padding is attention-masked. Bucket the LM/attention length to 64 bit-exactly: pad ids, additive −inf mask, AND zero padded keys/values so exactness doesn't depend on bf16 exp(-inf) (verify PCC = 1.0, maxdiff = 0). Do NOT bucket a dimension that changes a matmul's contraction size — in bf16 that isn't exact across sizes and shifts RMSD (~0.065 Å); leave those at the smaller existing bucket (32). Don't over-bucket — bucketing every dim was tried and reverted.
• best-of-N / multi-sample: the device is underutilized at small L, so run the N samples as one batched B=N pass (≈ free: B=4 ≈ 1.20× B=1). BUT it replicates the pair state to [N,L,L,c] in both the sampler and the confidence head, so high N OOMs past short lengths. Chunk both stages over the sample axis at a calibrated B·L² budget (distinct per-chunk seeds + shrink-on-OOM net), and make the budget grid-aware (halve on small Wormhole grids vs Blackhole); expose an env override. B=1 path stays untouched.
• Measure every optimization; revert what doesn't pay. Real reverts: trimul matmul-tiling tuning, and complete-everything bucketing. Measured DEAD ENDS, do not retry: parallel from_torch (ttnn serializes the queue → 1.0×); and row-major upload + on-device tilize (net slower than host TILE_LAYOUT).
• Feasibility-microbench before committing, then re-verify accuracy.

Accuracy methodology

• Keep random-weight reference modules for regression (tests/<model>_*.py); per-module PCC > 0.98 is the gate.
• End-to-end Cα-RMSD via Kabsch vs examples/ground_truth_structures/.
• RNG confound — critical. The reference may draw from the global torch RNG for things the seed arg doesn't cover (e.g. initial pair state), so per-fold RMSD across code changes is confounded and unreliable for proving exactness. Use a bit-identical intermediate test (hidden-state PCC = 1.0 / maxdiff 0), not fold RMSD, to prove an optimization is lossless.
• Predicted-cif gotcha: the ESM-style .cif omits _atom_site.occupancy, so BioPython MMCIFParser KeyErrors — parse Cα by column (label_atom_id idx 2, Cartn_x/y/z idx 14/15/16) + Kabsch directly.
• best-of-N only helps when confidence is informative. On low-confidence / out-of-distribution targets pLDDT can be flat or anti-correlated with RMSD, so best-of-N can pick a worse structure. Validate the benefit on a confident (high-pLDDT) target, not a hard one.

Simplification (Phase 6 — be radical, but know where to stop)

After it works, delete aggressively (one ESMFold2 cleanup was −354 lines, same output): parallel host orchestrations nothing uses, standalone ttnn modules never wired in, leftover fields. Collapse near-identical wrappers into ONE generic adapter driven by a declarative spec table. Do NOT force-merge genuinely distinct framework pieces just to merge — name what you intentionally leave separate.

Hardware & process discipline

• Use ALL the free devices to parallelize bring-up. The machine usually has several TT cards (e.g. 4 on a Quiet Box — ls /dev/tenstorrent/). Bring-up is embarrassingly parallel: fan independent experiments across the cards instead of serializing on one. Launch each in its own process pinned to one card with TT_VISIBLE_DEVICES=<id> (run them in the background), so per-module parity tests, optimization variants, accuracy/seed sweeps, and size sweeps all run concurrently → ~N× faster iteration on an N-card box. One job per card. The same multi-device machinery serves production — --device_ids / --num_devices fans targets across cards via the shared worker.
• Pin the device (TT_VISIBLE_DEVICES/--device_ids); confirm a card is free before using it (a held card shows as a CHIP_IN_USE lock / hang — check tt-smi). Respect any user constraint to a specific device.
• NEVER SIGTERM a running device job mid-op — it can crash the driver. Wait for it (and for someone else's job to finish) rather than killing it.
• Account for the grid: Wormhole (8×8) has ~55% of Blackhole's L1, so memory budgets (chunk sizes, etc.) must be grid-aware (_IS_SMALL_GRID).

Definition of done

• Entire model on device; host is glue only.
• Per-module PCC > 0.98; end-to-end RMSD good on device (proven lossless via a bit-identical test, not just fold RMSD).
• All sizes (ceiling, no OOM), all args/variants/checkpoints, MSA if present.
• --fast (block-fp8) works; accuracy acceptable; high matmul accum kept.
• Load + predict time measured (warm) and not regressed.
• Routed through shared CLI (--model <name>) + scheduler/worker; multi- device works; I/O unified; output matches in normal and --debug --log.
• No training code, no Lightning, no runtime clone (reference vendored + licensed; deps in pyproject; vendored dirs not gitignored).
• Dead/redundant code pruned; README folded into existing sections.
• Tests under tests/; commit messages end with the project trailer.

Key files to study / mirror

• tt_bio/tenstorrent.py — shared ttnn primitives (Module, TorchWrapper, torch_to_tt, _lin, set_fast_mode, _dtype, get_device, WeightScope, triangle/attention blocks, bucketing + _IS_SMALL_GRID).
• tt_bio/boltz2.py — vendored-model + built-in-flag integration style.
• tt_bio/esmfold2.py — ttnn module wrappers for a folding model.
• tt_bio/esmfold2_runtime.py — external-model patch style (_SPEC, generic _Adapter, resident trunk loop, resident sampler context, best-of-N chunking).
• tt_bio/_vendor/ — vendored reference (host fold path + model files + licenses).
• tt_bio/main.py — predict CLI; tt_bio/worker.py — per-device dispatch.
• tests/test_esmfold2.py, tests/esmfold2_reference.py — the PCC test idiom.