fix: Add PT 2.1 as a supported framework for the smdistributed distribution (#4400)

* Allow smdistributed to be used on PyTorch 2.1

* Add PT 2.1 as supported framework for smdistributed

* Add unit test to ensure smdistributed works with PT-2.1.0

Author: ruhanprasad

src/sagemaker/estimator.py

@@ -3285,7 +3285,6 @@ class Framework(EstimatorBase):
     UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM = (
         "2.0.1-gpu-py310-cu121",
         "2.0-gpu-py310-cu121",
-        "2.1.0-gpu-py310",
     )
 
     def __init__(
@@ -3959,7 +3958,7 @@ def _distribution_configuration(self, distribution):
             for unsupported_image in Framework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM:
                 if (
                     unsupported_image in img_uri and not torch_distributed_enabled
-                ):  # disabling DLC images with CUDA12
+                ):  # disabling DLC images without SMDataParallel or SMModelParallel
                     raise ValueError(
                         f"SMDistributed is currently incompatible with DLC image: {img_uri}. "
                         "(Could be due to CUDA version being greater than 11.)"

tests/unit/test_fw_utils.py

@@ -932,6 +932,7 @@ def test_validate_smdataparallel_args_not_raises():
         ("ml.p3.16xlarge", "pytorch", "1.13.1", "py39", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "2.0.0", "py310", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "2.0.1", "py310", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "2.1.0", "py310", smdataparallel_enabled),
         ("ml.p3.16xlarge", "tensorflow", "2.4.1", "py3", smdataparallel_enabled_custom_mpi),
         ("ml.p3.16xlarge", "tensorflow", "2.4.1", "py37", smdataparallel_enabled_custom_mpi),
         ("ml.p3.16xlarge", "tensorflow", "2.4.3", "py3", smdataparallel_enabled_custom_mpi),
@@ -955,6 +956,7 @@ def test_validate_smdataparallel_args_not_raises():
         ("ml.p3.16xlarge", "pytorch", "1.13.1", "py39", smdataparallel_enabled_custom_mpi),
         ("ml.p3.16xlarge", "pytorch", "2.0.0", "py310", smdataparallel_enabled_custom_mpi),
         ("ml.p3.16xlarge", "pytorch", "2.0.1", "py310", smdataparallel_enabled_custom_mpi),
+        ("ml.p3.16xlarge", "pytorch", "2.1.0", "py310", smdataparallel_enabled_custom_mpi),
     ]
     for instance_type, framework_name, framework_version, py_version, distribution in good_args:
         fw_utils._validate_smdataparallel_args(