Reduce inference peak memory with chunking and fine-grained lifecycle control - #321
Open
heathcliff233 wants to merge 7 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.
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.
5 tasks
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
Second PR in the #318 breakdown, stacked on #320. This one cuts inference peak memory with opt-in chunk caps and shorter N² lifetimes — no fused pair kernels. Offload stays off in the numbers below.
Until #320 merges, the GitHub diff against
mainalso includes that PR. The three commits unique to this branch are:perf: release trunk pair after confidence cloneperf: add opt-in triangle and transition chunk capsperf: chunk diffusion conditioning and stream templatesWhy N² features and intermediates dominate
Define
1U = N² × c_z × 4— the storage of one fp32 trunk pair tensor. At long targets, almost every large allocation is pair-shaped: the workingz, transition / triangle scratch, template distogram and unit-vector features, and diffusion relpos concatenations. Singles, atom tensors, and MSA rows are small next to that (fractions of a U). So peak memory is mostly how many N² tensors coexist, and how large the pair-op scratch is while they do.This PR attacks that two ways: bound the pair-op working sets with opt-in chunk caps, and shorten lifetimes so fewer full N² tensors sit resident at once. Fused kernels that avoid materializing those intermediates are left for later.
Primary metric: peak CUDA alloc above live model params+buffers, reported in U.
Measured peak (homo_1200, S=1, cuEq, offload off)
Warm forward with caps+lifecycle is ~108s vs ~87s uncapped default; eager trimul chunking is the main wall cost.
Caps alone bound PairBlock / transition / OPM scratch. They do not touch diffusion's full
[N,N,267]relpos+concat (~10U), the 4-wide template embed (~9.5U), or input-embedder relpos staging (~6.6U). The drop from ~10U to 6.63U is the lifetime work on top of those caps:After that, the global peak is
input_embedderat 6.63U. Template distogram (~1.2U) stays inbatchfor the whole forward — no CPU park in this PR.What this PR changes
OPENFOLD3_TRI_ATTN_CHUNK_CAP,OPENFOLD3_TRIMUL_CHUNK_CAP,OPENFOLD3_TRANSITION_CHUNK_CAP(unset = upstream default)_inference_forwardwith that chunkrow_sliceonrelpos_complex)inplace_safeandn_templ > 1scripts/dev/profile_inference_stages.pyfor warm / overall / hooked stage peaksOut of scope
Fused SwiGLU / trimul / tri-attn / input-relpos, template CPU-park or O(N) coordinate templates, and dual
z_init+zresidency. Those are what the original tip used to reach ~4.2U; this PR stops at ~6.6U on homo_1200.Test plan
test_heads,test_diffusion_conditioning,test_template_module,test_utils(cap envs)--triangle-chunk-cap 128 --transition-chunk-cap 128; expect ~6.6U, input_embedder peak