How to implement model parallelism using PyTorch on an HPC environment?

[Akshara211]

Hello,
I am trying to implement model parallelism using PyTorch on my HPC environment, which has 4 GPUs available. My goal is to split a neural network model across these GPUs to improve training efficiency.

Here's what I've tried so far:

Followed the PyTorch documentation on model parallelism
Implemented a basic split of the model across GPUs
However, I am encountering performance bottlenecks and underutilization of the GPUs. Can someone guide me on how to implement this in my HPC setup?

Any advice or pointers to resources would be greatly appreciated!

[qubvel]

Hi @Akshara211, I'm not familiar with that topic that much, but I'm curious what is the motivation of doing it? What model configuration are you using?

[Akshara211]

Hi @qubvel ,
I am currently working on a computer vision project involving a dataset with 32,000 images and their corresponding annotations. While using ResNet-50, I have been able to set the batch size up to 32 without any issues. However, I am encountering difficulties when attempting to train a deeper model, ResNet-152. When using ResNet-152, I am only able to set a batch size of 8. Any attempt to increase the batch size results in a memory error.

Suggestions from various sources recommend implementing model parallelism by distributing the model across different GPUs. How are you training the ResNet-152 model, and what GPU specifications are you using for ResNet-152? I have 4 GPUs available.

[qubvel]

I would probably recommend using

• lower precision. e.g. float16/bfloat16
• gradient accumulation
• distributed data parallel (for multi-gpu training)

all these things are really easy to incorporate with Pytorch Lightning, just with a few flags provided to the Trainer class

[Akshara211]

Thank you for the suggestions!

We are already using Distributed Data Parallel (DDP), so data parallelism is applied. However, we are still facing issues with ResNet152. Our HPC GPU specification is 4x NVIDIA A100-SXM4-40GB. Could you please advise on the specific flags to use for enabling lower precision (float16/bfloat16) and gradient accumulation in PyTorch Lightning?

Is there anything else I could do? Any other suggestions from your side would be appreciated.

Thank you!

[qubvel]

Lower precision training
https://lightning.ai/docs/pytorch/stable/common/precision_intermediate.html

Gradient accumulation
https://lightning.ai/docs/pytorch/stable/advanced/training_tricks.html

[plo97]

How do I implement distributed data parallelism (for multi-GPU training)?

[qubvel]

@plo97 You can use Pytorch-Lightning for that too
https://lightning.ai/docs/pytorch/stable/accelerators/gpu_intermediate.html

All you need is:

# train on 8 GPUs (same machine (ie: node))
trainer = Trainer(accelerator="gpu", devices=8, strategy="ddp")

See how to train SMP model with PyTorch-Lightning here:
https://github.com/qubvel-org/segmentation_models.pytorch/blob/main/examples/binary_segmentation_intro.ipynb

[Akshara211]

Hi @qubvel
When using lower precision (bfloat16) and gradient accumulation, will it take more time to complete one epoch compared to normal training? In my case, the training already took 14 hours and only 7 epochs have been completed.

Additionally, will using these techniques cause any change in accuracy compared to normal training? Here is my trainer configuration:

trainer = pl.Trainer(accelerator='gpu', max_epochs=epochs,callbacks=[checkpoint_callback, early_stopping_callback], devices=4, strategy="ddp",precision="bf16-mixed",accumulate_grad_batches= 4)

[qubvel]

Hi @Akshara211 sorry for the late response, I probably missed the notification, were you able to solve your problem?

When using lower precision (bfloat16) and gradient accumulation, will it take more time to complete one epoch compared to normal training?

Regarding bfloat16, it depends on your GPU, for some GPUs you will find it faster, while for others it might be slower. Gradient accumulation should not slow down your training.

Additionally, will using these techniques cause any change in accuracy compared to normal training?

It depends on the careful setup, but in general, it should not change the accuracy. However, you might want to set up sync batchnorm in case you have a small batch size per one GPU (alternatively you can freeze batchnorms).

[Patataman]

As far as I know, no library provide automatic tools for model parallelism (FastAI, PyTorch, TensorFlow...), it's up to you to divide and send the appropriate layers to each GPU and synchronize.

For the specific case of smp, I have been able to do both, Data and Model parallelism using the Unet with resnet34 as encoder.
In my case I had 3 nodes with 1 GPU on each (3 GPUs in total). For DP I used PyTorch's DistributedDataParallelism and for MP I used PyTorch's RPC module.

I was also able to test it with 2 GPUs on the same machine, but briefly.

Both launched using torchrun.

About the execution time you mentioned, when doing DP training time should be reduced, for MP it should be similar (if your GPUs are in the same machine) or higher if (as my case) you have a distributed environment due to the overhead

[qubvel]

Hi @Patataman!

Thanks a lot for sharing your experience. In case you have time, is it possible to share any code examples how this can be implemented? I would appreciate any details and contributions to make this question more clear for the community, thanks!

[Patataman]

I think I can arrange a minimal example for DP and MP (this one when the GPUs are on the same node) . On the other hand, distributed MP require more changes and knowledge about the library (in this case RPC), but I can share an example with pseudo code that should help to understand how to do it

[qubvel]

Sounds great!