Skip to content

Account for mpirun use case in get_rank #8429

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.

Sign up for GitHub

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
janekl merged 1 commit into from
Feb 22, 2024
Merged

Account for mpirun use case in get_rank #8429

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions nemo/utils/get_rank.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,14 @@ def is_global_rank_zero():
if slurm_rank is not None:
return slurm_rank == 0

# if neither pytorch and SLURM env vars are set
# Try to get the MPI global rank env var
mpi_rank = get_envint("OMPI_COMM_WORLD_RANK", None)
if mpi_rank is not None:
return mpi_rank == 0

# if neither pytorch, SLURM nor MPI env vars are set
# check NODE_RANK/GROUP_RANK and LOCAL_RANK env vars
# asume global_rank is zero if undefined
# assume global_rank is zero if undefined
node_rank = get_envint("NODE_RANK", get_envint("GROUP_RANK", 0))
local_rank = get_envint("LOCAL_RANK", 0)
return node_rank == 0 and local_rank == 0
Expand Down