Distributed model parallel

Amazon SageMaker Distributed Model Parallel (SMP) is a model parallelism library for training large deep learning models that were previously difficult to train due to GPU memory limitations. SMP automatically and efficiently splits a model across multiple GPUs and instances and coordinates model training, allowing you to increase prediction accuracy by creating larger models with more parameters.

You can use SMP to automatically partition your existing TensorFlow and PyTorch workloads across multiple GPUs with minimal code changes. The SMP API can be accessed through the Amazon SageMaker SDK.

Use the following sections to learn more about the model parallelism and the SMP library.

Important

SMP only supports training jobs using CUDA 11. When you define a PyTorch or TensorFlow Estimator with modelparallel parameter enabled set to True, it uses CUDA 11. When you extend or customize your own training image you must use a CUDA 11 base image. See Extend or Adapt A Docker Container that Contains SMP for more information.

It is recommended to use this documentation alongside SageMaker Distributed Model Parallel in the Amazon SageMaker developer guide. This developer guide documentation includes:

How to Use this Guide

The SMP library contains a Common API that is shared across frameworks, as well as APIs that are specific to supported frameworks, TensorFlow and PyTroch. To use SMP, reference the Common API documentation alongside framework specific API documentation.