Skip to content

Add Open MPI automatic distributed initialization #13929

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 12, 2023

Conversation

nvcastet
Copy link
Collaborator

@nvcastet nvcastet commented Jan 9, 2023

This PR let you initialize JAX simply with jax.distributed.initialize() when launching processes via mpirun/mpiexec.

@nvcastet nvcastet requested a review from skye January 9, 2023 17:59
Copy link
Member

@skye skye left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM overall, just a few comments + I like @nvlcambier's clarification ideas

@nvcastet nvcastet requested review from nvlcambier and skye and removed request for nvlcambier January 11, 2023 17:22
@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Jan 11, 2023
Copy link
Member

@skye skye left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry one more comment about the job_id :P

Copy link
Member

@skye skye left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants