Repository navigation
Conversation
Improve numerical stability of the density-pair gradient against changes in memory-dependent contraction sizes. For density pairs from test_jk_energy_per_atom_dm_pairs in gpu4pyscf/df/tests/test_df_tdrhf_grad.py, substituting 16.2 MB and 2.49 MB as get_avail_mem returns produces a maximum K gradient component discrepancy of 4e-12 Eh/bohr with the expanded metric. _jk_energies_by_dm_factors forms G = C M^-1 C.T and applies G X to assembled J vectors and K tensors. M is the contracted auxiliary Coulomb matrix; C expands its basis into the integral representation. Small eigenvalues of M produce large entries in G, so its output dot products can subtract large terms to obtain small coefficients. Compute t = A.T X, then A t, for each assembled array X, sharing A across J and K. Following gpu4pyscf.df.df._decompose_j2c, use A = C L^-T for Cholesky M = L L.T, or A = C U lambda^-1/2 with retained eigenvectors U and eigenvalues lambda. Preserve decomposition selection, fallback, and the eigenvalue cutoff. In this fixture, G entries reach about 20000, while A entries reach about 70. One coefficient near -0.0065 sums positive and negative contributions of about 50000 in G X. Along the factorized path, a contributing t = A.T X entry combines signed sums of about 200 to give an intermediate near -0.3; the final A t combines signed sums of about 20 to give the coefficient near -0.0065. The first product thus performs part of the cancellation before forming the output coefficients. The final product reuses these smaller t values across outputs. An error in one t value enters each output with the corresponding weight from A; when those contributions cancel in the gradient, their shared error cancels too. Dense G X computes a separate large-term sum for each output, with its own rounding residual. Both factorized products still round; the gradient comparison below measures their combined effect. For the same memory probes, the factorized application reduces the maximum K gradient component discrepancy to 2e-14 Eh/bohr. Multiplying the same stored G in extended precision also reduces the discrepancy to 1e-14, isolating rounding in the dense product. The improved consistency trades one dense metric multiplication for two factor multiplications. With n expanded and m contracted auxiliary functions, Cholesky application uses 2*m/n times the leading arithmetic of G X; eigenvalue truncation replaces m with the retained rank.
Memory-derived auxiliary sizes can exceed the evaluator's largest tile. Cap integral and contraction scratch at that tile size in RHF, UHF, and the TD summed-density and density-pair paths. The same cap also bounds derivative batching. Hermitian exchange scratch also reserves excess AO-sized capacity. Size it by the density-factor width and AO count in RHF, UHF, and summed-density TD, retaining both spin channels in UHF. Preserve the transpose workspace required by non-Hermitian inputs.
RHF, UHF and TD-RHF derivative drivers filled one auxiliary group per launch despite buffers sized for several groups, limiting grid parallelism. Pack consecutive groups within the existing auxiliary-function capacity and pass the group count through the native driver's existing second grid axis. This exposes more independent blocks per packed launch while preserving group boundaries and function ordering. To keep RHF Hessian packed work distributable across devices, target the auxiliary-function count per device when limiting tile size, with whole groups setting the minimum granularity. Assign packed tiles through the existing shared iterator.
Int3c2eOpt.int3c2e_evaluator evaluates three-center Coulomb integrals. Its aux_batch_size argument specifies a requested auxiliary-function capacity, but setting it produced one shell group per call even when several groups fit. This left buffer capacity and GPU parallelism unused. Use splits_by_blocksize to pack consecutive groups using that capacity. Apply the same rule to _int3c2e_ip1_evaluator, which evaluates first derivatives of the three-center integrals, and pass the packed group count to its native driver. Preserve auxiliary-function ordering; both evaluators return offsets describing their packed batches. Larger tiles expose an existing buffer-sizing bug in tdscf.ris.get_Tpq. It allocated the integral-output buffer using the smaller inner orbital-contraction block size, causing buffer assertions under limited memory. Fix this by allocating for the largest returned integral tile, then contracting it in smaller orbital blocks. Add test_get_Tpq_limited_memory in df/tests/test_df_tddft_ris.py. It compares J/K orbital transforms under normal and constrained memory, covering both host and GPU result storage, to guard the full-tile buffer allocation.
zc62
force-pushed
the
df-derivative-batching
branch
from
October 16, 2026 01:47
c18624f to
4364bea
Compare
Author
|
Force-pushed to satisfy linter. |
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
1 participant
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.
DF derivative drivers often evaluate one auxiliary shell group per launch, even when several groups fit in the allocated buffer. Pack consecutive groups into larger launches to improve GPU utilization. Each CUDA block continues to process a homogeneous angular-momentum/primitive group.
Changes
1e-12to1e-11at these test lines:gpu4pyscf/gpu4pyscf/df/tests/test_df_tdrhf_grad.py
Lines 309 to 313 in e1d6778
tdscf.ris.get_Tpq: the integral evaluator writes a complete auxiliary tile before orbital contractions split it into smaller blocks. The original code sized the integral-output buffer for the smaller contraction block, causing a buffer-size assertion when a returned tile exceeded that capacity. Allocate integral scratch for the largest returned tile while retaining smaller contraction blocks. Larger packed tiles expose this bug more readily; a new low-memory test covers host and GPU result storage.Performance
RTX 5070 Ti (16 GB, 70 SMs), Ryzen 7 9700X, 8 OpenMP/OpenBLAS threads; CUDA 13.0, CuPy 13.6.0 with the CuPy contraction backend. PBE0, def2-TZVP/def2-TZVP-JKFIT, level-3 grids, auxiliary-basis response enabled. Times cover complete gradient or Hessian
kernel()calls using identical saved SCF/TD states: one warm-up followed by three synchronized serial measurements. The TDDFT gradient uses the first singlet root.The two batching commits contribute separately:
For water10, restricted Hermitian exchange scratch falls from 2,451 MB to 255 MB; the corresponding unrestricted allocation falls to 511 MB.
AI Usage Disclaimer
OpenAI Codex with GPT-6.1 assisted with implementation, benchmark execution and analysis, and drafting this PR.