(args, rank, world_size)
| 30 | |
| 31 | |
| 32 | def setup(args, rank, world_size): |
| 33 | os.environ['MASTER_ADDR'] = args.master_addr |
| 34 | os.environ['MASTER_PORT'] = args.master_port |
| 35 | |
| 36 | # initialize the process group |
| 37 | dist.init_process_group("nccl", rank=rank, world_size=world_size) |
| 38 | |
| 39 | def cleanup(): |
| 40 | dist.destroy_process_group() |