Huggingface Sagemaker-sdk - Distributed Training Demo¶

Distributed Question Answering with transformers scripts + Trainer and squad dataset¶

1. Introduction
2. Development Environment and Permissions
   1. Installation
   2. Development environment
   3. Permissions
3. Fine-tuning & starting Sagemaker Training Job
   1. Creating an Estimator and start a training job
   2. Estimator Parameters
   3. Download fine-tuned model from s3
   4. Attach to old training job to an estimator
4. Coming soon:Push model to the Hugging Face hub

Introduction¶

Welcome to our end-to-end distributed Question-Answering example. In this demo, we will use the Hugging Face transformers and datasets library together with a custom Amazon sagemaker-sdk extension to fine-tune a pre-trained transformer for question-answering on multiple-gpus. In particular, the pre-trained model will be fine-tuned using the squad dataset. The demo will use the new smdistributed library to run training on multiple gpus as training scripting we are going to use one of the transformers example scripts from the repository.

To get started, we need to set up the environment with a few prerequisite steps, for permissions, configurations, and so on.

NOTE: You can run this demo in Sagemaker Studio, your local machine or Sagemaker Notebook Instances

Development Environment and Permissions¶

Installation¶

Note: we only install the required libraries from Hugging Face and AWS. You also need PyTorch or Tensorflow, if you haven´t it installed

Development environment¶

Permissions¶

If you are going to use Sagemaker in a local environment. You need access to an IAM Role with the required permissions for Sagemaker. You can find here more about it.

Fine-tuning & starting Sagemaker Training Job¶

In order to create a sagemaker training job we need an HuggingFace Estimator. The Estimator handles end-to-end Amazon SageMaker training and deployment tasks. In a Estimator we define, which fine-tuning script should be used as entry_point, which instance_type should be used, which hyperparameters are passed in .....

huggingface_estimator = HuggingFace(entry_point='train.py',
                            source_dir='./scripts',
                            sagemaker_session=sess,
                            base_job_name='huggingface-sdk-extension',
                            instance_type='ml.p3.2xlarge',
                            instance_count=1,
                            transformers_version='4.4',
                            pytorch_version='1.6.0',
                            py_version='py36',
                            role=role,
                            hyperparameters = {'epochs': 1,
                                               'train_batch_size': 32,
                                               'model_name':'distilbert-base-uncased'
                                                })

When we create a SageMaker training job, SageMaker takes care of starting and managing all the required ec2 instances for us with the huggingface container, uploads the provided fine-tuning script train.py and downloads the data from our sagemaker_session_bucket into the container at /opt/ml/input/data. Then, it starts the training job by running.

/opt/conda/bin/python train.py --epochs 1 --model_name distilbert-base-uncased --train_batch_size 32

The hyperparameters you define in the HuggingFace estimator are passed in as named arguments.

Sagemaker is providing 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_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.

To run your training job locally you can define instance_type='local' or instance_type='local_gpu' for gpu usage. Note: this does not working within SageMaker Studio

Creating an Estimator and start a training job¶

In this example we are going to use the capability to download/use a fine-tuning script from a git- repository. We are using the run_qa.py from the transformers example scripts. You can find the code here.

Deploying the endpoint¶

To deploy our endpoint, we call deploy() on our HuggingFace estimator object, passing in our desired number of instances and instance type.

Then, we use the returned predictor object to call the endpoint.

Finally, we delete the endpoint again.