Repository: pytorch
Issue: Unexpected Gradient Results From conv1d on MacOS
Labels: module: convolution,triaged,module: correctness (silent),module: mps,module: functorch

### 🐛 Describe the bug

Obtained unexpected gradient values from conv1d for large tensors (on both CPU and MPS):

```python
import torch
def fn(kernel, input):
    return torch.mean(torch.nn.functional.conv1d(input, kernel, None, 1, 0, 1))

gradient_fn = torch.func.grad(fn)

def calculate_fn_and_grad(N, device = "cpu"):
    input = torch.ones(N, 3, 100, device = device, dtype=torch.float32)
    kernel = torch.ones(2, 3, 1, device = device, dtype=torch.float32)

    print(f"Output with {device}:", fn(kernel, input))
    print(f"Kernel grad with {device}:", gradient_fn(kernel, input))

N=1000
print(f"\nCompare results with N={N}:")
calculate_fn_and_grad(N, "cpu")
calculate_fn_and_grad(N, "mps")

N=100000
print(f"\nCompare results with N={N}:")
calculate_fn_and_grad(N, "cpu")
calculate_fn_and_grad(N, "mps")

N=1000000
print(f"\nCompare esults with N={N}:")
calculate_fn_and_grad(N, "cpu")
calculate_fn_and_grad(N, "mps")
```

Results:

```
Compare results with N=1000:
Output with cpu: tensor(3.)
Kernel grad with cpu: tensor([[[0.5000],
         [0.5000],
         [0.5000]],

        [[0.5000],
         [0.5000],
         [0.5000]]])
Output with mps: tensor(3., device='mps:0')
Kernel grad with mps: tensor([[[0.5000],
         [0.5000],
         [0.5000]],

        [[0.5000],
         [0.5000],
         [0.5000]]], device='mps:0')

Compare results with N=100000:
Output with cpu: tensor(3.)
Kernel grad with cpu: tensor([[[0.5005],
         [0.5005],
         [0.5005]],

        [[0.5005],
         [0.5005],
         [0.5005]]])
Output with mps: tensor(3., device='mps:0')
Kernel grad with mps: tensor([[[0.4912],
         [0.4912],
         [0.4912]],

        [[0.4912],
         [0.4912],
         [0.4912]]], device='mps:0')

Compare results with N=1000000:
Output with cpu: tensor(3.)
Kernel grad with cpu: tensor([[[0.5045],
         [0.5045],
         [0.5045]],

        [[0.5045],
         [0.5045],
         [0.5045]]])
Output with mps: tensor(3., device='mps:0')
Kernel grad with mps: tensor([[[0.2597],
         [0.2597],
         [0.2597]],

        [[0.2597],
         [0.2597],
         [0.2597]]], device='mps:0')
```

Similar behavior is observed when using conv1d without mean (gradient calcuation with vjp)

```
import torch
def fn(kernel, input):
    return torch.nn.functional.conv1d(input, kernel, None, 1, 0, 1)

gradient_fn = torch.func.grad(fn)

def calculate_fn_and_grad(N, device = "cpu"):
    input = torch.ones(N, 3, 100, device = device, dtype=torch.float32)
    kernel = torch.ones(2, 3, 1, device = device, dtype=torch.float32)
    (_, vjpfunc) = torch.func.vjp(fn, kernel, input)
    kernel_gradients, _ = vjpfunc(torch.ones_like(fn(kernel, input)))
    print(f"Kernel grad with {device}:", kernel_gradients)

N=1000
print(f"\nCompare results with N={N}:")
calculate_fn_and_grad(N, "cpu")
calculate_fn_and_grad(N, "mps")

N=100000
print(f"\nCompare results with N={N}:")
calculate_fn_and_grad(N, "cpu")
calculate_fn_and_grad(N, "mps")

N=1000000
print(f"\nCompare results with N={N}:")
calculate_fn_and_grad(N, "cpu")
calculate_fn_and_grad(N, "mps")
```

Results:
```
Compare results with N=1000:
Kernel grad with cpu: tensor([[[100000.],
         [100000.],
         [100000.]],

        [[100000.],
         [100000.],
         [100000.]]])
Kernel grad with mps: tensor([[[100000.],
         [100000.],
         [100000.]],

        [[100000.],
         [100000.],
         [100000.]]], device='mps:0')

Compare results with N=100000:
Kernel grad with cpu: tensor([[[10000000.],
         [10000000.],
         [10000000.]],

        [[10000000.],
         [10000000.],
         [10000000.]]])
Kernel grad with mps: tensor([[[9830400.],
         [9830400.],
         [9830400.]],

        [[9830400.],
         [9830400.],
         [9830400.]]], device='mps:0')

Compare results with N=1000000:
Kernel grad with cpu: tensor([[[98684352.],
         [98684352.],
         [98684352.]],

        [[98684352.],
         [98684352.],
         [98684352.]]])
Kernel grad with mps: tensor([[[52428800.],
         [52428800.],
         [52428800.]],

        [[52428800.],
         [52428800.],
         [52428800.]]], device='mps:0')
```

Same experiment works perfectly fine (for both CPU and CUDA) when running on Ubuntu.

### Versions

PyTorch version: 2.2.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.3.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.5 (main, Sep 11 2023, 08:31:25) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.3.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2 Pro

Versions of relevant libraries:
[pip3] numpy==1.26.3
[pip3] torch==2.2.0
[conda] numpy                     1.26.3          py311he598dae_0  
[conda] numpy-base                1.26.3          py311hfbfe69c_0  
[conda] torch                     2.2.0                    pypi_0    pypi


cc @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen @zou3519 @Chillee @samdow @kshitij12345 @razarmehr @janeyx99