Sagemaker-python-sdk: Distributed training with PyTorch container

[tothniki]

System Information

• Framework (e.g. TensorFlow) / Algorithm (e.g. KMeans):PyTorch
• Framework Version:1.2.0
• Python Version: py3
• CPU or GPU:GPU
• Are you using a custom image: No

Describe the problem

I want to use HuggingFace's run_squad.py script to fine-tune BERT on SQuAD. For that I use the PyTorch container. This script is designed to lunch by the torch.distributed.lunch and so it uses --local_rank arg for setting up distributed training.

My questions are:

1. I don't understand completely the SageMaker distributed training and the required modifications. I have to set up distributed training in my code only if I want to use multiple instances, like 2 x ml.p3.2xlarge, right?

1. So when I use one instance (e.g. ml.p3.8xlarge) which has 4 GPUs, do I have to modify somehow my training code to use all the GPUs for the training, or that is done automatically in the background?

2. Does the training with PyTorch container uses also the torch.distributed.lunch and so it can be that run_squad.py will work just fine during distributed training without modification?

3. What is the expected behaviour when I don't set up distributed training but I set 2 instances in the container? It will run the training on both instances one-by-one, or what will happen?

It would be really great if somebody could demystify these topics for me!
Thank you a lot!

[nadiaya]

Hi @tothniki,

PyTorch container provides a necessary libraries and allows you to run your training script in SageMaker similar to how you would run it outside with just minimal modifications.

In case of distributed training SageMaker provided libraries help to start user provided script on multiple instances at the same time and make sure instances can communicate with each other, but it is up to the training script provided by the user to use pytorch distributed functionality.
The same true about GPUs - it's up to the script provided by the user to choose how to use GPUs.

Here is SageMaker provided examples for distributed training:
https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/pytorch_mnist/pytorch_mnist.ipynb

And specifically an example of how distributed and single instance training can be defined in the same user script:
https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/pytorch_mnist/mnist.py#L78-L87

[laurenyu]

closing due to inactivity. feel free to reopen if necessary