From d7a9d0b09b0f19ea97a9bf8e7da3ad753f0639f8 Mon Sep 17 00:00:00 2001 From: Lorenzo Monacelli Date: Mon, 20 Jul 2026 23:47:44 +0200 Subject: [PATCH] Fix MPI broadcast 2 GB message-size limit (#36) Replace pickle-based comm.bcast() for numpy arrays with chunked raw MPI Bcast, avoiding the 2 GB MPI message-size limit that causes production failures at large tensor sizes. The broadcast function now: - Detects numpy arrays automatically and uses _broadcast_ndarray() with metadata sync (shape/dtype via pickle) + chunked raw Bcast (2M elements per chunk = 16 MB for float64, well under 2 GB limit) - Falls through to pickle-based comm.bcast() for non-array types - Removes the enforce_double = False override that made the original raw-Bcast path dead code New production-mimic regression test exercises the exact broadcast code path from the issue traceback: Tensor2.SetupFromPhonons and Center, both of which broadcast tensor/x_r_vector2/r_vector2 arrays across MPI ranks. Verified with mpirun -np 2. --- cellconstructor/Settings.py | 85 +++++-- .../TestForceTensor/test_tensor_broadcast.py | 221 ++++++++++++++++++ 2 files changed, 285 insertions(+), 21 deletions(-) create mode 100644 tests/TestForceTensor/test_tensor_broadcast.py diff --git a/cellconstructor/Settings.py b/cellconstructor/Settings.py index 7119bafb..01e4a6a3 100644 --- a/cellconstructor/Settings.py +++ b/cellconstructor/Settings.py @@ -64,40 +64,83 @@ def get_rank(): else: raise NotImplementedError("Error, I do not know what is the rank with the {} parallelization".format(__PARALLEL_TYPE__)) -def broadcast(list_of_values, enforce_double = False, other_type = None): +def broadcast(list_of_values, enforce_double=False, other_type=None): """ - Broadcast the list to all the processors from the master. - It returns a list equal for all the processors from the master. + Broadcast the data to all the processors from the master. + It returns an equal copy for all the processors from the master. - If you are broadcasting a numpy array, use enforce_double. If the array is not a C double - type, specify the other_type (must be an MPI type). + For numpy arrays, chunked raw MPI Bcast is used to avoid the 2 GB + message-size limit of pickle-based broadcast. Non-array types + (ints, lists, etc.) use the pickle-based ``comm.bcast`` path. - NOTE: Now enforce_double is just a dumb variable, as it is always overridded. - It seems that Bcast does not work as expected + Parameters + ---------- + list_of_values : any + The value to broadcast. Numpy arrays are detected + automatically and broadcast via raw Bcast in chunks. + enforce_double : bool + Ignored for numpy arrays (type is inferred automatically). + For non-array types the pickle path is always used. + other_type : MPI type + Ignored (kept for backward compatibility). """ - # STRONG OVERRIDE OVER ENFORCE DOUBLE THAT IS NOT WORKING - enforce_double = False - if __PARALLEL_TYPE__ == "mpi4py": comm = mpi4py.MPI.COMM_WORLD if comm.Get_size() == 1: return list_of_values - if not enforce_double: - return comm.bcast(list_of_values, root = 0) - else: - total_shape = list_of_values.shape - mpitype = mpi4py.MPI.DOUBLE - if other_type is not None: - mpitype = other_type - new_data = list_of_values.ravel() - comm.Bcast([new_data, np.prod(total_shape), mpitype], root = 0) - return new_data.reshape(total_shape) + if isinstance(list_of_values, np.ndarray): + return _broadcast_ndarray(comm, list_of_values) + + return comm.bcast(list_of_values, root=0) + elif __PARALLEL_TYPE__ == "serial": return list_of_values else: - raise NotImplementedError("broadcast not implemented for {} parallelization.".format(__PARALLEL_TYPE__)) + raise NotImplementedError( + "broadcast not implemented for {} parallelization.".format( + __PARALLEL_TYPE__)) + + +def _broadcast_ndarray(comm, array): + """ + Broadcast a numpy array using chunked raw MPI Bcast. + + The array metadata (shape and dtype) is first synchronized via + pickle, then the raw data is sent in fixed-size chunks to stay + well within MPI's 2 GB single-message limit. + """ + if am_i_the_master(): + meta = (array.shape, str(array.dtype)) + else: + meta = None + meta = comm.bcast(meta, root=0) + shape, dtype_str = meta + dtype = np.dtype(dtype_str) + + if not am_i_the_master(): + array = np.empty(shape, dtype=dtype) + + flat = array.ravel() + total = flat.size + + if total == 0: + return array + + CHUNK = 2 * 1024 * 1024 + + for start in range(0, total, CHUNK): + end = min(start + CHUNK, total) + n = end - start + if am_i_the_master(): + chunk = flat[start:end].copy() + else: + chunk = np.empty(n, dtype=dtype) + comm.Bcast(chunk, root=0) + flat[start:end] = chunk + + return array def barrier(): diff --git a/tests/TestForceTensor/test_tensor_broadcast.py b/tests/TestForceTensor/test_tensor_broadcast.py new file mode 100644 index 00000000..57a65f95 --- /dev/null +++ b/tests/TestForceTensor/test_tensor_broadcast.py @@ -0,0 +1,221 @@ +from __future__ import print_function +from __future__ import division + +import numpy as np + +import cellconstructor as CC +import cellconstructor.Settings +import cellconstructor.Phonons +import cellconstructor.ForceTensor +import cellconstructor.Units +import cellconstructor.symmetries + +import pytest +import sys, os + +try: + import mpi4py.MPI + _HAS_MPI = True +except ImportError: + _HAS_MPI = False + + +def _get_comm(): + if _HAS_MPI: + return mpi4py.MPI.COMM_WORLD + return None + + +def _rank(): + return CC.Settings.get_rank() + + +def _nproc(): + return CC.Settings.GetNProc() + + +def _allclose_across_ranks(arr): + """ + Return True if arr is identical on every MPI rank. + MUST be called by ALL ranks simultaneously. + """ + n = _nproc() + if n == 1: + return True + comm = _get_comm() + hashes = comm.gather(arr.ravel().tobytes()[:32], root=0) + ok = True + if _rank() == 0: + ref = hashes[0] + for h in hashes[1:]: + if ref != h: + ok = False + ok = comm.bcast(ok, root=0) + return ok + + +class TestTensorBroadcast: + + def test_broadcast_ndarray(self): + """ + Master creates a numpy array, broadcasts it via + Settings.broadcast(enforce_double=True). All ranks verify + they received identical data. + """ + shape = (27, 6, 6) + if _rank() == 0: + rng = np.random.RandomState(123) + data = rng.randn(*shape).astype(np.float64) + else: + data = np.zeros(1) + + result = CC.Settings.broadcast(data, enforce_double=True) + + assert result.shape == shape + assert result.dtype == np.float64 + assert np.max(np.abs(result)) > 0, "Broadcast produced all zeros" + + if _rank() == 0: + assert np.allclose(result, data) + + assert _allclose_across_ranks(result), \ + "Broadcast data differs across MPI ranks!" + + def test_broadcast_ndarray_medium(self): + """ + Broadcast a medium-sized array that exercises the broadcast + path with a realistic data volume. + """ + shape = (500, 100, 1) # 50K float64 elements + if _rank() == 0: + rng = np.random.RandomState(456) + data = rng.randn(*shape).astype(np.float64) + else: + data = np.zeros(1) + + result = CC.Settings.broadcast(data, enforce_double=True) + + assert result.shape == shape + assert result.dtype == np.float64 + assert np.max(np.abs(result)) > 0 + + if _rank() == 0: + assert np.allclose(result, data) + + assert _allclose_across_ranks(result), \ + "Broadcast data differs across MPI ranks!" + + def _build_valid_dynmat(self, q_list): + """ + Build dynamical matrices that satisfy D(-q)=D(q) so the + real-space force constant is purely real. + """ + nq = len(q_list) + nat3 = 3 * 2 + rng = np.random.RandomState(42) + + dm = np.zeros((nq, nat3, nat3), dtype=np.complex128) + paired = [False] * nq + + for i, qi in enumerate(q_list): + if paired[i]: + continue + a = rng.randn(nat3, nat3) + mat_real = 10.0 * 0.5 * (a + a.T) + + partner = None + for j, qj in enumerate(q_list): + if j == i: + continue + if np.linalg.norm(np.array(qi) + np.array(qj)) < 1e-10: + partner = j + break + + if partner is not None and not paired[partner]: + dm[i] = mat_real + dm[partner] = mat_real + paired[i] = paired[partner] = True + else: + dm[i] = mat_real + paired[i] = True + + return dm + + def test_tensor2_setup_and_center(self): + """ + Production-mimic test: builds a Tensor2 from a realistic Si + supercell, calls SetupFromPhonons (broadcast #1) and Center + (broadcast #2). Verifies data is consistent across all ranks. + """ + s = CC.Structure.Structure(2) + s.atoms = ["Si", "Si"] + s.coords[0, :] = [0.0, 0.0, 0.0] + s.coords[1, :] = [0.25, 0.25, 0.25] + s.unit_cell = np.eye(3) * 5.43 + s.has_unit_cell = True + s.masses = {"Si": 28.0855} + + sc = (3, 3, 3) + sc_struct = s.generate_supercell(sc) + + q_tot = CC.symmetries.GetQGrid(s.unit_cell, sc) + nq = len(q_tot) + + phon = CC.Phonons.Phonons(s, nqirr=nq) + phon.q_tot = [q.copy() for q in q_tot] + phon.q_stars = [[q.copy() for q in q_tot]] + + dm = self._build_valid_dynmat(q_tot) + for iq in range(nq): + phon.dynmats[iq] = dm[iq] + + tensor = CC.ForceTensor.Tensor2(s, sc_struct, phon.GetSupercell()) + tensor.SetupFromPhonons(phon) + + before = tensor.tensor.copy() + assert np.max(np.abs(before)) > 0 + + assert _allclose_across_ranks(tensor.tensor), \ + "Tensor differs after SetupFromPhonons" + + tensor.Center() + + assert tensor.tensor.shape == before.shape + assert np.max(np.abs(tensor.tensor)) > 0, \ + "Centered tensor is all zeros" + assert tensor.x_r_vector2.shape[1] == tensor.n_R + assert tensor.r_vector2.shape[1] == tensor.n_R + + assert _allclose_across_ranks(tensor.tensor), \ + "Tensor differs after Center" + assert _allclose_across_ranks(tensor.x_r_vector2), \ + "x_r_vector2 differs after Center" + assert _allclose_across_ranks(tensor.r_vector2), \ + "r_vector2 differs after Center" + + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--pytest", action="store_true", + help="Run via pytest instead of standalone") + parser.add_argument("-k", "--keyword", default=None, + help="pytest -k filter") + args = parser.parse_args() + + if args.pytest or args.keyword: + argv = [__file__, "-v"] + if args.keyword: + argv += ["-k", args.keyword] + pytest.main(argv) + else: + runner = TestTensorBroadcast() + r = _rank() + n = _nproc() + print("[rank {}/{}] === test_broadcast_ndarray ===".format(r, n)) + runner.test_broadcast_ndarray() + print("[rank {}/{}] === test_broadcast_ndarray_medium ===".format(r, n)) + runner.test_broadcast_ndarray_medium() + print("[rank {}/{}] === test_tensor2_setup_and_center ===".format(r, n)) + runner.test_tensor2_setup_and_center() + print("[rank {}/{}] All tests passed.".format(r, n))