Inference determinism: atom aggregation, feature seeding, and chunk-size reuse - #320
Open
heathcliff233 wants to merge 4 commits into
Open
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.
This was referenced Jul 18, 2026
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
First PR in a planned breakdown of #318. This lands only the inference-determinism and chunk-reuse foundation.
Why current inference can disagree across identical runs
Two independent sources of nondeterminism show up even with a fixed sample seed and production math (TF32 on):
seedduring featurization, so the same query can produce differentref_pos(and related features) depending on prior process state.scatter_add_, whose CUDA atomic reductions are not order-stable. Small differences are amplified through diffusion and can change the final pose.These are independent: seeding features alone does not make model outputs repeatable under scatter aggregation, and segmented aggregation alone does not make ligand features repeatable under polluted ambient RNG.
What this PR changes
seed(save/restore Python, NumPy, and CPU Torch RNG) so conformers and other feature RNG no longer depend on ambient process statesegment_reducefor atom→token aggregation in inference / no-grad, keepingscatter_add_for training/autogradChunkSizeTunerresults per argument shape and max chunk size so alternating targets reuse prior probes instead of retuning every revisit