Getting Started: Train a HuggingFace Transformers Model

To train a 🤗 Transformers model by using the HuggingFace SageMaker Python SDK you need to:

• Prepare a training script
• Create a HuggingFace Estimator
• Run training by calling the fit method
• Access you model

Setup & Installation

Before you can train a transformers models with Amazon SageMaker you need to sign up for an AWS account.

After you complete these tasks you can get started using either SageMaker Studio, SageMaker Notebook Instances, or a local environment.

Upgrade to the latest sagemaker version.

pip install sagemaker --upgrade

import sagemaker
sess = sagemaker.Session()
role = sagemaker.get_execution_role()

Prepare a 🤗 Transformers fine-tuning script.

The training script is very similar to a training script you might run outside of SageMaker, but you can access useful properties about the training environment through various environment variables, including the following:

• SM_MODEL_DIR: A string that represents the path where the training job writes the model artifacts to. After training, artifacts in this directory are uploaded to S3 for model hosting. SM_MODEL_DIR is always set to /opt/ml/model.

• SM_NUM_GPUS: An integer representing the number of GPUs available to the host.

• SM_CHANNEL_XXXX: A string that represents the path to the directory that contains the input data for the specified channel. For example, if you specify two input channels in the HuggingFace estimator’s fit call, named train and test, the environment variables SM_CHANNEL_TRAIN and SM_CHANNEL_TEST are set.

Later we define hyperparameters in the HuggingFace Estimator, which are passed in as named arguments and and can be processed with the ArgumentParser().

import transformers
import datasets
import argparse
import os

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--per_device_train_batch_size", type=int, default=32)
    parser.add_argument("--model_name_or_path", type=str)

    # Data, model, and output directories
    parser.add_argument("--model-dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])

Create an HuggingFace Estimator

You run 🤗 Transformers training scripts on SageMaker by creating HuggingFace Estimators. The Estimator handles end-to-end Amazon SageMaker training. The training of your script is invoked when you call fit on a HuggingFace Estimator. In the Estimator you define, which fine-tuning script should be used as entry_point, which instance_type should be used, which hyperparameters are passed in.

from sagemaker.huggingface import HuggingFace


# hyperparameters, which are passed into the training job
hyperparameters={'epochs': 1,
                 'per_device_train_batch_size': 32,
                 'model_name_or_path': 'distilbert-base-uncased'
                 }

# create the Estimator
huggingface_estimator = HuggingFace(
        entry_point='train.py',
        source_dir='./scripts',
        instance_type='ml.p3.2xlarge',
        instance_count=1,
        role=role,
        transformers_version='4.4',
        pytorch_version='1.6',
        py_version='py36',
        hyperparameters = hyperparameters
)

To run the TrainingJob locally you can define instance_type='local' or instance_type='local-gpu' for gpu usage.

Note: this does not working within SageMaker Studio

Execute Training

You start your TrainingJob by calling fit on a HuggingFace Estimator. In the fit method you specify your input training data, like a string S3 URI s3://my-bucket/my-training-data or a FileSystemInput for EFS or FSx Lustre.

huggingface_estimator.fit(
  {'train': 's3://sagemaker-us-east-1-558105141721/samples/datasets/imdb/train',
   'test': 's3://sagemaker-us-east-1-558105141721/samples/datasets/imdb/test'}
)

Access trained model

After training is done you can access your model either through the AWS console or downloading it directly from S3.

from sagemaker.s3 import S3Downloader

S3Downloader.download(
    s3_uri=huggingface_estimator.model_data, # s3 uri where the trained model is located
    local_path='.', # local path where *.targ.gz is saved
    sagemaker_session=sess # sagemaker session used for training the model
)

References