Use the SMDDP library in your PyTorch training script

Starting from the SageMaker distributed data parallelism (SMDDP) library v1.4.0, you can use the library as a backend option for the PyTorch distributed package. To use the SMDDP AllReduce and AllGather collective operations, you only need to import the SMDDP library at the beginning of your training script and set SMDDP as the the backend of PyTorch distributed modules during process group initialization. With the single line of backend specification, you can keep all the native PyTorch distributed modules and the entire training script unchanged. The following code snippets show how to use the SMDDP library as the backend of PyTorch-based distributed training packages: PyTorch distributed data parallel (DDP), PyTorch fully sharded data parallelism (FSDP), DeepSpeed, and Megatron-DeepSpeed.

For PyTorch DDP or FSDP

Initialize the process group as follows.

import torch.distributed as dist
import smdistributed.dataparallel.torch.torch_smddp

dist.init_process_group(backend="smddp")

For DeepSpeed or Megatron-DeepSpeed

Initialize the process group as follows.

import deepspeed
import smdistributed.dataparallel.torch.torch_smddp

deepspeed.init_distributed(dist_backend="smddp")

export SMDATAPARALLEL_OPTIMIZE_SDP=true

For general guidance on writing a PyTorch FSDP training script, see Advanced Model Training with Fully Sharded Data Parallel (FSDP) in the PyTorch documentation.

For general guidance on writing a PyTorch DDP training script, see Getting started with distributed data parallel in the PyTorch documentation.

After you have completed adapting your training script, proceed to Step 2: Launch a distributed training job using the SageMaker Python SDK.