dmlc / dmlc/dgl

[Performance] GSpMM regression with FP16

Open
#4,270 10 comments 0 reactions 1 assignee Claimed by @yaox12 View on GitHub
topic: system performance
Dominant language
Python
Stars
14.3k
Forks
3.1k
PR merge metrics
No merged PRs in 30d

Description

I spent some time profiling the GAT example with AMP in https://docs.dgl.ai/en/0.9.x/guide/mixed_precision.html and want to know why we didn't obtain performance gain from FP16. I observed regression in both the forward and backward phases.

## Performance on A100

| | AMP | FP32 |
|--------------|--------|--------|
| forward | 254 ms | 160 ms |
| backward | 482 ms | 369 ms |
| step | 738 ms | 529 ms |

Take the forward phase for example, there are three `EdgeSoftmax-GSpMM` pairs.
- AMP
![image](https://user-images.githubusercontent.com/3831900/179916055-e7c714f9-60ab-4577-88c5-b26e25c4d7b2.png)
- FP32
![image](https://user-images.githubusercontent.com/3831900/179916242-38875886-859d-4649-b62a-dea5a756f104.png)

Elapsed times (ms) of the `EdgeSoftmax-GSpMM` pairs x3.
| | AMP | FP32 |
|--------------|--------|--------|
| EdgeSoftmax1 | 20.896 | 22.055 |
| GSpMM1 | 73.851 | 31.890 |
| EdgeSoftmax2 | 20.865 | 22.080 |
| GSpMM2 | 73.797 | 31.882 |
| EdgeSoftmax3 | 20.869 | 22.154 |
| GSpMM3 | 27.528 | 13.422 |

We can see that the `EdgeSoftmax` is a bit faster with AMP while it's `GSpMM` that results in the regression.

Dig it a bit deeper, `GSpMM` with FP32 will go through the `cusparse::csrmm` code path while adopting `dgl::aten::cuda::SpMMCsrKernel` with FP16 inputs. I'll look into the performance issue on cusparse with FP16.

### ENV
- CUDA_VERSION: 11.7.0.022
- CUDA_DRIVER_VERSION: 515.43.04
- CUSPARSE_VERSION: 11.7.3.50
- GPU: NVIDIA A100 80GB PCIe
- PyTorch: 1.12.0a0+8a1a93a

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.