Compile and Train the GPT2 Model using the Transformers Trainer API with the SST2 Dataset for Single-Node Multi-GPU Training

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

1. Introduction
2. Development Environment and Permissions
   1. Installation
   2. Permissions
3. SageMaker Training Job
   1. Training with Native TensorFlow
   2. Training with Optimized TensorFlow
   3. Analysis

SageMaker Training Compiler Overview

SageMaker Training Compiler is a capability of SageMaker that makes these hard-to-implement optimizations to reduce training time on GPU instances. The compiler optimizes Deep Learning (DL) models to accelerate training by more efficiently using SageMaker machine learning (ML) GPU instances. SageMaker Training Compiler is available at no additional charge within SageMaker and can help reduce total billable time as it accelerates training.

SageMaker Training Compiler is integrated into the AWS Deep Learning Containers (DLCs). Using the SageMaker Training Compiler enabled AWS DLCs, you can compile and optimize training jobs on GPU instances with minimal changes to your code. Bring your deep learning models to SageMaker and enable SageMaker Training Compiler to accelerate the speed of your training job on SageMaker ML instances for accelerated computing.

For more information, see SageMaker Training Compiler in the Amazon SageMaker Developer Guide.

Introduction

In this demo, you'll use Hugging Face's transformers and datasets libraries with Amazon SageMaker Training Compiler to train the gpt-2 model on the Stanford Sentiment Treebank v2 (SST2) dataset. Please note that by using this notebook you will be downloading SST2 from https://huggingface.co/datasets/sst2 and can check dataset information and terms there. 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, SageMaker notebook instances, or your local machine with AWS CLI set up. If using SageMaker Studio or SageMaker notebook instances, make sure you choose one of the TensorFlow-based kernels, Python 3 (TensorFlow x.y Python 3.x CPU Optimized) or conda_tensorflow_p36 respectively.

NOTE: This notebook uses two ml.p3.8xlarge instances that have multiple GPUs. If you don't have enough quota, see Request a service quota increase for SageMaker resources.

Development Environment

Installation

This example notebook requires the SageMaker Python SDK v2.115.0 and transformers v4.21.

SageMaker environment

SageMaker Training Job

To create a SageMaker training job, we use a HuggingFace estimator. Using the estimator, you can define which fine-tuning script should SageMaker use through entry_point, which instance_type to use for training, which hyperparameters to pass, and so on.

When a SageMaker training job starts, SageMaker takes care of starting and managing all the required machine learning instances, picks up the HuggingFace Deep Learning Container, uploads your training script, and downloads the data from sagemaker_session_bucket into the container at /opt/ml/input/data.

In the following section, you learn how to set up two versions of the SageMaker HuggingFace estimator, a native one without the compiler and an optimized one with the compiler.

Training Setup

First, we define some basic parameters common to all estimators.

This example uses a modified version of HuggingFace training script run_clm.py, which you can find inside the scripts folder.

Note: We recommend you to turn the SageMaker Debugger's profiling and debugging tools off to avoid additional overheads.

Next, we define some basic arguments to be passed to the training script.

Training with Native TensorFlow

The batch size below is the maximum batch we could fit into the memory of an Nvidia V100 GPU (P3). If you change the model, instance type or sequence length etc., please experiment to find the largest batch size that will fit into memory.

Training with Optimized TensorFlow

Compilation through Training Compiler changes the memory footprint of the model. Most commonly, this manifests as a reduction in memory utilization and a consequent increase in the largest batch size that can fit on the GPU. But in some cases the compiler intelligently promotes caching which leads to a decrease in the largest batch size that can fit on the GPU. Note that if you want to change the batch size, you must adjust the learning rate appropriately.

Wait for training jobs to complete

Analysis

Note: If the estimator object is no longer available due to a kernel break or refresh, you need to directly use the training job name and manually attach the training job to a new HuggingFace estimator. For example:

huggingface_estimator = HuggingFace.attach("<your_huggingface_training_job_name>")

Load logs of the training job with SageMaker Training Compiler

Load logs of the training job without SageMaker Training Compiler

Create helper functions for analysis

Convergence of Training Loss

SageMaker Training Compiler does not affect the model convergence behavior. Here, we see the decrease in training loss is similar with and without SageMaker Training Compiler

Total Billable Time

Finally, the decrease in total training time results in a decrease in the billable seconds from SageMaker

Clean up

Stop all training jobs launched if the jobs are still running.

Also, to find instructions on cleaning up resources, see Clean Up in the Amazon SageMaker Developer Guide.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.