A template repository for submitting a job to the Slurm Cluster installed at the DISI - University of Bologna

Sto provando a inizializzare un gruppo con torch distributed, ma non riesco a venirne a capo.

Prendendo spunto da qui, ho scritto questo codice minimale:

import torch
import os
import argparse


def main():
    parser = argparse.ArgumentParser("Test")
    args = parser.parse_args()

    args.ngpus_per_node = torch.cuda.device_count()
    args.rank = int(os.getenv('SLURM_NODEID')) * args.ngpus_per_node
    args.world_size = int(os.getenv('SLURM_NNODES')) * args.ngpus_per_node
    args.dist_url = f'env://'
    torch.multiprocessing.spawn(do_something, (args,), args.ngpus_per_node)  

def do_something(device, args):
    args.rank += device
    torch.distributed.init_process_group(
        backend='nccl', init_method=args.dist_url,
        world_size=args.world_size, rank=args.rank
)  
    
    local_tensor = torch.rand(10, device=device)
    torch.distributed.all_reduce(local_tensor)

if __name__ == "__main__":
    main()

Usando questo sbatch:

#!/bin/bash
#SBATCH --job-name=tests
#SBATCH --mail-type=ALL
#SBATCH --mail-user=[myemail]
#SBATCH --time=05:00:00
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=1
#SBATCH --output=output.txt
#SBATCH --gres=gpu:1

source cvenv/bin/activate

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

python3 main.py

Il codice funziona se lo lancio su un singolo nodo, altrimenti l'inizializzazione di torch distributed rimane bloccata e ricevo questo errore dopo il time out:

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/public.hpc/oyounis/powersgd/cvenv/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/public.hpc/oyounis/powersgd/main.py", line 21, in do_something
    world_size=args.world_size, rank=args.rank
  File "/public.hpc/oyounis/powersgd/cvenv/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 627, in init_process_group
    _store_based_barrier(rank, store, timeout)
  File "/public.hpc/oyounis/powersgd/cvenv/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 258, in _store_based_barrier
    rank, store_key, world_size, worker_count, timeout
RuntimeError: Timed out initializing process group in store based barrier on rank: 0, for key: store_based_barrier_key:1 (world_size=4, worker_count=1, timeout=0:30:00)

Penso sia un problema sul MASTER_ADDR (che con questo codice viene impostato al nome del primo nodo, per esempio ws2gpu2), ma non saprei in che altro modo settarlo; ho provato localhost ma non funziona su più nodi.

Buongiorno, sono uno studente della magistrale di AI, e dal momento che molti miei colleghi sono abituati ad eseguire task di ML/DL usando notebook Jupyter, ho pensato potesse essere utile aggiungere le istruzioni per eseguire un server Jupyter sul cluster e connettersi da remoto.