usrbinenv python this a torch distributed diagnostics script that checks that all gpus in the cluster one or many nodes can talk to each other via nccl and allocate gpu memory to run first adjust the number of processes and nodes python m torch distributed run nprocpernode 2 nnodes 1 torchdistributedgputest py you may need to add masteraddr masteraddr masterport masterport if using a custom addr port you can also use the rdzv api rdzvendpoint masteraddr masterport rdzvbackend c10d use torch distributed launch instead of torch distributed run for torch 1 9 if you get a hanging in barrier calls you have some network issues you may try to debug this with nccldebuginfo python m torch distributed run nprocpernode 2 nnodes 1 torchdistributedgputest py which should tell you what s going on behind the scenes this script can be run via srun in the slurm environment as well here is a slurm script that runs on 2 nodes of 4 gpus per node sbatch jobnametestnodes name sbatch nodes2 nodes sbatch ntaskspernode1 crucial only 1 task per dist per node sbatch cpuspertask10 number of cores per tasks sbatch gresgpu 4 number of gpus sbatch time 0 05 00 maximum execution time hh mm ss sbatch outputxj out output file name gpuspernode4 masteraddrscontrol show hostnames slurmjobnodelist head n 1 masterport6000 srun jobid slurmjobid bash c python m torch distributed run nprocpernode gpuspernode nnodes slurmnnodes noderank slurmprocid masteraddr masteraddr masterport masterport torchdistributedgputest py solves multiprocess interleaved print problem with openfile r as fh fcntl flockfh fcntl lockex try printmsgs finally fcntl flockfh fcntl lockun localrank intos environlocalrank torch cuda setdevicelocalrank device torch devicecuda localrank hostname socket gethostname gpu fhostnamelocalrank try test distributed dist initprocessgroupnccl dist allreducetorch ones1 todevice opdist reduceop sum dist barrier test cuda is available and can allocate memory torch cuda isavailable torch ones1 cudalocalrank global rank rank dist getrank worldsize dist getworldsize printflockfgpu is ok global rank rankworldsize dist barrier if rank 0 printflockfpttorch version cudatorch version cuda nccltorch cuda nccl version except exception printflockfgpu is broken raise usr bin env python this a torch distributed diagnostics script that checks that all gpus in the cluster one or many nodes can talk to each other via nccl and allocate gpu memory to run first adjust the number of processes and nodes python m torch distributed run nproc_per_node 2 nnodes 1 torch distributed gpu test py you may need to add master_addr master_addr master_port master_port if using a custom addr port you can also use the rdzv api rdzv_endpoint master_addr master_port rdzv_backend c10d use torch distributed launch instead of torch distributed run for torch 1 9 if you get a hanging in barrier calls you have some network issues you may try to debug this with nccl_debug info python m torch distributed run nproc_per_node 2 nnodes 1 torch distributed gpu test py which should tell you what s going on behind the scenes this script can be run via srun in the slurm environment as well here is a slurm script that runs on 2 nodes of 4 gpus per node sbatch job name test nodes name sbatch nodes 2 nodes sbatch ntasks per node 1 crucial only 1 task per dist per node sbatch cpus per task 10 number of cores per tasks sbatch gres gpu 4 number of gpus sbatch time 0 05 00 maximum execution time hh mm ss sbatch output x j out output file name gpus_per_node 4 master_addr scontrol show hostnames slurm_job_nodelist head n 1 master_port 6000 srun jobid slurm_jobid bash c python m torch distributed run nproc_per_node gpus_per_node nnodes slurm_nnodes node_rank slurm_procid master_addr master_addr master_port master_port torch distributed gpu test py solves multi process interleaved print problem test distributed test cuda is available and can allocate memory global rank
import fcntl
import os
import socket

import torch
import torch.distributed as dist


def printflock(*msgs):
    
    with open(__file__, "r") as fh:
        fcntl.flock(fh, fcntl.LOCK_EX)
        try:
            print(*msgs)
        finally:
            fcntl.flock(fh, fcntl.LOCK_UN)


local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
hostname = socket.gethostname()

gpu = f"[{hostname}-{local_rank}]"

try:
    
    dist.init_process_group("nccl")
    dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
    dist.barrier()

    
    torch.cuda.is_available()
    torch.ones(1).cuda(local_rank)

    
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    printflock(f"{gpu} is OK (global rank: {rank}/{world_size})")

    dist.barrier()
    if rank == 0:
        printflock(f"pt={torch.__version__}, cuda={torch.version.cuda}, nccl={torch.cuda.nccl.version()}")

except Exception:
    printflock(f"{gpu} is broken")
    raise