fix(train): pass the total rank count to mpirun -np in the MPI driver

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,
         )

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,
         )

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()