feature: Adding Training Compiler support for TensorFlow estimator starting TF 2.9

doc/frameworks/tensorflow/sagemaker.tensorflow.rst

@@ -10,6 +10,14 @@ TensorFlow Estimator
     :undoc-members:
     :show-inheritance:
 
+TensorFlow Training Compiler Configuration
+------------------------------------------
+
+.. autoclass:: sagemaker.tensorflow.TrainingCompilerConfig
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 TensorFlow Serving Model
 ------------------------
 

src/sagemaker/huggingface/__init__.py

@@ -17,4 +17,4 @@
 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
+from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig  # noqa: F401

src/sagemaker/huggingface/estimator.py

@@ -26,7 +26,7 @@
 from sagemaker.huggingface.model import HuggingFaceModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
-from sagemaker.training_compiler.config import TrainingCompilerConfig
+from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig
 
 logger = logging.getLogger("sagemaker")
 
@@ -190,6 +190,8 @@ def __init__(
             entry_point, source_dir, hyperparameters, image_uri=image_uri, **kwargs
         )
 
+        self.distribution = distribution or {}
+
         if compiler_config is not None:
             if not isinstance(compiler_config, TrainingCompilerConfig):
                 error_string = (
@@ -199,13 +201,7 @@ def __init__(
                 )
                 raise ValueError(error_string)
             if compiler_config:
-                compiler_config.validate(
-                    image_uri=image_uri,
-                    instance_type=instance_type,
-                    distribution=distribution,
-                )
-
-        self.distribution = distribution or {}
+                compiler_config.validate(self)
         self.compiler_config = compiler_config
 
     def _validate_args(self, image_uri):

src/sagemaker/huggingface/training_compiler/config.py

@@ -0,0 +1,105 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Configuration for the SageMaker Training Compiler."""
+from __future__ import absolute_import
+import logging
+
+from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TrainingCompilerConfig(BaseConfig):
+    """The SageMaker Training Compiler configuration class."""
+
+    SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "g4dn", "p4"]
+
+    def __init__(
+        self,
+        enabled=True,
+        debug=False,
+    ):
+        """This class initializes a ``TrainingCompilerConfig`` instance.
+
+        `Amazon SageMaker Training Compiler
+        <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_
+        is a feature of SageMaker Training
+        and speeds up training jobs by optimizing model execution graphs.
+
+        You can compile Hugging Face models
+        by passing the object of this configuration class to the ``compiler_config``
+        parameter of the :class:`~sagemaker.huggingface.HuggingFace`
+        estimator.
+
+        Args:
+            enabled (bool): Optional. Switch to enable SageMaker Training Compiler.
+                The default is ``True``.
+            debug (bool): Optional. Whether to dump detailed logs for debugging.
+                This comes with a potential performance slowdown.
+                The default is ``False``.
+
+        **Example**: The following code shows the basic usage of the
+        :class:`sagemaker.huggingface.TrainingCompilerConfig()` class
+        to run a HuggingFace training job with the compiler.
+
+        .. code-block:: python
+
+            from sagemaker.huggingface import HuggingFace, TrainingCompilerConfig
+
+            huggingface_estimator=HuggingFace(
+                ...
+                compiler_config=TrainingCompilerConfig()
+            )
+
+        .. seealso::
+
+            For more information about how to enable SageMaker Training Compiler
+            for various training settings such as using TensorFlow-based models,
+            PyTorch-based models, and distributed training,
+            see `Enable SageMaker Training Compiler
+            <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-enable.html>`_
+            in the `Amazon SageMaker Training Compiler developer guide
+            <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_.
+
+        """
+
+        super(TrainingCompilerConfig, self).__init__(enabled=enabled, debug=debug)
+
+    @classmethod
+    def validate(
+        cls,
+        estimator,
+    ):
+        """Checks if SageMaker Training Compiler is configured correctly.
+
+        Args:
+            estimator (str): A estimator object
+                If SageMaker Training Compiler is enabled, it will validate whether
+                the estimator is configured to be compatible with Training Compiler.
+
+        Raises:
+            ValueError: Raised if the requested configuration is not compatible
+                        with SageMaker Training Compiler.
+        """
+
+        super(TrainingCompilerConfig, cls).validate(estimator)
+
+        if estimator.image_uri:
+            error_helper_string = (
+                "Overriding the image URI is currently not supported "
+                "for SageMaker Training Compiler."
+                "Specify the following parameters to run the Hugging Face training job "
+                "with SageMaker Training Compiler enabled: "
+                "transformer_version, tensorflow_version or pytorch_version, and compiler_config."
+            )
+            raise ValueError(error_helper_string)

src/sagemaker/image_uris.py

@@ -134,21 +134,18 @@ def retrieve(
             tolerate_vulnerable_model,
             tolerate_deprecated_model,
         )
-    if training_compiler_config is None:
+
+    if training_compiler_config and (framework == HUGGING_FACE_FRAMEWORK):
+        config = _config_for_framework_and_scope(
+            framework + "-training-compiler", image_scope, accelerator_type
+        )
+    else:
         _framework = framework
         if framework == HUGGING_FACE_FRAMEWORK:
             inference_tool = _get_inference_tool(inference_tool, instance_type)
             if inference_tool == "neuron":
                 _framework = f"{framework}-{inference_tool}"
         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)

src/sagemaker/tensorflow/__init__.py

@@ -16,3 +16,5 @@
 from sagemaker.tensorflow.estimator import TensorFlow  # noqa: F401 (imported but unused)
 from sagemaker.tensorflow.model import TensorFlowModel, TensorFlowPredictor  # noqa: F401
 from sagemaker.tensorflow.processing import TensorFlowProcessor  # noqa: F401
+
+from sagemaker.tensorflow.training_compiler.config import TrainingCompilerConfig  # noqa: F401

src/sagemaker/tensorflow/estimator.py

@@ -26,6 +26,7 @@
 from sagemaker.transformer import Transformer
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow import is_pipeline_variable
+from sagemaker.tensorflow.training_compiler.config import TrainingCompilerConfig
 
 logger = logging.getLogger("sagemaker")
 
@@ -45,7 +46,8 @@ def __init__(
         model_dir=None,
         image_uri=None,
         distribution=None,
-        **kwargs
+        compiler_config=None,
+        **kwargs,
     ):
         """Initialize a ``TensorFlow`` estimator.
 
@@ -157,6 +159,8 @@ def __init__(
 
                     To learn more, see `Training with parameter servers
                     <https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/using_tf.html#training-with-parameter-servers>`_.
+            compiler_config (:class:`~sagemaker.tensorflow.TrainingCompilerConfig`):
+                Configures SageMaker Training Compiler to accelerate training.
 
             **kwargs: Additional kwargs passed to the Framework constructor.
 
@@ -202,6 +206,17 @@ def __init__(
         self.distribution = distribution or {}
 
         self._validate_args(py_version=py_version)
+        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(self)
+        self.compiler_config = compiler_config
 
     def _validate_args(self, py_version):
         """Placeholder docstring"""
@@ -301,7 +316,7 @@ def create_model(
         entry_point=None,
         source_dir=None,
         dependencies=None,
-        **kwargs
+        **kwargs,
     ):
         """Creates ``TensorFlowModel`` object to be used for creating SageMaker model entities.
 
@@ -352,7 +367,7 @@ def create_model(
             entry_point=entry_point,
             source_dir=source_dir,
             dependencies=dependencies,
-            **kwargs
+            **kwargs,
         )
 
     def hyperparameters(self):
@@ -369,6 +384,13 @@ def hyperparameters(self):
         hyperparameters.update(
             EstimatorBase._json_encode_hyperparameters(additional_hyperparameters)
         )
+
+        if self.compiler_config:
+            training_compiler_hyperparameters = self.compiler_config._to_hyperparameter_dict()
+            hyperparameters.update(
+                EstimatorBase._json_encode_hyperparameters(training_compiler_hyperparameters)
+            )
+
         return hyperparameters
 
     def _default_s3_path(self, directory, mpi=False):

src/sagemaker/tensorflow/training_compiler/config.py

@@ -0,0 +1,111 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Configuration for the SageMaker Training Compiler."""
+from __future__ import absolute_import
+import logging
+from packaging.specifiers import SpecifierSet
+from packaging.version import Version
+
+from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TrainingCompilerConfig(BaseConfig):
+    """The SageMaker Training Compiler configuration class."""
+
+    SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "g4dn", "p4", "g5"]
+    MIN_SUPPORTED_VERSION = "2.9"
+
+    def __init__(
+        self,
+        enabled=True,
+        debug=False,
+    ):
+        """This class initializes a ``TrainingCompilerConfig`` instance.
+
+        `Amazon SageMaker Training Compiler
+        <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_
+        is a feature of SageMaker Training
+        and speeds up training jobs by optimizing model execution graphs.
+
+        You can compile TensorFlow models
+        by passing the object of this configuration class to the ``compiler_config``
+        parameter of the :class:`~sagemaker.tensorflow.TensorFlow`
+        estimator.
+
+        Args:
+            enabled (bool): Optional. Switch to enable SageMaker Training Compiler.
+                The default is ``True``.
+            debug (bool): Optional. Whether to dump detailed logs for debugging.
+                This comes with a potential performance slowdown.
+                The default is ``False``.
+
+        **Example**: The following code shows the basic usage of the
+        :class:`sagemaker.tensorflow.TrainingCompilerConfig()` class
+        to run a TensorFlow training job with the compiler.
+
+        .. code-block:: python
+
+            from sagemaker.tensorflow import TensorFlow, TrainingCompilerConfig
+
+            tensorflow_estimator=TensorFlow(
+                ...
+                compiler_config=TrainingCompilerConfig()
+            )
+
+        .. seealso::
+
+            For more information about how to enable SageMaker Training Compiler
+            for various training settings such as using TensorFlow-based models,
+            PyTorch-based models, and distributed training,
+            see `Enable SageMaker Training Compiler
+            <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-enable.html>`_
+            in the `Amazon SageMaker Training Compiler developer guide
+            <https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_.
+
+        """
+
+        super(TrainingCompilerConfig, self).__init__(enabled=enabled, debug=debug)
+
+    @classmethod
+    def validate(
+        cls,
+        estimator,
+    ):
+        """Checks if SageMaker Training Compiler is configured correctly.
+
+        Args:
+            estimator (str): A estimator object
+                If SageMaker Training Compiler is enabled, it will validate whether
+                the estimator is configured to be compatible with Training Compiler.
+
+        Raises:
+            ValueError: Raised if the requested configuration is not compatible
+                        with SageMaker Training Compiler.
+        """
+
+        super(TrainingCompilerConfig, cls).validate(estimator)
+
+        if estimator.framework_version:
+            if Version(estimator.framework_version) in SpecifierSet(
+                f"< {cls.MIN_SUPPORTED_VERSION}"
+            ):
+                error_helper_string = (
+                    "SageMaker Training Compiler only supports TensorFlow version "
+                    ">= {} but received {}"
+                )
+                error_helper_string = error_helper_string.format(
+                    cls.MIN_SUPPORTED_VERSION, estimator.framework_version
+                )
+                raise ValueError(error_helper_string)