From 0b41d197b4af4c318c64a0aa4401c37dc0cd0d0d Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Mon, 12 Oct 2026 12:02:38 -0500 Subject: [PATCH] fix(train): pass the total rank count to mpirun -np in the MPI driver mpi_driver.main() passed the per-host process count as num_processes, which get_mpirun_command emits as 'mpirun -np'. -np is the total number of ranks, and the hosts are listed with per-host slots (algo-1:8,algo-2:8), so Open MPI placed all ranks on the first host: a 2-node job with process_count_per_node=8 ran 8 ranks on algo-1 and none on algo-2, and with one process per node only a single rank ran in total. The sagemaker-training-toolkit this logic follows uses processes_per_host * len(hosts). Pass process_count * host_count, in the sagemaker-train and sagemaker-core copies of the driver. --- .../train/container_drivers/distributed_drivers/mpi_driver.py | 3 ++- .../train/container_drivers/distributed_drivers/mpi_driver.py | 3 ++- .../tests/unit/train/container_drivers/test_mpi_driver.py | 3 +++ 3 files changed, 7 insertions(+), 2 deletions(-) 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()