deprecation: Blocking submission of distribution configs that try to use CUDA 12.1 enabled DLC containers and/or p5 instance with smdistributed enabled and torch-distributed disabled.

Author: rnadimp

src/sagemaker/estimator.py

@@ -3414,6 +3414,7 @@ def __init__(
         self.checkpoint_s3_uri = checkpoint_s3_uri
         self.checkpoint_local_path = checkpoint_local_path
         self.enable_sagemaker_metrics = enable_sagemaker_metrics
+        self.unsupported_dlc_image_for_sm_parallelism = ["2.0.1-gpu-py310-cu121"]
 
     def _prepare_for_training(self, job_name=None):
         """Set hyperparameters needed for training. This method will also validate ``source_dir``.
@@ -3846,8 +3847,18 @@ def _distribution_configuration(self, distribution):
             # smdistributed strategy selected
             if get_mp_parameters(distribution):
                 distribution_config["mp_parameters"] = get_mp_parameters(distribution)
+            # first make sure torch_distributed is enabled if instance type is p5
+            torch_distributed_enabled = distribution.get("torch_distributed").get("enabled", False)
             smdistributed = distribution["smdistributed"]
             smdataparallel_enabled = smdistributed.get("dataparallel", {}).get("enabled", False)
+            p5_enabled = bool("p5.48xlarge" in self.instance_type)
+            img_uri = "" if self.image_uri is None else self.image_uri
+            for unsupported_image in self.unsupported_dlc_image_for_sm_parallelism:
+                if unsupported_image in img_uri and not torch_distributed_enabled: #disabling DLC images with CUDA12
+                    raise ValueError(f"SMDistributed is currently incompatible with DLC image: {img_uri}. (Could be due to CUDA version being greater than 11.)")
+            if not torch_distributed_enabled and p5_enabled: #disabling p5 when torch distributed is disabled
+                raise ValueError("SMModelParallel and SMDataParallel currently do not support p5 instances.")
+            # smdistributed strategy selected with supported instance type
             distribution_config[self.LAUNCH_SM_DDP_ENV_NAME] = smdataparallel_enabled
             distribution_config[self.INSTANCE_TYPE] = self.instance_type
             if smdataparallel_enabled:

tests/unit/test_estimator.py

@@ -169,7 +169,20 @@
     "mpi": {"enabled": True, "custom_mpi_options": "options", "processes_per_host": 2}
 }
 DISTRIBUTION_SM_DDP_ENABLED = {
-    "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "options"}}
+    "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "options"}},
+    "torch_distributed": {"enabled": False}
+}
+DISTRIBUTION_SM_DDP_DISABLED = {
+    "smdistributed": {"enabled": True},
+    "torch_distributed": {"enabled": False}
+}
+DISTRIBUTION_SM_TORCH_DIST_AND_DDP_ENABLED = {
+    "smdistributed": {"dataparallel": {"enabled": True, "custom_mpi_options": "options"}},
+    "torch_distributed": {"enabled": True}
+}
+DISTRIBUTION_SM_TORCH_DIST_AND_DDP_DISABLED = {
+    "smdistributed": {"enabled": True},
+    "torch_distributed": {"enabled": True}
 }
 MOCKED_S3_URI = "s3://mocked_s3_uri_from_source_dir"
 _DEFINITION_CONFIG = PipelineDefinitionConfig(use_custom_job_prefix=False)
@@ -309,6 +322,60 @@ def training_job_description(sagemaker_session):
     sagemaker_session.describe_training_job = mock_describe_training_job
     return returned_job_description
 
+def test_validate_smdistributed_p5_raises(sagemaker_session):
+    # supported DLC image
+    f = DummyFramework(
+        "some_script.py", 
+        role="DummyRole",
+        instance_type="ml.p5.48xlarge", 
+        sagemaker_session=sagemaker_session,
+        output_path="outputpath",
+        image_uri="some_acceptable_image"
+    )
+    #both fail because instance type is p5 and torch_distributed is off
+    with pytest.raises(ValueError):
+        f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED) 
+    with pytest.raises(ValueError):
+        f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED) 
+    # unsupported DLC image
+    f = DummyFramework(
+        "some_script.py",
+        role="DummyRole",
+        instance_type="ml.p5.48xlarge",
+        sagemaker_session=sagemaker_session,
+        output_path="outputpath",
+        image_uri="ecr-url/2.0.1-gpu-py310-cu121-ubuntu20.04-sagemaker-pr-3303"
+    )
+    #both fail due to unsupported CUDA12 DLC image
+    with pytest.raises(ValueError):
+        f._distribution_configuration(DISTRIBUTION_SM_DDP_ENABLED)
+    with pytest.raises(ValueError):
+        f._distribution_configuration(DISTRIBUTION_SM_DDP_DISABLED)
+
+def test_validate_smdistributed_p5_not_raises(sagemaker_session):
+    f = DummyFramework(
+        "some_script.py",
+        role="DummyRole",
+        instance_type="ml.p5.48xlarge",
+        sagemaker_session=sagemaker_session,
+        output_path="outputpath",
+        image_uri="ecr-url/2.0.1-gpu-py310-cu121-ubuntu20.04-sagemaker-pr-3303"
+    )
+    #testing with p5 instance and torch_distributed enabled
+    f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_ENABLED)
+    f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_DISABLED)
+    f = DummyFramework(
+        "some_script.py",
+        role="DummyRole",
+        instance_type="ml.p4.24xlarge",
+        sagemaker_session=sagemaker_session,
+        output_path="outputpath",
+        image_uri="some_acceptable_image"
+    )
+    #testing backwards compatability with p4d instances
+    f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_ENABLED)
+    f._distribution_configuration(DISTRIBUTION_SM_TORCH_DIST_AND_DDP_DISABLED)
+
 
 def test_framework_all_init_args(sagemaker_session):
     f = DummyFramework(