Skip to content

Commit 169dffd

Browse files
authored
documentation: update sagemaker training compiler docstring (#2969)
1 parent d610bfb commit 169dffd

File tree

2 files changed

+35
-17
lines changed

2 files changed

+35
-17
lines changed

src/sagemaker/huggingface/estimator.py

+7-6
Original file line numberDiff line numberDiff line change
@@ -50,14 +50,15 @@ def __init__(
5050
compiler_config=None,
5151
**kwargs,
5252
):
53-
"""This ``Estimator`` executes a HuggingFace script in a managed execution environment.
53+
"""This estimator runs a Hugging Face training script in a SageMaker training environment.
5454
55-
The managed HuggingFace environment is an Amazon-built Docker container that executes
56-
functions defined in the supplied ``entry_point`` Python script within a SageMaker
57-
Training Job.
55+
The estimator initiates the SageMaker-managed Hugging Face environment
56+
by using the pre-built Hugging Face Docker container and runs
57+
the Hugging Face training script that user provides through
58+
the ``entry_point`` argument.
5859
59-
Training is started by calling
60-
:meth:`~sagemaker.amazon.estimator.Framework.fit` on this Estimator.
60+
After configuring the estimator class, use the class method
61+
:meth:`~sagemaker.amazon.estimator.Framework.fit()` to start a training job.
6162
6263
Args:
6364
py_version (str): Python version you want to use for executing your model training

src/sagemaker/training_compiler/config.py

+28-11
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,7 @@
1818

1919

2020
class TrainingCompilerConfig(object):
21-
"""The configuration class for accelerating SageMaker training jobs through compilation.
22-
23-
SageMaker Training Compiler speeds up training by optimizing the model execution graph.
24-
25-
"""
21+
"""The SageMaker Training Compiler configuration class."""
2622

2723
DEBUG_PATH = "/opt/ml/output/data/compiler/"
2824
SUPPORTED_INSTANCE_CLASS_PREFIXES = ["p3", "g4dn", "p4"]
@@ -37,9 +33,15 @@ def __init__(
3733
):
3834
"""This class initializes a ``TrainingCompilerConfig`` instance.
3935
40-
Pass the output of it to the ``compiler_config``
36+
`Amazon SageMaker Training Compiler
37+
<https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_
38+
is a feature of SageMaker Training
39+
and speeds up training jobs by optimizing model execution graphs.
40+
41+
You can compile Hugging Face models
42+
by passing the object of this configuration class to the ``compiler_config``
4143
parameter of the :class:`~sagemaker.huggingface.HuggingFace`
42-
class.
44+
estimator.
4345
4446
Args:
4547
enabled (bool): Optional. Switch to enable SageMaker Training Compiler.
@@ -48,13 +50,28 @@ def __init__(
4850
This comes with a potential performance slowdown.
4951
The default is ``False``.
5052
51-
**Example**: The following example shows the basic ``compiler_config``
52-
parameter configuration, enabling compilation with default parameter values.
53+
**Example**: The following code shows the basic usage of the
54+
:class:`sagemaker.huggingface.TrainingCompilerConfig()` class
55+
to run a HuggingFace training job with the compiler.
5356
5457
.. code-block:: python
5558
56-
from sagemaker.huggingface import TrainingCompilerConfig
57-
compiler_config = TrainingCompilerConfig()
59+
from sagemaker.huggingface import HuggingFace, TrainingCompilerConfig
60+
61+
huggingface_estimator=HuggingFace(
62+
...
63+
compiler_config=TrainingCompilerConfig()
64+
)
65+
66+
.. seealso::
67+
68+
For more information about how to enable SageMaker Training Compiler
69+
for various training settings such as using TensorFlow-based models,
70+
PyTorch-based models, and distributed training,
71+
see `Enable SageMaker Training Compiler
72+
<https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-enable.html>`_
73+
in the `Amazon SageMaker Training Compiler developer guide
74+
<https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler.html>`_.
5875
5976
"""
6077

0 commit comments

Comments
 (0)