Skip to content

[BUG REPORT] nao mismatch causes RSJK PBC DFT failures with contracted basis sets #846

Open
Open
[BUG REPORT] nao mismatch causes RSJK PBC DFT failures with contracted basis sets#846
Assignees
sunqm

Description

@lhan67
lhan67
opened on Aug 3, 2026

Describe the bug
For periodic systems using contracted basis sets such as aug-cc-pVDZ, GPU4PySCF RSJK calculations fail because the number of atomic orbitals (nao) of the contracted basis is mixed with that of the internally uncontracted basis. The same AO-dimension mismatch leads to different failures depending on the J-engine used:

  1. With PBCJMatrixOpt, the calculation fails with an assertion error.

  2. With PBCJKMatrixOpt, the nao inconsistency sometimes causes the intermediate Gpq to become unexpectedly large, especially when the box is sufficiently big, eventually resulting in a GPU out-of-memory error.

To Reproduce
Case 1: PBCJMatrixOpt

import numpy as np
from pyscf.pbc import gto
from gpu4pyscf.pbc.dft import BeckeGrids
from gpu4pyscf.pbc.scf.j_engine import PBCJMatrixOpt

atoms = """
        N         16.23000       14.39000       12.87000
        C         15.94000       15.71000       12.83000
        H         16.78000       16.40000       12.90000
        C         14.66000       16.13000       12.69000
        H         14.39000       17.17000       12.65000
        C         13.66000       15.10000       12.60000
        N         12.36000       15.42000       12.46000
        H         12.10000       16.40000       12.43000
        H         11.66000       14.71000       12.40000
        N         13.97000       13.80000       12.64000
        C         15.26000       13.41000       12.78000
        O         15.60000       12.21000       12.82000
        H         17.19000       14.08000       12.97000
        H          9.09000        7.95000       12.19000
        N          9.10000        8.95000       12.18000
        C          8.00000        9.78000       12.06000
        H          7.01000        9.38000       11.96000
        N          8.32000       11.05000       12.09000
        C          9.70000       11.06000       12.23000
        C         10.63000       12.14000       12.31000
        O         10.38000       13.35000       12.28000
        N         11.95000       11.70000       12.45000
        H         12.66000       12.42000       12.52000
        C         12.33000       10.37000       12.50000
        N         13.65000       10.15000       12.64000
        H         13.95000        9.19000       12.67000
        H         14.33000       10.91000       12.70000
        N         11.48000        9.36000       12.42000
        C         10.20000        9.76000       12.29000
        """


cell = gto.M(
    atom=atoms,
    a=np.eye(3) * 25.0,
    basis="aug-cc-pvdz",
    verbose=3,
)
mf = cell.RKS(xc="pbe").to_gpu().density_fit()
mf.grids = BeckeGrids(cell)
mf.grids.atom_grid = (150, 974)
mf.j_engine = PBCJMatrixOpt(cell)

mf.kernel()

This calculation fails with:

