internode communication with torchrun

[pchliu]

Hi,

I'm having some issues with torchrun on >1 node. The following code works fine if I ask for 1 node and 1 task, but it does not work on >1 nodes (but it works on other clusters):

#!/bin/bash
#SBATCH -N 2               
#SBATCH -n 2         
...

export NUM_NODES=${SLURM_NNODES}
export GPUS_PER_NODE=4  

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29500
export LOGLEVEL=INFO
export NCCL_DEBUG=INFO

srun torchrun \
    --nnodes      $NUM_NODES \
    --node_rank   $SLURM_PROCID \
    --nproc_per_node $GPUS_PER_NODE \
    --rdzv_backend c10d \
    --rdzv_endpoint "$MASTER_ADDR:$MASTER_PORT" \
    --rdzv_id      "$SLURM_JOB_ID" \
    train.py \
        --num_nodes $NUM_NODES \
        ....

See logs attached below from 1 node (ok, 1node_ok_slurm-190800.txt), 2 nodes (failed,
2nodes_fail_slurm-190624.txt), and 2 nodes elsewhere (works,
2nodes_othercluster_ok_slurm-6731888.txt). I think it comes down to something about the rendez-vous ports. I was hoping someone could let me know what's the solution here? Thanks!

---

Additionally, I saw on the website that

special requests that extend beyond the above queue limits may potentially be accommodated on a case-by-case basis

I was wondering how can I ask to extend beyond the queue limit? We're trying to train some big models and it'd be great if we could access longer queue limits and/or more nodes. Thanks!

---

Brief discussion of the errors below:

1. The port I pass with --rdzv_endpoint ( 29500 ) is only used for the initial rendez‑vous.
2. Rank 0 then asks the OS for any free port (in my log it picked 36609) and publishes that to the other ranks via the rendez‑vous store .
3. All workers must be able to reach t008‑004.hpcfund:36609. When they cannot, the connect() call in the C10d store times out exactly after the 900 s default, giving the trace.

I0505 22:21:50.411871 1794449 torch/distributed/run.py:649] Using nproc_per_node=4.
I0505 22:21:50.412315 1794449 torch/distributed/launcher/api.py:195] Starting elastic_operator with launch configs:
I0505 22:21:50.412315 1794449 torch/distributed/launcher/api.py:195]   entrypoint       : train.py
I0505 22:21:50.412315 1794449 torch/distributed/launcher/api.py:195]   nproc_per_node   : 4
I0505 22:21:50.412315 1794449 torch/distributed/launcher/api.py:195]   rdzv_backend     : c10d
I0505 22:21:50.412315 1794449 torch/distributed/launcher/api.py:195]   rdzv_endpoint    : t008-004:29500
I0505 22:21:50.412315 1794449 torch/distributed/launcher/api.py:195]   rdzv_configs     : {'timeout': 900}
[repeated for the second task/node]

I0505 22:21:50.436545 2283825 torch/distributed/elastic/agent/server/api.py:860] [default] starting workers for entrypoint: python
I0505 22:21:50.436951 2283825 torch/distributed/elastic/agent/server/api.py:677] [default] Rendezvous'ing worker group
I0505 22:21:51.593055 1794449 torch/distributed/elastic/agent/server/api.py:525] [default] Rendezvous complete for workers. Result:
I0505 22:21:51.593055 1794449 torch/distributed/elastic/agent/server/api.py:525]   master_addr=t008-004.hpcfund
I0505 22:21:51.593055 1794449 torch/distributed/elastic/agent/server/api.py:525]   master_port=36609
I0505 22:21:51.593055 1794449 torch/distributed/elastic/agent/server/api.py:525]   group_rank=0
[repeated for the second task/node]
...

[W505 22:22:02.546732811 socket.cpp:755] [c10d] The IPv6 network addresses of (t008-004.hpcfund, 36609) cannot be retrieved (gai error: -2 - Name or service not known).
[repeated many times]
...

[E505 22:30:57.854404041 socket.cpp:1019] [c10d] The client socket has timed out after 600000ms while trying to connect to (t008-004.hpcfund, 36609).
Exception raised from throwTimeoutError at /pytorch/torch/csrc/distributed/c10d/socket.cpp:1021 (most recent call first):
...

[tom-papatheodore]

