Build template pair features from coordinates online - #322
Open
heathcliff233 wants to merge 10 commits into
Open
Build template pair features from coordinates online#322heathcliff233 wants to merge 10 commits into
heathcliff233 wants to merge 10 commits into
Conversation
Gather token features through precomputed atom-to-token indices in the decoder, and use segmented reductions for packed inference aggregation to avoid CUDA atomic scatter nondeterminism. Keep the scatter path under autograd. Tests: pytest openfold3/tests/test_atomize_utils.py -q Tests: pytest openfold3/tests/test_sequence_local_atom_attention.py -q
Keep separate chunk-size entries keyed by argument structure and max_chunk_size so alternating inference targets can reuse prior OOM probes instead of re-tuning on every revisit. Tests: pytest openfold3/tests/utils/test_utils.py -k chunk_size_tuner -q
Scope on-the-fly feature RNG (e.g. ligand conformers) to each sample's seed with save/restore so featurization no longer depends on ambient Python/NumPy/Torch RNG state. Tests: pytest openfold3/tests/core/data/framework/single_datasets/test_inference_seeding.py -q
Factorial check of feature seeding and segmented atom aggregation under TF32 production math, with ambient RNG pollution between feature and forward repeats.
Hand ownership of the inference trunk pair into rollout so confidence can drop it after cloning, removing ~1U of co-resident pair tensor from the confidence peak.
Bound shared PairBlock working sets after the chunk-size tuner via OPENFOLD3_TRI_ATTN_CHUNK_CAP, OPENFOLD3_TRIMUL_CHUNK_CAP, and OPENFOLD3_TRANSITION_CHUNK_CAP. Unset keeps upstream defaults. Add a stage profiler to attribute peaks under these caps.
Row-chunk inference diffusion pair+relpos embedding to avoid the full N² concat transient, and process templates one-at-a-time on GPU so the four-wide stack activation is never materialized.
Avoid resident N² template distogram/unit-vector features in inference by projecting from compact O(N) coordinates with a length-generic Triton path.
Enable coordinate template features in training with a length-generic Triton backward, bf16 activations with fp32 accumulate, and keep the inference in-place path.
Describe use_coordinate_pair_features for inference/training and the Triton OPENFOLD3_FUSED_TEMPLATE_COORD path.
Author
|
Latest commits enable the usage of the fused template coordinate kernel in training. The training kernel disabled in-place operations, so it roughly reduces the peak here by 1U. |
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Next PR in the breakdown of #318, stacked on #321. This slice implements the online template coordinate path.
After #321 chunk caps and lifecycle work, a ~1.2U template distogram (+ unit vectors) still sat in
batchfor the whole forward. This PR replaces those quadratic template pair inputs with compact O(N) coordinates and builds/projects the pair features online during the template embedder, so the N² tensors never land in HBM.Until #321 merges, the GitHub diff against
mainalso includes that PR (and #320). The commit unique to this branch is:feat: build template pair features from coordinates onlineMeasured peak (S=1, cuEq, offload off, caps=128)
Model forward wall is essentially unchanged. The ~1.3U drop is the resident distogram/unit-vector payload leaving
batch.Microbench host payload for one template falls from O(N^2) to O(N) versus precomputed distogram+UV; embed transient stays ~0.6U either way. Triton compile is length-generic (
do_not_specializeon N and strides): one filesystem compile serves N=16…96 in tests.What this PR changes
template.use_coordinate_pair_features(inference only; default distogram bins 3.25/50.75/39)template_pseudo_beta_coords/template_frame_atom_coordsinstead of N² distogram / unit-vectorOPENFOLD3_FUSED_TEMPLATE_COORD, default on)_forward_streaming; coordinate path swaps only the per-template embedder and skips template offload (coords are O(N))scripts/dev/bench_fused_template_coordinate_embed.pyandscripts/dev/check_template_coordinate_parity.pyOnline local-frame math (normalize/cross) is algebraically the same as the Rigid path but not bitwise identical (~1e-6 UV, amplifies through the trunk). E2E structure coords stay within ~2e-4 abs on ubiquitin / 7cnx under deterministic CUDA.
Out of scope
Test plan
pytest openfold3/tests/test_fused_template_coordinate_embed.py— pair-embedder parity, Triton vs reference, compile reuse across lengthsTemplateEmbedderAllAtommulti-template coords vs precomputedscripts/dev/check_template_coordinate_parity.py --query 7cnx/--query ubiquitinscripts/dev/bench_fused_template_coordinate_embed.py --n 384,590,1264