Scale deep learning with 2 new libraries for distributed training on Amazon SageMaker | by Emily Webber | Towards Data Science

By Olivier Cruchant, Emily Webber, and the AWS team
The clouds of Hawaii. Beautiful reminders of peace, serenity, and potentially experiences awaiting us as we exit the global pandemic of 2020 and move back into a world allowing global travel. But until then, why not spend some time learning about how Amazon SageMaker’s new data parallelism and model parallelism libraries can dramatically improve our model performance on distributed training for deep learning!   As a machine learning practitioner, you might find yourself in the following situations. You might have found just the perfect state of the art transformer-based model, only to find that when you try to fine-tune it you run into memory issues. You might have just added billions of parameters to your model, which should improve your model performance, but this too only gets you an out of memory issue. You might be able to comfortably fit a model on a single GPU, but are struggling to take advantage of all your GPU’s and find that your model still takes days to train.  Should you just accept the status quo and limit your applications to models that fit within the existing hardware capacity or that train within an acceptable time? Of course not! Enter model parallelism and data parallelism on Amazon SageMaker. While customers have always been able to take advantage of SageMaker’s built-in capabilities for distributed training, such as multi-node jobs, the parameter server and support for open source frameworks, we are excited to announce two new toolkits, model parallelism and data parallelism, that allow ML developers to scale distributed training by fitting larger models into RAM and speeding up their training jobs. In this post we first describe the theory behind those two techniques and then explain how to use them in Amazon SageMaker. Right off the bat we’re seeing impressive results for PyTorch and TensorFlow. Using SageMaker model parallelism, we were able to train models much bigger than a single GPU’s memory. Using SageMaker data parallelism, we witnessed throughput improvements vs PyTorch DistributedDataParallel between 15% and 41% on BERT-Large training (cluster size 16–64 GPUs), and between 4% and 13% on MaskRCNN training (cluster size 16–64 GPUs).
This means two things. First, you can train bigger models. Second, you can do this faster than ever. Part of these improvements come from our taking full advantage of the AWS network infrastructure and EC2 instance topology. We’re introducing these to developers in less than 10 lines of code.
In this post, we’ll explain how to use SageMaker distributed training toolkits with PyTorch, and the rationale is similar with TensorFlow. You can step through our examples on GitHub right here. If you like, you can dive into the full training guides right here for TensorFlow and PyTorch.
If you’re unfamiliar with mini-batch stochastic gradient descent (SGD), we’ll quickly recap that here, but you’d be better off getting caught up on the nuts and bolts of it elsewhere such as the Optimization chapter (https://d2l.ai/chapter_optimization/index.html) of Dive Into Deep Learning (D2L.ai) co-authored by Amazon scientists Zhang, Lipton, Li and Smola.   Simply put, we train neural networks by updating the parameters of the model in a direction that reduces its error. These changes happen iteratively, conducted on equally-sized sub-samples of the training dataset called mini-batches. We get predictions from the model per record for each of the mini-batches, measure the error, compute the partial derivative of the loss function with respect to the model inputs, and update the parameters. One pass over the training dataset is called an epoch, and full model trainings commonly consist of dozens to hundreds of epochs.   As it turns out, mini-batch SGD by itself works really nicely in a parallel world. Passing the records of a single mini-batch over different computing devices is called data parallel distributed training. Data-parallel training scales the mini-batch size and allows even faster mini-batch processing. However, data parallel training comes with the extra complexity of having to compute the mini-batch gradient average with gradients coming from all the workers and communicating it to all the workers, a step called AllReduce.   The problem with generic AllReduce algorithms is that they tend to slow down the training process as datasets grow. This can dramatically penalize training time if not properly implemented.
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 scriptWrap your model with SDP’s DistributedDataParallel methodUse 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.DistributedSamplerWhich 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.
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 initializeimport smdistributed.dataparallel.torch.distributed as distfrom smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDPdist.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 PyTorchjob = 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)!
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.
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 scriptWrap your model with SMP’s DistributedModel methodPut the forward and backward logic in a function and decorate it with .smp.stepReplace 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.
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 smpsmp.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 lossfor 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 PyTorchmpi_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.
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