diff --git a/src/sagemaker/huggingface/estimator.py b/src/sagemaker/huggingface/estimator.py index 9d154d7183..81b24b5aa3 100644 --- a/src/sagemaker/huggingface/estimator.py +++ b/src/sagemaker/huggingface/estimator.py @@ -50,14 +50,15 @@ def __init__( compiler_config=None, **kwargs, ): - """This ``Estimator`` executes a HuggingFace script in a managed execution environment. + """This estimator runs a Hugging Face training script in a SageMaker training environment. - The managed HuggingFace environment is an Amazon-built Docker container that executes - functions defined in the supplied ``entry_point`` Python script within a SageMaker - Training Job. + The estimator initiates the SageMaker-managed Hugging Face environment + by using the pre-built Hugging Face Docker container and runs + the Hugging Face training script that user provides through + the ``entry_point`` argument. - Training is started by calling - :meth:`~sagemaker.amazon.estimator.Framework.fit` on this Estimator. + After configuring the estimator class, use the class method + :meth:`~sagemaker.amazon.estimator.Framework.fit()` to start a training job. Args: py_version (str): Python version you want to use for executing your model training diff --git a/src/sagemaker/training_compiler/config.py b/src/sagemaker/training_compiler/config.py index 0659c43507..c45fa4cdaf 100644 --- a/src/sagemaker/training_compiler/config.py +++ b/src/sagemaker/training_compiler/config.py @@ -18,11 +18,7 @@ class TrainingCompilerConfig(object): - """The configuration class for accelerating SageMaker training jobs through compilation. - - SageMaker Training Compiler speeds up training by optimizing the model execution graph. - - """ + """The SageMaker Training Compiler configuration class.""" DEBUG_PATH = "/opt/ml/output/data/compiler/" SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "g4dn", "p4"] @@ -37,9 +33,15 @@ def __init__( ): """This class initializes a ``TrainingCompilerConfig`` instance. - Pass the output of it to the ``compiler_config`` + `Amazon SageMaker Training Compiler + `_ + 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` - class. + estimator. Args: enabled (bool): Optional. Switch to enable SageMaker Training Compiler. @@ -48,13 +50,28 @@ def __init__( This comes with a potential performance slowdown. The default is ``False``. - **Example**: The following example shows the basic ``compiler_config`` - parameter configuration, enabling compilation with default parameter values. + **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 TrainingCompilerConfig - compiler_config = TrainingCompilerConfig() + 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 + `_ + in the `Amazon SageMaker Training Compiler developer guide + `_. """