`torchrun` and `srun`: `rank 0 != master_addr`

Jerome_Ku: <p>I’m using SLURM <code>srun</code>  to launch <code>torchrun</code> on each node with the following command:</p>
<pre><code class="lang-auto">LAUNCHER_CMD="torchrun \
--nproc_per_node $GPUS_PER_NODE \
--nnodes $NNODES \
--master_addr $MASTER_NODE \
--master_port $MASTER_PORT \
--rdzv_endpoint $MASTER_NODE:$MASTER_PORT \
--rdzv_backend c10d \
--max-restarts 0 \
$TRAIN_SCRIPT \$(&lt;$TRAIN_ARGS_OUTPUT_PATH)"
LAUNCHER_CMD="srun $SRUN_ARGS bash -c '${LAUNCHER_CMD}'"
</code></pre>
<p>All bash variables (MASTER_NODE, MASTER_PORT, etc.) are set in the script <strong>prior</strong> to <code>srun</code> execution and NOT during <code>srun</code> execution.  That is, all processes receive the same MASTER_NODE and MASTER_PORT.</p>
<p>However, if I print the environment variables set by <code>torchrun</code> in the TRAIN_SCRIPT, I get:</p>
<pre><code class="lang-auto">{'HOSTNAME': 'GPU36DC', 'MASTER_ADDR': 'GPU3B62', 'MASTER_PORT': '15472', 'LOCAL_RANK': '0', 'RANK': '0', 'WORLD_SIZE': '4', 'LOCAL_WORLD_SIZE': '2'}
</code></pre>
<p>Specifically, host <code>GPU36DC</code> is being assigned <code>rank 0</code>, even though the <code>MASTER_ADDR</code> is <code>GPU3B62</code>.</p>
<p>How to ensure that <code>MASTER_ADDR</code> is assigned global <code>rank=0</code>?  It seems the assignment is random, on some runs <code>MASTER_ADDR</code> will be assigned global rank 0, others (as above) not.</p>
<p>Here is a minimal repro:</p>
<p>The SLURM <code>sbatch</code> script:</p>
<pre data-code-wrap="bash"><code class="lang-bash">#!/bin/bash

#SBATCH --job-name=torchrun-test
#SBATCH --partition=gpu_batch
#SBATCH --nodes=2
#SBATCH --gres=gpu:1
#SBATCH --time=00:10:00
#SBATCH --ntasks-per-node=1
#SBATCH --output=sbatch-%N-%j.out
#SBATCH --error=sbatch-%N-%j.err

# Environment info
GPUS_PER_NODE=$SLURM_GPUS_ON_NODE
NNODES=$SLURM_NNODES
echo "NNODES: $NNODES"
echo "GPUS_PER_NODE: $GPUS_PER_NODE"
echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST"
scontrol show hostnames $SLURM_JOB_NODELIST
MASTER_NODE=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MASTER_NODE
MASTER_NODE_IP=$(srun --nodes=1 --ntasks=1 -w "$MASTER_NODE" hostname --ip-address)
export MASTER_NODE_IP
echo "MASTER_NODE: $MASTER_NODE MASTER_NODE_IP: $MASTER_NODE_IP"
MASTER_PORT=15472

export MASTER_PORT

TRAIN_SCRIPT=torch-dist-check.py

LAUNCHER_CMD="torchrun \
--nproc_per_node $GPUS_PER_NODE \
--nnodes $NNODES \
--master_addr $MASTER_NODE \
--master_port $MASTER_PORT \
--rdzv_endpoint $MASTER_NODE:$MASTER_PORT \
--rdzv_backend c10d \
--max-restarts 0 \
$TRAIN_SCRIPT"
LAUNCHER_CMD="srun --output srun-%N-%J.out --error srun-%N-%J.err bash -c '${LAUNCHER_CMD}'"
echo $LAUNCHER_CMD
eval $LAUNCHER_CMD
</code></pre>
<p>Here is <code>torch-dist-check.py</code> that <code>srun</code> calls:</p>
<pre data-code-wrap="python"><code class="lang-python">#!/usr/bin/env python
import os
import socket


def get_distributed_vars():
    TORCHRUN_VARS = ["MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK", "RANK", "WORLD_SIZE", "LOCAL_WORLD_SIZE"]
    assert all(var in os.environ for var in TORCHRUN_VARS)

    env_vars = {k: os.environ[k] for k in TORCHRUN_VARS}
    env_vars["HOSTNAME"] = socket.gethostname()
    return env_vars

if __name__ == "__main__":
    
    env_vars = get_distributed_vars()
    print(f"TORCHRUN_VARS:\n{env_vars}")
</code></pre>
<p>Submit the SLURM script using <code>sbatch</code> the examine the per-node <code>srun</code> output logs.</p>