pytorch - 💡(How to fix) Fix [MPS] Non-deterministic backward pass for F.linear [1 participants]

🐛 Describe the bug

The backward pass for F.linear(x, w) produces different x.grad across consecutive calls.

This is a follow-up to #180776, which was determined to be specific to the M5. Testing the latest nightly build containing @malfet's fix from #181466, I noticed that the forward pass seems fine, but I am now seeing the same non-deterministic behavior for the backward pass. My assumption is that this occurs under the same conditions as the original issue (BFloat16/Float16, >2D input, M5).

The fix from #181466 should probably be applied to _mps_linear_backward_input to resolve this.

MRE

import torch
import torch.nn.functional as F

x = torch.randn(2, 13, 1024, device="mps", dtype=torch.bfloat16)
w = torch.randn(1024, 1024, device="mps", dtype=torch.bfloat16)
grad = torch.randn(2, 13, 1024, device="mps", dtype=torch.bfloat16)

def backward_test():
    x0, w0 = x.clone().requires_grad_(), w.clone()
    y = F.linear(x0, w0)
    y.backward(grad)
    return x0.grad.clone()

first = backward_test()
second = backward_test()
print("Diff between runs:", (second - first).abs().max().item())

Output: Diff between runs: 130.0

Versions

PyTorch version: 2.13.0.dev20260429 Is debug build: False CUDA used to build PyTorch: None ROCM used to build PyTorch: N/A

OS: macOS 26.4.1 (arm64) GCC version: Could not collect Clang version: 21.0.0 (clang-2100.0.123.102) CMake version: Could not collect Libc version: N/A

Python version: 3.14.3 (v3.14.3:323c59a5e34, Feb 3 2026, 11:41:37) [Clang 16.0.0 (clang-1600.0.26.6)] (64-bit runtime) Python platform: macOS-26.4.1-arm64-arm-64bit-Mach-O 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 Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

CPU: Apple M5

Versions of relevant libraries: [pip3] numpy==2.4.4 [pip3] torch==2.13.0.dev20260429 [pip3] torchvision==0.27.0.dev20260429 [conda] Could not collect

cc @ezyang @albanD @gqchen @nikitaved @soulitzer @Varal7 @bobrenjc93 @mruberry @kurtamohler @snadampal @milpuz01 @aditew01 @nikhil-arm @fadara01 @robert-hardwick @nWEIdia @kulinseth @malfet @DenisVieriu97 @jhavukainen @aditvenk