feature: SM Training Compiler with an UI to enable/disable compilation for HuggingFace DLCs to speedup training

* feature: SM Training Compiler with an UI to enable/disable compilation for HuggingFace DLCs to speedup training

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

* change: New repo for training compiler

Co-authored-by: Sergey Togulev <[email protected]>

Co-authored-by: Loki <[email protected]>
Co-authored-by: Miyoung <[email protected]>
Co-authored-by: Sergey Togulev <[email protected]>
Co-authored-by: Sergey Togulev <[email protected]>

Author: 5 people

doc/frameworks/huggingface/sagemaker.huggingface.rst

@@ -4,7 +4,15 @@ Hugging Face
 Hugging Face Estimator
 ----------------------
 
-.. autoclass:: sagemaker.huggingface.estimator.HuggingFace
+.. autoclass:: sagemaker.huggingface.HuggingFace
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+Hugging Face Training Compiler Configuration
+--------------------------------------------
+
+.. autoclass:: sagemaker.huggingface.TrainingCompilerConfig
     :members:
     :undoc-members:
     :show-inheritance:
@@ -17,8 +25,8 @@ Hugging Face Model
     :undoc-members:
     :show-inheritance:
 
-HuggingFace Predictor
----------------------
+Hugging Face Predictor
+----------------------
 
 .. autoclass:: sagemaker.huggingface.model.HuggingFacePredictor
     :members:

src/sagemaker/estimator.py

@@ -2432,10 +2432,12 @@ def training_image_uri(self):
             distribution = self.distribution  # pylint: disable=no-member
         else:
             distribution = None
+        compiler_config = getattr(self, "compiler_config", None)
 
         if hasattr(self, "tensorflow_version") or hasattr(self, "pytorch_version"):
             processor = image_uris._processor(self.instance_type, ["cpu", "gpu"])
-            container_version = "cu110-ubuntu18.04" if processor == "gpu" else None
+            is_native_huggingface_gpu = processor == "gpu" and not compiler_config
+            container_version = "cu110-ubuntu18.04" if is_native_huggingface_gpu else None
             if self.tensorflow_version is not None:  # pylint: disable=no-member
                 base_framework_version = (
                     f"tensorflow{self.tensorflow_version}"  # pylint: disable=no-member
@@ -2458,6 +2460,7 @@ def training_image_uri(self):
             distribution=distribution,
             base_framework_version=base_framework_version,
             container_version=container_version,
+            training_compiler_config=compiler_config,
         )
 
     @classmethod

src/sagemaker/fw_utils.py

