Initializes the default process group. Args: local_rank (int): the rank on the current local machine. local_world_size (int): the world size (number of processes running) on the current local machine. shard_id (int): the shard index (machine rank) of the curr
(
local_rank,
local_world_size,
shard_id,
num_shards,
init_method,
dist_backend="nccl",
)
| 54 | |
| 55 | |
| 56 | def init_process_group( |
| 57 | local_rank, |
| 58 | local_world_size, |
| 59 | shard_id, |
| 60 | num_shards, |
| 61 | init_method, |
| 62 | dist_backend="nccl", |
| 63 | ): |
| 64 | """ |
| 65 | Initializes the default process group. |
| 66 | Args: |
| 67 | local_rank (int): the rank on the current local machine. |
| 68 | local_world_size (int): the world size (number of processes running) on |
| 69 | the current local machine. |
| 70 | shard_id (int): the shard index (machine rank) of the current machine. |
| 71 | num_shards (int): number of shards for distributed training. |
| 72 | init_method (string): supporting three different methods for |
| 73 | initializing process groups: |
| 74 | "file": use shared file system to initialize the groups across |
| 75 | different processes. |
| 76 | "tcp": use tcp address to initialize the groups across different |
| 77 | dist_backend (string): backend to use for distributed training. Options |
| 78 | includes gloo, mpi and nccl, the details can be found here: |
| 79 | https://pytorch.org/docs/stable/distributed.html |
| 80 | """ |
| 81 | # Sets the GPU to use. |
| 82 | torch.cuda.set_device(local_rank) |
| 83 | # Initialize the process group. |
| 84 | proc_rank = local_rank + shard_id * local_world_size |
| 85 | world_size = local_world_size * num_shards |
| 86 | dist.init_process_group( |
| 87 | backend=dist_backend, |
| 88 | init_method=init_method, |
| 89 | world_size=world_size, |
| 90 | rank=proc_rank, |
| 91 | ) |
| 92 | |
| 93 | |
| 94 | def is_master_proc(num_gpus=8, debug=False): |
nothing calls this directly
no outgoing calls
no test coverage detected