fix: add PT 2.2 support for smdistributed, pytorchddp, and torch_distributed distributions (aws#4480)

* Add support for smdistributed, pytorchddp, torch_distributed for PT 2.2

* formatting

* formatting

---------

Co-authored-by: liujiaor <[email protected]>

Author: 2 people

src/sagemaker/estimator.py

@@ -3285,10 +3285,7 @@ class Framework(EstimatorBase):
     """
 
     _framework_name = None
-    UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM = (
-        "2.0.1-gpu-py310-cu121",
-        "2.0-gpu-py310-cu121",
-    )
+    UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM = ("2.0.1-gpu-py310-cu121",)
 
     def __init__(
         self,

src/sagemaker/fw_utils.py

@@ -161,7 +161,6 @@
     "2.2.0",
 ]
 
-
 TORCH_DISTRIBUTED_GPU_SUPPORTED_FRAMEWORK_VERSIONS = [
     "1.13.1",
     "2.0.0",

tests/unit/test_fw_utils.py

@@ -933,6 +933,7 @@ def test_validate_smdataparallel_args_not_raises():
         ("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", "pytorch", "2.2.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),
@@ -957,6 +958,7 @@ def test_validate_smdataparallel_args_not_raises():
         ("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),
+        ("ml.p3.16xlarge", "pytorch", "2.2.0", "py310", smdataparallel_enabled_custom_mpi),
     ]
     for instance_type, framework_name, framework_version, py_version, distribution in good_args:
         fw_utils._validate_smdataparallel_args(