Background
The MCMM sampler (src/mcmm.py, src/concerted_rotation.py) already batches its two
GPU-bound stages across all N walkers: MMFF minimisation runs in one
nvmolkit.mmffOptimization.MMFFOptimizeMoleculesConfs call and MACE scoring goes through the
batched _mace_batch_energies path. Coordinates, RMSD, inertia, and basin memory are already
torch tensors. The ParallelMCMMDriver / ReplicaExchangeMCMMDriver orchestration is in place.
The one stage still running CPU-sequential per walker is the DBT concerted-rotation move
proposal — specifically its ring-closure solve (scipy.optimize.least_squares, one walker at a
time) and the finite-difference Wu-Deem |det J| estimate. Now that MMFF is batched, this serial
proposal is the dominant per-step cost at large walker counts. This is a quality-of-life /
throughput improvement: it does not change what the sampler finds, only how fast it runs.
Proposed approach
Move the move-proposal stage onto the GPU as a batched torch batch_propose_fn, drop-in for the
interface ParallelMCMMDriver already expects
(batch_propose_fn(coords_list) -> [(new_coords, energy, det_j, success)]). No changes to the
driver, replica exchange, or basin memory.
- Batched dihedral rotation.
apply_dihedral_changes is pure rigid rotation — vectorise over
walkers as [N, W, 3] with [N] batched Rodrigues matrices (the _kabsch_rmsd_pairwise batched
pattern is the template). No solver.
- Batched ring closure. Replace the per-walker
least_squares with either a batched
Gauss-Newton / Levenberg-Marquardt (all N walkers' 6-residual/3-unknown systems solved in
lockstep via torch.linalg.solve, warm-started from the previous step) or the analytical
Coutsias/Plucker closure (closed-form, trivially batched).
- Batched Wu-Deem |det J| via
torch.autograd.functional.jacobian / torch.func.vmap on the
batched closure, instead of finite-difference re-solves.
- Uniform control flow. Per-walker branching (move type, accept/reject, convergence) handled by
compute-all-then-mask, keeping the whole MC step on-device (no per-step host<->device transfers).
Key questions / unknowns
- Profile first: confirm the CPU proposal is now the per-step ceiling vs the already-batched MMFF,
and quantify the achievable speedup as a function of N.
- Does the batched closure reproduce the scipy path's closed geometries and |det J| to
tolerance? Detailed balance depends on it — needs a parity regression test against the current
proposer.
- Numerical precision: closure Jacobians likely need float64 (as the existing Kabsch/inertia code
already uses).
- Convergence robustness of batched Gauss-Newton near the closure-manifold boundary (where the
current code already returns success=False).
Relationship to other work
Other notes
Phased path: (1) profile to confirm the ceiling; (2) batched dihedral rotation; (3) batched
Gauss-Newton closure as a drop-in batch_propose_fn with scipy parity; (4) autograd |det J|;
(5) swap in the analytical #24 closure for the final speedup + shared code. Warrants its own branch.
Background
The MCMM sampler (
src/mcmm.py,src/concerted_rotation.py) already batches its twoGPU-bound stages across all N walkers: MMFF minimisation runs in one
nvmolkit.mmffOptimization.MMFFOptimizeMoleculesConfscall and MACE scoring goes through thebatched
_mace_batch_energiespath. Coordinates, RMSD, inertia, and basin memory are alreadytorch tensors. The
ParallelMCMMDriver/ReplicaExchangeMCMMDriverorchestration is in place.The one stage still running CPU-sequential per walker is the DBT concerted-rotation move
proposal — specifically its ring-closure solve (
scipy.optimize.least_squares, one walker at atime) and the finite-difference Wu-Deem |det J| estimate. Now that MMFF is batched, this serial
proposal is the dominant per-step cost at large walker counts. This is a quality-of-life /
throughput improvement: it does not change what the sampler finds, only how fast it runs.
Proposed approach
Move the move-proposal stage onto the GPU as a batched torch
batch_propose_fn, drop-in for theinterface
ParallelMCMMDriveralready expects(
batch_propose_fn(coords_list) -> [(new_coords, energy, det_j, success)]). No changes to thedriver, replica exchange, or basin memory.
apply_dihedral_changesis pure rigid rotation — vectorise overwalkers as
[N, W, 3]with[N]batched Rodrigues matrices (the_kabsch_rmsd_pairwisebatchedpattern is the template). No solver.
least_squareswith either a batchedGauss-Newton / Levenberg-Marquardt (all N walkers' 6-residual/3-unknown systems solved in
lockstep via
torch.linalg.solve, warm-started from the previous step) or the analyticalCoutsias/Plucker closure (closed-form, trivially batched).
torch.autograd.functional.jacobian/torch.func.vmapon thebatched closure, instead of finite-difference re-solves.
compute-all-then-mask, keeping the whole MC step on-device (no per-step host<->device transfers).
Key questions / unknowns
and quantify the achievable speedup as a function of N.
tolerance? Detailed balance depends on it — needs a parity regression test against the current
proposer.
already uses).
current code already returns
success=False).Relationship to other work
src/ring_geometry.py's NeRF +closure is pure-numpy and torch-portable by design; porting
place_atom/close_ringto batchedtorch yields both the GPU move proposer here and the analytical-closure speedup for Plucker Coordinate Embedding #24
seeding. One implementation, two payoffs.
concerted_rotation.pyalready flags the analyticalclosure as the deferred "Option A".
get_mol_PE_mcmm(the issue Benchmark non-ETKDG samplers against exhaustive ETKDG on cyclic peptides #10 sampler benchmark) and the coverage harness.Other notes
Phased path: (1) profile to confirm the ceiling; (2) batched dihedral rotation; (3) batched
Gauss-Newton closure as a drop-in
batch_propose_fnwith scipy parity; (4) autograd |det J|;(5) swap in the analytical #24 closure for the final speedup + shared code. Warrants its own branch.