Enable opt-in autodetection of distributed configuration for mpi4py #20174

coreyjadams · 2024-03-11T17:00:29Z

This PR is in response to the discussion on #19409.

It does the following:

First, it adds an additional cluster environment for jax.distributed that is based on autodetection of rank, size, local_rank/devices, and coordination address based on mpi4py. The motivation here is that for clusters with mpi4py, the parameters needed for jax.distributed.initialize can be entirely inferred from mpi4py provided the job is launched in a way compatible with MPI.
Because of the constraint above, compatibility with MPI, this autodetect method is exclusively opt-in. Users must pass spec_detection_method="mpi4py" in a call to jax.distributed.initialize.
Consistent with the behavior of all other initialization methods, options passed to jax.distributed.initialize for coordinator_address, etc., will override and auto-detected settings from mpi4py.
Because of the new option in the arguments to jax.distributed.initialize, I have updated the documentation accordingly.
Lastly, related to using jax.distributed.initialize on HPC systems, there is sometimes a hang (See Usage of `jax.distributed.initialize` on HPC clusters #9582). I have included a warning if any of the suspect variables are detected, and updated the documentation to point out that the user may want to unset these variables if they are on an HPC cluster. I suspect the warning will be viewed as to noisy, since I don't know the default log level in JAX off the top of my head. In this case, perhaps it can only be emitted if the TimeOut occurs, or at the very least maintained in the documentation.

I hope this is helpful, and it would simplify our lives using JAX on supercomputers :).

Corey

google-cla · 2024-03-11T17:00:35Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

nvcastet

Sorry for the late review.
Thanks a lot for the contribution. I left some comments.

nvcastet · 2024-03-27T17:19:19Z

jax/_src/clusters/cluster.py

+
+    # First, we check the spec detection method because it will ignore submitted values
+    # If if succeeds.
+    if spec_detection_method == "mpi4py":


Could we add a "name" attribute under ClusterEnv and set it in the derived classes: aka ompi, slurm, mpi4py, gketpu, multislicegcetpu, singleslicegcetpu and also adding a new bool "opt_in_only_method" and have here:

if spec_detection_method is not None: env = next((env for env in cls._cluster_types if env.name == spec_detection_method), None) if env is None: # print fatal error that method {spec_detection_method} is not supported. if not env.is_env_present(): # print fatal error requested env is not present. else: env = next((env for env in cls._cluster_types if env.opt_in_only_method == False and env.is_env_present()), None)

Sure, good idea, this makes it more flexible to add other opt-in methods in the future. I've implemented this with the names you suggested, every type except mpi4py has opt_in_only_method = False as a class attribute.

nvcastet · 2024-03-27T17:24:24Z

jax/_src/clusters/mpi4py_cluster.py

+    from mpi4py import MPI
+    COMM_WORLD = MPI.COMM_WORLD
+
+    hostname = socket.gethostname()


Could you use COMM_WORLD.Split_type(MPI.COMM_TYPE_SHARED) instead of doing the hostname query?

In principle, that seems possible. I am trying to imagine scenarios on some clusters where processors that can access the same shared memory space might be over or under restrictive compared to matching against hostname...but at worst it seems like it might lead to NUMA-induced slowdowns but not anything wrong.

For example, on one system I am thinking of, the hosts will have two CPUs with two sockets each, attached to 6 GPUs and each CPU has its own set of shared memory. Using shared memory accesses here might be ambiguous, depending on the job launcher, and it's not a COMM-split method I see used much.

If you prefer COMM_TYPE_SHARED I won't object, though I'm not sure what the downside of a host query is.

Even in multi-NUMA systems, it works fine since the processes can still communicate via shared-memory.
I am pretty confident that MPI.COMM_TYPE_SHARED works well for this use case.
We have been using it in Horovod (distributed layer that can be used on top of PyTorch/TF/MXNet) for years with a wide variety of users:
https://github.com/horovod/horovod/blob/master/horovod/common/mpi/mpi_context.cc#L126

I would prefer the split over implementing it ourself by broadcasting the hostname since MPI vendors may be able to optimize this specific comm split under the cover.

jax/_src/distributed.py

jax/_src/clusters/mpi4py_cluster.py

jax/_src/distributed.py

