Fix MPI broadcast 2 GB message-size limit (#36)#124
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #36
Problem
Settings.broadcast()usescomm.bcast()(pickle-based MPI broadcast) for all data types, including numpy arrays. This has a 2 GB MPI message-size limit, causing cryptic failures in production when broadcasting large tensor arrays. The original raw Bcast code path was dead (enforce_doublealways overridden toFalseat line 80).Fix
comm.Bcast:enforce_doubleoverride is removed; parameter kept for backward compatibilityBenchmark: pickle vs chunked Bcast
Measured with
mpirun -np 2across varying array sizes (5 repeats each):Chunked Bcast is consistently faster (~1.6x at scale) and has no 2 GB limit.
New test
tests/TestForceTensor/test_tensor_broadcast.py— production-mimic regression test that exercises the exact code path from the issue traceback (Tensor2.SetupFromPhononsthenCenter), both of which broadcast tensor data across MPI ranks. Verified withmpirun -np 2.Regression verification
broadcast()call sites in ForceTensor.py (50+): pass