You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit b1c36c1
Browse filesBrowse the repository at this point in the historyBrowse files
Batch DF derivative work and reduce gradient workspace (#967)
* df/grad/td: stabilize metric application with factor products
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.
* df/grad: size scratch for auxiliary tiles and density factors
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.
* df/grad,hessian: pack auxiliary groups in derivative launches
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.
* df/int3c2e_bdiv: pack auxiliary batches
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.
0 commit comments