practice on ParallelMPO for mem and speed management.

[jiangtong1000]

Hi Huanchen @hczhai

I’m running a ~large-scale DMRG calculation on a 78-orbital system. The MPO bond dimension is 6320.

On a SerialMPO run, using a single node with ~1000 GB memory, I can push the MPS bond dimension up to around 4000 before hitting OOM, and the calculation becomes very slow at that point. I have not yet tried the MPO reloading trick in this setup, but I plan to test that separately.

To speed things up, I tried using ParallelMPO:

• I built a global serial MPO and split it into 78 bins (Pham = 78).
• Then I launched a 78-rank job, where each rank reads its own mpo_rank.bin.
• I also used the minimal=True flag when loading the MPO.

However, in this parallel setup the calculation hits OOM already at M ~ 500.

My initial intuition was that sub-Hamiltonian parallelization would keep the total memory usage roughly similar to the serial case, since each sub-Hamiltonian would have its own environment and we’d just add their contributions. Clearly I’m misunderstanding something about how memory scales here, especially when Pham is large (I noticed in your 2021 JCP paper that the largest Pham used was about 16).

So I was wondering that

1. what is the best practice for choosing Pham and the MPI layout in this regime?
2. does ParallelMPO inherently increase per-node memory usage at large Pham (due to replicated MPS/environments or additional per-rank data), or is this likely a misconfiguration on my side?
3. Are there any specific options you recommend to keep memory under control when combining ParallelMPO?

Any guidance or rule-of-thumb setup would be very helpful. Thanks a lot for your time and for block2!

Below are my running scripts:

from pyblock2.driver.core import DMRGDriver, SymmetryTypes
import h5py
with h5py.File("../0_Gen_Itg/dmrg_input.h5", "r") as f:
    h1e = f["h1e"][:]
    g2e = f["g2e"][:]
    occs = f["occs"][:]
    ncas = h1e.shape[0]

from block2 import *
from block2.su2 import *
import os
    
nelec = 48
mpi_size = 78
spin = 0
ecore = 0
orb_sym = None
scratch_dir = "/n/netscratch/tmp"
mpo_dir = "/n/netscratch/mpo_tmp"

print('ncas', ncas, 'n_elec', nelec, 'spin', spin, 'ecore', ecore, 'h1e', h1e.shape, 'g2e', g2e.shape, 'orb_sym', orb_sym)
driver = DMRGDriver(scratch=scratch_dir, symm_type=SymmetryTypes.SU2, n_threads=12, stack_mem=int(300*1024*1024*1024))

print('initializing system')
driver.initialize_system(n_sites=ncas, n_elec=nelec, spin=spin, orb_sym=orb_sym)
print('getting mpo')
mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, iprint=1)
mpo.reduce_data()

# size, rank, root
comm = ParallelCommunicator(mpi_size, 0, 0)
prule = ParallelRuleQC(comm)

for irank in range(comm.size):
    comm.rank = irank
    para_mpo = ParallelMPO(mpo, prule)
    para_mpo.save_data(f'{mpo_dir}/mpo.bin.{irank}')
    fsize = os.path.getsize(f'{mpo_dir}/mpo.bin.{irank}')
    print("mpo.%d size = %10s" % (irank, Parsing.to_size_string(fsize)))

import h5py
from pyblock2.driver.core import DMRGDriver, SymmetryTypes
import psutil
import os

nelec = 48
spin = 0
ecore = 0
orb_sym = None

with h5py.File("../0_Gen_Itg/dmrg_input.h5", "r") as f:
    h1e = f["h1e"][:]
    occs = f["occs"][:]
    ncas = h1e.shape[0]

# Custom interface """
from block2 import *
from block2.su2 import *
MPI = MPICommunicator()

mpo_dir = "/n/netscratch/mpo_tmp"
scratch_dir = "/n/netscratch/tmp"

n_threads = Global.threading.n_threads_global // MPI.size
driver = DMRGDriver(scratch=scratch_dir, symm_type=SymmetryTypes.SU2, n_threads=n_threads, stack_mem=int(12*1024*1024*1024), mpi=True)
bond_dims = [500]*2 + [800]*2 + [2000] * 2 + [4000] * 2 + [5000] * 20
noises = [1e-4] * 2 + [5e-5] * 2 + [0] * 22
thrds = [1e-6] * 4 + [1e-7] * 8 + [1e-9] * 8

prule = ParallelRuleQC(MPI)

mem = psutil.Process(os.getpid()).memory_info().rss
print(" pre-load-mpo memory usage = %10s" % Parsing.to_size_string(mem))

mpo = ParallelMPO(0, prule)
mpo.load_data(f'{mpo_dir}/mpo.bin.{MPI.rank}', minimal=True)

mem = psutil.Process(os.getpid()).memory_info().rss
print("post-load-mpo memory usage = %10s" % Parsing.to_size_string(mem))
# """

driver.initialize_system(n_sites=ncas, n_elec=nelec, spin=spin, orb_sym=orb_sym)
ket = driver.get_random_mps(tag="GS", bond_dim=500, nroots=1, occs=occs)
# ket = driver.load_mps(tag="GS", nroots=1)
# sweep_start = 0
# forward = sweep_start % 2 == 0
energy = driver.dmrg(mpo, ket, n_sweeps=28, bond_dims=bond_dims, noises=noises, thrds=thrds, iprint=2, twosite_to_onesite=16) #forward=forward, sweep_start=sweep_start)
print('DMRG energy (variational) = %20.15f' % energy)

