Compile and Tune a Vision Transformer Model using HyperParameter Tuner on a Single Node

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. SageMaker Training Compiler Overview
   1. Introduction
2. Working with the Caltech-256 dataset
   1. Installation
   2. SageMaker environment
3. How effective is SageMaker Training Compiler?
   1. SageMaker Training Job
   2. Training Setup
   3. Experimenting with Native TensorFlow
   4. Experimenting with Optimized TensorFlow
   5. Wait for tuning jobs to complete
   6. Fastest Training Job
4. Continue tuning with SageMaker Training Compiler
   1. Wait for tuning jobs to complete
   2. Fastest Convergence
5. Conclusion
6. Clean up

SageMaker Training Compiler Overview

SageMaker Training Compiler is a capability of SageMaker that makes hard-to-implement optimizations to reduce training time on GPU instances. The compiler optimizes 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 SageMaker Training Compiler and SageMaker Hyperparameter Tuner to speed up training the Vision Transformer model on the Caltech-256 dataset. 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_p39 respectively.

NOTE: This notebook uses 20 ml.p3.2xlarge instances, each with a single GPU. However, it can easily be extended to multiple GPUs on a single node. If you don't have enough quota, see Request a service quota increase for SageMaker resources.

Development Environment

Installation

This example notebook requires SageMaker Python SDK v2.115.0 or later

SageMaker environment

Working with the Caltech-256 dataset

We have hosted the Caltech-256 dataset in S3 in us-west-2. We will transfer this dataset to your account and region for use with SageMaker Training.

The dataset consists of JPEG images organized into directories with each directory representing an object category.

How effective is SageMaker Training Compiler?

The effectiveness of SageMaker Training Compiler depends on the model architecture, model size, input shape, and the training loop. Please refer to our Best Practices documentation to understand how to get the most out of your training job using SageMaker Training Compiler. In this section, we will compare and contrast a training job with and without SageMaker Training Compiler.

SageMaker Training Job

To create a SageMaker training job, we use a TensorFlow estimator. Using the estimator, you can define which training 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 TensorFlow 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 TensorFlow estimator, a native one without the compiler and an optimized one with the compiler.

Training Setup

In this section, we set our hyperparameters to a naive first guess. Notice the low value for EPOCHS - this is because we are just experimenting with our hyperparameters to find the best setting that will lead to the fastest training. The effectiveness of SageMaker Training Compiler is often apparent within the first few steps. In the example below we will inspect the speed of the training job after every epoch.

Experimenting with Native TensorFlow

We attempt to find the largest BATCH_SIZE that can fit into the memory of a ml.p3.2xlarge instance. This will consequently give us the fastest training speed.

SageMaker Hyperparameter Tuning Job

We use the sagemaker.tuner.HyperparameterTuner object to define a Hyperparameter Tuning Job. It will import the training job configuration specified in the estimator. We additionally specify some metric_definitions to extract training metrics from the training logs. From these metric_definitions we select a single metric as the objective_metric_name and configure the tuning job to Minimize or Maximize it. We further provide a constrained search space through the hyperparameter_ranges argument.

We can limit the number of training jobs spawned concurrently in the max_parallel_jobs argument and limit the total number of training jobs spawned in the max_jobs argument.

For more information regarding SageMaker Hyperparameter Tuner refer to Perform Automatic Model Tuning with SageMaker

In the example below, we are trying to find the best batch size between 32 and 80 that will result in the smallest possible epoch latency, by launching 40 training jobs, 10 at a time. The range for batch sizes is our best guess. You can always reuse and restart a tuning job with an extended range, as explained in Run a Warm Start Hyperparameter Tuning Job.

Tip: You can reduce the cost of tuning by restricting the batch size to be multiple of 8. Refer to Nvidia's article on the significance of the number 8 when training with Automatic Mixed Precision.

from sagemaker.tuner import CategoricalParameter
hyperparameter_ranges={
                        'BATCH_SIZE': CategoricalParameter(list(range(32, 80, 8))),
                    }

This can restrict the search space to just 6 training jobs as opposed to 40!

Experimenting with Optimized TensorFlow

Compilation through SageMaker 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 increased memory utilization and a consequent decrease in the largest batch size that can fit on the GPU. In the example below we will find the new batch size with SageMaker Training Compiler enabled and the resultant latency per epoch.

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

Wait for tuning jobs to complete

The tuning jobs described above typically take around 50 mins to complete

Note: If the tuner 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 tuning job to a new tuner. For example:

native_tuner = HyperparameterTuner.attach("<your_tuning_job_name>")

Fastest Training Job

Let us collate and analyze the results from the tuning jobs. The tuner provides the results as a Pandas dataframe. We combine the results from both the tuners, sort them according to the epoch latency and display the top 5 results.

Continue tuning with SageMaker Training Compiler

Now that we have the fastest batch size and compiler configuration, we need to tune the associated hyperparameters to get the fastest convergence.

Remember Total_Training_Time ~= Latency_per_epoch * Number_of_epochs

First, we tuned to reduce the Latency_per_epoch. Now we will tune to reduce the number of epochs required for convergence. Since, hyperparameters that directly affect convergence (like learning rate, weight decay, learning schedule, etc.) are dependent on batch size, we decouple the 2 steps as described.

We now train for a higher number of epochs since we are testing the speed of convergence. Ideally, you should tune learning rate and weight decay to minimize validation loss, but for the sake of example let's minimize the training loss.

Wait for tuning jobs to complete

The tuning jobs described above typically take around 2 hours to complete

Note: If the tuner 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 tuning job to a new tuner. For example:

tuner = HyperparameterTuner.attach("<your_tuning_job_name>")

Fastest Convergence

Let us analyze the results from the tuning jobs. The tuner provides the results as a Pandas dataframe. We sort by training loss and display the top 5 results.

Having obtained the best configuration for your training job, you can now train to completion. Please consider check-pointing in order to resume training from the best performing job indicated by the tuner.

Conclusion

In conclusion, we first arrived at the batch size and compiler configuration that leads to the highest training throughput. Then, we tuned the associated hyperparameters to arrive at the configuration that leads to the fastest convergence. The resultant combinations lead to maximum savings !

Clean up

Stop all tuning 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.