Merge branch 'master' into env_support_training

Author: icywang86rui

doc/frameworks/pytorch/using_pytorch.rst

@@ -154,7 +154,7 @@ directories ('train' and 'test').
     pytorch_estimator = PyTorch('pytorch-train.py',
                                 instance_type='ml.p3.2xlarge',
                                 instance_count=1,
-                                framework_version='1.5.0',
+                                framework_version='1.8.0',
                                 py_version='py3',
                                 hyperparameters = {'epochs': 20, 'batch-size': 64, 'learning-rate': 0.1})
     pytorch_estimator.fit({'train': 's3://my-data-bucket/path/to/my/training/data',
@@ -248,7 +248,7 @@ operation.
     pytorch_estimator = PyTorch(entry_point='train_and_deploy.py',
                                 instance_type='ml.p3.2xlarge',
                                 instance_count=1,
-                                framework_version='1.5.0',
+                                framework_version='1.8.0',
                                 py_version='py3')
     pytorch_estimator.fit('s3://my_bucket/my_training_data/')
 

src/sagemaker/fw_utils.py

@@ -59,8 +59,8 @@
     "local_gpu",
 )
 SM_DATAPARALLEL_SUPPORTED_FRAMEWORK_VERSIONS = {
-    "tensorflow": ["2.3.1", "2.3.2", "2.4.1"],
-    "pytorch": ["1.6.0", "1.7.1", "1.8.0"],
+    "tensorflow": ["2.3", "2.3.1", "2.3.2", "2.4", "2.4.1"],
+    "pytorch": ["1.6", "1.6.0", "1.7", "1.7.1", "1.8", "1.8.0"],
 }
 SMDISTRIBUTED_SUPPORTED_STRATEGIES = ["dataparallel", "modelparallel"]
 

tests/unit/test_fw_utils.py

@@ -632,10 +632,15 @@ def test_validate_smdataparallel_args_not_raises():
         (None, None, None, None, smdataparallel_disabled),
         ("ml.p3.16xlarge", "tensorflow", "2.3.1", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "tensorflow", "2.3.2", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "tensorflow", "2.3", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "tensorflow", "2.4.1", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "tensorflow", "2.4", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.6.0", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "1.6", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.7.1", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "1.7", "py3", smdataparallel_enabled),
         ("ml.p3.16xlarge", "pytorch", "1.8.0", "py3", smdataparallel_enabled),
+        ("ml.p3.16xlarge", "pytorch", "1.8", "py3", smdataparallel_enabled),
     ]
     for instance_type, framework_name, framework_version, py_version, distribution in good_args:
         fw_utils._validate_smdataparallel_args(