Lightning-AI / Lightning-AI/lightning-thunder

Support RN50 BatchNorm fusions with cudnn

Open
#487 1 comment 2 reactions 1 assignee View on GitHub

@vedaanta is already working on this.

Since Jun 3, 2024.

cudnn enhancement
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature
CuDNN can support RN50 batchnorm fusions (and their symmetric backward counterparts) like:
- BatchNorm + Relu
- BatchNorm + Add + Relu

Above fusions also support fp8 quantization around them.

### Motivation
RN50 benchmark that is being added in #443, can benefit from cudnn fusions at:
```
# BN-relu fusion
self.bn1 = norm_layer(planes)
self.relu = nn.ReLU(inplace=True)

# BN-Add-relu fusion
out = self.bn3(out)
...
out = out + identity
out = self.relu(out)

```

### Pitch
Add support for BN fusions via cudnn executor.

CC @IvanYashchuk @kshitij12345 @Anerudhan

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.