diff --git a/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/distributed_drivers/mpi_driver.py b/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/distributed_drivers/mpi_driver.py index 9ab8f55a1a..291e867a2e 100644 --- a/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/distributed_drivers/mpi_driver.py +++ b/sagemaker-core/src/sagemaker/core/modules/train/container_drivers/distributed_drivers/mpi_driver.py @@ -103,7 +103,8 @@ def main(): mpi_command = get_mpirun_command( host_count=host_count, host_list=host_list, - num_processes=process_count, + # -np is the total number of ranks: process_count per host on every host. + num_processes=process_count * host_count, additional_options=distributed_config["mpi_additional_options"] or [], entry_script_path=entry_script, ) diff --git a/sagemaker-train/src/sagemaker/train/container_drivers/distributed_drivers/mpi_driver.py b/sagemaker-train/src/sagemaker/train/container_drivers/distributed_drivers/mpi_driver.py index 6ed3fe0d86..a65feae1d0 100644 --- a/sagemaker-train/src/sagemaker/train/container_drivers/distributed_drivers/mpi_driver.py +++ b/sagemaker-train/src/sagemaker/train/container_drivers/distributed_drivers/mpi_driver.py @@ -89,7 +89,8 @@ def main(): mpi_command = get_mpirun_command( host_count=host_count, host_list=host_list, - num_processes=process_count, + # -np is the total number of ranks: process_count per host on every host. + num_processes=process_count * host_count, additional_options=distributed_config["mpi_additional_options"] or [], entry_script_path=entry_script, ) diff --git a/sagemaker-train/tests/unit/train/container_drivers/test_mpi_driver.py b/sagemaker-train/tests/unit/train/container_drivers/test_mpi_driver.py index eccfa3c816..13eaf087b5 100644 --- a/sagemaker-train/tests/unit/train/container_drivers/test_mpi_driver.py +++ b/sagemaker-train/tests/unit/train/container_drivers/test_mpi_driver.py @@ -141,6 +141,9 @@ def test_mpi_driver_master( mock_start_sshd_daemon.assert_called_once() mock_bootstrap_master_node.assert_called_once() mock_get_mpirun_command.assert_called_once() + # mpirun -np is the total number of ranks across all hosts, not the per-host count. + assert mock_get_mpirun_command.call_args.kwargs["host_list"] == ["algo-1:2", "algo-2:2"] + assert mock_get_mpirun_command.call_args.kwargs["num_processes"] == 4 mock_execute_commands.assert_called_once_with(DUMMY_MPI_COMMAND) mock_write_status_file_to_workers.assert_called_once()