[hczhai]

For standard quantum chemistry Hamiltonians, block2 has specialized optimizations for MPO construction and parallelization, and there is no need for the user to manually set the quantum chemistry MPO distribution. When you use the Python API interface, you need some additional arguments to activate these optimizations, see below. See also #156.

In case you are interested, the recommended approach for manually setting the MPO distribution in Python API is given in #145 (reply in thread). But this is mostly targeting custom Hamiltonians and it will not give you the speed and memory cost comparable to the "1" and "2" given below.

1. To get the best performance and memory consumption for single node calculation, the recommended script is (using MPOAlgorithmTypes.Conventional for specialized construction of quantum chemistry MPO, details are given in https://block2.readthedocs.io/en/latest/theory/su2.html):

import numpy as np
from pyblock2._pyscf.ao2mo import integrals as itg
from pyblock2.driver.core import DMRGDriver, SymmetryTypes, MPOAlgorithmTypes

bond_dims = [250] * 4 + [500] * 4
noises = [1e-4] * 4 + [1e-5] * 4 + [0]
thrds = [1e-10] * 8

from pyscf import gto, scf

mol = gto.M(atom="N 0 0 0; N 0 0 1.1", basis="sto3g", symmetry="d2h", verbose=0)
mf = scf.RHF(mol).run(conv_tol=1E-14)
ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(mf, ncore=0, ncas=None, g2e_symm=8)

driver = DMRGDriver(scratch="./tmp", symm_type=SymmetryTypes.SU2, n_threads=64, mpi=False)
driver.initialize_system(n_sites=ncas, n_elec=n_elec, spin=spin, orb_sym=orb_sym)

mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, algo_type=MPOAlgorithmTypes.Conventional, iprint=1)
ket = driver.get_random_mps(tag="GS", bond_dim=250, nroots=1)
energy = driver.dmrg(mpo, ket, n_sweeps=20, bond_dims=bond_dims, noises=noises, lowmem_noise=True,
    dav_def_max_size=30, thrds=thrds, iprint=1)
print('DMRG energy = %20.15f' % energy)

2. To get the best performance and memory consumption for multi-node calculation, the recommended script is (using MPOAlgorithmTypes.Conventional for specialized parallelization strategy of quantum chemistry MPO, according to J. Chem. Phys. 120, 3172–3178 (2004), which is conceptually equivalent, but manually optimized version of the sub-Hamiltonian parallelism described in J. Chem. Phys. 154, 224116 (2021)):

import numpy as np
from pyblock2._pyscf.ao2mo import integrals as itg
from pyblock2.driver.core import DMRGDriver, SymmetryTypes, MPOAlgorithmTypes, ParallelTypes

bond_dims = [250] * 4 + [500] * 4
noises = [1e-4] * 4 + [1e-5] * 4 + [0]
thrds = [1e-10] * 8

from pyscf import gto, scf

mol = gto.M(atom="N 0 0 0; N 0 0 1.1", basis="sto3g", symmetry="d2h", verbose=0)
mf = scf.RHF(mol).run(conv_tol=1E-14)
ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(mf, ncore=0, ncas=None, g2e_symm=8)

driver = DMRGDriver(scratch="./tmp", symm_type=SymmetryTypes.SU2, n_threads=64, mpi=True)
driver.prule = driver.bw.bs.ParallelRuleQC(driver.mpi)

driver.initialize_system(n_sites=ncas, n_elec=n_elec, spin=spin, orb_sym=orb_sym)

mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, algo_type=MPOAlgorithmTypes.Conventional,
    para_type=ParallelTypes.Nothing, iprint=2)
ket = driver.get_random_mps(tag="GS", bond_dim=250, nroots=1)
energy = driver.dmrg(mpo, ket, n_sweeps=20, bond_dims=bond_dims, noises=noises, lowmem_noise=True,
    dav_def_max_size=30, thrds=thrds, iprint=1)
print('DMRG energy = %20.15f' % energy)

and executed as

export OMP_NUM_THREADS=64
mpirun --map-by ppr:$SLURM_TASKS_PER_NODE:node:pe=$OMP_NUM_THREADS python3 -u dmrg.py > dmrg.out

3. Unoptimized implementation for sub-Hamiltonian parallelization as described in J. Chem. Phys. 145, 014102 (2016) can be activated by modifying the above script as:

# delete: driver.prule = driver.bw.bs.ParallelRuleQC(driver.mpi)
mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, algo_type=MPOAlgorithmTypes.FastBipartite,
    para_type=ParallelTypes.I, iprint=2)

4. Unoptimized implementation for sub-Hamiltonian parallelization as described in J. Chem. Phys. 154, 224116 (2021) can be activated using (not recommended for typical parallel DMRG cases, but maybe useful when the number of nodes is larger than 50)

# delete: driver.prule = driver.bw.bs.ParallelRuleQC(driver.mpi)
mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore, algo_type=MPOAlgorithmTypes.FastBipartite,
    para_type=ParallelTypes.SIJ, iprint=2)