Hi @pchliu. This appears to be a known issue with PyTorch Lightning (https://hychiang.info/blog/2024/pytorch_distributed_training/#hostname-not-known-issue), but there appear to be potential workarounds.

@koomie Can you please take a look at the linked workaround and let me and Chenghao know if the workaround is something we could consider on our cluster?

@pchliu Would it be possible to use something other than PyTorch Lightning in the meantime so you can make progress?
NOTE: I'll be out on PTO from today until June 9.

[tom-papatheodore]

I should have mentioned that I was able to reproduce this issue on the cluster myself and I wanted to share in case it's helpful. To reproduce...

Create a Python virtual environment:

## Move to WORK directory, create and activate virtual environment
$ cd $WORK/<somewhere-in-work-dir>
$ python -m venv –-upgrade-deps <name-of-venv>
$ source <name-of-venv>/bin/activate

## Install ROCm-enabled PyTorch
$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.3
$ pip3 install importlib-metadata
$ pip3 install lightning

Example training script:

$ cat train.py
"""
train.py — PyTorch Lightning distributed fine-tune example
Target  : 2 nodes × 4 AMD MI210 GPUs each  (ROCm 6.3  +  PyTorch 2.7)
Launch  :
    export MASTER_ADDR=$(srun --nodes=1 --ntasks=1 hostname)
    export MASTER_PORT=29500
    srun -N2 -n2 torchrun --nnodes=2 --nproc_per_node=4 \
         --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
         train.py
"""

import os, pytorch_lightning as pl, torch, torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import models, transforms
from torchvision.datasets import CIFAR10
from typing import Optional

# ───────────────────────────────────────────────────────────────
# 1. LightningModule – fine-tune ResNet-18 on CIFAR-10
# ───────────────────────────────────────────────────────────────
class FineTuneModel(pl.LightningModule):
    def __init__(self, num_classes: int = 10, lr: float = 1e-3):
        super().__init__()
        self.save_hyperparameters()
        self.model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):                       # inference / forward pass
        return self.model(x)

    def training_step(self, batch, _):
        x, y = batch
        loss = self.criterion(self(x), y)
        self.log("train_loss", loss, prog_bar=True, sync_dist=True)
        return loss

    def validation_step(self, batch, _):
        x, y = batch
        loss = self.criterion(self(x), y)
        self.log("val_loss", loss, prog_bar=True, sync_dist=True)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)

# ───────────────────────────────────────────────────────────────
# 2. LightningDataModule – Lightning handles one-time download
# ───────────────────────────────────────────────────────────────
class CIFAR10Data(pl.LightningDataModule):
    def __init__(self,
                 root: str = "./data",
                 batch_size: int = 256,
                 num_workers: int = 8):
        super().__init__()
        self.root, self.batch_size, self.num_workers = root, batch_size, num_workers
        self.tfm = {
            "train": transforms.Compose([
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
            ]),
            "val": transforms.Compose([transforms.ToTensor()])
        }

    # runs **once** on rank-0 before DDP init
    def prepare_data(self):
        CIFAR10(self.root, train=True,  download=True)
        CIFAR10(self.root, train=False, download=True)

    # runs on every rank (each GPU) after DDP init
    def setup(self, stage: Optional[str] = None):
        self.train_set = CIFAR10(self.root, train=True,  download=False,
                                 transform=self.tfm["train"])
        self.val_set   = CIFAR10(self.root, train=False, download=False,
                                 transform=self.tfm["val"])

    def train_dataloader(self):
        return DataLoader(self.train_set, batch_size=self.batch_size,
                          shuffle=True, num_workers=self.num_workers,
                          pin_memory=True)

    def val_dataloader(self):
        return DataLoader(self.val_set,   batch_size=self.batch_size,
                          shuffle=False, num_workers=self.num_workers,
                          pin_memory=True)

# ───────────────────────────────────────────────────────────────
# 3. Main
# ───────────────────────────────────────────────────────────────
def main():
    # dataset root (absolute path is safest for multi-node runs)
    data_root = os.environ.get("CIFAR10_ROOT",
                               "/work1/staff/papatheodore/ai-hpc-cluster/"
                               "user-support/github-issues/issue-0047/data")

    datamodule = CIFAR10Data(root=data_root)

    trainer = pl.Trainer(
        accelerator="gpu",
        strategy="ddp",          # Lightning handles RCCL (NCCL backend) on ROCm
        devices=4,               # GPUs *per node*
        num_nodes=2,
        max_epochs=5,
        log_every_n_steps=25,
    )

    trainer.fit(FineTuneModel(), datamodule=datamodule)


if __name__ == "__main__":
    # optional: give each rank a few OpenMP threads (else torchrun sets OMP_NUM_THREADS=1)
    os.environ.setdefault("OMP_NUM_THREADS", "8")  # 8 threads × 4 ranks = 32 cores
    main()

Batch script:

$ cat submit_torchrun.sh
#!/bin/bash
#SBATCH -N 2
#SBATCH -n 2
#SBATCH --ntasks-per-node=1    # one torchrun launcher per node
#SBATCH --cpus-per-task=8      # e.g. 16 cores for 4 workers (4 cores/worker)
#SBATCH -t 02:00:00
#SBATCH -p mi2104x

source <path-to-venv>/venv/bin/activate

