Repository navigation
Keep ETKDG retry batches full and bound retries per molecule - #353
Conversation
The ETKDG scheduler allowed at most one attempt in flight per conformer still needed. Once only a few molecules remain unfinished - typically ones that keep failing embedding - batches shrink to a handful of attempts, a single worker runs them while the others wait, and the remaining retry budget is spent serially on a mostly idle GPU. The default retry limit was also 10x the atom count of the largest molecule in the input, applied to every molecule. A molecule's budget depended on what else was in the call, and small molecules that cannot be embedded got a budget sized for the largest one. - When the regular dispatch pass leaves a batch short, fill the remaining slots round-robin with extra attempts for unfinished molecules that still have budget. These attempts count against the per-molecule limit and get stable attempt IDs; waiting on in-flight results (NVIDIA-BioNeMo#299) is unchanged. Surplus successes are discarded at writeback, and completed counts are clamped to the requested number of conformers. - Compute the default limit per molecule as 10x its own atom count, matching RDKit. An explicit maxIterations still applies to all molecules.
5b3274a to
5d80c26
Compare
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
For problematic datasets - e.g. platinum, which has 2 mols that RDKIt and nvmolkit fail on 100% of the time, with default max iterations, this cuts down on runtime by 50% |
Speculative top-up attempts can yield more successes than a molecule needs, and writeback kept whichever arrived first, making retained conformers timing-dependent across workers. Tag conformers with their attempt ID and keep the N lowest per molecule on both the host and device paths, matching the set the in-flight-capped scheduler produced before.
evasnow1992
left a comment
evasnow1992
left a comment
There was a problem hiding this comment.
Thank you for addressing my comment!
|
Holding off until i can double-check the new greptile comment |
pruneDeviceConformers now walks each molecule's conformers in conformer-index order (attempt-ID rank for ETKDG output) and numbers survivors in that order, so greedy RMS pruning no longer depends on which worker collected which attempt. finalizeOnTarget ranks attempts with an O(n) per-molecule bucketing instead of a global sort.
Dropping surplus conformers split the per-collector bulk copy into one copy call per surviving run. When anything is dropped, gather positions with a single kernel instead.
evasnow1992
left a comment
evasnow1992
left a comment
There was a problem hiding this comment.
Updated changes also look good to me
Stream-ordered allocations on another GPU are not covered by cudaDeviceEnablePeerAccess, so the gather kernel must not read them directly. Copy each such buffer to the target GPU with one peer copy and gather from there.
| stagedPositions[c] = AsyncDeviceVector<double>(collector.positions.size(), targetStream.stream()); | ||
| copyDeviceToDeviceAsync(stagedPositions[c].data(), | ||
| collector.positions.data(), | ||
| collector.positions.size() * sizeof(double), | ||
| collector.gpuId, | ||
| collector.stream, | ||
| targetGpu, | ||
| targetStream.stream()); |
There was a problem hiding this comment.
When multi-GPU device output drops surplus successful attempts, this copies every remote collector’s entire positions buffer onto the target GPU, including coordinates that will be discarded. Those copies coexist with the retained result. On a large run, finalization can run out of target-GPU memory even when the final result would fit.
There was a problem hiding this comment.
Really just making up stuff now
There was a problem hiding this comment.
The concern is specifically about peak target-GPU memory, not the final result size. In the filtered path, result.positions is allocated first, then stagedPositions[c] allocates a full copy of every remote collector’s positions buffer on the target GPU (lines 167–177), even though the gather later keeps only a subset.
For example, if remote collectors hold 100 attempts and only 10 are retained, the target needs space for the retained result plus all 100 staged attempts. The retained result fitting does not imply those staging allocations fit. This is a potential allocation failure for retry-heavy, multi-GPU inputs—not a claim that every run will fail.
A bounded approach would stage only the source spans needed by kept entries, or gather in chunks.
2 participants
The ETKDG scheduler allowed at most one attempt in flight per conformer still needed. Once only a few molecules remain unfinished - typically ones that keep failing embedding - batches shrink to a handful of attempts, a single worker runs them while the others wait, and the remaining retry budget is spent serially on a mostly idle GPU.
The default retry limit was also 10x the atom count of the largest molecule in the input, applied to every molecule. A molecule's budget depended on what else was in the call, and small molecules that cannot be embedded got a budget sized for the largest one.