nvcastet

Thanks Corey for following up.
I left a few comments.
Do you mind re-testing the final code on your cluster and adding an unit test similar to the ompi one https://github.com/google/jax/blob/main/tests/multiprocess_gpu_test.py#L184 ?

nvcastet · 2024-04-23T21:30:35Z

jax/_src/clusters/cluster.py

@@ -43,7 +44,11 @@ def auto_detect_unset_distributed_params(cls,
                                           num_processes: int | None,
                                           process_id: int | None,
                                           local_device_ids: Sequence[int] | None,
+<<<<<<< HEAD


some merge conflict left over here?

nvcastet · 2024-04-23T21:33:05Z

jax/_src/distributed.py

@@ -215,7 +225,11 @@ def initialize(coordinator_address: str | None = None,
    raise RuntimeError("jax.distributed.initialize() must be called before "
                        "any JAX computations are executed.")
  global_state.initialize(coordinator_address, num_processes, process_id,
+<<<<<<< HEAD


nvcastet · 2024-04-23T21:46:01Z

jax/_src/clusters/mpi4py_cluster.py

+    from mpi4py import MPI
+    COMM_WORLD = MPI.COMM_WORLD
+
+    hostname = socket.gethostname()


Even in multi-NUMA systems, it works fine since the processes can still communicate via shared-memory.
I am pretty confident that MPI.COMM_TYPE_SHARED works well for this use case.
We have been using it in Horovod (distributed layer that can be used on top of PyTorch/TF/MXNet) for years with a wide variety of users:
https://github.com/horovod/horovod/blob/master/horovod/common/mpi/mpi_context.cc#L126

I would prefer the split over implementing it ourself by broadcasting the hostname since MPI vendors may be able to optimize this specific comm split under the cover.

coreyjadams · 2024-05-28T16:38:00Z

Hey, I'm back, sorry for the delay. It's been a busy time here - I'm about to transition jobs and I'm wrapping up some open things.

I've made the changes requested, and added the test. I'll point out that on my system, mpirun does not allow the following arguments:

          '--oversubscribe',
          '--allow-run-as-root',

So I ran the test without these and it passes. I left them in the final code, however - I'm not sure what test system you have that allows launching mpirun jobs as root but it sure isn't a DOE super computer :).

Let me know if you want any more changes. I'll be able to test things for about 1 month more.

nvcastet

@skye @hawkinsp Do you mind giving another look at this PR?

jax/_src/clusters/mpi4py_cluster.py

jax/_src/distributed.py

jax/_src/clusters/mpi4py_cluster.py

coreyjadams · 2024-06-27T17:12:40Z

Requested changes to comments and documentation is done - thanks for the feedback :)

nvcastet · 2024-06-27T17:51:07Z

@coreyjadams You would need to sign the CLA too: https://cla.developers.google.com/

nvcastet · 2024-06-27T18:01:21Z

@coreyjadams could you update the year at the top of your new file and also add it to https://github.com/google/jax/blob/98b87540a76bc3b4f9f414bb763304f1cc474544/jax/BUILD#L953 so that it gets properly packaged.

coreyjadams · 2024-06-27T18:17:06Z

I did the BUILD file and I'm working on the cla. It's signed (Since April!) but looks like my commits were coming in on a different email address. Working on it...

coreyjadams · 2024-06-27T18:18:22Z

I did the BUILD file and I'm working on the cla. It's signed (Since April!) but looks like my commits were coming in on a different email address. Working on it...

OK, I see the green check mark now, just had to link my work email to github. Thanks!

nouiz

I didn't look at the code, but this looks useful.
Can you update the title to help understand more what it does? SOmething line:
Enable opt-in autodetection of distributed configuration for mpi4py

hawkinsp

This mostly looks good to me: one api naming suggestion.

jax/_src/distributed.py

hawkinsp · 2024-06-27T21:24:27Z

Please fix the lint errors and squash your commits.

coreyjadams · 2024-07-01T19:21:26Z

The lint errors are fixed locally - just waiting on the CI to confirm then I'll deal with squashing the commits. Thanks for the feedback.

coreyjadams · 2024-07-02T16:32:31Z

I made an inadvertent force push here, so I created a branch on the commit just before that: #22235 . Suggestion is if you can force the squash, please do so? Then we just close this.

Appreciate working with you on this and sorry my git skills aren't up to par on this one :).

Update: @felker below was able to salvage the git snafu I created, I think. Should be ok to merge?

commit 79b8cbf Author: Corey Adams <corey.adams@anl.gov> Date: Mon Jul 1 14:14:15 2024 -0500 Fix mypy issues; change variable name to more universally known name commit 10edc86 Author: Corey Adams <corey.adams@anl.gov> Date: Thu Jun 27 13:25:32 2024 -0500 Change copyright year to the year this was authored commit f7086cb Author: Corey Adams <corey.adams@anl.gov> Date: Thu Jun 27 13:15:32 2024 -0500 Update build file to include mpi4py cluster. commit 6235eb3 Author: Corey adams <coreyjadams@gmail.com> Date: Thu Jun 27 12:11:48 2024 -0500 Update distributed.py Clean up documentation slightly. commit ef3a2e2 Author: Corey adams <coreyjadams@gmail.com> Date: Thu Jun 27 12:09:37 2024 -0500 Update mpi4py_cluster.py Further clean up unneeded comments. commit 6cc07a9 Author: Corey adams <coreyjadams@gmail.com> Date: Thu Jun 27 12:08:38 2024 -0500 Update mpi4py_cluster.py Remove unneeded commented code. commit 6701bd1 Merge: 5a91ac3 98b8754 Author: Corey adams <coreyjadams@gmail.com> Date: Thu Jun 27 12:07:25 2024 -0500 Merge branch 'google:main' into main commit 5a91ac3 Merge: 301bbc6 6c51234 Author: Corey adams <coreyjadams@gmail.com> Date: Tue May 28 22:14:08 2024 -0500 Merge branch 'google:main' into main commit 301bbc6 Author: Corey Adams <corey.adams@anl.gov> Date: Tue May 28 11:34:51 2024 -0500 Add test to verify mpi4py based distributed initialization commit 19e6694 Author: Corey Adams <corey.adams@anl.gov> Date: Tue May 28 11:14:40 2024 -0500 Unify variable naming and fix function argument ordering commit 72fe093 Author: Corey Adams <corey.adams@anl.gov> Date: Tue May 28 10:56:25 2024 -0500 Remove unmerged code commit 3a96e73 Merge: e4fd97e ff3db9b Author: Corey adams <coreyjadams@gmail.com> Date: Tue May 28 10:51:41 2024 -0500 Merge branch 'google:main' into main commit e4fd97e Merge: a697299 72a81e5 Author: Corey adams <coreyjadams@gmail.com> Date: Mon May 13 16:01:35 2024 -0500 Merge branch 'google:main' into main commit a697299 Merge: 85bcf42 1e48adc Author: Corey adams <coreyjadams@gmail.com> Date: Mon May 13 14:21:32 2024 -0500 Merge branch 'google:main' into main commit 85bcf42 Merge: af1a4f0 06cd05d Author: Corey Adams <corey.adams@anl.gov> Date: Tue Apr 16 09:09:31 2024 -0500 Merge branch 'main' of https://github.com/google/jax commit af1a4f0 Author: Corey Adams <corey.adams@anl.gov> Date: Tue Apr 16 08:58:33 2024 -0500 update documentation and elaborate on spec_detect_method variable commit 01f4709 Author: Corey Adams <corey.adams@anl.gov> Date: Tue Apr 16 08:45:38 2024 -0500 Address feedback and comments on PR 20174; fix typo in documentation. commit 4f22d86 Merge: 900a037 71ec6e3 Author: Corey adams <coreyjadams@gmail.com> Date: Mon Mar 11 11:51:30 2024 -0500 Merge branch 'google:main' into main commit 900a037 Author: Corey Adams <corey.adams@anl.gov> Date: Mon Mar 11 11:50:48 2024 -0500 Auto-detect of mpi4py-based configuration is now strictly opt-in. commit 1992969 Author: Corey Adams <corey.adams@anl.gov> Date: Thu Mar 7 12:27:43 2024 -0600 Enable automatic detection of distrbuted variables with any configuration of MPI, as long as mpi4py is available

coreyjadams mentioned this pull request Mar 11, 2024

Proposal for generic mpi4py initialization of jax distributed module #19409

Closed

mjsML requested review from nouiz and nvcastet March 11, 2024 18:20

nvcastet reviewed Mar 27, 2024

View reviewed changes

nvcastet reviewed Apr 23, 2024

View reviewed changes

nvcastet reviewed Jun 10, 2024

View reviewed changes

jax/_src/clusters/mpi4py_cluster.py Outdated Show resolved Hide resolved

jax/_src/distributed.py Outdated Show resolved Hide resolved

jax/_src/clusters/mpi4py_cluster.py Outdated Show resolved Hide resolved

nvcastet approved these changes Jun 27, 2024

View reviewed changes

google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Jun 27, 2024

kokoro-team removed the kokoro:force-run label Jun 27, 2024

nvcastet added pull ready Ready for copybara import and testing and removed pull ready Ready for copybara import and testing labels Jun 27, 2024

nouiz reviewed Jun 27, 2024

View reviewed changes

hawkinsp requested changes Jun 27, 2024

View reviewed changes

jax/_src/distributed.py Outdated Show resolved Hide resolved

coreyjadams changed the title ~~Enable opt-in autodetection of distributed configuration~~ Enable opt-in autodetection of distributed configuration for mpi4py Jul 1, 2024

coreyjadams force-pushed the main branch from 9b1c28d to 569f92c Compare July 2, 2024 15:57

coreyjadams mentioned this pull request Jul 2, 2024

Enable opt-in autodetection of distributed configuration for mpi4py, attempt 2 #22235

Draft

felker force-pushed the main branch from 569f92c to ffc9292 Compare July 2, 2024 18:19

hawkinsp approved these changes Jul 3, 2024

View reviewed changes

google-ml-butler bot added the kokoro:force-run label Jul 3, 2024

kokoro-team removed the kokoro:force-run label Jul 3, 2024

copybara-service bot merged commit 0d57c72 into jax-ml:main Jul 8, 2024
15 of 16 checks passed

nvcastet mentioned this pull request Aug 27, 2024

Make automatic distributed initialization work with Open MPI 5.x, PMIx, and PRRTE #14576

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable opt-in autodetection of distributed configuration for mpi4py #20174

Enable opt-in autodetection of distributed configuration for mpi4py #20174

coreyjadams commented Mar 11, 2024

google-cla bot commented Mar 11, 2024

nvcastet left a comment

nvcastet Mar 27, 2024

coreyjadams Apr 16, 2024

nvcastet Mar 27, 2024

coreyjadams Apr 16, 2024

nvcastet Apr 23, 2024

nvcastet left a comment

nvcastet Apr 23, 2024

nvcastet Apr 23, 2024

nvcastet Apr 23, 2024

coreyjadams commented May 28, 2024

nvcastet left a comment

coreyjadams commented Jun 27, 2024

nvcastet commented Jun 27, 2024

nvcastet commented Jun 27, 2024

coreyjadams commented Jun 27, 2024

coreyjadams commented Jun 27, 2024

nouiz left a comment

hawkinsp left a comment

hawkinsp commented Jun 27, 2024

coreyjadams commented Jul 1, 2024

coreyjadams commented Jul 2, 2024 •

edited

Loading

Enable opt-in autodetection of distributed configuration for mpi4py #20174

Enable opt-in autodetection of distributed configuration for mpi4py #20174

Conversation

coreyjadams commented Mar 11, 2024

google-cla bot commented Mar 11, 2024

nvcastet left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nvcastet left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

coreyjadams commented May 28, 2024

nvcastet left a comment

Choose a reason for hiding this comment

coreyjadams commented Jun 27, 2024

nvcastet commented Jun 27, 2024

nvcastet commented Jun 27, 2024

coreyjadams commented Jun 27, 2024

coreyjadams commented Jun 27, 2024

nouiz left a comment

Choose a reason for hiding this comment

hawkinsp left a comment

Choose a reason for hiding this comment

hawkinsp commented Jun 27, 2024

coreyjadams commented Jul 1, 2024

coreyjadams commented Jul 2, 2024 • edited Loading

coreyjadams commented Jul 2, 2024 •

edited

Loading