export MASTER_ADDR=$(srun --nodes=1 --ntasks=1 hostname -s)
export MASTER_PORT=29500

# Use torchrun to launch 8 processes on each node (16 total).
# --nnodes and --nproc_per_node specify the world size.
srun -N 2 -n 2 torchrun --nnodes=2 --nproc_per_node=4 \
  --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  train.py

[pchliu]

Hi @koomie and @tom-papatheodore , has anyone looked at the workarounds? Thank you in advance.

[Schwidola0607]

Hi, I am also running into this issue with torchrun / deepspeed. Below is simple reproduce-able script.

Training script:

#!/usr/bin/env python3
"""
Minimal DDP script.
* One Linear layer, synthetic data
* Prints rank-info and first loss value
* Runs 5 iterations then exits
"""

import os
import time
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def main():
    local_rank  = int(os.environ["LOCAL_RANK"])
    world_size  = int(os.environ["WORLD_SIZE"])
    global_rank = int(os.environ["RANK"])

    torch.cuda.set_device(local_rank)
    dist.init_process_group("nccl")

    if global_rank == 0:
        print(f"\n[+] world_size={world_size}, MASTER_ADDR={os.getenv('MASTER_ADDR')},"
              f" MASTER_PORT={os.getenv('MASTER_PORT')}\n")

    model = torch.nn.Linear(16, 16).cuda()
    model = DDP(model, device_ids=[local_rank])

    # 16 × 16 random data, same on every rank
    x = torch.randn(32, 16).cuda()
    y = torch.randn(32, 16).cuda()
    opt = torch.optim.SGD(model.parameters(), lr=1e-1)

    for step in range(5):
        opt.zero_grad()
        loss = torch.nn.functional.mse_loss(model(x), y)
        loss.backward()
        opt.step()
        
        if step == 0:
            print(f"[rank {global_rank}/{world_size}] first loss = {loss.item():.4f}")

    if global_rank == 0:
        print("\n[✓] Training finished — shutting down\n")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

Batch script:

#!/bin/bash
#SBATCH -J dummy-ddp
#SBATCH -o dummy_train.out
#SBATCH -e dummy_train.err
#SBATCH -N 2
#SBATCH -t 00:10:00
#SBATCH -p mi2508x

DATETIME=$(date +"%Y-%m-%d_%H-%M-%S")
echo $DATETIME

module load pytorch
ifconfig
export NCCL_SOCKET_IFNAME=eth0
export GLOO_SOCKET_IFNAME=eth0
export ROCM_HOME=/opt/rocm-6.3.1

NNODES=$SLURM_NNODES
GPUS=4
NODE_RANK=$SLURM_NODEID
MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n1)
MASTER_PORT=29500

srun --nodes=${NNODES} \
     --ntasks-per-node=1 \
     --exclusive \
     --export=ALL,MASTER_ADDR,MASTER_PORT \
     bash -c "
       torchrun \
         --nnodes=${NNODES} \
         --nproc_per_node=${GPUS} \
         --node_rank=${NODE_RANK} \
         --rdzv_backend=c10d \
         --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
         ./dummy_train.py
     "

[Schwidola0607]

@tom-papatheodore @pchliu @koomie
Hi I got internode communication with torchrun to work by avoid using --rdzv_endpoint and instead speciying --master_addr and --master_port directly. For my example above, the bash script become

#!/bin/bash
#SBATCH -J dummy-ddp
#SBATCH -o dummy_train.out
#SBATCH -e dummy_train.err
#SBATCH -N 2
#SBATCH -t 00:10:00
#SBATCH -p mi2104x

DATETIME=$(date +"%Y-%m-%d_%H-%M-%S")
echo $DATETIME

module load pytorch
export ROCM_HOME=/opt/rocm-6.3.1

srun --nodes=$SLURM_NNODES \
     --ntasks-per-node=1 \
     --exclusive \
     --export=ALL \
     bash -c '
       NODE_RANK=$SLURM_NODEID            
       MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n1)
       MASTER_PORT=29500

       torchrun \
         --nnodes=$SLURM_NNODES \
         --nproc_per_node=4 \
         --node_rank=$NODE_RANK \
         --master_addr=$MASTER_ADDR \
         --master_port=$MASTER_PORT \
         ./dummy_train.py
     '

Edit: The issue was that NODE_RANK was previously retrieved from $SLURM_NODEID at login node which is always 0. Moving it inside the bash script being executed helps.

[tom-papatheodore]

Hi @Schwidola0607. I cannot reproduce your results. When I use your exact training script and latest Slurm batch script, I appear to just hang until the job times out. If I SSH directly to the nodes I have access to and run rocm-smi, I show no GPU activity.

[Schwidola0607]

Hi @tom-papatheodore, sorry for the late reply. I edited the latest script. Please take a look and see if its reproducible for you. Thanks!