Trivial Multi-Node Training With Pytorch-Lightning | by William Falcon | Towards Data Science

So you have this awesome HPC cluster but still train your model on only 1 GPU?
I know. The struggle is real. 😭
Pytorch-lightning, the Pytorch Keras for AI researchers, makes this trivial. In this guide I’ll cover:
Running a single model on multiple-GPUs on the same machine.Running a single model on multiple machines with multiple GPUs.
Running a single model on multiple-GPUs on the same machine.
Running a single model on multiple machines with multiple GPUs.
Disclaimer: This tutorial assumes your cluster is managed by SLURM.
Let’s first define a PyTorch-Lightning (PTL) model. This will be the simple MNIST example from the PTL docs.
Notice that this model has NOTHING specific about GPUs, .cuda or anything like that. The PTL workflow is to define an arbitrarily complex model and PTL will run it on whatever GPUs you specify.
Let’s train our CoolModel on the CPU alone to see how it’s done.
To train on a single GPU simply pass in the GPU id
And for multiple-GPUs, just add more ids!
So, the full script you’d run on your fancy cluster GPU machine is:
Let’s take a quick detour to the land of NVIDIA GPU flags. When you set the gpus [0,1] in the trainer, you’re saying “use GPU 0 and GPU 1 on this machine.” The order corresponds to the output of your NVIDIA-SMI command
Under the hood, PTL does the following to get your model working on GPUs
Sets the CUDA_VISIBLE_DEVICES to those ids and CUDA_DEVICE_ORDER = “PCI_BUS_ID” which guarantees the ids will correspond to the physical GPU layout on the machine (ie: output of NVIDIA-SMI).Adds the appropriate DistributedDataParallel or DataParallel wrapper on your model.
Sets the CUDA_VISIBLE_DEVICES to those ids and CUDA_DEVICE_ORDER = “PCI_BUS_ID” which guarantees the ids will correspond to the physical GPU layout on the machine (ie: output of NVIDIA-SMI).
Adds the appropriate DistributedDataParallel or DataParallel wrapper on your model.
SINGLE NODE SLURM
Let’s say you submit a SLURM job with 2 GPUs. You should still pass GPU ids 0,1 to Trainer even though you might actually be using the physical last 2 GPUs on an 8-GPU node. That’s because SLURM scopes out the correct GPUs for your job which to you, start indexed at 0.
Now, if you’re in an interactive job for instance, and the node has 8 GPUs, then you’ll want to change the indices accordingly.
For example, you could run this same model 4 times on a single 8-GPU node by launching the script 4 times with different GPU ids, or running 4 processes as shown below:
Of course, you’d probably only want to do this for debugging purposes or because you’re just ssh-ing directly into a single machine.
Test-tube makes running a grid-search on a single node trivial. Here’s a grid search example when you’re not submitting through SLURM.
To train the PTL model across multiple-nodes just set the number of nodes in the trainer:
If you create the appropriate SLURM submit script and run this file, your model will train on 80 GPUs.
Remember, the original model you coded IS STILL THE SAME. The underlying model has no knowledge of the distributed complexity.
That’s one of the strongest points of using PTL, defer the engineering to the framework so you can focus on research.
An example SLURM script will look like:
# SLURM SUBMIT SCRIPT#SBATCH --gres=gpu:8#SBATCH --nodes=10#SBATCH --ntasks-per-node=8#SBATCH --mem=0#SBATCH --time=02:00:00# activate conda envconda activate my_env  # run script from abovepython my_test_script_above.py
There are other flags which might be specific to your cluster, such as:
partitionconstraint
partition
constraint
In addition, PTL only supports NCCL backend at the moment. This is the recommended backend by the PyTorch team and the one with the fastest library. If a new, faster one comes out, PTL will add that option as soon as it’s available.
Writing SLURM scripts can be a PITA. This is especially true if you want to run a grid search.
For that, you can use the SlurmCluster object with PTL which will:
Auto-submit scripts for you.Split off a grid-search into individual scripts and submit them all.Requeue your job a few minutes before the wall time.PTL will also save and load your model to make the re-queue seamless (you don’t have to do anything to save or load the model, PTL will do that automatically).
Auto-submit scripts for you.
Split off a grid-search into individual scripts and submit them all.
Requeue your job a few minutes before the wall time.
PTL will also save and load your model to make the re-queue seamless (you don’t have to do anything to save or load the model, PTL will do that automatically).
Let’s adapt the script above to use the SlurmCluster example.
First, let’s configure the SLURM job using the SlurmManager
When the optimizer_parallel_cluster_gpu gets called, the SlurmManager will submit a script for each set of hyperparameters passed in via the HyperOptArgumentParser.
So, the full updated script now looks like:
To run this grid search on the cluster simply:
ssh into your login nodeActivate your conda env with lightning installedRUN the python script above
ssh into your login node
Activate your conda env with lightning installed
RUN the python script above
ssh some_nodeconda activate my_env_with_ptl# run the above scriptpython above_script.py
This may not be the workflow you’re used to, but when you run the script, it will ONLY submit each slurm job with a set of hyperparameters. Then when the script runs again, it will actually run the function you passed into optimize_parallel_cluster_gpu.
SLURM will save all out, err and submit script files to the directory you passed into the SlurmCluster object.
The actual training logs will be written by the Experiment object. Make sure to set the experiment version to the cluster.hpc_exp_number which has the correct experiment version in the overall scheme of the cluster jobs.
Distributed training on multiple nodes, unfortunately, requires a bit more work because of the underlying SLURM system. PTL abstracts as much out as possible and gives you an optional way of reducing the cluster complexity using the SlurmCluster object.