Added memory sharing for most rank-identical objects

dgamore/DGAmore.py

@@ -33,6 +33,7 @@
 import dgamore.eliashberg_solver as eliashberg_solver
 import dgamore.local_sde as local_sde
 import dgamore.memory_estimator as memory_estimator
+import dgamore.mpi_utils as mpi_utils
 import dgamore.nonlocal_sde as nonlocal_sde
 import dgamore.plotting as plotting
 from dgamore import max_ent
@@ -47,42 +48,7 @@
 
 # Fraction of a node's memory budget a run may occupy at any branch's peak. The remainder is the sole overhead
 # margin (OS/allocator/transients); the estimator's OVERHEAD_FACTOR is 1.0, so the two do not compound.
-NODE_MEMORY_FRACTION: float = 0.97
-
-
-def _cgroup_memory_limit() -> int | None:
-    """
-    Returns this process's effective cgroup memory limit in bytes, or ``None`` when no limit is set (or none is
-    readable). Batch schedulers such as slurm enforce a job's memory request through a cgroup, which can be far
-    below the node's physical memory, so the node budget must honor it. The cgroup path is taken from
-    ``/proc/self/cgroup``; on cgroup v2 every ancestor's ``memory.max`` is read up to the root (the limit may sit on
-    the job level rather than the process's own leaf) and the smallest set value wins, on v1 the controller's
-    ``memory.limit_in_bytes`` is read directly. Values at or above ``2**62`` mean "unlimited" and are ignored.
-
-    :return: The smallest configured limit in bytes, or ``None`` if unlimited or undeterminable.
-    """
-    limits = []
-    try:
-        entries = dict(
-            (line.split(":", 2)[1], line.split(":", 2)[2].strip()) for line in open("/proc/self/cgroup", "r")
-        )
-        if "" in entries:  # cgroup v2: one unified hierarchy, limits possibly on an ancestor
-            path = os.path.normpath("/sys/fs/cgroup" + entries[""])
-            while path.startswith("/sys/fs/cgroup"):
-                try:
-                    value = open(os.path.join(path, "memory.max"), "r").read().strip()
-                    if value != "max":
-                        limits.append(int(value))
-                except OSError:
-                    pass
-                path = os.path.dirname(path)
-        elif "memory" in entries:  # cgroup v1: the memory controller's own hierarchy
-            value = open("/sys/fs/cgroup/memory" + entries["memory"] + "/memory.limit_in_bytes", "r").read()
-            limits.append(int(value))
-    except (OSError, ValueError, IndexError):
-        return None
-    limits = [limit for limit in limits if limit < 2**62]
-    return min(limits) if limits else None
+NODE_MEMORY_FRACTION: float = 0.95
 
 
 def main():
@@ -591,7 +557,7 @@ def autodetect_memory_settings(comm: MPI.Comm) -> None:
     transient is held by every rank at once, a *single-rank* transient by one rank while the others idle), minus
     ``(r - 1) * giwk_shareable`` for the branch's node-shared ``giwk_full`` window, and this must not exceed
     ``NODE_MEMORY_FRACTION`` times the node budget - ``psutil.virtual_memory().available``, capped by the cgroup
-    memory limit when the scheduler sets one (see :func:`_cgroup_memory_limit`). Each
+    memory limit when the scheduler sets one (see :func:`dgamore.mpi_utils.cgroup_memory_limit`). Each
     node's rank count and available memory are collected with a single ``allgather`` of
     ``(hostname, available_bytes)``; a branch's path is judged to "fit" only if it fits on **every** node (the flags
     are process-wide, so the tightest node governs, and a single-rank transient may land on any node). The
@@ -608,7 +574,7 @@ def autodetect_memory_settings(comm: MPI.Comm) -> None:
     # Gather (hostname, available bytes) from every rank in a single collective and reduce to one entry per node:
     # the rank count and the (minimum, conservative) available memory on that node.
     node_available = psutil.virtual_memory().available
-    cgroup_limit = _cgroup_memory_limit()
+    cgroup_limit = mpi_utils.cgroup_memory_limit()
     if cgroup_limit is not None:
         node_available = min(node_available, cgroup_limit)
     hostname = socket.gethostname()
@@ -630,6 +596,7 @@ def autodetect_memory_settings(comm: MPI.Comm) -> None:
         niv_core=config.box.niv_core,
         niv_full=config.box.niv_full,
         niv_cut=niv_cut,
