参考https://theaisummer.com/distributed-training-pytorch/
一般来说,最好先用非分布式的模型下debug,然后包装成分布式用命令行运行。
1、在__main__ 中parse args之前init distributed
#============================Example 2=================================
def init_distributed():
# Initializes the distributed backend which will take care of synchronizing nodes/GPUs
dist_url = "env://" # default
# # only works with torch.distributed.launch // torch.run
rank = int(os.environ["RANK"])
world_size = int(os.environ['WORLD_SIZE'])
local_rank = int(os.environ['LOCAL_RANK'])
dist.init_process_group(
backend="nccl",
init_method=dist_url,
world_size=world_size,
rank=rank)
# this will make all .cuda() calls work properly
torch.cuda.set_device(local_rank)
# synchron