Conversation
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
3 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.
Executive summary
This PR correlates to #261
This report combines two independent GPU evaluations of the reweighted Triton
ball-query implementation of smooth lDDT:
pr_reweighting.md;pr_reweighting_h800.md.Both studies use the same pinned set of 51 full protein–RNA structures, the same
coordinate perturbations, four reservoir sizes, three reservoir seeds, and the
same dense smooth-lDDT definition. Together they evaluate numerical fidelity,
training and inference runtime, GPU memory, large-structure behavior, and the
trade-off controlled by the ball-query reservoir size
K.The central numerical finding is reproduced across both machines. At the
recommended default
K=512, type-aware reweighting improves gradient agreementwith dense smooth lDDT:
0.26738 → 0.25021;0.95665 → 0.96995.At
K=256andK=1024, reweighting also improves both gradient metrics. AtK=2048, both reweighting factors equal one, so weighted and unweightedreductions become equivalent apart from floating-point reduction order.
The two machines expose complementary deployment properties:
20,016-atom
8h2hcase. Every sparse configuration completes.8h2h, butits peak incremental training allocation reaches 24.45 GiB, compared with
0.60 GiB for reweighted
K=512on the same case set.H800 dense training path is approximately 60% faster than the RTX 4090.
K; dense memory growsquadratically with atom count.
A critical methodological distinction must be retained when interpreting the
runtime and memory comparisons:
matched unweighted reduction running on the same current compacted
ball-query kernel. This isolates the incremental cost of reweighting.
778f30aawith theexact historical ball-query implementation at
074d75c2. This measuresthe complete historical-to-current change, including kernel and compaction
changes as well as reweighting.
Consequently, the RTX study is the authoritative measurement of the isolated
reweighting overhead, while the H800 study is the authoritative end-to-end
comparison with the exact original implementation requested for regression
validation.
1. Background and motivation
The default dense smooth-lDDT loss constructs pairwise atom-distance data with
quadratic
O(A²)scaling forAresolved atoms. This gives an exact reference,but its memory cost becomes prohibitive for large all-atom structures.
The Triton ball-query backend instead scans candidates and retains at most
Kin-radius neighbors for each center. Its dominant saved neighborhood storage is
therefore
O(AK)rather thanO(A²).The two center types use different physical cutoffs:
A 30 Å sphere has eight times the volume of a 15 Å sphere under a uniform-density
approximation. Using one fixed reservoir size therefore retains a smaller
fraction of an RNA-centered neighborhood than of a protein-centered
neighborhood. A simple unweighted truncated reduction can consequently
underrepresent nucleotide-centered rows.
The reweighted implementation corrects this imbalance for truncated rows while
preserving the normalization of the smooth-lDDT reduction. It applies the same
weight to each row's score sum and pair-count denominator.
For reservoir size
K, the lower-bounded row weights are:The tested schedule is:
Rows whose complete neighborhood fits in the reservoir retain weight one. The
additional weighted reduction uses
O(A)metadata and does not add anotheratom-by-
Ktensor.2. Implementations under test
2.1 Dense reference
The dense reference is the default
smooth_lddt_loss. It computes the completepairwise neighborhood and serves as the numerical target for scalar loss and
predicted-coordinate gradients.
2.2 Current reweighted ball query
The current implementation is commit:
It includes type-aware row reweighting and the current compacted ball-query
kernel/reduction implementation.
2.3 Matched current unweighted reduction
The RTX 4090 study disables row reweighting while retaining the current query,
compaction, and reduction machinery. This is the correct controlled comparison
for answering:
2.4 Exact historical ball query
The H800 study loads the exact historical implementation from commit:
It is executed from an isolated source snapshot in a fresh process, rather than
through the compatibility function retained in current code. This is the
correct comparison for answering:
Because this historical comparison also includes intervening kernel and
compaction changes, it cannot isolate the cost of reweighting by itself.
3. Benchmark dataset
Both studies use the same pinned 51-structure protein–RNA validation set.
7vtn,7r9f,7v93,8h2h,7ozs,7r7cThe dataset setup verifies:
The six large cases provide the strongest scaling test. The largest,
8h2h, contains 20,016 resolved atoms and is the only case on which the RTX4090 dense reference runs out of memory.
4. Shared benchmark protocol
For accuracy, each sparse scalar loss and predicted-coordinate gradient is
compared with dense smooth lDDT computed on the same structure and perturbed
coordinates.
The reported metrics are:
5. Hardware and software environments
H800 isolation controls
The H800 host had two GPUs. One was occupied throughout the campaign and was
excluded. All measurements used only:
The campaign enforced:
No overlap was detected. The selected GPU returned to 4 MiB used and 0%
utilization after the campaign.
6. Validation coverage
RTX 4090
The four K runs produced 1,428 result rows. All 1,224 sparse rows completed.
Dense completed 50 structures and OOMed only on
8h2h.H800
The four K runs also produced 1,428 result rows:
The H800 therefore adds a dense reference for
8h2hwhile reproducing the50-case numerical comparison used by the RTX report.
7. Cross-hardware numerical reproducibility
The following table uses the 50-structure common subset because
8h2hlacks adense reference on the RTX 4090. The values from both runs agree to the shown
precision.
Interpretation
K=256, reweighting improves scalar loss and both gradient metrics.K=512, reweighting slightly worsens mean scalar-loss error but improvesboth gradient metrics.
K=1024, scalar loss is nearly unchanged while both gradient metricsimprove.
K=2048, both weights equal one; differences reduce to floating-pointorder.
For training, gradient agreement is the more directly relevant signal. The
scalar loss and its gradient measure different properties, so the small scalar
regressions at
K=512and1024should not be hidden or conflated with theclear gradient improvement.
H800 all-51 accuracy
With the H800 dense reference for
8h2h, the complete 51-case means are:The complete-set result supports the same conclusion as the common subset.
8. Runtime on RTX 4090
The RTX runtime table is the controlled current-kernel comparison and therefore
the primary estimate of isolated reweighting overhead.
On RTX 4090:
reduction;
This small training difference is consistent with the extra
O(A)type-grouprow reduction.
9. Runtime on H800
The H800 table compares the exact historical implementation with the current
reweighted implementation over all 51 cases.
Relative to the exact historical commit, the current reweighted implementation
is:
These percentages include all historical-to-current implementation changes.
They are not estimates of reweighting-only overhead. For that question, use the
RTX matched-current comparison above.
10. Cross-hardware runtime comparison
10.1 Dense
On the RTX-common 50 structures, H800 dense training averages 6.35–6.57 ms,
compared with 16.18–16.54 ms on RTX 4090. This is a reduction of approximately
60%.
Dense forward averages 3.72–3.76 ms on H800 versus 8.35–8.43 ms on RTX 4090,
roughly 55% lower.
10.2 Current reweighted sparse backend
Across all 51 H800 cases versus the corresponding RTX aggregate:
Sparse forward timing is more launch- and shape-sensitive than dense timing.
The H800 advantage becomes clearer as K increases and each launch performs more
work. At low K, fixed launch overhead and aggregation across many small cases
can dominate.
10.3 Scope warning
The RTX and H800 runtime rows are not perfectly symmetric experimental objects:
8h2h, unless stated otherwise;Numerical comparisons use the common subset where needed. Runtime tables retain
each machine's complete successful coverage and explicitly state that scope.
11. RTX 4090 memory
The matched current unweighted and reweighted reductions have identical peak
memory in all 612 paired RTX comparisons. This establishes that reweighting
does not add another atom-by-
Kallocation.Sparse maxima include
8h2h; dense maxima do not, because RTX dense OOMs onthat case.
12. H800 memory
The H800 exact-historical comparison shows the effect of the complete current
implementation, including compaction changes.
Relative to the exact historical implementation, the current implementation
uses approximately:
This does not conflict with the RTX finding that matched current weighted and
unweighted reductions use identical memory. The two statements answer
different questions:
13. Dense memory and the
8h2hcaseThe H800 completes dense smooth lDDT on all structures. Across all 51 cases,
dense incremental memory is:
The maxima come from
8h2h. On that structure alone:At the default
K=512, sparse training uses about 39.5× less incrementalmemory than dense on
8h2h, while completing substantially faster.The RTX 4090 OOM is therefore expected: the measured H800 incremental dense
training peak alone is approximately 23.9 GiB, before accounting for resident
inputs, runtime context, allocator fragmentation, and other process memory.
14. Large-structure behavior
14.1 RTX 4090 large cases
The original RTX table reports reweighted accuracy over the five large cases
with dense references and runtime/memory over all six sparse cases.
14.2 H800 large cases
The H800 table includes dense accuracy for all six cases.
The accuracy difference between these tables primarily reflects inclusion of
8h2hin the H800 dense-reference set. Memory maxima are hardware-independentallocator results and match across the current sparse implementation.
15. Scaling analysis
15.1 Sparse memory scales with K
Doubling K approximately doubles sparse memory:
The relationship is not perfectly proportional because resident inputs and
O(A)metadata do not scale with K, but the dominant reservoir allocation does.15.2 Dense memory scales with atom-pair count
Dense memory is independent of K and strongly dependent on structure size. Its
maximum rises from 12.84 GiB on the RTX-executable subset to 24.45 GiB when the
20,016-atom case is included.
15.3 Accuracy improves with K
Increasing K reduces truncation:
K=1024 and 0.024 at K=2048 for the reweighted common subset;
K therefore exposes a direct memory-versus-fidelity control.
16. Recommended default and operating points
K=512: recommended default
K=512remains the recommended default when training memory and throughput areimportant:
over this demanding 51-case set;
allocation on H800.
K=1024: higher-fidelity middle ground
Choose
K=1024when roughly doubling K=512 reservoir memory is acceptable:K=2048: near-dense gradients
Choose
K=2048when near-dense gradients justify the memory cost:K=256: minimum-memory mode
K=256minimizes reservoir memory and benefits most strongly from reweighting,but gradient approximation remains less accurate than at larger K.
17. Conclusions
The combined RTX 4090 and H800 evidence supports the following conclusions.
4090 accuracy table to displayed precision on the common 50 structures.
cosine improve at K=256, 512, and 1024.
improves at K=256 but is slightly worse at K=512 and K=1024.
RTX comparison, training overhead is approximately 3% and forward overhead
is inconsistent/negligible.
and unweighted peaks are identical in all 612 RTX pairs.
H800 exact-commit comparison shows about 9% lower training memory and 25%
lower forward memory.
completes on both GPUs. Dense fails on the largest case on RTX 4090 and
consumes 24.45 GiB incremental training memory on H800.
low overhead, and a large memory reduction while preserving configurability
for higher-fidelity K values.
Overall, the reweighted ball-query smooth-lDDT backend provides a strong
replacement for dense smooth lDDT in memory-constrained all-atom training. It
preserves controllable fidelity, improves the bias introduced by a shared
truncated reservoir, and remains practical from RTX-class accelerators through
large-memory data-center GPUs.
18. Validation and artifact inventory
RTX 4090 source report
H800 source report
H800 benchmark bundle
The bundle contains:
results/: raw JSON, CSV, and streaming JSONL rows;analysis.json: validated aggregates;logs/: environment, campaign, per-block, and GPU-monitor logs;dataset/protein_rna_51_cases_full/: immutable uncropped benchmark cases;historical-074d75c2/: exact historical source snapshot;benchmark_one_backend.py: isolated-process benchmark harness;run_campaign.sh: sequential single-GPU campaign runner;analyze_results.py: row validation and aggregation.The OpenFold3 source checkout remained clean after testing.