PyTorch and OpenMP/MPI-enabled AMReX don't get along in load_state_dict

[RTSandberg]

On my local machine, PyTorch has some internal multithreaded functionality that doesn't get along with AMReX. Unless I set PyTorch.set_num_threads(1 or 2), then the attached script will hang when the neural network tries to set its initial parameters.

This script downloads some neural network parameters from Zenodo archive to then load them, and the load_state_dict function is the specific point of failure.

pytorch_amrex_hang_reproducer_v2.py.txt

[ax3l]

Thank you, @RTSandberg !

For reproducibility, can you please add the OS you used, versions of Python, pyAMReX, PyTorch, MPI flavor and version, and mpi4py version?

[ax3l]

If we can reduce this problem to a pure mpi4py + PyTorch issue, then we could also report this upstream in PyTorch: https://github.com/pytorch/pytorch/issues

[ax3l]

The issue seems to be more general, also showing sometimes later, after model load: ECP-WarpX/impactx#773 (comment)

Could also be a mixing of OpenMP libraries (gomp and llvm omp)... Especially when downloading PyTorch from one source (say: Pip) and building the rest of the stack from another (system, conda, etc.).

One should check that using the same dependencies & compilers for ImpactX/AMReX and PyTorch does show the same issue, e.g., using exclusively Spack or Conda-Forge for everything.