Repository: pytorch
Issue: aten::cudnn_convolution chooses different conv implementation given the same inputs. 
Labels: module: cudnn,module: cuda,module: convolution,triaged

## Issue description

I'm using PyTorch DDP, DeepSpeed ZeRO, and PyTorch FSDP to train my model that contains several `nn.Conv1d` layers. `Conv1d.forward` eventually calls `aten::cudnn_convolution` to do the actual computation. I noticed that when using DDP, `aten::cudnn_convolution` launches one CUDA kernel (`cudnn::cnn::conv1D_NCHW_general`) and is quite fast (0.1ms). However, when using ZeRO or FSDP, `aten::cudnn_convolution` launches multiple smaller CUDA kernels (`implicit_convolve_sgemm`) and takes much longer time (4ms).  

The model weights and input data have exactly same shapes (and values) but cudnn behaves differently with big performance difference. 

PyTorch DDP
![image](https://user-images.githubusercontent.com/24818503/218893716-b8ae2c4c-2e9d-4de9-a71c-79d85579cd21.png)
![image](https://user-images.githubusercontent.com/24818503/218893734-cd49eeaf-2c98-41bd-a7f6-b1e6009c74d0.png)

PyTorch FSDP NO_SHARD
![image](https://user-images.githubusercontent.com/24818503/218893761-930d0c0e-c1cb-4520-931b-abd4fce33cbc.png)
![image](https://user-images.githubusercontent.com/24818503/218893774-783e5f70-eede-4ba0-bd2f-1c58950b5f09.png)

DeepSpeed ZeRO2
![image](https://user-images.githubusercontent.com/24818503/218893800-755d9d18-bb74-454f-8d01-9eb3058a111b.png)
![image](https://user-images.githubusercontent.com/24818503/218893811-6da63415-f1df-4199-b0b2-af889ffb9046.png)

## Code example


## System Info

PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1085-azure-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.3.109
CUDA_MODULE_LOADING set to: 
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB
GPU 6: Tesla V100-SXM2-32GB
GPU 7: Tesla V100-SXM2-32GB

Nvidia driver version: 510.73.08
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          40
On-line CPU(s) list:             0-39
Thread(s) per core:              1
Core(s) per socket:              20
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel(R) Xeon(R) Platinum 8168 CPU @ 2.70GHz
Stepping:                        4
CPU MHz:                         2693.672
BogoMIPS:                        5387.34
Hypervisor vendor:               Microsoft
Virtualization type:             full
L1d cache:                       1.3 MiB
L1i cache:                       1.3 MiB
L2 cache:                        40 MiB
L3 cache:                        66 MiB
NUMA node0 CPU(s):               0-19
NUMA node1 CPU(s):               20-39
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Mitigation; PTE Inversion
Vulnerability Mds:               Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:          Mitigation; PTI
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, STIBP disabled, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Mitigation; Clear CPU buffers; SMT Host state unknown
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti fsgsbase bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves md_clear arch_capabilities

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.4
[pip3] pytorch-lightning==1.6.3
[pip3] torch==1.12.1
[pip3] torch-nebula==0.15.1
[pip3] torch-ort==1.12.0
[pip3] torch-tb-profiler==0.4.0
[pip3] torchmetrics==0.7.1
[pip3] torchvision==0.13.1+cu113
[conda] magma-cuda113             2.5.2                         1    pytorch
[conda] mkl                       2023.0.0                 pypi_0    pypi
[conda] mkl-fft                   1.3.1                    pypi_0    pypi
[conda] mkl-include               2021.4.0                 pypi_0    pypi
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] pytorch-lightning         1.6.3                    pypi_0    pypi
[conda] torch                     1.12.1                   pypi_0    pypi
[conda] torch-nebula              0.15.1                   pypi_0    pypi
[conda] torch-ort                 1.12.0                   pypi_0    pypi
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchmetrics              0.7.1                    pypi_0    pypi
[conda] torchvision               0.13.1+cu113             pypi_0    pypi

cc @csarofeen @ptrblck @xwang233 @ngimel