+        niv_dmft=config.box.niv_dmft,
         niv_pp=niv_pp,
         n_ranks=comm.size,
         with_eliashberg=config.eliashberg.perform_eliashberg,

dgamore/eliashberg_solver.py

@@ -953,9 +953,19 @@ def solve_eliashberg_lanczos_grid(
     block_mat = _gather_grid_vertex_block(gamma_r_pp, comm, row_slice, col_slice, in_grid)
     gamma_r_pp.free()
 
-    chi0_mat = mpi_utils.bcast_rows(
-        comm, gchi0_q0_pp.mat if comm.rank == bubble_rank else np.empty(0), root=bubble_rank
-    )
+    # the ranks beyond the rows x cols grid never touch the bubble: ship it inside the grid only, so the idle
+    # ranks do not each receive a full-BZ copy (Split keeps the rank order, so the grid root index is bubble_rank)
+    grid_comm = comm.Split(0 if in_grid else 1, comm.rank)
+    if bubble_rank < rows * cols:
+        chi0_mat = (
+            mpi_utils.bcast_rows(grid_comm, gchi0_q0_pp.mat if comm.rank == bubble_rank else None, root=bubble_rank)
+            if in_grid
+            else None
+        )
+    else:
+        chi0_mat = mpi_utils.bcast_rows(
+            comm, gchi0_q0_pp.mat if comm.rank == bubble_rank else np.empty(0), root=bubble_rank
+        )
     if not in_grid:
         return None
 
@@ -1652,6 +1662,7 @@ def dispatch_full_vertex_calculation(
     niv_pp: int,
     mpi_dist: MpiDistributor,
     chunk_bytes: int | None = None,
+    node_comm: MPI.Comm | None = None,
 ) -> FourPoint:
     r"""
     Loads the local irreducible vertex for ``channel`` and builds the full ladder pp vertex through the slice-direct
@@ -1680,14 +1691,21 @@ def dispatch_full_vertex_calculation(
     :param niv_pp: Number of positive fermionic frequencies of the pp vertex.
     :param mpi_dist: MPI distributor over the irreducible BZ q-points.
     :param chunk_bytes: Chunk byte budget of the build (``None`` uses the floor).
+    :param node_comm: Optional node-local communicator; when given, the multi-GB local vertex is loaded once per
+        node into an MPI shared-memory window instead of once per rank (the build only reads it).
     :return: The full ladder pp vertex :math:`F^{\mathrm{q}}_{r}` as a :class:`FourPoint`.
     """
-    gamma_r = LocalFourPoint.load(os.path.join(config.output.output_path, f"gamma_{channel.value}_loc.npy"), channel)
+    gamma_r, gamma_win = nonlocal_sde._load_node_shared_local_vertex(
+        node_comm, os.path.join(config.output.output_path, f"gamma_{channel.value}_loc.npy"), channel
+    )
     if config.eliashberg.save_fq:
         f_q_r = create_pairing_vertex_streaming_fq(u_loc, v_nonloc, gamma_r, niv_pp, mpi_dist, chunk_bytes)
     else:
         f_q_r = create_pairing_vertex_slice_q_r(u_loc, v_nonloc, gamma_r, niv_pp, mpi_dist, chunk_bytes)
-    gamma_r.free()
+    gamma_r.mat = None
+    if gamma_win is None:
+        gamma_r.free()
+    nonlocal_sde._free_shared_window(gamma_win, node_comm)
     mpi_dist.barrier()
     return f_q_r
 
@@ -1824,11 +1842,16 @@ def solve(
         config.eliashberg.n_eig,
         comm.size,
     )
+    # the node budget honors the scheduler's cgroup memory limit (e.g. slurm --mem), like the driver's fit check
+    node_budget = psutil.virtual_memory().available
+    cgroup_limit = mpi_utils.cgroup_memory_limit()
+    if cgroup_limit is not None:
+        node_budget = min(node_budget, cgroup_limit)
+
     # one full-BZ sector residency per rank picks the in-memory solve, otherwise the grid takes over; rank 0
     # decides and broadcasts, so ranks on differently loaded nodes can never pick different solvers
     use_grid = FORCE_GRID_SOLVER or (
-        comm.size > 1
-        and per_sector_bytes + giwk_dga.mat.nbytes > psutil.virtual_memory().available * NODE_MEMORY_FRACTION
+        comm.size > 1 and per_sector_bytes + giwk_dga.mat.nbytes > node_budget * NODE_MEMORY_FRACTION
     )
     use_grid = comm.bcast(use_grid, root=0)
     if use_grid:
@@ -1850,7 +1873,7 @@ def solve(
         )
     else:
         sing_ranks, trip_ranks = get_ranks_for_lanczos(
-            comm, len(parities), psutil.virtual_memory().available, per_sector_bytes, giwk_dga.mat.nbytes
+            comm, len(parities), node_budget, per_sector_bytes, giwk_dga.mat.nbytes
         )
         bubble_rank = sing_ranks[0]
         n_concurrent = len(set(sing_ranks) | set(trip_ranks))
@@ -1867,11 +1890,15 @@ def solve(
 
     node_comm = comm.Split_type(MPI.COMM_TYPE_SHARED) if comm.size > 1 else None
     chunk_bytes = memory_estimator.dynamic_chunk_budget(
-        psutil.virtual_memory().total, node_comm.size if node_comm is not None else 1
+        mpi_utils.job_memory_total(), node_comm.size if node_comm is not None else 1
     )
 
-    f_dens_pp = dispatch_full_vertex_calculation(SpinChannel.DENS, u_loc, v_nonloc, niv_pp, mpi_dist_irrk, chunk_bytes)
-    f_magn_pp = dispatch_full_vertex_calculation(SpinChannel.MAGN, u_loc, v_nonloc, niv_pp, mpi_dist_irrk, chunk_bytes)
+    f_dens_pp = dispatch_full_vertex_calculation(
+        SpinChannel.DENS, u_loc, v_nonloc, niv_pp, mpi_dist_irrk, chunk_bytes, node_comm
+    )
+    f_magn_pp = dispatch_full_vertex_calculation(
+        SpinChannel.MAGN, u_loc, v_nonloc, niv_pp, mpi_dist_irrk, chunk_bytes, node_comm
+    )
 
     delete_files(config.output.eliashberg_path, f"gchi0_q_inv_rank_{comm.rank}.npy")
 
@@ -1957,7 +1984,7 @@ def solve(
 
 
 # Fraction of a node's available host memory the sector packing may occupy (mirrors DGAmore.NODE_MEMORY_FRACTION).
-NODE_MEMORY_FRACTION: float = 0.97
+NODE_MEMORY_FRACTION: float = 0.95
 
 
 def get_ranks_for_lanczos(

dgamore/greens_function.py

@@ -101,6 +101,34 @@ def root_fun(
     return get_total_fill(mu, ek, sigma_mat, beta, smom0) - target_filling
 
 
+# Largest filling residual accepted from a "converged" Newton result; genuine roots sit orders of magnitude below
+# this, while a secant that stalled in a flat filling region leaves a residual of order the filling itself.
+_FILL_RESIDUAL_TOL: float = 1e-3
+
+
+def _find_mu_bracket(
+    mu0: float, args: tuple, initial_width: float = 0.05, max_width: float = 64.0
+) -> tuple[float, float] | None:
+    r"""
+    Expands a symmetric interval around ``mu0``, doubling its half-width each step, until the filling residual
+    :func:`root_fun` changes sign across it. Starting narrow makes the search return an interval around the root
+    closest to ``mu0``, which keeps a self-consistency trajectory in its current basin.
+
+    :param mu0: Center of the search interval.
+    :param args: The :func:`root_fun` arguments after ``mu`` (target filling, dispersion, self-energy, beta, moment).
+    :param initial_width: Half-width of the first interval.
+    :param max_width: Half-width beyond which the search gives up.
+    :return: A bracketing interval ``(lo, hi)``, or ``None`` if no sign change was found.
+    """
+    width = initial_width
+    while width <= max_width:
+        lo, hi = mu0 - width, mu0 + width
+        if root_fun(lo, *args) * root_fun(hi, *args) < 0:
+            return lo, hi
+        width *= 2
+    return None
+
+
 def update_mu(
     mu0: float,
     target_filling: float,
@@ -113,26 +141,39 @@ def update_mu(
 ) -> float:
     r"""
     Updates the chemical potential to match the target filling by using Newton's method to find the optimal
-    :math:`\mu`. On failure to converge the starting value is returned unchanged.
+    :math:`\mu`. A Newton result is only accepted if its filling residual is small; when Newton fails or lands
+    away from an actual root, the root nearest the starting value is found instead with a bracketed Brent search
+    (see :func:`_find_mu_bracket`). The starting value is returned unchanged only if no bracket exists.
 
     :param mu0: Initial guess for the chemical potential.
     :param target_filling: Desired total filling.
     :param ek: Band dispersion :math:`\varepsilon(\mathbf{k})`.
     :param sigma_mat: Self-energy array, shape ``[k, o1, o2, v]``.
     :param beta: Inverse temperature :math:`\beta`.
     :param smom0: Zeroth moment :math:`\Sigma_\infty` of the self-energy.
-    :param logger: Optional logger; if given, a failed root search is logged at debug level.
-    :param tol: Newton tolerance for the root search.
-    :return: The updated (real) chemical potential, or ``mu0`` if the root search did not converge.
+    :param logger: Optional logger; if given, the bracketed fallback is logged at info level and a fully failed
+        root search at warning level.
+    :param tol: Root search tolerance for the chemical potential.
+    :return: The updated (real) chemical potential, or ``mu0`` if no root was found.
     :raises ValueError: If the converged chemical potential has a non-negligible imaginary part.
     """
     mu = mu0
+    args = (target_filling, ek, sigma_mat, beta, smom0)
     try:
-        mu = opt.newton(root_fun, mu, args=(target_filling, ek, sigma_mat, beta, smom0), tol=tol)
+        mu = opt.newton(root_fun, mu, args=args, tol=tol)
+        # the secant step criterion can also "converge" inside a flat filling region far from any root, so the
+        # residual is verified before the value is accepted
+        if np.abs(root_fun(mu, *args)) > _FILL_RESIDUAL_TOL:
+            raise RuntimeError
     except RuntimeError:
+        bracket = _find_mu_bracket(mu0, args)
+        if bracket is None:
+            if logger is not None:
+                logger.warning("Root finding for chemical potential failed; keeping the previous value.")
+            return mu0
         if logger is not None:
-            logger.debug("Root finding for chemical potential failed.")
-        return mu0
+            logger.info("Newton did not find a chemical potential root; using a bracketed root search.")
+        mu = opt.brentq(root_fun, *bracket, args=args, xtol=tol)
 
     if np.abs(mu.imag) < 1e-8:
         mu = mu.real
@@ -233,7 +274,9 @@ def occ_k(self) -> np.ndarray:
     def get_g_full(siw: SelfEnergy, mu: float, ek: np.ndarray, beta: float):
         r"""
         Builds the full momentum-dependent Green's function :math:`G^{\mathrm{k}} = [(\imath\nu + \mu) -
-        \varepsilon(\mathbf{k}) - \Sigma^{\mathrm{k}}]^{-1}`.
+        \varepsilon(\mathbf{k}) - \Sigma^{\mathrm{k}}]^{-1}`. The Dyson matrix is assembled and inverted in bounded
+        fermionic-frequency chunks, so beyond the result only one chunk-sized transient is alive at a time (large
+        boxes such as the full DMFT one would otherwise triple the peak).
 
         :param siw: The :class:`SelfEnergy` :math:`\Sigma`.
         :param mu: Chemical potential :math:`\mu`.
@@ -243,15 +286,17 @@ def get_g_full(siw: SelfEnergy, mu: float, ek: np.ndarray, beta: float):
         """
         eye_bands = np.eye(siw.n_bands, siw.n_bands)
         iv = 1j * MFHelper.vn(siw.niv, beta)
-        iv_bands = iv[None, None, :] * eye_bands[..., None]
-        mu_bands = mu * eye_bands[:, :, None]
-        mat = (
-            iv_bands[None, None, None, ...]
-            + mu_bands[None, None, None, ...]
-            - ek[..., None]
-            - siw.decompress_q_dimension().mat
-        )
-        mat = GreensFunction._invert_last_orbital_block(mat)
+        static_bands = (mu * eye_bands)[None, None, None, :, :, None] - ek[..., None]
+        sigma_mat = siw.decompress_q_dimension().mat
+
+        # a momentum-local sigma ([1, 1, 1, ...]) broadcasts against the dispersion, so the result is always full-k
+        mat = np.empty((*ek.shape[:3], siw.n_bands, siw.n_bands, 2 * siw.niv), dtype=sigma_mat.dtype)
+        step = max(1, _MODEL_EPOT_CHUNK_ELEMENTS // (int(np.prod(ek.shape[:3])) * siw.n_bands**2))
+        for start in range(0, 2 * siw.niv, step):
+            stop = min(2 * siw.niv, start + step)
+            dyson = (iv[start:stop][None, None, :] * eye_bands[..., None])[None, None, None, ...] + static_bands
+            dyson -= sigma_mat[..., start:stop]
+            mat[..., start:stop] = GreensFunction._invert_last_orbital_block(dyson)
         return GreensFunction(mat, siw, ek, siw.full_niv_range, False, False, nk=ek.shape[:3], beta=beta, mu=mu)
 
     @staticmethod

dgamore/max_ent.py

@@ -79,21 +79,24 @@ def perform_maxent_giwk(giwk: GreensFunction, name: str, comm: MPI.Comm):
     logger = config.logger
 
     logger.info(f"Starting analytic continuation of the {name} Green's function using the maximum entropy method.")
-    giwk_maxent = giwk.cut_niv(config.box.niv_core).to_half_niv_range()
-
-    # Rotate G into the band (H(k)-eigen) basis first, so the band-diagonal below is the band-resolved spectral function
-    # (invariant under lattice symmetries, so the naive irrk_inv unfold is correct); done on the full BZ.
-    hk = config.lattice.hamiltonian.get_ek(config.lattice.k_grid)
-    giwk_maxent = giwk_maxent.decompress_q_dimension()
-    orbital_to_band_basis(hk, giwk_maxent.mat)
-
     irrq_list = config.lattice.k_grid.get_irrq_list()
-
     mpi_dist = MpiDistributor(ntasks=len(irrq_list), comm=comm, name="Maxent_G", output_path=config.output.output_path)
 
-    giwk_maxent = giwk_maxent.reduce_q(irrq_list)
+    # The full-BZ preparation runs on rank 0 only: every rank used to build the identical core-cut, band-rotated,
+    # irr-reduced Green's function just to keep its own scatter slice - multi-GB transients times the rank count.
+    g_irr_mat = None
+    if comm.rank == 0:
+        giwk_maxent = giwk.cut_niv(config.box.niv_core).to_half_niv_range()
+
+        # Rotate G into the band (H(k)-eigen) basis first, so the band-diagonal below is the band-resolved spectral
+        # function (invariant under lattice symmetries, so the naive irrk_inv unfold is correct); on the full BZ.
+        hk = config.lattice.hamiltonian.get_ek(config.lattice.k_grid)
+        giwk_maxent = giwk_maxent.decompress_q_dimension()
+        orbital_to_band_basis(hk, giwk_maxent.mat)
+        g_irr_mat = giwk_maxent.reduce_q(irrq_list).mat
+
     logger.info("Scattering Green's function in the IBZ to all ranks.")
-    giwk_maxent.mat = mpi_dist.scatter(giwk_maxent.mat)  # each rank now has a slice of the irr BZ
+    g_irr_slice = mpi_dist.scatter(g_irr_mat)  # each rank now has a slice of the irr BZ
 
     wn = np.pi / config.sys.beta * (2 * np.arange(config.box.niv_core) + 1)
     w = (
@@ -109,7 +112,7 @@ def perform_maxent_giwk(giwk: GreensFunction, name: str, comm: MPI.Comm):
 
     for band in range(config.sys.n_bands):
         logger.info(f"Processing analytic continuation of band {band+1}.")
-        for k in range(giwk_maxent.mat.shape[0]):
+        for k in range(g_irr_slice.shape[0]):
             # Capture the vendored solver's stdout so its print() diagnostics go through the logger instead of
             # leaking to the output; re-logged (prefixed) below whether the continuation succeeds or fails.
             captured_output = io.StringIO()
@@ -121,7 +124,7 @@ def perform_maxent_giwk(giwk: GreensFunction, name: str, comm: MPI.Comm):
                     # The alpha-fit curve_fit inside the solver harmlessly fails to estimate its covariance; mute it.
                     warnings.simplefilter("ignore", OptimizeWarning)
                     probl_maxent = AnalyticContinuationProblem(
-                        im_axis=wn, re_axis=w, im_data=giwk_maxent[k, band, band], beta=config.sys.beta
+                        im_axis=wn, re_axis=w, im_data=g_irr_slice[k, band, band], beta=config.sys.beta
                     )
                     result = probl_maxent.solve(model=model, stdev=stdev)[0]
                     spectral_function[k, band] = result.A_opt.astype(np.float32)