pytorch-lightning: bfloat16 running 4x slower than fp32 (conv)

I ran the imagenet example using TORCH_CUDNN_V8_API_ENABLED=1 with PyTorch 1.12 as suggested in the forum post linked above, and recorded new values. The benchmark was run on a single A100 for 2500 iterations.

This fixes the performance regression I was seeing with imagenet bfloat16, and probably will with conformer as well!

Till CUDNN v8 becomes standard, I suggest people use the env variable to see the expected performance. Still slightly slower than FP16 but that might just be down to the underlying cudnn optimizations.

Pure PyTorch imagenet benchmark

I made 4 different versions of the script, using AMP+scaler/Bfloat16 calls

PyTorch Lightning Imagenet benchmark

Obviously throughput can be much better as we’re not compute-bound at all, but shows performance regression is considerably less.

Closing this issue for now. Can be re-opened if TORCH_CUDNN_V8_API_ENABLED=1 with PyTorch 1.12 does not fix the issue.

Here are my results running a pytorch-lightning minGPT script (here). This was a transformer model, with the below cmd:

python min_gpt.py --accelerator gpu --devices 1 --precision X

I also used torch.set_float32_matmul_precision('medium') which gave roughly the same results.

We can see that the performance is roughly the same when using AMP, but when converting the model we get a small speedup + a significant memory drop as expected.

It seems for transformer architectures BFloat16 works as expected.

I found this forum post and this comment describing slowdowns when using Conv2d + Bfloat16.

The model type in OP’s post is based on convs, and the above reference issue could explain the reason why we see a 4x slowdown.

Next steps are to re-run the imagenet comparison, and optionally the xformers example for @blefaudeux.

Hi, so I was finally able to find some time to benchmark for possible precisions (apologies for the delay).

Note:

1. Tested on: Image Net dataset
2. Scripts used:
   • PL: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/domain_templates/imagenet.py
   • PyTorch: https://github.com/pytorch/examples/blob/main/imagenet/main.py (used AMP for float16 and bf16 precisions)
3. Tested on ~2500 iterations. Numbers are averaged out.
4. Environment details:
   • PyTorch: 1.11.0
   • PyTorch Lightning: 1.6.0
   • CUDA Version: 11.3 (with PyTorch); (NVIDIA GeForce RTX 3090)
   • DDP: No (trained on a single device/gpu)

Analysis:

• From the numbers, it looks like PyTorch Lightning is around 1.6 times slower than PyTorch for bfloat16 precision, while for other precisions - there doesn’t seem to be a huge difference. PL performs a little faster, but I assume it will be near to equal if tested for >1 epochs.
• Also, I agree that bf16 is running slower than fp32 in PyTorch Lightning, which is not expected.

We’ll try investigating why performance is slower for bf16 in PyTorch Lightning and will get back to this issue once we have some numbers. 😃 Thank you for your patience and for reporting this issue to us!

cc: @carmocca @SeanNaren @Borda for visibility! 😃

ping @SeanNaren @ananthsub, any news on that ?

ok, synced offline and seems that it’s not a priority 😦

I’m going to revisit this. I do wonder if it’s a PyTorch thing or a PyTorch Lightning thing. If anyone has any related issues they could point me to would help! If not, going to redo the benchmarks with PyTorch 1.12, and try a different model.

ping @SeanNaren @ananthsub, any news on that ?

ok, synced offline and seems that it’s not a priority 😦

That’s great, thanks so much @SeanNaren ! The memory benefits are huge, for a similar speed it’s still very useful. Cc @matthieutphr @quentindrx

Hello, I meet the same problem with PL 1.6.4 on A100 with cuda 11.6 pytorch 1.12.0

Working with @krshrimali to try to investigate!

I’m excluding Lightning and using the pytorch imagenet example to gather the below results on an A100. I ran for more than 2500 iterations.

I also tried the latest NGC container which improved results marginally but showed a massive slowdown for BFloat16. I also tried resnet101 and still showed a slowdown.

I also tried adding torch.cuda.synchronize() before measurements were taken, and a negligible difference was made.

I think it will be worth checking other model types with a minimal script to see if we see the same slow down without PyTorch Lightning before investigating.

Note: no TF32 meant adding this:

# The flag below controls whether to allow TF32 on matmul. This flag defaults to True.
torch.backends.cuda.matmul.allow_tf32 = False

# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = False