Amazon SageMaker Distributed Training Notebook Examples

The following case studies and notebooks provide examples of implementing the SageMaker distributed training libraries for the supported deep learning frameworks (PyTorch, TensorFlow, and HuggingFace) and models, such as CNN and MaskRCNN for vision, and BERT for natural language processing.

These notebooks are provided in the SageMaker examples GitHub repository. You can also browse them on the SageMaker examples website.

Blogs and Case Studies

The following blogs discuss case studies about using the SageMaker distributed training libraries.

The SageMaker data parallelism library

• Enable faster training with Amazon SageMaker data parallel library, AWS Machine Learning Blog (December 05, 2023)

• How I trained 10TB for Stable Diffusion on SageMaker in Medium (November 29, 2022)

• Run PyTorch Lightning and native PyTorch DDP on Amazon SageMaker Training, featuring Amazon Search , AWS Machine Learning Blog (August 18, 2022)

• Training YOLOv5 on AWS with PyTorch and the SageMaker distributed data parallel library, Medium (May 6, 2022)

• Speed up EfficientNet model training on SageMaker with PyTorch and the SageMaker distributed data parallel library, Medium (March 21, 2022)

• Speed up EfficientNet training on AWS with the SageMaker distributed data parallel library, Towards Data Science (January 12, 2022)

• Hyundai reduces ML model training time for autonomous driving models using Amazon SageMaker, AWS Machine Learning Blog (June 25, 2021)

• Distributed Training: Train BART/T5 for Summarization using Transformers and Amazon SageMaker, the Hugging Face website (April 8, 2021)

The SageMaker model parallelism library

• New performance improvements in the Amazon SageMaker model parallelism library, AWS Machine Learning Blog (December 16, 2022)

• Train gigantic models with near-linear scaling using sharded data parallelism on Amazon SageMaker, AWS Machine Learning Blog (October 31, 2022)

PyTorch Examples

The SageMaker data parallelism library

• CNN with PyTorch and the SageMaker data parallelism library

• BERT with PyTorch and the SageMaker data parallelism library

The SageMaker model parallelism library

• Train GPT-2 with near-linear scaling using the sharded data parallelism technique in the SageMaker model parallelism library

• Fine-tune GPT-2 with near-linear scaling using sharded data parallelism technique in the SageMaker model parallelism library

• Train GPT-NeoX-20B with near-linear scaling using the sharded data parallelism technique in the SageMaker model parallelism library

• Train GPT-J 6B using the sharded data parallelism and tensor parallelism techniques in the SageMaker model parallelism library

• Train FLAN-T5 with near-linear scaling using sharded data parallelism technique in the SageMaker model parallelism library

• Train Falcon with near-linear scaling using sharded data parallelism technique in the SageMaker model parallelism library

TensorFlow Examples

The SageMaker data parallelism library

• CNN with TensorFlow 2.3.1 and the SageMaker data parallelism library

• BERT with TensorFlow 2.3.1 and the SageMaker data parallelism library

The SageMaker model parallelism library

• CNN with TensorFlow 2.3.1 and the SageMaker model parallelism library

HuggingFace Examples

The following HuggingFace on SageMaker examples are available in the HuggingFace notebooks repository.

The SageMaker data parallelism library

• HuggingFace Distributed Data Parallel Training in PyTorch on SageMaker - Distributed Question Answering

• HuggingFace Distributed Data Parallel Training in PyTorch on SageMaker - Distributed Text Summarization

• HuggingFace Distributed Data Parallel Training in TensorFlow on SageMaker

The SageMaker model parallelism library

• HuggingFace with TensorFlow Distributed model parallelism library Training on SageMaker

How to Access or Download the SageMaker Distributed Training Notebook Examples

Follow instructions to access or download the SageMaker distributed training example notebooks.

Option 1: Use a SageMaker notebook instance

To use the aforementioned examples, we recommend that you use an Amazon SageMaker notebook instance. A notebook instance runs Jupyter Notebook and JupyterServer apps on Amazon EC2 instances, which are optimized for machine learning. If you do not have an active notebook instance, follow the instructions in Create a Notebook Instance in the SageMaker developer guide to create one.

After you have created an instance, in the Notebook instances page of the SageMaker console, do the following:

1. Open JupyterLab.

2. Select the examples icon ( ) in the left tray.

3. Browse the examples for Training and look for notebooks titled Distributed Data Parallel or Distributed Model Parallel.

Option 2: Clone the SageMaker example repository to SageMaker Studio or notebook instance

To download and use the aforementioned example notebooks, do the following to clone the example GitHub repositories:

1. Open a terminal.

2. In the command line, navigate to the SageMaker folder.

   cd SageMaker

3. Clone the SageMaker examples GitHub repository.

   git clone https://github.com/aws/amazon-sagemaker-examples.git

   Note

   To download the HuggingFace example notebooks, clone the HuggingFace notebooks GitHub repository:

   git clone https://github.com/huggingface/notebooks huggingface-notebooks

4. In the JupyterLab interface, navigate into the amazon-sagemaker-examples folder.

5. In the training/distributed_training folder, there are folders for frameworks, and in each of these, there are folders for data_parallel and model_parallel. Choose the example of your choice and follow the instructions to launch distributed training with an SageMaker distributed training library.