PyTorch aims to make code non-intrusive in that you can take existing code which trains a neural net on a single processer an easily scale out when hitting resource limitations with minimal code changes. With this spirit in mind, let’s first point out all the code that can be reused. Below are a set of functions and import statements that we will not need to modify that were created previously.

Next, let’s highlight the modifications that need to be made to scale our DesignSafe Classifier code in the following order. As a reminder below are the 8 modifications we highlighted in the this and the previous tutorial.

Serial to Parallel Modifications

1. Create a process group
2. Use PyTorch’s DistributedSampler to ensure that data passed to each GPU is different
3. Wrap Model with PyTorch’s DistributedDataParallel
4. Modify Training Loop to write model from one GPU
5. Close process group

Torchrun Modifications

6. Modify code for environment variables set by torchrun:
   1. Remove code that sets environment variables as this done for you automatically with torchrun.
   2. Instead, use these environment variables set by pytorch and instead of explicitly defining them.
7. Add code for writing checkpoints and resuming training after failure
   1. Create location to store checkpoints
   2. Read checkpoints if they exist and resume training at epoch checkpoint was written
   3. Write checkpoints periodically during training
8. Remove using the mp.spawn to parallelize code and replace this with a function call, as this is done automatically by torchrun