I Contributed to PyTorch. Here's What I Learned

The Issue Must Not Be That Bad

That's what I thought when I encountered a PyTorch problem during one of my college assignments. Jupyter kernel was dying because of some bug in the LSTM implementation for MPS.

After a quick investigation, I discovered that this happens because of the batch_first flag. MPS's backend did not work correctly with it and crushed the entire kernel.

"Easy fix"
P.S. After that phrase, Alex spend the next two days fixing what looked like an "easy fix"

PR was merged pretty quickly. Thanks, PyTorch team, for that! And the story could've ended here, but I discovered a funny detail in MPS tests.

@unittest.skipIf(True, "Backward of lstm returns wrong result")
def test_lstm_2(self, device="mps", dtype=torch.float32)

And LSTM was really bad. It got a whole lot worse score than when trained on CUDA or CPU.

It Was Bad. Really Bad

It turned out that LSTM on MPS was completely broken. The forward pass had a bug with the batch_first flag and hidden cell initialization.

Backward pass used first layers weights for the last layers, mixing up all gradients. It did not calculate gradients for hidden states. And my favourite: the backward function returned initialized with garbage tensors, screwing up all subsequent training. It was a mess that I kept investigating for several days.

Eventually, I fixed LSTM and its tests in a massive PR, ensuring that it is consistent with the CPU.

What I Learned

• Big projects also have garbage code. Broken implementation lived in stable releases for almost a year, generating several related GitHub issues.
• Contributing to a big project is fun and challenging. And it eventually helps a lot of developers, which keeps me warm during cold winter nights. Specifically, contributing to PyTorch is extremely simple. Thanks, PyTorch team, for arranging that!
• Deploying untested code that looks right is extremely dangerous. I listed pretty severe mistakes that I found scrutinizing LSTM sources for several days. There's no way they could have been discovered without extensive testing. Even though the issues were severe, they were also subtle. The code looked right.

Finally

I was able to complete the college PyTorch assignment even though it required rewriting PyTorch's LSTM MPS implementation. Consider also solving open issues of your favourite framework or project. At the end of the day, it is a lot more fun than Leetcode problems.

Appendix: Common PyTorch MPS Errors

Since publishing this story, this post turns up in searches for specific MPS
errors — so here is a quick reference for the ones related to my PRs.

AttributeError: module 'torch.backends' has no attribute 'mps'

Your PyTorch build predates the MPS backend, which landed in PyTorch 1.12.
Upgrade with:

pip install --upgrade torch torchvision torchaudio

Cannot execute empty_cache() without MPS backend

You are calling torch.mps.empty_cache() on a machine (or build) without MPS.
Guard the call:

if torch.backends.mps.is_available():
    torch.mps.empty_cache()

Note the asymmetry: torch.backends.mps.is_available() checks whether your
PyTorch build and hardware support MPS, while torch.mps.* functions assume
it unconditionally.

LSTM crashes or silently wrong results on MPS (batch_first=True)

This is exactly the bug from my story. Symptoms included the training process
dying outright, LSTM scoring dramatically worse than on CUDA/CPU, gradients
computed with wrong layer transpositions, missing hidden-state gradients, and
backward passes returning garbage-initialized tensors.

All of it was fixed across early-2023 pull requests:

• #95137 — LSTM backward and forward pass fixes
• #95563 — bidirectional & one-direction LSTM fixes
• #96601 — grad_y fix
• #95091 — LogSoftmax numerical stability on MPS

If you hit any of these symptoms, upgrade past these fixes before suspecting
your own code — and if the results still differ between devices, that itself
is worth a bug report.

See My Work

[MPS] Fix LSTM backward and forward pass by AlexRoar · Pull Request #95137 · pytorch/pytorch

Fixes #91694 Fixes #92615 Several transpositions were missing for backward graph in case of batch_first=True. The #91694 is not reproduced with batch_first=False. After fixing transpose issue, I fi...

GitHubpytorch

[MPS] Fix bidirectional LSTM & small one-direction LSTM fix by AlexRoar · Pull Request #95563 · pytorch/pytorch

Fixes #94754 With this PR I hope to finish my breathtaking journey of fixing MPS LSTM. Here, I enable bidirectional on MPS. Also, I’ve noticed that cache key did not account for all parameters, so ...

GitHubpytorch

[MPS] LSTM grad_y missing fix by AlexRoar · Pull Request #96601 · pytorch/pytorch

Fixes #96416 Added tests that do not use LSTM output simalarly to the issue Seems like this fix once again introduces backward incompatibility.

GitHubpytorch

[MPS] LogSoftmax numerical stability by AlexRoar · Pull Request #95091 · pytorch/pytorch

Fixes #94043 Calculations are now consistent with numericaly stable formula and CPU: $LogSoftmax(X, \dim) = X - \max(X, \dim) - \log(sum(X - \max(X, \dim), \dim))$ @malfet

GitHubpytorch

WRITTEN BY

Alex Dremov

I'm an AI researcher and ML engineer. I hold an M.Sc. in Data Science from EPFL and a B.Sc. in CS from MIPT. Over the last couple of years, I've built real-time speech recognition systems from the ground up at Yandex and researched compute-optimal quantization at Apple. My academic work on efficient training dynamics and scaling laws has been published at ICLR and TMLR. When I'm not doing research or contributing to open-source projects, I write here about whatever I'm currently exploring—from optimizing GPU kernels with Triton to the scaling laws behind efficient training.

View all articles by Alex Dremov →