Multi-GPU execution

Sharding is one value on the builder: mesh() spreads the XC quadrature over the device mesh and builds the density-fitting 3-center tensor directly in per-device slabs, so no device ever holds more than its naux/ndev slice. The numbers match single-device execution; the collectives differentiate, so SCF and forces are unchanged.

The streamed backend, df(chunk=...), shards as well and holds no tensor at all; that is the combination that reaches protein scale, since the slabs themselves stop fitting long before. RI-J, RI-K and range-separated exchange all shard, so hybrids run on a mesh too.

Executed on a 4×A100 node; outputs are committed with the notebook.

import jax; jax.config.update("jax_enable_x64", True)
print(jax.devices())

[CudaDevice(id=0), CudaDevice(id=1), CudaDevice(id=2), CudaDevice(id=3)]

E0725 14:15:54.218694 1279283 numa_hwloc.cc:121] Call to hwloc_set_cpubind() failed: Invalid argument [22]

from dftax import KS, Molecule, becke, df, mesh, scf
from dftax.energy.xc import PBE0

water = Molecule.from_xyz("O 0 0 0; H 0.7586 0 0.5043; H 0.7586 0 -0.5043", "sto-3g")
AUX = "def2-universal-jkfit"

ks1 = KS(water, PBE0(), grid=becke(50, 110), coulomb=df(AUX))
ksm = KS(water, PBE0(), grid=becke(50, 110), coulomb=df(AUX), mesh=mesh())
print("int3c global:", ksm.coulomb.int3c.shape)
print("per-device slabs:", sorted({s.data.shape for s in ksm.coulomb.int3c.addressable_shards}))

int3c global: (7, 7, 136)
per-device slabs: [(7, 7, 34)]

r1, rm = scf(ks1), scf(ksm)
print(f"single device: {r1.e_tot:.12f}")
print(f"4-GPU mesh   : {rm.e_tot:.12f}")
print(f"difference   : {abs(r1.e_tot - rm.e_tot):.2e}")

single device: -75.167040839943
4-GPU mesh   : -75.167047968814
difference   : 7.13e-06