Scale deep learning with 2 new libraries for distributed training on Amazon SageMaker
By: Emily Webber
Tags: AWS, Deep Learning, Machine Learning, Distributed Systems, Sagemaker

SageMaker data parallelism

Using a smart parameter-server-based AllReduce, SageMaker data parallelism (SDP) trains faster than PyTorch DistributedDataParallel and Horovod in a number of specific NLP and vision tasks while being compatible with their popular APIs. SDP communicates parameters via equally-sized Balanced Fusion Buffers (BFB), which contribute to a better use of available bandwidth than pure parameter-based communication. SDP requires minor edits to be used on existing compatible code, and you’ll get faster distributed jobs as a result.



Using SageMaker data parallelism in PyTorch



At launch SDP supports both PyTorch and TensorFlow, but here we’ll step through a PyTorch example. It’s actually pretty accessible! First you’ll import the SDK into your script. Then you’re using the SDK to tell your data loader which node it’s on, and the rank of that node. You’ll load the model itself into the SDP framework. Finally, you’ll also want to call out in the estimator for your job that you are going to use this method of distribution.

Import and initialize the SDP client in your script Wrap your model with SDP’s DistributedDataParallel method Use SDP’s methods to manage which nodes in your cluster are handling certain elements: Which chunk of data is used by each worker — this can be controlled with torch.utils.data.distributed.DistributedSampler Which worker saves the model — most of the time we’ll just save this on the 0th node. If need be, scale the batch size and adapt hyperparameters to fully occupy the training cluster while maintaining convergence.

The below snippet of code summarizes edits in a model-agnostic template. For simplicity we’ll just focus on the main training function itself.

# SDP: Import SDP PyTorch API, DDP and initialize

import smdistributed.dataparallel.torch.distributed as dist

from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP



dist.init_process_group()



def main():



# Prepare dataset

train_dataset = torchvision.datasets.MNIST(...)



# SDP: Set num_replicas and rank in DistributedSampler

train_sampler = torch.utils.data.distributed.DistributedSampler(

train_dataset,

num_replicas=dist.get_world_size(),

rank=dist. get_rank() )



train_loader = torch.utils.data.DataLoader(train_sampler=..., ...)



# SDP: Wrap the PyTorch model with SDP’s DDP

model = DDP(Net().to(device))



# SDP: Pin each GPU to a single SDP process.

torch.cuda.set_device(local_rank)

model.cuda(local_rank)



# Train

optimizer = optim.Adadelta(...) # use your own optimizer and hyperparameters

scheduler = StepLR(...)

for epoch in range(1, args.epochs + 1):

train(...)

if rank == 0:

test(...)

scheduler.step()



# SDP: Save model on master node.

if dist.get_rank() == 0:

torch.save(...)

In order to launch the distributed training in SageMaker, you can execute the job using the SageMaker Python SDK:

from sagemaker.pytorch import PyTorch



job = PyTorch(entry_point="my_sdp_script.py",

role=role, # IAM role for the training cluster

framework_version='1.6.0',

py_version="py36",

train_instance_count=2,

train_instance_type="ml.p3.16xlarge",

distributions={"smdistributed": {"dataparallel": {"enabled": True}}} # SDP distribution method enabled here

)



job.fit({"mychannel":"s3://bucket/path/to/training/data"})

# In this case the s3 data will by copied locally at opt/ml/input/data/mychannel

Now that we’ve learned how to use SDP, let’s check out SageMaker model parallelism (SMP)!

SageMaker model parallelism

Current data parallelism tools such as PyTorch DistributedDataParallel and Uber’s Horovod still require developers to fit at least the model and a single record within a computing worker. These current methods of distribution block innovation within deep learning, because recent results have shown that larger models tend to be associated with more accurate results.



When training very large models — such as large transformers in NLP or segmentation models over high-resolution images — this may not be feasible. An alternative way to fraction the workload to accommodate models is to partition the model over multiple computing devices. This is called model parallel distributed training.



The primary benefit of model parallel training is the ability to train models much larger than the memory of a single GPU. Model parallel distributed training has lighter device-to-device communication requirements and uses memory more frugally; it does not create one model copy per worker like with data parallelism.

Historically, implementing model parallel training on your own has been prohibitively difficult. Even for experts in deep learning, naive implementations are filled with custom code to manage the distributed computation and inter-worker communication. These end up using GPU unevenly, or even fail to truly parallelize the computation.

