Accelerate SageMaker-PyTorch FSDP Training of Mixtral on P4 instances.

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.

In this notebook, you will learn how to accelerate distributed training of the Hugging Face Transformers Mixtral model on P4 instances.

You can either launch this notebook from an Amazon SageMaker notebook instance which handles all credentials automatically, or by running it locally and setting credentials manually.

The notebook is accompanied by the following files:

• train.py: The entry point script that'll be passed to the SageMaker PyTorch estimator later in this notebook when launching the training job.
• arguments.py: This file has functions for argument parsing (i.e. hyperparameters).
• checkpoints.py: This file has functions for saving and loading checkpoints.
• data_utils: This file has functions for handling S3 URLs.
• data: This directory has scripts for preparing and loading data.
• fsdp_utils.py: This file has util functions for fully sharded data parallelism.
• learning_rates.py: This file has functions for learning rate schedule.
• logging_utils.py: This file has functions to handle logging.
• memory_tracker.py: This file has functions to track memory usage.
• requirements.txt: This file installs the dependencies, including HuggingFace transformers.
• train_lib.py: This file has functions for running an end-to-end training of the GPT-NeoX or Llama-v2 model with SMP FSDP, settings for hybrid sharding applied, and implemented with code lines to save, load, and fine-tune the model.
• train_utils.py: This file has utility functions for training.

Additional Resources

• To learn more about launching a multi-node distributed PyTorch training job, see Launching a Distributed Training Job.
• To learn more about using the SageMaker Python SDK with PyTorch, see Using PyTorch with the SageMaker Python SDK.
• To learn more about launching a training job in Amazon SageMaker with your own training image, see Use Your Own Training Algorithms.

Prerequisites

You need to create an S3 bucket to store the input data for training. This bucket must be located in the same AWS Region that you choose to launch your training job. To learn how to create a S3 bucket, see Create your first S3 bucket in the Amazon S3 documentation.

Launching Environment

Amazon SageMaker Notebook

You can run the notebook with an Amazon SageMaker notebook instance without manually setting your aws credentials.

1. Create a new SageMaker notebook instance and open it.
2. Zip the contents of this folder & upload to the instance with the Upload button on the top-right.
3. Open a new terminal with New -> Terminal.
4. Within the terminal, enter the correct directory and unzip the file.
   1. cd SageMaker && unzip <your-zip-name-here>.zip

Locally

You can run locally by launching a Jupyter notebook server with jupyter notebook. This requires you to set your aws credentials in the environment manually. See Configure the AWS CLI for more details.

Amazon SageMaker Initialization

Run the following cell to import SageMaker modules and retrieve information of your current SageMaker work environment, such as your AWS account ID, the AWS Region, and the ARN of your Amazon SageMaker execution role. Upgrade SageMaker SDK to the latest version.

NOTE: This step might require a kernel restart.

AWS Setup

Download and Prepare the GLUE/SST2 Dataset

Here you will download, prepare the GLUE/SST2 dataset and then copy the files to S3.

Install the Hugging Face Transformers and Datasets libraries

Choose Model

Load data

This section loads the GLUE/SST2 dataset and splits it to training and validation datasets. You can update this section to load any HuggingFace dataset you want.

Load tokenizer

Nearly every NLP task begins with a tokenizer, which converts your text data into a format (token) that can be processed by the NLP model.

The following cell loads a tokenizer for GPT-NeoX-7B using AutoTokenizer.from_pretrained().

Preprocess data

Set up a function to run the tokenizer and group texts into chunks smaller than the block size.

Set additional hyperparameters and S3 paths for mapping the train and validation datasets properly depending on the phase (training or validation) of the training job in each epoch.

Specify Amazon S3 Bucket Paths

Here you need to specify the paths for training data to be used by your job. The bucket used must be in the same region as where training will run. In the cells above you downloaded the GLUE/SST2 training and validation split datasets and uploaded the json files in an S3 bucket in your account. This example will train on those json files.

After you successfully run this example expert parallel + fully sharded data parallel training job, you can modify the S3 bucket to where your own dataset is stored.

The following S3 bucket will store the output artifacts of the training job. You can modify this as needed.

Define Data Channels for SageMaker Training Using Amazon S3

In this step, define SageMaker training data channels to the S3 buckets.

(Optional) Set Up and Use Amazon FSx for Data Channels and Checkpoints

While the previous option of using Amazon S3 is easier to setup, using an FSx can be beneficial for performance when dealing with large input sizes and large model sizes and is more stable. In general, checkpointing should be done using FSx.

Please see the instructions from Distributed Training of Mask-RCNN in Amazon SageMaker Using FSx to create an FSx Lustre file system and import the dataset from the S3 bucket to your FSx file system. Note that the FSx file system must be created in a private subnet with internet gateway to ensure that training job has access to the internet. For general guidance on setting an FSx Lustre file system as data input channel, see Configure Data Input Channel to Use Amazon FSx for Lustre.

Set Hyperparameters, Metric Definitions, and MPI Options

Specify Essential Parameters for a SageMaker Training Job

Next, you use the SageMaker Estimator class to define a SageMaker Training Job, passing values through the following parameters for training job name, the number of EC2 instances, the instance type, and the size of the volume attached to the instances.

• instance_count
• instance_type
• volume_size
• base_job_name

Update the Type and Number of EC2 Instance to Use

The instance type and the number of instances you specify to the instance_type and instance_count parameters, respectively, determine the total number of GPUs (world_size).

\text{world_size = (the number of GPUs on a single instance)} \times \text{(the number of instances)}

Specify a Base Job Name

Start Training

Finally, run the estimator.fit method to launch the SageMaker training job of the model with hybrid sharding and activation offloading.

Resume from Checkpoint

You can continue a previous SageMaker training job from a saved checkpoint. Set your checkpoint directory to the appropriate checkpoint you want to resume from and pass in the checkpoint_s3_uri from a previous job.

Access the Launched SM Training Job

You can access the launched training job from SageMaker.
Go to Amazon SageMaker -> Training -> Training jobs.
You can also access the training logs from here with View Logs which opens CloudWatch directly.

Access the Training Logs

You can access the training logs from Amazon CloudWatch.

You can use CloudWatch to track SageMaker GPU and memory utilization during training and inference. To view the metrics and logs that SageMaker writes to CloudWatch, see SageMaker Jobs and Endpoint Metrics in the Amazon SageMaker Developer Guide.

If you are a new user of CloudWatch, see Getting Started with Amazon CloudWatch.

For additional information on monitoring and analyzing Amazon SageMaker training jobs, see Monitor and Analyze Training Jobs Using Metrics.

Deploy Trained Model for Inference

In most cases, a trained model can be deployed on a single device for inference because inference only requires a small amount of memory.

After you build and train your models, you can deploy them to get predictions in one of two ways:

• To set up a persistent endpoint to get predictions from your models, use SageMaker hosting services. For an overview on deploying a single model or multiple models with SageMaker hosting services, see Deploy a Model on SageMaker Hosting Services.
• To get predictions for an entire dataset, use SageMaker batch transform. For an overview on deploying a model with SageMaker Batch Transform, see Get Inferences for an Entire Dataset with Batch Transform.

To learn more about deploying models for inference using SageMaker, see Deploy Models for Inference.

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.