change: allow smdistributed to be enabled with torch_distributed.

Author: viclzhu

src/sagemaker/estimator.py

@@ -3822,14 +3822,10 @@ def _distribution_configuration(self, distribution):
                 "custom_mpi_options", ""
             )
 
-            if get_mp_parameters(distribution):
-                distribution_config["mp_parameters"] = get_mp_parameters(distribution)
-
-        elif "modelparallel" in distribution.get("smdistributed", {}):
-            raise ValueError("Cannot use Model Parallelism without MPI enabled!")
-
         if "smdistributed" in distribution:
             # smdistributed strategy selected
+            if get_mp_parameters(distribution):
+                distribution_config["mp_parameters"] = get_mp_parameters(distribution)
             smdistributed = distribution["smdistributed"]
             smdataparallel_enabled = smdistributed.get("dataparallel", {}).get("enabled", False)
             distribution_config[self.LAUNCH_SM_DDP_ENV_NAME] = smdataparallel_enabled

src/sagemaker/pytorch/estimator.py

@@ -326,6 +326,9 @@ def _pytorch_distribution_configuration(self, distribution):
             if self.instance_type is not None:
                 distribution_config[self.INSTANCE_TYPE_ENV_NAME] = self.instance_type
         elif torch_distributed_enabled:
+            if "smdistributed" in distribution:
+                # Enable torch_distributed for smdistributed.
+                distribution_config = self._distribution_configuration(distribution=distribution)
             distribution_config[self.LAUNCH_TORCH_DISTRIBUTED_ENV_NAME] = torch_distributed_enabled
             if self.instance_type is not None:
                 distribution_config[self.INSTANCE_TYPE_ENV_NAME] = self.instance_type