SageMaker model parallelism (SMP) handles model parallelism for you in a few lines of code. It includes a model graph partitioning algorithm that optimizes for speed or memory, and supports manual partitioning. It trains in a pipelined fashion over micro-batches to maximize GPU usage even under the model-parallel design. With SMP, we were able to fit an 20-record batch of BERT-MEGA on the 8 V100 GPUs of a single p3dn.24xlarge instance, and we were able to fit a 16-record batch of GPT-2 XL on the 4 V100 GPUs of a single p3.8xlarge instance, two achievements difficult to complete without model parallelism.

Using SageMaker model parallelism

In order to distribute your PyTorch model training code with SMP, you’ll need to do the following steps. Most of these are similar to the steps for the data parallel toolkit above, but we’re going to add decorators over train and test steps in the model script itself.

Import and initialize the SMP client in your script Wrap your model with SMP’s DistributedModel method Put the forward and backward logic in a function and decorate it with .smp.step Replace torch.Tensor.backward / torch.autograd.backward with DistributedModelObj.backward . Post-processing on the outputs across microbatches using StepOutput methods such as reduce_mean . You also need to write or edit model-agnostic distribution management code if necessary, for example which worker saves the model, and any changes to your batch size and hyperparameters to fully occupy the training cluster while maintaining convergence.

The below snippet of code summarizes typical necessary edits in a fictional PyTorch training script:

# SMP: Import and initialize SMP API.

import smdistributed.modelparallel.torch as smp



smp.init()



model = load_my_model()



# Download dataset and create dataloader.

train_loader = load_my_data()



# SMP: Instantiate DistributedModel object using the model.

# This handles distributing the model among multiple ranks

# behind the scenes

# If horovod is enabled this will do an overlapping_all_reduce by

# default.

model = smp.DistributedModel(model)



# SMP: Define the smp.step consisting of forward and backward passes.

@smp.step()

def forward_backward(model, inputs, labels):

loss = my_loss_function(model(inputs), labels)

# SMP: Call backward on the model instead of the output tensor.

# If loss is not a scalar or a 0d tensor, the backward call requires

# out grad tensors in addition to the output tensors,

# similar to torch.autograd.backward call.

model.backward(loss)

return loss



for input, label in train_loader:

image, label = image.to(device), label.to(device)



optimizer.zero_grad()

loss_mb = forward_backward(model, inputs, labels)



# SMP: Average the loss across microbatches.

loss = loss_mb.reduce_mean()

optimizer.step()



# SMP: Print the loss only at rank 0.

if smp.rank() == 0:

print(f"Loss: {loss}")

In order to launch the distributed execution of this training script, you can use the high-level SageMaker SDK:

from sagemaker.pytorch.estimator import PyTorch



mpi_options = {

"enabled": True,

"processes_per_host": 4,

"custom_mpi_options":"--mca btl_vader_single_copy_mechanism none "

}



dist_options = {

"modelparallel":{

"enabled": True,

"parameters": {

"partitions": 4, # we'll partition the model among the 4 GPUs

"microbatches": 8, # Mini-batchs are split in micro-batch to increase parallelism

"optimize": "memory" # The automatic model partitioning can optimize speed or memory

}

}

}





job = PyTorch(

entry_point="my_script.py",

role=role, # IAM role for the training cluster

instance_type="ml.p3.8xlarge",

instance_count=1,

framework_version=1.6,

py_version="py36",

distribution={"mpi": mpi_options, "smdistributed": dist_options}

)



job.fit() # Can optionally use fit(wait=False) for asynchronous execution.

We enable MPI because SMP uses MPI as a communication library under the hood. The number of processes_per_host must be set up to the number of GPUs available in your chosen instance type. Finally, the custom MPI flag "--mca btl_vader_single_copy_mechanism none" helps suppress certain warnings that show up when MPI is used with SageMaker containers, and keeps your training logs clean.

Closing Thoughts on new launches for SageMaker distributed training

Distributed training means you can train on bigger batches and bigger models. One thing to keep in mind is that increases in batch sizes generally impact SGD convergence; you may need to adjust the learning rate or other hyperparameters to maintain convergence quality.

SageMaker data parallelism improves the development experience and the performance of multi-machine data parallel training — this means you’re not going to see those speed improvements unless you’re truly using multiple instances in your job.

At launch we expect customers who need to use both methods of parallelism for the same jobs to use open source data distribution methods while

You can cruise over to the toolkit documentation for a deeper dive. Don’t hesitate to take a look, and let us know what you build!

All images by the authors