Allow one GPU multi-node distributed training
@@ -291,17 +291,17 @@ def train_model(
# passed the wrong thing - cuda_devices are required.
device_ids = distributed_params.pop("cuda_devices", None)
multi_device = isinstance(device_ids, list) and len(device_ids) > 1
+ num_nodes = distributed_params.pop("num_nodes", 1)
- if not multi_device:
+ if not (multi_device or num_nodes > 1):
raise ConfigurationError(
- "Multiple cuda devices need to be configured to run distributed training."
+ "Multiple cuda devices/nodes need to be configured to run distributed training."
)
check_for_gpu(device_ids)
master_addr = distributed_params.pop("master_address", "127.0.0.1")
master_port = distributed_params.pop("master_port", 29500)
num_procs = len(device_ids)
- num_nodes = distributed_params.pop("num_nodes", 1)
world_size = num_nodes * num_procs
os.environ["MASTER_ADDR"] = master_addr