Traceback (most recent call last):
  File "/var/lib/condor/execute/slot1/dir_2456861/scratch/./PBC_DFT.py", line 79, in <module>
    main()
  File "/var/lib/condor/execute/slot1/dir_2456861/scratch/./PBC_DFT.py", line 67, in main
    e = mf.kernel(dm0=dm0)
        ^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/scf/hf.py", line 364, in scf
    _kernel(mf, mf.conv_tol, mf.conv_tol_grad,
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/scf/hf.py", line 222, in _kernel
    vhf = mf.get_veff(mol, dm)
          ^^^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/dft/rks.py", line 104, in get_veff
    vj, vk, vj_sr, vk_sr = _get_jk(
                           ^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/dft/krks.py", line 112, in _get_jk
    vj = mf.j_engine._get_j_lr(dm, hermi, kpts, kpts_band)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/scf/j_engine.py", line 351, in _get_j_lr
    return get_j_kpts(self, dm, hermi, kpts, kpts_band)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/df/aft_jk.py", line 56, in get_j_kpts
    for Gpq, p0, p1 in mydf.ft_loop(q=kpt_allow, kpts=kpts):
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/df/aft.py", line 251, in ft_loop
    dat = ft_kern(Gv[p0:p1], q, kpts)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/df/ft_ao.py", line 504, in ft_kernel
    fill_triu_bvk(out.view(np.float64), nao, self.bvk_kmesh,
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/df/int3c2e.py", line 191, in fill_triu_bvk
    assert a.shape[0] == nao*bvk_ncells*nao
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

Case 2: PBCJKMatrixOpt
Using the same system and calculation settings, replace the J-engine with:

from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt

mf.j_engine = PBCJKMatrixOpt(cell)
mf.kernel()

On one H200 with over 140GB memory, This calculation fails with:

Traceback (most recent call last):
  File "/var/lib/condor/execute/slot2/dir_837936/scratch/./PBC_DFT.py", line 79, in <module>
    main()
  File "/var/lib/condor/execute/slot2/dir_837936/scratch/./PBC_DFT.py", line 67, in main
    e = mf.kernel(dm0=dm0)
        ^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/scf/hf.py", line 364, in scf
    _kernel(mf, mf.conv_tol, mf.conv_tol_grad,
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/scf/hf.py", line 222, in _kernel
    vhf = mf.get_veff(mol, dm)
          ^^^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/dft/rks.py", line 104, in get_veff
    vj, vk, vj_sr, vk_sr = _get_jk(
                           ^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/dft/krks.py", line 112, in _get_jk
    vj = mf.j_engine._get_j_lr(dm, hermi, kpts, kpts_band)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/scf/rsjk.py", line 883, in _get_j_lr
    aft_jk._update_vj_(vj, Gpq, dms, wcoulG[p0:p1])
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/pbc/df/aft_jk.py", line 69, in _update_vj_
    rho = contract('nkij,kgij->ng', dms, Gpq.conj())
                                         ^^^^^^^^^^
  File "cupy/_core/core.pyx", line 1478, in cupy._core.core._ndarray_base.conj
  File "cupy/_core/core.pyx", line 1479, in cupy._core.core._ndarray_base.conj
  File "cupy/_core/_routines_math.pyx", line 35, in cupy._core._routines_math._ndarray_conj
  File "cupy/_core/_kernel.pyx", line 1349, in cupy._core._kernel.ufunc.__call__
  File "cupy/_core/_kernel.pyx", line 645, in cupy._core._kernel._get_out_args_from_optionals
  File "cupy/_core/core.pyx", line 2884, in cupy._core.core._ndarray_init
  File "cupy/_core/core.pyx", line 257, in cupy._core.core._ndarray_base._init_fast
  File "cupy/cuda/memory.pyx", line 738, in cupy.cuda.memory.alloc
  File "/home/pyscf_user/gpu4pyscf/gpu4pyscf/lib/cupy_helper.py", line 1175, in malloc
    return cuda_malloc(size)
           ^^^^^^^^^^^^^^^^^
  File "cupy/cuda/memory.pyx", line 633, in cupy.cuda.memory._malloc
  File "cupy/cuda/memory.pyx", line 634, in cupy.cuda.memory._malloc
  File "cupy/cuda/memory.pyx", line 101, in cupy.cuda.memory.Memory.__init__
  File "cupy_backends/cuda/api/runtime.pyx", line 499, in cupy_backends.cuda.api.runtime.malloc
  File "cupy_backends/cuda/api/runtime.pyx", line 146, in cupy_backends.cuda.api.runtime.check_status
cupy_backends.cuda.api.runtime.CUDARuntimeError: cudaErrorMemoryAllocation: out of memory

Version and hardware (please complete the following information):

  • GPU4PySCF version [latest development version, commit 848bf11288104ba44617240fffbf307f56056067]
  • PySCF install type [built from source]
  • Operating system [Ubuntu]
  • GPU model [H200]

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions