Support launch() on CPU machines.
Summary: Pull Request resolved:
@@ -26,6 +26,7 @@ def _find_free_port():
def launch(
main_func,
+ # Should be num_processes_per_machine, but kept for compatibility.
num_gpus_per_machine,
num_machines=1,
machine_rank=0,
@@ -34,13 +35,14 @@ def launch(
timeout=DEFAULT_TIMEOUT,
):
"""
- Launch multi-gpu or distributed training.
+ Launch multi-process or distributed training.
This function must be called on all machines involved in the training.
It will spawn child processes (defined by ``num_gpus_per_machine``) on each machine.
Args:
main_func: a function that will be called by `main_func(*args)`
- num_gpus_per_machine (int): number of GPUs per machine
+ num_gpus_per_machine (int): number of processes per machine. When
+ using GPUs, this should be the number of GPUs.
num_machines (int): the total number of machines
machine_rank (int): the rank of this machine
dist_url (str): url to connect to for distributed jobs, including protocol
@@ -64,7 +66,7 @@ def launch(
"file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://"
)
- mp.spawn(
+ mp.start_processes(
_distributed_worker,
nprocs=num_gpus_per_machine,
args=(
@@ -92,11 +94,13 @@ def _distributed_worker(
args,
timeout=DEFAULT_TIMEOUT,
):
- assert torch.cuda.is_available(), "cuda is not available. Please check your installation."
+ has_gpu = torch.cuda.is_available()
+ if has_gpu:
+ assert num_gpus_per_machine <= torch.cuda.device_count()
global_rank = machine_rank * num_gpus_per_machine + local_rank
try:
dist.init_process_group(
- backend="NCCL",
+ backend="NCCL" if has_gpu else "GLOO",
init_method=dist_url,
world_size=world_size,
rank=global_rank,
@@ -116,7 +120,7 @@ def _distributed_worker(
if i == machine_rank:
comm._LOCAL_PROCESS_GROUP = pg
- assert num_gpus_per_machine <= torch.cuda.device_count()
+ if has_gpu:
torch.cuda.set_device(local_rank)
# synchronize is needed here to prevent a possible timeout after calling init_process_group