@@ -296,7 +296,8 @@ def framework_name_from_image(image_uri):
     name_pattern = re.compile(
         r"""^(?:sagemaker(?:-rl)?-)?
         (tensorflow|mxnet|chainer|pytorch|scikit-learn|xgboost
-        |huggingface-tensorflow|huggingface-pytorch)(?:-)?
+        |huggingface-tensorflow|huggingface-pytorch
+        |huggingface-tensorflow-trcomp|huggingface-pytorch-trcomp)(?:-)?
         (scriptmode|training)?
         :(.*)-(.*?)-(py2|py3\d*)(?:.*)$""",
         re.VERBOSE,

src/sagemaker/huggingface/__init__.py

@@ -16,3 +16,5 @@
 from sagemaker.huggingface.estimator import HuggingFace  # noqa: F401
 from sagemaker.huggingface.model import HuggingFaceModel, HuggingFacePredictor  # noqa: F401
 from sagemaker.huggingface.processing import HuggingFaceProcessor  # noqa:F401
+
+from sagemaker.training_compiler.config import TrainingCompilerConfig  # noqa: F401

src/sagemaker/huggingface/estimator.py

@@ -26,6 +26,8 @@
 from sagemaker.huggingface.model import HuggingFaceModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
+from sagemaker.training_compiler.config import TrainingCompilerConfig
+
 logger = logging.getLogger("sagemaker")
 
 
@@ -45,7 +47,8 @@ def __init__(
         hyperparameters=None,
         image_uri=None,
         distribution=None,
-        **kwargs
+        compiler_config=None,
+        **kwargs,
     ):
         """This ``Estimator`` executes a HuggingFace script in a managed execution environment.
 
@@ -135,6 +138,8 @@ def __init__(
                             }
                         }
                     }
+            compiler_config (:class:`~sagemaker.huggingface.TrainingCompilerConfig`):
+                Configures SageMaker Training Compiler to accelerate training.
 
             **kwargs: Additional kwargs passed to the :class:`~sagemaker.estimator.Framework`
                 constructor.
@@ -152,16 +157,16 @@ def __init__(
 
         self._validate_args(image_uri=image_uri)
 
-        if distribution is not None:
-            instance_type = renamed_kwargs(
-                "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
-            )
+        instance_type = renamed_kwargs(
+            "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
+        )
 
-            base_framework_name = "tensorflow" if tensorflow_version is not None else "pytorch"
-            base_framework_version = (
-                tensorflow_version if tensorflow_version is not None else pytorch_version
-            )
+        base_framework_name = "tensorflow" if tensorflow_version is not None else "pytorch"
+        base_framework_version = (
+            tensorflow_version if tensorflow_version is not None else pytorch_version
+        )
 
+        if distribution is not None:
             validate_smdistributed(
                 instance_type=instance_type,
                 framework_name=base_framework_name,
@@ -183,7 +188,24 @@ def __init__(
         super(HuggingFace, self).__init__(
             entry_point, source_dir, hyperparameters, image_uri=image_uri, **kwargs
         )
+
+        if compiler_config is not None:
+            if not isinstance(compiler_config, TrainingCompilerConfig):
+                error_string = (
+                    f"Expected instance of type {TrainingCompilerConfig}"
+                    f"for argument compiler_config. "
+                    f"Instead got {type(compiler_config)}"
+                )
+                raise ValueError(error_string)
+            if compiler_config:
+                compiler_config.validate(
+                    image_uri=image_uri,
+                    instance_type=instance_type,
+                    distribution=distribution,
+                )
+
         self.distribution = distribution or {}
+        self.compiler_config = compiler_config
 
     def _validate_args(self, image_uri):
         """Placeholder docstring"""
@@ -220,10 +242,19 @@ def _validate_args(self, image_uri):
     def hyperparameters(self):
         """Return hyperparameters used by your custom PyTorch code during model training."""
         hyperparameters = super(HuggingFace, self).hyperparameters()
-        additional_hyperparameters = self._distribution_configuration(
+        distributed_training_hyperparameters = self._distribution_configuration(
             distribution=self.distribution
         )
-        hyperparameters.update(Framework._json_encode_hyperparameters(additional_hyperparameters))
+        hyperparameters.update(
+            Framework._json_encode_hyperparameters(distributed_training_hyperparameters)
+        )
+
+        if self.compiler_config:
+            training_compiler_hyperparameters = self.compiler_config._to_hyperparameter_dict()
+            hyperparameters.update(
+                Framework._json_encode_hyperparameters(training_compiler_hyperparameters)
+            )
+
         return hyperparameters
 
     def create_model(
@@ -234,7 +265,7 @@ def create_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
-        **kwargs
+        **kwargs,
     ):
         """Create a SageMaker ``HuggingFaceModel`` object that can be deployed to an ``Endpoint``.
 
@@ -286,7 +317,7 @@ def create_model(
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
             dependencies=(dependencies or self.dependencies),
-            **kwargs
+            **kwargs,
         )
 
     @classmethod
@@ -311,7 +342,7 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         if tag is None:
             framework_version = None
         else:
-            framework, pt_or_tf = framework.split("-")
+            framework, pt_or_tf = framework.split("-")[:2]
             tag_pattern = re.compile(r"^(.*)-transformers(.*)-(cpu|gpu)-(py2|py3\d*)$")
             tag_match = tag_pattern.match(tag)
             pt_or_tf_version = tag_match.group(1)

src/sagemaker/image_uri_config/huggingface-training-compiler.json

@@ -0,0 +1,38 @@
+{
+    "training": {
+        "processors": ["gpu"],
+        "version_aliases": {
+            "4.10": "4.10.2"
+        },
+        "versions": {
+            "4.10.2": {
+                "version_aliases": {
+                    "pytorch1.9": "pytorch1.9.0",
+                    "tensorflow2.5": "tensorflow2.5.1"
+                },
+                "pytorch1.9.0": {
+                    "py_versions": ["py38"],
+                    "registries": {
+                        "eu-west-1": "763104351884",
+                        "us-east-1": "763104351884",
+                        "us-east-2": "763104351884",
+                        "us-west-2": "763104351884"
+                    },
+                    "repository": "huggingface-pytorch-trcomp-training",
+                    "container_version": {"gpu":"cu111-ubuntu20.04"}
+                },
+                "tensorflow2.5.1": {
+                    "py_versions": ["py37"],
+                    "registries": {
+                        "eu-west-1": "763104351884",
+                        "us-east-1": "763104351884",
+                        "us-east-2": "763104351884",
+                        "us-west-2": "763104351884"
+                    },
+                    "repository": "huggingface-tensorflow-trcomp-training",
+                    "container_version": {"gpu":"cu112-ubuntu18.04"}
+                }
+            }
+        }
+    }
+}

src/sagemaker/image_uris.py

@@ -38,6 +38,7 @@ def retrieve(
     container_version=None,
     distribution=None,
     base_framework_version=None,
+    training_compiler_config=None,
 ):
     """Retrieves the ECR URI for the Docker image matching the given arguments.
 
@@ -65,6 +66,8 @@ def retrieve(
             https://github.com/aws/deep-learning-containers/blob/master/available_images.md
             (default: None).
         distribution (dict): A dictionary with information on how to run distributed training
+        training_compiler_config (:class:`~sagemaker.training_compiler.TrainingCompilerConfig`):
+            A configuration class for the SageMaker Training Compiler
             (default: None).
 
     Returns:
@@ -73,8 +76,16 @@ def retrieve(
     Raises:
         ValueError: If the combination of arguments specified is not supported.
     """
-
-    config = _config_for_framework_and_scope(framework, image_scope, accelerator_type)
+    if training_compiler_config is None:
+        config = _config_for_framework_and_scope(framework, image_scope, accelerator_type)
+    elif framework == HUGGING_FACE_FRAMEWORK:
+        config = _config_for_framework_and_scope(
+            framework + "-training-compiler", image_scope, accelerator_type
+        )
+    else:
+        raise ValueError(
+            "Unsupported Configuration: Training Compiler is only supported with HuggingFace"
+        )
     original_version = version
     version = _validate_version_and_set_if_needed(version, config, framework)
     version_config = config["versions"][_version_for_config(version, config)]
@@ -108,7 +119,6 @@ def retrieve(
             re.compile("^(pytorch|tensorflow)(.*)$").match(base_framework_version).group(2)
         )
         tag_prefix = f"{pt_or_tf_version}-transformers{original_version}"
-
     else:
         tag_prefix = version_config.get("tag_prefix", version)
 
@@ -134,9 +144,7 @@ def retrieve(
             "pytorch-1.6-gpu-py3": "cu110-ubuntu18.04-v3",
             "pytorch-1.6.0-gpu-py3": "cu110-ubuntu18.04",
         }
-
         key = "-".join([framework, tag])
-
         if key in container_versions:
             tag = "-".join([tag, container_versions[key]])