Initialize a NCCL group for one rank of the spawned demo.
(rank, world_size)
| 737 | WORLD_SIZE = 2 |
| 738 | |
| 739 | def init_pg(rank, world_size): |
| 740 | """Initialize a NCCL group for one rank of the spawned demo.""" |
| 741 | os.environ["MASTER_ADDR"] = "127.0.0.1" |
| 742 | os.environ["MASTER_PORT"] = "29500" |
| 743 | os.environ["RANK"] = str(rank) |
| 744 | os.environ["WORLD_SIZE"] = str(world_size) |
| 745 | # Use loopback interface for single-node setup |
| 746 | os.environ["NCCL_SOCKET_IFNAME"] = "lo" |
| 747 | dist.init_process_group("nccl", rank=rank, world_size=world_size) |
| 748 | torch.cuda.set_device(rank) |
| 749 | |
| 750 | def _comm_worker(rank, world_size): |
| 751 | """Per-rank worker: build, capture, profile, and (on rank 0) post-process.""" |